201 lines
7.9 KiB
Python
201 lines
7.9 KiB
Python
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"
|
|
}
|
|
|