File size: 8,528 Bytes
708238a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 |
import torch
import comfy.model_management
import comfy.utils
import folder_paths
import os
import logging
from tqdm import tqdm
device = comfy.model_management.get_torch_device()
CLAMP_QUANTILE = 0.99
def extract_lora(diff, key, rank, algorithm, lora_type, lowrank_iters=7, adaptive_param=1.0):
"""
Extracts LoRA weights from a weight difference tensor using SVD.
"""
conv2d = (len(diff.shape) == 4)
kernel_size = None if not conv2d else diff.size()[2:4]
conv2d_3x3 = conv2d and kernel_size != (1, 1)
out_dim, in_dim = diff.size()[0:2]
if conv2d:
if conv2d_3x3:
diff = diff.flatten(start_dim=1)
else:
diff = diff.squeeze()
diff_float = diff.float()
if algorithm == "svd_lowrank":
U, S, V = torch.svd_lowrank(diff_float, q=min(rank, in_dim, out_dim), niter=lowrank_iters)
U = U @ torch.diag(S)
Vh = V.t()
else:
#torch.linalg.svdvals()
U, S, Vh = torch.linalg.svd(diff_float)
# Flexible rank selection logic like locon: https://github.com/KohakuBlueleaf/LyCORIS/blob/main/tools/extract_locon.py
if "adaptive" in lora_type:
if lora_type == "adaptive_ratio":
min_s = torch.max(S) * adaptive_param
lora_rank = torch.sum(S > min_s).item()
elif lora_type == "adaptive_energy":
energy = torch.cumsum(S**2, dim=0)
total_energy = torch.sum(S**2)
threshold = adaptive_param * total_energy # e.g., adaptive_param=0.95 for 95%
lora_rank = torch.sum(energy < threshold).item() + 1
elif lora_type == "adaptive_quantile":
s_cum = torch.cumsum(S, dim=0)
min_cum_sum = adaptive_param * torch.sum(S)
lora_rank = torch.sum(s_cum < min_cum_sum).item()
print(f"{key} Extracted LoRA rank: {lora_rank}")
else:
lora_rank = rank
lora_rank = max(1, lora_rank)
lora_rank = min(out_dim, in_dim, lora_rank)
U = U[:, :lora_rank]
S = S[:lora_rank]
U = U @ torch.diag(S)
Vh = Vh[:lora_rank, :]
dist = torch.cat([U.flatten(), Vh.flatten()])
if dist.numel() > 100_000:
# Sample 100,000 elements for quantile estimation
idx = torch.randperm(dist.numel(), device=dist.device)[:100_000]
dist_sample = dist[idx]
hi_val = torch.quantile(dist_sample, CLAMP_QUANTILE)
else:
hi_val = torch.quantile(dist, CLAMP_QUANTILE)
low_val = -hi_val
U = U.clamp(low_val, hi_val)
Vh = Vh.clamp(low_val, hi_val)
if conv2d:
U = U.reshape(out_dim, lora_rank, 1, 1)
Vh = Vh.reshape(lora_rank, in_dim, kernel_size[0], kernel_size[1])
return (U, Vh)
def calc_lora_model(model_diff, rank, prefix_model, prefix_lora, output_sd, lora_type, algorithm, lowrank_iters, out_dtype, bias_diff=False, adaptive_param=1.0):
comfy.model_management.load_models_gpu([model_diff], force_patch_weights=True)
model_diff.model.diffusion_model.cpu()
sd = model_diff.model_state_dict(filter_prefix=prefix_model)
del model_diff
comfy.model_management.soft_empty_cache()
for k, v in sd.items():
if isinstance(v, torch.Tensor):
sd[k] = v.cpu()
# Get total number of keys to process for progress bar
total_keys = len([k for k in sd if k.endswith(".weight") or (bias_diff and k.endswith(".bias"))])
# Create progress bar
progress_bar = tqdm(total=total_keys, desc=f"Extracting LoRA ({prefix_lora.strip('.')})")
comfy_pbar = comfy.utils.ProgressBar(total_keys)
for k in sd:
if k.endswith(".weight"):
weight_diff = sd[k]
if weight_diff.ndim == 5:
logging.info(f"Skipping 5D tensor for key {k}") #skip patch embed
progress_bar.update(1)
comfy_pbar.update(1)
continue
if lora_type != "full":
if weight_diff.ndim < 2:
if bias_diff:
output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().to(out_dtype).cpu()
progress_bar.update(1)
comfy_pbar.update(1)
continue
try:
out = extract_lora(weight_diff.to(device), k, rank, algorithm, lora_type, lowrank_iters=lowrank_iters, adaptive_param=adaptive_param)
output_sd["{}{}.lora_up.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[0].contiguous().to(out_dtype).cpu()
output_sd["{}{}.lora_down.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[1].contiguous().to(out_dtype).cpu()
except Exception as e:
logging.warning(f"Could not generate lora weights for key {k}, error {e}")
else:
output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().to(out_dtype).cpu()
progress_bar.update(1)
comfy_pbar.update(1)
elif bias_diff and k.endswith(".bias"):
output_sd["{}{}.diff_b".format(prefix_lora, k[len(prefix_model):-5])] = sd[k].contiguous().to(out_dtype).cpu()
progress_bar.update(1)
comfy_pbar.update(1)
progress_bar.close()
return output_sd
class LoraExtractKJ:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(s):
return {"required":
{
"finetuned_model": ("MODEL",),
"original_model": ("MODEL",),
"filename_prefix": ("STRING", {"default": "loras/ComfyUI_extracted_lora"}),
"rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1}),
"lora_type": (["standard", "full", "adaptive_ratio", "adaptive_quantile", "adaptive_energy"],),
"algorithm": (["svd_linalg", "svd_lowrank"], {"default": "svd_linalg", "tooltip": "SVD algorithm to use, svd_lowrank is faster but less accurate."}),
"lowrank_iters": ("INT", {"default": 7, "min": 1, "max": 100, "step": 1, "tooltip": "The number of subspace iterations for lowrank SVD algorithm."}),
"output_dtype": (["fp16", "bf16", "fp32"], {"default": "fp16"}),
"bias_diff": ("BOOLEAN", {"default": True}),
"adaptive_param": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "For ratio mode, this is the ratio of the maximum singular value. For quantile mode, this is the quantile of the singular values."}),
},
}
RETURN_TYPES = ()
FUNCTION = "save"
OUTPUT_NODE = True
CATEGORY = "KJNodes/lora"
def save(self, finetuned_model, original_model, filename_prefix, rank, lora_type, algorithm, lowrank_iters, output_dtype, bias_diff, adaptive_param):
if algorithm == "svd_lowrank" and lora_type != "standard":
raise ValueError("svd_lowrank algorithm is only supported for standard LoRA extraction.")
dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp16_fast": torch.float16, "fp32": torch.float32}[output_dtype]
m = finetuned_model.clone()
kp = original_model.get_key_patches("diffusion_model.")
for k in kp:
m.add_patches({k: kp[k]}, - 1.0, 1.0)
model_diff = m
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
output_sd = {}
if model_diff is not None:
output_sd = calc_lora_model(model_diff, rank, "diffusion_model.", "diffusion_model.", output_sd, lora_type, algorithm, lowrank_iters, dtype, bias_diff=bias_diff, adaptive_param=adaptive_param)
if "adaptive" in lora_type:
rank_str = f"{lora_type}_{adaptive_param:.2f}"
else:
rank_str = rank
output_checkpoint = f"{filename}_rank_{rank_str}_{output_dtype}_{counter:05}_.safetensors"
output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
comfy.utils.save_torch_file(output_sd, output_checkpoint, metadata=None)
return {}
NODE_CLASS_MAPPINGS = {
"LoraExtractKJ": LoraExtractKJ
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LoraExtractKJ": "LoraExtractKJ"
}
|