import torch import os import folder_paths import comfy.utils import comfy.model_management from safetensors.torch import save_file, safe_open from tqdm import tqdm 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": return torch.float32 if dtype_str == "fp16": return torch.float16 if dtype_str == "bf16": return torch.bfloat16 raise ValueError(f"Unsupported dtype: {dtype_str}") 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 class ExtractLoRA: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "model_a": (folder_paths.get_filename_list("diffusion_models"), ), "model_b": (folder_paths.get_filename_list("diffusion_models"), ), "rank": ("INT", {"default": 64, "min": 1, "max": 1024, "step": 1}), "method": (["rsvd", "svd"], {"default": "rsvd"}), "precision": (["fp32", "fp16", "bf16"], {"default": "fp16"}), "device": (["auto", "cuda", "cpu"], {"default": "auto"}), "save_name": ("STRING", {"default": "extracted_lora"}), }, "optional": { "oversamples": ("INT", {"default": 10, "min": 0, "max": 100}), } } RETURN_TYPES = () FUNCTION = "extract" OUTPUT_NODE = True CATEGORY = "Advanced/LoRA" def extract(self, model_a, model_b, rank, method, precision, device, save_name, oversamples=10): # Resolve paths path_a = folder_paths.get_full_path("diffusion_models", model_a) path_b = folder_paths.get_full_path("diffusion_models", model_b) output_dir = folder_paths.get_output_directory() output_path = os.path.join(output_dir, f"{save_name}.safetensors") # Determine Device if device == "auto": compute_device = comfy.model_management.get_torch_device() else: compute_device = torch.device(device) print(f"Loading base model A: {model_a}") print(f"Loading finetuned model B: {model_b}") print(f"Using decomposition method: {method}") lora_tensors = {} dtype = get_torch_dtype(precision) # Open files safely without loading everything to RAM at once with safe_open(path_a, framework="pt", device="cpu") as f_a, \ safe_open(path_b, framework="pt", device="cpu") as f_b: keys_a_original = set(f_a.keys()) keys_b_original = set(f_b.keys()) 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} common_normalized_keys = set(normalized_keys_a.keys()).intersection(set(normalized_keys_b.keys())) 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: raise ValueError("No common weight or bias keys found to process. Check if models are compatible.") print(f"Found {len(processable_keys)} common keys to process.") # Sort keys for consistent processing order sorted_keys = sorted(list(processable_keys)) pbar = comfy.utils.ProgressBar(len(sorted_keys)) for i, norm_key in enumerate(sorted_keys): try: original_key_a = normalized_keys_a[norm_key] original_key_b = normalized_keys_b[norm_key] # Load tensors to the compute device tensor_a = f_a.get_tensor(original_key_a).to(device=compute_device, dtype=dtype) tensor_b = f_b.get_tensor(original_key_b).to(device=compute_device, dtype=dtype) if tensor_a.shape != tensor_b.shape: print(f"Skipping key {norm_key} due to shape mismatch") pbar.update(1) continue delta = tensor_b - tensor_a # Process Weights if norm_key.endswith('.weight'): delta_w = delta # Skip 1D tensors (biases or LayerNorm usually handled separately, but script logic skips non-2D weights) if delta_w.dim() < 2: # If it's not a matrix, we can't SVD it easily for LoRA in the standard way. # Standard LoRA usually only targets Linear/Conv layers (2D+). # If it's a small 1D weight difference, some scripts save it as full diff, # but this script logic skips it. We adhere to the script. pass else: # Flatten if > 2 dims (like Conv2d: out, in, k, k -> out, in*k*k) if delta_w.dim() > 2: delta_w = delta_w.view(delta_w.shape[0], -1) if method == 'rsvd': U, S, Vh = randomized_svd(delta_w, rank, n_oversamples=oversamples) else: U, S, Vh = torch.linalg.svd(delta_w, full_matrices=False) current_rank = min(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', '') # Re-add prefix for correct LoRA loading format 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) # Note: The original script effectively skips biases in the SVD block # because of `if norm_key.endswith('.weight')`. # It does not appear to save bias diffs. We maintain this behavior. except Exception as e: print(f"Failed to process key {norm_key}: {e}") pbar.update(1) if not lora_tensors: raise ValueError("No tensors were processed. Output file will not be created.") print(f"\nSaving {len(lora_tensors)} tensors to {output_path}...") save_file(lora_tensors, output_path) print("✅ Done!") return () NODE_CLASS_MAPPINGS = { "ExtractLoRA": ExtractLoRA } NODE_DISPLAY_NAME_MAPPINGS = { "ExtractLoRA": "Extract LoRA from Checkpoints" }