chore: import upstream snapshot with attribution

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

@@ -0,0 +1,76 @@
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import (
+    LTX2ConnectorConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _list_safetensors_files,
+    set_default_torch_dtype,
+    skip_init_modules,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.precision import resolve_precision
+
+
+class AdapterLoader(ComponentLoader):
+    """Loader for small adapter-style modules (e.g., LTX-2 connectors).
+
+    This loader intentionally avoids FSDP sharding and just:
+    1) Instantiates the module from `config.json`.
+    2) Loads a single safetensors state_dict.
+    """
+
+    component_names = ["connectors"]
+    expected_library = "diffusers"
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, *args
+    ):
+        config = get_diffusers_component_config(component_path=component_model_path)
+
+        cls_name = config.pop("_class_name", None)
+        if cls_name is None:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported."
+            )
+
+        config.pop("_diffusers_version", None)
+        config.pop("_name_or_path", None)
+
+        server_args.model_paths["connectors"] = component_model_path
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+
+        target_device = get_local_torch_device()
+        default_dtype = resolve_precision(
+            server_args, "connectors", precision_attr="dit_precision"
+        )
+
+        with set_default_torch_dtype(default_dtype), skip_init_modules():
+            connector_cfg = LTX2ConnectorConfig()
+            connector_cfg.update_model_arch(config)
+            model = model_cls(connector_cfg).to(
+                device=target_device, dtype=default_dtype
+            )
+
+        safetensors_list = _list_safetensors_files(component_model_path)
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {component_model_path}")
+        if len(safetensors_list) != 1:
+            raise ValueError(
+                f"Found {len(safetensors_list)} safetensors files in {component_model_path}, expected 1"
+            )
+
+        loaded = safetensors_load_file(safetensors_list[0])
+        model.load_state_dict(loaded, strict=False)
+
+        return model

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

@@ -0,0 +1,112 @@
+from copy import deepcopy
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
+from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.precision import resolve_precision
+
+logger = init_logger(__name__)
+
+
+class BridgeLoader(ComponentLoader):
+    """Loader for MOVA dual tower bridge with FSDP support."""
+
+    pipeline_bridge_config_attr: str = "bridge_config"
+
+    component_names = ["dual_tower_bridge"]
+    expected_library = "diffusers"
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ):
+        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:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported."
+            )
+        server_args.model_paths[component_name] = component_model_path
+
+        # Try to get bridge config from pipeline config, fallback to creating one
+        bridge_config = getattr(
+            server_args.pipeline_config, self.pipeline_bridge_config_attr, None
+        )
+        if bridge_config is not None:
+            bridge_config.update_model_arch(config)
+        else:
+            # Create a minimal config from hf_config
+            from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
+                MOVADualTowerConfig,
+            )
+
+            bridge_config = MOVADualTowerConfig()
+            bridge_config.update_model_arch(config)
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        # Find all safetensors files
+        safetensors_list = _list_safetensors_files(component_model_path)
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {component_model_path}")
+
+        default_dtype = resolve_precision(
+            server_args, component_name, precision_attr="dit_precision"
+        )
+
+        logger.info(
+            "Loading %s from %s safetensors files, default_dtype: %s",
+            class_name,
+            len(safetensors_list),
+            default_dtype,
+        )
+
+        # Use the FSDP loader when FSDP is requested or shard rules are declared.
+        fsdp_shard_conditions = getattr(model_cls, "_fsdp_shard_conditions", None)
+        if server_args.use_fsdp_inference or (
+            server_args.hsdp_shard_dim is not None and fsdp_shard_conditions
+        ):
+            local_torch_device = get_local_torch_device()
+            # Load with FSDP support
+            model = maybe_load_fsdp_model(
+                model_cls=model_cls,
+                init_params={"config": bridge_config, "hf_config": hf_config},
+                weight_dir_list=safetensors_list,
+                device=local_torch_device,
+                hsdp_replicate_dim=server_args.hsdp_replicate_dim,
+                hsdp_shard_dim=server_args.hsdp_shard_dim,
+                cpu_offload=server_args.dit_cpu_offload,
+                pin_cpu_memory=server_args.pin_cpu_memory,
+                fsdp_inference=server_args.use_fsdp_inference,
+                param_dtype=default_dtype,
+                reduce_dtype=torch.float32,
+                output_dtype=None,
+                strict=False,
+                weight_load_plan=WeightLoadPlan(
+                    checkpoint_load_device=local_torch_device
+                ),
+            )
+        else:
+            # Fallback to simple loading (for non-FSDP or legacy models)
+            model = model_cls.from_pretrained(
+                component_model_path, torch_dtype=default_dtype
+            )
+            model = model.to(device=get_local_torch_device(), dtype=default_dtype)
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded bridge model with %.2fM parameters", total_params / 1e6)
+
+        return model

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

@@ -0,0 +1,587 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+import os
+import pkgutil
+import traceback
+from abc import ABC
+from typing import Any, Type
+
+import torch
+from diffusers import AutoModel
+from torch import nn
+from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
+
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    component_attn_backend_context_manager,
+    get_component_attn_backend_context,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _normalize_component_type,
+    component_name_to_loader_cls,
+    get_memory_usage_of_component,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    configure_layerwise_offload_modules,
+    is_layerwise_offloaded_module,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
+    LAYERWISE_OFFLOAD_ALL_COMPONENTS,
+    LAYERWISE_OFFLOAD_DIT_GROUP,
+    layerwise_component_matches_any_selection,
+    normalize_layerwise_offload_components,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_hf_config,
+    prepare_diffusers_component_path_for_loading,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision
+
+logger = init_logger(__name__)
+
+
+def _load_auto_tokenizer_with_roberta_processing_compat(*args, **kwargs):
+    from tokenizers import processors
+
+    roberta_processing = processors.RobertaProcessing
+
+    def roberta_processing_compat(*processor_args, **processor_kwargs):
+        if "sep" in processor_kwargs and "cls" in processor_kwargs:
+            sep = processor_kwargs.pop("sep")
+            cls_token = processor_kwargs.pop("cls")
+            return roberta_processing(
+                sep, cls_token, *processor_args, **processor_kwargs
+            )
+        return roberta_processing(*processor_args, **processor_kwargs)
+
+    processors.RobertaProcessing = roberta_processing_compat
+    try:
+        return AutoTokenizer.from_pretrained(*args, **kwargs)
+    finally:
+        processors.RobertaProcessing = roberta_processing
+
+
+class ComponentLoader(ABC):
+    """Base class for loading a specific type of model component."""
+
+    # the list of possible name of the component in model_index.json, e.g., scheduler
+    component_names: list[str] = []
+
+    # diffusers or transformers
+    expected_library: str = ""
+
+    _loaders_registered = False
+
+    def __init_subclass__(cls, **kwargs):
+        """
+        register loaders, called when subclass is imported
+        """
+        super().__init_subclass__(**kwargs)
+        for component_name in cls.component_names:
+            component_name_to_loader_cls[component_name] = cls
+
+    def __init__(self, device=None) -> None:
+        self.device = device
+        self.component_architecture: str | None = None
+
+    def should_offload(
+        self, server_args: ServerArgs, model_config: ModelConfig | None = None
+    ):
+        # not offload by default
+        return False
+
+    def target_device(self, should_offload):
+        if should_offload:
+            return (
+                torch.device("mps")
+                if current_platform.is_mps()
+                else torch.device("cpu")
+            )
+        else:
+            return get_local_torch_device()
+
+    def customized_load_kwargs_for_component(
+        self, _server_args: ServerArgs, _component_name: str
+    ) -> dict[str, Any]:
+        return {}
+
+    def should_raise_customized_load_error(
+        self, _server_args: ServerArgs, _component_name: str
+    ) -> bool:
+        return False
+
+    @staticmethod
+    def _is_component_set_as_layerwise_load(
+        server_args: ServerArgs, component_name: str
+    ) -> bool:
+        """if a component should be loaded in a layerwise-fashion"""
+        selected_component_names = normalize_layerwise_offload_components(
+            server_args.layerwise_offload_components
+        )
+        if selected_component_names is None:
+            return False
+        selected_component_names = set(selected_component_names)
+        if LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names:
+            return True
+        explicit_component_names = selected_component_names - {
+            LAYERWISE_OFFLOAD_DIT_GROUP
+        }
+        return layerwise_component_matches_any_selection(
+            component_name, explicit_component_names
+        )
+
+    def _maybe_configure_layerwise_after_startup_cpu_staging(
+        self,
+        component: AutoModel,
+        server_args: ServerArgs,
+        component_name: str,
+        load_kwargs: dict[str, Any],
+    ) -> AutoModel:
+        if not load_kwargs.get("cpu_offload_flag"):
+            return component
+        if not isinstance(component, nn.Module):
+            return component
+
+        # try to configure layerwise-offload with the component
+        configured_components = configure_layerwise_offload_modules(
+            {component_name: component},
+            server_args,
+            component_names=server_args.layerwise_offload_components,
+            warn_missing=False,
+        )
+        if is_layerwise_offloaded_module(component):
+            logger.info(
+                "Configured layerwise offload for %s immediately after startup CPU staging",
+                component_name,
+            )
+            return component
+
+        logger.warning(
+            "Layerwise startup CPU staging was requested for %s, but the loaded "
+            "module did not enable layerwise offload. Moving it to GPU.",
+            component_name,
+        )
+        # ensures the module is on GPU
+        if component_name in configured_components:
+            return component
+        return component.to(get_local_torch_device())
+
+    def _load_customized_with_context(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+        attn_backend: Any,
+        component_attn_name: str | None,
+    ) -> AutoModel:
+        with component_attn_backend_context_manager(
+            attn_backend, component_name=component_attn_name
+        ):
+            load_kwargs = self.customized_load_kwargs_for_component(
+                server_args, component_name
+            )
+            component = self.load_customized(
+                component_model_path, server_args, component_name, **load_kwargs
+            )
+            return self._maybe_configure_layerwise_after_startup_cpu_staging(
+                component, server_args, component_name, load_kwargs
+            )
+
+    def _load_native_with_context(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+        transformers_or_diffusers: str,
+        attn_backend: Any,
+        component_attn_name: str | None,
+    ) -> AutoModel:
+        with component_attn_backend_context_manager(
+            attn_backend, component_name=component_attn_name
+        ):
+            component = self.load_native(
+                component_model_path,
+                server_args,
+                transformers_or_diffusers,
+                component_name,
+            )
+        should_offload = self.should_offload(server_args)
+        target_device = self.target_device(should_offload)
+        return component.to(device=target_device)
+
+    def load(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+        transformers_or_diffusers: str,
+    ) -> tuple[AutoModel, float]:
+        """
+        Template method that standardizes logging around the core load implementation.
+        The priority of loading method is:
+            1. load customized component
+            2. load native diffusers/transformers component
+        If all of the above methods failed, an error will be thrown
+
+        """
+        gpu_mem_before_loading = current_platform.get_available_gpu_memory()
+        logger.info(
+            "Loading %s from %s. avail mem: %.2f GB",
+            component_name,
+            component_model_path,
+            gpu_mem_before_loading,
+        )
+        attn_backend = None
+        component_attn_name = None
+        if get_component_attn_backend_context() is None:
+            attn_backend, matched_backend_key = (
+                server_args.resolve_component_attention_backend(component_name)
+            )
+            component_attn_name = matched_backend_key or component_name
+            if attn_backend is not None:
+                logger.info(
+                    "Using %s backend for component: %s",
+                    attn_backend.name.lower(),
+                    matched_backend_key,
+                )
+        try:
+            component = self._load_customized_with_context(
+                component_model_path,
+                server_args,
+                component_name,
+                attn_backend,
+                component_attn_name,
+            )
+            source = "sgl-diffusion"
+        except Exception as e:
+            if self.should_raise_customized_load_error(server_args, component_name):
+                traceback.print_exc()
+                raise RuntimeError(
+                    f"Failed to load customized {component_name}; native fallback "
+                    "is disabled for this component configuration."
+                ) from e
+            if "Unsupported model architecture" in str(e):
+                logger.info(
+                    f"Component: {component_name} doesn't have a customized version yet, using native version"
+                )
+            else:
+                traceback.print_exc()
+                logger.error(
+                    f"Error while loading customized {component_name}, falling back to native version"
+                )
+            # fallback to native version
+            component = self._load_native_with_context(
+                component_model_path,
+                server_args,
+                component_name,
+                transformers_or_diffusers,
+                attn_backend,
+                component_attn_name,
+            )
+            source = "native"
+            logger.warning(
+                "Native component %s: %s is loaded, performance may be sub-optimal",
+                component_name,
+                component.__class__.__name__,
+            )
+
+        if component is None:
+            logger.error("Load %s failed", component_name)
+            consumed = 0.0
+        else:
+            if isinstance(component, nn.Module):
+                component = component.eval()
+            current_gpu_mem = current_platform.get_available_gpu_memory()
+            model_size = get_memory_usage_of_component(component) or "NA"
+            consumed = gpu_mem_before_loading - current_gpu_mem
+            logger.info(
+                f"Loaded %s: %s ({source} version). model size: %s GB, consumed GPU mem: %.2f GB, avail GPU mem: %.2f GB",
+                component_name,
+                component.__class__.__name__,
+                model_size,
+                consumed,
+                current_gpu_mem,
+            )
+        return component, consumed
+
+    def load_native(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        transformers_or_diffusers: str,
+        component_name: str | None = None,
+    ) -> AutoModel:
+        """
+        Load the component using the native library (transformers/diffusers).
+        """
+        precision = (
+            resolve_component_precision(server_args, component_name)
+            if component_name is not None
+            else None
+        )
+        load_kwargs = {}
+        if precision is not None:
+            load_kwargs["torch_dtype"] = precision
+
+        if transformers_or_diffusers == "transformers":
+            from transformers import AutoModel
+
+            config = get_hf_config(
+                component_model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+            )
+            return AutoModel.from_pretrained(
+                component_model_path,
+                config=config,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+                **load_kwargs,
+            )
+        elif transformers_or_diffusers == "diffusers":
+            from diffusers import AutoModel
+
+            component_model_path = prepare_diffusers_component_path_for_loading(
+                component_model_path
+            )
+            return AutoModel.from_pretrained(
+                component_model_path,
+                revision=server_args.revision,
+                trust_remote_code=server_args.trust_remote_code,
+                **load_kwargs,
+            )
+        else:
+            raise ValueError(f"Unsupported library: {transformers_or_diffusers}")
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ):
+        """
+        Load the customized version component, implemented and optimized in SGL-diffusion
+        """
+        raise NotImplementedError(
+            f"load_customized not implemented for {self.__class__.__name__}"
+        )
+
+    @classmethod
+    def _ensure_loaders_registered(cls):
+        """
+        avoid multiple registration
+        """
+        if cls._loaders_registered:
+            return
+
+        package_dir = os.path.dirname(__file__)
+        package_name = (
+            __package__
+            or "sglang.multimodal_gen.runtime.loader.component_loaders.component_loaders"
+        )
+
+        for _, name, _ in pkgutil.iter_modules([package_dir]):
+            # skip importing self to avoid circular dependency issues
+            if name == "component_loader":
+                continue
+            try:
+                importlib.import_module(f".{name}", package=package_name)
+            except ImportError as e:
+                logger.warning(f"Failed to import loader component {name}: {e}")
+
+        cls._loaders_registered = True
+
+    @classmethod
+    def resolve_transformers_or_diffusers(
+        self, transformers_or_diffusers: str, component_name: str
+    ) -> str:
+        # NOTE(FlamingoPg): special for LTX-2 models
+        if component_name == "vocoder" or component_name == "connectors":
+            transformers_or_diffusers = "diffusers"
+
+        # NOTE(CloudRipple): special for MOVA models
+        # TODO(CloudRipple): remove most of these special cases after unifying the loading logic
+        if component_name in [
+            "audio_vae",
+            "audio_dit",
+            "dual_tower_bridge",
+            "video_dit",
+        ]:
+            transformers_or_diffusers = "diffusers"
+
+        if (
+            component_name == "scheduler"
+            and transformers_or_diffusers == "mova.diffusion.schedulers.flow_match_pair"
+        ):
+            transformers_or_diffusers = "diffusers"
+
+        return transformers_or_diffusers
+
+    @classmethod
+    def for_component_type(
+        cls,
+        component_name: str,
+        transformers_or_diffusers: str,
+        component_architecture: str | None = None,
+    ) -> "ComponentLoader":
+        """
+        Factory method to create a component loader for a specific component type.
+
+        Args:
+            component_name: Type of component (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            transformers_or_diffusers: Whether the component is from transformers or diffusers
+        """
+        cls._ensure_loaders_registered()
+
+        # Map of component types to their loader classes and expected library
+        component_name = _normalize_component_type(component_name)
+
+        transformers_or_diffusers = cls.resolve_transformers_or_diffusers(
+            transformers_or_diffusers, component_name
+        )
+
+        if component_name in component_name_to_loader_cls:
+            loader_cls: Type[ComponentLoader] = component_name_to_loader_cls[
+                component_name
+            ]
+            expected_library = loader_cls.expected_library
+            # Assert that the library matches what's expected for this component type
+            assert (
+                transformers_or_diffusers == expected_library
+            ), f"{component_name} must be loaded from {expected_library}, got {transformers_or_diffusers}"
+            loader = loader_cls()
+            loader.component_architecture = component_architecture
+            return loader
+
+        # For unknown component types, use a generic loader
+        logger.warning(
+            "No specific loader found for component type: %s. Using generic loader.",
+            component_name,
+        )
+        return GenericComponentLoader(transformers_or_diffusers, component_architecture)
+
+
+class ImageProcessorLoader(ComponentLoader):
+    """Loader for image processor."""
+
+    component_names = ["image_processor"]
+    expected_library = "transformers"
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ) -> Any:
+        return AutoImageProcessor.from_pretrained(component_model_path, use_fast=True)
+
+
+class AutoProcessorLoader(ComponentLoader):
+    """Loader for auto processor."""
+
+    component_names = ["processor"]
+    expected_library = "transformers"
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ) -> Any:
+        return AutoProcessor.from_pretrained(component_model_path)
+
+
+class TokenizerLoader(ComponentLoader):
+    """Loader for tokenizers."""
+
+    component_names = ["tokenizer", "text_tokenizer"]
+    expected_library = "transformers"
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ) -> Any:
+        # Some pipelines keep the slot name `tokenizer` in model_index.json even
+        # when the declared class is a processor. e.g. FLUX.2:
+        # `tokenizer: ["transformers", "PixtralProcessor"]`.
+        # Honor the declared component class instead of guessing from the slot name.
+        if (
+            self.component_architecture is not None
+            and self.component_architecture.endswith("Processor")
+        ):
+            return AutoProcessor.from_pretrained(component_model_path)
+
+        # Qwen-Image's model_index declares Qwen2Tokenizer; using the fast class
+        # changes text preprocessing and shifts official GT comparisons.
+        use_fast = self.component_architecture != "Qwen2Tokenizer"
+        try:
+            return AutoTokenizer.from_pretrained(
+                component_model_path,
+                padding_side="right",
+                use_fast=use_fast,
+            )
+        except TypeError as e:
+            # tokenizers>=0.21 removed the `cls` kwarg from RobertaProcessing,
+            # but some transformers CLIPTokenizer builds still pass it. Fall back
+            # to the pure-Python (slow) tokenizer which avoids the rust path.
+            if "RobertaProcessing" in str(e) and use_fast:
+                logger.warning(
+                    "Fast tokenizer failed (%s), retrying with use_fast=False", e
+                )
+                return _load_auto_tokenizer_with_roberta_processing_compat(
+                    component_model_path,
+                    padding_side="right",
+                    use_fast=False,
+                )
+            raise
+
+
+class GenericComponentLoader(ComponentLoader):
+    """Generic loader for components that don't have a specific loader."""
+
+    def __init__(
+        self, library="transformers", component_architecture: str | None = None
+    ) -> None:
+        super().__init__()
+        self.library = library
+        self.component_architecture = component_architecture
+
+
+class PipelineComponentLoader:
+    """
+    Utility class for loading the components in a pipeline.
+    """
+
+    @staticmethod
+    def load_component(
+        component_name: str,
+        component_model_path: str,
+        transformers_or_diffusers: str,
+        server_args: ServerArgs,
+        component_architecture: str | None = None,
+    ):
+        """
+        Load a pipeline component.
+
+        Args:
+            component_name: Name of the component (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            component_model_path: Path to the component model
+            transformers_or_diffusers: Whether the component is from transformers or diffusers
+            component_architecture: the class name of the module
+        """
+
+        # Get the appropriate loader for this component type
+        loader = ComponentLoader.for_component_type(
+            component_name, transformers_or_diffusers, component_architecture
+        )
+
+        try:
+            # Load the component
+            return loader.load(
+                component_model_path,
+                server_args,
+                component_name,
+                transformers_or_diffusers,
+            )
+        except Exception as e:
+            logger.error(
+                f"Error while loading component: {component_name}, {component_model_path=}"
+            )
+            raise e

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

