From 795dd7abcec560a8aa329c6470e53801e56db75e Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 4 Sep 2026 09:15:05 +0800 Subject: [PATCH] [diffusion] feat: compose third-party component bundles safely (#37816) --- .../runtime/layers/lora/linear.py | 11 ++ .../loader/component_loaders/vae_loader.py | 135 ++++++++++++++++- .../runtime/loader/minimax_h3_weights.py | 11 +- .../runtime/pipelines_core/lora/pipeline.py | 14 ++ .../runtime/utils/hf_diffusers_utils.py | 23 ++- .../test/server/perf_baselines/h100.json | 4 +- .../test/unit/test_hf_diffusers_utils.py | 40 +++++ .../test/unit/test_lora_pipeline.py | 27 ++++ .../test/unit/test_transformer_quant.py | 13 ++ .../test/unit/test_vae_loader.py | 143 ++++++++++++++++++ 10 files changed, 408 insertions(+), 13 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 2bba15051..b0ccd2af0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -27,6 +27,7 @@ from sglang.multimodal_gen.runtime.layers.linear import ( QKVParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, @@ -116,6 +117,16 @@ class BaseLayerWithLoRA(nn.Module): def bias(self): return getattr(self.base_layer, "bias", None) + @property + def can_merge_base_weight(self) -> bool: + """Whether a LoRA delta may safely replace the stored base weight.""" + weight = self.weight + if not (weight.dtype.is_floating_point or weight.dtype.is_complex): + return False + if isinstance(self.base_layer, LinearBase): + return isinstance(self.base_layer.quant_method, UnquantizedLinearMethod) + return True + @torch.compile() def forward(self, x: torch.Tensor) -> torch.Tensor: lora_A = self.lora_A diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 32f754376..c74805aaf 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -1,13 +1,17 @@ import hashlib import importlib.util import os +from collections.abc import Iterable import torch import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file +from safetensors.torch import safe_open from safetensors.torch import save_file as safetensors_save_file +from torch.nn.utils import parametrize from sglang.multimodal_gen import envs +from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, @@ -52,6 +56,7 @@ from sglang.srt.model_loader.checkpoint_quantization import ( logger = init_logger(__name__) VAE_CHANNELS_LAST_3D_ENV = "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D" +_VAE_CHECKPOINT_ARCH_METADATA = ("latents_mean", "latents_std") def _require_native_loader_for_quantized_vae( @@ -293,6 +298,101 @@ def _match_checkpoint_dtypes(loaded: dict, target_state: dict) -> dict: return loaded +def _adopt_plain_weight_norm_state( + module: nn.Module, loaded_names: Iterable[str] +) -> int: + """Make deparameterized checkpoint weights native module state. + + PyTorch's weight-norm load hook accepts legacy ``weight_g``/``weight_v`` + tensors, while inference exports commonly fold those tensors into one + plain ``weight``. Removing only the matching parametrizations preserves + that already-computed weight exactly and leaves every other parameterized + module untouched. + """ + state_names = set(module.state_dict()) + module_by_name = dict(module.named_modules()) + owners: set[str] = set() + for name in loaded_names: + if name == "weight": + owner_name = "" + elif name.endswith(".weight"): + owner_name = name.removesuffix(".weight") + else: + continue + state_prefix = f"{owner_name}." if owner_name else "" + if { + f"{state_prefix}parametrizations.weight.original0", + f"{state_prefix}parametrizations.weight.original1", + }.issubset(state_names): + owners.add(owner_name) + + for owner_name in sorted(owners): + parametrize.remove_parametrizations( + module_by_name[owner_name], "weight", leave_parametrized=True + ) + return len(owners) + + +def _vae_checkpoint_arch_metadata_names( + vae_config: VAEConfig, + target_state: dict[str, torch.Tensor], +) -> tuple[str, ...]: + arch_values = vars(vae_config.arch_config) + return tuple( + name + for name in _VAE_CHECKPOINT_ARCH_METADATA + if name not in target_state and name in arch_values + ) + + +def _consume_vae_checkpoint_arch_metadata( + loaded: dict[str, torch.Tensor], + vae_config: VAEConfig, + target_state: dict[str, torch.Tensor], +) -> tuple[str, ...]: + """Move checkpoint-carried latent statistics into the VAE config.""" + arch_values = vars(vae_config.arch_config) + consumed = [] + for name in _vae_checkpoint_arch_metadata_names(vae_config, target_state): + tensor = loaded.get(name) + if tensor is None: + continue + if tensor.ndim != 1: + raise ValueError( + f"VAE checkpoint metadata {name!r} must be one-dimensional, " + f"got shape {tuple(tensor.shape)}" + ) + arch_values[name] = tensor.tolist() + del loaded[name] + consumed.append(name) + if consumed: + vae_config.post_init() + return tuple(consumed) + + +def _vae_checkpoint_tensor_names(weight_files: list[str]) -> set[str]: + names: set[str] = set() + for path in weight_files: + with safe_open(path, framework="pt", device="cpu") as checkpoint: + names.update(checkpoint.keys()) + return names + + +def _log_vae_checkpoint_adaptations( + num_deparameterized: int, consumed_metadata: tuple[str, ...] +) -> None: + if num_deparameterized: + logger.info( + "VAE: adopted %d deparameterized weight-normalized layers", + num_deparameterized, + ) + if consumed_metadata: + logger.info( + "VAE: loaded architecture metadata from checkpoint: %s", + ", ".join(consumed_metadata), + ) + + def _direct_gpu_vae_state_slots( vae: nn.Module, component_name: str ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[nn.Module, str, bool]]]: @@ -341,15 +441,24 @@ def _assign_direct_gpu_vae_state( *, component_name: str, device: torch.device, -) -> None: + vae_config: VAEConfig, +) -> tuple[int, tuple[str, ...]]: """Stream a complete standard VAE state directly onto its target device.""" + num_deparameterized = _adopt_plain_weight_norm_state( + vae, _vae_checkpoint_tensor_names(weight_files) + ) target_state, slots = _direct_gpu_vae_state_slots(vae, component_name) + metadata_names = _vae_checkpoint_arch_metadata_names(vae_config, target_state) loaded_names: set[str] = set() + metadata: dict[str, torch.Tensor] = {} with torch.no_grad(): for raw_name, tensor in safetensors_weights_iterator( weight_files, to_cpu=device.type == "cpu" ): name = raw_name + if name in metadata_names: + metadata[name] = tensor + continue if name in loaded_names: raise ComponentCheckpointUnsupportedError( f"Direct GPU VAE checkpoint maps multiple tensors to {name!r}" @@ -378,6 +487,9 @@ def _assign_direct_gpu_vae_state( module._buffers[local_name] = tensor loaded_names.add(name) + consumed_metadata = _consume_vae_checkpoint_arch_metadata( + metadata, vae_config, target_state + ) missing = sorted(set(slots) - loaded_names) if missing: raise ComponentCheckpointUnsupportedError( @@ -390,6 +502,7 @@ def _assign_direct_gpu_vae_state( raise RuntimeError( f"Direct GPU VAE loading left meta tensors: {remaining_meta}" ) + return num_deparameterized, consumed_metadata class VAELoader(WeightOverrideComponentLoader): @@ -506,12 +619,16 @@ class VAELoader(WeightOverrideComponentLoader): auto_map = config.get("auto_map", {}) auto_model_map = auto_map.get("AutoModel") - if direct_gpu_weight_loading and auto_model_map: + if direct_gpu_weight_loading and auto_model_map and not native_only: raise ComponentCheckpointUnsupportedError( f"Direct GPU loading for {component_name!r} requires a native " "ModelRegistry VAE; custom Diffusers auto_map code is unsupported" ) - if auto_model_map and component_weights_path != component_model_path: + if ( + auto_model_map + and not native_only + and component_weights_path != component_model_path + ): raise ComponentCheckpointUnsupportedError( f"{component_name!r} uses a custom Diffusers class that cannot " "consume a weights-only override" @@ -591,12 +708,14 @@ class VAELoader(WeightOverrideComponentLoader): f"Found no safetensors files in {component_weights_path}" ) if direct_gpu_weight_loading: - _assign_direct_gpu_vae_state( + adaptations = _assign_direct_gpu_vae_state( vae, safetensors_list, component_name=component_name, device=target_device, + vae_config=vae_config, ) + _log_vae_checkpoint_adaptations(*adaptations) if _should_use_channels_last_3d(server_args, component_name): n = _convert_conv3d_weights_to_channels_last_3d(vae) if n > 0: @@ -610,6 +729,12 @@ class VAELoader(WeightOverrideComponentLoader): for sf_path in safetensors_list: loaded.update(safetensors_load_file(sf_path)) _backfill_ltx2_audio_vae_latent_stats(loaded, component_type) + num_deparameterized = _adopt_plain_weight_norm_state(vae, loaded) + target_state = vae.state_dict() + consumed_metadata = _consume_vae_checkpoint_arch_metadata( + loaded, vae_config, target_state + ) + _log_vae_checkpoint_adaptations(num_deparameterized, consumed_metadata) strict_load = native_only # `loaded` holds views into the safetensors mapping. When the component # starts on the CPU and the host cannot afford copies of the whole @@ -635,7 +760,7 @@ class VAELoader(WeightOverrideComponentLoader): ) ) if keep_mapping: - _match_checkpoint_dtypes(loaded, vae.state_dict()) + _match_checkpoint_dtypes(loaded, target_state) vae.load_state_dict( loaded, strict=strict_load, diff --git a/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py b/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py index d314b8dc4..0a376ca0f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py +++ b/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py @@ -85,13 +85,14 @@ def validate_minimax_h3_checkpoint_variant( checkpoint_paths: list[str], selected_variant: str ) -> None: names = " ".join(path.lower() for path in checkpoint_paths) - checkpoint_variant = next( - (variant for variant in ("fl2va", "ref2va") if variant in names), None - ) + checkpoint_variants = { + variant for variant in ("fl2va", "ref2va") if variant in names + } if ( - checkpoint_variant is not None - and checkpoint_variant != selected_variant.lower() + len(checkpoint_variants) == 1 + and selected_variant.lower() not in checkpoint_variants ): + (checkpoint_variant,) = checkpoint_variants raise ValueError( f"MiniMax-H3 checkpoint variant {checkpoint_variant!r} does not match " f"--model-variant {selected_variant!r}" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index f4412c0f1..462307d7b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -600,6 +600,9 @@ class LoRAPipeline(ComposedPipelineBase): if merge_mode == "dynamic": return False uses_dtensor_weights = self._uses_dtensor_weights(lora_layers) + has_unmergeable_weights = any( + not layer.can_merge_base_weight for layer in lora_layers.values() + ) if merge_mode == "auto": if uses_dtensor_weights: logger.info( @@ -607,7 +610,18 @@ class LoRAPipeline(ComposedPipelineBase): module_name, ) return False + if has_unmergeable_weights: + logger.info( + "Using dynamic LoRA for %s because its quantized weights cannot be merged in place.", + module_name, + ) + return False return True + if has_unmergeable_weights: + raise ValueError( + f"LoRA merge mode is unavailable for {module_name} because its " + "quantized weights cannot be updated in place; use merge mode 'dynamic'" + ) if uses_dtensor_weights: logger.warning( "Merging LoRA for %s with FSDP-sharded weights may require full-gather and can OOM.", diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index 37b1ead5b..60695e263 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -956,7 +956,28 @@ def maybe_download_model( f"Cached ref for {model_name_or_path} is corrupt: resolved to the " f"snapshots parent {local_path!r} instead of a revision directory." ) - if not force_diffusers_model: + required_files = [ + pattern for pattern in allow_patterns or () if not glob.has_magic(pattern) + ] + missing_required_files = [ + path + for path in required_files + if not os.path.isfile(os.path.join(local_path, path)) + ] + if missing_required_files: + if not download: + raise ValueError( + f"Model {model_name_or_path} is cached but is missing requested " + f"files: {missing_required_files}." + ) + logger.info( + "Cached snapshot for %s is missing requested files %s; " + "will download them from %s", + model_name_or_path, + missing_required_files, + _model_hub_name(), + ) + elif not force_diffusers_model: # maybe_download_model_index's model_index.json fetch materializes a full # cache entry, so this resolve reports that stub as a hit; returning it # would skip the download. LoRA repos declare no components. diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 634f9c186..2570c6bd9 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -2713,7 +2713,7 @@ "expected_e2e_ms": 5648.49, "expected_avg_denoise_ms": 477.49, "expected_median_denoise_ms": 56.74, - "load_peak_vram_mb": 32892.0, + "load_peak_vram_mb": 34022.0, "runtime_peak_vram_mb": 61190.0, "estimated_full_test_time_s": 153.1 }, @@ -2731,7 +2731,7 @@ "expected_e2e_ms": 9086.81, "expected_avg_denoise_ms": 492.02, "expected_median_denoise_ms": 151.9, - "load_peak_vram_mb": 32892.0, + "load_peak_vram_mb": 34022.0, "runtime_peak_vram_mb": 62510.0, "estimated_full_test_time_s": 149.4 }, diff --git a/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py b/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py index bffa1e2f8..577c6ff15 100644 --- a/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py @@ -298,6 +298,46 @@ def test_metadata_only_cached_lora_snapshot_is_a_usable_hit( assert calls == ["probe"] +def test_cached_lora_snapshot_downloads_missing_selected_weight(monkeypatch, tmp_path): + calls = [] + selected_file = "loras/adapter.safetensors" + + def fake_snapshot_download(**kwargs): + calls.append("probe" if kwargs.get("local_files_only") else "download") + if not kwargs.get("local_files_only"): + target = tmp_path / selected_file + target.parent.mkdir() + target.write_bytes(b"weights") + return str(tmp_path) + + monkeypatch.setattr(hf_diffusers_utils, "snapshot_download", fake_snapshot_download) + + result = maybe_download_model( + "org/repo", + is_lora=True, + allow_patterns=["*.json", selected_file], + ) + + assert result == str(tmp_path) + assert calls == ["probe", "download"] + + +def test_cached_lora_snapshot_reports_missing_selected_weight_offline( + recording_snapshot_download, tmp_path +): + calls = recording_snapshot_download(tmp_path) + + with pytest.raises(ValueError, match="loras/adapter.safetensors"): + maybe_download_model( + "org/repo", + download=False, + is_lora=True, + allow_patterns=["*.json", "loras/adapter.safetensors"], + ) + + assert calls == ["probe"] + + def test_force_diffusers_model_stub_keeps_its_existing_path( recording_snapshot_download, tmp_path ): diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py index 28b2a2340..546f3ea5c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py @@ -3,13 +3,16 @@ from contextlib import contextmanager, nullcontext from types import SimpleNamespace from unittest.mock import patch +import pytest import torch +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.lora.linear import ( BaseLayerWithLoRA, _use_owned_base_snapshot, wrap_with_lora_layer, ) +from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config from sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_lora @@ -94,6 +97,30 @@ def test_zero_copy_snapshot_is_limited_to_cpu_backed_layers(): assert meta_layer._base_is_view +def test_quantized_base_uses_dynamic_lora_in_auto_mode(): + with patch( + "sglang.multimodal_gen.runtime.layers.quantization.fp8." + "get_tensor_model_parallel_world_size", + return_value=1, + ): + base_layer = ReplicatedLinear( + 2, + 2, + bias=False, + quant_config=Fp8Config(is_checkpoint_fp8_serialized=True), + ) + layer = BaseLayerWithLoRA(base_layer) + pipeline = _make_pipeline(layer) + + assert not pipeline._should_merge_lora_for_layers( + "transformer", {"linear": layer}, "auto" + ) + with pytest.raises(ValueError, match="use merge mode 'dynamic'"): + pipeline._should_merge_lora_for_layers( + "transformer", {"linear": layer}, "merge" + ) + + def test_dynamic_lora_reactivates_cached_layers_without_weight_update_context(): layer = _make_layer() pipeline = _make_pipeline(layer) diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index c1f07cd83..39ac4992d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -93,6 +93,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader i from sglang.multimodal_gen.runtime.loader.minimax_h3_weights import ( inspect_minimax_h3_safetensors, resolve_minimax_h3_checkpoint_quantization, + validate_minimax_h3_checkpoint_variant, ) from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( TransformerQuantLoadSpec, @@ -691,6 +692,18 @@ class TestTransformerQuantHelpers(unittest.TestCase): ) ) + def test_minimax_h3_hybrid_checkpoint_accepts_selected_partition(self): + checkpoint = "/cache/minimax_h3_hybrid_fl2va_ref2va_b25-49.safetensors" + + validate_minimax_h3_checkpoint_variant([checkpoint], "fl2va") + validate_minimax_h3_checkpoint_variant([checkpoint], "ref2va") + + def test_minimax_h3_single_partition_checkpoint_rejects_mismatch(self): + with self.assertRaisesRegex(ValueError, "does not match"): + validate_minimax_h3_checkpoint_variant( + ["/cache/minimax_h3_fl2va.safetensors"], "ref2va" + ) + def test_inspect_minimax_h3_safetensors_detects_curve_and_comfy_format(self): marker = json.dumps( { diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index 6dfae598f..4bd8cbe9a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -7,6 +7,9 @@ import torch import torch.nn as nn from safetensors.torch import save_file as safetensors_save_file +from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import ( + MiniMaxH3AudioVAEConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, @@ -24,8 +27,10 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp ComponentCheckpointUnsupportedError, ) from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import ( + _adopt_plain_weight_norm_state, _assign_direct_gpu_vae_state, _backfill_ltx2_audio_vae_latent_stats, + _consume_vae_checkpoint_arch_metadata, _direct_gpu_vae_state_slots, _match_checkpoint_dtypes, _require_native_loader_for_quantized_vae, @@ -147,6 +152,54 @@ class TestMatchCheckpointDtypes(unittest.TestCase): self.assertIs(loaded["extra"], before) +class TestPlainWeightNormCheckpoint(unittest.TestCase): + def test_adopts_a_folded_weight_without_reconstructing_it(self): + module = nn.Sequential( + torch.nn.utils.parametrizations.weight_norm( + nn.Conv1d(2, 3, kernel_size=3, bias=False) + ) + ) + expected = torch.arange(18, dtype=torch.float32).reshape(3, 2, 3) / 19 + loaded = {"0.weight": expected} + + self.assertEqual(_adopt_plain_weight_norm_state(module, loaded), 1) + module.load_state_dict(loaded, strict=True) + + self.assertEqual(set(module.state_dict()), {"0.weight"}) + self.assertTrue(torch.equal(module[0].weight, expected)) + + def test_keeps_legacy_weight_norm_state_parameterized(self): + module = nn.Sequential( + torch.nn.utils.parametrizations.weight_norm( + nn.Conv1d(2, 3, kernel_size=3, bias=False) + ) + ) + original_state = module.state_dict() + loaded = { + "0.weight_g": original_state["0.parametrizations.weight.original0"].clone(), + "0.weight_v": original_state["0.parametrizations.weight.original1"].clone(), + } + + self.assertEqual(_adopt_plain_weight_norm_state(module, loaded), 0) + module.load_state_dict(loaded, strict=True) + + self.assertIn("0.parametrizations.weight.original0", module.state_dict()) + + def test_moves_checkpoint_latent_stats_into_arch_config(self): + config = MiniMaxH3AudioVAEConfig() + loaded = { + "latents_mean": torch.arange(32, dtype=torch.float32), + "latents_std": torch.arange(1, 33, dtype=torch.float32), + } + + consumed = _consume_vae_checkpoint_arch_metadata(loaded, config, {}) + + self.assertEqual(consumed, ("latents_mean", "latents_std")) + self.assertEqual(config.arch_config.latents_mean, list(range(32))) + self.assertEqual(config.arch_config.latents_std, list(range(1, 33))) + self.assertEqual(loaded, {}) + + class TestDirectGPUVAEState(unittest.TestCase): class _StandardVAE(nn.Module): def __init__(self, *_args, **_kwargs): @@ -169,12 +222,49 @@ class TestDirectGPUVAEState(unittest.TestCase): [str(checkpoint)], component_name="vae", device=torch.device("cpu"), + vae_config=QwenImagePipelineConfig().vae_config, ) self.assertTrue(torch.equal(vae.proj.weight, expected_weight)) self.assertTrue(torch.equal(vae.scale, expected_scale)) self.assertFalse(any(tensor.is_meta for tensor in vae.state_dict().values())) + def test_direct_load_adopts_folded_weight_norm_and_checkpoint_metadata(self): + class _WeightNormVAE(nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.utils.parametrizations.weight_norm( + nn.Conv1d(2, 3, kernel_size=3, bias=False) + ) + + expected = torch.arange(18, dtype=torch.float32).reshape(3, 2, 3) / 19 + config = MiniMaxH3AudioVAEConfig() + with TemporaryDirectory() as root: + checkpoint = pathlib.Path(root) / "model.safetensors" + safetensors_save_file( + { + "proj.weight": expected, + "latents_mean": torch.arange(32, dtype=torch.float32), + "latents_std": torch.arange(1, 33, dtype=torch.float32), + }, + checkpoint, + ) + with torch.device("meta"): + vae = _WeightNormVAE() + adaptations = _assign_direct_gpu_vae_state( + vae, + [str(checkpoint)], + component_name="audio_vae", + device=torch.device("cpu"), + vae_config=config, + ) + + self.assertEqual(adaptations, (1, ("latents_mean", "latents_std"))) + self.assertEqual(set(vae.state_dict()), {"proj.weight"}) + self.assertTrue(torch.equal(vae.proj.weight, expected)) + self.assertEqual(config.arch_config.latents_mean, list(range(32))) + self.assertEqual(config.arch_config.latents_std, list(range(1, 33))) + def test_rejects_nonstandard_state_lifecycle(self): class _CustomVAE(self._StandardVAE): def state_dict(self, *args, **kwargs): @@ -243,6 +333,59 @@ class TestDirectGPUVAEState(unittest.TestCase): self.assertTrue(torch.equal(loaded.proj.weight, expected_weight)) self.assertTrue(torch.equal(loaded.scale, expected_scale)) + def test_native_vae_ignores_diffusers_auto_map_for_weight_override(self): + loader = vae_loader.VAELoader() + expected_weight = torch.arange(4, dtype=torch.bfloat16).reshape(2, 2) + expected_scale = torch.tensor([3.0], dtype=torch.bfloat16) + + with TemporaryDirectory() as root: + checkpoint = pathlib.Path(root) / "override.safetensors" + safetensors_save_file( + {"proj.weight": expected_weight, "scale": expected_scale}, checkpoint + ) + + for direct_gpu_loading in (False, True): + with self.subTest(direct_gpu_loading=direct_gpu_loading): + pipeline_config = QwenImagePipelineConfig() + pipeline_config.native_only_components = ("vae",) + server_args = _FakeServerArgs(pipeline_config) + server_args.component_direct_gpu_weight_loading = { + "vae": direct_gpu_loading + } + + with ( + patch.object( + vae_loader, + "get_diffusers_component_config", + return_value={ + "_class_name": "TestVAE", + "auto_map": {"AutoModel": "custom.TestVAE"}, + }, + ), + patch.object( + loader, + "resolve_component_weights_path", + return_value=str(checkpoint), + ), + patch.object( + vae_loader.ModelRegistry, + "resolve_model_cls", + return_value=(self._StandardVAE, None), + ), + patch.object( + loader, "target_device", return_value=torch.device("cpu") + ), + patch.object( + vae_loader.current_platform, + "optimize_vae", + side_effect=lambda vae: vae, + ), + ): + loaded = loader.load_customized(root, server_args, "vae") + + self.assertTrue(torch.equal(loaded.proj.weight, expected_weight)) + self.assertTrue(torch.equal(loaded.scale, expected_scale)) + def test_quantized_checkpoint_does_not_fall_back_from_direct_loading(self): loader = vae_loader.VAELoader() server_args = _FakeServerArgs(QwenImagePipelineConfig())