Upload extract_lora.py

extract_lora.py

@@ -5,6 +5,14 @@ from tqdm import tqdm
 import sys
 
 
+def normalize_key(key):
+    """Strips the 'model.diffusion_model.' prefix from a key if it exists."""
+    prefix = 'model.diffusion_model.'
+    if key.startswith(prefix):
+        return key[len(prefix):]
+    return key
+
+
 def get_torch_dtype(dtype_str: str):
     """Converts a string to a torch.dtype object."""
     if dtype_str == "fp32":
@@ -16,124 +24,148 @@ def get_torch_dtype(dtype_str: str):
     raise ValueError(f"Unsupported dtype: {dtype_str}")
 
 
-def extract_and_svd_lora(model_a_path: str, model_b_path: str, output_path: str, rank: int, device: str, alpha: float,
-                         dtype: torch.dtype):
-    """
-    Extracts the difference between two models, applies SVD to reduce the rank,
-    and saves the result as a LoRA file.
-    """
-    print(f"Loading base model A: {model_a_path}")
-    print(f"Loading finetuned model B: {model_b_path}")
+def randomized_svd(matrix, rank, n_oversamples=10):
+    """Performs Randomized SVD for a faster approximation."""
+    max_rank = min(matrix.shape)
+    if rank >= max_rank:
+        rank = max_rank
+        n_oversamples = 0
+
+    target_rank = min(rank + n_oversamples, max_rank)
+
+    P = torch.randn((matrix.shape[1], target_rank), device=matrix.device, dtype=matrix.dtype)
+    Y = matrix @ P
+
+    Q, _ = torch.linalg.qr(Y.float())
+
+    B = Q.T @ matrix.float()
+
+    U_b, S, Vh = torch.linalg.svd(B, full_matrices=False)
+    U = Q @ U_b
+
+    U = U[:, :rank]
+    S = S[:rank]
+    Vh = Vh[:rank, :]
+
+    return U, S, Vh
+
+
+def extract_and_svd_lora(args):
+    """Main function to extract, decompose, and save the LoRA."""
+    print(f"Loading base model A: {args.model_a}")
+    print(f"Loading finetuned model B: {args.model_b}")
+    print(f"Using decomposition method: {args.method}")
 
     lora_tensors = {}
+    dtype = get_torch_dtype(args.precision)
+
+    with safe_open(args.model_a, framework="pt", device="cpu") as f_a, \
+            safe_open(args.model_b, framework="pt", device="cpu") as f_b:
 
-    with safe_open(model_a_path, framework="pt", device="cpu") as f_a, \
-            safe_open(model_b_path, framework="pt", device="cpu") as f_b:
+        keys_a_original = set(f_a.keys())
+        keys_b_original = set(f_b.keys())
+        print(f"\nFound {len(keys_a_original)} keys in model A.")
+        print(f"Found {len(keys_b_original)} keys in model B.")
 
-        keys_a = set(f_a.keys())
-        keys_b = set(f_b.keys())
-        common_keys = keys_a.intersection(keys_b)
+        normalized_keys_a = {normalize_key(k): k for k in keys_a_original}
+        normalized_keys_b = {normalize_key(k): k for k in keys_b_original}
 
-        # Filter for processable layers (typically linear and conv weights)
-        # We exclude biases and non-weight tensors.
-        weight_keys = {k for k in common_keys if k.endswith('.weight') and 'lora_' not in k}
+        common_normalized_keys = set(normalized_keys_a.keys()).intersection(set(normalized_keys_b.keys()))
+        print(f"Found {len(common_normalized_keys)} common keys after normalization.\n")
 
-        if not weight_keys:
-            print("No common weight keys found between the two models. Exiting.")
+        processable_keys = {k for k in common_normalized_keys if
+                            (k.endswith('.weight') or k.endswith('.bias')) and 'lora_' not in k}
+
+        if not processable_keys:
+            print("No common weight or bias keys found to process. Check if models are compatible.")
             sys.exit(1)
 
