| import comfy.model_management
|
| import safetensors.torch
|
| import torch, os, comfy, json
|
|
|
|
|
|
|
| CLAMP_QUANTILE = 0.99
|
| def extract_lora(diff, rank):
|
| 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]
|
| rank = min(rank, in_dim, out_dim)
|
|
|
| if conv2d:
|
| if conv2d_3x3:
|
| diff = diff.flatten(start_dim=1)
|
| else:
|
| diff = diff.squeeze()
|
|
|
| U, S, Vh = torch.linalg.svd(diff.float())
|
| U = U[:, :rank]
|
| S = S[:rank]
|
| U = U @ torch.diag(S)
|
| Vh = Vh[:rank, :]
|
|
|
| dist = torch.cat([U.flatten(), Vh.flatten()])
|
| 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, rank, 1, 1)
|
| Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
|
| return (U, Vh)
|
|
|
|
|
| def do_lora_handle(base_data, other_data, rank, prefix, require, do_bias, callback):
|
| out_data = {}
|
| device = comfy.model_management.get_torch_device()
|
| for key in base_data.keys():
|
| callback()
|
| if key not in other_data:
|
| continue
|
| base_tensor = base_data[key].float()
|
| other_tensor = other_data[key].float()
|
| if key.startswith("clip_g"):
|
| key = "1." + key[len("clip_g."):]
|
| elif key.startswith("clip_l"):
|
| key = "0." + key[len("clip_l."):]
|
| if require:
|
| if not key.startswith(require):
|
| print(f"Ignore unmatched key {key} (doesn't match {require})")
|
| continue
|
| key = key[len(require):]
|
| if base_tensor.shape != other_tensor.shape:
|
| continue
|
| diff = other_tensor.to(device) - base_tensor.to(device)
|
| other_tensor = other_tensor.cpu()
|
| base_tensor = base_tensor.cpu()
|
| max_diff = float(diff.abs().max())
|
| if max_diff < 1e-5:
|
| print(f"discard unaltered key {key} ({max_diff})")
|
| continue
|
| if key.endswith(".weight"):
|
| fixed_key = key[:-len(".weight")].replace('.', '_')
|
| name = f"lora_{prefix}_{fixed_key}"
|
| if len(base_tensor.shape) >= 2:
|
| print(f"extract key {name} ({max_diff})")
|
| out = extract_lora(diff, rank)
|
| out_data[f"{name}.lora_up.weight"] = out[0].contiguous().half().cpu()
|
| out_data[f"{name}.lora_down.weight"] = out[1].contiguous().half().cpu()
|
| else:
|
| print(f"ignore valid raw pass-through key {name} ({max_diff})")
|
|
|
| elif key.endswith(".bias") and do_bias:
|
| fixed_key = key[:-len(".bias")].replace('.', '_')
|
| name = f"lora_{prefix}_{fixed_key}"
|
| print(f"extract bias key {name} ({max_diff})")
|
| out_data[f"{name}.diff_b"] = diff.contiguous().half().cpu()
|
|
|
|
|
| return out_data
|
|
|
| class SwarmExtractLora:
|
| def __init__(self):
|
| self.loaded_lora = None
|
|
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {
|
| "required": {
|
| "base_model": ("MODEL", ),
|
| "base_model_clip": ("CLIP", ),
|
| "other_model": ("MODEL", ),
|
| "other_model_clip": ("CLIP", ),
|
| "rank": ("INT", {"default": 16, "min": 1, "max": 320}),
|
| "save_rawpath": ("STRING", {"multiline": False}),
|
| "save_filename": ("STRING", {"multiline": False}),
|
| "save_clip": ("BOOLEAN", {"default": True}),
|
| "metadata": ("STRING", {"multiline": True}),
|
| }
|
| }
|
|
|
| CATEGORY = "SwarmUI/models"
|
| RETURN_TYPES = ()
|
| FUNCTION = "extract_lora"
|
| OUTPUT_NODE = True
|
| DESCRIPTION = "Internal node, do not use directly - extracts a LoRA from the difference between two models. This is used by SwarmUI Utilities tab."
|
|
|
| def extract_lora(self, base_model, base_model_clip, other_model, other_model_clip, rank, save_rawpath, save_filename, save_clip, metadata):
|
| base_data = base_model.model_state_dict()
|
| other_data = other_model.model_state_dict()
|
| key_count = len(base_data.keys())
|
| if save_clip:
|
| key_count += len(base_model_clip.get_sd().keys())
|
| pbar = comfy.utils.ProgressBar(key_count)
|
| class Helper:
|
| steps = 0
|
| def callback(self):
|
| self.steps += 1
|
| pbar.update_absolute(self.steps, key_count, None)
|
| helper = Helper()
|
| out_data = do_lora_handle(base_data, other_data, rank, "unet", "diffusion_model.", True, lambda: helper.callback())
|
| if save_clip:
|
|
|
| out_clip = do_lora_handle(base_model_clip.get_sd(), other_model_clip.get_sd(), rank, "te_text_model_encoder_layers", "0.transformer.text_model.encoder.layers.", False, lambda: helper.callback())
|
| out_clip = do_lora_handle(base_model_clip.get_sd(), other_model_clip.get_sd(), rank, "te2_text_model_encoder_layers", "1.transformer.text_model.encoder.layers.", False, lambda: helper.callback())
|
| out_data.update(out_clip)
|
|
|
|
|
| out_metadata = {
|
| "modelspec.title": f"(Extracted LoRA) {save_filename}",
|
| "modelspec.description": f"LoRA extracted in SwarmUI"
|
| }
|
| if metadata:
|
| out_metadata.update(json.loads(metadata))
|
| path = f"{save_rawpath}{save_filename}.safetensors"
|
| print(f"saving to path {path}")
|
| safetensors.torch.save_file(out_data, path, metadata=out_metadata)
|
| return ()
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "SwarmExtractLora": SwarmExtractLora,
|
| }
|
|
|
