diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py index 4c9cf7a5f..30647c760 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py index 523fc6a67..9550e92a8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py +++ b/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()