-        print(f"Found {len(weight_keys)} common weight keys to process.")
+        print(f"Found {len(processable_keys)} common keys to process.")
 
-        # Main processing loop with progress bar
-        for key in tqdm(sorted(list(weight_keys)), desc="Processing Layers"):
+        for norm_key in tqdm(sorted(list(processable_keys)), desc="Processing Layers"):
             try:
-                # Load tensors and move to the selected device and dtype
-                tensor_a = f_a.get_tensor(key).to(device=device, dtype=dtype)
-                tensor_b = f_b.get_tensor(key).to(device=device, dtype=dtype)
+                original_key_a = normalized_keys_a[norm_key]
+                original_key_b = normalized_keys_b[norm_key]
+
+                tensor_a = f_a.get_tensor(original_key_a).to(device=args.device, dtype=dtype)
+                tensor_b = f_b.get_tensor(original_key_b).to(device=args.device, dtype=dtype)
 
                 if tensor_a.shape != tensor_b.shape:
-                    print(f"Skipping key {key} due to shape mismatch: A={tensor_a.shape}, B={tensor_b.shape}")
+                    tqdm.write(f"Skipping key {norm_key} due to shape mismatch")
                     continue
 
-                # Calculate the difference (delta weight)
-                delta_w = tensor_b - tensor_a
-
-                # SVD works on 2D matrices. Reshape conv layers and other ND tensors.
-                original_shape = delta_w.shape
-                if delta_w.dim() > 2:
-                    delta_w = delta_w.view(original_shape[0], -1)
-
-                # --- Core SVD Logic ---
-                # ΔW ≈ U * S * Vh
-                # U: Left singular vectors
-                # S: Singular values (a 1D vector)
-                # Vh: Right singular vectors (transposed)
-                U, S, Vh = torch.linalg.svd(delta_w, full_matrices=False)
-
-                # Truncate to the desired rank
-                current_rank = min(rank, S.size(0))  # Ensure rank is not > possible rank
-                U = U[:, :current_rank]
-                S = S[:current_rank]
-                Vh = Vh[:current_rank, :]
-
-                # --- Decompose into LoRA A and B matrices ---
-                # LoRA A (lora_down) is Vh
-                # LoRA B (lora_up) is U * S
-                # We scale lora_up by the singular values to retain the magnitude
-                lora_down = Vh
-                lora_up = U @ torch.diag(S)
-
-                # Reshape back to original conv format if necessary
-                if len(original_shape) > 2:
-                    # For Conv2D, lora_down is (rank, in_channels * k_h * k_w)
-                    # and lora_up is (out_channels, rank). No reshape needed for up.
-                    pass  # The matrix form is standard for LoRA conv layers
-
-                # Create LoRA tensor names
-                base_name = key.replace('.weight', '')
-                lora_down_name = f"{base_name}.lora_down.weight"
-                lora_up_name = f"{base_name}.lora_up.weight"
-                alpha_name = f"{base_name}.alpha"
-
-                # Store tensors, moving them to CPU for saving
-                lora_tensors[lora_down_name] = lora_down.contiguous().cpu().to(torch.float32)
-                lora_tensors[lora_up_name] = lora_up.contiguous().cpu().to(torch.float32)
-                lora_tensors[alpha_name] = torch.tensor(alpha).to(torch.float32)
+                delta = tensor_b - tensor_a
+
+                if norm_key.endswith('.weight'):
+                    delta_w = delta
+                    if delta_w.dim() < 2:
+                        tqdm.write(f"Skipping weight key {norm_key} as it's not a 2D matrix.")
+                        continue
+                    if delta_w.dim() > 2:
+                        delta_w = delta_w.view(delta_w.shape[0], -1)
+
+                    if args.method == 'rsvd':
+                        # Use the new oversamples argument
+                        U, S, Vh = randomized_svd(delta_w, args.rank, n_oversamples=args.oversamples)
+                    else:
+                        U, S, Vh = torch.linalg.svd(delta_w, full_matrices=False)
+                        current_rank = min(args.rank, S.size(0))
+                        U = U[:, :current_rank]
+                        S = S[:current_rank]
+                        Vh = Vh[:current_rank, :]
+
+                    lora_down = Vh
+                    lora_up = U @ torch.diag(S)
+
+                    base_name = norm_key.replace('.weight', '')
+                    prefixed_base_name = f"diffusion_model.{base_name}"
+                    lora_down_name = f"{prefixed_base_name}.lora_down.weight"
+                    lora_up_name = f"{prefixed_base_name}.lora_up.weight"
+
+                    lora_tensors[lora_down_name] = lora_down.contiguous().cpu().to(torch.float32)
+                    lora_tensors[lora_up_name] = lora_up.contiguous().cpu().to(torch.float32)
+
+                elif norm_key.endswith('.bias'):
+                    delta_b = delta
+                    base_name = norm_key.replace('.bias', '')
+                    prefixed_base_name = f"diffusion_model.{base_name}"
+                    diff_b_name = f"{prefixed_base_name}.diff_b"
+                    lora_tensors[diff_b_name] = delta_b.contiguous().cpu().to(torch.float32)
 
             except Exception as e:
