[diffusion] feat: reject unsupported quantized component checkpoints (#35873)

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

@@ -7,7 +7,7 @@ from sglang.multimodal_gen.configs.models.adapter.ltx_2_duration_head import (
     LTX2DurationHeadConfig,
 )
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
-    ComponentLoader,
+    PlainStateDictComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.utils import (
     load_safetensors_state_dict,
@@ -16,13 +16,10 @@ from sglang.multimodal_gen.runtime.loader.utils import (
 )
 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):
+class AdapterLoader(PlainStateDictComponentLoader):
     """Loader for small adapter-style modules (e.g., LTX-2 connectors).
 
     This loader intentionally avoids FSDP sharding and just:
@@ -46,7 +43,7 @@ class AdapterLoader(ComponentLoader):
         component_name: str = "connectors",
         *args,
     ):
-        config = get_diffusers_component_config(component_path=component_model_path)
+        config = self.load_component_config(component_model_path, component_name)
 
         cls_name = config.pop("_class_name", None)
         if cls_name is None:

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

@@ -4,7 +4,7 @@ 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,
+    PlainStateDictComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
 from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
@@ -14,16 +14,13 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency
 )
 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):
+class BridgeLoader(PlainStateDictComponentLoader):
     """Loader for MOVA dual tower bridge with FSDP support."""
 
     pipeline_bridge_config_attr: str = "bridge_config"
@@ -34,7 +31,7 @@ class BridgeLoader(ComponentLoader):
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
-        config = get_diffusers_component_config(component_path=component_model_path)
+        config = self.load_component_config(component_model_path, component_name)
         hf_config = deepcopy(config)
         class_name = config.pop("_class_name", None)
         if class_name is None:

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

@@ -34,11 +34,15 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_
 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_diffusers_component_config,
     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
+from sglang.srt.model_loader.checkpoint_quantization import (
+    resolve_checkpoint_quant_spec,
+)
 
 logger = init_logger(__name__)
 
@@ -441,6 +445,36 @@ class ComponentLoader(ABC):
         return GenericComponentLoader(transformers_or_diffusers, component_architecture)
 
 
+class PlainStateDictComponentLoader(ComponentLoader):
+    """Base for native loaders whose current materializer expects plain weights."""
+
+    @staticmethod
+    def ensure_plain_state_dict_checkpoint(config: object, component_name: str) -> None:
+        try:
+            quant_spec = resolve_checkpoint_quant_spec(config)
+        except (TypeError, ValueError) as error:
+            raise ComponentCheckpointUnsupportedError(
+                f"Cannot parse checkpoint quantization metadata for "
+                f"{component_name!r}: {error}"
+            ) from error
+        if quant_spec is None:
+            return
+
+        method = quant_spec.declared_method or "unspecified"
+        raise ComponentCheckpointUnsupportedError(
+            f"{component_name!r} checkpoint declares quantization metadata in "
+            f"{quant_spec.source} (quant_method={method!r}), which its current "
+            "plain state-dict materializer cannot restore."
+        )
+
+    def load_component_config(
+        self, component_model_path: str, component_name: str
+    ) -> dict[str, Any]:
+        config = get_diffusers_component_config(component_path=component_model_path)
+        self.ensure_plain_state_dict_checkpoint(config, component_name)
+        return config
+
+
 class ImageProcessorLoader(ComponentLoader):
     """Loader for image processor."""
 

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

@@ -4,7 +4,7 @@ from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder imp
     LTX25DiffusionDecoderConfig,
 )
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
-    ComponentLoader,
+    PlainStateDictComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.utils import (
     load_safetensors_state_dict,
@@ -13,13 +13,10 @@ from sglang.multimodal_gen.runtime.loader.utils import (
 )
 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 DiffusionDecoderLoader(ComponentLoader):
+class DiffusionDecoderLoader(PlainStateDictComponentLoader):
     """Loader for the standalone, replicated LTX-2.5 diffusion decoder."""
 
     component_names = ["diffusion_decoder"]
@@ -32,7 +29,7 @@ class DiffusionDecoderLoader(ComponentLoader):
         component_name: str = "diffusion_decoder",
         *args,
     ):
-        config = get_diffusers_component_config(component_path=component_model_path)
+        config = self.load_component_config(component_model_path, component_name)
         class_name = config.pop("_class_name", None)
         if class_name is None:
             raise ValueError(

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

@@ -2,7 +2,7 @@
 from safetensors.torch import load_file as safetensors_load_file
 
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
-    ComponentLoader,
+    PlainStateDictComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.utils import (
     _list_safetensors_files,
@@ -11,23 +11,20 @@ from sglang.multimodal_gen.runtime.loader.utils import (
 )
 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):
+class SoundTokenizerLoader(PlainStateDictComponentLoader):
     component_names = ["sound_tokenizer"]
     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)
+        config = self.load_component_config(component_model_path, component_name)
         class_name = config.pop("_class_name", None) or self.component_architecture
         assert (
             class_name is not None

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

@@ -8,7 +8,7 @@ import torch
 from safetensors.torch import load_file as safetensors_load_file
 
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
-    ComponentLoader,
+    PlainStateDictComponentLoader,
 )
 from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import (
     LatentUpsampler,
@@ -158,28 +158,26 @@ def _infer_config_from_state_dict(state_dict: dict[str, torch.Tensor]) -> dict:
     return config
 
 
-def _load_config(
+def _load_explicit_config(
     safetensors_path: str,
     original_path: str,
-    state_dict: dict[str, torch.Tensor],
-) -> dict:
-    """Load upsampler config with fallback chain:
+) -> dict | None:
+    """Load an explicit upsampler config with this 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"]))
+            return 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))
+            return json.load(fp)
 
     hf = _parse_hf_url(original_path)
     if hf:
@@ -189,15 +187,14 @@ def _load_config(
             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))
+                return json.load(fp)
         except Exception:
             pass
 
-    logger.info("No explicit config found, inferring from state dict")
-    return _infer_config_from_state_dict(state_dict)
+    return None
 
 
-class UpsamplerLoader(ComponentLoader):
+class UpsamplerLoader(PlainStateDictComponentLoader):
     component_names = ["spatial_upsampler"]
     expected_library = "diffusers"
 
@@ -208,8 +205,16 @@ class UpsamplerLoader(ComponentLoader):
         component_name: str,
     ):
         safetensors_path = _find_safetensors_file(component_model_path)
+        raw_config = _load_explicit_config(safetensors_path, component_model_path)
+        if raw_config is not None:
+            self.ensure_plain_state_dict_checkpoint(raw_config, component_name)
+
         state_dict = safetensors_load_file(safetensors_path)
-        config = _load_config(safetensors_path, component_model_path, state_dict)
+        if raw_config is None:
+            logger.info("No explicit config found, inferring from state dict")
+            config = _infer_config_from_state_dict(state_dict)
+        else:
+            config = _normalize_config(raw_config)
 
         logger.info("Loading LatentUpsampler with config: %s", config)
 

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

@@ -3,7 +3,7 @@ import re
 from safetensors.torch import load_file as safetensors_load_file
 
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
-    ComponentLoader,
+    PlainStateDictComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.utils import (
     _list_safetensors_files,
@@ -12,9 +12,6 @@ from sglang.multimodal_gen.runtime.loader.utils import (
 )
 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
@@ -22,14 +19,14 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 logger = init_logger(__name__)
 
 
-class VocoderLoader(ComponentLoader):
+class VocoderLoader(PlainStateDictComponentLoader):
     component_names = ["vocoder"]
     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)
+        config = self.load_component_config(component_model_path, component_name)
         class_name = config.pop("_class_name", None) or self.component_architecture
         assert (
             class_name is not None

python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+import unittest
+from unittest.mock import patch
+
+from sglang.multimodal_gen.runtime.loader.component_loaders.adapter_loader import (
+    AdapterLoader,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.bridge_loader import (
+    BridgeLoader,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentCheckpointUnsupportedError,
+    PlainStateDictComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.diffusion_decoder_loader import (
+    DiffusionDecoderLoader,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.sound_tokenizer_loader import (
+    SoundTokenizerLoader,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.upsampler_loader import (
+    UpsamplerLoader,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.vocoder_loader import (
+    VocoderLoader,
+)
+
+
+class _TestLoader(PlainStateDictComponentLoader):
+    pass
+
+
+class TestComponentQuantizationAdmission(unittest.TestCase):
+    def test_plain_checkpoint_config_is_accepted(self):
+        config = {"_class_name": "TestModel"}
+
+        with patch(
+            "sglang.multimodal_gen.runtime.loader.component_loaders."
+            "component_loader.get_diffusers_component_config",
+            return_value=config,
+        ):
+            loaded = _TestLoader().load_component_config("/model/component", "test")
+
+        self.assertIs(loaded, config)
+
+    def test_all_quantization_metadata_layouts_fail_closed(self):
+        configs = {
+            "quantization_config": {
+                "quantization_config": {"quant_method": "bitsandbytes"}
+            },
+            "text_config.quantization_config": {
+                "text_config": {"quantization_config": {"quant_method": "fp8"}}
+            },
+            "compression_config": {
+                "compression_config": {"quant_method": "compressed-tensors"}
+            },
+        }
+
+        for source, config in configs.items():
+            with (
+                self.subTest(source=source),
+                self.assertRaisesRegex(
+                    ComponentCheckpointUnsupportedError,
+                    rf"{re.escape(source)}.*quant_method=.*cannot restore",
+                ),
+            ):
+                _TestLoader.ensure_plain_state_dict_checkpoint(config, "test_component")
+
+        with self.assertRaisesRegex(
+            ComponentCheckpointUnsupportedError,
+            "Cannot parse checkpoint quantization metadata",
+        ):
+            _TestLoader.ensure_plain_state_dict_checkpoint(
+                {"quantization_config": "invalid"}, "test_component"
+            )
+
+    def test_native_raw_state_loaders_share_the_admission_boundary(self):
+        loader_classes = (
+            AdapterLoader,
+            BridgeLoader,
+            DiffusionDecoderLoader,
+            SoundTokenizerLoader,
+            UpsamplerLoader,
+            VocoderLoader,
+        )
+
+        for loader_class in loader_classes:
+            with self.subTest(loader=loader_class.__name__):
+                self.assertTrue(issubclass(loader_class, PlainStateDictComponentLoader))
+
+    def test_adapter_rejects_quantization_before_model_construction(self):
+        config = {
+            "_class_name": "LTX2ConnectorModel",
+            "quantization_config": {"quant_method": "bitsandbytes"},
+        }
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.loader.component_loaders."
+                "component_loader.get_diffusers_component_config",
+                return_value=config,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.loader.component_loaders."
+                "adapter_loader.ModelRegistry.resolve_model_cls"
+            ) as resolve_model,
+            self.assertRaises(ComponentCheckpointUnsupportedError),
+        ):
+            AdapterLoader().load_customized("/model/connectors", None, "connectors")
+
+        resolve_model.assert_not_called()
+
+    def test_upsampler_rejects_quantization_before_loading_weights(self):
+        config = {
+            "_class_name": "LatentUpsampler",
+            "quantization_config": {"quant_method": "bitsandbytes"},
+        }
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.loader.component_loaders."
+                "upsampler_loader._find_safetensors_file",
+                return_value="/model/spatial_upsampler/model.safetensors",
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.loader.component_loaders."
+                "upsampler_loader._load_explicit_config",
+                return_value=config,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.loader.component_loaders."
+                "upsampler_loader.safetensors_load_file"
+            ) as load_weights,
+            self.assertRaises(ComponentCheckpointUnsupportedError),
+        ):
+            UpsamplerLoader().load_customized(
+                "/model/spatial_upsampler", None, "spatial_upsampler"
+            )
+
+        load_weights.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()