[diffusion] feat: respect component weight overrides for upsamplers (#36874)

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

@@ -197,6 +197,7 @@ def _load_explicit_config(
 class UpsamplerLoader(PlainStateDictComponentLoader):
     component_names = ["spatial_upsampler"]
     expected_library = "diffusers"
+    supports_component_weight_override = True
 
     def load_customized(
         self,
@@ -204,7 +205,10 @@ class UpsamplerLoader(PlainStateDictComponentLoader):
         server_args: ServerArgs,
         component_name: str,
     ):
-        safetensors_path = _find_safetensors_file(component_model_path)
+        component_weights_path = self.resolve_component_weights_path(
+            component_model_path, server_args, component_name
+        )
+        safetensors_path = _find_safetensors_file(component_weights_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)

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

@@ -141,6 +141,7 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
             "_class_name": "LatentUpsampler",
             "quantization_config": {"quant_method": "bitsandbytes"},
         }
+        server_args = SimpleNamespace(component_weights_paths={})
 
         with (
             patch(
@@ -160,11 +161,38 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
             self.assertRaises(ComponentCheckpointUnsupportedError),
         ):
             UpsamplerLoader().load_customized(
-                "/model/spatial_upsampler", None, "spatial_upsampler"
+                "/model/spatial_upsampler", server_args, "spatial_upsampler"
             )
 
         load_weights.assert_not_called()
 
+    def test_upsampler_uses_exact_component_weight_override(self):
+        self.assertTrue(UpsamplerLoader.supports_component_weight_override)
+        server_args = SimpleNamespace(
+            component_weights_paths={"spatial_upsampler": "owner/repo/upsampler"}
+        )
+        with (
+            patch.object(
+                UpsamplerLoader,
+                "resolve_component_weights_path",
+                return_value="/cache/upsampler.safetensors",
+            ) as resolve_weights,
+            patch(
+                "sglang.multimodal_gen.runtime.loader.component_loaders."
+                "upsampler_loader._find_safetensors_file",
+                side_effect=RuntimeError("stop after routing"),
+            ) as find_weights,
+            self.assertRaisesRegex(RuntimeError, "stop after routing"),
+        ):
+            UpsamplerLoader().load_customized(
+                "/base/spatial_upsampler", server_args, "spatial_upsampler"
+            )
+
+        resolve_weights.assert_called_once_with(
+            "/base/spatial_upsampler", server_args, "spatial_upsampler"
+        )
+        find_weights.assert_called_once_with("/cache/upsampler.safetensors")
+
 
 if __name__ == "__main__":
     unittest.main()