-                print(f"Failed to process key {key}: {e}")
+                tqdm.write(f"Failed to process key {norm_key}: {e}")
 
-    # Save the final LoRA file
     if not lora_tensors:
         print("No tensors were processed. Output file will not be created.")
         return
 
-    print(f"\nSaving {len(lora_tensors)} tensors to {output_path}...")
-    save_file(lora_tensors, output_path)
+    print(f"\nSaving {len(lora_tensors)} tensors to {args.output}...")
+    save_file(lora_tensors, args.output)
     print("✅ Done!")
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Extract and SVD a LoRA from two SafeTensors checkpoints.")
-
-    parser.add_argument("model_a", type=str, help="Path to the base model (A) checkpoint in .safetensors format.")
-    parser.add_argument("model_b", type=str, help="Path to the finetuned model (B) checkpoint in .safetensors format.")
-    parser.add_argument("output", type=str, help="Path to save the output LoRA file in .safetensors format.")
-
-    parser.add_argument("--rank", type=int, required=True, help="The target rank for the SVD.")
+    parser = argparse.ArgumentParser(description="Extract a LoRA/LoRA+ from two SafeTensors checkpoints.")
+    parser.add_argument("model_a", type=str, help="Path to the base model (A) checkpoint.")
+    parser.add_argument("model_b", type=str, help="Path to the finetuned model (B) checkpoint.")
+    parser.add_argument("output", type=str, help="Path to save the output file.")
+
+    parser.add_argument("--rank", type=int, required=True, help="The target rank for the decomposition.")
+    parser.add_argument("--alpha", type=float, default=1.0,
+                        help="Informational alpha value for scaling. This value is NOT saved in the file.")
+    parser.add_argument("--method", type=str, default="rsvd", choices=["svd", "rsvd"], help="Decomposition method.")
     parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"],
-                        help="Device to use for computation ('cuda' or 'cpu').")
-    parser.add_argument("--alpha", type=float, default=1.0, help="The alpha (scaling) factor for the LoRA.")
+                        help="Device to use for computation.")
     parser.add_argument("--precision", type=str, default="fp32", choices=["fp32", "fp16", "bf16"],
-                        help="Precision to use for calculations.")
+                        help="Precision for calculations.")
+    # --- NEW ARGUMENT ---
+    parser.add_argument("--oversamples", type=int, default=10,
+                        help="Oversampling parameter for Randomized SVD for better accuracy.")
 
     args = parser.parse_args()
 
-    # Device check
     if args.device == "cuda" and not torch.cuda.is_available():
         print("CUDA is not available. Falling back to CPU.")
         args.device = "cpu"
 
-    dtype = get_torch_dtype(args.precision)
-
-    extract_and_svd_lora(args.model_a, args.model_b, args.output, args.rank, args.device, args.alpha, dtype)
+    extract_and_svd_lora(args)