load_arm_lora: handle mixed-rank LoRAs via zero padding

training/v2/model.py

@@ -113,14 +113,19 @@ def load_arm_lora(model, path):
     """Load an ai-toolkit/kohya edit LoRA (e.g. armV5) into a frozen LoRA module on `model`.
     Handles lora_down/lora_up and lora_A/lora_B keys; folds file alpha/rank into B."""
     sd = load_file(path)
-    rank = None
+    ranks = set()
     for k, v in sd.items():
         if k.endswith(('.lora_down.weight', '.lora_A.weight')):
-            rank = v.shape[0]
-            break
-    if rank is None:
+            ranks.add(v.shape[0])
+    if not ranks:
         raise ValueError('no LoRA tensors found')
+    rank = max(ranks)  # mixed-rank LoRAs: pad smaller deltas with zeros (exact)
+    if len(ranks) > 1:
+        print(f'load_arm_lora: mixed ranks {sorted(ranks)}, padding to {rank}', flush=True)
     lora = LoRA(model.dims, model.depth, rank, rank).requires_grad_(False)
+    with torch.no_grad():
+        for p in lora.parameters():
+            p.zero_()
     for i in range(model.depth):
         for name in PROJ:
             p = f'diffusion_model.layers.{i}.proj.{name}'
@@ -130,8 +135,9 @@ def load_arm_lora(model, path):
                 down, up = sd[p + '.lora_A.weight'], sd[p + '.lora_B.weight']
             alpha_t = sd.get(p + '.alpha')
             alpha = float(alpha_t) if alpha_t is not None else float(down.shape[0])
-            lora.A[name][i].copy_(down.float())
-            lora.B[name][i].copy_(up.float() * (alpha / down.shape[0]) / lora.scale)
+            r = down.shape[0]
+            lora.A[name][i, :r].copy_(down.float())
+            lora.B[name][i, :, :r].copy_(up.float() * (alpha / r) / lora.scale)
     return lora.to(next(model.tr.parameters()).device)