[diffusion] feat: support passing component path via server args (#19108)

docs/diffusion/api/cli.md

@@ -12,7 +12,6 @@ The SGLang-diffusion CLI provides a quick way to access the inference pipeline f
 ### Server Arguments
 
 - `--model-path {MODEL_PATH}`: Path to the model or model ID
-- `--vae-path {VAE_PATH}`: Path to a custom VAE model or HuggingFace model ID (e.g., `fal/FLUX.2-Tiny-AutoEncoder`). If not specified, the VAE will be loaded from the main model path.
 - `--lora-path {LORA_PATH}`: Path to a LoRA adapter (local path or HuggingFace model ID). If not specified, LoRA will not be applied.
 - `--lora-nickname {NAME}`: Nickname for the LoRA adapter. (default: `default`).
 - `--num-gpus {NUM_GPUS}`: Number of GPUs to use
@@ -218,6 +217,32 @@ Once the generation task has finished, the server will shut down automatically.
 > [!NOTE]
 > The HTTP server-related arguments are ignored in this subcommand.
 
+## Component Path Overrides
+
+SGLang diffusion allows you to override any pipeline component (e.g., `vae`, `transformer`, `text_encoder`) by specifying a custom checkpoint path. This is useful for:
+
+### Example: FLUX.2-dev with Tiny AutoEncoder
+
+You can override **any** component by using `--<component>-path`, where `<component>` matches the key in the model's `model_index.json`:
+
+For example, replace the default VAE with a distilled tiny autoencoder for ~3x faster decoding:
+
+```bash
+sglang serve \
+  --model-path=black-forest-labs/FLUX.2-dev \
+  # with a Huggingface Repo ID
+  --vae-path=fal/FLUX.2-Tiny-AutoEncoder
+  # or use a local path
+  --vae-path=~/.cache/huggingface/hub/models--fal--FLUX.2-Tiny-AutoEncoder/snapshots/.../vae
+```
+
+**Important:**
+- The component key must match the one in your model's `model_index.json` (e.g., `vae`).
+- The path must:
+    - either be a Huggingface Repo ID (e.g., fal/FLUX.2-Tiny-AutoEncoder)
+    - or point to a **complete component folder**, containing `config.json` and safetensors files
+
+
 ## Diffusers Backend
 
 SGLang diffusion supports a **diffusers backend** that allows you to run any diffusers-compatible model through SGLang's infrastructure using vanilla diffusers pipelines. This is useful for running models without native SGLang implementations or models with custom pipeline classes.

python/sglang/cli/generate.py

@@ -25,8 +25,8 @@ def generate(args, extra_argv):
 
         parser = argparse.ArgumentParser(description="SGLang Multimodal Generation")
         add_multimodal_gen_generate_args(parser)
-        parsed_args = parser.parse_args(extra_argv)
+        parsed_args, unknown_args = parser.parse_known_args(extra_argv)
-        generate_cmd(parsed_args)
+        generate_cmd(parsed_args, unknown_args)
     else:
         raise Exception(
             f"Generate subcommand is not yet supported for model: {model_path}"

python/sglang/multimodal_gen/runtime/entrypoints/cli/cli_types.py

@@ -13,7 +13,9 @@ class CLISubcommand:
 
     name: str
 
-    def cmd(self, args: argparse.Namespace) -> None:
+    def cmd(
+        self, args: argparse.Namespace, unknown_args: list[str] | None = None
+    ) -> None:
         """Execute the command with the given arguments"""
         raise NotImplementedError
 

python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py

@@ -106,11 +106,11 @@ def maybe_dump_performance(
     )
 
 
-def generate_cmd(args: argparse.Namespace):
+def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None):
     """The entry point for the generate command."""
     args.request_id = "mocked_fake_id_for_offline_generate"
 
-    server_args = ServerArgs.from_cli_args(args)
+    server_args = ServerArgs.from_cli_args(args, unknown_args)
 
     sampling_params_kwargs = SamplingParams.get_cli_args(args)
     sampling_params_kwargs["request_id"] = generate_request_id()
@@ -158,8 +158,10 @@ class GenerateSubcommand(CLISubcommand):
         """Get names of arguments for generate_video method"""
         return [field.name for field in dataclasses.fields(SamplingParams)]
 
-    def cmd(self, args: argparse.Namespace) -> None:
+    def cmd(
-        generate_cmd(args)
+        self, args: argparse.Namespace, unknown_args: list[str] | None = None
+    ) -> None:
+        generate_cmd(args, unknown_args)
 
     def validate(self, args: argparse.Namespace) -> None:
         """Validate the arguments for this command"""

python/sglang/multimodal_gen/runtime/entrypoints/cli/main.py

@@ -30,12 +30,12 @@ def main() -> None:
     for cmd in cmd_init():
         cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
         cmds[cmd.name] = cmd
-    args = parser.parse_args()
+    args, unknown_args = parser.parse_known_args()
     if args.subparser in cmds:
         cmds[args.subparser].validate(args)
 
     if hasattr(args, "dispatch_function"):
-        args.dispatch_function(args)
+        args.dispatch_function(args, unknown_args=unknown_args)
     else:
         parser.print_help()
 

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

@@ -31,7 +31,7 @@ class AdapterLoader(ComponentLoader):
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, *args
     ):
