[diffusion] refactor: split component_loader into component-wise files (#17820)

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

@@ -0,0 +1,70 @@
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.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.utils import PRECISION_TO_TYPE
+
+
+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(model_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 = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        from types import SimpleNamespace
+
+        with set_default_torch_dtype(default_dtype), skip_init_modules():
+            connector_cfg = SimpleNamespace(**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/bridge_loader.py

@@ -0,0 +1,105 @@
+from copy import deepcopy
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.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.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 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(model_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 = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        logger.info(
+            "Loading %s from %s safetensors files, default_dtype: %s",
+            class_name,
+            len(safetensors_list),
+            default_dtype,
+        )
+
+        # Check if FSDP loading is available
+        if (
+            server_args.hsdp_shard_dim is not None
+            and hasattr(model_cls, "_fsdp_shard_conditions")
+            and model_cls._fsdp_shard_conditions
+        ):
+            # 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=get_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,
+                default_dtype=default_dtype,
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                output_dtype=None,
+                strict=False,
+            )
+        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/image_encoder_loader.py

@@ -0,0 +1,58 @@
+import json
+import os
+
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.runtime.loader.text_encoder_loader import TextEncoderLoader
+from sglang.multimodal_gen.runtime.loader.utils import _clean_hf_config_inplace
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+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, *args
+    ):
+        """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,
+        # )
+        with open(os.path.join(component_model_path, "config.json")) as f:
+            model_config = json.load(f)
+        _clean_hf_config_inplace(model_config)
+        logger.debug("HF model config: %s", model_config)
+
+        encoder_config = server_args.pipeline_config.image_encoder_config
+        encoder_config.update_model_arch(model_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=server_args.image_encoder_cpu_offload,
+        )

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

@@ -0,0 +1,35 @@
+from sglang.multimodal_gen.runtime.loader.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(model_path=component_model_path)
+
+        class_name = config.pop("_class_name")
+        assert (
+            class_name is not None
+        ), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+
+        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/text_encoder_loader.py

@@ -0,0 +1,285 @@
+import dataclasses
+import glob
+import os
+from collections.abc import Generator, Iterable
+from typing import Generator, Iterable, cast
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+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
+from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
+from sglang.multimodal_gen.runtime.loader.fsdp_load import shard_model
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _clean_hf_config_inplace,
+    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.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.utils import PRECISION_TO_TYPE
+
+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 _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}`"
+            )
+
+        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
+    ):
+        """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,
+        # )
+        diffusers_pretrained_config = get_config(
+            component_model_path, trust_remote_code=True
+        )
+        model_config = get_diffusers_component_config(model_path=component_model_path)
+        _clean_hf_config_inplace(model_config)
+        logger.debug("HF model config: %s", model_config)
+
+        def is_not_first_encoder(module_name):
+            return "2" in module_name
+
+        # TODO(mick): had to throw an exception for different text-encoder arch
+        if not is_not_first_encoder(component_name):
+            encoder_config = server_args.pipeline_config.text_encoder_configs[0]
+            encoder_config.update_model_arch(model_config)
+            for key, value in diffusers_pretrained_config.__dict__.items():
+                setattr(encoder_config.arch_config, key, value)
+            encoder_dtype = server_args.pipeline_config.text_encoder_precisions[0]
+        else:
+            assert len(server_args.pipeline_config.text_encoder_configs) == 2
+            encoder_config = server_args.pipeline_config.text_encoder_configs[1]
+            encoder_config.update_model_arch(model_config)
+            encoder_dtype = server_args.pipeline_config.text_encoder_precisions[1]
+        # TODO(will): add support for other dtypes
+        return self.load_model(
+            component_model_path,
+            encoder_config,
+            server_args,
+            encoder_dtype,
+        )
+
+    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()
+        should_offload = self.should_offload(server_args, model_config)
+
+        if should_offload and not current_platform.is_mps():
+            model_device = torch.device("cpu")
+        else:
+            model_device = local_torch_device
+
+        with 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)
+            )
+
+            # Explicitly move model to target device after loading weights
+            if not should_offload:
+                model = model.to(local_torch_device)
+
+            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)
+                else:
+                    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(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:
+                raise ValueError(
+                    "Following model weights were not initialized from "
+                    f"checkpoint: {weights_not_loaded}"
+                )
+
+        return model

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

@@ -0,0 +1,120 @@
+import os
+from copy import deepcopy
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.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,
+    _normalize_component_type,
+)
+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 TransformerLoader(ComponentLoader):
+    """Shared loader for (video/audio) DiT transformers."""
+
+    component_names = ["transformer", "audio_dit", "video_dit"]
+    expected_library = "diffusers"
+
+    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."""
+        config = get_diffusers_component_config(model_path=component_model_path)
+        hf_config = deepcopy(config)
+        cls_name = config.pop("_class_name")
+        if cls_name is None:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported."
+            )
+
+        component_name = _normalize_component_type(component_name)
+        server_args.model_paths[component_name] = component_model_path
+
+        if component_name in ("transformer", "video_dit"):
+            pipeline_dit_config_attr = "dit_config"
+        elif component_name in ("audio_dit",):
+            pipeline_dit_config_attr = "audio_dit_config"
+        else:
+            raise ValueError(f"Invalid module name: {component_name}")
+        # Config from Diffusers supersedes sgl_diffusion's model config
+        dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr)
+        dit_config.update_model_arch(config)
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+
+        # 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}")
+
+        # Check if we should use custom initialization weights
+        custom_weights_path = getattr(
+            server_args, "init_weights_from_safetensors", None
+        )
+        use_custom_weights = False
+
+        if use_custom_weights:
+            logger.info(
+                "Using custom initialization weights from: %s", custom_weights_path
+            )
+            assert (
+                custom_weights_path is not None
+            ), "Custom initialization weights must be provided"
+            if os.path.isdir(custom_weights_path):
+                safetensors_list = _list_safetensors_files(custom_weights_path)
+            else:
+                assert custom_weights_path.endswith(
+                    ".safetensors"
+                ), "Custom initialization weights must be a safetensors file"
+                safetensors_list = [custom_weights_path]
+
+        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        logger.info(
+            "Loading %s from %s safetensors files, default_dtype: %s",
+            cls_name,
+            len(safetensors_list),
+            default_dtype,
+        )
+
+        # Load the model using FSDP loader
+        assert server_args.hsdp_shard_dim is not None
+        model = maybe_load_fsdp_model(
+            model_cls=model_cls,
+            init_params={"config": dit_config, "hf_config": hf_config},
+            weight_dir_list=safetensors_list,
+            device=get_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,
+            # TODO(will): make these configurable
+            default_dtype=default_dtype,
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            output_dtype=None,
+            strict=False,
+        )
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
+
+        assert (
+            next(model.parameters()).dtype == default_dtype
+        ), "Model dtype does not match default dtype"
+
+        return model

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

@@ -3,12 +3,15 @@
 # SPDX-License-Identifier: Apache-2.0
 """Utilities for selecting and loading models."""
 import contextlib
+import glob
+import os
 import re
 from collections import defaultdict
 from collections.abc import Callable, Iterator
-from typing import Any
+from typing import Any, Dict, Type
 
 import torch
+from torch import nn
 
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -102,3 +105,65 @@ def hf_to_custom_state_dict(
                 continue
         custom_param_sd[target_param_name] = full_tensor
     return custom_param_sd, reverse_param_names_mapping
+
+
+class skip_init_modules:
+    def __enter__(self):
+        # Save originals
+        self._orig_reset = {}
+        for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d):
+            self._orig_reset[cls] = cls.reset_parameters
+            cls.reset_parameters = lambda self: None  # skip init
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # restore originals
+        for cls, orig in self._orig_reset.items():
+            cls.reset_parameters = orig
+
+
+def _normalize_component_type(module_type: str) -> str:
+    """Normalize module types like 'text_encoder_2' -> 'text_encoder'."""
+    if module_type.endswith("_2"):
+        return module_type[:-2]
+    return module_type
+
+
+def _clean_hf_config_inplace(model_config: dict) -> None:
+    """Remove common extraneous HF fields if present."""
+    for key in (
+        "_name_or_path",
+        "transformers_version",
+        "model_type",
+        "tokenizer_class",
+        "torch_dtype",
+    ):
+        model_config.pop(key, None)
+
+
+def _list_safetensors_files(model_path: str) -> list[str]:
+    """List all .safetensors files under a directory."""
+    return sorted(glob.glob(os.path.join(str(model_path), "*.safetensors")))
+
+
+def get_memory_usage_of_component(module) -> float | None:
+    """
+    returned value is in GB, rounded to 2 decimal digits
+    """
+    if not isinstance(module, nn.Module):
+        return None
+    BYTES_PER_GB = 1024**3
+    if hasattr(module, "get_memory_footprint"):
+        usage = module.get_memory_footprint() / BYTES_PER_GB
+    else:
+        # manually
+        param_size = sum(p.numel() * p.element_size() for p in module.parameters())
+        buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
+
+        total_size_bytes = param_size + buffer_size
+        usage = total_size_bytes / (1024**3)
+
+    return round(usage, 2)
+
+
+# component name ->  ComponentLoader class
+component_name_to_loader_cls: Dict[str, Type[Any]] = {}

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

@@ -0,0 +1,112 @@
+import importlib.util
+import os
+
+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_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 VAELoader(ComponentLoader):
+    """Shared loader for (video/audio) VAE modules."""
+
+    component_names = ["vae", "audio_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(model_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
+
+        logger.debug("HF model config: %s", config)
+        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)
+        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)
+            vae_dtype = PRECISION_TO_TYPE[vae_precision]
+            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)
+            return vae
+
+        # Load from ModelRegistry (standard VAE classes)
+        with (
+            set_default_torch_dtype(PRECISION_TO_TYPE[vae_precision]),
+            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)
+        assert (
+            len(safetensors_list) == 1
+        ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        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)
+
+        return vae

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

@@ -0,0 +1,39 @@
+from typing import Any
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.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
+
+
+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":
+            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/vocoder_loader.py

@@ -0,0 +1,86 @@
+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_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 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(model_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
+
+        from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import (
+            LTXVocoderConfig,
+        )
+
+        vocoder_config = LTXVocoderConfig()
+        vocoder_config.update_model_arch(config)
+
+        try:
+            vocoder_precision = server_args.pipeline_config.audio_vae_precision
+        except AttributeError:
+            vocoder_precision = "fp32"
+        vocoder_dtype = PRECISION_TO_TYPE[vocoder_precision]
+
+        should_offload = self.should_offload(server_args)
+        target_device = self.target_device(should_offload)
+
+        with set_default_torch_dtype(vocoder_dtype), skip_init_modules():
+            vocoder_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            vocoder = vocoder_cls(vocoder_config).to(target_device)
+
+        safetensors_list = _list_safetensors_files(component_model_path)
+        assert (
+            len(safetensors_list) == 1
+        ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        incompatible = vocoder.load_state_dict(loaded, strict=False)
+        missing_keys = []
+        unexpected_keys = []
+        try:
+            missing_keys = incompatible.missing_keys
+            unexpected_keys = incompatible.unexpected_keys
+        except AttributeError:
+            # Best-effort fallback in case older torch returns a tuple-like.
+            try:
+                missing_keys = incompatible[0]
+                unexpected_keys = incompatible[1]
+            except Exception:
+                pass
+
+        if missing_keys or unexpected_keys:
+            logger.warning(
+                "Loaded vocoder with missing_keys=%d unexpected_keys=%d",
+                len(missing_keys),
+                len(unexpected_keys),
+            )
+        return vocoder

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

@@ -248,7 +248,7 @@ class ComposedPipelineBase(ABC):
         required_modules = self.required_config_modules
         logger.info("Loading required components: %s", required_modules)
 
-        components = {}
+        loaded_components = {}
         for module_name, (
             transformers_or_diffusers,
             architecture,
@@ -266,7 +266,7 @@ class ComposedPipelineBase(ABC):
                 continue
             if loaded_modules is not None and module_name in loaded_modules:
                 logger.info("Using module %s already provided", module_name)
-                components[module_name] = loaded_modules[module_name]
+                loaded_components[module_name] = loaded_modules[module_name]
                 continue
 
             # we load the module from the extra config module map if it exists
@@ -288,8 +288,8 @@ class ComposedPipelineBase(ABC):
                 )
             else:
                 component_model_path = os.path.join(self.model_path, load_module_name)
-            module, memory_usage = PipelineComponentLoader.load_module(
+            module, memory_usage = PipelineComponentLoader.load_component(
-                module_name=load_module_name,
+                component_name=load_module_name,
                 component_model_path=component_model_path,
                 transformers_or_diffusers=transformers_or_diffusers,
                 server_args=server_args,
@@ -297,20 +297,23 @@ class ComposedPipelineBase(ABC):
 
             self.memory_usages[load_module_name] = memory_usage
 
-            if module_name in components:
+            if module_name in loaded_components:
                 logger.warning("Overwriting module %s", module_name)
-            components[module_name] = module
+            loaded_components[module_name] = module
 
         # Check if all required modules were loaded
         for module_name in required_modules:
-            if module_name not in components or components[module_name] is None:
+            if (
+                module_name not in loaded_components
+                or loaded_components[module_name] is None
+            ):
                 raise ValueError(
-                    f"Required module key: {module_name} value: {components.get(module_name)} was not found in loaded modules {components.keys()}"
+                    f"Required module: {module_name} was not found in loaded modules: {list(loaded_components.keys())}"
                 )
 
         logger.debug("Memory usage of loaded modules: %s", self.memory_usages)
 
-        return components
+        return loaded_components
 
     def add_stage(self, stage_name: str, stage: PipelineStage):
         assert self.modules is not None, "No modules are registered"

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -10,7 +10,7 @@ import weakref
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
-from sglang.multimodal_gen.runtime.loader.component_loader import VAELoader
+from sglang.multimodal_gen.runtime.loader.vae_loader import VAELoader
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -44,7 +44,7 @@ from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
     configure_sta,
     save_mask_search_results,
 )
-from sglang.multimodal_gen.runtime.loader.component_loader import TransformerLoader
+from sglang.multimodal_gen.runtime.loader.transformer_loader import TransformerLoader
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (