[diffusion] refactor: rename quantized model path server arg (#19142)

python/sglang/multimodal_gen/configs/quantization.py

@@ -28,23 +28,25 @@ class NunchakuSVDQuantArgs:
     """
 
     enable_svdquant: bool = False
-    quantized_model_path: str | None = None
+    transformer_weights_path: str | None = None
     quantization_precision: str | None = None  # "int4" or "nvfp4"
     quantization_rank: int | None = None
     quantization_act_unsigned: bool = False
 
     def _adjust_config(self) -> None:
         """infer precision and rank from filename if not provided"""
-        if self.quantized_model_path and not self.enable_svdquant:
+        if self.transformer_weights_path and not self.enable_svdquant:
+            filename = os.path.basename(self.transformer_weights_path)
+            if re.search(r"svdq-(int4|fp4)_r(\d+)", filename):
                 self.enable_svdquant = True
 
-        if not self.enable_svdquant or not self.quantized_model_path:
+        if not self.enable_svdquant or not self.transformer_weights_path:
             return
 
         inferred_precision = None
         inferred_rank = None
 
-        filename = os.path.basename(self.quantized_model_path)
+        filename = os.path.basename(self.transformer_weights_path)
         # Expected pattern: svdq-{precision}_r{rank}-...
         # e.g., svdq-int4_r32-qwen-image.safetensors
         match = re.search(r"svdq-(int4|fp4)_r(\d+)", filename)
@@ -59,7 +61,7 @@ class NunchakuSVDQuantArgs:
             if inferred_precision:
                 logger.info(
                     f"inferred --quantization-precision: {self.quantization_precision} "
-                    f"from --quantized-model-path: {self.quantized_model_path}"
+                    f"from --transformer-weights-path: {self.transformer_weights_path}"
                 )
 
         if self.quantization_rank is None:
@@ -67,7 +69,7 @@ class NunchakuSVDQuantArgs:
             if inferred_rank:
                 logger.info(
                     f"inferred --quantization-rank: {self.quantization_rank} "
-                    f"from --quantized-model-path: {self.quantized_model_path}"
+                    f"from --transformer-weights-path: {self.transformer_weights_path}"
                 )
 
     def validate(self) -> None:
@@ -101,9 +103,9 @@ class NunchakuSVDQuantArgs:
                 "Disable it with --enable-svdquant false."
             )
 
-        if not self.quantized_model_path:
+        if not self.transformer_weights_path:
             raise ValueError(
-                "--enable-svdquant requires --quantized-model-path to be set"
+                "--enable-svdquant requires --transformer-weights-path to be set"
             )
 
         if not is_nunchaku_available():
@@ -131,12 +133,12 @@ class NunchakuSVDQuantArgs:
             help="Enable Nunchaku SVDQuant (W4A4-style) inference.",
         )
         parser.add_argument(
-            "--quantized-model-path",
+            "--transformer-weights-path",
             type=str,
-            default=NunchakuSVDQuantArgs.quantized_model_path,
+            default=NunchakuSVDQuantArgs.transformer_weights_path,
             help=(
-                "Path to pre-quantized Nunchaku weights. Can be a single .safetensors "
+                "Path to pre-quantized transformer weights. Can be a single .safetensors "
-                "file or a directory containing .safetensors."
+                "file, a directory, or a HuggingFace repo ID. Used by Nunchaku (SVDQuant) and quantized single-file checkpoints."
             ),
         )
         parser.add_argument(
@@ -161,11 +163,14 @@ class NunchakuSVDQuantArgs:
     @classmethod
     def from_dict(cls, kwargs: dict[str, Any]) -> "NunchakuSVDQuantArgs":
         # Map CLI/config keys to dataclass fields (keep backwards compatibility).
+        path = (
+            kwargs.get("transformer_weights_path")
+            or kwargs.get("transformer_quantized_path")
+            or kwargs.get("quantized_model_path")
+        )
         return cls(
             enable_svdquant=bool(kwargs.get("enable_svdquant", cls.enable_svdquant)),
-            quantized_model_path=kwargs.get(
+            transformer_weights_path=path,
-                "quantized_model_path", cls.quantized_model_path
-            ),
             quantization_precision=kwargs.get("quantization_precision"),
             quantization_rank=kwargs.get("quantization_rank"),
             quantization_act_unsigned=bool(

python/sglang/multimodal_gen/docs/quantization.md

@@ -70,7 +70,7 @@ sglang generate \
   --model-path Qwen/Qwen-Image \
   --prompt "change the raccoon to a cute cat" \
   --save-output \
-  --quantized-model-path /path/to/svdq-int4_r32-qwen-image.safetensors
+  --transformer-weights-path /path/to/svdq-int4_r32-qwen-image.safetensors
 ```
 
 **Manual Override (If needed):**
@@ -89,7 +89,7 @@ sglang generate \
   --model-path Qwen/Qwen-Image \
   --prompt "a beautiful sunset" \
   --enable-svdquant \
-  --quantized-model-path /path/to/custom_model.safetensors \
+  --transformer-weights-path /path/to/custom_model.safetensors \
   --quantization-precision int4 \
   --quantization-rank 128
 ```
@@ -108,7 +108,7 @@ Choose the appropriate configuration based on your hardware and requirements:
 
 ### Notes
 
-1.  Model Path Correspondence: `--model-path` should point to the original non-quantized model (for loading config and tokenizer, etc.), while `--quantized-model-path` points to the quantized weight file.
+1.  Model Path Correspondence: `--model-path` should point to the original non-quantized model (for loading config and tokenizer, etc.), while `--transformer-weights-path` points to the quantized weight file / folder / Huggingface Repo ID.
 
 2.  Auto-Detection Requirements: For auto-detection to work, the filename must contain the pattern `svdq-{precision}_r{rank}` (e.g., `svdq-int4_r32`).
 
@@ -122,26 +122,24 @@ Choose the appropriate configuration based on your hardware and requirements:
 
 If you want to quantize your own models, you can use the [DeepCompressor](https://github.com/mit-han-lab/deepcompressor) tool. For detailed instructions, please refer to the Nunchaku official documentation.
 
-## FP8 Quantization
+## Quantization
 
 ### Usage
 
-#### Option 1: Use Pre-quantized Models (Recommended)
+#### Option 1: Pre-quantized folder (has `config.json`)
 
-If available, you can directly use pre-quantized FP8 models from Hugging Face or other sources. Simply load them with SGLang:
+For quantized checkpoints that include a `config.json` with a `quantization_config` field (e.g., models converted via `convert_hf_to_fp8.py`), where the transformer's `config.json` already encodes the `quantization_config`, use the component override:
 
 ```bash
 sglang generate \
-  --model-path /path/to/FLUX.1-dev-FP8/ \
+  --model-path /path/to/FLUX.1-dev \
+  --transformer-path /path/to/FLUX.1-dev/transformer-FP8 \
   --prompt "A Logo With Bold Large Text: SGL Diffusion" \
   --save-output
 ```
 
-#### Option 2: Convert Your Own Models
 
-If you need to convert a model to FP8 format, use the provided conversion script:
+If you need to convert a model to FP8 format yourself, use the provided conversion script:
-
-**Step 1: Convert the Model**
 
 ```bash
 # convert transformer to FP8 with block quantization
@@ -152,13 +150,20 @@ python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \
   --block-size 128 128
 ```
 
-**Step 2: Run Inference**
+#### Option 2: Pre-quantized single-file checkpoint (no `config.json`)
+
+
+
+Some providers (e.g., [black-forest-labs/FLUX.2-klein-9b-fp8](https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8)) distribute a single `.safetensors` file without a companion `config.json`. Use `--transformer-weights-path` to point to this file (or HuggingFace repo ID) while keeping `--model-path` for the base model:
 
 ```bash
 sglang generate \
-  --model-path /path/to/FLUX.1-dev/
+  --model-path black-forest-labs/FLUX.2-klein-9B \
-  # override transformer component with path to converted model
+  --transformer-weights-path black-forest-labs/FLUX.2-klein-9b-fp8 \
-  --transformer-path /path/to/FLUX.1-dev/transformer-FP8
   --prompt "A Logo With Bold Large Text: SGL Diffusion" \
   --save-output
 ```
+
+SGLang-Diffusion will automatically read the `quantization_config` metadata embedded in the safetensors file header (if present). For the quant config to be auto-detected, the file's metadata must contain a JSON-encoded `quantization_config` key with at least a `quant_method` field (e.g. `"fp8"`).
+
+Note: this feature is a WIP

python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py

@@ -40,7 +40,7 @@ class NunchakuConfig(QuantizationConfig):
         rank: SVD low-rank dimension for absorbing outliers
         group_size: Quantization group size (automatically set based on precision)
         act_unsigned: Use unsigned activation quantization
-        quantized_model_path: Path to pre-quantized model weights (.safetensors)
+        transformer_weights_path: Path to pre-quantized transformer weights (.safetensors)
         model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules()
     """
 
@@ -48,7 +48,7 @@ class NunchakuConfig(QuantizationConfig):
     rank: int = 32
     group_size: Optional[int] = None
     act_unsigned: bool = False
-    quantized_model_path: Optional[str] = None
+    transformer_weights_path: Optional[str] = None
     model_cls: Optional[type] = None
 
     @classmethod
@@ -75,7 +75,7 @@ class NunchakuConfig(QuantizationConfig):
             rank=int(config.get("rank", 32)),
             group_size=config.get("group_size"),
             act_unsigned=bool(config.get("act_unsigned", False)),
-            quantized_model_path=config.get("quantized_model_path"),
+            transformer_weights_path=config.get("transformer_weights_path"),
         )
 
     def get_quant_method(
@@ -158,7 +158,7 @@ class NunchakuConfig(QuantizationConfig):
             "rank": self.rank,
             "group_size": self.group_size,
             "act_unsigned": self.act_unsigned,
-            "quantized_model_path": self.quantized_model_path,
+            "transformer_weights_path": self.transformer_weights_path,
         }
 
     @classmethod

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -2,13 +2,13 @@ import inspect
 import json
 import logging
 import os