-        config = get_diffusers_component_config(model_path=component_model_path)
+        config = get_diffusers_component_config(component_path=component_model_path)
 
         cls_name = config.pop("_class_name", None)
         if cls_name is None:

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

@@ -30,7 +30,7 @@ class BridgeLoader(ComponentLoader):
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
-        config = get_diffusers_component_config(model_path=component_model_path)
+        config = get_diffusers_component_config(component_path=component_model_path)
         hf_config = deepcopy(config)
         class_name = config.pop("_class_name", None)
         if class_name is None:

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

@@ -21,7 +21,7 @@ class SchedulerLoader(ComponentLoader):
         self, component_model_path: str, server_args: ServerArgs, *args
     ):
         """Load the scheduler based on the model path, and inference args."""
-        config = get_diffusers_component_config(model_path=component_model_path)
+        config = get_diffusers_component_config(component_path=component_model_path)
 
         class_name = config.pop("_class_name")
         assert (

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

@@ -182,7 +182,9 @@ class TextEncoderLoader(ComponentLoader):
         diffusers_pretrained_config = get_config(
             component_model_path, trust_remote_code=True
         )
-        model_config = get_diffusers_component_config(model_path=component_model_path)
+        model_config = get_diffusers_component_config(
+            component_path=component_model_path
+        )
         _clean_hf_config_inplace(model_config)
         logger.debug("HF model config: %s", model_config)
 

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

@@ -69,7 +69,7 @@ class TransformerLoader(ComponentLoader):
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
         """Load the transformer based on the model path, and inference args."""
-        config = get_diffusers_component_config(model_path=component_model_path)
+        config = get_diffusers_component_config(component_path=component_model_path)
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
         if cls_name is None:

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

@@ -60,7 +60,7 @@ class VAELoader(ComponentLoader):
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
         """Load the VAE based on the model path, and inference args."""
-        config = get_diffusers_component_config(model_path=component_model_path)
+        config = get_diffusers_component_config(component_path=component_model_path)
         class_name = config.pop("_class_name", None)
         assert (
             class_name is not None

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

@@ -32,7 +32,7 @@ class VocoderLoader(ComponentLoader):
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
-        config = get_diffusers_component_config(model_path=component_model_path)
+        config = get_diffusers_component_config(component_path=component_model_path)
         class_name = config.pop("_class_name", None)
         assert (
             class_name is not None

python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py

@@ -391,7 +391,7 @@ class DiffusersPipeline(ComposedPipelineBase):
         """
 
         original_model_path = model_path  # Keep original for custom_pipeline
-        model_path = maybe_download_model(model_path)
+        model_path = maybe_download_model(model_path, force_diffusers_model=True)
         self.model_path = model_path
 
         dtype = self._get_dtype(server_args)

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -126,9 +126,8 @@ class ComposedPipelineBase(ABC):
         self.modules[module_name] = module
 
     def _load_config(self) -> dict[str, Any]:
-        model_path = maybe_download_model(self.model_path)
+        model_path = maybe_download_model(self.model_path, force_diffusers_model=True)
         self.model_path = model_path
-        # server_args.downloaded_model_path = model_path
         logger.info("Model path: %s", model_path)
         config = verify_model_config_and_directory(model_path)
         return cast(dict[str, Any], config)
@@ -171,6 +170,19 @@ class ComposedPipelineBase(ABC):
         """
         return
 
+    def _resolve_component_path(
+        self, server_args: ServerArgs, module_name: str, load_module_name: str
+    ) -> str:
+        override_path = server_args.component_paths.get(module_name)
+        if override_path is not None:
+            # overridden with args like --vae-path
+            component_model_path = maybe_download_model(override_path)
+        else:
+            component_model_path = os.path.join(self.model_path, load_module_name)
+
+        logger.debug("Resolved component path: %s", component_model_path)
+        return component_model_path
+
     def load_modules(
         self,
         server_args: ServerArgs,
@@ -285,19 +297,9 @@ class ComposedPipelineBase(ABC):
             else:
                 load_module_name = module_name
 
-            # Use custom VAE path if provided, otherwise use default path
+            component_model_path = self._resolve_component_path(
-            if module_name == "vae" and server_args.vae_path is not None:
+                server_args, module_name, load_module_name
-                component_model_path = server_args.vae_path
+            )
-                # Download from HuggingFace Hub if path doesn't exist locally
-                if not os.path.exists(component_model_path):
-                    component_model_path = maybe_download_model(component_model_path)
-                logger.info(
-                    "Using custom VAE path: %s instead of default path: %s",
-                    component_model_path,
-                    os.path.join(self.model_path, load_module_name),
-                )
-            else:
-                component_model_path = os.path.join(self.model_path, load_module_name)
             module, memory_usage = PipelineComponentLoader.load_component(
                 component_name=load_module_name,
                 component_model_path=component_model_path,

python/sglang/multimodal_gen/runtime/server_args.py

@@ -287,8 +287,8 @@ class ServerArgs:
     lora_nickname: str = "default"  # for swapping adapters in the pipeline
     lora_scale: float = 1.0  # LoRA scale for merging (e.g., 0.125 for Hyper-SD)
 
-    # VAE parameters
+    # Component path overrides (key = model_index.json component name, value = path)
-    vae_path: str | None = None  # Custom VAE path (e.g., for distilled autoencoder)
+    component_paths: dict[str, str] = field(default_factory=dict)
     # 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
@@ -623,13 +623,6 @@ class ServerArgs:
             type=str,
             help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
         )
-        parser.add_argument(
-            "--vae-path",
-            type=str,
-            default=ServerArgs.vae_path,
-            help="Custom path to VAE model (e.g., for distilled autoencoder). If not specified, VAE will be loaded from the main model path.",
-        )
-
         # attention
         parser.add_argument(
             "--attention-backend",
@@ -963,24 +956,66 @@ class ServerArgs:
             f"(started from port {original_port})"
         )
 
+    @staticmethod
+    def _extract_component_paths(
+        unknown_args: list[str],
+    ) -> tuple[dict[str, str], list[str]]:
+        """
+        Extract dynamic ``--<component>-path`` args from unrecognised CLI args.
+        """
+        component_paths: dict[str, str] = {}
+        remaining: list[str] = []
+        i = 0
+        while i < len(unknown_args):
+            arg = unknown_args[i]
+            key_part = arg.split("=", 1)[0] if "=" in arg else arg
+            if key_part.startswith("--") and key_part.endswith("-path"):
+                component = key_part[2:-5].replace("-", "_")
+                if "=" in arg:
+                    component_paths[component] = arg.split("=", 1)[1]
+                elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
+                    "-"
+                ):
+                    i += 1
+                    component_paths[component] = unknown_args[i]
+                else:
+                    remaining.append(arg)
+                    i += 1
+                    continue
+            else:
+                remaining.append(arg)
+            i += 1
+
+        # canonicalize and validate
+        for component, path in component_paths.items():
+            path = os.path.expanduser(path)
+            component_paths[component] = path
+        return component_paths, remaining
+
     @classmethod
     def from_cli_args(
         cls, args: argparse.Namespace, unknown_args: list[str] | None = None
     ) -> "ServerArgs":
         if unknown_args is None:
             unknown_args = []
+
+        # extract dynamic --<component>-path from unknown args
+        dynamic_paths, remaining = cls._extract_component_paths(unknown_args)
+        if remaining:
+            raise SystemExit(f"error: unrecognized arguments: {' '.join(remaining)}")
+
         provided_args = cls.get_provided_args(args, unknown_args)
 
         # Handle config file
         config_file = provided_args.get("config")
         if config_file:
             config_args = cls.load_config_file(config_file)
-            # Provided args override config file args
             provided_args = {**config_args, **provided_args}
 
-        # Handle special cases
+        if dynamic_paths:
-        # if "tp_size" in provided_args:
+            existing = dict(provided_args.get("component_paths") or {})
-        #     provided_args["tp"] = provided_args.pop("tp_size")
+            existing.update(dynamic_paths)
+            provided_args["component_paths"] = existing
 
         return cls.from_dict(provided_args)
 
@@ -990,6 +1025,10 @@ class ServerArgs:
         attrs = [attr.name for attr in dataclasses.fields(cls)]
         server_args_kwargs: dict[str, Any] = {}
 
+        component_paths = dict(kwargs.get("component_paths") or {})
+        if component_paths:
+            server_args_kwargs["component_paths"] = component_paths
+
         for attr in attrs:
             if attr == "pipeline_config":
                 pipeline_config = PipelineConfig.from_kwargs(kwargs)
@@ -1239,8 +1278,8 @@ def prepare_server_args(argv: list[str]) -> ServerArgs:
     """
     parser = FlexibleArgumentParser()
     ServerArgs.add_cli_args(parser)
-    raw_args = parser.parse_args(argv)
+    raw_args, unknown_args = parser.parse_known_args(argv)
-    server_args = ServerArgs.from_cli_args(raw_args)
+    server_args = ServerArgs.from_cli_args(raw_args, unknown_args)
     return server_args
 
 

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

@@ -69,9 +69,6 @@ def _check_index_files_for_missing_shards(
     missing_files = []
     checked_subdirs = []
 
-    # Check the root directory and all subdirectories that might contain model weights
-    dirs_to_check = [model_path]
-
     # Add common subdirectories for diffusers models
     try:
         subdirs = os.listdir(model_path)
@@ -79,6 +76,9 @@ def _check_index_files_for_missing_shards(
         logger.warning("Failed to list model directory %s: %s", model_path, e)
         return True, [], []  # Assume valid if we can't check
 
+    # Check the root directory and all subdirectories that might contain model weights
+    dirs_to_check = [model_path]
+
     for subdir in subdirs:
         subdir_path = os.path.join(model_path, subdir)
         if os.path.isdir(subdir_path):
@@ -176,7 +176,6 @@ def _ci_validate_diffusers_model(model_path: str) -> tuple[bool, bool]:
     """
     if not is_in_ci():
         return True, False
-
     is_valid, missing_files, checked_subdirs = _check_index_files_for_missing_shards(
         model_path
     )
@@ -295,24 +294,23 @@ def load_dict(file_path):
 
 
 def get_diffusers_component_config(
-    model_path: str,
+    component_path: str,
 ) -> dict[str, Any]:
     """Gets a configuration of a submodule for the given diffusers model."""
-
     # Download from HuggingFace Hub if path doesn't exist locally
-    if not os.path.exists(model_path):
+    if not os.path.exists(component_path):
-        model_path = maybe_download_model(model_path)
+        component_path = maybe_download_model(component_path)
 
     # tokenizer
     config_names = ["generation_config.json"]
     # By default, we load config.json, but scheduler_config.json for scheduler
-    if "scheduler" in model_path:
+    if "scheduler" in component_path:
         config_names.append("scheduler_config.json")
     else:
         config_names.append("config.json")
 
     config_file_paths = [
-        os.path.join(model_path, config_name) for config_name in config_names
+        os.path.join(component_path, config_name) for config_name in config_names
     ]
 
     combined_config = reduce(
@@ -544,6 +542,7 @@ def maybe_download_model(
     download: bool = True,
     is_lora: bool = False,
     allow_patterns: list[str] | None = None,
+    force_diffusers_model: bool = False,
 ) -> str:
     """
     Check if the model path is a Hugging Face Hub model ID and download it if needed.
@@ -553,13 +552,13 @@ def maybe_download_model(
         local_dir: Local directory to save the model
         download: Whether to download the model from Hugging Face Hub
         is_lora: If True, skip model completeness verification (LoRA models don't have transformer/vae directories)
-
+        force_diffusers_model: If True, apply diffusers model check. Otherwise it should be a component model
     Returns:
         Local path to the model
     """
 
-    def _verify_model_complete(path: str) -> bool:
+    def _verify_diffusers_model_complete(path: str) -> bool:
-        """Check if model directory has required subdirectories."""
+        """Check if model directory (of a diffusers model, not a component) has required subdirectories."""
         config_path = os.path.join(path, "model_index.json")
         if not os.path.exists(config_path):
             return False
@@ -592,7 +591,10 @@ def maybe_download_model(
 
     # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA)
     if os.path.exists(model_name_or_path):
-        if is_lora or _verify_model_complete(model_name_or_path):
+        # TODO: lots of duplication here
+        if not force_diffusers_model:
+            return model_name_or_path
+        elif is_lora or _verify_diffusers_model_complete(model_name_or_path):
             # CI validation: check all subdirectories for missing shards
             if not is_lora:
                 is_valid, cleanup_performed = _ci_validate_diffusers_model(
@@ -640,7 +642,9 @@ def maybe_download_model(
             local_files_only=True,
             max_workers=8,
         )
-        if is_lora or _verify_model_complete(local_path):
+        if not force_diffusers_model:
+            return str(local_path)
+        elif is_lora or _verify_diffusers_model_complete(local_path):
             # CI validation: check all subdirectories for missing shards
             if not is_lora:
                 is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path)
@@ -710,8 +714,10 @@ def maybe_download_model(
                     max_workers=8,
                 )
 
+            if not force_diffusers_model:
+                return str(local_path)
             # Verify downloaded model is complete (skip for LoRA)
-            if not is_lora and not _verify_model_complete(local_path):
+            elif not is_lora and not _verify_diffusers_model_complete(local_path):
                 logger.warning(
                     "Downloaded model at %s is incomplete, retrying with force_download=True",
                     local_path,
@@ -724,7 +730,7 @@ def maybe_download_model(
                         max_workers=8,
                         force_download=True,
                     )
-                if not _verify_model_complete(local_path):
+                if not _verify_diffusers_model_complete(local_path):
                     raise ValueError(
                         f"Downloaded model at {local_path} is still incomplete after forced re-download. "
                         "The model repository may be missing required components (model_index.json, transformer/, or vae/)."

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

@@ -243,20 +243,24 @@ Consider updating perf_baselines.json with the snippets below:
 
         summary = validator.collect_metrics(perf_record)
 
-        if is_baseline_generation_mode or missing_scenario:
+        if case.run_perf_check:
-            self._dump_baseline_for_testcase(case, summary, missing_scenario)
+            if is_baseline_generation_mode or missing_scenario:
-            if missing_scenario:
+                self._dump_baseline_for_testcase(case, summary, missing_scenario)
-                pytest.fail(f"Testcase '{case.id}' not found in perf_baselines.json")
+                if missing_scenario:
-            return
+                    pytest.fail(
+                        f"Testcase '{case.id}' not found in perf_baselines.json"
+                    )
+                return
 
-        self._check_for_improvement(case, summary, scenario)
+            self._check_for_improvement(case, summary, scenario)
 
-        try:
+            # only run performance validation if run_perf_check is True
-            validator.validate(perf_record, case.sampling_params.num_frames)
+            try:
-        except AssertionError as e:
+                validator.validate(perf_record, case.sampling_params.num_frames)
-            logger.error(f"Performance validation failed for {case.id}:\n{e}")
+            except AssertionError as e:
-            self._dump_baseline_for_testcase(case, summary, missing_scenario)
+                logger.error(f"Performance validation failed for {case.id}:\n{e}")
-            raise
+                self._dump_baseline_for_testcase(case, summary, missing_scenario)
+                raise
 
         result = {
             "test_name": case.id,

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

@@ -226,6 +226,7 @@ class DiffusionTestCase:
     id: str  # pytest test id and scenario name
     server_args: DiffusionServerArgs
     sampling_params: DiffusionSamplingParams
+    run_perf_check: bool = True
 
 
 def sample_step_indices(
@@ -507,6 +508,16 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
         ),
         TI2I_sampling_params,
     ),
+    DiffusionTestCase(
+        "flux_2_t2i_customized_vae_path",
+        DiffusionServerArgs(
+            model_path="black-forest-labs/FLUX.2-dev",
+            modality="image",
+            extras=["--vae-path=fal/FLUX.2-Tiny-AutoEncoder"],
+        ),
+        T2I_sampling_params,
+        run_perf_check=False,
+    ),
     DiffusionTestCase(
         "fast_hunyuan_video",
         DiffusionServerArgs(