@@ -0,0 +1,69 @@
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
+    TextEncoderLoader,
+)
+from sglang.multimodal_gen.runtime.models.encoders.base import finalize_encoder_folding
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class ImageEncoderLoader(TextEncoderLoader):
+    component_names = ["image_encoder"]
+    expected_library = "transformers"
+
+    def should_offload(self, server_args, model_config: ModelConfig | None = None):
+        should_offload = server_args.image_encoder_cpu_offload
+        if not should_offload:
+            return False
+        # _fsdp_shard_conditions is in arch_config, not directly on model_config
+        arch_config = (
+            getattr(model_config, "arch_config", model_config) if model_config else None
+        )
+        fsdp_shard_conditions = (
+            getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
+        )
+        use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
+        return use_cpu_offload
+
+    def load_customized(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str = "image_encoder",
+        cpu_offload_flag: bool | None = None,
+    ):
+        """Load the text encoders based on the model path, and inference args."""
+        # model_config: PretrainedConfig = get_hf_config(
+        #     model=model_path,
+        #     trust_remote_code=server_args.trust_remote_code,
+        #     revision=server_args.revision,
+        #     model_override_args=None,
+        # )
+        model_config = get_diffusers_component_config(
+            component_path=component_model_path
+        )
+
+        encoder_config = server_args.pipeline_config.image_encoder_config
+        encoder_config.update_model_arch(model_config)
+        # Keep the proposed fold group only if the encoder is wide enough
+        # (image encoders are small, so this normally reverts to replicated).
+        finalize_encoder_folding(encoder_config)
+
+        # Always start with local device; load_model will adjust for offload if needed
+        # TODO(will): add support for other dtypes
+        return self.load_model(
+            component_model_path,
+            encoder_config,
+            server_args,
+            server_args.pipeline_config.image_encoder_precision,
+            cpu_offload_flag=(
+                cpu_offload_flag
+                if cpu_offload_flag is not None
+                else server_args.image_encoder_cpu_offload
+            ),
+        )

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

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def _read_model_max_length(model_path: str) -> int | None:
+    """Read model_max_length from tokenizer_config.json in the given directory."""
+    config_path = os.path.join(model_path, "tokenizer_config.json")
+    if os.path.exists(config_path):
+        try:
+            with open(config_path, encoding="utf-8") as f:
+                config = json.load(f)
+            val = config.get("model_max_length")
+            if val is not None:
+                return int(val)
+        except Exception as e:
+            logger.warning(
+                "Failed to read tokenizer_config.json from %s: %s", model_path, e
+            )
+    return None
+
+
+class PEModelWrapper:
+
+    def __init__(self, model, tokenizer, device, model_max_length: int):
+        self.model = model
+        self.pe_tokenizer = tokenizer
+        self.device = device
+        self.model_max_length = model_max_length
+
+    def generate(self, prompt: str, sampling_params: dict) -> dict:
+        inputs = self.pe_tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=self.model_max_length,
+        ).to(self.device)
+
+        input_len = inputs["input_ids"].shape[1]
+
+        generate_kwargs = dict(
+            **inputs,
+            max_new_tokens=sampling_params.get("max_new_tokens", self.model_max_length),
+            do_sample=True,
+        )
+        temperature = sampling_params.get("temperature")
+        top_p = sampling_params.get("top_p")
+        if temperature is not None:
+            generate_kwargs["temperature"] = temperature
+        if top_p is not None:
+            generate_kwargs["top_p"] = top_p
+
+        with torch.no_grad():
+            output_ids = self.model.generate(**generate_kwargs)
+
+        new_tokens = output_ids[0, input_len:]
+        text = self.pe_tokenizer.decode(new_tokens, skip_special_tokens=True)
+        return {"text": text}
+
+    def to(self, *args, **kwargs):
+        """Move underlying model to device."""
+        self.model = self.model.to(*args, **kwargs)
+        if args:
+            device = args[0]
+            if isinstance(device, (str, torch.device)):
+                self.device = torch.device(device)
+        return self
+
+
+class PELoader(ComponentLoader):
+    """Loader for prompt-enhancement causal LM (Ministral-3 based)."""
+
+    component_names = ["pe"]
+    expected_library = "transformers"
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ):
+        logger.info("Loading PE model from %s ...", component_model_path)
+
+        pe_tokenizer_dir = os.path.join(
+            os.path.dirname(component_model_path), "pe_tokenizer"
+        )
+        if not os.path.exists(
+            os.path.join(component_model_path, "tokenizer_config.json")
+        ) and os.path.exists(os.path.join(pe_tokenizer_dir, "tokenizer_config.json")):
+            tokenizer_path = pe_tokenizer_dir
+            logger.info(
+                "PE tokenizer files not found in %s, using %s",
+                component_model_path,
+                tokenizer_path,
+            )
+        else:
+            tokenizer_path = component_model_path
+
+        model_max_length = _read_model_max_length(tokenizer_path)
+        if model_max_length is None:
+            raise RuntimeError(
+                f"Cannot load PE model: 'model_max_length' not found in "
+                f"{os.path.join(tokenizer_path, 'tokenizer_config.json')}. "
+                "Please ensure the PE component directory (or its sibling "
+                "pe_tokenizer/ directory) contains a valid tokenizer_config.json "
+                "with a 'model_max_length' field."
+            )
+        logger.info(
+            "PE model_max_length=%d (from tokenizer_config.json)", model_max_length
+        )
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path,
+            trust_remote_code=server_args.trust_remote_code,
+        )
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+
+        attn_impl = "flash_attention_2"
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                component_model_path,
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=server_args.trust_remote_code,
+                attn_implementation=attn_impl,
+            )
+            logger.info("PE model: using Flash Attention 2")
+        except (ValueError, ImportError):
+            logger.warning("Flash Attention 2 not available, falling back to SDPA")
+            attn_impl = "sdpa"
+            model = AutoModelForCausalLM.from_pretrained(
+                component_model_path,
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=server_args.trust_remote_code,
+                attn_implementation=attn_impl,
+            )
+
+        device = get_local_torch_device()
+        model = model.to(device).eval()
+
+        logger.info(
+            "PE model loaded on %s: %s (attn=%s)",
+            device,
+            model.__class__.__name__,
+            attn_impl,
+        )
+
+        return PEModelWrapper(
+            model=model,
+            tokenizer=tokenizer,
+            device=device,
+            model_max_length=model_max_length,
+        )

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

