[LoRA] apply .alpha to Qwen-Image lora_A/lora_B LoRAs

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()=}")

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."""