1514 lines
64 KiB
Python
1514 lines
64 KiB
Python
"""
|
|
LoRA Resize - Resize existing LoRAs to different ranks.
|
|
|
|
Based on kohya_ss resize_lora.py. Merges LoRA weights and re-extracts via SVD.
|
|
Supports fixed rank and dynamic methods (sv_ratio, sv_fro, sv_cumulative).
|
|
"""
|
|
import os
|
|
import re
|
|
import torch
|
|
import torch.linalg as linalg
|
|
import folder_paths
|
|
import comfy.utils
|
|
from tqdm import tqdm
|
|
from comfy_api.latest import io
|
|
from safetensors.torch import save_file
|
|
from typing import Optional, Dict, Tuple, List
|
|
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
|
|
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
|
|
|
|
# Reuse rank functions from extraction module
|
|
from .lora_extract_svd import (
|
|
_index_sv_ratio,
|
|
_index_sv_cumulative,
|
|
_index_sv_fro,
|
|
_svd_extract_linear,
|
|
_svd_extract_conv,
|
|
_compile_patterns,
|
|
_matches_any_pattern,
|
|
)
|
|
|
|
|
|
MIN_SV = 1e-6
|
|
|
|
|
|
# =============================================================================
|
|
# LoRA Format Detection - Support various LoRA sources
|
|
# =============================================================================
|
|
|
|
def detect_lora_format(keys: List[str]) -> Dict:
|
|
"""
|
|
Detect LoRA format and key patterns from various sources.
|
|
|
|
Supports:
|
|
- Kohya/A1111: lora_unet_*, lora_te_*
|
|
- Diffusers: *.lora_A.weight, *.lora_B.weight
|
|
- PEFT: base_model.model.*.lora_A/B
|
|
- ComfyUI native: *.lora_down.weight, *.lora_up.weight
|
|
- Flux/SDXL: Various transformer key patterns
|
|
|
|
Returns:
|
|
Dict with format info and key patterns
|
|
"""
|
|
format_info = {
|
|
"format": "unknown",
|
|
"down_suffix": ".lora_down.weight",
|
|
"up_suffix": ".lora_up.weight",
|
|
"alpha_suffix": ".alpha",
|
|
"has_alpha": False,
|
|
"key_count": 0,
|
|
}
|
|
|
|
# Check for alpha keys
|
|
alpha_keys = [k for k in keys if '.alpha' in k or k.endswith('alpha')]
|
|
format_info["has_alpha"] = len(alpha_keys) > 0
|
|
|
|
# Detect format by key patterns
|
|
sample_keys = keys[:20] # Check first 20 keys
|
|
|
|
# Kohya/A1111 format: lora_unet_down_blocks_0_*.lora_down.weight
|
|
if any('lora_unet_' in k or 'lora_te' in k for k in sample_keys):
|
|
format_info["format"] = "kohya"
|
|
format_info["down_suffix"] = ".lora_down.weight"
|
|
format_info["up_suffix"] = ".lora_up.weight"
|
|
|
|
# Diffusers format: transformer.*.lora_A.weight
|
|
elif any('.lora_A.weight' in k for k in sample_keys):
|
|
format_info["format"] = "diffusers"
|
|
format_info["down_suffix"] = ".lora_A.weight"
|
|
format_info["up_suffix"] = ".lora_B.weight"
|
|
format_info["alpha_suffix"] = ".alpha"
|
|
|
|
# PEFT format: base_model.model.*.lora_A.default.weight
|
|
elif any('base_model.model.' in k and 'lora_' in k for k in sample_keys):
|
|
format_info["format"] = "peft"
|
|
# Handle nested structure
|
|
format_info["down_suffix"] = ".lora_A.default.weight"
|
|
format_info["up_suffix"] = ".lora_B.default.weight"
|
|
|
|
# ComfyUI native / standard
|
|
elif any('.lora_down.weight' in k for k in sample_keys):
|
|
format_info["format"] = "comfy"
|
|
format_info["down_suffix"] = ".lora_down.weight"
|
|
format_info["up_suffix"] = ".lora_up.weight"
|
|
|
|
# Count LoRA layers
|
|
down_keys = [k for k in keys if format_info["down_suffix"] in k]
|
|
format_info["key_count"] = len(down_keys)
|
|
|
|
return format_info
|
|
|
|
|
|
def extract_lora_pairs(keys: List[str], format_info: Dict) -> Dict[str, Dict[str, str]]:
|
|
"""
|
|
Group LoRA keys into down/up/alpha pairs.
|
|
|
|
Returns:
|
|
Dict[block_name, {"down": key, "up": key, "alpha": key}]
|
|
"""
|
|
down_suffix = format_info["down_suffix"]
|
|
up_suffix = format_info["up_suffix"]
|
|
alpha_suffix = format_info["alpha_suffix"]
|
|
|
|
pairs = {}
|
|
|
|
for key in keys:
|
|
if down_suffix in key:
|
|
block_name = key.replace(down_suffix, "")
|
|
if block_name not in pairs:
|
|
pairs[block_name] = {}
|
|
pairs[block_name]["down"] = key
|
|
elif up_suffix in key:
|
|
block_name = key.replace(up_suffix, "")
|
|
if block_name not in pairs:
|
|
pairs[block_name] = {}
|
|
pairs[block_name]["up"] = key
|
|
elif alpha_suffix in key or key.endswith('.alpha'):
|
|
# Handle various alpha key formats
|
|
block_name = key.replace(alpha_suffix, "").replace(".alpha", "")
|
|
if block_name not in pairs:
|
|
pairs[block_name] = {}
|
|
pairs[block_name]["alpha"] = key
|
|
|
|
return pairs
|
|
|
|
|
|
def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int, float]:
|
|
"""
|
|
Detect the rank and alpha of an existing LoRA.
|
|
|
|
Returns:
|
|
(network_dim, network_alpha)
|
|
"""
|
|
network_dim = None
|
|
network_alpha = None
|
|
|
|
for block_name, block_keys in pairs.items():
|
|
if "down" not in block_keys:
|
|
continue
|
|
|
|
# Get dim from down weight shape
|
|
if network_dim is None:
|
|
down_key = block_keys["down"]
|
|
shape = handler.header[down_key]["shape"]
|
|
# Linear: [rank, in_features] or Conv: [rank, in_ch, k, k]
|
|
network_dim = shape[0]
|
|
|
|
# Get alpha if present
|
|
if network_alpha is None and "alpha" in block_keys:
|
|
alpha_tensor = handler.get_tensor(block_keys["alpha"])
|
|
network_alpha = float(alpha_tensor.item())
|
|
|
|
if network_dim is not None and network_alpha is not None:
|
|
break
|
|
|
|
# Default alpha to dim if not found
|
|
if network_alpha is None:
|
|
network_alpha = float(network_dim) if network_dim else 1.0
|
|
if network_dim is None:
|
|
network_dim = 1
|
|
|
|
return network_dim, network_alpha
|
|
|
|
|
|
# =============================================================================
|
|
# Merge Functions
|
|
# =============================================================================
|
|
|
|
def _merge_linear(lora_down: torch.Tensor, lora_up: torch.Tensor, device: str) -> torch.Tensor:
|
|
"""Merge linear LoRA weights: lora_up @ lora_down."""
|
|
lora_down = lora_down.to(device=device, dtype=torch.float32)
|
|
lora_up = lora_up.to(device=device, dtype=torch.float32)
|
|
return lora_up @ lora_down
|
|
|
|
|
|
def _merge_conv(lora_down: torch.Tensor, lora_up: torch.Tensor, device: str) -> torch.Tensor:
|
|
"""Merge conv LoRA weights."""
|
|
in_rank, in_size, kernel_size, k_ = lora_down.shape
|
|
out_size, out_rank, _, _ = lora_up.shape
|
|
assert in_rank == out_rank, f"rank mismatch: {in_rank} vs {out_rank}"
|
|
|
|
lora_down = lora_down.to(device=device, dtype=torch.float32)
|
|
lora_up = lora_up.to(device=device, dtype=torch.float32)
|
|
|
|
merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
|
|
return merged.reshape(out_size, in_size, kernel_size, kernel_size)
|
|
|
|
|
|
# =============================================================================
|
|
# Extract Functions (after merge)
|
|
# =============================================================================
|
|
|
|
def _extract_linear(
|
|
weight: torch.Tensor,
|
|
max_rank: int,
|
|
dynamic_method: Optional[str],
|
|
dynamic_param: Optional[float],
|
|
scale: float,
|
|
niter: int = 2,
|
|
) -> Dict:
|
|
"""Extract LoRA from merged linear weight."""
|
|
out_size, in_size = weight.shape
|
|
|
|
if dynamic_method is None:
|
|
# Fixed rank - use svd_lowrank for 10x speedup
|
|
rank = min(max_rank, min(out_size, in_size) - 1)
|
|
U, S, Vh = torch.svd_lowrank(weight, q=rank, niter=niter)
|
|
Vh = Vh.T # svd_lowrank returns V, not Vh
|
|
new_rank = rank
|
|
new_alpha = float(scale * new_rank)
|
|
stats = {"sum_retained": 1.0, "fro_retained": 1.0} # Not computed for lowrank
|
|
else:
|
|
# Dynamic methods need full SVD to compute rank from all singular values
|
|
U, S, Vh = linalg.svd(weight, full_matrices=False)
|
|
new_rank, new_alpha, stats = _compute_resize(S, max_rank, dynamic_method, dynamic_param, scale)
|
|
U = U[:, :new_rank]
|
|
S = S[:new_rank]
|
|
Vh = Vh[:new_rank, :]
|
|
|
|
lora_up = U @ torch.diag(S)
|
|
lora_down = Vh
|
|
|
|
return {
|
|
"lora_down": lora_down.cpu().contiguous(),
|
|
"lora_up": lora_up.cpu().contiguous(),
|
|
"new_rank": new_rank,
|
|
"new_alpha": new_alpha,
|
|
**stats
|
|
}
|
|
|
|
|
|
def _extract_conv(
|
|
weight: torch.Tensor,
|
|
max_rank: int,
|
|
dynamic_method: Optional[str],
|
|
dynamic_param: Optional[float],
|
|
scale: float,
|
|
niter: int = 2,
|
|
) -> Dict:
|
|
"""Extract LoRA from merged conv weight."""
|
|
out_ch, in_ch, kh, kw = weight.shape
|
|
mat = weight.reshape(out_ch, -1)
|
|
|
|
if dynamic_method is None:
|
|
# Fixed rank - use svd_lowrank for 10x speedup
|
|
rank = min(max_rank, min(mat.shape) - 1)
|
|
U, S, Vh = torch.svd_lowrank(mat, q=rank, niter=niter)
|
|
Vh = Vh.T # svd_lowrank returns V, not Vh
|
|
new_rank = rank
|
|
new_alpha = float(scale * new_rank)
|
|
stats = {"sum_retained": 1.0, "fro_retained": 1.0} # Not computed for lowrank
|
|
else:
|
|
# Dynamic methods need full SVD to compute rank from all singular values
|
|
U, S, Vh = linalg.svd(mat, full_matrices=False)
|
|
new_rank, new_alpha, stats = _compute_resize(S, max_rank, dynamic_method, dynamic_param, scale)
|
|
U = U[:, :new_rank]
|
|
S = S[:new_rank]
|
|
Vh = Vh[:new_rank, :]
|
|
|
|
lora_up = (U @ torch.diag(S)).reshape(out_ch, new_rank, 1, 1)
|
|
lora_down = Vh.reshape(new_rank, in_ch, kh, kw)
|
|
|
|
return {
|
|
"lora_down": lora_down.cpu().contiguous(),
|
|
"lora_up": lora_up.cpu().contiguous(),
|
|
"new_rank": new_rank,
|
|
"new_alpha": new_alpha,
|
|
**stats
|
|
}
|
|
|
|
|
|
def _compute_resize(
|
|
S: torch.Tensor,
|
|
max_rank: int,
|
|
dynamic_method: Optional[str],
|
|
dynamic_param: Optional[float],
|
|
scale: float
|
|
) -> Tuple[int, float, Dict]:
|
|
"""Compute new rank and alpha based on resize method."""
|
|
|
|
if dynamic_method == "sv_ratio" and dynamic_param is not None:
|
|
# Use kohya convention: S[0]/ratio
|
|
min_sv = S[0] / dynamic_param
|
|
new_rank = max(1, int(torch.sum(S > min_sv).item()))
|
|
elif dynamic_method == "sv_cumulative" and dynamic_param is not None:
|
|
new_rank = _index_sv_cumulative(S, dynamic_param, max_rank)
|
|
elif dynamic_method == "sv_fro" and dynamic_param is not None:
|
|
new_rank = _index_sv_fro(S, dynamic_param, max_rank)
|
|
else:
|
|
new_rank = max_rank
|
|
|
|
# Clamp rank
|
|
if S[0] < MIN_SV:
|
|
new_rank = 1
|
|
else:
|
|
new_rank = max(1, min(new_rank, max_rank, len(S) - 1))
|
|
|
|
new_alpha = float(scale * new_rank)
|
|
|
|
# Compute retention stats
|
|
s_sum = float(torch.sum(torch.abs(S)))
|
|
s_rank = float(torch.sum(torch.abs(S[:new_rank]))) if new_rank <= len(S) else s_sum
|
|
|
|
S_sq = S.pow(2)
|
|
s_fro = float(torch.sqrt(torch.sum(S_sq)))
|
|
s_red_fro = float(torch.sqrt(torch.sum(S_sq[:new_rank]))) if new_rank <= len(S) else s_fro
|
|
|
|
stats = {
|
|
"sum_retained": s_rank / s_sum if s_sum > 0 else 1.0,
|
|
"fro_retained": s_red_fro / s_fro if s_fro > 0 else 1.0,
|
|
}
|
|
|
|
return new_rank, new_alpha, stats
|
|
|
|
|
|
# =============================================================================
|
|
# Main Resize Function
|
|
# =============================================================================
|
|
|
|
def resize_lora_file(
|
|
lora_path: str,
|
|
new_rank: int,
|
|
dynamic_method: Optional[str],
|
|
dynamic_param: Optional[float],
|
|
device: str,
|
|
save_dtype: torch.dtype,
|
|
output_filename: str,
|
|
verbose: bool = True,
|
|
svd_niter: int = 2,
|
|
) -> str:
|
|
"""
|
|
Resize a LoRA file to a new rank.
|
|
|
|
Args:
|
|
lora_path: Path to input LoRA
|
|
new_rank: Target rank (max rank for dynamic methods)
|
|
dynamic_method: None, "sv_ratio", "sv_fro", "sv_cumulative"
|
|
dynamic_param: Parameter for dynamic method
|
|
device: Processing device
|
|
save_dtype: Output dtype
|
|
output_filename: Output filename (without extension)
|
|
verbose: Print progress info
|
|
svd_niter: Power iterations for SVD accuracy
|
|
|
|
Returns:
|
|
Path to saved resized LoRA
|
|
"""
|
|
# Prepare memory
|
|
lora_size_gb = estimate_model_size(lora_path)
|
|
prepare_for_large_operation(lora_size_gb * 2, torch.device(device))
|
|
|
|
handler = MemoryEfficientSafeOpen(lora_path)
|
|
|
|
try:
|
|
metadata = handler.metadata().copy()
|
|
all_keys = handler.keys()
|
|
|
|
# Detect format and extract pairs
|
|
format_info = detect_lora_format(all_keys)
|
|
pairs = extract_lora_pairs(all_keys, format_info)
|
|
network_dim, network_alpha = detect_lora_rank(handler, pairs)
|
|
|
|
scale = network_alpha / network_dim if network_dim > 0 else 1.0
|
|
|
|
if verbose:
|
|
method_str = f"{dynamic_method}: {dynamic_param}" if dynamic_method else "fixed"
|
|
print(f"[LoRA Resize] Format: {format_info['format']}, layers: {format_info['key_count']}")
|
|
print(f"[LoRA Resize] Original dim={network_dim}, alpha={network_alpha:.1f}, scale={scale:.3f}")
|
|
print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}")
|
|
|
|
output_sd = {}
|
|
fro_list = []
|
|
pbar = comfy.utils.ProgressBar(len(pairs))
|
|
|
|
with torch.no_grad():
|
|
for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"):
|
|
if "down" not in block_keys or "up" not in block_keys:
|
|
pbar.update(1)
|
|
continue
|
|
|
|
lora_down = handler.get_tensor(block_keys["down"])
|
|
lora_up = handler.get_tensor(block_keys["up"])
|
|
|
|
is_conv = len(lora_down.shape) == 4
|
|
|
|
# Transfer to GPU with pinned memory if available
|
|
if device == 'cuda':
|
|
lora_down = transfer_to_gpu_pinned(lora_down, device, torch.float32)
|
|
lora_up = transfer_to_gpu_pinned(lora_up, device, torch.float32)
|
|
|
|
# Merge and re-extract
|
|
if is_conv:
|
|
weight = _merge_conv(lora_down, lora_up, device)
|
|
result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
|
|
else:
|
|
weight = _merge_linear(lora_down, lora_up, device)
|
|
result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
|
|
|
|
del weight, lora_down, lora_up
|
|
|
|
fro_list.append(result['fro_retained'])
|
|
|
|
# Store using same format suffixes as input
|
|
down_suffix = format_info["down_suffix"]
|
|
up_suffix = format_info["up_suffix"]
|
|
alpha_suffix = format_info["alpha_suffix"]
|
|
|
|
output_sd[f"{block_name}{down_suffix}"] = result["lora_down"].to(save_dtype)
|
|
output_sd[f"{block_name}{up_suffix}"] = result["lora_up"].to(save_dtype)
|
|
output_sd[f"{block_name}{alpha_suffix}"] = torch.tensor([result["new_alpha"]]).to(save_dtype)
|
|
|
|
pbar.update(1)
|
|
|
|
if verbose and fro_list:
|
|
import numpy as np
|
|
avg_fro = np.mean(fro_list)
|
|
std_fro = np.std(fro_list)
|
|
print(f"[LoRA Resize] Average Frobenius retention: {avg_fro:.1%} ± {std_fro:.3f}")
|
|
|
|
# Update metadata
|
|
if dynamic_method:
|
|
metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}"
|
|
metadata["ss_network_dim"] = "Dynamic"
|
|
metadata["ss_network_alpha"] = "Dynamic"
|
|
else:
|
|
metadata["ss_training_comment"] = f"Resized from dim {network_dim} to {new_rank}"
|
|
metadata["ss_network_dim"] = str(new_rank)
|
|
metadata["ss_network_alpha"] = str(scale * new_rank)
|
|
|
|
# Save
|
|
output_dir = folder_paths.get_folder_paths("loras")[0]
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
|
|
|
|
save_file(output_sd, output_path, metadata)
|
|
print(f"[LoRA Resize] Saved to {output_path}")
|
|
|
|
return output_path
|
|
|
|
finally:
|
|
handler.__exit__(None, None, None)
|
|
cleanup_after_operation()
|
|
|
|
|
|
# =============================================================================
|
|
# Node Definitions
|
|
# =============================================================================
|
|
|
|
class LoRAResizeFixed(io.ComfyNode):
|
|
"""Resize LoRA to a fixed rank."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeFixed",
|
|
display_name="LoRA Resize (Fixed Rank)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Resize existing LoRA to a specific rank by merging and re-extracting via SVD.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Int.Input("new_rank", default=64, min=1, max=1024,
|
|
tooltip="Target rank"),
|
|
io.Int.Input("svd_niter", default=2, min=0, max=10,
|
|
tooltip="SVD power iterations (higher = more accurate but slower)"),
|
|
io.String.Input("output_filename", default="resized_lora"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, new_rank, svd_niter, output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_file(lora_path, new_rank, None, None, device, dtype, output_filename, svd_niter=svd_niter)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
class LoRAResizeRatio(io.ComfyNode):
|
|
"""Resize LoRA keeping SVs above ratio threshold."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeRatio",
|
|
display_name="LoRA Resize (SV Ratio)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Dynamically resize LoRA, keeping singular values where S[i] > S[0]/ratio.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Int.Input("max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum allowed rank"),
|
|
io.Float.Input("ratio", default=2.0, min=1.0, max=100.0, step=0.1,
|
|
tooltip="Keep SVs where S[i] > S[0]/ratio"),
|
|
io.String.Input("output_filename", default="resized_lora_ratio"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, max_rank, ratio, output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_file(lora_path, max_rank, "sv_ratio", ratio, device, dtype, output_filename)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
class LoRAResizeFrobenius(io.ComfyNode):
|
|
"""Resize LoRA to preserve Frobenius norm target."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeFrobenius",
|
|
display_name="LoRA Resize (Frobenius)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Dynamically resize LoRA to preserve target fraction of Frobenius norm.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Int.Input("max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum allowed rank"),
|
|
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
|
|
tooltip="Target Frobenius norm retention (0.9 = 90%)"),
|
|
io.String.Input("output_filename", default="resized_lora_fro"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_file(lora_path, max_rank, "sv_fro", target, device, dtype, output_filename)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
class LoRAResizeCumulative(io.ComfyNode):
|
|
"""Resize LoRA to preserve cumulative SV target."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeCumulative",
|
|
display_name="LoRA Resize (Cumulative)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Dynamically resize LoRA to preserve target fraction of cumulative singular values.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Int.Input("max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum allowed rank"),
|
|
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
|
|
tooltip="Target cumulative SV retention (0.9 = 90%)"),
|
|
io.String.Input("output_filename", default="resized_lora_cumulative"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_file(lora_path, max_rank, "sv_cumulative", target, device, dtype, output_filename)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
# =============================================================================
|
|
# Fast Resize via Base Model - Uses extraction code path for ~91x speedup
|
|
# =============================================================================
|
|
|
|
def resize_lora_via_base(
|
|
lora_path: str,
|
|
base_model_path: str,
|
|
mode: str,
|
|
linear_param: float,
|
|
conv_param: float,
|
|
device: str,
|
|
save_dtype: torch.dtype,
|
|
output_filename: str,
|
|
linear_max_rank: int = None,
|
|
conv_max_rank: int = None,
|
|
clamp_quantile: float = 0.99,
|
|
mismatch_mode: str = "skip",
|
|
skip_patterns_str: str = "",
|
|
svd_niter: int = 2,
|
|
verbose: bool = True,
|
|
) -> str:
|
|
"""
|
|
Resize LoRA by merging to base model weights and re-extracting.
|
|
|
|
This uses the fast extraction code path for ~91x speed improvement over
|
|
the original merge-then-SVD approach.
|
|
|
|
Args:
|
|
lora_path: Path to input LoRA
|
|
base_model_path: Path to base model the LoRA was trained on
|
|
mode: "fixed", "ratio", "sv_fro", "sv_cumulative"
|
|
linear_param: Rank for fixed mode, or mode parameter for dynamic
|
|
conv_param: Same for conv layers
|
|
device: Processing device
|
|
save_dtype: Output dtype
|
|
output_filename: Output filename (without extension)
|
|
linear_max_rank: Max rank for linear layers (dynamic modes)
|
|
conv_max_rank: Max rank for conv layers (dynamic modes)
|
|
clamp_quantile: Quantile for weight clamping
|
|
mismatch_mode: How to handle mismatches ("skip", "zeros", "error")
|
|
skip_patterns_str: Regex patterns for layers to skip
|
|
svd_niter: Power iterations for SVD accuracy
|
|
verbose: Print progress info
|
|
|
|
Returns:
|
|
Path to saved resized LoRA
|
|
"""
|
|
# Prepare memory
|
|
lora_size_gb = estimate_model_size(lora_path)
|
|
base_size_gb = estimate_model_size(base_model_path)
|
|
total_size_gb = lora_size_gb + base_size_gb
|
|
|
|
if verbose:
|
|
print(f"[LoRA Resize Fast] Preparing memory for {total_size_gb:.2f}GB operation...")
|
|
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
|
|
|
|
# Open files with memory-efficient handlers
|
|
lora_handler = MemoryEfficientSafeOpen(lora_path)
|
|
base_handler = MemoryEfficientSafeOpen(base_model_path)
|
|
|
|
try:
|
|
lora_metadata = lora_handler.metadata().copy()
|
|
lora_keys = lora_handler.keys()
|
|
base_keys = set(base_handler.keys())
|
|
|
|
# Detect LoRA format
|
|
format_info = detect_lora_format(lora_keys)
|
|
lora_pairs = extract_lora_pairs(lora_keys, format_info)
|
|
network_dim, network_alpha = detect_lora_rank(lora_handler, lora_pairs)
|
|
scale = network_alpha / network_dim if network_dim > 0 else 1.0
|
|
|
|
if verbose:
|
|
print(f"[LoRA Resize Fast] Format: {format_info['format']}, layers: {format_info['key_count']}")
|
|
print(f"[LoRA Resize Fast] Original dim={network_dim}, alpha={network_alpha:.1f}")
|
|
print(f"[LoRA Resize Fast] Mode: {mode}, linear_param={linear_param}, conv_param={conv_param}")
|
|
|
|
# Compile skip patterns
|
|
skip_patterns = _compile_patterns(skip_patterns_str)
|
|
|
|
# Build mapping from base model keys to LoRA block names
|
|
# LoRA keys are like: lora_unet_transformer_blocks_0_attn_to_q.lora_down.weight
|
|
# Base keys are like: transformer_blocks.0.attn.to_q.weight
|
|
down_suffix = format_info["down_suffix"]
|
|
up_suffix = format_info["up_suffix"]
|
|
alpha_suffix = format_info["alpha_suffix"]
|
|
|
|
output_sd = {}
|
|
stats = {"extracted": 0, "skipped": 0, "no_match": 0, "full": 0}
|
|
|
|
# Common prefixes to strip from base model keys
|
|
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."]
|
|
|
|
# Common prefixes to strip from LoRA block names (order matters - longer prefixes first)
|
|
LORA_PREFIXES = [
|
|
"lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_",
|
|
"lycoris_", "diffusion_model.", "transformer.", "unet."
|
|
]
|
|
|
|
def extract_core_layer_base(key: str) -> str:
|
|
"""Extract core layer name from base model key (keeps dots)."""
|
|
# Strip .weight/.bias suffix
|
|
if key.endswith(".weight"):
|
|
key = key[:-7]
|
|
elif key.endswith(".bias"):
|
|
key = key[:-5]
|
|
# Strip common prefixes
|
|
for prefix in BASE_PREFIXES:
|
|
if key.startswith(prefix):
|
|
key = key[len(prefix):]
|
|
break
|
|
return key
|
|
|
|
def extract_core_layer_lora(block_name: str) -> str:
|
|
"""Extract core layer name from LoRA block name and normalize to underscored format."""
|
|
# Strip common prefixes
|
|
for prefix in LORA_PREFIXES:
|
|
if block_name.startswith(prefix):
|
|
block_name = block_name[len(prefix):]
|
|
break
|
|
# Convert dots to underscores (diffusers format uses dots)
|
|
return block_name.replace(".", "_")
|
|
|
|
# Build lookup table: core layer name (underscored) -> original base key
|
|
base_key_lookup = {}
|
|
for bk in base_handler.keys():
|
|
if bk.endswith(".weight"):
|
|
core = extract_core_layer_base(bk)
|
|
# Convert dots to underscores for matching with LoRA format
|
|
core_underscored = core.replace(".", "_")
|
|
base_key_lookup[core_underscored] = bk
|
|
|
|
if verbose:
|
|
print(f"[LoRA Resize Fast] Base model has {len(base_key_lookup)} weight keys")
|
|
# Debug: show first 3 samples from each side with their normalized forms
|
|
sample_lora = list(lora_pairs.keys())[:3]
|
|
sample_lora_normalized = [extract_core_layer_lora(b) for b in sample_lora]
|
|
sample_base = list(base_key_lookup.items())[:3]
|
|
print(f"[DEBUG] Sample LoRA blocks: {sample_lora}")
|
|
print(f"[DEBUG] Sample LoRA normalized: {sample_lora_normalized}")
|
|
print(f"[DEBUG] Sample base (underscored->original): {sample_base}")
|
|
|
|
pbar = comfy.utils.ProgressBar(len(lora_pairs))
|
|
|
|
with torch.no_grad():
|
|
for block_name, block_keys in tqdm(lora_pairs.items(), desc="Resizing via extraction", unit="layers"):
|
|
if "down" not in block_keys or "up" not in block_keys:
|
|
pbar.update(1)
|
|
continue
|
|
|
|
# Check skip patterns on the LoRA block name
|
|
if _matches_any_pattern(block_name, skip_patterns):
|
|
stats["skipped"] += 1
|
|
pbar.update(1)
|
|
continue
|
|
|
|
# Extract core layer name from LoRA block (already underscored)
|
|
core_lora = extract_core_layer_lora(block_name)
|
|
|
|
# Look up the base model key
|
|
if core_lora not in base_key_lookup:
|
|
if verbose and stats["no_match"] < 3:
|
|
print(f"[DEBUG] No match for LoRA core: {core_lora}")
|
|
stats["no_match"] += 1
|
|
pbar.update(1)
|
|
continue
|
|
|
|
base_key = base_key_lookup[core_lora]
|
|
|
|
# Load tensors with pinned memory for CUDA
|
|
# Load base weight to RAM first
|
|
cpu_base = base_handler.get_tensor(base_key)
|
|
|
|
# Transfer to GPU if cuda selected
|
|
if device == 'cuda':
|
|
base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32)
|
|
else:
|
|
base_weight = cpu_base.to(device=device, dtype=torch.float32)
|
|
del cpu_base
|
|
|
|
# Load LoRA weights
|
|
cpu_down = lora_handler.get_tensor(block_keys["down"])
|
|
cpu_up = lora_handler.get_tensor(block_keys["up"])
|
|
if device == 'cuda':
|
|
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
|
|
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
|
|
else:
|
|
lora_down = cpu_down.to(device=device, dtype=torch.float32)
|
|
lora_up = cpu_up.to(device=device, dtype=torch.float32)
|
|
del cpu_down, cpu_up
|
|
|
|
# Get alpha for this layer
|
|
if "alpha" in block_keys:
|
|
alpha_tensor = lora_handler.get_tensor(block_keys["alpha"])
|
|
layer_alpha = float(alpha_tensor.item())
|
|
else:
|
|
layer_alpha = float(network_dim)
|
|
layer_scale = layer_alpha / network_dim if network_dim > 0 else 1.0
|
|
|
|
# Compute LoRA delta: scale * (up @ down)
|
|
is_conv = len(lora_down.shape) == 4
|
|
|
|
if is_conv:
|
|
# Conv merge: scale * (up @ down reshaped)
|
|
in_rank, in_size, kernel_size, k_ = lora_down.shape
|
|
out_size, out_rank, _, _ = lora_up.shape
|
|
lora_delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
|
|
lora_delta = lora_delta.reshape(out_size, in_size, kernel_size, kernel_size)
|
|
else:
|
|
# Linear merge: scale * (up @ down)
|
|
lora_delta = lora_up @ lora_down
|
|
del lora_down, lora_up
|
|
|
|
# Apply scale and merge into base to get finetuned weight
|
|
finetuned_weight = base_weight + layer_scale * lora_delta
|
|
del lora_delta
|
|
|
|
# Subtract base to get the weight difference for extraction
|
|
weight_diff = finetuned_weight - base_weight
|
|
del base_weight, finetuned_weight
|
|
|
|
# Extract new LoRA from the weight difference
|
|
if is_conv:
|
|
result, mode_str = _svd_extract_conv(
|
|
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
|
|
)
|
|
else:
|
|
result, mode_str = _svd_extract_linear(
|
|
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
|
|
)
|
|
|
|
# Cleanup GPU tensors
|
|
del weight_diff
|
|
|
|
if mode_str == "full":
|
|
# Rank too high, store as full diff (rare)
|
|
stats["full"] += 1
|
|
pbar.update(1)
|
|
continue
|
|
|
|
new_lora_down, new_lora_up, _ = result
|
|
new_rank = new_lora_down.shape[0]
|
|
new_alpha = float(new_rank) # Standard: alpha = rank for 1.0 scale
|
|
|
|
# Transfer extracted LoRA tensors back to RAM until saving
|
|
output_sd[f"{block_name}{down_suffix}"] = new_lora_down.to(save_dtype).cpu().contiguous()
|
|
output_sd[f"{block_name}{up_suffix}"] = new_lora_up.to(save_dtype).cpu().contiguous()
|
|
output_sd[f"{block_name}{alpha_suffix}"] = torch.tensor([new_alpha]).to(save_dtype)
|
|
|
|
# Clear remaining tensors from this layer
|
|
del new_lora_down, new_lora_up, result
|
|
|
|
stats["extracted"] += 1
|
|
pbar.update(1)
|
|
|
|
if verbose:
|
|
print(f"[LoRA Resize Fast] Done: {stats['extracted']} extracted, {stats['skipped']} skipped, {stats['full']} full")
|
|
|
|
if not output_sd:
|
|
raise ValueError("No LoRA weights extracted. Check that the base model matches the LoRA.")
|
|
|
|
# Update metadata
|
|
lora_metadata["ss_training_comment"] = f"Resized via base model extraction, mode={mode}"
|
|
if mode == "fixed":
|
|
lora_metadata["ss_network_dim"] = str(int(linear_param))
|
|
lora_metadata["ss_network_alpha"] = str(int(linear_param))
|
|
else:
|
|
lora_metadata["ss_network_dim"] = "Dynamic"
|
|
lora_metadata["ss_network_alpha"] = "Dynamic"
|
|
|
|
# Save
|
|
output_dir = folder_paths.get_folder_paths("loras")[0]
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
|
|
|
|
save_file(output_sd, output_path, lora_metadata)
|
|
print(f"[LoRA Resize Fast] Saved to {output_path}")
|
|
|
|
return output_path
|
|
|
|
finally:
|
|
lora_handler.__exit__(None, None, None)
|
|
base_handler.__exit__(None, None, None)
|
|
cleanup_after_operation()
|
|
|
|
|
|
# =============================================================================
|
|
# Fast Resize Node Definitions
|
|
# =============================================================================
|
|
|
|
class LoRAResizeViaBaseFixed(io.ComfyNode):
|
|
"""Fast LoRA resize using base model + extraction path."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeViaBaseFixed",
|
|
display_name="LoRA Resize Fast (Fixed Rank)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Fast LoRA resize by merging to base model and re-extracting. ~91x faster than standard resize.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
|
|
tooltip="Base model the LoRA was trained on"),
|
|
io.Int.Input("linear_dim", default=64, min=1, max=4096,
|
|
tooltip="Target rank for linear/attention layers"),
|
|
io.Int.Input("conv_dim", default=32, min=1, max=4096,
|
|
tooltip="Target rank for conv layers"),
|
|
io.Int.Input("svd_niter", default=2, min=0, max=10,
|
|
tooltip="SVD power iterations (higher = more accurate)"),
|
|
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
|
|
tooltip="Clamp outlier singular values"),
|
|
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip",
|
|
tooltip="How to handle key mismatches"),
|
|
io.String.Input("skip_patterns", default="", multiline=True,
|
|
tooltip="Regex patterns for layers to skip"),
|
|
io.String.Input("output_filename", default="resized_lora_fast"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, base_model, linear_dim, conv_dim, svd_niter,
|
|
clamp_quantile, mismatch_mode, skip_patterns,
|
|
output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_via_base(
|
|
lora_path=lora_path,
|
|
base_model_path=base_path,
|
|
mode="fixed",
|
|
linear_param=linear_dim,
|
|
conv_param=conv_dim,
|
|
device=device,
|
|
save_dtype=dtype,
|
|
output_filename=output_filename,
|
|
linear_max_rank=linear_dim,
|
|
conv_max_rank=conv_dim,
|
|
clamp_quantile=clamp_quantile,
|
|
mismatch_mode=mismatch_mode,
|
|
skip_patterns_str=skip_patterns,
|
|
svd_niter=svd_niter,
|
|
)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
class LoRAResizeViaBaseRatio(io.ComfyNode):
|
|
"""Fast LoRA resize keeping SVs above ratio threshold."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeViaBaseRatio",
|
|
display_name="LoRA Resize Fast (SV Ratio)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Fast dynamic resize keeping singular values where S[i] > S[0]/ratio.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
|
|
tooltip="Base model the LoRA was trained on"),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum rank for linear layers"),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum rank for conv layers"),
|
|
io.Float.Input("ratio", default=2.0, min=1.0, max=100.0, step=0.1,
|
|
tooltip="Keep SVs where S[i] > S[0]/ratio"),
|
|
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
|
|
tooltip="Clamp outlier singular values"),
|
|
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
|
|
io.String.Input("skip_patterns", default="", multiline=True,
|
|
tooltip="Regex patterns for layers to skip"),
|
|
io.String.Input("output_filename", default="resized_lora_ratio_fast"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, base_model, linear_max_rank, conv_max_rank, ratio,
|
|
clamp_quantile, mismatch_mode, skip_patterns,
|
|
output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_via_base(
|
|
lora_path=lora_path,
|
|
base_model_path=base_path,
|
|
mode="ratio",
|
|
linear_param=ratio,
|
|
conv_param=ratio,
|
|
device=device,
|
|
save_dtype=dtype,
|
|
output_filename=output_filename,
|
|
linear_max_rank=linear_max_rank,
|
|
conv_max_rank=conv_max_rank,
|
|
clamp_quantile=clamp_quantile,
|
|
mismatch_mode=mismatch_mode,
|
|
skip_patterns_str=skip_patterns,
|
|
)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
class LoRAResizeViaBaseFrobenius(io.ComfyNode):
|
|
"""Fast LoRA resize preserving Frobenius norm target."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeViaBaseFrobenius",
|
|
display_name="LoRA Resize Fast (Frobenius)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Fast dynamic resize preserving target fraction of Frobenius norm.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
|
|
tooltip="Base model the LoRA was trained on"),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum rank for linear layers"),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum rank for conv layers"),
|
|
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
|
|
tooltip="Target Frobenius norm retention (0.9 = 90%)"),
|
|
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
|
|
tooltip="Clamp outlier singular values"),
|
|
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
|
|
io.String.Input("skip_patterns", default="", multiline=True,
|
|
tooltip="Regex patterns for layers to skip"),
|
|
io.String.Input("output_filename", default="resized_lora_fro_fast"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, base_model, linear_max_rank, conv_max_rank, target,
|
|
clamp_quantile, mismatch_mode, skip_patterns,
|
|
output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_via_base(
|
|
lora_path=lora_path,
|
|
base_model_path=base_path,
|
|
mode="sv_fro",
|
|
linear_param=target,
|
|
conv_param=target,
|
|
device=device,
|
|
save_dtype=dtype,
|
|
output_filename=output_filename,
|
|
linear_max_rank=linear_max_rank,
|
|
conv_max_rank=conv_max_rank,
|
|
clamp_quantile=clamp_quantile,
|
|
mismatch_mode=mismatch_mode,
|
|
skip_patterns_str=skip_patterns,
|
|
)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
class LoRAResizeViaBaseCumulative(io.ComfyNode):
|
|
"""Fast LoRA resize preserving cumulative SV target."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAResizeViaBaseCumulative",
|
|
display_name="LoRA Resize Fast (Cumulative)",
|
|
category="ModelUtils/LoRA/Resize",
|
|
description="Fast dynamic resize preserving target fraction of cumulative singular values.",
|
|
inputs=[
|
|
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to resize"),
|
|
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
|
|
tooltip="Base model the LoRA was trained on"),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum rank for linear layers"),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
|
tooltip="Maximum rank for conv layers"),
|
|
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
|
|
tooltip="Target cumulative SV retention (0.9 = 90%)"),
|
|
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
|
|
tooltip="Clamp outlier singular values"),
|
|
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
|
|
io.String.Input("skip_patterns", default="", multiline=True,
|
|
tooltip="Regex patterns for layers to skip"),
|
|
io.String.Input("output_filename", default="resized_lora_cumulative_fast"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, lora_name, base_model, linear_max_rank, conv_max_rank, target,
|
|
clamp_quantile, mismatch_mode, skip_patterns,
|
|
output_filename, save_dtype, device) -> io.NodeOutput:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
path = resize_lora_via_base(
|
|
lora_path=lora_path,
|
|
base_model_path=base_path,
|
|
mode="sv_cumulative",
|
|
linear_param=target,
|
|
conv_param=target,
|
|
device=device,
|
|
save_dtype=dtype,
|
|
output_filename=output_filename,
|
|
linear_max_rank=linear_max_rank,
|
|
conv_max_rank=conv_max_rank,
|
|
clamp_quantile=clamp_quantile,
|
|
mismatch_mode=mismatch_mode,
|
|
skip_patterns_str=skip_patterns,
|
|
)
|
|
return io.NodeOutput(path)
|
|
|
|
|
|
# =============================================================================
|
|
# Multi-LoRA Merge to Single LoRA
|
|
# =============================================================================
|
|
|
|
def merge_multi_loras_via_base(
|
|
lora_paths: List[str],
|
|
lora_weights: List[float],
|
|
base_model_path: str,
|
|
mode: str,
|
|
linear_param: float,
|
|
conv_param: float,
|
|
device: str,
|
|
save_dtype: torch.dtype,
|
|
output_filename: str,
|
|
linear_max_rank: int = None,
|
|
conv_max_rank: int = None,
|
|
clamp_quantile: float = 0.99,
|
|
skip_patterns_str: str = "",
|
|
svd_niter: int = 2,
|
|
verbose: bool = True,
|
|
) -> str:
|
|
"""
|
|
Merge multiple LoRAs to base model and re-extract as a single LoRA.
|
|
|
|
Args:
|
|
lora_paths: List of paths to LoRA files
|
|
lora_weights: List of weight strengths (0.0-1.0) for each LoRA
|
|
base_model_path: Path to base model
|
|
mode: "fixed", "ratio", "sv_fro", "sv_cumulative"
|
|
linear_param: Rank for fixed mode, or mode parameter for dynamic
|
|
conv_param: Same for conv layers
|
|
device: Processing device
|
|
save_dtype: Output dtype
|
|
output_filename: Output filename (without extension)
|
|
linear_max_rank: Max rank for linear layers (dynamic modes)
|
|
conv_max_rank: Max rank for conv layers (dynamic modes)
|
|
clamp_quantile: Clamp outlier singular values
|
|
skip_patterns_str: Regex patterns for layers to skip
|
|
svd_niter: Power iterations for SVD accuracy
|
|
verbose: Print progress info
|
|
|
|
Returns:
|
|
Path to saved merged LoRA
|
|
"""
|
|
from .lora_extract_svd import _svd_extract_linear, _svd_extract_conv
|
|
|
|
# Estimate memory and prepare
|
|
total_size_gb = estimate_model_size(base_model_path)
|
|
for lp in lora_paths:
|
|
total_size_gb += estimate_model_size(lp)
|
|
|
|
if verbose:
|
|
print(f"[Multi-LoRA Merge] Preparing memory for {total_size_gb:.2f}GB operation...")
|
|
print(f"[Multi-LoRA Merge] Merging {len(lora_paths)} LoRAs with weights: {lora_weights}")
|
|
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
|
|
|
|
# Open all files
|
|
base_handler = MemoryEfficientSafeOpen(base_model_path)
|
|
lora_handlers = [MemoryEfficientSafeOpen(lp) for lp in lora_paths]
|
|
|
|
try:
|
|
# Detect format and extract pairs for each LoRA
|
|
lora_infos = []
|
|
all_block_names = set()
|
|
|
|
for i, handler in enumerate(lora_handlers):
|
|
keys = handler.keys()
|
|
format_info = detect_lora_format(keys)
|
|
pairs = extract_lora_pairs(keys, format_info)
|
|
network_dim, network_alpha = detect_lora_rank(handler, pairs)
|
|
lora_infos.append({
|
|
"handler": handler,
|
|
"format_info": format_info,
|
|
"pairs": pairs,
|
|
"network_dim": network_dim,
|
|
"network_alpha": network_alpha,
|
|
"weight": lora_weights[i],
|
|
})
|
|
all_block_names.update(pairs.keys())
|
|
if verbose:
|
|
print(f"[Multi-LoRA Merge] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}")
|
|
|
|
# Use first LoRA's format for output
|
|
output_format = lora_infos[0]["format_info"]
|
|
down_suffix = output_format["down_suffix"]
|
|
up_suffix = output_format["up_suffix"]
|
|
alpha_suffix = output_format["alpha_suffix"]
|
|
|
|
# Common prefixes
|
|
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."]
|
|
LORA_PREFIXES = [
|
|
"lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_",
|
|
"lycoris_", "diffusion_model.", "transformer.", "unet."
|
|
]
|
|
|
|
def extract_core_layer_base(key: str) -> str:
|
|
if key.endswith(".weight"):
|
|
key = key[:-7]
|
|
elif key.endswith(".bias"):
|
|
key = key[:-5]
|
|
for prefix in BASE_PREFIXES:
|
|
if key.startswith(prefix):
|
|
key = key[len(prefix):]
|
|
break
|
|
return key
|
|
|
|
def extract_core_layer_lora(block_name: str) -> str:
|
|
for prefix in LORA_PREFIXES:
|
|
if block_name.startswith(prefix):
|
|
block_name = block_name[len(prefix):]
|
|
break
|
|
return block_name.replace(".", "_")
|
|
|
|
# Build base key lookup
|
|
base_key_lookup = {}
|
|
for bk in base_handler.keys():
|
|
if bk.endswith(".weight"):
|
|
core = extract_core_layer_base(bk)
|
|
core_underscored = core.replace(".", "_")
|
|
base_key_lookup[core_underscored] = bk
|
|
|
|
if verbose:
|
|
print(f"[Multi-LoRA Merge] Base model has {len(base_key_lookup)} weight keys")
|
|
|
|
# Compile skip patterns
|
|
skip_patterns = _compile_patterns(skip_patterns_str)
|
|
|
|
output_sd = {}
|
|
stats = {"extracted": 0, "skipped": 0, "no_match": 0, "full": 0}
|
|
pbar = comfy.utils.ProgressBar(len(all_block_names))
|
|
|
|
with torch.no_grad():
|
|
for block_name in tqdm(sorted(all_block_names), desc="Merging LoRAs", unit="layers"):
|
|
# Check skip patterns
|
|
if _matches_any_pattern(block_name, skip_patterns):
|
|
stats["skipped"] += 1
|
|
pbar.update(1)
|
|
continue
|
|
|
|
# Extract core layer name
|
|
core_lora = extract_core_layer_lora(block_name)
|
|
|
|
# Find base model key
|
|
if core_lora not in base_key_lookup:
|
|
if verbose and stats["no_match"] < 3:
|
|
print(f"[DEBUG] No match for: {core_lora}")
|
|
stats["no_match"] += 1
|
|
pbar.update(1)
|
|
continue
|
|
|
|
base_key = base_key_lookup[core_lora]
|
|
|
|
# Load base weight
|
|
cpu_base = base_handler.get_tensor(base_key)
|
|
if device == 'cuda':
|
|
base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32)
|
|
else:
|
|
base_weight = cpu_base.to(device=device, dtype=torch.float32)
|
|
del cpu_base
|
|
|
|
# Accumulate deltas from all LoRAs that have this block
|
|
combined_delta = None
|
|
is_conv = None
|
|
|
|
for info in lora_infos:
|
|
if block_name not in info["pairs"]:
|
|
continue
|
|
|
|
block_keys = info["pairs"][block_name]
|
|
if "down" not in block_keys or "up" not in block_keys:
|
|
continue
|
|
|
|
# Load LoRA weights
|
|
cpu_down = info["handler"].get_tensor(block_keys["down"])
|
|
cpu_up = info["handler"].get_tensor(block_keys["up"])
|
|
if device == 'cuda':
|
|
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
|
|
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
|
|
else:
|
|
lora_down = cpu_down.to(device=device, dtype=torch.float32)
|
|
lora_up = cpu_up.to(device=device, dtype=torch.float32)
|
|
del cpu_down, cpu_up
|
|
|
|
# Get alpha
|
|
if "alpha" in block_keys:
|
|
alpha_tensor = info["handler"].get_tensor(block_keys["alpha"])
|
|
layer_alpha = float(alpha_tensor.item())
|
|
else:
|
|
layer_alpha = float(info["network_dim"])
|
|
layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0
|
|
|
|
# Apply LoRA weight strength
|
|
effective_scale = layer_scale * info["weight"]
|
|
|
|
# Compute delta
|
|
is_conv = len(lora_down.shape) == 4
|
|
if is_conv:
|
|
in_rank, in_size, kernel_size, k_ = lora_down.shape
|
|
out_size, out_rank, _, _ = lora_up.shape
|
|
delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
|
|
delta = delta.reshape(out_size, in_size, kernel_size, kernel_size)
|
|
else:
|
|
delta = lora_up @ lora_down
|
|
del lora_down, lora_up
|
|
|
|
# Accumulate
|
|
if combined_delta is None:
|
|
combined_delta = effective_scale * delta
|
|
else:
|
|
combined_delta = combined_delta + effective_scale * delta
|
|
del delta
|
|
|
|
if combined_delta is None:
|
|
# No LoRA had this block
|
|
del base_weight
|
|
pbar.update(1)
|
|
continue
|
|
|
|
# Compute finetuned weight
|
|
finetuned_weight = base_weight + combined_delta
|
|
|
|
# Compute weight difference for extraction
|
|
weight_diff = finetuned_weight - base_weight
|
|
del base_weight, finetuned_weight, combined_delta
|
|
|
|
# Extract new LoRA
|
|
if is_conv:
|
|
result, mode_str = _svd_extract_conv(
|
|
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
|
|
)
|
|
else:
|
|
result, mode_str = _svd_extract_linear(
|
|
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
|
|
)
|
|
del weight_diff
|
|
|
|
if mode_str == "full":
|
|
stats["full"] += 1
|
|
pbar.update(1)
|
|
continue
|
|
|
|
new_lora_down, new_lora_up, _ = result
|
|
new_rank = new_lora_down.shape[0]
|
|
new_alpha = float(new_rank)
|
|
|
|
# Store result
|
|
output_sd[f"{block_name}{down_suffix}"] = new_lora_down.to(save_dtype).cpu().contiguous()
|
|
output_sd[f"{block_name}{up_suffix}"] = new_lora_up.to(save_dtype).cpu().contiguous()
|
|
output_sd[f"{block_name}{alpha_suffix}"] = torch.tensor([new_alpha]).to(save_dtype)
|
|
|
|
del new_lora_down, new_lora_up, result
|
|
stats["extracted"] += 1
|
|
pbar.update(1)
|
|
|
|
if verbose:
|
|
print(f"[Multi-LoRA Merge] Done: {stats['extracted']} extracted, {stats['skipped']} skipped, {stats['no_match']} no_match")
|
|
|
|
if not output_sd:
|
|
raise ValueError("No LoRA weights extracted. Check that base model matches the LoRAs.")
|
|
|
|
# Build metadata
|
|
metadata = {
|
|
"ss_training_comment": f"Merged {len(lora_paths)} LoRAs via base model, mode={mode}",
|
|
"ss_network_dim": str(int(linear_param)) if mode == "fixed" else "Dynamic",
|
|
"ss_network_alpha": str(int(linear_param)) if mode == "fixed" else "Dynamic",
|
|
}
|
|
|
|
# Save
|
|
output_dir = folder_paths.get_folder_paths("loras")[0]
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
|
|
|
|
save_file(output_sd, output_path, metadata)
|
|
print(f"[Multi-LoRA Merge] Saved to {output_path}")
|
|
|
|
return output_path
|
|
|
|
finally:
|
|
base_handler.__exit__(None, None, None)
|
|
for handler in lora_handlers:
|
|
handler.__exit__(None, None, None)
|
|
cleanup_after_operation()
|
|
|
|
|
|
class LoRAMultiMerge(io.ComfyNode):
|
|
"""Merge multiple LoRAs to base model and extract as single LoRA."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoRAMultiMerge",
|
|
display_name="LoRA Multi-Merge",
|
|
category="ModelUtils/LoRA/Merge",
|
|
description="Merge 1-4 LoRAs to base model and re-extract as a single LoRA.",
|
|
inputs=[
|
|
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
|
|
tooltip="Base model the LoRAs were trained on"),
|
|
io.Combo.Input("lora_count", options=["1", "2", "3", "4"], default="2",
|
|
tooltip="Number of LoRAs to merge"),
|
|
# LoRA 1
|
|
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"),
|
|
tooltip="First LoRA"),
|
|
io.Float.Input("weight_1", default=1.0, min=0.0, max=2.0, step=0.05,
|
|
tooltip="Weight strength for LoRA 1"),
|
|
# LoRA 2
|
|
io.Combo.Input("lora_2", options=["None"] + folder_paths.get_filename_list("loras"),
|
|
default="None", tooltip="Second LoRA"),
|
|
io.Float.Input("weight_2", default=1.0, min=0.0, max=2.0, step=0.05,
|
|
tooltip="Weight strength for LoRA 2"),
|
|
# LoRA 3
|
|
io.Combo.Input("lora_3", options=["None"] + folder_paths.get_filename_list("loras"),
|
|
default="None", tooltip="Third LoRA"),
|
|
io.Float.Input("weight_3", default=1.0, min=0.0, max=2.0, step=0.05,
|
|
tooltip="Weight strength for LoRA 3"),
|
|
# LoRA 4
|
|
io.Combo.Input("lora_4", options=["None"] + folder_paths.get_filename_list("loras"),
|
|
default="None", tooltip="Fourth LoRA"),
|
|
io.Float.Input("weight_4", default=1.0, min=0.0, max=2.0, step=0.05,
|
|
tooltip="Weight strength for LoRA 4"),
|
|
# Extraction settings
|
|
io.Combo.Input("mode", options=["fixed", "sv_ratio", "sv_fro", "sv_cumulative"], default="fixed",
|
|
tooltip="Rank selection mode"),
|
|
io.Int.Input("linear_rank", default=64, min=1, max=1024,
|
|
tooltip="Target rank for linear layers (fixed) or max rank (dynamic)"),
|
|
io.Int.Input("conv_rank", default=32, min=1, max=512,
|
|
tooltip="Target rank for conv layers (fixed) or max rank (dynamic)"),
|
|
io.Float.Input("dynamic_param", default=0.9, min=0.1, max=1.0, step=0.01,
|
|
tooltip="Parameter for dynamic modes (ratio/target)"),
|
|
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
|
|
tooltip="Clamp outlier singular values"),
|
|
io.String.Input("skip_patterns", default="", multiline=True,
|
|
tooltip="Regex patterns for layers to skip"),
|
|
io.String.Input("output_filename", default="merged_lora"),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
|
],
|
|
outputs=[io.String.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, base_model, lora_count,
|
|
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
|
|
mode, linear_rank, conv_rank, dynamic_param,
|
|
clamp_quantile, skip_patterns, output_filename, save_dtype, device) -> io.NodeOutput:
|
|
|
|
# Build LoRA list based on count
|
|
count = int(lora_count)
|
|
lora_names = [lora_1, lora_2, lora_3, lora_4][:count]
|
|
lora_weights = [weight_1, weight_2, weight_3, weight_4][:count]
|
|
|
|
# Filter out "None" entries
|
|
valid_loras = [(name, weight) for name, weight in zip(lora_names, lora_weights) if name != "None"]
|
|
if not valid_loras:
|
|
raise ValueError("At least one LoRA must be selected")
|
|
|
|
lora_names, lora_weights = zip(*valid_loras)
|
|
lora_paths = [folder_paths.get_full_path_or_raise("loras", name) for name in lora_names]
|
|
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
|
|
|
# Determine parameters based on mode
|
|
if mode == "fixed":
|
|
linear_param = float(linear_rank)
|
|
conv_param = float(conv_rank)
|
|
linear_max_rank = None
|
|
conv_max_rank = None
|
|
else:
|
|
linear_param = dynamic_param
|
|
conv_param = dynamic_param
|
|
linear_max_rank = linear_rank
|
|
conv_max_rank = conv_rank
|
|
|
|
path = merge_multi_loras_via_base(
|
|
lora_paths=list(lora_paths),
|
|
lora_weights=list(lora_weights),
|
|
base_model_path=base_path,
|
|
mode=mode,
|
|
linear_param=linear_param,
|
|
conv_param=conv_param,
|
|
device=device,
|
|
save_dtype=dtype,
|
|
output_filename=output_filename,
|
|
linear_max_rank=linear_max_rank,
|
|
conv_max_rank=conv_max_rank,
|
|
clamp_quantile=clamp_quantile,
|
|
skip_patterns_str=skip_patterns,
|
|
)
|
|
return io.NodeOutput(path)
|