fix: prevent silent double-application of LoRA weights

[WangXukang-cypher]

Summary

• Wire up the previously dead num_fused_loras counter in LoRABaseMixin to track merged LoRA identities
• Before each merge, the loader checks whether the same (path, weight_name) was already applied and skips with a warning
• Covers all three NNX loaders: Wan2_1NNXLoraLoader, Wan2_2NNXLoraLoader, LTX2NNXLoraLoader

Problem

merge_lora unconditionally adds delta to model weights (kernel += delta). Calling load_lora_weights twice with the same LoRA — via duplicate config entries or notebook cell re-execution — silently doubles the LoRA effect. The existing num_fused_loras = 0 counter in LoRABaseMixin was never incremented or checked.

Test plan

• Pre-commit passes (ruff, pyink, pylint ≥ 7)
• Verified: first merge allowed, duplicate blocked, different LoRA allowed, different weight name allowed, independent loader instances isolated, class-level counter unchanged
• Existing CI tests (require TPU environment)

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to prevent duplicate LoRA weight applications by tracking merged LoRA keys within the LoRABaseMixin class and checking them in the LTX2 and Wan LoRA loaders. The review feedback highlights two important issues: first, defining __init__ in a mixin class can cause runtime issues, so lazy initialization of the tracking set is recommended; second, checking and recording the LoRA key before verifying if the pipeline actually has the target transformer attributes can lead to incorrect tracking state, so the check should be deferred until after those attributes are verified.

[gemini-code-assist]

Defining __init__ in a mixin class like LoRABaseMixin can easily lead to AttributeError at runtime if subclasses or multiple inheritance hierarchies do not properly call super().__init__().\n\nA much safer and more idiomatic approach for Python mixins is to use lazy initialization inside the helper method (e.g., checking hasattr(self, '_fused_lora_keys')). This completely avoids the need for __init__ in the mixin.

Suggested change

	def __init__(self):
	self._fused_lora_keys = set()
	def _check_and_record_lora(self, lora_key):
	"""Return True if this LoRA was already merged (duplicate). Records it otherwise."""
	if lora_key in self._fused_lora_keys:
	return True
	self._fused_lora_keys.add(lora_key)
	self.num_fused_loras += 1
	return False
	def _check_and_record_lora(self, lora_key):
	"""Return True if this LoRA was already merged (duplicate). Records it otherwise."""
	if not hasattr(self, "_fused_lora_keys"):
	self._fused_lora_keys = set()
	if lora_key in self._fused_lora_keys:
	return True
	self._fused_lora_keys.add(lora_key)
	self.num_fused_loras += 1
	return False

[gemini-code-assist]

Checking and recording the LoRA key before verifying if transformer_weight_name is provided and if the pipeline has the transformer attribute causes num_fused_loras to be incorrectly incremented and a dummy key to be recorded even when no LoRA is actually merged.\n\nWe should only check and record the LoRA key if we are actually going to attempt to merge it.

Suggested change

	lora_key = (lora_model_path, transformer_weight_name)
	if self._check_and_record_lora(lora_key):
	max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged — skipping to avoid double-application.")
	return pipeline
	if hasattr(pipeline, "transformer") and transformer_weight_name:
	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)
	if hasattr(pipeline, "transformer") and transformer_weight_name:
	lora_key = (lora_model_path, transformer_weight_name)
	if self._check_and_record_lora(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)

[Perseus14]

Thanks for tackling this, @WangXukang-cypher!

The general direction makes a lot of sense, but there are a few edge cases around state ownership, error handling, and Python mixin design that we should address before merging:

