fix: prevent silent double-application of LoRA weights - #473
fix: prevent silent double-application of LoRA weights#473WangXukang-cypher wants to merge 2 commits into
Conversation
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 if so. This prevents accidental weight corruption when users list duplicate LoRA paths in config or re-run merge calls in notebooks.
|
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[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
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.
| 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 |
| 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) |
There was a problem hiding this comment.
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.
| 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) |
1 participant
Summary
num_fused_lorascounter inLoRABaseMixinto track merged LoRA identities(path, weight_name)was already applied and skips with a warningWan2_1NNXLoraLoader,Wan2_2NNXLoraLoader,LTX2NNXLoraLoaderProblem
merge_loraunconditionally addsdeltato model weights (kernel += delta). Callingload_lora_weightstwice with the same LoRA — via duplicate config entries or notebook cell re-execution — silently doubles the LoRA effect. The existingnum_fused_loras = 0counter inLoRABaseMixinwas never incremented or checked.Test plan