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