diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py index 2361071d3..050fc080b 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py @@ -44,6 +44,9 @@ class AdapterLoader(PlainStateDictComponentLoader): *args, ): config = self.load_component_config(component_model_path, component_name) + component_weights_path = self.resolve_component_weights_path( + component_model_path, server_args, component_name + ) cls_name = config.pop("_class_name", None) if cls_name is None: @@ -74,7 +77,7 @@ class AdapterLoader(PlainStateDictComponentLoader): adapter_cfg.update_model_arch(config) model = model_cls(adapter_cfg).to(device=target_device, dtype=default_dtype) - loaded = load_safetensors_state_dict(component_model_path) + loaded = load_safetensors_state_dict(component_weights_path) mapping = adapter_cfg.arch_config.param_names_mapping loaded = {_remap_connector_key(k, mapping): v for k, v in loaded.items()} @@ -84,7 +87,7 @@ class AdapterLoader(PlainStateDictComponentLoader): # else uninitialized would surface later as garbage embeddings. if missing or unexpected: raise ValueError( - f"Adapter weights at '{component_model_path}' do not match the " + f"Adapter weights at '{component_weights_path}' do not match the " f"instantiated {cls_name}. Missing: {sorted(missing)}. " f"Unexpected: {sorted(unexpected)}. This usually means the " "adapter config or its weight-name mapping is wrong." diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 931ac03bc..84d9171bc 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -46,6 +46,10 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( ) 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.runtime.weights.source import ( + materialize_weight, + resolve_weight, +) from sglang.srt.model_loader.checkpoint_quantization import ( resolve_checkpoint_quant_spec, ) @@ -518,6 +522,19 @@ class PlainStateDictComponentLoader(ComponentLoader): self.ensure_plain_state_dict_checkpoint(config, component_name) return config + def resolve_component_weights_path( + self, + component_model_path: str, + server_args: ServerArgs, + component_name: str, + ) -> str: + override = server_args.component_weights_paths.get(component_name) + if override is None: + return component_model_path + weights_path = materialize_weight(resolve_weight(override)) + logger.info("Using weight override for %s: %s", component_name, weights_path) + return weights_path + class ImageProcessorLoader(ComponentLoader): """Loader for image processor.""" diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py index f78c55e28..b4b7926b6 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py @@ -30,6 +30,9 @@ class DiffusionDecoderLoader(PlainStateDictComponentLoader): *args, ): config = self.load_component_config(component_model_path, component_name) + component_weights_path = self.resolve_component_weights_path( + component_model_path, server_args, component_name + ) class_name = config.pop("_class_name", None) if class_name is None: raise ValueError( @@ -54,6 +57,6 @@ class DiffusionDecoderLoader(PlainStateDictComponentLoader): model = model_cls(decoder_config).to(device=target_device, dtype=dtype) model.load_state_dict( - load_safetensors_state_dict(component_model_path), strict=True + load_safetensors_state_dict(component_weights_path), strict=True ) return model diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py index 6ebf7926d..5d236b221 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( - _list_safetensors_files, + load_safetensors_state_dict, set_default_torch_dtype, skip_init_modules, ) @@ -25,6 +24,9 @@ class SoundTokenizerLoader(PlainStateDictComponentLoader): self, component_model_path: str, server_args: ServerArgs, component_name: str ): config = self.load_component_config(component_model_path, component_name) + component_weights_path = self.resolve_component_weights_path( + component_model_path, server_args, component_name + ) class_name = config.pop("_class_name", None) or self.component_architecture assert ( class_name is not None @@ -45,11 +47,7 @@ class SoundTokenizerLoader(PlainStateDictComponentLoader): 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]) + loaded = load_safetensors_state_dict(component_weights_path) incompatible = model.load_state_dict(loaded, strict=False) missing = getattr(incompatible, "missing_keys", []) # The tokenizer is decoder-only; the checkpoint's encoder weights are diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py index e91499bbf..b21ccaaae 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py @@ -1,12 +1,10 @@ import re -from safetensors.torch import load_file as safetensors_load_file - from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( - _list_safetensors_files, + load_safetensors_state_dict, set_default_torch_dtype, skip_init_modules, ) @@ -27,6 +25,9 @@ class VocoderLoader(PlainStateDictComponentLoader): self, component_model_path: str, server_args: ServerArgs, component_name: str ): config = self.load_component_config(component_model_path, component_name) + component_weights_path = self.resolve_component_weights_path( + component_model_path, server_args, component_name + ) class_name = config.pop("_class_name", None) or self.component_architecture assert ( class_name is not None @@ -57,11 +58,7 @@ class VocoderLoader(PlainStateDictComponentLoader): 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]) + loaded = load_safetensors_state_dict(component_weights_path) mapping = vocoder_config.arch_config.param_names_mapping loaded = {_remap_vocoder_key(k, mapping): v for k, v in loaded.items()} @@ -69,7 +66,7 @@ class VocoderLoader(PlainStateDictComponentLoader): # A half-loaded vocoder produces plausible but wrong audio. if missing_keys or unexpected_keys: raise ValueError( - f"Vocoder weights at '{component_model_path}' do not match the " + f"Vocoder weights at '{component_weights_path}' do not match the " f"instantiated {class_name}. Missing: {sorted(missing_keys)}. " f"Unexpected: {sorted(unexpected_keys)}." ) diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index 4306860eb..a91623674 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -297,6 +297,9 @@ def _list_safetensors_files(model_path: str) -> list[str]: automatically via HuggingFace Hub (if the path is an HF cache entry); if repair fails a clear RuntimeError is raised. """ + if os.path.isfile(model_path): + return [str(model_path)] if str(model_path).endswith(".safetensors") else [] + found = sorted(glob.glob(os.path.join(str(model_path), "*.safetensors"))) index_path = os.path.join( 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 180f54905..523fc6a67 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 @@ -2,6 +2,7 @@ import re import unittest +from types import SimpleNamespace from unittest.mock import patch from sglang.multimodal_gen.runtime.loader.component_loaders.adapter_loader import ( @@ -33,6 +34,29 @@ class _TestLoader(PlainStateDictComponentLoader): class TestComponentQuantizationAdmission(unittest.TestCase): + def test_plain_loader_resolves_weights_separately_from_config(self): + server_args = SimpleNamespace( + component_weights_paths={"vocoder": "owner/repo/vocoder.safetensors"} + ) + with ( + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.resolve_weight", + return_value="resolved", + ), + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.materialize_weight", + return_value="/cache/vocoder.safetensors", + ), + ): + self.assertEqual( + _TestLoader().resolve_component_weights_path( + "/base/vocoder", server_args, "vocoder" + ), + "/cache/vocoder.safetensors", + ) + def test_plain_checkpoint_config_is_accepted(self): config = {"_class_name": "TestModel"}