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

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:

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