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