From 04a316f9a03c1b66aa2e991a2f39ecc4a7d87ed8 Mon Sep 17 00:00:00 2001 From: dcaayushd Date: Mon, 12 Oct 2026 17:16:19 +0545 Subject: [PATCH 1/2] [LoRA] apply .alpha in lora_A/lora_B Qwen-Image LoRAs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../loaders/lora_conversion_utils.py | 56 ++++++++++++------ tests/pipelines/qwenimage/test_qwenimage.py | 58 +++++++++++++++++++ 2 files changed, 97 insertions(+), 17 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index bfaec3b0b5a7..cdefb1589cf5 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2323,19 +2323,22 @@ 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, None) + if alpha is None: + return 1.0, 1.0 + scale = ( + alpha.item() / 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,13 +2351,32 @@ 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), apply alpha scaling and pop. 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: - state_dict.pop(k) + if not k.endswith(a_key): + continue + + diffusers_up_key = k.replace(a_key, b_key) + if diffusers_up_key not in state_dict: + raise ValueError(f"Missing lora_B weight for {k.removesuffix(a_key)}") + + alpha_key = k.replace(a_key, ".alpha") + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(diffusers_up_key) + scale_down, scale_up = get_alpha_scales(down_weight, alpha_key) + converted_state_dict[k] = down_weight * scale_down + converted_state_dict[diffusers_up_key] = up_weight * scale_up + + for k in all_keys: + if k.endswith(b_key) and k in state_dict: + raise ValueError(f"Missing lora_A weight for {k.removesuffix(b_key)}") + if ".alpha" in k: + state_dict.pop(k, None) + + for k in all_keys: + if ".alpha" in k: + state_dict.pop(k, None) if len(state_dict) > 0: raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}") diff --git a/tests/pipelines/qwenimage/test_qwenimage.py b/tests/pipelines/qwenimage/test_qwenimage.py index 44e3d87a18e9..066cc1b7e62d 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,63 @@ class TestQwenImagePipelineMemory(QwenImagePipelineTesterConfig, MemoryTesterMix class TestQwenImagePipelineLoRA(QwenImagePipelineTesterConfig, LoraTesterMixin): """LoRA tests for the QwenImage pipeline.""" + @pytest.mark.parametrize("alpha", [None, 2, 4, 8]) + def test_lora_state_dict_applies_alpha(self, alpha): + rank = 4 + down = torch.arange(rank * 16, dtype=torch.float32).reshape(rank, 16) / 100 + up = torch.arange(16 * rank, dtype=torch.float32).reshape(16, rank) / 100 + state_dict = { + "diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight": down.clone(), + "diffusion_model.transformer_blocks.0.attn.to_q.lora_B.weight": up.clone(), + } + if alpha is not None: + state_dict["diffusion_model.transformer_blocks.0.attn.to_q.alpha"] = torch.tensor(float(alpha)) + + converted = self.pipeline_class.lora_state_dict(state_dict) + + assert set(converted) == { + "transformer.transformer_blocks.0.attn.to_q.lora_A.weight", + "transformer.transformer_blocks.0.attn.to_q.lora_B.weight", + } + expected_scale = alpha / rank if alpha is not None else 1.0 + assert_tensors_close( + converted["transformer.transformer_blocks.0.attn.to_q.lora_B.weight"] + @ converted["transformer.transformer_blocks.0.attn.to_q.lora_A.weight"], + expected_scale * (up @ down), + atol=1e-6, + rtol=1e-6, + ) + + def test_lora_state_dict_missing_lora_b(self): + state_dict = { + "diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight": torch.randn(4, 16), + } + with pytest.raises(ValueError, match="Missing lora_B weight"): + self.pipeline_class.lora_state_dict(state_dict) + + def test_lora_state_dict_missing_lora_a(self): + state_dict = { + "diffusion_model.transformer_blocks.0.attn.to_q.lora_B.weight": torch.randn(16, 4), + } + with pytest.raises(ValueError, match="Missing lora_A weight"): + self.pipeline_class.lora_state_dict(state_dict) + + def test_lora_state_dict_drops_stray_alpha(self): + state_dict = { + "diffusion_model.transformer_blocks.0.attn.to_q.alpha": torch.tensor(8.0), + } + assert self.pipeline_class.lora_state_dict(state_dict) == {} + + def test_lora_state_dict_preserves_dtype(self): + state_dict = { + "diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight": torch.randn(4, 16, dtype=torch.bfloat16), + "diffusion_model.transformer_blocks.0.attn.to_q.lora_B.weight": torch.randn(16, 4, dtype=torch.bfloat16), + "diffusion_model.transformer_blocks.0.attn.to_q.alpha": torch.tensor(8.0), + } + converted = self.pipeline_class.lora_state_dict(state_dict) + assert converted["transformer.transformer_blocks.0.attn.to_q.lora_A.weight"].dtype == torch.bfloat16 + assert converted["transformer.transformer_blocks.0.attn.to_q.lora_B.weight"].dtype == torch.bfloat16 + class TestQwenImagePipelineLoRAMemory(QwenImagePipelineTesterConfig, LoraMemoryTesterMixin): """LoRA x memory-optimization tests for the QwenImage pipeline.""" From 76d9987441077ca4d9609c1598d8517812bfcbbe Mon Sep 17 00:00:00 2001 From: dcaayushd Date: Tue, 20 Oct 2026 19:28:48 +0545 Subject: [PATCH 2/2] [LoRA] remove unnecessary Qwen-Image alpha tests --- tests/pipelines/qwenimage/test_qwenimage.py | 58 --------------------- 1 file changed, 58 deletions(-) diff --git a/tests/pipelines/qwenimage/test_qwenimage.py b/tests/pipelines/qwenimage/test_qwenimage.py index 066cc1b7e62d..44e3d87a18e9 100644 --- a/tests/pipelines/qwenimage/test_qwenimage.py +++ b/tests/pipelines/qwenimage/test_qwenimage.py @@ -12,7 +12,6 @@ # 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 @@ -196,63 +195,6 @@ class TestQwenImagePipelineMemory(QwenImagePipelineTesterConfig, MemoryTesterMix class TestQwenImagePipelineLoRA(QwenImagePipelineTesterConfig, LoraTesterMixin): """LoRA tests for the QwenImage pipeline.""" - @pytest.mark.parametrize("alpha", [None, 2, 4, 8]) - def test_lora_state_dict_applies_alpha(self, alpha): - rank = 4 - down = torch.arange(rank * 16, dtype=torch.float32).reshape(rank, 16) / 100 - up = torch.arange(16 * rank, dtype=torch.float32).reshape(16, rank) / 100 - state_dict = { - "diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight": down.clone(), - "diffusion_model.transformer_blocks.0.attn.to_q.lora_B.weight": up.clone(), - } - if alpha is not None: - state_dict["diffusion_model.transformer_blocks.0.attn.to_q.alpha"] = torch.tensor(float(alpha)) - - converted = self.pipeline_class.lora_state_dict(state_dict) - - assert set(converted) == { - "transformer.transformer_blocks.0.attn.to_q.lora_A.weight", - "transformer.transformer_blocks.0.attn.to_q.lora_B.weight", - } - expected_scale = alpha / rank if alpha is not None else 1.0 - assert_tensors_close( - converted["transformer.transformer_blocks.0.attn.to_q.lora_B.weight"] - @ converted["transformer.transformer_blocks.0.attn.to_q.lora_A.weight"], - expected_scale * (up @ down), - atol=1e-6, - rtol=1e-6, - ) - - def test_lora_state_dict_missing_lora_b(self): - state_dict = { - "diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight": torch.randn(4, 16), - } - with pytest.raises(ValueError, match="Missing lora_B weight"): - self.pipeline_class.lora_state_dict(state_dict) - - def test_lora_state_dict_missing_lora_a(self): - state_dict = { - "diffusion_model.transformer_blocks.0.attn.to_q.lora_B.weight": torch.randn(16, 4), - } - with pytest.raises(ValueError, match="Missing lora_A weight"): - self.pipeline_class.lora_state_dict(state_dict) - - def test_lora_state_dict_drops_stray_alpha(self): - state_dict = { - "diffusion_model.transformer_blocks.0.attn.to_q.alpha": torch.tensor(8.0), - } - assert self.pipeline_class.lora_state_dict(state_dict) == {} - - def test_lora_state_dict_preserves_dtype(self): - state_dict = { - "diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight": torch.randn(4, 16, dtype=torch.bfloat16), - "diffusion_model.transformer_blocks.0.attn.to_q.lora_B.weight": torch.randn(16, 4, dtype=torch.bfloat16), - "diffusion_model.transformer_blocks.0.attn.to_q.alpha": torch.tensor(8.0), - } - converted = self.pipeline_class.lora_state_dict(state_dict) - assert converted["transformer.transformer_blocks.0.attn.to_q.lora_A.weight"].dtype == torch.bfloat16 - assert converted["transformer.transformer_blocks.0.attn.to_q.lora_B.weight"].dtype == torch.bfloat16 - class TestQwenImagePipelineLoRAMemory(QwenImagePipelineTesterConfig, LoraMemoryTesterMixin): """LoRA x memory-optimization tests for the QwenImage pipeline."""