From dcabf3f9fd4b4c37ff113a7945afd21dba868108 Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Sat, 3 Oct 2026 17:28:16 -0500 Subject: [PATCH] [LoRA] apply `.alpha` in lora_A/lora_B Qwen-Image LoRAs The Qwen converter dropped per-module `.alpha` when keys were already named lora_A/lora_B, so LoRAs with alpha != rank loaded at the wrong strength (e.g. e-n-v-y/Qwen-Image-2.1-Fix-v2.0 at about a third). Fold it into the weights as the lora_down/lora_up path already does. Co-Authored-By: Claude Opus 5.5 --- .../loaders/lora_conversion_utils.py | 43 ++++++++++++------- tests/pipelines/qwenimage/test_qwenimage.py | 22 ++++++++++ 2 files changed, 49 insertions(+), 16 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index cef2e88454a8..067f7f143f14 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2323,19 +2323,18 @@ def convert_key(key: str) -> str: has_non_diffusers_lora_id = any(down_key in k or up_key in k for k in all_keys) has_diffusers_lora_id = any(a_key in k or b_key in k for k in all_keys) - if has_non_diffusers_lora_id: - - def get_alpha_scales(down_weight, alpha_key): - rank = down_weight.shape[0] - alpha = state_dict.pop(alpha_key).item() - scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - return scale_down, scale_up + def get_alpha_scales(down_weight, alpha_key): + rank = down_weight.shape[0] + alpha = state_dict.pop(alpha_key).item() + scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here + scale_down = scale + scale_up = 1.0 + while scale_down * 2 < scale_up: + scale_down *= 2 + scale_up /= 2 + return scale_down, scale_up + if has_non_diffusers_lora_id: for k in all_keys: if k.endswith(down_key): diffusers_down_key = k.replace(down_key, ".lora_A.weight") @@ -2348,12 +2347,24 @@ def get_alpha_scales(down_weight, alpha_key): converted_state_dict[diffusers_down_key] = down_weight * scale_down converted_state_dict[diffusers_up_key] = up_weight * scale_up - # Already in diffusers format (lora_A/lora_B), just pop + # Already in diffusers format (lora_A/lora_B). ai-toolkit and ComfyUI LoRAs can still carry a per-module + # `.alpha`, which has to be folded into the weights like above or the LoRA loads at `rank / alpha` strength. elif has_diffusers_lora_id: for k in all_keys: - if a_key in k or b_key in k: - converted_state_dict[k] = state_dict.pop(k) - elif ".alpha" in k: + if k.endswith(a_key): + diffusers_up_key = k.replace(a_key, b_key) + alpha_key = k.replace(a_key, ".alpha") + + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(diffusers_up_key) + if alpha_key in state_dict: + scale_down, scale_up = get_alpha_scales(down_weight, alpha_key) + down_weight = down_weight * scale_down + up_weight = up_weight * scale_up + converted_state_dict[k] = down_weight + converted_state_dict[diffusers_up_key] = up_weight + for k in list(state_dict): + if k.endswith(".alpha"): state_dict.pop(k) if len(state_dict) > 0: diff --git a/tests/pipelines/qwenimage/test_qwenimage.py b/tests/pipelines/qwenimage/test_qwenimage.py index 44e3d87a18e9..257523223055 100644 --- a/tests/pipelines/qwenimage/test_qwenimage.py +++ b/tests/pipelines/qwenimage/test_qwenimage.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer @@ -195,6 +196,27 @@ class TestQwenImagePipelineMemory(QwenImagePipelineTesterConfig, MemoryTesterMix class TestQwenImagePipelineLoRA(QwenImagePipelineTesterConfig, LoraTesterMixin): """LoRA tests for the QwenImage pipeline.""" + @pytest.mark.parametrize("alpha", [None, 4, 8]) + def test_lora_state_dict_folds_alpha_into_lora_A_lora_B(self, alpha): + """A per-module `.alpha` in a lora_A/lora_B checkpoint must scale the LoRA delta by `alpha / rank`.""" + rank = 4 + to_q = self.get_dummy_components()["transformer"].transformer_blocks[0].attn.to_q + generator = torch.Generator("cpu").manual_seed(0) + lora_A = torch.randn(rank, to_q.in_features, generator=generator) + lora_B = torch.randn(to_q.out_features, rank, generator=generator) + + key = "diffusion_model.transformer_blocks.0.attn.to_q" + state_dict = {f"{key}.lora_A.weight": lora_A, f"{key}.lora_B.weight": lora_B} + if alpha is not None: + state_dict[f"{key}.alpha"] = torch.tensor(float(alpha)) + converted = self.pipeline_class.lora_state_dict(state_dict) + + prefix = "transformer.transformer_blocks.0.attn.to_q" + assert set(converted) == {f"{prefix}.lora_A.weight", f"{prefix}.lora_B.weight"} + scale = 1.0 if alpha is None else alpha / rank + delta = converted[f"{prefix}.lora_B.weight"] @ converted[f"{prefix}.lora_A.weight"] + assert_tensors_close(delta, scale * (lora_B @ lora_A), atol=1e-6, rtol=1e-6) + class TestQwenImagePipelineLoRAMemory(QwenImagePipelineTesterConfig, LoraMemoryTesterMixin): """LoRA x memory-optimization tests for the QwenImage pipeline."""