-from copy import deepcopy
+from typing import Any, Dict, List, Optional
-from typing import Any
 
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    NunchakuConfig,
     _patch_nunchaku_scales,
 )
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
@@ -25,6 +25,8 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
     get_metadata_from_safetensors_file,
     get_quant_config,
+    get_quant_config_from_safetensors_metadata,
+    maybe_download_model,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
@@ -44,35 +46,105 @@ class TransformerLoader(ComponentLoader):
         """
         get list of safetensors to load.
 
-        For some quantization framework, if --quantized-model-path is provided, load from this path instead of main model
+        If --transformer-weights-path is provided, load weights from that path
+        instead of the base model's component directory.
         """
-        nunchaku_config = server_args.nunchaku_config
+        quantized_path = server_args.transformer_weights_path
 
-        if nunchaku_config is not None and nunchaku_config.quantized_model_path:
+        if quantized_path:
-            # load from quantized_model_path if applicable
+            quantized_path = maybe_download_model(quantized_path)
-            weights_path = nunchaku_config.quantized_model_path
+            logger.info("using quantized transformer weights from: %s", quantized_path)
-
+            if os.path.isfile(quantized_path) and quantized_path.endswith(
-            logger.info("Using quantized model weights from: %s", weights_path)
+                ".safetensors"
-            if os.path.isfile(weights_path) and weights_path.endswith(".safetensors"):
+            ):
-                safetensors_list = [weights_path]
+                safetensors_list = [quantized_path]
             else:
-                safetensors_list = _list_safetensors_files(weights_path)
+                safetensors_list = _list_safetensors_files(quantized_path)
         else:
-            weights_path = component_model_path
+            safetensors_list = _list_safetensors_files(component_model_path)
-            safetensors_list = _list_safetensors_files(weights_path)
 
         if not safetensors_list:
-            raise ValueError(f"No safetensors files found in {weights_path}")
+            raise ValueError(
+                f"no safetensors files found in "
+                f"{quantized_path or component_model_path}"
+            )
 
         return safetensors_list
 
+    def _resolve_quant_config(
+        self,
+        hf_config: Dict[str, List[str]],
+        server_args: ServerArgs,
+        safetensors_list: list[str],
+    ) -> Optional[dict]:
+        # priority: model config.json → safetensors metadata → nunchaku config
+        quant_config = get_quant_config(hf_config)
+        if quant_config is None and server_args.transformer_weights_path:
+            # try to read quantization_config from the safetensors metadata header
+            for safetensors_file in safetensors_list:
+                quant_config = get_quant_config_from_safetensors_metadata(
+                    safetensors_file
+                )
+                if quant_config:
+                    break
+        return quant_config
+
+    def _resolve_target_param_dtype(
+        self,
+        quant_config: Optional[dict],
+        nunchaku_config: Optional[NunchakuConfig],
+        model_cls,
+        server_args: ServerArgs,
+    ) -> Optional[torch.dtype]:
+        if quant_config is not None or nunchaku_config is not None:
+            # TODO: improve the condition
+            # respect dtype from checkpoint
+            param_dtype = None
+        else:
+            param_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        if nunchaku_config is not None:
+            nunchaku_config.model_cls = model_cls
+            # verify that the nunchaku checkpoint matches the selected model class
+            original_dit_cls_name = json.loads(
+                get_metadata_from_safetensors_file(
+                    nunchaku_config.transformer_weights_path
+                ).get("config")
+            )["_class_name"]
+            specified_dit_cls_name = str(model_cls.__name__)
+            if original_dit_cls_name != specified_dit_cls_name:
+                raise Exception(
+                    f"Class name of DiT specified in nunchaku transformer_weights_path: {original_dit_cls_name} does not match that of specified DiT name: {specified_dit_cls_name}"
+                )
+
+        return param_dtype
+
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
         """Load the transformer based on the model path, and inference args."""