@@ -0,0 +1,48 @@
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SchedulerLoader(ComponentLoader):
+    """Loader for scheduler."""
+
+    component_names = ["scheduler"]
+    expected_library = "diffusers"
+
+    def load_customized(
+        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(component_path=component_model_path)
+
+        checkpoint_class_name = config.pop("_class_name", None)
+        class_name = (
+            getattr(server_args.pipeline_config, "scheduler_class_override", None)
+            or checkpoint_class_name
+        )
+        assert (
+            class_name is not None
+        ), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+
+        if checkpoint_class_name is not None and class_name != checkpoint_class_name:
+            logger.info(
+                "Overriding scheduler class from %s to %s",
+                checkpoint_class_name,
+                class_name,
+            )
+
+        scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        scheduler = scheduler_cls(**config)
+        if server_args.pipeline_config.flow_shift is not None:
+            scheduler.set_shift(server_args.pipeline_config.flow_shift)
+
+        return scheduler

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

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _list_safetensors_files,
+    set_default_torch_dtype,
+    skip_init_modules,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class SoundTokenizerLoader(ComponentLoader):
+    component_names = ["sound_tokenizer"]
+    expected_library = "diffusers"
+
+    def should_offload(
+        self, server_args: ServerArgs, model_config: ModelConfig | None = None
+    ) -> bool:
+        return server_args.vae_cpu_offload
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ):
+        config = get_diffusers_component_config(component_path=component_model_path)
+        class_name = config.pop("_class_name", None) or self.component_architecture
+        assert (
+            class_name is not None
+        ), "Sound tokenizer class name must be available from component config."
+
+        server_args.model_paths[component_name] = component_model_path
+
+        try:
+            precision = server_args.pipeline_config.vae_precision
+        except AttributeError:
+            precision = "bf16"
+        dtype = PRECISION_TO_TYPE[precision]
+        target_device = self.target_device(self.should_offload(server_args))
+
+        with set_default_torch_dtype(dtype), skip_init_modules():
+            model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            model = model_cls(config).to(target_device)
+
+        safetensors_list = _list_safetensors_files(component_model_path)
+        assert (
+            len(safetensors_list) == 1
+        ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        incompatible = model.load_state_dict(loaded, strict=False)
+        missing = getattr(incompatible, "missing_keys", [])
+        # The tokenizer is decoder-only; the checkpoint's encoder weights are
+        # expected leftovers, so they're excluded from the load warning.
+        unexpected = [
+            k
+            for k in getattr(incompatible, "unexpected_keys", [])
+            if not k.startswith("encoder.")
+        ]
+        if missing or unexpected:
+            logger.warning(
+                "Loaded sound_tokenizer with missing_keys=%d unexpected_keys=%d",
+                len(missing),
+                len(unexpected),
+            )
+        model.eval()
+        return model

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

@@ -0,0 +1,489 @@
+import dataclasses
+import glob
+import os
+import re
+from collections.abc import Generator, Iterable
+from contextlib import nullcontext
+from typing import cast
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed import init_device_mesh
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
+
+from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig
+from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+    QwenImageEditPipelineConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_tp_group,
+)
+from sglang.multimodal_gen.runtime.distributed.group_coordinator import GroupCoordinator
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    patch_tensor_parallel_group,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.fsdp_load import shard_model
+from sglang.multimodal_gen.runtime.loader.utils import (
+    set_default_torch_dtype,
+    skip_init_modules,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    filter_duplicate_safetensors_files,
+    filter_files_not_needed_for_inference,
+    pt_weights_iterator,
+    safetensors_weights_iterator,
+)
+from sglang.multimodal_gen.runtime.models.encoders.base import (
+    finalize_encoder_folding,
+    get_folding_tp_group,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_config,
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.precision import precision_to_dtype
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+from sglang.srt.environ import envs
+
+logger = init_logger(__name__)
+
+
+class TextEncoderLoader(ComponentLoader):
+    """Loader for text encoders."""
+
+    component_names = ["text_encoder"]
+    expected_library = "transformers"
+
+    @dataclasses.dataclass
+    class Source:
+        """A source for weights."""
+
+        model_or_path: str
+        """The model ID or path."""
+
+        prefix: str = ""
+        """A prefix to prepend to all weights."""
+
+        fall_back_to_pt: bool = True
+        """Whether .pt weights can be used."""
+
+        allow_patterns_overrides: list[str] | None = None
+        """If defined, weights will load exclusively using these patterns."""
+
+    def should_offload(self, server_args, model_config: ModelConfig | None = None):
+        should_offload = server_args.text_encoder_cpu_offload
+        if not should_offload:
+            return False
+        # _fsdp_shard_conditions is in arch_config, not directly on model_config
+        arch_config = (
+            getattr(model_config, "arch_config", model_config) if model_config else None
+        )
+        fsdp_shard_conditions = (
+            getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
+        )
+        use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
+        return use_cpu_offload
+
+    def customized_load_kwargs_for_component(
+        self, server_args: ServerArgs, component_name: str
+    ) -> dict[str, bool]:
+        if ComponentLoader._is_component_set_as_layerwise_load(
+            server_args, component_name
+        ):
+            logger.info(
+                "Loading %s on CPU first because it is selected for layerwise offload",
+                component_name,
+            )
+            return {"cpu_offload_flag": True}
+        return {}
+
+    def load_native(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        transformers_or_diffusers: str,
+        component_name: str | None = None,
+    ):
+        if transformers_or_diffusers != "transformers":
+            return super().load_native(
+                component_model_path,
+                server_args,
+                transformers_or_diffusers,
+                component_name,
+            )
+
+        encoder_idx = (
+            self._extract_encoder_index(component_name or "text_encoder_2")
+            if component_name
+            else 1 if component_model_path.rstrip("/").endswith("text_encoder_2") else 0
+        )
+        encoder_dtype = server_args.pipeline_config.text_encoder_precisions[encoder_idx]
+        dtype = precision_to_dtype(
+            encoder_dtype,
+            f"text_encoder_precisions[{encoder_idx}]",
+        )
+        transformers_model_class = self._resolve_transformers_text_encoder_class(
+            component_model_path, server_args
+        )
+        return transformers_model_class.from_pretrained(
+            component_model_path,
+            trust_remote_code=server_args.trust_remote_code,
+            revision=server_args.revision,
+            torch_dtype=dtype,
+        )
+
+    @staticmethod
+    def _resolve_transformers_text_encoder_class(component_model_path, server_args):
+        """Resolve the concrete transformers class for a text encoder.
+
+        AutoModel maps encoder-decoder model types (e.g. T5/UMT5) to full
+        seq2seq classes, whose forward expects decoder inputs and raises when
+        the module is used purely as a text encoder. For such checkpoints,
+        prefer the encoder-only class from the config architectures or map the
+        full seq2seq architecture to its encoder-only counterpart. Encoders that
+        are not encoder-decoder keep using AutoModel unchanged.
+        """
+        import transformers
+        from transformers import AutoConfig, AutoModel
+
+        try:
+            config = AutoConfig.from_pretrained(
+                component_model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+            )
+        except Exception:
+            return AutoModel
+        if getattr(config, "is_encoder_decoder", False):
+            encoder_only_map = {
+                "T5Model": "T5EncoderModel",
+                "T5ForConditionalGeneration": "T5EncoderModel",
+                "UMT5Model": "UMT5EncoderModel",
+                "UMT5ForConditionalGeneration": "UMT5EncoderModel",
+                "MT5Model": "MT5EncoderModel",
+                "MT5ForConditionalGeneration": "MT5EncoderModel",
+            }
+            for arch in getattr(config, "architectures", None) or []:
+                encoder_arch = encoder_only_map.get(arch, arch)
+                transformers_model_class = getattr(transformers, encoder_arch, None)
+                if isinstance(transformers_model_class, type):
+                    return transformers_model_class
+        return AutoModel
+
+    def _prepare_weights(
+        self,
+        model_name_or_path: str,
+        fall_back_to_pt: bool,
+        allow_patterns_overrides: list[str] | None,
+    ) -> tuple[str, list[str], bool]:
+        """Prepare weights for the model.
+
+        If the model is not local, it will be downloaded."""
+        # model_name_or_path = (self._maybe_download_from_modelscope(
+        #     model_name_or_path, revision) or model_name_or_path)
+
+        is_local = os.path.isdir(model_name_or_path)
+        assert is_local, "Model path must be a local directory"
+
+        use_safetensors = False
+        index_file = SAFE_WEIGHTS_INDEX_NAME
+        allow_patterns = ["*.safetensors", "*.bin"]
+
+        if fall_back_to_pt:
+            allow_patterns += ["*.pt"]
+
+        if allow_patterns_overrides is not None:
+            allow_patterns = allow_patterns_overrides
+
+        hf_folder = model_name_or_path
+
+        hf_weights_files: list[str] = []
+        for pattern in allow_patterns:
+            hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
+            if len(hf_weights_files) > 0:
+                if pattern == "*.safetensors":
+                    use_safetensors = True
+                break
+
+        if use_safetensors:
+            hf_weights_files = filter_duplicate_safetensors_files(
+                hf_weights_files, hf_folder, index_file
+            )
+        else:
+            hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
+
+        if len(hf_weights_files) == 0:
+            raise RuntimeError(
+                f"Cannot find any model weights with `{model_name_or_path}`"
+            )
+
+        # Sort weight files when SGLANG_SORT_WEIGHT_FILES >= 0 (default).
+        # Staggering is not applicable to text-encoder loading (no TP split).
+        if envs.SGLANG_SORT_WEIGHT_FILES.get() >= 0:
+            hf_weights_files.sort()
+
+        return hf_folder, hf_weights_files, use_safetensors
+
+    def _get_weights_iterator(
+        self,
+        source: "Source",
+        to_cpu: bool,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        """get an iterator for the model weights based on the load format."""
+        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
+            source.model_or_path,
+            source.fall_back_to_pt,
+            source.allow_patterns_overrides,
+        )
+        if use_safetensors:
+            weights_iterator = safetensors_weights_iterator(
+                hf_weights_files,
+                to_cpu=to_cpu,
+            )
+        else:
+            weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
+
+        # apply the prefix.
+        return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
+
+    def _get_all_weights(
+        self,
+        model: nn.Module,
+        model_path: str,
+        to_cpu: bool,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        primary_weights = TextEncoderLoader.Source(
+            model_path,
+            prefix="",
+            fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
+            allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
+        )
+        yield from self._get_weights_iterator(
+            primary_weights,
+            to_cpu,
+        )
+
+        secondary_weights = cast(
+            Iterable[TextEncoderLoader.Source],
+            getattr(model, "secondary_weights", ()),
+        )
+        for source in secondary_weights:
+            yield from self._get_weights_iterator(
+                source,
+                to_cpu,
+            )
+
+    def load_customized(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+        cpu_offload_flag: bool | None = None,
+    ):
+        """Load the text encoders based on the model path, and inference args."""
+        diffusers_pretrained_config = get_config(
+            component_model_path, trust_remote_code=True
+        )
+        model_config = get_diffusers_component_config(
+            component_path=component_model_path
+        )
+
+        # TODO(mick): had to throw an exception for different text-encoder arch
+        encoder_index = self._extract_encoder_index(component_name)
+        assert encoder_index < len(
+            server_args.pipeline_config.text_encoder_configs
+        ) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions)
+
+        encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index]
+        encoder_config.update_model_arch(model_config)
+
+        if encoder_index == 0:
+            for key, value in diffusers_pretrained_config.__dict__.items():
+                setattr(encoder_config.arch_config, key, value)
+        post_diffusers_config_update = getattr(
+            encoder_config, "post_diffusers_config_update", None
+        )
+        if post_diffusers_config_update is not None:
+            post_diffusers_config_update()
+        # Real dims are populated now; keep the proposed fold group only if this
+        # encoder is actually wide enough to benefit at its real size.
+        finalize_encoder_folding(encoder_config)
+        encoder_dtype = server_args.pipeline_config.text_encoder_precisions[
+            encoder_index
+        ]
+        # TODO(will): add support for other dtypes
+        return self.load_model(
+            component_model_path,
+            encoder_config,
+            server_args,
+            encoder_dtype,
+            cpu_offload_flag=cpu_offload_flag,
+        )
+
+    @staticmethod
+    def _extract_encoder_index(component_name: str) -> int:
+        """
+        Map text encoder component names to zero-based indices.
+
+        Examples:
+        - text_encoder -> 0
+        - text_encoder_2 -> 1
+        - text_encoder_3 -> 2
+        """
+        match = re.search(r"_(\d+)$", component_name)
+        if match is None:
+            return 0
+
+        suffix_num = int(match.group(1))
+        if suffix_num <= 0:
+            raise ValueError(
+                f"Invalid text encoder component name '{component_name}': "
+                "numeric suffix must be >= 1."
+            )
+        return suffix_num - 1
+
+    def load_model(
+        self,
+        model_path: str,
+        model_config: EncoderConfig,
+        server_args: ServerArgs,
+        dtype: str = "fp16",
+        cpu_offload_flag: bool | None = None,
+    ):
+        # Determine CPU offload behavior and target device
+
+        local_torch_device = get_local_torch_device()
+
+        if not current_platform.is_cpu():
+            fsdp_cpu_offload = self.should_offload(server_args, model_config)
+            should_offload = (
+                cpu_offload_flag if cpu_offload_flag is not None else fsdp_cpu_offload
+            )
+        else:
+            fsdp_cpu_offload = False
+            should_offload = False
+
+        if (
+            getattr(
+                model_config.arch_config, "requires_gpu_resident_text_encoder", False
+            )
+            and should_offload
+        ):
+            logger.warning(
+                "Keeping bitsandbytes 4-bit text encoder GPU-resident; CUDA "
+                "weights and quant states are required for this checkpoint."
+            )
+            should_offload = False
+
+        if should_offload and not current_platform.is_mps():
+            model_device = torch.device("cpu")
+        else:
+            model_device = local_torch_device
+
+        # Parallel folding: build + shard the encoder over the folding group (the
+        # idle DiT replica during the encoding stage) instead of the default TP
+        # group, so every encoder folds without threading the group through each layer.
+        fold_ctx = nullcontext()
+        if getattr(model_config, "parallel_folding_mode", None) is not None:
+            folding_group = get_folding_tp_group(model_config)
+            if (
+                isinstance(folding_group, GroupCoordinator)
+                and folding_group is not get_tp_group()
+            ):
+                fold_ctx = patch_tensor_parallel_group(folding_group)
+
+        # patch tp group with folding group to achieve TP among folding group
+        with fold_ctx, set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
+            with model_device, skip_init_modules():
+                architectures = getattr(model_config, "architectures", [])
+                model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+                enable_image_understanding = (
+                    True
+                    if isinstance(
+                        server_args.pipeline_config, QwenImageEditPipelineConfig
+                    )
+                    else False
+                )
+                model_config.enable_image_understanding = enable_image_understanding
+                model = model_cls(model_config)
+
+            weights_to_load = {name for name, _ in model.named_parameters()}
+            loaded_weights = model.load_weights(
+                self._get_all_weights(
+                    model,
+                    model_path,
+                    to_cpu=should_offload,
+                )
+            )
+
+            if should_offload:
+                # Disable FSDP for MPS as it's not compatible
+                if current_platform.is_mps():
+                    logger.info(
+                        "Disabling FSDP sharding for MPS platform as it's not compatible"
+                    )
+                    model = model.to(local_torch_device)
+                elif fsdp_cpu_offload:
+                    mesh = init_device_mesh(
+                        current_platform.device_type,
+                        mesh_shape=(1, dist.get_world_size()),
+                        mesh_dim_names=("offload", "replicate"),
+                    )
+                    shard_model(
+                        model,
+                        cpu_offload=True,
+                        reshard_after_forward=True,
+                        mesh=mesh["offload"],
+                        fsdp_shard_conditions=model_config.arch_config._fsdp_shard_conditions
+                        or getattr(model, "_fsdp_shard_conditions", None),
+                        pin_cpu_memory=server_args.pin_cpu_memory,
+                    )
+                else:
+                    model = model.to("cpu")
+            else:
+                model = model.to(local_torch_device)
+            # We only enable strict check for non-quantized models
+            # that have loaded weights tracking currently.
+            # if loaded_weights is not None:
+            weights_not_loaded = weights_to_load - loaded_weights
+            if weights_not_loaded:
+                # NOTE:
+                # If we silently continue with uninitialized weights, the text encoder can
+                # produce NaNs/garbage embeddings that later fail stage verification in a
+                # hard-to-debug way (e.g., `prompt_embeds` fails the NaN check).
+                #
+                # We allow a small set of known-optional parameters to be missing, but
+                # default to strict behavior for the rest.
+                allowed_missing_patterns = (
+                    getattr(model, "_allowed_missing_weights_patterns", []) or []
+                )
+                unexpected_missing = {
+                    n
+                    for n in weights_not_loaded
+                    if not any(pat in n for pat in allowed_missing_patterns)
+                }
+                if unexpected_missing:
+                    raise ValueError(
+                        "Following text encoder weights were not initialized from checkpoint: "
+                        f"{sorted(unexpected_missing)}. "
+                        "This usually indicates a checkpoint/model-arch mismatch or a broken "
+                        "weight-name mapping. If these are truly optional, set "
+                        "`model._allowed_missing_weights_patterns` to whitelist patterns."
+                    )
+                logger.warning(
+                    "Following (allowed) text encoder weights were not initialized from "
+                    "checkpoint: %s (allowed patterns: %s)",
+                    sorted(weights_not_loaded),
+                    allowed_missing_patterns,
+                )
+
+        return model

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