1. Where state should live (Pipeline vs. Loader): In MaxDiffusion, loaders like Wan2_1NNXLoraLoader are typically instantiated as transient helpers (e.g. Wan2_1NNXLoraLoader().load_lora_weights(...)). If a user re-runs a notebook cell that creates a new loader instance, self._fused_lora_keys resets and the weights can still be doubly merged. Conversely, sharing one loader across different pipelines might inadvertently block the second pipeline from loading its LoRA. Attaching the tracked keys to the pipeline (or module) instead of the loader solves both issues cleanly.
2. Atomic registration: Currently, _check_and_record_lora records the LoRA as merged before we know whether downloading/reading the weights succeeds or if the model has matching layers. Splitting this into a pre-check (to skip if already present) and a post-merge record (only after merge_fn finishes) makes the loader resilient to network glitches and file errors.
3. Input handling: In maxdiffusion, LoRA sources can sometimes be passed as in-memory state dicts (dict), which aren't hashable in a set. We should handle this case safely.
4. Tests: Could you add a couple of unit tests in src/maxdiffusion/tests/ to verify duplicate skipping and ensure retry/fresh pipeline behavior works as expected?

I’ve left some inline suggestions with code snippets to help guide the refactor. Happy to discuss further!

[Perseus14]

Defining an __init__ in a mixin class can be tricky in Python because mixins participate in multiple inheritance. If a subclass doesn't explicitly call super().__init__(), self._fused_lora_keys won't be initialized and might raise an AttributeError.

More importantly, since users often instantiate loaders on the fly (e.g., lora_loader = Wan2_1NNXLoraLoader()), storing the set on self means creating a new loader in a notebook cell will reset the set to empty and won't guard against cell re-runs.

A clean alternative is to store the tracking set on the pipeline instance being modified (for example, getattr(pipeline, "_fused_lora_keys", set())). That way, the tracked state always travels with the actual weights being mutated.

[Perseus14]

It’s safer to split the "check" and "record" steps into two separate methods here:

1. Check first: Before loading, check if the LoRA is already merged. If so, log and return early.
2. Record after success: Only add to the set and increment num_fused_loras after the weights are successfully loaded and merged via merge_fn.

If we record beforehand and lora_state_dict() raises an exception (like a network timeout or file not found), the key remains marked as fused. If the user catches the error or retries, it would be skipped on the second attempt.

Also, keep in mind that lora_model_path can sometimes be a dict (in-memory weights). Since dictionaries are unhashable, adding a tuple containing a dict to a set will raise a TypeError. We could fall back to id(lora_model_path) or a custom string tag if it's a dict.

[Perseus14]

We should move this check inside the if hasattr(pipeline, "transformer") and transformer_weight_name: block.

If transformer_weight_name is empty or pipeline doesn't have a transformer, we don't want to record the key or increment the fused count since no weights were actually modified.

Minor tip: It's also helpful to run os.path.normpath or os.path.abspath on lora_model_path (when it's a string path) so that ./checkpoints/my_lora and checkpoints/my_lora are recognized as the same file.

[Perseus14]

Check the order of conditions here: if high_noise_weight_name is None (for instance, if a user only wants to condition the low-noise stage), high_key will be (path, None, "high_noise"). Calling _check_and_record_lora will record this key and increment num_fused_loras, and then fall into the else branch without merging anything.

If we nest the duplicate check inside the if hasattr(...) and weight_name: branch, we'll ensure we only check and record when there are valid weights to apply.

[Perseus14]

Same as above

[Perseus14]

Similar to Wan, let's defer recording the key until after the weights have been applied to pipeline.transformer and/or pipeline.connectors. That keeps the tracking state accurate even if a file load fails midway.

[WangXukang-cypher]

Thanks for the detailed review, @Perseus14! I've updated the PR in 573fe4c based on your suggestions:

• Moved fused-LoRA tracking from transient loader instances to the pipeline that owns the mutated weights. This preserves duplicate protection across newly created loaders while keeping separate pipelines independent.
• Split duplicate checking from registration and now record a LoRA only after its load and merge complete successfully.
• Moved the checks inside the valid-target branches, so missing transformers or weight names are not recorded.
• Normalized path-based keys and added safe identity-based keys for in-memory state dictionaries.
• Applied the same post-merge registration behavior to WAN 2.1, WAN 2.2, and LTX2.
• Added unit tests covering duplicate skipping across loader instances, fresh-pipeline behavior, retry after load/merge failures, in-memory dictionaries, and missing WAN 2.2 weight names.

Thanks again for pointing out these edge cases. Please take another look when you have a chance!