Compare commits

..
Author SHA1 Message Date
Silver f6aac8cedf Merge branch 'main' into refactor/incremental-writer 2026-05-25 14:21:24 +02:00
silveroxides 7aa70d8173 Bump version from 0.4.0 to 0.4.1 2026-05-25 14:20:07 +02:00
silveroxides 26c8dffff9 Bump version from 0.3.8 to 0.4.0 2026-05-25 14:18:59 +02:00
silveroxides f3497658fb update unifiedefficientloader requirements 2026-05-25 14:17:58 +02:00
silveroxides 6747f2e8dc Fix block indices error 2026-05-25 14:16:59 +02:00
Silver 54bb6621d8 Bump version from 0.3.9 to 0.4.0 2026-05-21 11:14:38 +02:00
Silver d0d1063964 Update unifiedefficientloader version to 0.5.0 2026-05-21 11:14:19 +02:00
Silver 518e8e0b6d Bump version from 0.3.8 to 0.3.9 2026-05-16 16:33:59 +02:00
Silver fee30e36f5 Merge pull request #12 from silveroxides/refactor/incremental-writer
Refactor/incremental writer
2026-05-16 16:33:37 +02:00
silveroxides 94d2a93b91 fix layer renaming issues 2026-05-16 16:32:44 +02:00
silveroxides bb915f0599 refactor all saves to use IncrementalSafetensorsWriter 2026-05-16 11:07:20 +02:00
silveroxides f9a96d78a2 Bump version from 0.3.7 to 0.3.8 2026-05-10 18:52:38 +02:00
silveroxides 8ee4654c80 Improve layer filters in the merger nodes 2026-05-10 18:52:10 +02:00
silveroxides 73ce12cf9f Bump version from 0.3.6 to 0.3.7 2026-05-10 18:22:23 +02:00
silveroxides 21c0972a5b Refine LoRA extraction to output accurate layer formats across the board and add experimental new DARE+Ties merge operation variant 2026-05-10 18:21:54 +02:00
silveroxides 3bcd2bf3f8 Bump version from 0.3.5 to 0.3.6 2026-05-09 03:09:30 +02:00
silveroxides 29f03fe979 Add new 'enhanced' PowerUp/DARE-Ties merge modes 2026-05-09 03:08:57 +02:00
11 changed files with 1256 additions and 762 deletions
+2 -2
View File
@@ -1,7 +1,7 @@
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE
from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced
from .nodes.metakeys import (
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys
@@ -65,7 +65,7 @@ class ModelUtilsExtension(ComfyExtension):
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
# LoRA Multi-Merge
LoRAMultiMerge, LoRAMultiMergeDARE,
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
# LoRA Merge To Model
LoRAMergeToModel,
# Downloaders
+117 -130
View File
@@ -4,6 +4,7 @@ LoRA Extraction via SVD decomposition.
Native implementation of Low-Rank Adapter extraction with multiple rank selection modes.
No external dependencies required.
"""
import fnmatch
import os
import re
import torch
@@ -12,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
@@ -399,14 +399,22 @@ def _extract_chunked_layer(
# Pattern Matching
# =============================================================================
def _compile_patterns(pattern_string: str) -> list:
"""Compile regex patterns from whitespace-separated string."""
def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
"""Compile patterns from whitespace-separated string.
In regex mode (default) returns compiled re.Pattern objects.
In glob mode returns plain strings; fnmatch handles matching.
"""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for p in pattern_string.split():
p = p.strip()
if p:
if not p:
continue
if glob_mode:
patterns.append(p)
else:
try:
patterns.append(re.compile(p))
except re.error:
@@ -414,8 +422,15 @@ def _compile_patterns(pattern_string: str) -> list:
return patterns
def _matches_any_pattern(key: str, patterns: list) -> bool:
"""Check if key matches any pattern."""
def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
"""Check if key matches any pattern.
Glob mode: each pattern is matched as a substring glob using fnmatch
(dots are literal, * matches any sequence of characters).
Regex mode: each pattern is a compiled re.Pattern, matched via search().
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
return any(p.search(key) for p in patterns)
@@ -431,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,
@@ -441,9 +457,10 @@ def extract_lora_from_files(
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
) -> dict[str, torch.Tensor]:
glob_skip_patterns: bool = False,
) -> 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
@@ -453,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
@@ -460,16 +478,16 @@ def extract_lora_from_files(
skip_patterns_str: Patterns for layers to skip
mismatch_mode: How to handle mismatches
chunk_large_layers: Enable chunked extraction
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,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
@@ -489,7 +507,7 @@ def extract_lora_from_files(
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns):
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
# Load tensors with pinned memory for CUDA
@@ -570,9 +588,8 @@ def extract_lora_from_files(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
@@ -586,27 +603,31 @@ def extract_lora_from_files(
status = "full"
else:
lora_down, lora_up, _ = result
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype)
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
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")
@@ -616,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")
# =============================================================================
@@ -647,85 +658,49 @@ 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.
Standardizes to 'diffusion_model.' or 'transformer.' prefix.
Always standardizes to 'diffusion_model.' prefix for maximum compatibility.
ComfyUI maps 'diffusion_model.<layer>' generically for all model types, making
it universally compatible. The 'transformer.' prefix only works for a subset of
model types via model-specific mapping code and is intentionally not used here.
"""
if key.endswith(".weight"):
key = key[:-7]
# Handle ComfyUI Checkpoint wrapper
if key.startswith("model.diffusion_model.") or key.startswith("net."):
# Check if it's a Diffusers-format transformer inside a checkpoint
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # len("model.diffusion_model.")
else:
inner_key = key[4:] # len("net.")
if inner_key.startswith("transformer_blocks") or inner_key.startswith("single_transformer_blocks"):
return f"transformer.{inner_key}"
# Standard checkpoint or Original Flux
# Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*)
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # strip len("model.diffusion_model.")
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys
# Handle net.* wrapper
if key.startswith("net."):
inner_key = key[4:] # strip len("net.")
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys (transformer_blocks.* / single_transformer_blocks.* without prefix)
if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"):
return f"transformer.{key}"
return f"diffusion_model.{key}"
# Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"):
core = key[10:]
dotted = _reconstruct_dots(core)
# Check if reconstructed key implies transformer format
if dotted.startswith("transformer_blocks") or dotted.startswith("single_transformer_blocks"):
return f"transformer.{dotted}"
return f"diffusion_model.{dotted}"
# Handle already correct prefixes
if key.startswith("diffusion_model.") or key.startswith("transformer."):
return key
# Handle known Unet blocks (Diffusers format UNet)
# Handle already-prefixed keys
if key.startswith("diffusion_model."):
return key
# Absorb now-invalid transformer. prefix into diffusion_model. for compatibility
if key.startswith("transformer."):
return f"diffusion_model.{key[12:]}"
# Handle known Diffusers UNet blocks
if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]):
return f"diffusion_model.{key}"
# Default fallback to diffusion_model
# Default fallback
return f"diffusion_model.{key}"
@@ -744,7 +719,10 @@ def _get_common_inputs():
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip"),
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
]
@@ -775,19 +753,20 @@ class LoRAExtractFixed(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
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
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):
@@ -819,19 +798,21 @@ class LoRAExtractRatio(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
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
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):
@@ -863,19 +844,21 @@ class LoRAExtractQuantile(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
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
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):
@@ -905,19 +888,21 @@ class LoRAExtractKnee(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
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
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):
@@ -949,16 +934,18 @@ class LoRAExtractFrobenius(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
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
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)
+664 -256
View File
File diff suppressed because it is too large Load Diff
+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:
+129 -92
View File
@@ -1,3 +1,4 @@
import fnmatch
import os
import re
import torch
@@ -5,9 +6,7 @@ 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 .merger_utils import parse_custom_weights, get_custom_parameters
from .device_utils import (
estimate_model_size, prepare_for_large_operation, cleanup_after_operation
)
@@ -26,26 +25,38 @@ def load_documentation_from_file(filename):
return f"# Documentation File Not Found\n\nPlease ensure `{filename}` exists in the `docs` directory of the custom node."
def _compile_patterns(pattern_string):
"""Compiles whitespace-separated regex patterns into a list of compiled patterns.
def _compile_patterns(pattern_string, glob_mode=False):
"""Compiles whitespace-separated patterns into a list.
Returns empty list if pattern_string is empty or whitespace-only.
Raises ValueError if any pattern is invalid regex.
In regex mode (default): returns compiled re.Pattern objects.
Raises ValueError if any pattern is invalid regex.
In glob mode: returns plain strings; fnmatch handles matching.
"""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for pattern in pattern_string.split():
try:
patterns.append(re.compile(pattern))
except re.error as e:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}")
if not pattern:
continue
if glob_mode:
patterns.append(pattern)
else:
try:
patterns.append(re.compile(pattern))
except re.error as e:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}")
return patterns
def _matches_any_pattern(key, patterns):
"""Returns True if key matches any of the compiled regex patterns (substring search)."""
def _matches_any_pattern(key, patterns, glob_mode=False):
"""Returns True if key matches any pattern.
Glob mode: fnmatch substring match — dots are literal, * matches any sequence.
Regex mode: compiled re.Pattern substring search.
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
for pattern in patterns:
if pattern.search(key):
return True
@@ -101,11 +112,9 @@ class MergerLogic:
alignment_mode = recipe_params.get('alignment_mode', 'pad/crop')
# Compile filter patterns
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''))
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''))
# Parse custom per-layer weights
parsed_custom_weights = parse_custom_weights(recipe_params.get('custom_weights', ''))
glob_mode = recipe_params.get('glob_patterns', False)
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''), glob_mode=glob_mode)
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''), glob_mode=glob_mode)
# Pre-compute key differences for logging
primary_keys = set(all_keys)
@@ -166,7 +175,7 @@ class MergerLogic:
with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode):
discarded_keys += 1
continue
@@ -179,7 +188,7 @@ class MergerLogic:
del cpu_tensor
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
if _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode):
t = tensor_a.detach().to(save_torch_dtype).cpu()
batch_buffer[key] = t
@@ -197,9 +206,7 @@ class MergerLogic:
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
try:
# Apply per-layer parameter overrides (regex/substring matching)
layer_params = get_custom_parameters(key, parsed_custom_weights, recipe_params)
recipe = calc_mode_class.create_recipe(key=key, **layer_params)
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
result = recipe.merge()
except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR:
@@ -302,14 +309,18 @@ class CheckpointTwoMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -322,10 +333,9 @@ class CheckpointTwoMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -333,11 +343,11 @@ class CheckpointTwoMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -364,14 +374,18 @@ class ModelTwoMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -384,10 +398,9 @@ class ModelTwoMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -395,11 +408,11 @@ class ModelTwoMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -426,14 +439,18 @@ class TextEncoderTwoMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -446,10 +463,9 @@ class TextEncoderTwoMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -457,11 +473,11 @@ class TextEncoderTwoMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -488,14 +504,18 @@ class LoRATwoMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -508,10 +528,9 @@ class LoRATwoMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -519,11 +538,11 @@ class LoRATwoMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -550,14 +569,18 @@ class EmbeddingTwoMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -570,10 +593,9 @@ class EmbeddingTwoMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -581,11 +603,11 @@ class EmbeddingTwoMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -615,14 +637,18 @@ class CheckpointThreeMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -635,10 +661,9 @@ class CheckpointThreeMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -646,11 +671,11 @@ class CheckpointThreeMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -678,14 +703,18 @@ class ModelThreeMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -698,10 +727,9 @@ class ModelThreeMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -709,11 +737,11 @@ class ModelThreeMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -741,14 +769,18 @@ class TextEncoderThreeMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -761,10 +793,9 @@ class TextEncoderThreeMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -772,11 +803,11 @@ class TextEncoderThreeMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -804,14 +835,18 @@ class LoRAThreeMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -824,10 +859,9 @@ class LoRAThreeMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -835,11 +869,11 @@ class LoRAThreeMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -867,14 +901,18 @@ class EmbeddingThreeMerger(io.ComfyNode):
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -887,10 +925,9 @@ class EmbeddingThreeMerger(io.ComfyNode):
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -898,11 +935,11 @@ class EmbeddingThreeMerger(io.ComfyNode):
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
+139 -1
View File
@@ -793,5 +793,143 @@ class DARETIESMode(CalcMode):
return op
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
class EnhancedPowerUpOp(Operation):
def __init__(self, key, alpha, beta, gamma, delta, seed, *sources):
super().__init__(key, *sources)
self.mask_power = alpha # Mask power
self.min_keep_prob = beta # Min keep prob
self.mask_smooth = gamma # Smoothness (0 or 1)
self.delta = delta # Final multiplier
self.seed = seed
def oper(self, a, b):
if a is None or b is None: return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta_val = (b - a).float()
abs_delta = torch.abs(delta_val)
max_delta = torch.max(abs_delta)
if max_delta == 0: return torch.zeros_like(delta_val)
prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(prob, generator=rng)
interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
rescaled_delta = delta_val * interpolated_mask
return rescaled_delta
class EnhancedDARETIESOp(Operation):
def __init__(self, key, alpha, beta, gamma, delta, epsilon, zeta, seed, *sources):
super().__init__(key, *sources)
self.alpha = alpha # Unused directly for drop scale
self.beta = beta # TIES trim quantile
self.gamma = gamma # TIES lambda scale
self.mask_power = delta
self.min_keep_prob = epsilon
self.mask_smooth = zeta
self.seed = seed
def oper(self, a, b):
if a is None or b is None: return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta_val = (b - a).float()
abs_delta = torch.abs(delta_val)
max_delta = torch.max(abs_delta)
if max_delta > 0:
prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(prob, generator=rng)
interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
delta_val = delta_val * interpolated_mask
if self.beta > 0:
flat_abs = delta_val.abs().flatten()
k = max(1, int(flat_abs.numel() * self.beta))
threshold = torch.kthvalue(flat_abs, k).values
delta_val = torch.where(delta_val.abs() < threshold, torch.zeros_like(delta_val), delta_val)
dominant_sign = torch.sign(delta_val)
delta_val = torch.where(torch.sign(delta_val) == dominant_sign, delta_val, torch.zeros_like(delta_val))
return (a + self.gamma * delta_val).to(a.dtype)
class EnhancedPowerUp(CalcMode):
name = 'Power-Up Enhanced (DARE)'
description = 'Adds capabilities of B to A using Enhanced Drop and Rescale based on relative tensor differences.'
models_used = ['A', 'B']
param_docs = {
'alpha': 'Mask power (Curve. 2.0 = quadratic).',
'beta': 'Minimum keep probability (Floor to prevent explosion).',
'gamma': 'Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale).',
'delta': 'Multiplier for the final difference.',
}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
deltahat = EnhancedPowerUpOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b)
deltahat.alignment_mode = alignment_mode
multiplied = Multiply(key, kwargs['delta'], deltahat)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class EnhancedDARETIESMode(CalcMode):
name = 'Power-Up Enhanced (DARE+TIES)'
description = 'Adds capabilities of B to A using Enhanced DARE mask followed by TIES.'
models_used = ['A', 'B']
param_docs = {
'alpha': 'Unused (Drop scale handled via mask interpolation).',
'beta': 'TIES trim quantile (fraction of smallest-magnitude delta zeroed).',
'gamma': 'TIES Lambda scale.',
'delta': 'Mask power (Curve. 2.0 = quadratic).',
'epsilon': 'Minimum keep probability (Floor to prevent explosion).',
'zeta': 'Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale).',
}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
op = EnhancedDARETIESOp(
key,
kwargs.get('alpha', 0.5),
kwargs.get('beta', 0.2),
kwargs.get('gamma', 1.0),
kwargs.get('delta', 0.5),
kwargs.get('epsilon', 0.5),
kwargs.get('zeta', 0.5),
kwargs['seed'],
a, b,
)
op.alignment_mode = alignment_mode
return op
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode(), EnhancedPowerUp(), EnhancedDARETIESMode()]
THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
+1 -83
View File
@@ -1,8 +1,7 @@
"""
Utilities for model merging operations.
"""
from typing import Dict, List, Tuple, Any
import re
from typing import Dict, List
def get_union_keys(*handlers) -> List[str]:
"""Returns the union of all tensor keys across multiple model handlers."""
@@ -22,84 +21,3 @@ def get_intersection_keys(*handlers) -> List[str]:
for handler in handlers[1:]:
all_keys &= set(handler.keys())
return sorted(all_keys)
def parse_custom_weights(custom_weights_str: str) -> List[Tuple[re.Pattern, Dict[str, float]]]:
"""Parses multiline string of custom weights.
Format: pattern:weight OR pattern:param=value,param2=value
Returns list of (compiled_regex, parameter_dict).
"""
if not custom_weights_str or not custom_weights_str.strip():
return []
results = []
lines = custom_weights_str.strip().split('\n')
for line in lines:
line = line.strip()
if not line or line.startswith('#'):
continue
if ':' not in line:
continue
pattern_str, weight_str = line.rsplit(':', 1)
pattern_str = pattern_str.strip()
weight_str = weight_str.strip()
try:
pattern = re.compile(pattern_str)
except re.error:
# Fallback to literal substring if regex invalid
pattern = re.compile(re.escape(pattern_str))
params = {}
if '=' in weight_str:
# param=value,param2=value
parts = weight_str.split(',')
for part in parts:
if '=' in part:
k, v = part.split('=', 1)
k = k.strip()
# Map common aliases
if k in ['a', 'alpha']: k = 'alpha'
elif k in ['b', 'beta']: k = 'beta'
elif k in ['g', 'gamma']: k = 'gamma'
elif k in ['d', 'delta']: k = 'delta'
elif k in ['w', 'weight']: k = 'weight'
elif re.match(r'^w[1-8]$', k): pass # Keep w1-w8 as is
try:
params[k] = float(v.strip())
except ValueError:
pass
else:
# single float or comma-separated list of floats (for multi-LoRA)
if ',' in weight_str:
parts = weight_str.split(',')
for i, p in enumerate(parts):
try:
params[f'w{i+1}'] = float(p.strip())
except ValueError:
pass
else:
try:
val = float(weight_str)
params['alpha'] = val # default for merger
params['weight'] = val # default for lora_merger
except ValueError:
pass
if params:
results.append((pattern, params))
return results
def get_custom_parameters(key: str, parsed_weights: List[Tuple[re.Pattern, Dict[str, float]]], base_params: Dict[str, Any]) -> Dict[str, Any]:
"""Returns updated parameters if key matches any pattern."""
new_params = base_params.copy()
for pattern, custom_dict in parsed_weights:
if pattern.search(key):
new_params.update(custom_dict)
break # First match wins
return new_params
+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
+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
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.3.5"
version = "0.4.1"
license = { file = "LICENSE" }
classifiers = [
+1 -1
View File
@@ -1,5 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader==0.4.4
unifiedefficientloader>=0.5.1
av