fix: prevent silent double-application of LoRA weights

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

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

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.")
 

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