@@ -0,0 +1,199 @@
+import copy
+import logging
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
+from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
+    resolve_transformer_quant_load_spec,
+    resolve_transformer_safetensors_to_load,
+)
+from sglang.multimodal_gen.runtime.loader.utils import _normalize_component_type
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger
+from sglang.srt.utils import is_npu
+
+_is_npu = is_npu()
+
+logger = init_logger(__name__)
+
+
+def _server_args_for_transformer_component(
+    server_args: ServerArgs, component_name: str
+) -> ServerArgs:
+    """Mask global quantized override flags for secondary transformer components."""
+    if component_name not in ("transformer_2", "unconditional_transformer"):
+        return server_args
+
+    # Some pipelines have secondary DiT components with their own quantized
+    # weight file. Keep the mapping model-owned and the loader generic.
+    component_weights_paths = getattr(
+        server_args, "component_transformer_weights_paths", {}
+    )
+    component_weights_path = component_weights_paths.get(component_name)
+    if component_weights_path is not None:
+        component_server_args = copy.copy(server_args)
+        component_server_args.transformer_weights_path = component_weights_path
+        component_server_args.nunchaku_config = None
+        logger.info(
+            "Using transformer_weights_path override for %s: %s",
+            component_name,
+            component_weights_path,
+        )
+        return component_server_args
+
+    if (
+        server_args.transformer_weights_path is None
+        and server_args.nunchaku_config is None
+    ):
+        return server_args
+
+    component_server_args = copy.copy(server_args)
+    component_server_args.transformer_weights_path = None
+    component_server_args.nunchaku_config = None
+    logger.info(
+        "Ignoring global transformer_weights_path for %s; keep it on the base "
+        "checkpoint unless a per-component override path is provided.",
+        component_name,
+    )
+    return component_server_args
+
+
+class TransformerLoader(ComponentLoader):
+    """Shared loader for (video/audio) DiT transformers."""
+
+    component_names = [
+        "transformer",
+        "unconditional_transformer",
+        "audio_dit",
+        "video_dit",
+    ]
+    expected_library = "diffusers"
+
+    def should_raise_customized_load_error(
+        self, server_args: ServerArgs, component_name: str
+    ) -> bool:
+        component_server_args = _server_args_for_transformer_component(
+            server_args, component_name
+        )
+        # Don't let a quantized load quietly fall back to the unquantized native
+        # model. That would drop the requested precision and bury the real error.
+        return (
+            component_server_args.transformer_weights_path is not None
+            or component_server_args.quantization is not None
+        )
+
+    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."""
+        component_server_args = _server_args_for_transformer_component(
+            server_args, component_name
+        )
+
+        # 1. hf config
+        config = get_diffusers_component_config(component_path=component_model_path)
+
+        safetensors_list = resolve_transformer_safetensors_to_load(
+            component_server_args, component_model_path
+        )
+
+        # 2. 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", "unconditional_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")
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+
+        quant_spec = resolve_transformer_quant_load_spec(
+            hf_config=config,
+            server_args=component_server_args,
+            safetensors_list=safetensors_list,
+            component_model_path=component_model_path,
+            model_cls=model_cls,
+            cls_name=cls_name,
+        )
+
+        logger.info(
+            "Loading %s from %s safetensors file(s) %s, param_dtype: %s",
+            cls_name,
+            len(safetensors_list),
+            f": {safetensors_list}" if get_log_level() == logging.DEBUG else "",
+            quant_spec.param_dtype,
+        )
+        # prepare init_param
+        init_params: dict[str, Any] = {
+            "config": dit_config,
+            "hf_config": config,
+            "quant_config": quant_spec.runtime_quant_config,
+        }
+        if (
+            init_params["quant_config"] is None
+            and component_server_args.transformer_weights_path is not None
+        ):
+            logger.warning(
+                "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"])
+
+        local_torch_device = get_local_torch_device()
+        weight_load_plan = WeightLoadPlan.for_component(
+            checkpoint_load_device=local_torch_device,
+            needs_device_weight_postprocess=quant_spec.needs_device_weight_postprocess,
+            component_cpu_offload=bool(component_server_args.dit_cpu_offload),
+        )
+
+        # Load the model using FSDP loader
+        model = maybe_load_fsdp_model(
+            model_cls=model_cls,
+            init_params=init_params,
+            weight_dir_list=safetensors_list,
+            device=local_torch_device,
+            hsdp_replicate_dim=server_args.hsdp_replicate_dim,
+            hsdp_shard_dim=server_args.hsdp_shard_dim,
+            cpu_offload=component_server_args.dit_cpu_offload,
+            pin_cpu_memory=component_server_args.pin_cpu_memory,
+            fsdp_inference=component_server_args.use_fsdp_inference,
+            param_dtype=quant_spec.param_dtype,
+            reduce_dtype=torch.float32,
+            output_dtype=None,
+            strict=False,
+            weight_load_plan=weight_load_plan,
+        )
+
+        # post-hooks (e.g., patch scales (nunchaku))
+        for post_load_hook in quant_spec.post_load_hooks:
+            post_load_hook(model)
+
+        # considering the existent of mixed-precision models (e.g., nunchaku)
+        if (
+            next(model.parameters()).dtype != quant_spec.param_dtype
+            and quant_spec.param_dtype
+        ):
+            logger.warning(
+                "Model dtype does not match expected param dtype, %s vs %s",
+                next(model.parameters()).dtype,
+                quant_spec.param_dtype,
+            )
+
+        return model

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

@@ -0,0 +1,223 @@
+import glob
+import json
+import os
+import re
+
+import safetensors
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import (
+    LatentUpsampler,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+UPSAMPLER_CONSTRUCTOR_KEYS = {
+    "in_channels",
+    "mid_channels",
+    "num_blocks_per_stage",
+    "dims",
+    "spatial_upsample",
+    "temporal_upsample",
+    "spatial_scale",
+    "rational_resampler",
+}
+
+_HF_BLOB_URL_RE = re.compile(
+    r"https?://huggingface\.co/([^/]+/[^/]+)/blob/([^/]+)/(.*)"
+)
+_HF_RESOLVE_URL_RE = re.compile(
+    r"https?://huggingface\.co/([^/]+/[^/]+)/resolve/([^/]+)/(.*)"
+)
+
+
+def _parse_hf_url(path: str):
+    m = _HF_BLOB_URL_RE.match(path) or _HF_RESOLVE_URL_RE.match(path)
+    if m:
+        return m.group(1), m.group(2), m.group(3)
+    return None
+
+
+def _download_hf_file(repo_id: str, filename: str, revision: str = "main") -> str:
+    from huggingface_hub import hf_hub_download
+
+    logger.info("Downloading %s from %s (revision=%s)", filename, repo_id, revision)
+    return hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
+
+
+def _find_safetensors_file(path: str) -> str:
+    """Resolve path to a single safetensors file (local path, directory, HF URL, or HF repo id)."""
+    if os.path.isfile(path) and path.endswith(".safetensors"):
+        return path
+
+    if os.path.isdir(path):
+        files = sorted(glob.glob(os.path.join(path, "*.safetensors")))
+        if len(files) == 1:
+            return files[0]
+        elif len(files) > 1:
+            raise ValueError(
+                f"Found {len(files)} safetensors files in {path}, expected 1"
+            )
+
+    hf = _parse_hf_url(path)
+    if hf:
+        repo_id, revision, filename = hf
+        return _download_hf_file(repo_id, filename, revision)
+
+    try:
+        maybe_downloaded = maybe_download_model(path)
+        if os.path.isdir(maybe_downloaded):
+            files = sorted(glob.glob(os.path.join(maybe_downloaded, "*.safetensors")))
+            if len(files) == 1:
+                return files[0]
+            elif len(files) > 1:
+                raise ValueError(
+                    f"Found {len(files)} safetensors files in {maybe_downloaded}, expected 1"
+                )
+    except Exception:
+        pass
+
+    raise FileNotFoundError(
+        f"No safetensors file found at {path}. "
+        "Provide a local .safetensors file, a directory containing one, "
+        "a HuggingFace URL (https://huggingface.co/<repo>/blob/main/<path>), "
+        "or a HuggingFace repo id."
+    )
+
+
+def _normalize_config(raw: dict) -> dict:
+    """Map diffusers / original-repo config fields to LatentUpsampler kwargs."""
+    config = {k: v for k, v in raw.items() if k in UPSAMPLER_CONSTRUCTOR_KEYS}
+
+    # diffusers uses rational_spatial_scale instead of rational_resampler + spatial_scale
+    if "rational_spatial_scale" in raw and "rational_resampler" not in config:
+        config["rational_resampler"] = True
+        config.setdefault("spatial_scale", raw["rational_spatial_scale"])
+
+    return config
+
+
+def _infer_config_from_state_dict(state_dict: dict[str, torch.Tensor]) -> dict:
+    """Infer LatentUpsampler kwargs from weight shapes and key names.
+
+    Works even when no config.json or safetensors metadata is available.
+    """
+    config: dict = {}
+
+    w = state_dict.get("initial_conv.weight")
+    if w is not None:
+        config["mid_channels"] = w.shape[0]
+        config["in_channels"] = w.shape[1]
+        config["dims"] = 3 if w.ndim == 5 else 2
+
+    num_blocks = sum(
+        1
+        for k in state_dict
+        if k.startswith("res_blocks.") and k.endswith(".conv1.weight")
+    )
+    if num_blocks > 0:
+        config["num_blocks_per_stage"] = num_blocks
+
+    # Detect upsampler type from key patterns
+    has_rational = any(k.startswith("upsampler.blur_down.") for k in state_dict)
+    if has_rational:
+        config["rational_resampler"] = True
+        config["spatial_upsample"] = True
+        config["temporal_upsample"] = False
+        config["spatial_scale"] = 2.0
+    else:
+        up_w = state_dict.get("upsampler.0.weight")
+        if up_w is not None and up_w.ndim == 5:
+            ratio = up_w.shape[0] // up_w.shape[1]
+            if ratio == 8:
+                config["spatial_upsample"] = True
+                config["temporal_upsample"] = True
+            elif ratio == 2:
+                config["spatial_upsample"] = False
+                config["temporal_upsample"] = True
+            else:
+                config["spatial_upsample"] = True
+                config["temporal_upsample"] = False
+        else:
+            config["spatial_upsample"] = True
+            config["temporal_upsample"] = False
+
+    return config
+
+
+def _load_config(
+    safetensors_path: str,
+    original_path: str,
+    state_dict: dict[str, torch.Tensor],
+) -> dict:
+    """Load upsampler config with fallback chain:
+    1. safetensors metadata ("config" key) - original LTX-2 repo format
+    2. sibling config.json - diffusers format
+    3. config.json from HF (if original_path was a URL)
+    4. infer from state dict shapes (always works)
+    """
+    with safetensors.safe_open(safetensors_path, framework="pt") as f:
+        meta = f.metadata()
+        if meta and "config" in meta:
+            logger.info("Using config from safetensors metadata")
+            return _normalize_config(json.loads(meta["config"]))
+
+    config_json_path = os.path.join(os.path.dirname(safetensors_path), "config.json")
+    if os.path.isfile(config_json_path):
+        with open(config_json_path) as fp:
+            logger.info("Using config from sibling config.json")
+            return _normalize_config(json.load(fp))
+
+    hf = _parse_hf_url(original_path)
+    if hf:
+        repo_id, revision, filename = hf
+        config_filename = os.path.dirname(filename) + "/config.json"
+        try:
+            local = _download_hf_file(repo_id, config_filename, revision)
+            with open(local) as fp:
+                logger.info("Using config from HF config.json")
+                return _normalize_config(json.load(fp))
+        except Exception:
+            pass
+
+    logger.info("No explicit config found, inferring from state dict")
+    return _infer_config_from_state_dict(state_dict)
+
+
+class UpsamplerLoader(ComponentLoader):
+    component_names = ["spatial_upsampler"]
+    expected_library = "diffusers"
+
+    def should_offload(self, server_args: ServerArgs, model_config=None):
+        return server_args.vae_cpu_offload
+
+    def load_customized(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+    ):
+        safetensors_path = _find_safetensors_file(component_model_path)
+        state_dict = safetensors_load_file(safetensors_path)
+        config = _load_config(safetensors_path, component_model_path, state_dict)
+
+        logger.info("Loading LatentUpsampler with config: %s", config)
+
+        should_offload = self.should_offload(server_args)
+        target_device = self.target_device(should_offload)
+
+        with torch.device("meta"):
+            model = LatentUpsampler(**config)
+
+        model.load_state_dict(state_dict, assign=True)
+        model = model.to(device=target_device, dtype=torch.bfloat16).eval()
+
+        logger.info("Loaded LatentUpsampler to %s", target_device)
+        return model

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

@@ -0,0 +1,211 @@
+import importlib.util
+import os
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+    QwenImagePipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _list_safetensors_files,
+    set_default_torch_dtype,
+    skip_init_modules,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+VAE_CHANNELS_LAST_3D_ENV = "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D"
+
+
+def _backfill_ltx2_audio_vae_latent_stats(
+    loaded: dict[str, torch.Tensor], component_name: str
+) -> None:
+    if component_name != "audio_vae":
+        return
+    mean_key = "per_channel_statistics.mean-of-means"
+    std_key = "per_channel_statistics.std-of-means"
+    if "latents_mean" not in loaded and mean_key in loaded:
+        loaded["latents_mean"] = loaded[mean_key]
+    if "latents_std" not in loaded and std_key in loaded:
+        loaded["latents_std"] = loaded[std_key]
+
+
+def _convert_conv3d_weights_to_channels_last_3d(module: nn.Module) -> int:
+    """
+    Convert Conv3d weights to channels_last_3d (NDHWC) memory format.
+    Returns the number of Conv3d modules converted.
+    """
+    if not hasattr(torch, "channels_last_3d"):
+        return 0
+    num_converted = 0
+    for m in module.modules():
+        if isinstance(m, nn.Conv3d):
+            try:
+                m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d)
+                num_converted += 1
+            except Exception:
+                # Best-effort; skip unsupported cases.
+                continue
+    return num_converted
+
+
+def _should_use_channels_last_3d(
+    server_args: ServerArgs | None, component_name: str
+) -> bool:
+    if component_name not in (
+        "vae",
+        "video_vae",
+    ) or not (current_platform.is_cuda() or current_platform.is_rocm()):
+        return False
+
+    override = os.getenv(VAE_CHANNELS_LAST_3D_ENV)
+    if override is not None and override.strip().lower() != "auto":
+        return get_bool_env_var(VAE_CHANNELS_LAST_3D_ENV)
+
+    if server_args is None:
+        return False
+
+    pipeline_config = server_args.pipeline_config
+    if isinstance(pipeline_config, QwenImagePipelineConfig):
+        return True
+    if (
+        isinstance(pipeline_config, (WanT2V480PConfig, LTX2PipelineConfig))
+        and server_args.num_gpus == 1
+    ):
+        return True
+    return False
+
+
+class VAELoader(ComponentLoader):
+    """Shared loader for (video/audio) VAE modules."""
+
+    component_names = ["vae", "audio_vae", "video_vae"]
+    expected_library = "diffusers"
+
+    def should_offload(
+        self, server_args: ServerArgs, model_config: ModelConfig | None = None
+    ):
+        return server_args.vae_cpu_offload
+
+    def load_customized(
+        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(component_path=component_model_path)
+        class_name = config.pop("_class_name", None)
+        assert (
+            class_name is not None
+        ), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+
+        server_args.model_paths[component_name] = component_model_path
+
+        if component_name in ("vae", "video_vae"):
+            pipeline_vae_config_attr = "vae_config"
+            pipeline_vae_precision = "vae_precision"
+        elif component_name in ("audio_vae",):
+            pipeline_vae_config_attr = "audio_vae_config"
+            pipeline_vae_precision = "audio_vae_precision"
+        else:
+            raise ValueError(
+                f"Unsupported module name for VAE loader: {component_name}"
+            )
+        vae_config = getattr(server_args.pipeline_config, pipeline_vae_config_attr)
+        vae_precision = getattr(server_args.pipeline_config, pipeline_vae_precision)
+        resolved_vae_dtype = resolve_component_precision(server_args, component_name)
+        vae_dtype = (
+            resolved_vae_dtype
+            if resolved_vae_dtype is not None
+            else PRECISION_TO_TYPE[vae_precision]
+        )
+        vae_config.update_model_arch(config)
+        if hasattr(vae_config, "post_init"):
+            # NOTE: some post init logics are only available after updated with config
+            vae_config.post_init()
+
+        should_offload = self.should_offload(server_args)
+        target_device = self.target_device(should_offload)
+
+        # Check for auto_map first (custom VAE classes)
+        auto_map = config.get("auto_map", {})
+        auto_model_map = auto_map.get("AutoModel")
+        if auto_model_map:
+            module_path, cls_name = auto_model_map.rsplit(".", 1)
+            custom_module_file = os.path.join(component_model_path, f"{module_path}.py")
+            spec = importlib.util.spec_from_file_location("_custom", custom_module_file)
+            custom_module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(custom_module)
+            vae_cls = getattr(custom_module, cls_name)
+            with set_default_torch_dtype(vae_dtype):
+                vae = vae_cls.from_pretrained(
+                    component_model_path,
+                    revision=server_args.revision,
+                    trust_remote_code=server_args.trust_remote_code,
+                )
+            vae = vae.to(device=target_device, dtype=vae_dtype)
+            if _should_use_channels_last_3d(server_args, component_name):
+                n = _convert_conv3d_weights_to_channels_last_3d(vae)
+                if n > 0:
+                    logger.info(
+                        "VAE: converted %d Conv3d weights to channels_last_3d", n
+                    )
+            vae = current_platform.optimize_vae(vae)
+            return vae
+
+        # Load from ModelRegistry (standard VAE classes)
+        with (
+            set_default_torch_dtype(vae_dtype),
+            skip_init_modules(),
+        ):
+            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            vae = vae_cls(vae_config).to(target_device)
+
+        safetensors_list = _list_safetensors_files(component_model_path)
+        safetensors_list = server_args.pipeline_config.select_vae_weight_files(
+            safetensors_list=safetensors_list,
+            component_model_path=component_model_path,
+            component_name=component_name,
+            vae_precision=vae_precision,
+        )
+
+        assert (
+            len(safetensors_list) >= 1
+        ), f"Found no safetensors files in {component_model_path}"
+        loaded = {}
+        for sf_path in safetensors_list:
+            loaded.update(safetensors_load_file(sf_path))
+        _backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
+        vae.load_state_dict(loaded, strict=False)
+
+        state_keys = set(vae.state_dict().keys())
+        loaded_keys = set(loaded.keys())
+        missing_keys = sorted(state_keys - loaded_keys)
+        unexpected_keys = sorted(loaded_keys - state_keys)
+        if missing_keys:
+            logger.warning("VAE missing keys: %s", missing_keys)
+        if unexpected_keys:
+            logger.warning("VAE unexpected keys: %s", unexpected_keys)
+
+        if _should_use_channels_last_3d(server_args, component_name):
+            n = _convert_conv3d_weights_to_channels_last_3d(vae)
+            if n > 0:
+                logger.info("VAE: converted %d Conv3d weights to channels_last_3d", n)
+
+        vae = current_platform.optimize_vae(vae)
+        return vae

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

@@ -0,0 +1,72 @@
+import logging
+from typing import Any
+
+import requests
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import get_hf_config
+
+logger = logging.getLogger(__name__)
+
+
+class VisionLanguageEncoderLoader(ComponentLoader):
+    """Loader for vision language encoder (typically Causal LM or Vision2Seq)."""
+
+    component_names = ["vision_language_encoder"]
+    expected_library = "transformers"
+
+    def load_customized(
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        transformers_or_diffusers: str = "vision_language_encoder",
+    ) -> Any:
+        if transformers_or_diffusers == "vision_language_encoder":
+
+            if server_args.srt_encoder_url is not None:
+                health_url = server_args.srt_encoder_url.rstrip("/") + "/health"
+                try:
+                    logger.info(f"Checking AR encoder server health at: {health_url}")
+                    response = requests.get(
+                        health_url, timeout=server_args.srt_encoder_connect_timeout
+                    )
+
+                    if response.status_code != 200:
+                        error_msg = (
+                            f"AR encoder server returned unhealthy status code: {response.status_code}. "
+                            f"Please ensure the server at {server_args.srt_encoder_url} is fully initialized and compatible."
+                        )
+                        logger.error(error_msg)
+                        raise RuntimeError(error_msg)
+                    logger.info("Successfully connected to AR encoder server.")
+                except requests.RequestException as e:
+                    error_msg = (
+                        f"Failed to reach AR encoder server at {server_args.srt_encoder_url}. "
+                        f"Error: {e}."
+                    )
+                    logger.error(error_msg)
+                    raise RuntimeError(error_msg) from e
+                return server_args.srt_encoder_url
+
+            from transformers import GlmImageForConditionalGeneration
+
+            config = get_hf_config(
+                component_model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+            )
+            model = GlmImageForConditionalGeneration.from_pretrained(
+                component_model_path,
+                config=config,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+            ).to(get_local_torch_device())
+            return model
+        else:
+            raise ValueError(
+                f"Unsupported library for VisionLanguageEncoder: {transformers_or_diffusers}"
+            )

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

@@ -0,0 +1,90 @@
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _list_safetensors_files,
+    set_default_torch_dtype,
+    skip_init_modules,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class VocoderLoader(ComponentLoader):
+    component_names = ["vocoder"]
+    expected_library = "diffusers"
+
+    def should_offload(
+        self, server_args: ServerArgs, model_config: ModelConfig | None = None
+    ):
+        return server_args.vae_cpu_offload
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ):
+        config = get_diffusers_component_config(component_path=component_model_path)
+        class_name = config.pop("_class_name", None) or self.component_architecture
+        assert (
+            class_name is not None
+        ), "Vocoder class name must be available from component config or pipeline config."
+
+        server_args.model_paths[component_name] = component_model_path
+
+        from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import (
+            LTXVocoderConfig,
+        )
+
+        vocoder_config = LTXVocoderConfig()
+        vocoder_config.update_model_arch(config)
+
+        resolved_vocoder_dtype = resolve_component_precision(server_args, "vocoder")
+        vocoder_dtype = (
+            resolved_vocoder_dtype
+            if resolved_vocoder_dtype is not None
+            else PRECISION_TO_TYPE["fp32"]
+        )
+
+        should_offload = self.should_offload(server_args)
+        target_device = self.target_device(should_offload)
+
+        with set_default_torch_dtype(vocoder_dtype), skip_init_modules():
+            vocoder_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            vocoder = vocoder_cls(vocoder_config).to(target_device)
+
+        safetensors_list = _list_safetensors_files(component_model_path)
+        assert (
+            len(safetensors_list) == 1
+        ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        incompatible = vocoder.load_state_dict(loaded, strict=False)
+        missing_keys = []
+        unexpected_keys = []
+        try:
+            missing_keys = incompatible.missing_keys
+            unexpected_keys = incompatible.unexpected_keys
+        except AttributeError:
+            # Best-effort fallback in case older torch returns a tuple-like.
+            try:
+                missing_keys = incompatible[0]
+                unexpected_keys = incompatible[1]
+            except Exception:
+                pass
+
+        if missing_keys or unexpected_keys:
+            logger.warning(
+                "Loaded vocoder with missing_keys=%d unexpected_keys=%d",
+                len(missing_keys),
+                len(unexpected_keys),
+            )
+        return vocoder

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -0,0 +1,769 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from torchtune
+# Copyright 2024 The TorchTune Authors.
+# Copyright 2025 The sglang-diffusion Authors.
+
+from collections import Counter, defaultdict
+from collections.abc import Callable, Generator
+from itertools import chain
+from typing import Any
+
+import torch
+from torch import nn
+from torch.distributed import DeviceMesh, init_device_mesh
+from torch.distributed._tensor import distribute_tensor
+from torch.distributed.fsdp import (
+    CPUOffloadPolicy,
+    FSDPModule,
+    MixedPrecisionPolicy,
+    fully_shard,
+)
+from torch.nn.modules.module import _IncompatibleKeys
+
+from sglang.multimodal_gen.configs.models.fsdp import is_module_list_entry_in
+from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
+from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
+    attach_bitsandbytes_4bit_quant_states,
+    build_bitsandbytes_4bit_quant_states,
+    split_bitsandbytes_4bit_state,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    get_param_names_mapping,
+    hf_to_custom_state_dict,
+    set_default_torch_dtype,
+)
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    safetensors_weights_iterator,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import set_mixed_precision_policy
+from sglang.srt.utils import is_npu
+
+_is_npu = is_npu()
+
+logger = init_logger(__name__)
+
+_QUANTIZED_DTYPES = (
+    torch.uint8,
+    torch.float8_e4m3fn,
+    torch.float8_e5m2,
+    torch.int8,
+)
+_DTYPE_MISMATCH_EXAMPLE_LIMIT = 3
+
+
+def _is_bitsandbytes_quant_config(quant_config: Any | None) -> bool:
+    if quant_config is None:
+        return False
+    quant_name_getter = getattr(type(quant_config), "get_name", None)
+    return bool(callable(quant_name_getter) and quant_name_getter() == "bitsandbytes")
+
+
+def _format_dtype_mismatch_summary(
+    mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]],
+    mismatch_examples: dict[tuple[torch.dtype, torch.dtype], list[str]],
+) -> str:
+    parts: list[str] = []
+    for (checkpoint_dtype, target_dtype), count in mismatch_counts.items():
+        examples = mismatch_examples[(checkpoint_dtype, target_dtype)]
+        part = f"{checkpoint_dtype}->{target_dtype} x{count}"
+        if examples:
+            part += f" (e.g. {', '.join(examples)})"
+        parts.append(part)
+    return "; ".join(parts)
+
+
+def _make_param_like(
+    actual_param: torch.nn.Parameter, tensor: torch.Tensor
+) -> torch.nn.Parameter:
+    cls = actual_param.__class__
+    # nn.Parameter defaults to requires_grad=True, which is illegal for non-floating/complex dtypes (e.g., int8/FP8
+    # quantized weights).
+    try:
+        new_param = cls.__new__(cls, tensor, requires_grad=False)
+    except TypeError:
+        try:
+            new_param = cls.__new__(cls, tensor)
+        except TypeError:
+            new_param = nn.Parameter(tensor, requires_grad=False)
+    new_param.__dict__.update(actual_param.__dict__)
+    new_param.requires_grad = False
+    return new_param
+
+
+def _get_param_for_weight_loading(
+    model: torch.nn.Module,
+    param_dict: dict[str, torch.nn.Parameter],
+    param_name: str,
+) -> torch.nn.Parameter | None:
+    actual_param = param_dict.get(param_name)
+    if actual_param is not None and getattr(actual_param, "weight_loader", None):
+        return actual_param
+
+    pre_fsdp_weight_loader_params = getattr(model, "_pre_fsdp_weight_loader_params", {})
+    pre_fsdp_param = pre_fsdp_weight_loader_params.get(param_name)
+    if pre_fsdp_param is not None:
+        return pre_fsdp_param
+
+    return actual_param
+
+
+def _make_class_name_shard_condition(class_names: set[str]):
+    def shard_condition(n: str, m: nn.Module) -> bool:
+        return type(m).__name__ in class_names
+
+    return shard_condition
+
+
+def _is_common_numbered_block(n: str, m: nn.Module) -> bool:
+    return is_module_list_entry_in(
+        n,
+        (
+            "blocks",
+            "layers",
+            "double_blocks",
+            "single_blocks",
+            "refiner_blocks",
+            "noise_refiner",
+            "context_refiner",
+            "transformer_blocks",
+            "single_transformer_blocks",
+        ),
+    )
+
+
+def _resolve_fsdp_shard_conditions(
+    model: torch.nn.Module,
+    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] | None,
+) -> tuple[list[Callable[[str, nn.Module], bool]], str]:
+    if fsdp_shard_conditions:
+        return fsdp_shard_conditions, "explicit"
+
+    block_class_names = set(getattr(model, "_repeated_blocks", []) or [])
+    block_class_names.update(getattr(model, "_no_split_modules", []) or [])
+    if block_class_names:
+        return [_make_class_name_shard_condition(block_class_names)], "block-class"
+
+    return [_is_common_numbered_block], "common-numbered-block"
+
+
+def _maybe_dequantize_fp8(
+    full_tensor: torch.Tensor,
+    target_dtype: torch.dtype,
+    target_param_name: str,
+    param_sd: dict[str, torch.Tensor],
+) -> torch.Tensor:
+    """Auto-dequantize an FP8 checkpoint weight when the model parameter expects a higher-precision type.
+
+    Some modules (e.g. AdaLayerNormZero) don't accept quant_config, so their
+    parameters remain in higher precision even when the checkpoint stores FP8
+    weights.  In that case we multiply by the per-tensor weight_scale to
+    recover the original unquantized value.
+    """
+    if not (
+        full_tensor.dtype == torch.float8_e4m3fn and target_dtype != torch.float8_e4m3fn
+    ):
+        return full_tensor
+
+    scale_key = target_param_name.rsplit(".", 1)[0] + ".weight_scale"
+    scale_tensor = param_sd.get(scale_key)
+    if scale_tensor is not None:
+        full_tensor = full_tensor.to(torch.float32) * scale_tensor.float()
+        logger.debug(
+            "Auto-dequantized FP8 weight %s using %s",
+            target_param_name,
+            scale_key,
+        )
+    return full_tensor
+
+
+# TODO(PY): add compile option
+def maybe_load_fsdp_model(
+    model_cls: type[nn.Module],
+    init_params: dict[str, Any],
+    weight_dir_list: list[str],
+    device: torch.device,
+    hsdp_replicate_dim: int,
+    hsdp_shard_dim: int,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
+    cpu_offload: bool = False,
+    fsdp_inference: bool = False,
+    output_dtype: torch.dtype | None = None,
+    pin_cpu_memory: bool = True,
+    strict: bool = True,
+    weight_load_plan: WeightLoadPlan | None = None,
+) -> torch.nn.Module:
+    """Load a model with optional FSDP (Fully Sharded Data Parallel) support.
+
+    Args:
+        param_dtype: Data type for model parameters, also used for:
+            - Model initialization context (set_default_torch_dtype)
+            - FSDP mixed precision policy
+            - Weight loading and casting
+        reduce_dtype: Data type for gradient reduction in FSDP mixed precision.
+        strict: If True, enforce strict state dict loading (all keys must match).
+        weight_load_plan: Optional checkpoint/postprocess device plan for this load.
+    """
+    # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
+    # manually casting the inputs to the model
+
+    # 1. prepare for loading
+    default_torch_dtype = param_dtype if param_dtype else torch.bfloat16
+    mp_policy = MixedPrecisionPolicy(
+        default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
+    )
+
+    set_mixed_precision_policy(
+        param_dtype=default_torch_dtype,
+        reduce_dtype=reduce_dtype,
+        output_dtype=output_dtype,
+        mp_policy=mp_policy,
+    )
+
+    with set_default_torch_dtype(default_torch_dtype), torch.device("meta"):
+        model = model_cls(**init_params)
+
+    # Check if we should use FSDP
+    use_fsdp = fsdp_inference
+
+    # Disable FSDP for MPS as it's not compatible
+    if current_platform.is_mps():
+        use_fsdp = False
+        logger.info("Disabling FSDP for MPS platform as it's not compatible")
+
+    weight_load_plan = weight_load_plan or WeightLoadPlan(checkpoint_load_device=device)
+    defer_cpu_offload = bool(
+        cpu_offload and weight_load_plan.defer_component_cpu_offload
+    )
+    if defer_cpu_offload and use_fsdp:
+        logger.warning(
+            "Ignoring deferred CPU offload for FSDP loading; keeping the existing "
+            "FSDP offload policy."
+        )
+        defer_cpu_offload = False
+    load_cpu_offload = bool(cpu_offload and not defer_cpu_offload)
+    weight_postprocess_device = weight_load_plan.weight_postprocess_device
+    if use_fsdp and weight_postprocess_device is not None:
+        logger.warning("Ignoring weight postprocess device override for FSDP loading.")
+        weight_postprocess_device = None
+
+    if use_fsdp:
+        model._pre_fsdp_weight_loader_params = {
+            n: p
+            for n, p in model.named_parameters()
+            if getattr(p, "weight_loader", None)
+        }
+        world_size = hsdp_replicate_dim * hsdp_shard_dim
+        if not fsdp_inference:
+            hsdp_replicate_dim = world_size
+            hsdp_shard_dim = 1
+
+        device_mesh = init_device_mesh(
+            current_platform.device_type,
+            # (Replicate(), Shard(dim=0))
+            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
+            mesh_dim_names=("replicate", "shard"),
+        )
+        shard_model(
+            model,
+            cpu_offload=load_cpu_offload,
+            reshard_after_forward=True,
+            mp_policy=mp_policy,
+            mesh=device_mesh,
+            fsdp_shard_conditions=getattr(model, "_fsdp_shard_conditions", None),
+            pin_cpu_memory=pin_cpu_memory,
+        )
+
+    param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
+
+    # 2. load model from disk
+    weight_iterator = safetensors_weights_iterator(weight_dir_list)
+    preprocess_loaded_state_dict = getattr(model, "preprocess_loaded_state_dict", None)
+    if preprocess_loaded_state_dict is not None:
+        weight_iterator = preprocess_loaded_state_dict(weight_iterator)
+    bnb_quant_states = None
+    if _is_bitsandbytes_quant_config(init_params.get("quant_config")):
+        normal_weights, raw_quant_state = split_bitsandbytes_4bit_state(weight_iterator)
+        bnb_quant_states = build_bitsandbytes_4bit_quant_states(
+            [name for name, _ in normal_weights],
+            raw_quant_state,
+            device,
+            param_names_mapping_fn,
+        )
+        weight_iterator = iter(normal_weights)
+    load_model_from_full_model_state_dict(
+        model,
+        weight_iterator,
+        weight_load_plan.checkpoint_load_device,
+        param_dtype,
+        strict=strict,
+        cpu_offload=load_cpu_offload,
+        param_names_mapping=param_names_mapping_fn,
+    )
+    if bnb_quant_states:
+        attach_bitsandbytes_4bit_quant_states(
+            dict(model.named_parameters()), bnb_quant_states
+        )
+
+    # 3. postprocessing
+    if weight_postprocess_device is not None:
+        # move to device to perform postprocessing
+        model.to(weight_postprocess_device)
+
+    for _, module in model.named_modules():
+        quant_method = getattr(module, "quant_method", None)
+        if quant_method is not None and hasattr(
+            quant_method, "process_weights_after_loading"
+        ):
+            if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod):
+                # Activate the NZ format for storing weights,
+                # which is a specific optimization for Ascend NPU
+                torch.npu.config.allow_internal_format = True
+            quant_method.process_weights_after_loading(module)
+            if _is_npu:
+                torch.npu.empty_cache()
+    model.post_load_weights()
+
+    for n, p in chain(model.named_parameters(), model.named_buffers()):
+        if p.is_meta:
+            raise RuntimeError(f"Unexpected param or buffer {n} on meta device.")
+        # Avoid unintended computation graph accumulation during inference
+        if isinstance(p, torch.nn.Parameter):
+            p.requires_grad = False
+
+    # 4. deferred cpu offload
+    if defer_cpu_offload:
+        model.to("cpu")
+
+    return model
+
+
+def shard_model(
+    model,
+    *,
+    cpu_offload: bool,
+    reshard_after_forward: bool = True,
+    mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(),  # noqa
+    mesh: DeviceMesh | None = None,
+    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] | None = None,
+    pin_cpu_memory: bool = True,
+) -> None:
+    """
+    Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.
+
+    This method will over the model's named modules from the bottom-up and apply shard modules
+    based on whether they meet any of the criteria from shard_conditions.
+
+    Args:
+        model (TransformerDecoder): Model to shard with FSDP.
+        cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer
+            states to CPU.
+        reshard_after_forward (bool): Whether to reshard parameters and buffers after
+            the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
+            from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
+        mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
+            Default to None.
+        fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
+            which modules to shard with FSDP.
+        pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters.
+
+    """
+    fsdp_shard_conditions, condition_source = _resolve_fsdp_shard_conditions(
+        model, fsdp_shard_conditions
+    )
+    if condition_source != "explicit":
+        logger.warning(
+            "Using %s FSDP shard condition fallback for %s",
+            condition_source,
+            type(model).__name__,
+        )
+
+    fsdp_kwargs = {
+        "reshard_after_forward": reshard_after_forward,
+        "mesh": mesh,
+        "mp_policy": mp_policy,
+    }
+    if cpu_offload:
+        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=pin_cpu_memory)
+
+    # iterating in reverse to start with
+    # lowest-level modules first
+    num_layers_sharded = 0
+    # TODO(will): don't reshard after forward for the last layer to save on the
+    # all-gather that will immediately happen Shard the model with FSDP,
+    for n, m in reversed(list(model.named_modules())):
+        if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):  # type: ignore
+            fully_shard(m, **fsdp_kwargs)
+            num_layers_sharded += 1
+
+    if num_layers_sharded == 0:
+        raise ValueError(
+            f"No layer modules were sharded in {type(model).__name__}. "
+            f"FSDP shard condition source: {condition_source}."
+        )
+
+    # Finally shard the entire model to account for any stragglers
+    fully_shard(model, **fsdp_kwargs)
+    logger.info(
+        "Applied FSDP to %d submodules in %s using %s shard conditions",
+        num_layers_sharded,
+        type(model).__name__,
+        condition_source,
+    )
+
+
+# TODO(mick): need refactor, to move out checkpoint-specific adjustments
+def load_model_from_full_model_state_dict(
+    model: FSDPModule | torch.nn.Module,
+    full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
+    checkpoint_load_device: torch.device,
+    param_dtype: torch.dtype | None,
+    strict: bool = False,
+    cpu_offload: bool = False,
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
+) -> _IncompatibleKeys:
+    """
+    Converting full state dict into a sharded state dict
+    and loading it into FSDP model (if training) or normal huggingface model
+    Args:
+        model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict
+        full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
+        checkpoint_load_device (torch.device): device used to move full state dict tensors
+        param_dtype (torch.dtype): dtype used to move full state dict tensors. If none, respect original dtype from checkpoint
+        strict (bool): flag to check if to load the model in strict mode
+        cpu_offload (bool): flag to check if FSDP offload is enabled
+        param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
+    Returns:
+        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+            * **missing_keys** is a list of str containing the missing keys
+            * **unexpected_keys** is a list of str containing the unexpected keys
+
+    """
+    meta_sd = model.state_dict()
+    param_dict = dict(model.named_parameters())
+
+    # map names from checkpoint to customized names
+    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
+        full_sd_iterator,
+        param_names_mapping,
+        valid_target_names=set(meta_sd.keys()),
+    )  # type: ignore
+
+    is_fsdp_model = isinstance(model, FSDPModule) or any(
+        hasattr(p, "device_mesh") for p in meta_sd.values()
+    )
+
+    # sort parameter names to ensure all ranks process parameters in the same order
+    sorted_param_names = sorted(custom_param_sd.keys())
+
+    sharded_sd = {}
+    skipped_checkpoint_keys: list[str] = []
+    non_quantized_dtype_mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]] = (
+        Counter()
+    )
+    non_quantized_dtype_mismatch_examples: dict[
+        tuple[torch.dtype, torch.dtype], list[str]
+    ] = defaultdict(list)
+    quantized_dtype_mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]] = (
+        Counter()
+    )
+    quantized_dtype_mismatch_examples: dict[
+        tuple[torch.dtype, torch.dtype], list[str]
+    ] = defaultdict(list)
+
+    # shard from loaded state_dict, custom_param_sd -> sharded_sd
+    for target_param_name in sorted_param_names:
+        full_tensor = custom_param_sd[target_param_name]
+        meta_sharded_param = meta_sd.get(target_param_name)
+
+        if meta_sharded_param is None:
+            # For FSDP models, ensure all ranks process parameters consistently
+            if strict or is_fsdp_model:
+                raise ValueError(
+                    f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
+                )
+            else:
+                skipped_checkpoint_keys.append(target_param_name)
+                continue
+
+        # use meta param dtype so quantized params (e.g. FP8) keep their dtype;
+        # for non-quantized models meta dtype equals param_dtype anyway
+        if meta_sharded_param is None:
+            # for nunchaku, some scales are patched later
+            target_dtype = full_tensor.dtype
+        else:
+            target_dtype = meta_sharded_param.dtype
+
+        full_tensor = _maybe_dequantize_fp8(
+            full_tensor, target_dtype, target_param_name, custom_param_sd
+        )
+
+        if full_tensor.dtype != target_dtype:
+            mismatch_key = (full_tensor.dtype, target_dtype)
+            if (
+                full_tensor.dtype in _QUANTIZED_DTYPES
+                or target_dtype in _QUANTIZED_DTYPES
+            ):
+                quantized_dtype_mismatch_counts[mismatch_key] += 1
+                if (
+                    len(quantized_dtype_mismatch_examples[mismatch_key])
+                    < _DTYPE_MISMATCH_EXAMPLE_LIMIT
+                ):
+                    quantized_dtype_mismatch_examples[mismatch_key].append(
+                        target_param_name
+                    )
+            else:
+                non_quantized_dtype_mismatch_counts[mismatch_key] += 1
+                if (
+                    len(non_quantized_dtype_mismatch_examples[mismatch_key])
+                    < _DTYPE_MISMATCH_EXAMPLE_LIMIT
+                ):
+                    non_quantized_dtype_mismatch_examples[mismatch_key].append(
+                        target_param_name
+                    )
+
+        if not hasattr(meta_sharded_param, "device_mesh"):
+            full_tensor = full_tensor.to(
+                device=checkpoint_load_device, dtype=target_dtype
+            )
+            actual_param = _get_param_for_weight_loading(
+                model, param_dict, target_param_name
+            )
+            weight_loader = (
+                getattr(actual_param, "weight_loader", None)
+                if actual_param is not None
+                else None
+            )
+            if weight_loader is not None:
+                assert actual_param is not None
+                sharded_tensor = torch.empty_like(
+                    meta_sharded_param,
+                    device=checkpoint_load_device,
+                    dtype=target_dtype,
+                )
+                # Preserve requires_grad flag to avoid errors with non-floating dtypes
+                requires_grad = getattr(meta_sharded_param, "requires_grad", False)
+                temp_param = _make_param_like(actual_param, sharded_tensor)
+                if not (
+                    sharded_tensor.is_floating_point() or sharded_tensor.is_complex()
+                ):
+                    requires_grad = False
+                temp_param.requires_grad = requires_grad
+                try:
+                    weight_loader(temp_param, full_tensor)
+                except AssertionError as exc:
+                    raise AssertionError(
+                        "Failed to shard/load parameter "
+                        f"{target_param_name}: full_tensor.shape={tuple(full_tensor.shape)}, "
+                        f"meta_sharded_param.shape={tuple(meta_sharded_param.shape)}, "
+                        f"temp_param.shape={tuple(temp_param.shape)}, "
+                        f"param_cls={type(actual_param).__name__}"
+                    ) from exc
+                sharded_tensor = temp_param.data
+            else:
+                # In cases where parts of the model aren't sharded, some parameters will be plain tensors
+                sharded_tensor = full_tensor
+
+            # Important: `cpu_offload` is intended for FSDP-managed parameter movement.
+            # If a parameter is not sharded into a DTensor (i.e., no `device_mesh`), FSDP
+            # will NOT manage it. Offloading it here would leave CPU parameters that
+            # later participate in GPU kernels (e.g., conv/embedding), causing device/dtype
+            # mismatches like "Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type)".
+            #
+            # Therefore:
+            # - For non-FSDP models, keep the historical behavior (allow CPU offload).
+            # - For FSDP models, do NOT offload non-sharded parameters here.
+            if cpu_offload and not is_fsdp_model:
+                sharded_tensor = sharded_tensor.cpu()
+        else:
+            full_tensor = full_tensor.to(
+                device=checkpoint_load_device, dtype=target_dtype
+            )
+            actual_param = _get_param_for_weight_loading(
+                model, param_dict, target_param_name
+            )
+            weight_loader = (
+                getattr(actual_param, "weight_loader", None)
+                if actual_param is not None
+                else None
+            )
+            if weight_loader is not None:
+                assert actual_param is not None
+                tp_sharded_tensor = torch.empty(
+                    tuple(actual_param.shape),
+                    device=checkpoint_load_device,
+                    dtype=target_dtype,
+                )
+                temp_param = _make_param_like(actual_param, tp_sharded_tensor)
+                if not (
+                    tp_sharded_tensor.is_floating_point()
+                    or tp_sharded_tensor.is_complex()
+                ):
+                    temp_param.requires_grad = False
+                try:
+                    weight_loader(temp_param, full_tensor)
+                except AssertionError as exc:
+                    raise AssertionError(
+                        "Failed to TP-shard/load FSDP parameter "
+                        f"{target_param_name}: full_tensor.shape={tuple(full_tensor.shape)}, "
+                        f"meta_sharded_param.shape={tuple(meta_sharded_param.shape)}, "
+                        f"temp_param.shape={tuple(temp_param.shape)}, "
+                        f"param_cls={type(actual_param).__name__}"
+                    ) from exc
+                full_tensor = temp_param.data
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                meta_sharded_param.device_mesh,
+                meta_sharded_param.placements,
+            )
+            if cpu_offload:
+                sharded_tensor = sharded_tensor.to("cpu")
+
+        actual_param = param_dict.get(target_param_name)
+        if actual_param is not None:
+            sharded_sd[target_param_name] = _make_param_like(
+                actual_param, sharded_tensor
+            )
+        else:
+            sharded_sd[target_param_name] = nn.Parameter(
+                sharded_tensor, requires_grad=False
+            )
+
+    model.reverse_param_names_mapping = reverse_param_names_mapping
+
+    if non_quantized_dtype_mismatch_counts:
+        logger.debug(
+            "Casting checkpoint tensors to target dtype during load: %s",
+            _format_dtype_mismatch_summary(
+                non_quantized_dtype_mismatch_counts,
+                non_quantized_dtype_mismatch_examples,
+            ),
+            main_process_only=True,
+            local_main_process_only=True,
+        )
+
+    if quantized_dtype_mismatch_counts:
+        logger.warning(
+            "Dtype mismatches detected for quantized parameters during load: %s",
+            _format_dtype_mismatch_summary(
+                quantized_dtype_mismatch_counts,
+                quantized_dtype_mismatch_examples,
+            ),
+            main_process_only=True,
+            local_main_process_only=True,
+        )
+
+    if skipped_checkpoint_keys:
+        logger.warning(
+            "Checkpoint keys not loaded (no matching model parameter) %s",
+            (
+                skipped_checkpoint_keys[:20]
+                if len(skipped_checkpoint_keys) > 20
+                else skipped_checkpoint_keys
+            ),
+        )
+        if len(skipped_checkpoint_keys) > 20:
+            logger.warning(
+                "... and %d more skipped keys.",
+                len(skipped_checkpoint_keys) - 20,
+            )
+
+    # parameters in nn.Module that doesn't exist in safetensor files
+    unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
+    if unused_keys:
+        logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys)
+
+    # Legacy allowlist for parameter families synthesized after loading.
+    # New formats should declare missing_param_init on the parameter instead.
+    LEGACY_ALLOWED_NEW_PARAM_PATTERNS = [
+        "gate_compress",
+        "wcscales",
+        "wtscale",
+        "input_scale",
+        "weight_scale",
+        "bias",
+        "norm_q",
+        "norm_k",
+        "weight_scale",
+    ]
+    for new_param_name in unused_keys:
+        meta_sharded_param = meta_sd.get(new_param_name)
+        meta_sharded_param_dtype = meta_sharded_param.dtype
+        actual_param = param_dict.get(new_param_name)
+        missing_param_init = (
+            getattr(actual_param, "missing_param_init", None)
+            if actual_param is not None
+            else None
+        )
+
+        if missing_param_init == "error":
+            raise ValueError(
+                f"Required checkpoint parameter '{new_param_name}' was not loaded. "
+                "This usually indicates a checkpoint/model-arch mismatch or a "
+                "broken weight-name mapping."
+            )
+
+        if missing_param_init is None and not any(
+            pattern in new_param_name for pattern in LEGACY_ALLOWED_NEW_PARAM_PATTERNS
+        ):
+            logger.error(
+                "Unsupported new parameter: %s. Allowed legacy patterns: %s",
+                new_param_name,
+                LEGACY_ALLOWED_NEW_PARAM_PATTERNS,
+            )
+            raise ValueError(
+                f"New parameter '{new_param_name}' is not supported. "
+                "Checkpoint-specific synthesized parameters should either match "
+                f"{LEGACY_ALLOWED_NEW_PARAM_PATTERNS} or declare missing_param_init."
+            )
+
+        if missing_param_init == "ones" or any(
+            p in new_param_name
+            for p in (
+                "wcscales",
+                "wtscale",
+                "input_scale",
+                "weight_scale",
+                "norm_q",
+                "norm_k",
+            )
+        ):
+            init_like = torch.ones_like
+        elif missing_param_init == "zeros" or missing_param_init is None:
+            init_like = torch.zeros_like
+        else:
+            raise ValueError(
+                f"Unsupported missing_param_init={missing_param_init!r} for {new_param_name}"
+            )
+
+        if not hasattr(meta_sharded_param, "device_mesh"):
+            sharded_tensor = init_like(
+                meta_sharded_param,
+                device=checkpoint_load_device,
+                dtype=meta_sharded_param_dtype,
+            )
+            if cpu_offload and not is_fsdp_model:
+                sharded_tensor = sharded_tensor.cpu()
+        else:
+            full_tensor = init_like(
+                meta_sharded_param,
+                device=checkpoint_load_device,
+                dtype=meta_sharded_param_dtype,
+            )
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                meta_sharded_param.device_mesh,
+                meta_sharded_param.placements,
+            )
+            if cpu_offload:
+                sharded_tensor = sharded_tensor.cpu()
+        sharded_sd[new_param_name] = nn.Parameter(sharded_tensor)
+
+    # choose `assign=True` since we cannot call `copy_` on meta tensor
+    return model.load_state_dict(sharded_sd, strict=strict, assign=True)

python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py

@@ -0,0 +1,664 @@
+"""Helpers and adapters for transformer quantized checkpoint loading.
+
+This module keeps format-specific loading quirks out of `TransformerLoader`.
+The loader should stay focused on the generic load flow, while special cases
+such as Nunchaku validation, NVFP4 fallback adjustments, and post-load patching
+are handled here behind a small helper/adapter layer.
+"""
+
+import json
+import os
+import re
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    NunchakuConfig,
+    _patch_nunchaku_scales,
+)
+from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    maybe_download_model,
+    snapshot_download,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.precision import resolve_precision
+from sglang.multimodal_gen.runtime.utils.quantization_utils import (
+    build_nvfp4_config_from_safetensors_list,
+    get_metadata_from_safetensors_file,
+    get_quant_config,
+    get_quant_config_from_safetensors_metadata,
+)
+from sglang.srt.layers.quantization import QuantizationConfig
+
+logger = init_logger(__name__)
+
+PostLoadHook = Callable[[nn.Module], None]
+
+_PRECISION_VARIANT_SUFFIX_RE = re.compile(
+    r"^(?P<stem>.+?)(?P<precision>\.(?:fp16|bf16|fp32))(?P<shard>-\d+-of-\d+)?(?P<ext>\.safetensors)$"
+)
+_MIXED_SAFETENSORS_RE = re.compile(r".*-mixed(?:-\d+-of-\d+)?\.safetensors$")
+
+
+def _get_quant_config_name(config: Optional[QuantizationConfig]) -> Optional[str]:
+    if config is None:
+        return None
+    quant_name_getter = getattr(type(config), "get_name", None)
+    return quant_name_getter() if callable(quant_name_getter) else None
+
+
+def _merge_modelopt_fp4_configs(
+    existing_config: Optional[QuantizationConfig],
+    inferred_config: Optional[QuantizationConfig],
+) -> Optional[QuantizationConfig]:
+    """Prefer safetensors-inferred NVFP4 layout over stale config.json ignores.
+
+    Some ModelOpt NVFP4 transformer repos ship a flat `quantization_config` in
+    `config.json`, but its `ignore` list can lag behind the actual checkpoint
+    contents. The safetensors shards are the source of truth for which modules
+    remain BF16 fallbacks, so when we can infer an NVFP4 config from the shards
+    we should use its exclude list while preserving explicit repo-level knobs
+    such as `swap_weight_nibbles`.
+    """
+    if inferred_config is None:
+        return existing_config
+
+    if _get_quant_config_name(inferred_config) != "modelopt_fp4":
+        return existing_config or inferred_config
+
+    if existing_config is None:
+        return inferred_config
+
+    if _get_quant_config_name(existing_config) != "modelopt_fp4":
+        return existing_config
+
+    existing_excludes = getattr(existing_config, "exclude_modules", []) or []
+    inferred_excludes = getattr(inferred_config, "exclude_modules", []) or []
+    if inferred_excludes != existing_excludes:
+        logger.warning(
+            "Overriding ModelOpt NVFP4 exclude_modules from config.json with "
+            "safetensors-inferred layout (%d -> %d entries).",
+            len(existing_excludes),
+            len(inferred_excludes),
+        )
+
+    inferred_config.packed_modules_mapping = getattr(
+        existing_config, "packed_modules_mapping", {}
+    )
+    inferred_config.checkpoint_uses_packed_qkv = getattr(
+        inferred_config, "checkpoint_uses_packed_qkv", False
+    ) or getattr(existing_config, "checkpoint_uses_packed_qkv", False)
+    inferred_config.swap_weight_nibbles = getattr(
+        inferred_config, "swap_weight_nibbles", False
+    ) or getattr(existing_config, "swap_weight_nibbles", False)
+    existing_scale_layout = getattr(
+        existing_config, "checkpoint_weight_scale_layout", "linear"
+    )
+    inferred_scale_layout = getattr(
+        inferred_config, "checkpoint_weight_scale_layout", "linear"
+    )
+    inferred_config.checkpoint_weight_scale_layout = (
+        existing_scale_layout
+        if inferred_scale_layout == "linear" and existing_scale_layout != "linear"
+        else inferred_scale_layout
+    )
+    if getattr(inferred_config, "group_size", None) is None:
+        inferred_config.group_size = getattr(existing_config, "group_size", None)
+
+    return inferred_config
+
+
+@dataclass
+class TransformerQuantLoadSpec:
+    """Resolved loading plan for a transformer checkpoint."""
+
+    safetensors_list: list[str]
+    quant_config: Optional[QuantizationConfig]
+    nunchaku_config: Optional[NunchakuConfig]
+    param_dtype: Optional[torch.dtype]
+    needs_device_weight_postprocess: bool = False
+    post_load_hooks: list[PostLoadHook] = field(default_factory=list)
+
+    @property
+    def runtime_quant_config(self) -> Optional[object]:
+        if self.quant_config is not None:
+            return self.quant_config
+        return self.nunchaku_config
+
+
+class _TransformerQuantAdapter:
+    def prepare(self) -> None:
+        """initialize"""
+        pass
+
+    def get_post_load_hooks(self) -> list[PostLoadHook]:
+        """post - fsdp load - hook"""
+        return []
+
+
+class _NunchakuQuantAdapter(_TransformerQuantAdapter):
+    """Adapter for Nunchaku checkpoints"""
+
+    def __init__(
+        self,
+        *,
+        nunchaku_config: NunchakuConfig,
+        model_cls: type[nn.Module],
+        safetensors_list: list[str],
+    ) -> None:
+        self.nunchaku_config = nunchaku_config
+        self.model_cls = model_cls
+        self.safetensors_list = safetensors_list
+
+    @staticmethod
+    def _validate_nunchaku_checkpoint_matches_model(
+        nunchaku_config: NunchakuConfig, model_cls: type[nn.Module]
+    ) -> None:
+        metadata = get_metadata_from_safetensors_file(
+            nunchaku_config.transformer_weights_path
+        )
+        original_dit_cls_name = json.loads(metadata.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: "
+                f"{original_dit_cls_name} does not match that of specified DiT name: "
+                f"{specified_dit_cls_name}"
+            )
+
+    def prepare(self) -> None:
+        self.nunchaku_config.model_cls = self.model_cls
+        _NunchakuQuantAdapter._validate_nunchaku_checkpoint_matches_model(
+            nunchaku_config=self.nunchaku_config,
+            model_cls=self.model_cls,
+        )
+
+    def get_post_load_hooks(self) -> list[PostLoadHook]:
+        return [partial(_patch_nunchaku_scales, safetensors_list=self.safetensors_list)]
+
+
+class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter):
+    """Adapter for black-forest-labs/FLUX.2-dev-NVFP4"""
+
+    def __init__(
+        self,
+        *,
+        cls_name: str,
+        server_args: ServerArgs,
+        quant_config: Optional[QuantizationConfig],
+    ) -> None:
+        self.cls_name = cls_name
+        self.server_args = server_args
+        self.quant_config = quant_config
+
+    @staticmethod
+    def _maybe_adjust_flux2_nvfp4_fallback_defaults(
+        cls_name: str,
+        server_args: ServerArgs,
+        quant_config: Optional[QuantizationConfig],
+    ) -> None:
+        if cls_name != "Flux2Transformer2DModel" or quant_config is None:
+            return
+
+        quant_name_getter = getattr(type(quant_config), "get_name", None)
+        quant_name = quant_name_getter() if callable(quant_name_getter) else None
+        if quant_name != "modelopt_fp4":
+            return
+
+        weights_path = os.path.basename(server_args.transformer_weights_path or "")
+        if not weights_path.endswith("-mixed.safetensors") or server_args.tp_size <= 1:
+            return
+
+        if server_args.dit_cpu_offload or server_args.text_encoder_cpu_offload:
+            server_args.dit_cpu_offload = False
+            server_args.text_encoder_cpu_offload = False
+            logger.warning(
+                "FLUX.2 mixed NVFP4 is using the ModelOpt FP4 path with tp_size=%d; "
+                "disabling dit/text-encoder CPU offload to avoid TP all-gather "
+                "launch failures. Override the offload flags explicitly if you need "
+                "the old behavior.",
+                server_args.tp_size,
+            )
+
+    def prepare(self) -> None:
+        _Flux2Nvfp4FallbackAdapter._maybe_adjust_flux2_nvfp4_fallback_defaults(
+            cls_name=self.cls_name,
+            server_args=self.server_args,
+            quant_config=self.quant_config,
+        )
+
+
+class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter):
+    """Adapter for diffusion ModelOpt FP8 checkpoints."""
+
+    def __init__(
+        self,
+        *,
+        server_args: ServerArgs,
+        quant_config: Optional[QuantizationConfig],
+    ) -> None:
+        self.server_args = server_args
+        self.quant_config = quant_config
+
+    @staticmethod
+    def _maybe_disable_incompatible_dit_offload_modes(
+        server_args: ServerArgs,
+        quant_config: Optional[QuantizationConfig],
+    ) -> None:
+        if quant_config is None:
+            return
+
+        quant_name_getter = getattr(type(quant_config), "get_name", None)
+        quant_name = quant_name_getter() if callable(quant_name_getter) else None
+
+        if quant_name != "modelopt_fp8":
+            return
+
+        if server_args.dit_cpu_offload:
+            server_args.dit_cpu_offload = False
+            logger.warning(
+                "ModelOpt FP8 diffusion checkpoints currently keep dit_cpu_offload "
+                "disabled. Layerwise DiT offload stays enabled because the runtime "
+                "now preserves the restored FP8 tensor strides.",
+            )
+
+    def prepare(self) -> None:
+        _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes(
+            server_args=self.server_args,
+            quant_config=self.quant_config,
+        )
+
+
+class _BitsAndBytes4BitAdapter(_TransformerQuantAdapter):
+    """Adapter for pre-quantized bitsandbytes 4-bit transformer checkpoints."""
+
+    def __init__(
+        self,
+        *,
+        server_args: ServerArgs,
+        quant_config: Optional[QuantizationConfig],
+    ) -> None:
+        self.server_args = server_args
+        self.quant_config = quant_config
+
+    @staticmethod
+    def _maybe_disable_incompatible_offload_modes(
+        server_args: ServerArgs,
+        quant_config: Optional[QuantizationConfig],
+    ) -> None:
+        if _get_quant_config_name(quant_config) != "bitsandbytes":
+            return
+
+        changed = []
+        if server_args.dit_cpu_offload:
+            server_args.dit_cpu_offload = False
+            changed.append("dit_cpu_offload=False")
+        if server_args.use_fsdp_inference:
+            server_args.use_fsdp_inference = False
+            changed.append("use_fsdp_inference=False")
+        if changed:
+            logger.warning(
+                "Keeping bitsandbytes 4-bit transformer GPU-resident: %s",
+                ", ".join(changed),
+            )
+
+    def prepare(self) -> None:
+        _BitsAndBytes4BitAdapter._maybe_disable_incompatible_offload_modes(
+            server_args=self.server_args,
+            quant_config=self.quant_config,
+        )
+
+
+def resolve_transformer_safetensors_to_load(
+    server_args: ServerArgs, component_model_path: str
+) -> list[str]:
+    """Resolve transformer weights from the base component path or an override."""
+    quantized_path = server_args.transformer_weights_path
+
+    if quantized_path:
+        original_quantized_path = quantized_path
+        quantized_path = maybe_download_model(original_quantized_path)
+        logger.info("using quantized transformer weights from: %s", quantized_path)
+        if os.path.isfile(quantized_path) and quantized_path.endswith(".safetensors"):
+            safetensors_list = [quantized_path]
+        else:
+            safetensors_list = _list_safetensors_files(quantized_path)
+            if not safetensors_list and not os.path.exists(original_quantized_path):
+                logger.warning(
+                    "No safetensors files found in cached transformer weights path "
+                    "%s; refreshing snapshot for %s",
+                    quantized_path,
+                    original_quantized_path,
+                )
+                quantized_path = snapshot_download(
+                    repo_id=original_quantized_path,
+                    ignore_patterns=["*.onnx", "*.msgpack"],
+                    allow_patterns=[
+                        "*.json",
+                        "*.safetensors",
+                        "*.safetensors.index.json",
+                    ],
+                    max_workers=8,
+                )
+                safetensors_list = _list_safetensors_files(quantized_path)
+    else:
+        safetensors_list = _list_safetensors_files(component_model_path)
+
+    safetensors_list = _prefer_mixed_safetensors_files(safetensors_list)
+    safetensors_list = _filter_duplicate_precision_variant_safetensors(safetensors_list)
+
+    if not safetensors_list:
+        raise ValueError(
+            f"no safetensors files found in {quantized_path or component_model_path}"
+        )
+
+    return safetensors_list
+
+
+def _prefer_mixed_safetensors_files(safetensors_list: list[str]) -> list[str]:
+    """Prefer mixed-precision transformer exports over sibling full exports.
+
+    Some raw ModelOpt NVFP4 repos ship both `foo-mixed.safetensors` and
+    `foo.safetensors`. They are alternative full transformer exports, not
+    shards, so loading both trips duplicate tensor-name validation.
+    """
+    mixed_files = [
+        path
+        for path in safetensors_list
+        if _MIXED_SAFETENSORS_RE.match(os.path.basename(path))
+    ]
+    if not mixed_files or len(mixed_files) == len(safetensors_list):
+        return safetensors_list
+
+    logger.info(
+        "Using %d mixed transformer safetensors file(s) and ignoring %d sibling "
+        "non-mixed file(s): %s",
+        len(mixed_files),
+        len(safetensors_list) - len(mixed_files),
+        mixed_files,
+    )
+    return mixed_files
+
+
+def _filter_duplicate_precision_variant_safetensors(
+    safetensors_list: list[str],
+) -> list[str]:
+    """Drop precision-specific duplicates when a canonical file is present.
+
+    Diffusers checkpoints sometimes ship both `foo.safetensors` and
+    `foo.fp16.safetensors` (and their sharded variants) in the same directory.
+    Loading both is unsafe because duplicate parameter names race and whichever
+    tensor arrives last wins, leading to non-deterministic behavior
+
+    If a canonical unsuffixed (non bf16|fp32) file exists, prefer it and drop the precision
+    variant from the same family. Precision-only families are left untouched.
+    """
+    canonical_paths = set(safetensors_list)
+    filtered: list[str] = []
+    removed: list[str] = []
+
+    for path in safetensors_list:
+        match = _PRECISION_VARIANT_SUFFIX_RE.match(path)
+        if match is None:
+            filtered.append(path)
+            continue
+
+        canonical_path = (
+            f"{match.group('stem')}{match.group('shard') or ''}{match.group('ext')}"
+        )
+        if canonical_path in canonical_paths:
+            removed.append(path)
+            continue
+
+        filtered.append(path)
+
+    if removed:
+        logger.info(
+            "Filtered %d duplicate transformer precision variant file(s): %s",
+            len(removed),
+            removed,
+        )
+
+    return filtered
+
+
+def resolve_transformer_quant_load_spec(
+    *,
+    hf_config: dict,
+    server_args: ServerArgs,
+    safetensors_list: list[str],
+    component_model_path: str,
+    model_cls: type[nn.Module],
+    cls_name: str,
+) -> TransformerQuantLoadSpec:
+    if getattr(model_cls, "handles_checkpoint_quantization", False):
+        quant_config = None
+    else:
+        quant_config = _resolve_quant_config(
+            hf_config=hf_config,
+            server_args=server_args,
+            safetensors_list=safetensors_list,
+            component_model_path=component_model_path,
+        )
+
+    if quant_config is not None:
+        packed = getattr(model_cls, "packed_modules_mapping", None)
+        if packed and hasattr(quant_config, "packed_modules_mapping"):
+            quant_config.packed_modules_mapping = packed
+
+    nunchaku_config = server_args.nunchaku_config
+
+    # resolve target param dtype
+    param_dtype = _resolve_target_param_dtype(
+        quant_config=quant_config,
+        nunchaku_config=nunchaku_config,
+        server_args=server_args,
+    )
+
+    adapters = _build_transformer_quant_adapters(
+        cls_name=cls_name,
+        server_args=server_args,
+        quant_config=quant_config,
+        nunchaku_config=nunchaku_config,
+        model_cls=model_cls,
+        safetensors_list=safetensors_list,
+    )
+    for adapter in adapters:
+        adapter.prepare()
+
+    # collect post-load hooks from built adapters
+    post_load_hooks: list[PostLoadHook] = []
+    for adapter in adapters:
+        post_load_hooks.extend(adapter.get_post_load_hooks())
+
+    return TransformerQuantLoadSpec(
+        safetensors_list=safetensors_list,
+        quant_config=quant_config,
+        nunchaku_config=nunchaku_config,
+        param_dtype=param_dtype,
+        needs_device_weight_postprocess=_needs_device_weight_postprocess(quant_config),
+        post_load_hooks=post_load_hooks,
+    )
+
+
+def _needs_device_weight_postprocess(
+    quant_config: Optional[QuantizationConfig],
+) -> bool:
+    """Return whether post-load weight processing needs CUDA/NPU tensors."""
+    quant_name = _get_quant_config_name(quant_config)
+    serialized_flag_by_quant_name = {
+        "fp8": "is_checkpoint_fp8_serialized",
+        "mxfp8": "is_checkpoint_fp8_serialized",
+        "mxfp4": "is_checkpoint_mxfp4_serialized",
+        "mxfp4_npu": "is_checkpoint_mxfp4_npu_serialized",
+    }
+    serialized_flag = serialized_flag_by_quant_name.get(quant_name)
+    if serialized_flag is None:
+        return False
+    return not getattr(quant_config, serialized_flag, False)
+
+
+def _build_transformer_quant_adapters(
+    *,
+    cls_name: str,
+    server_args: ServerArgs,
+    quant_config: Optional[QuantizationConfig],
+    nunchaku_config: Optional[NunchakuConfig],
+    model_cls: type[nn.Module],
+    safetensors_list: list[str],
+) -> list[_TransformerQuantAdapter]:
+    adapters: list[_TransformerQuantAdapter] = [
+        _Flux2Nvfp4FallbackAdapter(
+            cls_name=cls_name,
+            server_args=server_args,
+            quant_config=quant_config,
+        ),
+        _ModelOptFp8OffloadAdapter(
+            server_args=server_args,
+            quant_config=quant_config,
+        ),
+        _BitsAndBytes4BitAdapter(
+            server_args=server_args,
+            quant_config=quant_config,
+        ),
+    ]
+    if nunchaku_config is not None:
+        adapters.append(
+            _NunchakuQuantAdapter(
+                nunchaku_config=nunchaku_config,
+                model_cls=model_cls,
+                safetensors_list=safetensors_list,
+            )
+        )
+    return adapters
+
+
+def _resolve_quant_config_from_transformer_override(
+    transformer_weights_path: str,
+) -> Optional[QuantizationConfig]:
+    """Resolve quant config from an override transformer repo or directory."""
+    expanded_path = os.path.expanduser(transformer_weights_path)
+    if os.path.isfile(expanded_path):
+        return None
+
+    # A single local safetensors file does not carry a directory-level config.json.
+    # Let downstream metadata probing handle it instead of misrouting it through HF.
+    if expanded_path.endswith(".safetensors") and (
+        os.path.isabs(expanded_path)
+        or expanded_path.startswith(".")
+        or os.sep in expanded_path
+        or (os.path.altsep and os.path.altsep in expanded_path)
+    ):
+        return None
+
+    override_quantized_path = maybe_download_model(transformer_weights_path)
+    if not os.path.isdir(override_quantized_path):
+        return None
+
+    override_config_path = os.path.join(override_quantized_path, "config.json")
+    if not os.path.isfile(override_config_path):
+        return None
+
+    with open(override_config_path, encoding="utf-8") as f:
+        override_hf_config = json.load(f)
+
+    return get_quant_config(
+        override_hf_config,
+        override_quantized_path,
+    )
+
+
+def _resolve_quant_config(
+    *,
+    hf_config: dict,
+    server_args: ServerArgs,
+    safetensors_list: list[str],
+    component_model_path: str,
+) -> Optional[QuantizationConfig]:
+    """
+    resolve quant config from checkpoints' metadata
+    priority: explicit --quantization flag -> model config.json -> safetensors metadata -> format-specific fallback
+    """
+    # priority: explicit --quantization flag (e.g. mxfp8, mxfp4_npu, modelslim)
+    if server_args.quantization is not None:
+        from sglang.multimodal_gen.runtime.layers.quantization import (
+            get_quantization_config,
+        )
+
+        # modelslim requires a per-layer quant description file; load it from
+        # the component directory rather than constructing an empty config.
+        if server_args.quantization == "modelslim":
+            return get_quant_config(hf_config, component_model_path)
+
+        # Online-quant convention: for `fp8` and `mxfp4`, a no-arg
+        # QuantizationConfig() selects the post-load path -- weights load
+        # in source dtype and are quantized in
+        # process_weights_after_loading.
+        quant_cls = get_quantization_config(server_args.quantization)
+        return quant_cls()
+
+    quant_config = get_quant_config(hf_config, component_model_path)
+    if quant_config is None and server_args.transformer_weights_path:
+        for safetensors_file in safetensors_list:
+            quant_config = get_quant_config_from_safetensors_metadata(safetensors_file)
+            if quant_config is not None:
+                return quant_config
+
+    arch_config = server_args.pipeline_config.dit_config.arch_config
+    param_names_mapping_dict = arch_config.param_names_mapping
+    reverse_param_names_mapping_dict = getattr(
+        arch_config, "reverse_param_names_mapping", None
+    )
+    quant_ignore_remap_dict = getattr(arch_config, "quant_ignore_remap", None)
+    quant_config = get_quant_config(
+        hf_config,
+        component_model_path,
+        reverse_param_names_mapping=reverse_param_names_mapping_dict,
+        quant_ignore_remap=quant_ignore_remap_dict,
+    )
+    quant_config_name = _get_quant_config_name(quant_config)
+    inferred_nvfp4_config = None
+    if quant_config is None or quant_config_name == "modelopt_fp4":
+        fallback_group_size = None
+        if quant_config_name == "modelopt_fp4":
+            fallback_group_size = getattr(quant_config, "group_size", None)
+        inferred_nvfp4_config = build_nvfp4_config_from_safetensors_list(
+            safetensors_list,
+            param_names_mapping_dict,
+            reverse_param_names_mapping_dict,
+            fallback_group_size,
+        )
+    quant_config = _merge_modelopt_fp4_configs(quant_config, inferred_nvfp4_config)
+    if quant_config is not None or not server_args.transformer_weights_path:
+        return quant_config
+
+    quant_config = _resolve_quant_config_from_transformer_override(
+        server_args.transformer_weights_path
+    )
+    quant_config = _merge_modelopt_fp4_configs(quant_config, inferred_nvfp4_config)
+    if quant_config is not None:
+        return quant_config
+
+    for safetensors_file in safetensors_list:
+        quant_config = get_quant_config_from_safetensors_metadata(safetensors_file)
+        if quant_config is not None:
+            return quant_config
+
+    return inferred_nvfp4_config
+
+
+def _resolve_target_param_dtype(
+    *,
+    quant_config: Optional[QuantizationConfig],
+    nunchaku_config: Optional[NunchakuConfig],
+    server_args: ServerArgs,
+) -> Optional[torch.dtype]:
+    if quant_config is not None or nunchaku_config is not None:
+        return None
+    return resolve_precision(server_args, "dit", precision_attr="dit_precision")

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

@@ -0,0 +1,310 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for selecting and loading models."""
+
+import contextlib
+import glob
+import os
+import re
+from collections import defaultdict
+from collections.abc import Callable, Iterator
+from typing import Any, Dict, Type
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_QUANTIZED_DTYPES = {
+    torch.uint8,
+    torch.float8_e4m3fn,
+    torch.float8_e5m2,
+    torch.int8,
+}
+
+
+@contextlib.contextmanager
+def set_default_torch_dtype(dtype: torch.dtype):
+    """Sets the default torch dtype to the given dtype."""
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)
+
+
+def get_param_names_mapping(
+    mapping_dict: dict[str, str | tuple[str, int, int]],
+) -> Callable[[str], tuple[str, Any, Any]]:
+    """
+    Creates a mapping function that transforms parameter names using regex patterns.
+
+    Args:
+        mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns
+
+    Returns:
+        Callable[[str], str]: A function that maps parameter names from source to target format
+    """
+
+    def mapping_fn(name: str) -> tuple[str, Any, Any]:
+        # support chained conversions, e.g.:
+        # transformer.xxx.lora_down -> xxx.lora_down -> xxx.proj_down
+        merge_index = None
+        total_split_params = None
+        max_steps = max(8, len(mapping_dict) * 2)
+        applied_patterns: set[str] = set()
+        visited_names: set[str] = {name}
+
+        for _ in range(max_steps):
+            transformed = False
+            for pattern, replacement in mapping_dict.items():
+                # avoid re-applying the same rule on its own output
+                if pattern in applied_patterns:
+                    continue
+                if re.match(pattern, name) is None:
+                    continue
+
+                curr_merge_index = None
+                curr_total_split_params = None
+                if isinstance(replacement, tuple):
+                    curr_merge_index = replacement[1]
+                    curr_total_split_params = replacement[2]
+                    replacement = replacement[0]
+
+                new_name = re.sub(pattern, replacement, name)
+
+                if new_name != name:
+                    if curr_merge_index is not None:
+                        merge_index = curr_merge_index
+                        total_split_params = curr_total_split_params
+
+                    name = new_name
+                    applied_patterns.add(pattern)
+                    if name in visited_names:
+                        transformed = False
+                        break
+                    visited_names.add(name)
+                    transformed = True
+                    break
+
+            if not transformed:
+                break
+
+        return name, merge_index, total_split_params
+
+    return mapping_fn
+
+
+def hf_to_custom_state_dict(
+    hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]],
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]],
+    valid_target_names: set[str] | None = None,
+) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]:
+    """
+    Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary.
+
+    Args:
+        hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary
+        param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format
+
+    Returns:
+        custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict
+        reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf
+    """
+    custom_param_sd = {}
+    to_merge_params = defaultdict(dict)  # type: ignore
+    reverse_param_names_mapping = {}
+    if isinstance(hf_param_sd, dict):
+        hf_param_sd = hf_param_sd.items()  # type: ignore
+    for source_param_name, full_tensor in hf_param_sd:  # type: ignore
+        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
+            source_param_name
+        )
+        if (
+            valid_target_names is not None
+            and target_param_name != source_param_name
+            and source_param_name in valid_target_names
+            and target_param_name not in valid_target_names
+        ):
+            target_param_name = source_param_name
+            merge_index = None
+            num_params_to_merge = None
+        if target_param_name == "" or target_param_name is None:  # type: ignore[comparison-overlap]
+            continue
+        reverse_param_names_mapping[target_param_name] = (
+            source_param_name,
+            merge_index,
+            num_params_to_merge,
+        )
+        if merge_index is not None:
+            to_merge_params[target_param_name][merge_index] = full_tensor
+            if len(to_merge_params[target_param_name]) == num_params_to_merge:
+                # cat at output dim according to the merge_index order
+                sorted_tensors = [
+                    to_merge_params[target_param_name][i]
+                    for i in range(num_params_to_merge)
+                ]
+                full_tensor = torch.cat(sorted_tensors, dim=0)
+                del to_merge_params[target_param_name]
+            else:
+                continue
+        existing_tensor = custom_param_sd.get(target_param_name)
+        if existing_tensor is not None and existing_tensor.dtype != full_tensor.dtype:
+            existing_is_quantized = existing_tensor.dtype in _QUANTIZED_DTYPES
+            current_is_quantized = full_tensor.dtype in _QUANTIZED_DTYPES
+            if existing_is_quantized and not current_is_quantized:
+                logger.debug(
+                    "Keeping quantized duplicate for %s: existing=%s new=%s",
+                    target_param_name,
+                    existing_tensor.dtype,
+                    full_tensor.dtype,
+                )
+                continue
+            if current_is_quantized and not existing_is_quantized:
+                logger.debug(
+                    "Replacing non-quantized duplicate for %s: existing=%s new=%s",
+                    target_param_name,
+                    existing_tensor.dtype,
+                    full_tensor.dtype,
+                )
+        custom_param_sd[target_param_name] = full_tensor
+    return custom_param_sd, reverse_param_names_mapping
+
+
+class skip_init_modules:
+    def __enter__(self):
+        # Save originals
+        self._orig_reset = {}
+        for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Embedding):
+            self._orig_reset[cls] = cls.reset_parameters
+            cls.reset_parameters = lambda self: None  # skip init
+        from transformers.modeling_utils import PreTrainedModel
+
+        self._pretrained_model_cls = PreTrainedModel
+        self._orig_post_init = PreTrainedModel.post_init
+        PreTrainedModel.post_init = lambda self: None
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # restore originals
+        for cls, orig in self._orig_reset.items():
+            cls.reset_parameters = orig
+        self._pretrained_model_cls.post_init = self._orig_post_init
+
+
+def _normalize_component_type(module_type: str) -> str:
+    """Normalize module types like 'text_encoder_2' -> 'text_encoder'."""
+    return re.sub(r"_\d+$", "", module_type)
+
+
+def _clean_hf_config_inplace(model_config: dict) -> None:
+    """Remove common extraneous HF fields if present."""
+    for key in (
+        "_name_or_path",
+        "transformers_version",
+        "model_type",
+        "tokenizer_class",
+        "torch_dtype",
+    ):
+        model_config.pop(key, None)
+
+
+def _try_redownload_missing_shards(model_path: str, missing: list[str]) -> bool:
+    """Try to re-download missing safetensors shards from HuggingFace Hub.
+
+    Parses the repo_id and revision from the HF cache path structure
+    (models--{org}--{repo}/snapshots/{revision}) and calls hf_hub_download
+    for each missing shard. Returns True if all shards were recovered.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+
+        match = re.search(
+            r"models--([^/\\]+)--([^/\\]+)[/\\]snapshots[/\\]([^/\\]+)", model_path
+        )
+        if not match:
+            return False
+
+        repo_id = f"{match.group(1)}/{match.group(2)}"
+        revision = match.group(3)
+        logger.warning(
+            "Incomplete checkpoint for %s (revision %.8s) — missing shards: %s. "
+            "Attempting auto-repair via HuggingFace Hub...",
+            repo_id,
+            revision,
+            missing,
+        )
+        for shard in missing:
+            hf_hub_download(repo_id=repo_id, filename=shard, revision=revision)
+        logger.info("Auto-repair succeeded for %s.", repo_id)
+        return True
+    except Exception as e:
+        logger.warning("Auto-repair failed: %s", e)
+        return False
+
+
+def _list_safetensors_files(model_path: str) -> list[str]:
+    """List all .safetensors files under a directory.
+
+    If a safetensors index file is present, verifies that every shard listed
+    in the index actually exists on disk. Missing shards are first repaired
+    automatically via HuggingFace Hub (if the path is an HF cache entry);
+    if repair fails a clear RuntimeError is raised.
+    """
+    found = sorted(glob.glob(os.path.join(str(model_path), "*.safetensors")))
+
+    index_path = os.path.join(
+        str(model_path), "diffusion_pytorch_model.safetensors.index.json"
+    )
+    if os.path.exists(index_path):
+        import json
+
+        with open(index_path) as f:
+            index = json.load(f)
+        expected_shards = sorted(set(index.get("weight_map", {}).values()))
+        found_basenames = {os.path.basename(p) for p in found}
+        missing = [s for s in expected_shards if s not in found_basenames]
+        if missing:
+            repaired = _try_redownload_missing_shards(model_path, missing)
+            if repaired:
+                found = sorted(
+                    glob.glob(os.path.join(str(model_path), "*.safetensors"))
+                )
+            else:
+                raise RuntimeError(
+                    f"Checkpoint at '{model_path}' is incomplete — the following "
+                    f"shard(s) listed in the index are missing from disk: "
+                    f"{missing}. Re-download the checkpoint (e.g. "
+                    f"`huggingface-cli download {os.path.basename(model_path)}`)."
+                )
+
+    return found
+
+
+BYTES_PER_GB = 1024**3
+
+
+def get_memory_usage_of_component(module) -> float | None:
+    """
+    returned value is in GB, rounded to 2 decimal digits
+    """
+    if not isinstance(module, nn.Module):
+        return None
+    if hasattr(module, "get_memory_footprint"):
+        usage = module.get_memory_footprint() / BYTES_PER_GB
+    else:
+        # manually
+        param_size = sum(p.numel() * p.element_size() for p in module.parameters())
+        buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
+
+        total_size_bytes = param_size + buffer_size
+        usage = total_size_bytes / (1024**3)
+
+    return round(usage, 2)
+
+
+# component name ->  ComponentLoader class
+component_name_to_loader_cls: Dict[str, Type[Any]] = {}

python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class WeightLoadPlan:
+    """Device plan for checkpoint loading, before runtime residency takes over."""
+
+    # Device used while materializing checkpoint tensors from files.
+    checkpoint_load_device: torch.device
+    # Device required while running process_weights_after_loading; None means unchanged.
+    weight_postprocess_device: torch.device | None = None
+    # Delay non-FSDP component CPU offload until after weight postprocessing.
+    defer_component_cpu_offload: bool = False
+
+    @classmethod
+    def for_component(
+        cls,
+        *,
+        checkpoint_load_device: torch.device,
+        needs_device_weight_postprocess: bool,
+        component_cpu_offload: bool,
+    ) -> "WeightLoadPlan":
+        # if on-device weight postprocessing is required, load directly to device to speedup loading
+        weight_postprocess_device = (
+            checkpoint_load_device if needs_device_weight_postprocess else None
+        )
+        return cls(
+            checkpoint_load_device=checkpoint_load_device,
+            weight_postprocess_device=weight_postprocess_device,
+            defer_component_cpu_offload=(
+                needs_device_weight_postprocess and component_cpu_offload
+            ),
+        )

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -0,0 +1,418 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py
+"""Utilities for downloading, loading, initializing and verifying model weights."""
+
+import hashlib
+import json
+import os
+import tempfile
+from collections import defaultdict
+from collections.abc import Callable, Generator, Iterable
+from pathlib import Path
+
+import filelock
+import torch
+from safetensors.torch import safe_open
+from torch.distributed.tensor import DTensor
+from tqdm.auto import tqdm
+
+try:
+    from runai_model_streamer import SafetensorsStreamer
+
+    HAS_RUNAI_MODEL_STREAMER = True
+except ImportError:
+    HAS_RUNAI_MODEL_STREAMER = False
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+# use system-level temp directory for file locks, so that multiple users
+# can share the same lock without error.
+# lock files in the temp directory will be automatically deleted when the
+# system reboots, so users will not complain about annoying lock files
+temp_dir = tempfile.gettempdir()
+
+
+class DisabledTqdm(tqdm):
+
+    def __init__(self, *args, **kwargs):
+        kwargs["disable"] = True
+        super().__init__(*args, **kwargs)
+
+
+def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
+    lock_dir = cache_dir or temp_dir
+    model_name_or_path = str(model_name_or_path)
+    os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
+    model_name = model_name_or_path.replace("/", "-")
+    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
+    # add hash to avoid conflict with old users' lock files
+    lock_file_name = hash_name + model_name + ".lock"
+    # mode 0o666 is required for the filelock to be shared across users
+    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
+    return lock
+
+
+# For models like Mistral-7B-v0.3, there are both sharded
+# safetensors files and a consolidated safetensors file.
+# Passing both of these to the weight loader functionality breaks.
+# So, we use the index_file to
+# look up which safetensors files should be used.
+def filter_duplicate_safetensors_files(
+    hf_weights_files: list[str], hf_folder: str, index_file: str
+) -> list[str]:
+    # model.safetensors.index.json is a mapping from keys in the
+    # torch state_dict to safetensors file holding that weight.
+    index_file_name = os.path.join(hf_folder, index_file)
+    if not os.path.isfile(index_file_name):
+        return hf_weights_files
+
+    # Iterate through the weight_map (weight_name: safetensors files)
+    # to identify weights that we should use.
+    with open(index_file_name) as f:
+        weight_map = json.load(f)["weight_map"]
+    weight_files_in_index = set()
+    for weight_name in weight_map:
+        weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
+    # Filter out any fields that are not found in the index file.
+    hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
+    return hf_weights_files
+
+
+def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]:
+    """
+    Exclude files that are not needed for inference.
+
+    See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
+    """
+    blacklist = [
+        "training_args.bin",
+        "optimizer.bin",
+        "optimizer.pt",
+        "scheduler.pt",
+        "scaler.pt",
+    ]
+    hf_weights_files = [
+        f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist)
+    ]
+    return hf_weights_files
+
+
+# explicitly use pure text format, with a newline at the end
+# this makes it impossible to see the animation in the progress bar
+# but will avoid messing up with ray or multiprocessing, which wraps
+# each line of output with some prefix.
+_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501
+
+
+def _validate_safetensors_file(file_path: str) -> bool:
+    """
+    Validate that a safetensors file is readable and not corrupted.
+
+    Args:
+        file_path: Path to the safetensors file
+
+    Returns:
+        True if file is valid, False if corrupted
+    """
+    try:
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            _ = list(f.keys())
+        return True
+    except Exception as e:
+        logger.error(
+            "Corrupted safetensors file detected: %s - %s: %s",
+            file_path,
+            type(e).__name__,
+            str(e),
+        )
+        return False
+
+
+def _raise_if_duplicate_safetensors_keys(hf_weights_files: list[str]) -> None:
+    """Fail fast when multiple safetensors files define the same tensor name. Make sure runtime behavior is deterministic
+
+    Duplicate keys across files are almost always a packaging error for inference:
+    for example shipping both full and fp16 variants, or mixing consolidated and
+    sharded checkpoints. Continuing would make the final loaded value depend on
+    file iteration or streamer delivery order.
+    """
+    if len(hf_weights_files) <= 1:
+        return
+
+    key_to_file: dict[str, str] = {}
+    duplicate_files_by_key: dict[str, set[str]] = defaultdict(set)
+
+    for st_file in hf_weights_files:
+        with safe_open(st_file, framework="pt", device="cpu") as f:
+            for name in f.keys():  # noqa: SIM118
+                previous_file = key_to_file.get(name)
+                if previous_file is None:
+                    key_to_file[name] = st_file
+                    continue
+                if previous_file == st_file:
+                    continue
+                duplicate_files_by_key[name].update((previous_file, st_file))
+
+    if not duplicate_files_by_key:
+        return
+
+    examples = []
+    for key in sorted(duplicate_files_by_key)[:8]:
+        files = ", ".join(
+            sorted(os.path.basename(p) for p in duplicate_files_by_key[key])
+        )
+        examples.append(f"{key} [{files}]")
+
+    raise ValueError(
+        "Duplicate tensor names detected across safetensors files. Refusing to load "
+        "because final weights would depend on file or streamer ordering. "
+        f"Found {len(duplicate_files_by_key)} duplicate tensor name(s). "
+        f"Examples: {examples}. "
+        "This usually means multiple precision variants or consolidated+sharded "
+        "checkpoints were passed together."
+    )
+
+
+def safetensors_weights_iterator(
+    hf_weights_files: list[str],
+    to_cpu: bool = True,
+    use_runai_model_streamer: bool | None = None,
+    key_filter: Callable[[str], bool] | None = None,
+    clone_streamed_tensors: bool = True,
+    weight_load_plan: WeightLoadPlan | None = None,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files."""
+    enable_tqdm = (
+        not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+    )
+    if weight_load_plan is not None:
+        checkpoint_device = torch.device(weight_load_plan.checkpoint_load_device)
+        to_cpu = checkpoint_device.type == "cpu"
+        device = str(checkpoint_device)
+    else:
+        device = "cpu" if to_cpu else str(get_local_torch_device())
+    if use_runai_model_streamer is None:
+        use_runai_model_streamer = (
+            HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER
+        )
+
+    # Validate files before loading
+    corrupted_files = [
+        st_file
+        for st_file in hf_weights_files
+        if not _validate_safetensors_file(st_file)
+    ]
+
+    if corrupted_files:
+        # Delete corrupted files (both symlink and blob if applicable)
+        for file_path in corrupted_files:
+            try:
+                if os.path.islink(file_path):
+                    blob_path = os.path.realpath(file_path)
+                    os.remove(file_path)
+                    logger.info(
+                        "Removed corrupted symlink: %s", os.path.basename(file_path)
+                    )
+                    if os.path.exists(blob_path):
+                        os.remove(blob_path)
+                        logger.info(
+                            "Removed corrupted blob: %s", os.path.basename(blob_path)
+                        )
+                elif os.path.isfile(file_path):
+                    os.remove(file_path)
+                    logger.info(
+                        "Removed corrupted file: %s", os.path.basename(file_path)
+                    )
+            except Exception as e:
+                logger.warning("Failed to remove corrupted file %s: %s", file_path, e)
+
+        raise RuntimeError(
+            f"Found {len(corrupted_files)} corrupted safetensors file(s). "
+            f"Files have been removed: {[os.path.basename(f) for f in corrupted_files]}. "
+            "Please retry - the files will be re-downloaded automatically."
+        )
+
+    _raise_if_duplicate_safetensors_keys(hf_weights_files)
+
+    if use_runai_model_streamer:
+        logger.info(
+            "Loading safetensors with Run:ai Model Streamer to %s",
+            "cpu" if to_cpu else device,
+        )
+        with SafetensorsStreamer() as streamer:
+            if to_cpu:
+                streamer.stream_files(hf_weights_files)
+            else:
+                streamer.stream_files(hf_weights_files, device=device)
+            for name, tensor in streamer.get_tensors():
+                if key_filter is not None and not key_filter(name):
+                    continue
+                if to_cpu:
+                    yield name, tensor.clone().detach()
+                elif clone_streamed_tensors:
+                    yield name, tensor.clone().detach()
+                else:
+                    yield name, tensor
+    else:
+        for st_file in tqdm(
+            hf_weights_files,
+            desc="Loading safetensors checkpoint shards",
+            disable=not enable_tqdm,
+            bar_format=_BAR_FORMAT,
+        ):
+            with safe_open(st_file, framework="pt", device=device) as f:
+                for name in f.keys():  # noqa: SIM118
+                    if key_filter is not None and not key_filter(name):
+                        continue
+                    param = f.get_tensor(name)
+                    yield name, param
+
+
+def _load_pt_file(bin_file: str, device: str) -> dict:
+    """Load a PyTorch checkpoint file, handling legacy tar format.
+
+    PyTorch 2.6 changed the default of weights_only from False to True.
+    Legacy tar format files cannot be loaded with weights_only=True.
+    This function tries weights_only=True first, then falls back to False
+    for legacy tar format files from trusted sources (HuggingFace Hub).
+    """
+    try:
+        return torch.load(bin_file, map_location=device, weights_only=True)
+    except RuntimeError as e:
+        if "legacy .tar format" in str(e):
+            logger.warning(
+                "Loading %s with weights_only=False (legacy tar format)",
+                os.path.basename(bin_file),
+            )
+            return torch.load(bin_file, map_location=device, weights_only=False)
+        raise
+
+
+def pt_weights_iterator(
+    hf_weights_files: list[str],
+    to_cpu: bool = True,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model bin/pt files."""
+    device = "cpu" if to_cpu else str(get_local_torch_device())
+    enable_tqdm = (
+        not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+    )
+    for bin_file in tqdm(
+        hf_weights_files,
+        desc="Loading pt checkpoint shards",
+        disable=not enable_tqdm,
+        bar_format=_BAR_FORMAT,
+    ):
+        state = _load_pt_file(bin_file, device)
+        yield from state.items()
+        del state
+
+
+def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+    """Default weight loader."""
+    try:
+        if param.numel() == 1 and loaded_weight.numel() == 1:
+            # Sometimes scalar values aren't considered tensors with shapes
+            # so if both param and loaded_weight are a scalar,
+            # "broadcast" instead of copy
+            param.data.fill_(loaded_weight.item())
+        else:
+            assert param.size() == loaded_weight.size(), (
+                f"Attempted to load weight ({loaded_weight.size()}) "
+                f"into parameter ({param.size()})"
+            )
+
+            param.data.copy_(loaded_weight)
+    except Exception:
+        # NOTE: This exception is added for the purpose of setting breakpoint to
+        # debug weight loading issues.
+        raise
+
+
+def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
+    """Remap the name of FP8 k/v_scale parameters.
+
+    This function handles the remapping of FP8 k/v_scale parameter names.
+    It detects if the given name ends with a suffix and attempts to remap
+    it to the expected name format in the model. If the remapped name is not
+    found in the params_dict, a warning is printed and None is returned.
+
+    Args:
+        name (str): The original loaded checkpoint parameter name.
+        params_dict (dict): Dictionary containing the model's named parameters.
+
+    Returns:
+        str: The remapped parameter name if successful, or the original name
+             if no remapping is needed.
+        None: If the remapped name is not found in params_dict.
+    """
+    if name.endswith(".kv_scale"):
+        logger.warning_once(
+            "DEPRECATED. Found kv_scale in the checkpoint. "
+            "This format is deprecated in favor of separate k_scale and "
+            "v_scale tensors and will be removed in a future release. "
+            "Functionally, we will remap kv_scale to k_scale and duplicate "
+            "k_scale to v_scale"
+        )
+        # NOTE: we remap the deprecated kv_scale to k_scale
+        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
+        if remapped_name not in params_dict:
+            logger.warning_once(
+                f"Found kv_scale in the checkpoint (e.g. {name}), "
+                "but not found the expected name in the model "
+                f"(e.g. {remapped_name}). kv_scale is "
+                "not loaded."
+            )
+            return None
+        return remapped_name
+
+    possible_scale_names = [".k_scale", ".v_scale"]
+    modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"]
+    for scale_name in possible_scale_names:
+        if name.endswith(scale_name):
+            if any(mo_scale_name in name for mo_scale_name in modelopt_scale_names):
+                remapped_name = name.replace(
+                    f".self_attn.{scale_name[1]}_proj{scale_name}",
+                    f".self_attn.attn{scale_name}",
+                )
+            else:
+                remapped_name = name.replace(scale_name, f".attn{scale_name}")
+            if remapped_name not in params_dict:
+                logger.warning_once(
+                    f"Found {scale_name} in the checkpoint (e.g. {name}), "
+                    "but not found the expected name in the model "
+                    f"(e.g. {remapped_name}). {scale_name} is "
+                    "not loaded."
+                )
+                return None
+            return remapped_name
+
+    # If there were no matches, return the untouched param name
+    return name
+
+
+def compute_weights_checksum(
+    named_params: Iterable[tuple[str, torch.Tensor]],
+) -> str:
+    """Compute a SHA-256 checksum for a set of (name, tensor) pairs.
+
+    Used to verify the correctness of weight refitting. After a refit,
+    compare the checksum of the in-GPU model weights against the checksum
+    of the on-disk tensors or the tensors in the training engine.
+    """
+    hasher = hashlib.sha256()
+    for name, tensor in sorted(named_params, key=lambda x: x[0]):
+        hasher.update(name.encode())
+        t = tensor.detach()
+        # DTensor doesn't support .numpy(); extract the local tensor.
+        if isinstance(t, DTensor):
+            t = t._local_tensor
+        hasher.update(t.cpu().contiguous().reshape(-1).view(torch.uint8).numpy().data)
+    return hasher.hexdigest()