nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #12 from silveroxides/refactor/incremental-writer

Browse Source 
Refactor/incremental writer
This commit is contained in:
Silver
2026-05-16 16:33:37 +02:00
committed by GitHub
parent f9a96d78a2 94d2a93b91
commit fee30e36f5
7 changed files with 748 additions and 600 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/lora_extract_svd.py
+45 -84
View File
@@ -13,13 +13,12 @@ import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
@@ -447,6 +446,7 @@ def extract_lora_from_files(
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
@@ -458,9 +458,9 @@ def extract_lora_from_files(
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> dict[str, torch.Tensor]:
) -> None:
"""
Extract LoRA from difference between two models.
Extract LoRA from difference between two models, writing incrementally to disk.
Args:
model_a_path: Finetuned model path
@@ -470,6 +470,7 @@ def extract_lora_from_files(
conv_param: Mode parameter for conv layers
device: Computation device
save_dtype: Output dtype
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers
clamp_quantile: Quantile for weight clamping
@@ -480,8 +481,6 @@ def extract_lora_from_files(
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
"""
output_sd = {}
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
@@ -611,19 +610,24 @@ def extract_lora_from_files(
del weight_diff
return status, layer_results
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
output_sd.update(layer_sd)
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
@@ -633,22 +637,12 @@ def extract_lora_from_files(
handler_b.__exit__(None, None, None)
cleanup_after_operation()
return output_sd
def _save_lora(output_sd: dict, output_filename: str) -> str:
"""Save LoRA to file."""
if not output_sd:
raise ValueError("No LoRA weights extracted")
def _build_lora_output_path(output_filename: str) -> str:
"""Build output path for LoRA file."""
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)
print(f"[LoRA Extract] Saved to {output_path}")
return output_path
return os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
# =============================================================================
@@ -664,42 +658,6 @@ def _get_model_inputs():
]
def _reconstruct_dots(key: str) -> str:
"""Reconstruct dot structure from underscored key using heuristics."""
# Block indices: input_blocks_1 -> input_blocks.1
blocks = [
"input_blocks", "output_blocks", "middle_block",
"transformer_blocks", "single_transformer_blocks",
"double_blocks", "single_blocks",
"down_blocks", "up_blocks", "mid_block",
"attentions", "resnets", "upsamplers", "downsamplers"
]
for b in blocks:
key = re.sub(f"{b}_(\\d+)", f"{b}.\\1", key)
# Sequential numbers: .1_1 -> .1.1
key = re.sub(r"(\.\d+)_(\d+)", r"\1.\2", key)
# Separators that should be preceded by dot
separators = [
"transformer_blocks", "single_transformer_blocks",
"attn1", "attn2", "attn",
"img_attn", "txt_attn",
"to_q", "to_k", "to_v", "to_out",
"q_proj", "k_proj", "v_proj", "out_proj",
"qkv", "proj", "ff", "net",
"norm1", "norm2", "norm3", "norm",
"time_emb_proj"
]
for s in separators:
key = key.replace(f"_{s}", f".{s}")
# Special case: to_out_0 -> to_out.0
key = re.sub(r"to_out_(\d+)", r"to_out.\1", key)
return key
def _format_lora_key(key: str) -> str:
"""
Format the key for LoRA saving.
@@ -728,9 +686,7 @@ def _format_lora_key(key: str) -> str:
# Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"):
core = key[10:]
dotted = _reconstruct_dots(core)
return f"diffusion_model.{dotted}"
return key
# Handle already-prefixed keys
if key.startswith("diffusion_model."):
@@ -801,15 +757,16 @@ class LoRAExtractFixed(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files(
extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, linear_dim, conv_dim,
device, save_dtype, output_path, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
return io.NodeOutput(output_path)
class LoRAExtractRatio(io.ComfyNode):
@@ -845,16 +802,17 @@ class LoRAExtractRatio(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files(
extract_lora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, linear_max_rank, conv_max_rank,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
return io.NodeOutput(output_path)
class LoRAExtractQuantile(io.ComfyNode):
@@ -890,16 +848,17 @@ class LoRAExtractQuantile(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files(
extract_lora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, linear_max_rank, conv_max_rank,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
return io.NodeOutput(output_path)
class LoRAExtractKnee(io.ComfyNode):
@@ -933,16 +892,17 @@ class LoRAExtractKnee(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files(
extract_lora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, linear_max_rank, conv_max_rank,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
return io.NodeOutput(output_path)
class LoRAExtractFrobenius(io.ComfyNode):
@@ -978,13 +938,14 @@ class LoRAExtractFrobenius(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files(
extract_lora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, linear_max_rank, conv_max_rank,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
return io.NodeOutput(output_path)
nodes/lora_merger.py
+500 -318
View File
@@ -8,7 +8,6 @@ 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 List, Dict, Tuple, Optional
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from .lora_resize import (
@@ -20,7 +19,7 @@ from .lora_resize import (
_format_lora_key
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
@@ -31,6 +30,7 @@ def merge_multi_loras(
device: str,
save_dtype: torch.dtype,
output_filename: str,
base_model_path: Optional[str] = None,
verbose: bool = True,
) -> str:
"""
@@ -78,129 +78,193 @@ def merge_multi_loras(
if verbose:
print(f"[LoRA Multi-Merge] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}, format={fmt['format']}")
for block_name in pairs.keys():
core = _format_lora_key(block_name)
if core not in layer_map:
layer_map[core] = []
layer_map[core].append(i)
if base_model_path:
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."]
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:
result = key
if result.endswith(".weight"):
result = result[:-7]
elif result.endswith(".bias"):
result = result[:-5]
for prefix in BASE_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
return result
def extract_core_layer_lora(block_name: str) -> str:
result = block_name
for prefix in LORA_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
if result.endswith(".lora"):
result = result[:-5]
return result.replace(".", "_")
with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler:
base_keys = list(base_handler.keys())
lora_lookup = {}
for i, info in enumerate(lora_infos):
for block_name in info["pairs"].keys():
core_lora = extract_core_layer_lora(block_name)
if core_lora not in lora_lookup:
lora_lookup[core_lora] = []
lora_lookup[core_lora].append((i, block_name))
for base_key in base_keys:
if base_key.endswith(".weight"):
core_base = extract_core_layer_base(base_key)
core_underscored = core_base.replace(".", "_")
if core_underscored in lora_lookup:
layer_map[core_base] = lora_lookup[core_underscored] # List of (idx, block_name)
else:
for i, info in enumerate(lora_infos):
for block_name in info["pairs"].keys():
core = block_name # Keep block_name as-is without reference
if core not in layer_map:
layer_map[core] = []
layer_map[core].append((i, block_name))
if verbose:
print(f"[LoRA Multi-Merge] Total unique layers after resolving naming: {len(layer_map)}")
output_sd = {}
# Build metadata and output path before loop
metadata = {
"ss_training_comment": f"Merged {len(lora_paths)} LoRAs via concatenation",
"ss_network_module": "networks.lora",
}
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")
pbar = comfy.utils.ProgressBar(len(layer_map))
tensor_count = 0
# 2. Merge each core layer
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"):
downs = []
ups = []
is_conv = False
max_rank = 0
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
# 2. Merge each core layer
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"):
downs = []
ups = []
is_conv = False
max_rank = 0
# First pass: load and determine max rank
for idx in indices:
info = lora_infos[idx]
handler = handlers[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
# First pass: load and determine max rank
for idx, orig_block in indices:
info = lora_infos[idx]
handler = handlers[idx]
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys:
if "down" not in block_keys or "up" not in block_keys:
continue
t_down = handler.get_tensor(block_keys["down"])
t_up = handler.get_tensor(block_keys["up"])
if device == 'cuda':
t_down = transfer_to_gpu_pinned(t_down, device, torch.float32)
t_up = transfer_to_gpu_pinned(t_up, device, torch.float32)
else:
t_down = t_down.to(device=device, dtype=torch.float32)
t_up = t_up.to(device=device, dtype=torch.float32)
# Store with original info for padding/weighting
is_conv = len(t_down.shape) == 4
current_rank = t_down.shape[0]
max_rank = max(max_rank, current_rank)
# Apply scale immediately for concatenate mode, or keep for later
if merge_mode == "concatenate":
effective_weight = info["weight"] * info["scale"]
t_up = t_up * effective_weight
downs.append((t_down, info))
ups.append((t_up, info))
if not downs:
pbar.update(1)
continue
t_down = handler.get_tensor(block_keys["down"])
t_up = handler.get_tensor(block_keys["up"])
if device == 'cuda':
t_down = transfer_to_gpu_pinned(t_down, device, torch.float32)
t_up = transfer_to_gpu_pinned(t_up, device, torch.float32)
else:
t_down = t_down.to(device=device, dtype=torch.float32)
t_up = t_up.to(device=device, dtype=torch.float32)
# Store with original info for padding/weighting
is_conv = len(t_down.shape) == 4
current_rank = t_down.shape[0]
max_rank = max(max_rank, current_rank)
# Apply scale immediately for concatenate mode, or keep for later
if merge_mode == "concatenate":
effective_weight = info["weight"] * info["scale"]
t_up = t_up * effective_weight
# Stack ranks
merged_down = torch.cat([d[0] for d in downs], dim=0)
merged_up = torch.cat([u[0] for u in ups], dim=1)
new_rank = merged_down.shape[0]
new_alpha = float(new_rank)
else:
# weighted_sum (Kohya style)
# Result = sum( weight_i * Padded(B_i) ), sum( weight_i * Padded(A_i) )
# Note: We apply scale to weights here to normalize different alphas
merged_down = torch.zeros_like(downs[0][0])
# Need to handle different ranks via padding
target_down_shape = list(downs[0][0].shape)
target_down_shape[0] = max_rank
target_up_shape = list(ups[0][0].shape)
target_up_shape[1] = max_rank
downs.append((t_down, info))
ups.append((t_up, info))
merged_down = torch.zeros(target_down_shape, device=device, dtype=torch.float32)
merged_up = torch.zeros(target_up_shape, device=device, dtype=torch.float32)
for (t_d, info_d), (t_u, info_u) in zip(downs, ups):
r = t_d.shape[0]
w = info_d["weight"]
# We also apply sqrt(scale) to both to distribute the scale factor?
# Kohya just uses weights. But to be safe with different alphas,
# we apply scale to the final delta.
# For direct weight merge, we'll just use weights.
if r < max_rank:
# Pad down: [r, in...] -> [max_rank, in...]
pad_d = [0] * (len(t_d.shape) * 2)
pad_d[-1] = max_rank - r # last dim in pad is first dim in tensor (reversed)
# Wait, F.pad uses reverse order of dims.
# For [r, in], padding is (0,0, 0, max_rank-r)
padding_d = [0, 0] * (len(t_d.shape) - 1) + [0, max_rank - r]
t_d = torch.nn.functional.pad(t_d, tuple(padding_d))
# Pad up: [out, r...] -> [out, max_rank...]
# For [out, r], padding is (0, max_rank-r, 0, 0)
padding_u = [0, max_rank - r] + [0, 0] * (len(t_u.shape) - 1)
t_u = torch.nn.functional.pad(t_u, tuple(padding_u))
merged_down += w * t_d
merged_up += w * t_u
new_rank = max_rank
# Alpha is usually max of alphas or same as rank
new_alpha = float(max_rank)
# Write layer immediately
writer.write_dict({
f"{core}.lora_down.weight": merged_down.to(save_dtype).cpu().contiguous(),
f"{core}.lora_up.weight": merged_up.to(save_dtype).cpu().contiguous(),
f"{core}.alpha": torch.tensor(new_alpha, dtype=save_dtype),
})
tensor_count += 3
# Cleanup
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if not downs:
pbar.update(1)
continue
if merge_mode == "concatenate":
# Stack ranks
merged_down = torch.cat([d[0] for d in downs], dim=0)
merged_up = torch.cat([u[0] for u in ups], dim=1)
new_rank = merged_down.shape[0]
new_alpha = float(new_rank)
else:
# weighted_sum (Kohya style)
# Result = sum( weight_i * Padded(B_i) ), sum( weight_i * Padded(A_i) )
# Note: We apply scale to weights here to normalize different alphas
merged_down = torch.zeros_like(downs[0][0])
# Need to handle different ranks via padding
target_down_shape = list(downs[0][0].shape)
target_down_shape[0] = max_rank
target_up_shape = list(ups[0][0].shape)
target_up_shape[1] = max_rank
merged_down = torch.zeros(target_down_shape, device=device, dtype=torch.float32)
merged_up = torch.zeros(target_up_shape, device=device, dtype=torch.float32)
for (t_d, info_d), (t_u, info_u) in zip(downs, ups):
r = t_d.shape[0]
w = info_d["weight"]
# We also apply sqrt(scale) to both to distribute the scale factor?
# Kohya just uses weights. But to be safe with different alphas,
# we apply scale to the final delta.
# For direct weight merge, we'll just use weights.
if r < max_rank:
# Pad down: [r, in...] -> [max_rank, in...]
pad_d = [0] * (len(t_d.shape) * 2)
pad_d[-1] = max_rank - r # last dim in pad is first dim in tensor (reversed)
# Wait, F.pad uses reverse order of dims.
# For [r, in], padding is (0,0, 0, max_rank-r)
padding_d = [0, 0] * (len(t_d.shape) - 1) + [0, max_rank - r]
t_d = torch.nn.functional.pad(t_d, tuple(padding_d))
# Pad up: [out, r...] -> [out, max_rank...]
# For [out, r], padding is (0, max_rank-r, 0, 0)
padding_u = [0, max_rank - r] + [0, 0] * (len(t_u.shape) - 1)
t_u = torch.nn.functional.pad(t_u, tuple(padding_u))
merged_down += w * t_d
merged_up += w * t_u
new_rank = max_rank
# Alpha is usually max of alphas or same as rank
new_alpha = float(max_rank)
# Store
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(new_alpha, dtype=save_dtype)
# Cleanup
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# 3. Final Summary
if verbose:
@@ -212,21 +276,7 @@ def merge_multi_loras(
print(f"[LoRA Multi-Merge] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
print(f"[LoRA Multi-Merge] Output state dict has {len(output_sd)} tensors")
# 4. Metadata
metadata = {
"ss_training_comment": f"Merged {len(lora_paths)} LoRAs via concatenation",
"ss_network_module": "networks.lora",
}
# 4. 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)
if verbose:
print(f"[LoRA Multi-Merge] Output state dict has {tensor_count} tensors")
print(f"[LoRA Multi-Merge] Saved merged LoRA to {output_path}")
return output_path
@@ -248,6 +298,8 @@ class LoRAMultiMerge(io.ComfyNode):
category="ModelUtils/LoRA/Merge",
description="Merge 1-8 LoRAs into a single LoRA file. This resolves different ranks and naming conventions.",
inputs=[
io.Combo.Input("base_model", options=["None"] + folder_paths.get_filename_list("diffusion_models"), default="None",
tooltip="Optional reference model to resolve key naming issues across formats. If None, input keys are preserved verbatim."),
io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2",
tooltip="Number of LoRAs to merge"),
# LoRA 1
@@ -286,7 +338,7 @@ class LoRAMultiMerge(io.ComfyNode):
)
@classmethod
def execute(cls, lora_count,
def execute(cls, base_model, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
merge_mode, output_filename, save_dtype, device) -> io.NodeOutput:
@@ -308,6 +360,9 @@ class LoRAMultiMerge(io.ComfyNode):
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
base_path = None
if base_model and base_model != "None":
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
path = merge_multi_loras(
lora_paths=valid_paths,
@@ -315,7 +370,8 @@ class LoRAMultiMerge(io.ComfyNode):
merge_mode=merge_mode,
device=device,
save_dtype=dtype,
output_filename=output_filename
output_filename=output_filename,
base_model_path=base_path
)
return io.NodeOutput(path)
@@ -332,6 +388,8 @@ class LoRAMultiMergeDARE(io.ComfyNode):
category="ModelUtils/LoRA/Merge",
description="Merge 1-8 LoRAs into a single LoRA file using DARE-Ties. Drops small values and resolves sign conflicts.",
inputs=[
io.Combo.Input("base_model", options=["None"] + folder_paths.get_filename_list("diffusion_models"), default="None",
tooltip="Optional reference model to resolve key naming issues across formats. If None, input keys are preserved verbatim."),
io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2"),
# LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"),
@@ -364,7 +422,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
)
@classmethod
def execute(cls, lora_count,
def execute(cls, base_model, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput:
@@ -386,6 +444,9 @@ class LoRAMultiMergeDARE(io.ComfyNode):
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
base_path = None
if base_model and base_model != "None":
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
path = merge_multi_loras_dare(
lora_paths=valid_paths,
@@ -395,7 +456,8 @@ class LoRAMultiMergeDARE(io.ComfyNode):
seed=seed,
device=device,
save_dtype=dtype,
output_filename=output_filename
output_filename=output_filename,
base_model_path=base_path
)
return io.NodeOutput(path)
@@ -409,6 +471,7 @@ def merge_multi_loras_dare(
device: str,
save_dtype: torch.dtype,
output_filename: str,
base_model_path: Optional[str] = None,
verbose: bool = True,
) -> str:
"""
@@ -446,103 +509,164 @@ def merge_multi_loras_dare(
if verbose:
print(f"[LoRA Multi-Merge DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}")
for block_name in pairs.keys():
core = _format_lora_key(block_name)
if core not in layer_map: layer_map[core] = []
layer_map[core].append(i)
if base_model_path:
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."]
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:
result = key
if result.endswith(".weight"):
result = result[:-7]
elif result.endswith(".bias"):
result = result[:-5]
for prefix in BASE_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
return result
def extract_core_layer_lora(block_name: str) -> str:
result = block_name
for prefix in LORA_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
if result.endswith(".lora"):
result = result[:-5]
return result.replace(".", "_")
with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler:
base_keys = list(base_handler.keys())
lora_lookup = {}
for i, info in enumerate(lora_infos):
for block_name in info["pairs"].keys():
core_lora = extract_core_layer_lora(block_name)
if core_lora not in lora_lookup:
lora_lookup[core_lora] = []
lora_lookup[core_lora].append((i, block_name))
for base_key in base_keys:
if base_key.endswith(".weight"):
core_base = extract_core_layer_base(base_key)
core_underscored = core_base.replace(".", "_")
if core_underscored in lora_lookup:
layer_map[core_base] = lora_lookup[core_underscored] # List of (idx, block_name)
else:
for i, info in enumerate(lora_infos):
for block_name in info["pairs"].keys():
core = block_name
if core not in layer_map:
layer_map[core] = []
layer_map[core].append((i, block_name))
if verbose:
print(f"[LoRA Multi-Merge DARE] Total unique layers after resolving naming: {len(layer_map)}")
output_sd = {}
# Build output path before loop
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")
pbar = comfy.utils.ProgressBar(len(layer_map))
tensor_count = 0
# 2. Merge
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"):
downs = []
ups = []
max_rank = 0
writer = IncrementalSafetensorsWriter(output_path, metadata={"ss_training_comment": "Merged via DARE-Ties"})
writer.__enter__()
try:
# 2. Merge
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"):
downs = []
ups = []
max_rank = 0
for idx in indices:
info = lora_infos[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
for idx, orig_block in indices:
info = lora_infos[idx]
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
max_rank = max(max_rank, t_d.shape[0])
downs.append((t_d, info["weight"]))
ups.append((t_u, info["weight"]))
max_rank = max(max_rank, t_d.shape[0])
downs.append((t_d, info["weight"]))
ups.append((t_u, info["weight"]))
if not downs:
pbar.update(1)
continue
def process_ties_dare(tensors, weights, dim_to_pad):
# Pad all to max_rank
padded = []
for t, w in tensors:
if t.shape[dim_to_pad] < max_rank:
padding = [0] * (len(t.shape) * 2)
# dim_to_pad 0 (down) -> last pair in pad
# dim_to_pad 1 (up) -> second to last pair in pad
rev_dim = len(t.shape) - 1 - dim_to_pad
padding[rev_dim*2 + 1] = max_rank - t.shape[dim_to_pad]
t = torch.nn.functional.pad(t, tuple(padding))
padded.append(t * w)
# DARE
if drop_rate > 0:
for i in range(len(padded)):
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float()
padded[i] = (padded[i] * mask) / (1 - drop_rate)
# TIES
# 1. Trim
if trim_quantile > 0:
for i in range(len(padded)):
flat = padded[i].abs().flatten()
k = int(len(flat) * trim_quantile)
if k > 0:
threshold = torch.kthvalue(flat, k).values
padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i])
# 2. Elect & Merge
stacked = torch.stack(padded) # [N, ...]
signs = torch.sign(stacked)
sum_signs = signs.sum(dim=0)
dominant_sign = torch.sign(sum_signs)
# Filter those matching dominant sign
mask = (signs == dominant_sign) & (dominant_sign != 0)
filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
# Average matching signs
count = mask.sum(dim=0)
result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
return result
merged_down = process_ties_dare(downs, [1.0]*len(downs), 0)
merged_up = process_ties_dare(ups, [1.0]*len(ups), 1)
writer.write_dict({
f"{core}.lora_down.weight": merged_down.to(save_dtype).cpu().contiguous(),
f"{core}.lora_up.weight": merged_up.to(save_dtype).cpu().contiguous(),
f"{core}.alpha": torch.tensor(float(max_rank), dtype=save_dtype),
})
tensor_count += 3
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if not downs:
pbar.update(1)
continue
def process_ties_dare(tensors, weights, dim_to_pad):
# Pad all to max_rank
padded = []
for t, w in tensors:
if t.shape[dim_to_pad] < max_rank:
padding = [0] * (len(t.shape) * 2)
# dim_to_pad 0 (down) -> last pair in pad
# dim_to_pad 1 (up) -> second to last pair in pad
rev_dim = len(t.shape) - 1 - dim_to_pad
padding[rev_dim*2 + 1] = max_rank - t.shape[dim_to_pad]
t = torch.nn.functional.pad(t, tuple(padding))
padded.append(t * w)
# DARE
if drop_rate > 0:
for i in range(len(padded)):
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float()
padded[i] = (padded[i] * mask) / (1 - drop_rate)
# TIES
# 1. Trim
if trim_quantile > 0:
for i in range(len(padded)):
flat = padded[i].abs().flatten()
k = int(len(flat) * trim_quantile)
if k > 0:
threshold = torch.kthvalue(flat, k).values
padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i])
# 2. Elect & Merge
stacked = torch.stack(padded) # [N, ...]
signs = torch.sign(stacked)
sum_signs = signs.sum(dim=0)
dominant_sign = torch.sign(sum_signs)
# Filter those matching dominant sign
mask = (signs == dominant_sign) & (dominant_sign != 0)
filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
# Average matching signs
count = mask.sum(dim=0)
result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
return result
merged_down = process_ties_dare(downs, [1.0]*len(downs), 0)
merged_up = process_ties_dare(ups, [1.0]*len(ups), 1)
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype)
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Final Summary
if verbose:
@@ -554,13 +678,8 @@ def merge_multi_loras_dare(
print(f"[LoRA Multi-Merge DARE] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
print(f"[LoRA Multi-Merge DARE] Output state dict has {len(output_sd)} tensors")
print(f"[LoRA Multi-Merge DARE] Output state dict has {tensor_count} tensors")
# 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, {"ss_training_comment": "Merged via DARE-Ties"})
return output_path
finally:
@@ -579,6 +698,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
category="ModelUtils/LoRA/Merge",
description="Merge 1-8 LoRAs into a single LoRA file using Enhanced DARE-Ties. Uses dynamic probability masking based on value magnitudes.",
inputs=[
io.Combo.Input("base_model", options=["None"] + folder_paths.get_filename_list("diffusion_models"), default="None",
tooltip="Optional reference model to resolve key naming issues across formats. If None, input keys are preserved verbatim."),
io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2"),
# LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"),
@@ -613,7 +734,7 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
)
@classmethod
def execute(cls, lora_count,
def execute(cls, base_model, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
mask_power, min_keep_prob, mask_smooth, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput:
@@ -635,6 +756,9 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
base_path = None
if base_model and base_model != "None":
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
path = merge_multi_loras_dare_enhanced(
lora_paths=valid_paths,
@@ -646,7 +770,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
seed=seed,
device=device,
save_dtype=dtype,
output_filename=output_filename
output_filename=output_filename,
base_model_path=base_path
)
return io.NodeOutput(path)
@@ -663,6 +788,7 @@ def merge_multi_loras_dare_enhanced(
device: str,
save_dtype: torch.dtype,
output_filename: str,
base_model_path: Optional[str] = None,
verbose: bool = True,
) -> str:
"""
@@ -699,105 +825,166 @@ def merge_multi_loras_dare_enhanced(
if verbose:
print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}")
for block_name in pairs.keys():
core = _format_lora_key(block_name)
if core not in layer_map: layer_map[core] = []
layer_map[core].append(i)
if base_model_path:
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."]
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:
result = key
if result.endswith(".weight"):
result = result[:-7]
elif result.endswith(".bias"):
result = result[:-5]
for prefix in BASE_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
return result
def extract_core_layer_lora(block_name: str) -> str:
result = block_name
for prefix in LORA_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
if result.endswith(".lora"):
result = result[:-5]
return result.replace(".", "_")
with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler:
base_keys = list(base_handler.keys())
lora_lookup = {}
for i, info in enumerate(lora_infos):
for block_name in info["pairs"].keys():
core_lora = extract_core_layer_lora(block_name)
if core_lora not in lora_lookup:
lora_lookup[core_lora] = []
lora_lookup[core_lora].append((i, block_name))
for base_key in base_keys:
if base_key.endswith(".weight"):
core_base = extract_core_layer_base(base_key)
core_underscored = core_base.replace(".", "_")
if core_underscored in lora_lookup:
layer_map[core_base] = lora_lookup[core_underscored] # List of (idx, block_name)
else:
for i, info in enumerate(lora_infos):
for block_name in info["pairs"].keys():
core = block_name
if core not in layer_map:
layer_map[core] = []
layer_map[core].append((i, block_name))
if verbose:
print(f"[LoRA Multi-Merge Enhanced DARE] Total unique layers after resolving naming: {len(layer_map)}")
output_sd = {}
# Build output path before loop
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")
pbar = comfy.utils.ProgressBar(len(layer_map))
tensor_count = 0
# 2. Merge
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers (Enhanced DARE)", unit="layers"):
downs = []
ups = []
max_rank = 0
writer = IncrementalSafetensorsWriter(output_path, metadata={"ss_training_comment": "Merged via Enhanced DARE-Ties"})
writer.__enter__()
try:
# 2. Merge
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers (Enhanced DARE)", unit="layers"):
downs = []
ups = []
max_rank = 0
for idx in indices:
info = lora_infos[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
for idx, orig_block in indices:
info = lora_infos[idx]
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
max_rank = max(max_rank, t_d.shape[0])
downs.append((t_d, info["weight"]))
ups.append((t_u, info["weight"]))
max_rank = max(max_rank, t_d.shape[0])
downs.append((t_d, info["weight"]))
ups.append((t_u, info["weight"]))
if not downs:
pbar.update(1)
continue
def process_ties_dare_enhanced(tensors, weights, dim_to_pad):
processed = []
for t, w in tensors:
t_val = t * w
# Enhanced DARE
abs_t = torch.abs(t_val)
max_val = torch.max(abs_t)
if max_val > 0:
prob = torch.clamp((abs_t / max_val) ** max(mask_power, 0.001), min=min_keep_prob, max=1.0)
prob = torch.nan_to_num(prob)
random_mask = torch.bernoulli(prob, generator=rng)
interpolated_mask = torch.lerp(random_mask, prob, mask_smooth)
t_val = t_val * interpolated_mask
# TIES
if trim_quantile > 0:
flat = t_val.abs().flatten()
k = max(1, int(len(flat) * trim_quantile))
if k > 0:
threshold = torch.kthvalue(flat, k).values
t_val = torch.where(t_val.abs() < threshold, torch.zeros_like(t_val), t_val)
# Pad
if t_val.shape[dim_to_pad] < max_rank:
padding = [0] * (len(t_val.shape) * 2)
rev_dim = len(t_val.shape) - 1 - dim_to_pad
padding[rev_dim*2 + 1] = max_rank - t_val.shape[dim_to_pad]
t_val = torch.nn.functional.pad(t_val, tuple(padding))
processed.append(t_val)
stacked = torch.stack(processed)
signs = torch.sign(stacked)
sum_signs = signs.sum(dim=0)
dominant_sign = torch.sign(sum_signs)
mask = (signs == dominant_sign) & (dominant_sign != 0)
filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
count = mask.sum(dim=0)
result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
return result
merged_down = process_ties_dare_enhanced(downs, [1.0]*len(downs), 0)
merged_up = process_ties_dare_enhanced(ups, [1.0]*len(ups), 1)
writer.write_dict({
f"{core}.lora_down.weight": merged_down.to(save_dtype).cpu().contiguous(),
f"{core}.lora_up.weight": merged_up.to(save_dtype).cpu().contiguous(),
f"{core}.alpha": torch.tensor(float(max_rank), dtype=save_dtype),
})
tensor_count += 3
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if not downs:
pbar.update(1)
continue
def process_ties_dare_enhanced(tensors, weights, dim_to_pad):
processed = []
for t, w in tensors:
t_val = t * w
# Enhanced DARE
abs_t = torch.abs(t_val)
max_val = torch.max(abs_t)
if max_val > 0:
prob = torch.clamp((abs_t / max_val) ** max(mask_power, 0.001), min=min_keep_prob, max=1.0)
prob = torch.nan_to_num(prob)
random_mask = torch.bernoulli(prob, generator=rng)
interpolated_mask = torch.lerp(random_mask, prob, mask_smooth)
t_val = t_val * interpolated_mask
# TIES
if trim_quantile > 0:
flat = t_val.abs().flatten()
k = max(1, int(len(flat) * trim_quantile))
if k > 0:
threshold = torch.kthvalue(flat, k).values
t_val = torch.where(t_val.abs() < threshold, torch.zeros_like(t_val), t_val)
# Pad
if t_val.shape[dim_to_pad] < max_rank:
padding = [0] * (len(t_val.shape) * 2)
rev_dim = len(t_val.shape) - 1 - dim_to_pad
padding[rev_dim*2 + 1] = max_rank - t_val.shape[dim_to_pad]
t_val = torch.nn.functional.pad(t_val, tuple(padding))
processed.append(t_val)
stacked = torch.stack(processed)
signs = torch.sign(stacked)
sum_signs = signs.sum(dim=0)
dominant_sign = torch.sign(sum_signs)
mask = (signs == dominant_sign) & (dominant_sign != 0)
filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
count = mask.sum(dim=0)
result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
return result
merged_down = process_ties_dare_enhanced(downs, [1.0]*len(downs), 0)
merged_up = process_ties_dare_enhanced(ups, [1.0]*len(ups), 1)
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype)
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Final Summary
if verbose:
@@ -809,13 +996,8 @@ def merge_multi_loras_dare_enhanced(
print(f"[LoRA Multi-Merge Enhanced DARE] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
print(f"[LoRA Multi-Merge Enhanced DARE] Output state dict has {len(output_sd)} tensors")
print(f"[LoRA Multi-Merge Enhanced DARE] Output state dict has {tensor_count} tensors")
# 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, {"ss_training_comment": "Merged via Enhanced DARE-Ties"})
return output_path
finally:
nodes/lora_resize.py
+166 -160
View File
@@ -12,11 +12,10 @@ 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 .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
from typing import Optional, Dict, Tuple, List
@@ -419,66 +418,7 @@ def resize_lora_file(
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"]
# Standardize prefix using heuristics if possible
new_block_name = _format_lora_key(block_name)
output_sd[f"{new_block_name}{down_suffix}"] = result["lora_down"].to(save_dtype)
output_sd[f"{new_block_name}{up_suffix}"] = result["lora_up"].to(save_dtype)
if alpha_suffix:
output_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype)
del result
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
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
# Build metadata and output path before loop (all values known at this point)
if dynamic_method:
metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}"
metadata["ss_network_dim"] = "Dynamic"
@@ -488,14 +428,78 @@ def resize_lora_file(
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}")
fro_list = []
pbar = comfy.utils.ProgressBar(len(pairs))
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
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"]
# Preserve original key format from input LoRA verbatim
new_block_name = block_name
block_sd = {
f"{new_block_name}{down_suffix}": result["lora_down"].to(save_dtype),
f"{new_block_name}{up_suffix}": result["lora_up"].to(save_dtype),
}
if alpha_suffix:
block_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype)
writer.write_dict(block_sd)
del result
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
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}")
print(f"[LoRA Resize] Saved to {output_path}")
return output_path
finally:
@@ -715,7 +719,7 @@ def merge_loras_to_model(
print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}")
# Common prefixes
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."]
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."]
LORA_PREFIXES = [
"lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_",
"lycoris_", "diffusion_model.", "transformer.", "unet."
@@ -761,7 +765,11 @@ def merge_loras_to_model(
base_metadata = base_handler.metadata().copy() if base_handler.metadata() else {}
base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}"
output_sd = {}
# Build output path before loop
base_dir = os.path.dirname(base_model_path)
os.makedirs(base_dir, exist_ok=True)
output_path = os.path.join(base_dir, f"{output_filename.strip()}.safetensors")
stats = {"merged": 0, "copied": 0, "skipped": 0}
base_keys = list(base_handler.keys())
pbar = comfy.utils.ProgressBar(len(base_keys))
@@ -769,118 +777,116 @@ def merge_loras_to_model(
if verbose:
print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...")
with torch.no_grad():
for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"):
# Check skip patterns
if _matches_any_pattern(base_key, skip_patterns):
stats["skipped"] += 1
pbar.update(1)
continue
writer = IncrementalSafetensorsWriter(output_path, metadata=base_metadata)
writer.__enter__()
try:
with torch.no_grad():
for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"):
# Check skip patterns
if _matches_any_pattern(base_key, skip_patterns):
stats["skipped"] += 1
pbar.update(1)
continue
# Load base weight
cpu_base = base_handler.get_tensor(base_key)
# Load base weight
cpu_base = base_handler.get_tensor(base_key)
# Only process weight tensors for LoRA merging
if base_key.endswith(".weight"):
core = extract_core_layer_base(base_key)
core_underscored = core.replace(".", "_")
# Only process weight tensors for LoRA merging
if base_key.endswith(".weight"):
core = extract_core_layer_base(base_key)
core_underscored = core.replace(".", "_")
# Check if any LoRA contributes to this layer
if core_underscored in lora_lookup:
# Transfer to GPU for computation
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 contributing LoRAs
for info, block_keys in lora_lookup[core_underscored]:
is_full_diff = info["format_info"].get("is_full_diff", False)
if is_full_diff:
# Full diff format
if "diff" not in block_keys:
continue
cpu_diff = info["handler"].get_tensor(block_keys["diff"])
if device == 'cuda':
delta = transfer_to_gpu_pinned(cpu_diff, device, torch.float32)
else:
delta = cpu_diff.to(device=device, dtype=torch.float32)
del cpu_diff
effective_scale = info["weight"]
# Check if any LoRA contributes to this layer
if core_underscored in lora_lookup:
# Transfer to GPU for computation
if device == 'cuda':
base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32)
else:
# Standard LoRA format
if "down" not in block_keys or "up" not in block_keys:
continue
base_weight = cpu_base.to(device=device, dtype=torch.float32)
del cpu_base
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)
# Accumulate deltas from all contributing LoRAs
for info, block_keys in lora_lookup[core_underscored]:
is_full_diff = info["format_info"].get("is_full_diff", False)
if is_full_diff:
# Full diff format
if "diff" not in block_keys:
continue
cpu_diff = info["handler"].get_tensor(block_keys["diff"])
if device == 'cuda':
delta = transfer_to_gpu_pinned(cpu_diff, device, torch.float32)
else:
delta = cpu_diff.to(device=device, dtype=torch.float32)
del cpu_diff
effective_scale = info["weight"]
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
# Standard LoRA format
if "down" not in block_keys or "up" not in block_keys:
continue
# 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
effective_scale = layer_scale * info["weight"]
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
# 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
# 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
effective_scale = layer_scale * info["weight"]
# Apply delta to base weight
base_weight = base_weight + effective_scale * delta
del delta
# 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
# Store merged weight
output_sd[base_key] = base_weight.to(save_dtype).cpu().contiguous()
del base_weight
stats["merged"] += 1
# Apply delta to base weight
base_weight = base_weight + effective_scale * delta
del delta
# Write merged weight
writer.write(base_key, base_weight.to(save_dtype).cpu().contiguous())
del base_weight
stats["merged"] += 1
else:
# No LoRA contribution, write as-is
writer.write(base_key, cpu_base.to(save_dtype).contiguous())
del cpu_base
stats["copied"] += 1
else:
# No LoRA contribution, copy as-is
output_sd[base_key] = cpu_base.to(save_dtype).contiguous()
# Non-weight tensor (bias, norm, etc.), write as-is
writer.write(base_key, cpu_base.to(save_dtype).contiguous())
del cpu_base
stats["copied"] += 1
else:
# Non-weight tensor (bias, norm, etc.), copy as-is
output_sd[base_key] = cpu_base.to(save_dtype).contiguous()
del cpu_base
stats["copied"] += 1
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
pbar.update(1)
finally:
writer.__exit__(None, None, None)
if verbose:
print(f"[LoRA Merge To Model] Done: {stats['merged']} merged, {stats['copied']} copied, {stats['skipped']} skipped")
# Save to base model directory
base_dir = os.path.dirname(base_model_path)
os.makedirs(base_dir, exist_ok=True)
output_path = os.path.join(base_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path, base_metadata)
print(f"[LoRA Merge To Model] Saved to {output_path}")
return output_path
finally:
nodes/merger.py
-1
View File
@@ -6,7 +6,6 @@ import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
from .device_utils import (
estimate_model_size, prepare_for_large_operation, cleanup_after_operation
nodes/prunekeys.py
+22 -22
View File
@@ -4,12 +4,11 @@ import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
@@ -40,32 +39,33 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
if not patterns:
raise ValueError("No keys/patterns provided to prune.")
# Stream tensors and filter on the fly
pruned_tensors = {}
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
all_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(all_keys))
for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = False
if use_regex:
if any(re.search(pattern, key) for pattern in patterns):
is_match = True
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match:
pruned_tensors[key] = handler.get_tensor(key)
pbar.update(1)
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(pruned_tensors, output_path, metadata)
# Stream tensors, filter on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
all_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(all_keys))
for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = False
if use_regex:
if any(re.search(pattern, key) for pattern in patterns):
is_match = True
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match:
writer.write(key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
del pruned_tensors
cleanup_after_operation()
return output_path
nodes/renamekeys.py
+14 -14
View File
@@ -3,12 +3,11 @@ import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
@@ -43,24 +42,25 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
key_map = dict(zip(old_keys, new_keys))
# Stream tensors and rename on the fly
renamed_tensors = {}
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
for key in tqdm(original_keys, desc="Renaming keys", unit="keys"):
new_key = key_map.get(key, key)
renamed_tensors[new_key] = handler.get_tensor(key)
pbar.update(1)
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(renamed_tensors, output_path, metadata)
# Stream tensors, rename on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
for key in tqdm(original_keys, desc="Renaming keys", unit="keys"):
new_key = key_map.get(key, key)
writer.write(new_key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
del renamed_tensors
cleanup_after_operation()
return output_path
requirements.txt
+1 -1
View File
@@ -1,5 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader==0.4.4
unifiedefficientloader==0.4.5
av