Upload lora_redim.py

lora_redim.py

@@ -0,0 +1,171 @@
+import os
+import argparse
+import torch
+from safetensors.torch import load_file, save_file
+from safetensors import safe_open
+from tqdm import tqdm
+
+
+def resize_lora_model(model_path, output_path, new_dim, device):
+    """
+    Resizes the LoRA dimension of a model using SVD for optimal weight preservation.
+
+    Args:
+        model_path (str): Path to the LoRA model to resize.
+        output_path (str): Path to save the new resized model.
+        new_dim (int): The target new dimension for the LoRA weights.
+        device (str): The device to run calculations on ('cuda' or 'cpu').
+    """
+    print(f"Loading model from: {model_path}")
+    model = load_file(model_path)
+    new_model = {}
+
+    # --- Metadata & Weight Inspection ---
+    original_dim = None
+    alpha = None
+    try:
+        with safe_open(model_path, framework="pt", device="cpu") as f:
+            metadata = f.metadata()
+            if metadata:
+                if 'ss_network_dim' in metadata:
+                    original_dim = int(metadata['ss_network_dim'])
+                    print(f"Original dimension (from metadata): {original_dim}")
+                if 'ss_network_alpha' in metadata:
+                    alpha = float(metadata['ss_network_alpha'])
+                    print(f"Original alpha (from metadata): {alpha}")
+    except Exception as e:
+        print(f"Could not read metadata: {e}. Dimension and alpha will be inferred.")
+
+    # Infer original_dim from weights if not in metadata
+    if original_dim is None:
+        for key in model.keys():
+            if key.endswith((".lora_down.weight", ".lora_A.weight")):
+                original_dim = model[key].shape[0]
+                print(f"Inferred original dimension from weights: {original_dim}")
+                break
+
+    # Infer alpha from weights if not in metadata
+    if alpha is None:
+        for key in model.keys():
+            if key.endswith(".alpha"):
+                alpha = model[key].item()
+                print(f"Inferred alpha from weights: {alpha}")
+                break
+
+                # Fallback for alpha if still not found
+    if alpha is None and original_dim is not None:
+        alpha = float(original_dim)
+        print(f"Alpha not found, falling back to using dimension: {alpha}")
+
+    # --- Tensor Processing ---
+    lora_keys_to_process = set()
+    for key in model.keys():
+        if 'lora_' in key and key.endswith('.weight'):
+            # Get the base name (e.g., "lora_unet_down_blocks_0_attentions_0_proj_in")
+            base_key = key.split('.lora_')[0]
+            lora_keys_to_process.add(base_key)
+
+    if not lora_keys_to_process:
+        print("Error: No LoRA weights found in the model.")
+        return
+
+    print(f"\nFound {len(lora_keys_to_process)} LoRA modules to resize...")
+
+    for base_key in tqdm(sorted(list(lora_keys_to_process)), desc="Resizing modules"):
+        try:
+            down_key, up_key = None, None
+
+            # Determine naming convention
+            if base_key + ".lora_down.weight" in model:
+                down_key = base_key + ".lora_down.weight"
+                up_key = base_key + ".lora_up.weight"
+            elif base_key + ".lora_A.weight" in model:
+                down_key = base_key + ".lora_A.weight"
+                up_key = base_key + ".lora_B.weight"
+            else:
+                continue
+
+            # Move weights to the selected device for calculation
+            down_weight = model[down_key].to(device)
+            up_weight = model[up_key].to(device)
+
+            # --- SVD Resizing ---
+            original_dtype = up_weight.dtype
+
+            # Combine the two matrices to get the full weight update
+            conv2d = down_weight.ndim == 4
+            if conv2d:
+                # For conv layers, treat spatial dims as batch dims
+                down_weight = down_weight.flatten(1)
+                up_weight = up_weight.flatten(1)
+
+            full_weight = up_weight @ down_weight
+
+            # Always cast to float32 for SVD, as some devices (CPU, and some GPUs) don't support bfloat16
+            U, S, Vh = torch.linalg.svd(full_weight.to(torch.float32))
+
+            # Truncate or pad the SVD components
+            U = U[:, :new_dim]
+            S = S[:new_dim]
+            Vh = Vh[:new_dim, :]
+
+            # Reconstruct the new low-rank matrices
+            new_down = torch.diag(S) @ Vh
+            new_up = U
+
+            # Reshape back to original conv format if necessary
+            if conv2d:
+                new_down = new_down.reshape(new_dim, down_weight.shape[1], 1, 1)
+                new_up = new_up.reshape(up_weight.shape[0], new_dim, 1, 1)
+
+            # Move back to CPU and original dtype for saving
+            new_model[down_key] = new_down.contiguous().to(original_dtype).cpu()
+            new_model[up_key] = new_up.contiguous().to(original_dtype).cpu()
+
+            # Copy alpha tensor if it exists for this key
+            alpha_key = base_key + ".alpha"
+            if alpha_key in model:
+                new_model[alpha_key] = model[alpha_key]
+
+        except KeyError:
+            continue
+
+    # Copy non-LoRA tensors
+    for key, value in model.items():
+        if ".lora_" not in key:
+            new_model[key] = value
+
+    # --- Save New Model ---
+    new_metadata = {'ss_network_dim': str(new_dim)}
+    if alpha is not None and original_dim is not None and original_dim > 0:
+        new_alpha = alpha * (new_dim / original_dim)
+        new_metadata['ss_network_alpha'] = str(new_alpha)
+        print(f"\nNew alpha scaled to: {new_alpha:.2f}")
+
+    print(f"\nSaving resized model to: {output_path}")
+    save_file(new_model, output_path, metadata=new_metadata)
+    print("Done!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Resize a LoRA model to a new dimension using SVD.",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+    parser.add_argument("model_path", type=str, help="Path to the LoRA model (.safetensors).")
+    parser.add_argument("output_path", type=str, help="Path to save the resized LoRA model.")
+    parser.add_argument("new_dim", type=int, help="The new LoRA dimension (rank).")
+    parser.add_argument("--device", type=str, default=None,
+                        help="Device to use (e.g., 'cpu', 'cuda'). Autodetects if not specified.")
+
+    args = parser.parse_args()
+
+    if args.device:
+        device = args.device
+    else:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    print(f"Using device: {device}")
+
+    resize_lora_model(args.model_path, args.output_path, args.new_dim, device)
+