Skip to content

fix: prevent silent double-application of LoRA weights #473

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
WangXukang-cypher wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix: prevent silent double-application of LoRA weights #473

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ed85d6f
fix: prevent silent double-application of LoRA weights
WangXukang-cypher Sep 4, 2026
4962a21
chore: retrigger CLA check
WangXukang-cypher Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/maxdiffusion/loaders/lora_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,17 @@ class LoRABaseMixin:
_lora_lodable_modules = []
num_fused_loras = 0

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
Comment on lines +27 to +36

@gemini-code-assist gemini-code-assist Bot Sep 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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


def load_lora_weights(self, **kwargs):
raise NotImplementedError("`load_lora_weights()` is not implemented.")

Expand Down
5 changes: 5 additions & 0 deletions src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,11 @@ def translate_fn(nnx_path_str):
max_logging.log("No LoRA weight name provided; skipping LoRA load.")
return pipeline

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

h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs)
transformer_state_dict = {}
connector_state_dict = {}
Expand Down
15 changes: 13 additions & 2 deletions src/maxdiffusion/loaders/wan_lora_nnx_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ def load_lora_weights(
def translate_fn(nnx_path_str):
return lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)

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)
Comment on lines +53 to 60

@gemini-code-assist gemini-code-assist Bot Sep 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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)

Expand Down Expand Up @@ -91,7 +96,10 @@ def translate_fn(nnx_path_str: str):
return lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)

# Handle high noise model
if hasattr(pipeline, "high_noise_transformer") and high_noise_weight_name:
high_key = (lora_model_path, high_noise_weight_name, "high_noise")
if self._check_and_record_lora(high_key):
max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged into high_noise_transformer — skipping.")
elif 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)
Expand All @@ -100,7 +108,10 @@ def translate_fn(nnx_path_str: str):
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:
low_key = (lora_model_path, low_noise_weight_name, "low_noise")
if self._check_and_record_lora(low_key):
max_logging.log(f"WARNING: LoRA '{lora_model_path}' already merged into low_noise_transformer — skipping.")
elif 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)
Expand Down
Loading