diff --git a/src/maxdiffusion/loaders/lora_base.py b/src/maxdiffusion/loaders/lora_base.py index f22696d3c..ef35e9f85 100644 --- a/src/maxdiffusion/loaders/lora_base.py +++ b/src/maxdiffusion/loaders/lora_base.py @@ -12,17 +12,51 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ..models.modeling_utils import load_state_dict -from ..utils import _get_model_file +import os import safetensors +from ..models.modeling_utils import load_state_dict +from ..utils import _get_model_file + class LoRABaseMixin: """Utility class for handing LoRAs""" _lora_lodable_modules = [] - num_fused_loras = 0 + + @staticmethod + def _lora_key(lora_model_path, weight_name, target=None): + """Build a hashable key for a LoRA source and merge target.""" + if isinstance(lora_model_path, dict): + source_key = ("state_dict", id(lora_model_path)) + elif isinstance(lora_model_path, (str, os.PathLike)): + source_key = os.path.normpath(os.fspath(lora_model_path)) + else: + try: + hash(lora_model_path) + source_key = lora_model_path + except TypeError: + source_key = (type(lora_model_path).__name__, id(lora_model_path)) + + return (source_key, weight_name, target) + + @staticmethod + def _is_lora_fused(pipeline, lora_key): + """Return whether a LoRA has already been merged into this pipeline.""" + fused_lora_keys = getattr(pipeline, "_fused_lora_keys", None) + return fused_lora_keys is not None and lora_key in fused_lora_keys + + @staticmethod + def _record_lora_fused(pipeline, lora_key): + """Record a successful merge on the pipeline that owns the mutated weights.""" + fused_lora_keys = getattr(pipeline, "_fused_lora_keys", None) + if fused_lora_keys is None: + fused_lora_keys = set() + pipeline._fused_lora_keys = fused_lora_keys + if lora_key not in fused_lora_keys: + fused_lora_keys.add(lora_key) + pipeline.num_fused_loras = getattr(pipeline, "num_fused_loras", 0) + 1 def load_lora_weights(self, **kwargs): raise NotImplementedError("`load_lora_weights()` is not implemented.") @@ -46,6 +80,7 @@ def _fetch_state_dict( ): from .lora_pipeline import LORA_WEIGHT_NAME_SAFE + state_dict = pretrained_model_name_or_path_or_dict model_file = None if not isinstance(pretrained_model_name_or_path_or_dict, dict): # Let's first try to load .safetensors weights diff --git a/src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py b/src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py index a3c4d0d38..6bd30a21c 100644 --- a/src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py +++ b/src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py @@ -54,6 +54,16 @@ def translate_fn(nnx_path_str): max_logging.log("No LoRA weight name provided; skipping LoRA load.") return pipeline + merge_targets_available = hasattr(pipeline, "transformer") or hasattr(pipeline, "connectors") + if not merge_targets_available: + max_logging.log("Neither transformer nor connectors found; skipping LoRA load.") + return pipeline + + lora_key = self._lora_key(lora_model_path, transformer_weight_name, "ltx2_pipeline") + if self._is_lora_fused(pipeline, lora_key): + max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged; skipping to avoid double-application.") + return pipeline + h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs) transformer_state_dict = {} connector_state_dict = {} @@ -78,4 +88,5 @@ def translate_fn(nnx_path_str): if unmatched_keys: max_logging.log(f"{len(unmatched_keys)} key(s) in LoRA dictionary routed to no merge target: {unmatched_keys}") + self._record_lora_fused(pipeline, lora_key) return pipeline diff --git a/src/maxdiffusion/loaders/wan_lora_nnx_loader.py b/src/maxdiffusion/loaders/wan_lora_nnx_loader.py index a34c0f1a1..9ed77428c 100644 --- a/src/maxdiffusion/loaders/wan_lora_nnx_loader.py +++ b/src/maxdiffusion/loaders/wan_lora_nnx_loader.py @@ -51,10 +51,16 @@ def translate_fn(nnx_path_str): return lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers) if hasattr(pipeline, "transformer") and transformer_weight_name: + lora_key = self._lora_key(lora_model_path, transformer_weight_name, "transformer") + if self._is_lora_fused(pipeline, lora_key): + max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged; skipping to avoid double-application.") + return pipeline + max_logging.log(f"Merging LoRA into transformer with rank={rank}") h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs) h_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(h_state_dict) merge_fn(pipeline.transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype) + self._record_lora_fused(pipeline, lora_key) else: max_logging.log("transformer not found or no weight name provided for LoRA.") @@ -92,19 +98,29 @@ def translate_fn(nnx_path_str: str): # Handle high noise model if hasattr(pipeline, "high_noise_transformer") and high_noise_weight_name: - max_logging.log(f"Merging LoRA into high_noise_transformer with rank={rank}") - h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=high_noise_weight_name, **kwargs) - h_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(h_state_dict) - merge_fn(pipeline.high_noise_transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype) + high_key = self._lora_key(lora_model_path, high_noise_weight_name, "high_noise_transformer") + if self._is_lora_fused(pipeline, high_key): + max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged into high_noise_transformer; skipping.") + else: + max_logging.log(f"Merging LoRA into high_noise_transformer with rank={rank}") + h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=high_noise_weight_name, **kwargs) + h_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(h_state_dict) + merge_fn(pipeline.high_noise_transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype) + self._record_lora_fused(pipeline, high_key) else: max_logging.log("high_noise_transformer not found or no weight name provided for LoRA.") # Handle low noise model if hasattr(pipeline, "low_noise_transformer") and low_noise_weight_name: - max_logging.log(f"Merging LoRA into low_noise_transformer with rank={rank}") - l_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=low_noise_weight_name, **kwargs) - l_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(l_state_dict) - merge_fn(pipeline.low_noise_transformer, l_state_dict, rank, scale, translate_fn, dtype=dtype) + low_key = self._lora_key(lora_model_path, low_noise_weight_name, "low_noise_transformer") + if self._is_lora_fused(pipeline, low_key): + max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged into low_noise_transformer; skipping.") + else: + max_logging.log(f"Merging LoRA into low_noise_transformer with rank={rank}") + l_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=low_noise_weight_name, **kwargs) + l_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(l_state_dict) + merge_fn(pipeline.low_noise_transformer, l_state_dict, rank, scale, translate_fn, dtype=dtype) + self._record_lora_fused(pipeline, low_key) else: max_logging.log("low_noise_transformer not found or no weight name provided for LoRA.") diff --git a/src/maxdiffusion/tests/lora_nnx_loader_test.py b/src/maxdiffusion/tests/lora_nnx_loader_test.py new file mode 100644 index 000000000..ae09cc1f9 --- /dev/null +++ b/src/maxdiffusion/tests/lora_nnx_loader_test.py @@ -0,0 +1,131 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest import mock + +import pytest + +from maxdiffusion.loaders.ltx2_lora_nnx_loader import LTX2NNXLoraLoader +from maxdiffusion.loaders.wan_lora_nnx_loader import Wan2_1NNXLoraLoader, Wan2_2NNXLoraLoader + + +@pytest.fixture(name="wan_lora_mocks") +def fixture_wan_lora_mocks(): + with ( + mock.patch( + "maxdiffusion.loaders.wan_lora_nnx_loader.StableDiffusionLoraLoaderMixin.lora_state_dict", + return_value=({"lora.weight": object()}, None), + ) as state_dict_mock, + mock.patch( + "maxdiffusion.loaders.wan_lora_nnx_loader.lora_conversion_utils.preprocess_wan_lora_dict", + side_effect=lambda state_dict: state_dict, + ), + mock.patch("maxdiffusion.loaders.wan_lora_nnx_loader.lora_nnx.merge_lora") as merge_mock, + ): + yield state_dict_mock, merge_mock + + +def test_duplicate_is_skipped_across_loader_instances_but_not_pipelines(wan_lora_mocks): + state_dict_mock, merge_mock = wan_lora_mocks + pipeline = SimpleNamespace(transformer=object()) + + Wan2_1NNXLoraLoader().load_lora_weights(pipeline, "./checkpoints/lora", "weights.safetensors", rank=4) + Wan2_1NNXLoraLoader().load_lora_weights(pipeline, "checkpoints/lora", "weights.safetensors", rank=4) + + assert state_dict_mock.call_count == 1 + assert merge_mock.call_count == 1 + assert pipeline.num_fused_loras == 1 + + fresh_pipeline = SimpleNamespace(transformer=object()) + Wan2_1NNXLoraLoader().load_lora_weights(fresh_pipeline, "checkpoints/lora", "weights.safetensors", rank=4) + + assert state_dict_mock.call_count == 2 + assert merge_mock.call_count == 2 + assert fresh_pipeline.num_fused_loras == 1 + + +def test_failed_load_is_not_recorded_and_can_be_retried(wan_lora_mocks): + state_dict_mock, merge_mock = wan_lora_mocks + state_dict_mock.side_effect = [OSError("temporary download failure"), ({"lora.weight": object()}, None)] + pipeline = SimpleNamespace(transformer=object()) + loader = Wan2_1NNXLoraLoader() + + with pytest.raises(OSError, match="temporary download failure"): + loader.load_lora_weights(pipeline, "checkpoints/lora", "weights.safetensors", rank=4) + + assert getattr(pipeline, "_fused_lora_keys", set()) == set() + assert getattr(pipeline, "num_fused_loras", 0) == 0 + + loader.load_lora_weights(pipeline, "checkpoints/lora", "weights.safetensors", rank=4) + + assert state_dict_mock.call_count == 2 + assert merge_mock.call_count == 1 + assert pipeline.num_fused_loras == 1 + + +def test_failed_ltx2_merge_is_not_recorded_and_can_be_retried(): + pipeline = SimpleNamespace(transformer=object()) + loader = LTX2NNXLoraLoader() + with ( + mock.patch( + "maxdiffusion.loaders.ltx2_lora_nnx_loader.StableDiffusionLoraLoaderMixin.lora_state_dict", + return_value=({"diffusion_model.lora.weight": object()}, None), + ) as state_dict_mock, + mock.patch( + "maxdiffusion.loaders.ltx2_lora_nnx_loader.lora_nnx.merge_lora", + side_effect=[RuntimeError("merge failed"), None], + ) as merge_mock, + ): + with pytest.raises(RuntimeError, match="merge failed"): + loader.load_lora_weights(pipeline, "checkpoints/lora", "weights.safetensors", rank=4) + + assert getattr(pipeline, "_fused_lora_keys", set()) == set() + assert getattr(pipeline, "num_fused_loras", 0) == 0 + + loader.load_lora_weights(pipeline, "checkpoints/lora", "weights.safetensors", rank=4) + + assert state_dict_mock.call_count == 2 + assert merge_mock.call_count == 2 + assert pipeline.num_fused_loras == 1 + + +def test_in_memory_state_dict_has_a_stable_hashable_key(wan_lora_mocks): + state_dict_mock, merge_mock = wan_lora_mocks + lora_state_dict = {"lora.weight": object()} + pipeline = SimpleNamespace(transformer=object()) + + Wan2_1NNXLoraLoader().load_lora_weights(pipeline, lora_state_dict, "weights.safetensors", rank=4) + Wan2_1NNXLoraLoader().load_lora_weights(pipeline, lora_state_dict, "weights.safetensors", rank=4) + + assert state_dict_mock.call_count == 1 + assert merge_mock.call_count == 1 + + +def test_missing_wan22_weight_name_is_not_recorded(wan_lora_mocks): + state_dict_mock, merge_mock = wan_lora_mocks + pipeline = SimpleNamespace(high_noise_transformer=object(), low_noise_transformer=object()) + + Wan2_2NNXLoraLoader().load_lora_weights( + pipeline, + "checkpoints/lora", + high_noise_weight_name=None, + low_noise_weight_name=None, + rank=4, + ) + + state_dict_mock.assert_not_called() + merge_mock.assert_not_called() + assert getattr(pipeline, "_fused_lora_keys", set()) == set() + assert getattr(pipeline, "num_fused_loras", 0) == 0