276 lines
11 KiB
Python
276 lines
11 KiB
Python
import os
|
|
|
|
import torch
|
|
import folder_paths
|
|
import comfy.utils
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
LORA_A_SUFFIX = ".lora_A.weight"
|
|
LORA_B_SUFFIX = ".lora_B.weight"
|
|
|
|
|
|
def _lora_choices():
|
|
return folder_paths.get_filename_list("loras")
|
|
|
|
|
|
def _target_from_a_key(key):
|
|
if key.endswith(LORA_A_SUFFIX):
|
|
return key[:-len(LORA_A_SUFFIX)]
|
|
return None
|
|
|
|
|
|
def _alpha_for(sd, target, rank):
|
|
alpha = sd.get(f"{target}.alpha", None)
|
|
if alpha is None:
|
|
return float(rank)
|
|
return float(alpha.item())
|
|
|
|
|
|
def _is_gate_target(target):
|
|
return (
|
|
".attn.to_gate" in target or ".attn.gate" in target or
|
|
".ff.gate" in target or ".mlp.gate" in target
|
|
)
|
|
|
|
|
|
def _lora_factors(sd, target, strength):
|
|
a = sd[f"{target}{LORA_A_SUFFIX}"].float()
|
|
b = sd[f"{target}{LORA_B_SUFFIX}"].float()
|
|
rank = a.shape[0]
|
|
scale = _alpha_for(sd, target, rank) / rank
|
|
return a, b * (float(strength) * scale)
|
|
|
|
|
|
def _factor_dot(a_a, b_a, a_b, b_b):
|
|
return ((b_a.T @ b_b) * (a_a @ a_b.T)).sum()
|
|
|
|
|
|
def _factor_norm(a, b):
|
|
return torch.sqrt(torch.clamp(_factor_dot(a, b, a, b), min=0.0))
|
|
|
|
|
|
def _balanced_strengths(sd_a, sd_b, target, strength_a, strength_b, balance_mode):
|
|
if balance_mode != "per_layer_norm":
|
|
return strength_a, strength_b, False
|
|
|
|
a_a, b_a = _lora_factors(sd_a, target, strength_a)
|
|
a_b, b_b = _lora_factors(sd_b, target, strength_b)
|
|
norm_a = _factor_norm(a_a, b_a).item()
|
|
norm_b = _factor_norm(a_b, b_b).item()
|
|
if norm_a == 0.0 or norm_b == 0.0:
|
|
return strength_a, strength_b, False
|
|
|
|
target_norm = (norm_a + norm_b) * 0.5
|
|
return float(strength_a) * target_norm / norm_a, float(strength_b) * target_norm / norm_b, True
|
|
|
|
|
|
def _compress_factors(a, b, rank, device_mode):
|
|
current_rank = a.shape[0]
|
|
rank = max(1, min(int(rank), current_rank))
|
|
if rank >= current_rank:
|
|
return a.contiguous(), b.contiguous()
|
|
|
|
device = torch.device("cpu")
|
|
if device_mode == "cuda" and torch.cuda.is_available():
|
|
device = torch.device("cuda")
|
|
a_work = a.to(device)
|
|
b_work = b.to(device)
|
|
|
|
q, r = torch.linalg.qr(b_work, mode="reduced")
|
|
u, s, vh = torch.linalg.svd(r @ a_work, full_matrices=False)
|
|
u = u[:, :rank]
|
|
s = s[:rank]
|
|
vh = vh[:rank]
|
|
root = torch.sqrt(s)
|
|
b_out = q @ (u * root.unsqueeze(0))
|
|
a_out = root.unsqueeze(1) * vh
|
|
return a_out.cpu().contiguous(), b_out.cpu().contiguous()
|
|
|
|
|
|
def _concat_lora(sd_a, sd_b, target, strength_a, strength_b):
|
|
a_a = sd_a[f"{target}{LORA_A_SUFFIX}"].float()
|
|
b_a = sd_a[f"{target}{LORA_B_SUFFIX}"].float()
|
|
a_b = sd_b[f"{target}{LORA_A_SUFFIX}"].float()
|
|
b_b = sd_b[f"{target}{LORA_B_SUFFIX}"].float()
|
|
|
|
scale_a = float(strength_a) * _alpha_for(sd_a, target, a_a.shape[0]) / a_a.shape[0]
|
|
scale_b = float(strength_b) * _alpha_for(sd_b, target, a_b.shape[0]) / a_b.shape[0]
|
|
a = torch.cat([a_a, a_b], dim=0)
|
|
b = torch.cat([b_a * scale_a, b_b * scale_b], dim=1)
|
|
return a.contiguous(), b.contiguous()
|
|
|
|
|
|
def _max_input_rank(sd_a, sd_b, target):
|
|
return max(sd_a[f"{target}{LORA_A_SUFFIX}"].shape[0], sd_b[f"{target}{LORA_A_SUFFIX}"].shape[0])
|
|
|
|
|
|
def _single_lora(sd, target, strength):
|
|
a, b = _lora_factors(sd, target, strength)
|
|
return a.clone().contiguous(), b.clone().contiguous()
|
|
|
|
|
|
def _zero_lora(sd_a, sd_b, target, rank):
|
|
if rank <= 0:
|
|
rank = max(sd_a[f"{target}{LORA_A_SUFFIX}"].shape[0], sd_b[f"{target}{LORA_A_SUFFIX}"].shape[0])
|
|
in_dim = sd_a[f"{target}{LORA_A_SUFFIX}"].shape[1]
|
|
out_dim = sd_a[f"{target}{LORA_B_SUFFIX}"].shape[0]
|
|
a = torch.zeros((rank, in_dim), dtype=torch.float32)
|
|
b = torch.zeros((out_dim, rank), dtype=torch.float32)
|
|
return a, b
|
|
|
|
|
|
def _should_compress_gate(target, compression_mode):
|
|
if compression_mode == "all_gates":
|
|
return True
|
|
if compression_mode == "attn_to_gate_only":
|
|
return ".attn.to_gate" in target or ".attn.gate" in target
|
|
return False
|
|
|
|
|
|
def _gate_lora(sd_a, sd_b, target, strength_a, strength_b, opposite_strategy, opposite_threshold, gate_rank, compression_mode, compression_device):
|
|
a_a, b_a = _lora_factors(sd_a, target, strength_a)
|
|
a_b, b_b = _lora_factors(sd_b, target, strength_b)
|
|
norm_a = _factor_norm(a_a, b_a)
|
|
norm_b = _factor_norm(a_b, b_b)
|
|
if norm_a.item() == 0.0:
|
|
a, b = _single_lora(sd_b, target, strength_b)
|
|
conflict = False
|
|
elif norm_b.item() == 0.0:
|
|
a, b = _single_lora(sd_a, target, strength_a)
|
|
conflict = False
|
|
else:
|
|
cosine = _factor_dot(a_a, b_a, a_b, b_b) / (norm_a * norm_b)
|
|
conflict = bool(cosine.item() < -float(opposite_threshold))
|
|
if conflict:
|
|
if opposite_strategy == "zero_out":
|
|
a, b = _zero_lora(sd_a, sd_b, target, gate_rank)
|
|
else:
|
|
a, b = _single_lora(sd_a, target, strength_a) if norm_a >= norm_b else _single_lora(sd_b, target, strength_b)
|
|
else:
|
|
a, b = _concat_lora(sd_a, sd_b, target, strength_a, strength_b)
|
|
if _should_compress_gate(target, compression_mode):
|
|
a, b = _compress_factors(a, b, gate_rank, compression_device)
|
|
return a, b, conflict
|
|
|
|
|
|
def _copy_unpaired(out, sd, target, strength):
|
|
a, b = _single_lora(sd, target, strength)
|
|
out[f"{target}{LORA_A_SUFFIX}"] = a
|
|
out[f"{target}{LORA_B_SUFFIX}"] = b
|
|
|
|
|
|
def merge_krea2_loras(sd_a, sd_b, strength_a, strength_b, gate_strength, opposite_strategy, opposite_threshold, gate_rank, gate_compression="off", compression_device="cpu", normal_compression="off", balance_mode="off"):
|
|
targets_a = {_target_from_a_key(k) for k in sd_a if _target_from_a_key(k) is not None}
|
|
targets_b = {_target_from_a_key(k) for k in sd_b if _target_from_a_key(k) is not None}
|
|
out = {}
|
|
stats = {"normal": 0, "gate": 0, "gate_conflicts": 0, "balanced": 0, "only_a": 0, "only_b": 0}
|
|
|
|
for target in sorted(targets_a | targets_b):
|
|
has_a = target in targets_a and f"{target}{LORA_B_SUFFIX}" in sd_a
|
|
has_b = target in targets_b and f"{target}{LORA_B_SUFFIX}" in sd_b
|
|
if has_a and has_b:
|
|
target_strength_a, target_strength_b, balanced = _balanced_strengths(sd_a, sd_b, target, strength_a, strength_b, balance_mode)
|
|
stats["balanced"] += int(balanced)
|
|
if _is_gate_target(target):
|
|
a, b, conflict = _gate_lora(
|
|
sd_a, sd_b, target,
|
|
float(target_strength_a) * float(gate_strength),
|
|
float(target_strength_b) * float(gate_strength),
|
|
opposite_strategy,
|
|
opposite_threshold,
|
|
gate_rank,
|
|
gate_compression,
|
|
compression_device,
|
|
)
|
|
stats["gate"] += 1
|
|
stats["gate_conflicts"] += int(conflict)
|
|
else:
|
|
a, b = _concat_lora(sd_a, sd_b, target, target_strength_a, target_strength_b)
|
|
if normal_compression == "max_input_rank":
|
|
a, b = _compress_factors(a, b, _max_input_rank(sd_a, sd_b, target), compression_device)
|
|
stats["normal"] += 1
|
|
out[f"{target}{LORA_A_SUFFIX}"] = a
|
|
out[f"{target}{LORA_B_SUFFIX}"] = b
|
|
elif has_a:
|
|
_copy_unpaired(out, sd_a, target, strength_a)
|
|
stats["only_a"] += 1
|
|
elif has_b:
|
|
_copy_unpaired(out, sd_b, target, strength_b)
|
|
stats["only_b"] += 1
|
|
|
|
return out, stats
|
|
|
|
|
|
class Krea2GateAwareLoRAMerge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"lora_a": (_lora_choices(),),
|
|
"lora_b": (_lora_choices(),),
|
|
"strength_a": ("FLOAT", {"default": 0.5, "min": -4.0, "max": 4.0, "step": 0.01}),
|
|
"strength_b": ("FLOAT", {"default": 0.5, "min": -4.0, "max": 4.0, "step": 0.01}),
|
|
"gate_strength": ("FLOAT", {"default": 0.5, "min": -4.0, "max": 4.0, "step": 0.01}),
|
|
"opposite_gate_deltas": (["pick_strongest", "zero_out"], {"default": "pick_strongest"}),
|
|
"opposite_threshold": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"gate_rank": ("INT", {"default": 32, "min": 0, "max": 256}),
|
|
"balance_mode": (["off", "per_layer_norm"], {"default": "off"}),
|
|
"gate_compression": (["off", "attn_to_gate_only", "all_gates"], {"default": "attn_to_gate_only"}),
|
|
"normal_compression": (["off", "max_input_rank"], {"default": "off"}),
|
|
"compression_device": (["cpu", "cuda"], {"default": "cpu"}),
|
|
"filename_prefix": ("STRING", {"default": "loras/krea2_merged"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("info",)
|
|
FUNCTION = "merge"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "model/merging/krea2"
|
|
|
|
def merge(self, lora_a, lora_b, strength_a, strength_b, gate_strength, opposite_gate_deltas, opposite_threshold, gate_rank, balance_mode="off", gate_compression="attn_to_gate_only", normal_compression="off", compression_device="cpu", filename_prefix="loras/krea2_merged"):
|
|
path_a = folder_paths.get_full_path_or_raise("loras", lora_a)
|
|
path_b = folder_paths.get_full_path_or_raise("loras", lora_b)
|
|
sd_a = comfy.utils.load_torch_file(path_a, safe_load=True)
|
|
sd_b = comfy.utils.load_torch_file(path_b, safe_load=True)
|
|
|
|
merged, stats = merge_krea2_loras(
|
|
sd_a, sd_b, strength_a, strength_b, gate_strength,
|
|
opposite_gate_deltas, opposite_threshold, gate_rank, gate_compression, compression_device, normal_compression, balance_mode,
|
|
)
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
|
|
output_checkpoint = f"{filename}_{counter:05}_.safetensors"
|
|
output_path = os.path.join(full_output_folder, output_checkpoint)
|
|
save_file(merged, output_path, metadata={
|
|
"merge_node": "Krea2GateAwareLoRAMerge",
|
|
"lora_a": lora_a,
|
|
"lora_b": lora_b,
|
|
"strength_a": str(strength_a),
|
|
"strength_b": str(strength_b),
|
|
"gate_strength": str(gate_strength),
|
|
"opposite_gate_deltas": opposite_gate_deltas,
|
|
"balance_mode": balance_mode,
|
|
"gate_compression": gate_compression,
|
|
"normal_compression": normal_compression,
|
|
"compression_device": compression_device,
|
|
})
|
|
|
|
rel_path = os.path.join(subfolder, output_checkpoint) if subfolder else output_checkpoint
|
|
info = (
|
|
f"Saved {rel_path}. "
|
|
f"normal={stats['normal']}, gate={stats['gate']}, "
|
|
f"gate_conflicts={stats['gate_conflicts']}, balanced={stats['balanced']}, "
|
|
f"only_a={stats['only_a']}, only_b={stats['only_b']}."
|
|
)
|
|
return {"ui": {"text": [info]}, "result": (info,)}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Krea2GateAwareLoRAMerge": Krea2GateAwareLoRAMerge,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Krea2GateAwareLoRAMerge": "Krea2 Gate-Aware LoRA Merge",
|
|
}
|