+        # 1. hf config
         config = get_diffusers_component_config(component_path=component_model_path)
 
-        hf_config = deepcopy(config)
+        # 2. quant config
+        safetensors_list = self.get_list_of_safetensors_to_load(
+            server_args, component_model_path
+        )
+        quant_config = self._resolve_quant_config(config, server_args, safetensors_list)
+
+        # 3. dit config
+        # Config from Diffusers supersedes sgl_diffusion's model config
+        component_name = _normalize_component_type(component_name)
+        server_args.model_paths[component_name] = component_model_path
+        if component_name in ("transformer", "video_dit"):
+            pipeline_dit_config_attr = "dit_config"
+        elif component_name in ("audio_dit",):
+            pipeline_dit_config_attr = "audio_dit_config"
+        else:
+            raise ValueError(f"Invalid module name: {component_name}")
+        dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr)
+        dit_config.update_model_arch(config)
+
         cls_name = config.pop("_class_name")
         if cls_name is None:
             raise ValueError(
@@ -80,46 +152,11 @@ class TransformerLoader(ComponentLoader):
                 "Only diffusers format is supported."
             )
 
-        component_name = _normalize_component_type(component_name)
-        server_args.model_paths[component_name] = component_model_path
-
-        if component_name in ("transformer", "video_dit"):
-            pipeline_dit_config_attr = "dit_config"
-        elif component_name in ("audio_dit",):
-            pipeline_dit_config_attr = "audio_dit_config"
-        else:
-            raise ValueError(f"Invalid module name: {component_name}")
-        # Config from Diffusers supersedes sgl_diffusion's model config
-        dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr)
-        dit_config.update_model_arch(config)
-
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         nunchaku_config = server_args.nunchaku_config
-
+        param_dtype = self._resolve_target_param_dtype(
-        if nunchaku_config is not None:
+            quant_config, nunchaku_config, model_cls, server_args
-            nunchaku_config.model_cls = model_cls
-
-            # respect dtype from checkpoint
-            # TODO: improve the condition
-            param_dtype = None
-
-            # check if the specified nunchaku quantized model path matches with the specified model path
-            original_dit_cls_name = json.loads(
-                get_metadata_from_safetensors_file(
-                    nunchaku_config.quantized_model_path
-                ).get("config")
-            )["_class_name"]
-            specified_dit_cls_name = str(model_cls.__name__)
-            if original_dit_cls_name != specified_dit_cls_name:
-                raise Exception(
-                    f"Class name of DiT specified in nunchaku quantized model_path: {original_dit_cls_name} does not match that of specified DiT name: {specified_dit_cls_name}"
-                )
-        else:
-            param_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
-
-        safetensors_list = self.get_list_of_safetensors_to_load(
-            server_args, component_model_path
         )
 
         logger.info(
@@ -130,13 +167,20 @@ class TransformerLoader(ComponentLoader):
             param_dtype,
         )
 
-        init_params: dict[str, Any] = {"config": dit_config, "hf_config": hf_config}
+        init_params: dict[str, Any] = {"config": dit_config, "hf_config": config}
-
+        # prepare init_param
         if "quant_config" in inspect.signature(model_cls.__init__).parameters:
-            quant_config = get_quant_config(config)
+            init_params.update(
-            init_params["quant_config"] = (
+                {
-                quant_config if quant_config else nunchaku_config
+                    "quant_config": (quant_config if quant_config else nunchaku_config),
+                }
             )
+            if init_params["quant_config"] is None:
+                logger.warning(
+                    f"transformer_weights_path provided, but quantization config not resolved, which is unexpected and likely to cause errors"
+                )
+            else:
+                logger.debug("quantization config: %s", init_params["quant_config"])
 
         # Load the model using FSDP loader
         model = maybe_load_fsdp_model(

python/sglang/multimodal_gen/runtime/server_args.py

@@ -289,6 +289,9 @@ class ServerArgs:
 
     # Component path overrides (key = model_index.json component name, value = path)
     component_paths: dict[str, str] = field(default_factory=dict)
+
+    # path to pre-quantized transformer weights (single .safetensors or directory).
+    transformer_weights_path: str | None = None
     # can restrict layers to adapt, e.g. ["q_proj"]
     # Will adapt only q, k, v, o by default.
     lora_target_modules: list[str] | None = None
@@ -400,15 +403,19 @@ class ServerArgs:
         if ncfg is None or isinstance(ncfg, NunchakuConfig):
             return
         ncfg.validate()
-        if not ncfg.enable_svdquant or not ncfg.quantized_model_path:
+
-            # if nunchaku is not applied
+        # propagate the path to server_args
+        if ncfg.transformer_weights_path:
+            self.transformer_weights_path = ncfg.transformer_weights_path
+
+        if not ncfg.enable_svdquant or not ncfg.transformer_weights_path:
             self.nunchaku_config = None
         else:
             self.nunchaku_config = NunchakuConfig(
                 precision=self.nunchaku_config.quantization_precision,
                 rank=self.nunchaku_config.quantization_rank,
                 act_unsigned=self.nunchaku_config.quantization_act_unsigned,
-                quantized_model_path=self.nunchaku_config.quantized_model_path,
+                transformer_weights_path=self.nunchaku_config.transformer_weights_path,
             )
 
     def _adjust_offload(self):

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -913,48 +913,42 @@ def get_quant_config_from_safetensors_metadata(
     file_path: str,
 ) -> Optional[QuantizationConfig]:
     """Extract quantization config from a safetensors file's metadata header.
-
-    Safetensors files can embed a flat string→string metadata dict in their header.
-    We expect a ``quantization_config`` key containing a JSON-encoded dict with at
-    least a ``quant_method`` field (e.g. ``"fp8"``), matching the format written by
-    ``convert_hf_to_fp8.py`` when embedded into a config.json.
-
     Returns None if no recognizable quantization metadata is found.
     """
     metadata = get_metadata_from_safetensors_file(file_path)
     if not metadata:
         return None
 
-    quant_config_str = metadata.get("quantization_config")
+    quant_config_str = metadata.get("_quantization_metadata")
     if not quant_config_str:
         return None
-
     try:
         quant_config_dict = json.loads(quant_config_str)
-    except Exception as e:
+    except Exception as _e:
-        logger.warning(
-            "failed to parse quantization_config from safetensors metadata: %s", e
-        )
         return None
 
+    # handle diffusers fp8 safetensors metadata format
+    if (
+        "quant_method" not in quant_config_dict
+        and "format_version" in quant_config_dict
+        and "layers" in quant_config_dict
+    ):
+        layers = quant_config_dict.get("layers", {})
+        if any(
+            isinstance(v, dict) and "float8" in v.get("format", "")
+            for v in layers.values()
+        ):
+            quant_config_dict["quant_method"] = "fp8"
+            quant_config_dict["activation_scheme"] = "dynamic"
+
     quant_method = quant_config_dict.get("quant_method")
     if not quant_method:
-        logger.warning(
-            "quantization_config in safetensors metadata is missing 'quant_method'"
-        )
         return None
 
     try:
         quant_cls = get_quantization_config(quant_method)
         config = quant_cls.from_config(quant_config_dict)
-        logger.info(
+        logger.debug(f"Get quantization config from safetensors file: {file_path}")
-            "loaded quantization config (%s) from safetensors metadata: %s",
-            quant_method,
-            file_path,
-        )
         return config
-    except Exception as e:
+    except Exception as _e:
-        logger.warning(
-            "failed to build QuantizationConfig from safetensors metadata: %s", e
-        )
         return None

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -368,6 +368,16 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
         ),
         T2I_sampling_params,
     ),
+    # TODO: modeling of flux different from official flux, so weights can't be loaded
+    # consider opting for a different quantized hf-repo
+    # DiffusionTestCase(
+    #     "flux_image_t2i_override_transformer_weights_path_fp8",
+    #     DiffusionServerArgs(
+    #         model_path="black-forest-labs/FLUX.1-dev", modality="image",
+    #         extras=["--transformer-weights-path black-forest-labs/FLUX.1-dev-FP8"]
+    #     ),
+    #     T2I_sampling_params,
+    # ),
     DiffusionTestCase(
         "flux_2_image_t2i",
         DiffusionServerArgs(