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
10 changed files with 1263 additions and 615 deletions
+2 -2
View File
@@ -1,7 +1,7 @@
from typing_extensions import override from typing_extensions import override
from comfy_api.latest import ComfyExtension, io 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 ( from .nodes.metakeys import (
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys, ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys CheckpointMetaKeys, EmbeddingMetaKeys
@@ -65,7 +65,7 @@ class ModelUtilsExtension(ComfyExtension):
LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative, LoRAResizeFrobenius, LoRAResizeCumulative,
# LoRA Multi-Merge # LoRA Multi-Merge
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
# LoRA Merge To Model # LoRA Merge To Model
LoRAMergeToModel, LoRAMergeToModel,
# Downloaders # 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. Native implementation of Low-Rank Adapter extraction with multiple rank selection modes.
No external dependencies required. No external dependencies required.
""" """
import fnmatch
import os import os
import re import re
import torch import torch
@@ -12,13 +13,12 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file
from .device_utils import ( from .device_utils import (
estimate_model_size, prepare_for_large_operation, estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities 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 # Pattern Matching
# ============================================================================= # =============================================================================
def _compile_patterns(pattern_string: str) -> list: def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
"""Compile regex patterns from whitespace-separated string.""" """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(): if not pattern_string or not pattern_string.strip():
return [] return []
patterns = [] patterns = []
for p in pattern_string.split(): for p in pattern_string.split():
p = p.strip() p = p.strip()
if p: if not p:
continue
if glob_mode:
patterns.append(p)
else:
try: try:
patterns.append(re.compile(p)) patterns.append(re.compile(p))
except re.error: except re.error:
@@ -414,8 +422,15 @@ def _compile_patterns(pattern_string: str) -> list:
return patterns return patterns
def _matches_any_pattern(key: str, patterns: list) -> bool: def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
"""Check if key matches any pattern.""" """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) return any(p.search(key) for p in patterns)
@@ -431,6 +446,7 @@ def extract_lora_from_files(
conv_param: float, conv_param: float,
device: str, device: str,
save_dtype: str, save_dtype: str,
output_path: str,
linear_max_rank: int = None, linear_max_rank: int = None,
conv_max_rank: int = None, conv_max_rank: int = None,
clamp_quantile: float = 0.99, clamp_quantile: float = 0.99,
@@ -441,9 +457,10 @@ def extract_lora_from_files(
svd_niter: int = 2, svd_niter: int = 2,
lazy_load: bool = True, lazy_load: bool = True,
force_clear_cache: 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: Args:
model_a_path: Finetuned model path model_a_path: Finetuned model path
@@ -453,6 +470,7 @@ def extract_lora_from_files(
conv_param: Mode parameter for conv layers conv_param: Mode parameter for conv layers
device: Computation device device: Computation device
save_dtype: Output dtype save_dtype: Output dtype
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers conv_max_rank: Max rank for conv layers
clamp_quantile: Quantile for weight clamping clamp_quantile: Quantile for weight clamping
@@ -460,16 +478,16 @@ def extract_lora_from_files(
skip_patterns_str: Patterns for layers to skip skip_patterns_str: Patterns for layers to skip
mismatch_mode: How to handle mismatches mismatch_mode: How to handle mismatches
chunk_large_layers: Enable chunked extraction 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 = { save_torch_dtype = {
"fp32": torch.float32, "fp32": torch.float32,
"fp16": torch.float16, "fp16": torch.float16,
"bf16": torch.bfloat16, "bf16": torch.bfloat16,
}.get(save_dtype, torch.float16) }.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 # Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path) 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): def _process_layer(key):
lora_name = _format_lora_key(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 return "skipped", None
# Load tensors with pinned memory for CUDA # 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 weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
) )
if lora_up is not None: 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_B.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}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
del weight_diff del weight_diff
return "chunked", layer_results return "chunked", layer_results
@@ -586,27 +603,31 @@ def extract_lora_from_files(
status = "full" status = "full"
else: else:
lora_down, lora_up, _ = result 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_A.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}.lora_B.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)
status = "extracted" status = "extracted"
del weight_diff del weight_diff
return status, layer_results return status, layer_results
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"): writer = IncrementalSafetensorsWriter(output_path)
status, layer_sd = _process_layer(key) writer.__enter__()
stats[status] += 1 try:
if layer_sd: for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
output_sd.update(layer_sd) status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache: if force_clear_cache:
import gc import gc
gc.collect() gc.collect()
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.cuda.empty_cache() 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, " print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped") f"{stats['full']} full, {stats['skipped']} skipped")
@@ -616,22 +637,12 @@ def extract_lora_from_files(
handler_b.__exit__(None, None, None) handler_b.__exit__(None, None, None)
cleanup_after_operation() 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] output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") return 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
# ============================================================================= # =============================================================================
@@ -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: def _format_lora_key(key: str) -> str:
""" """
Format the key for LoRA saving. 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"): if key.endswith(".weight"):
key = key[:-7] key = key[:-7]
# Handle ComfyUI Checkpoint wrapper # Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*)
if key.startswith("model.diffusion_model.") or key.startswith("net."): if key.startswith("model.diffusion_model."):
# Check if it's a Diffusers-format transformer inside a checkpoint inner_key = key[22:] # strip len("model.diffusion_model.")
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
return f"diffusion_model.{inner_key}" 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"): 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) # Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"): 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 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"]): 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}" return f"diffusion_model.{key}"
# Default fallback to diffusion_model # Default fallback
return f"diffusion_model.{key}" 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("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True, 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 @classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename, 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_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) 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, 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, 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): class LoRAExtractRatio(io.ComfyNode):
@@ -819,19 +798,21 @@ class LoRAExtractRatio(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, 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, 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_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) 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, 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, 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): class LoRAExtractQuantile(io.ComfyNode):
@@ -863,19 +844,21 @@ class LoRAExtractQuantile(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, 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, 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_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) 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, 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, 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): class LoRAExtractKnee(io.ComfyNode):
@@ -905,19 +888,21 @@ class LoRAExtractKnee(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, 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, 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_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) 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, 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, 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): class LoRAExtractFrobenius(io.ComfyNode):
@@ -949,16 +934,18 @@ class LoRAExtractFrobenius(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, 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, 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_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) 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, 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, 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)
+668 -224
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 import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io 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 .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] Original dim={network_dim}, alpha={network_alpha:.1f}, scale={scale:.3f}")
print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}") print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}")
output_sd = {} # Build metadata and output path before loop (all values known at this point)
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
if dynamic_method: if dynamic_method:
metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}" metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}"
metadata["ss_network_dim"] = "Dynamic" metadata["ss_network_dim"] = "Dynamic"
@@ -488,14 +428,78 @@ def resize_lora_file(
metadata["ss_network_dim"] = str(new_rank) metadata["ss_network_dim"] = str(new_rank)
metadata["ss_network_alpha"] = str(scale * new_rank) metadata["ss_network_alpha"] = str(scale * new_rank)
# Save
output_dir = folder_paths.get_folder_paths("loras")[0] output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path, metadata) fro_list = []
print(f"[LoRA Resize] Saved to {output_path}") 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 return output_path
finally: 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}") print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}")
# Common prefixes # Common prefixes
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."] BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."]
LORA_PREFIXES = [ LORA_PREFIXES = [
"lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_",
"lycoris_", "diffusion_model.", "transformer.", "unet." "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 = base_handler.metadata().copy() if base_handler.metadata() else {}
base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}" 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} stats = {"merged": 0, "copied": 0, "skipped": 0}
base_keys = list(base_handler.keys()) base_keys = list(base_handler.keys())
pbar = comfy.utils.ProgressBar(len(base_keys)) pbar = comfy.utils.ProgressBar(len(base_keys))
@@ -769,118 +777,116 @@ def merge_loras_to_model(
if verbose: if verbose:
print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...") print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...")
with torch.no_grad(): writer = IncrementalSafetensorsWriter(output_path, metadata=base_metadata)
for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): writer.__enter__()
# Check skip patterns try:
if _matches_any_pattern(base_key, skip_patterns): with torch.no_grad():
stats["skipped"] += 1 for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"):
pbar.update(1) # Check skip patterns
continue if _matches_any_pattern(base_key, skip_patterns):
stats["skipped"] += 1
pbar.update(1)
continue
# Load base weight # Load base weight
cpu_base = base_handler.get_tensor(base_key) cpu_base = base_handler.get_tensor(base_key)
# Only process weight tensors for LoRA merging # Only process weight tensors for LoRA merging
if base_key.endswith(".weight"): if base_key.endswith(".weight"):
core = extract_core_layer_base(base_key) core = extract_core_layer_base(base_key)
core_underscored = core.replace(".", "_") core_underscored = core.replace(".", "_")
# Check if any LoRA contributes to this layer # Check if any LoRA contributes to this layer
if core_underscored in lora_lookup: if core_underscored in lora_lookup:
# Transfer to GPU for computation # Transfer to GPU for computation
if device == 'cuda': if device == 'cuda':
base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32) 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"]
else: else:
# Standard LoRA format base_weight = cpu_base.to(device=device, dtype=torch.float32)
if "down" not in block_keys or "up" not in block_keys: del cpu_base
continue
cpu_down = info["handler"].get_tensor(block_keys["down"]) # Accumulate deltas from all contributing LoRAs
cpu_up = info["handler"].get_tensor(block_keys["up"]) for info, block_keys in lora_lookup[core_underscored]:
if device == 'cuda': is_full_diff = info["format_info"].get("is_full_diff", False)
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32) 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: else:
lora_down = cpu_down.to(device=device, dtype=torch.float32) # Standard LoRA format
lora_up = cpu_up.to(device=device, dtype=torch.float32) if "down" not in block_keys or "up" not in block_keys:
del cpu_down, cpu_up continue
# Get alpha cpu_down = info["handler"].get_tensor(block_keys["down"])
if "alpha" in block_keys: cpu_up = info["handler"].get_tensor(block_keys["up"])
alpha_tensor = info["handler"].get_tensor(block_keys["alpha"]) if device == 'cuda':
layer_alpha = float(alpha_tensor.item()) lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
else: lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
layer_alpha = float(info["network_dim"]) else:
layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0 lora_down = cpu_down.to(device=device, dtype=torch.float32)
effective_scale = layer_scale * info["weight"] lora_up = cpu_up.to(device=device, dtype=torch.float32)
del cpu_down, cpu_up
# Compute delta # Get alpha
is_conv = len(lora_down.shape) == 4 if "alpha" in block_keys:
if is_conv: alpha_tensor = info["handler"].get_tensor(block_keys["alpha"])
in_rank, in_size, kernel_size, k_ = lora_down.shape layer_alpha = float(alpha_tensor.item())
out_size, out_rank, _, _ = lora_up.shape else:
delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1) layer_alpha = float(info["network_dim"])
delta = delta.reshape(out_size, in_size, kernel_size, kernel_size) layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0
else: effective_scale = layer_scale * info["weight"]
delta = lora_up @ lora_down
del lora_down, lora_up
# Apply delta to base weight # Compute delta
base_weight = base_weight + effective_scale * delta is_conv = len(lora_down.shape) == 4
del delta 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 # Apply delta to base weight
output_sd[base_key] = base_weight.to(save_dtype).cpu().contiguous() base_weight = base_weight + effective_scale * delta
del base_weight del delta
stats["merged"] += 1
# 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: else:
# No LoRA contribution, copy as-is # Non-weight tensor (bias, norm, etc.), write as-is
output_sd[base_key] = cpu_base.to(save_dtype).contiguous() writer.write(base_key, cpu_base.to(save_dtype).contiguous())
del cpu_base del cpu_base
stats["copied"] += 1 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 import gc
gc.collect() gc.collect()
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()
pbar.update(1) pbar.update(1)
finally:
writer.__exit__(None, None, None)
if verbose: if verbose:
print(f"[LoRA Merge To Model] Done: {stats['merged']} merged, {stats['copied']} copied, {stats['skipped']} skipped") 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}") print(f"[LoRA Merge To Model] Saved to {output_path}")
return output_path return output_path
finally: finally:
+128 -55
View File
@@ -1,3 +1,4 @@
import fnmatch
import os import os
import re import re
import torch import torch
@@ -5,7 +6,6 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io 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_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
from .device_utils import ( from .device_utils import (
estimate_model_size, prepare_for_large_operation, cleanup_after_operation estimate_model_size, prepare_for_large_operation, cleanup_after_operation
@@ -25,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." return f"# Documentation File Not Found\n\nPlease ensure `{filename}` exists in the `docs` directory of the custom node."
def _compile_patterns(pattern_string): def _compile_patterns(pattern_string, glob_mode=False):
"""Compiles whitespace-separated regex patterns into a list of compiled patterns. """Compiles whitespace-separated patterns into a list.
Returns empty list if pattern_string is empty or whitespace-only. In regex mode (default): returns compiled re.Pattern objects.
Raises ValueError if any pattern is invalid regex. 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(): if not pattern_string or not pattern_string.strip():
return [] return []
patterns = [] patterns = []
for pattern in pattern_string.split(): for pattern in pattern_string.split():
try: if not pattern:
patterns.append(re.compile(pattern)) continue
except re.error as e: if glob_mode:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}") 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 return patterns
def _matches_any_pattern(key, patterns): def _matches_any_pattern(key, patterns, glob_mode=False):
"""Returns True if key matches any of the compiled regex patterns (substring search).""" """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: for pattern in patterns:
if pattern.search(key): if pattern.search(key):
return True return True
@@ -100,8 +112,9 @@ class MergerLogic:
alignment_mode = recipe_params.get('alignment_mode', 'pad/crop') alignment_mode = recipe_params.get('alignment_mode', 'pad/crop')
# Compile filter patterns # Compile filter patterns
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', '')) glob_mode = recipe_params.get('glob_patterns', False)
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', '')) 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 # Pre-compute key differences for logging
primary_keys = set(all_keys) primary_keys = set(all_keys)
@@ -162,7 +175,7 @@ class MergerLogic:
with torch.no_grad(): with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"): for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely # 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 discarded_keys += 1
continue continue
@@ -175,7 +188,7 @@ class MergerLogic:
del cpu_tensor del cpu_tensor
# Check exclude patterns - use Model A only, no merge # 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() t = tensor_a.detach().to(save_torch_dtype).cpu()
batch_buffer[key] = t batch_buffer[key] = t
@@ -296,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_checkpoint"), io.String.Input("output_filename", default="merged_2_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -315,9 +333,9 @@ class CheckpointTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -325,10 +343,11 @@ class CheckpointTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} model_names = {"model_a": model_a, "model_b": model_b}
@@ -355,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_model"), io.String.Input("output_filename", default="merged_2_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -374,9 +398,9 @@ class ModelTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -384,10 +408,11 @@ class ModelTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} model_names = {"model_a": model_a, "model_b": model_b}
@@ -414,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_textencoder"), io.String.Input("output_filename", default="merged_2_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -433,9 +463,9 @@ class TextEncoderTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -443,10 +473,11 @@ class TextEncoderTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} model_names = {"model_a": model_a, "model_b": model_b}
@@ -473,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_lora"), io.String.Input("output_filename", default="merged_2_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -492,9 +528,9 @@ class LoRATwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -502,10 +538,11 @@ class LoRATwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} model_names = {"model_a": model_a, "model_b": model_b}
@@ -532,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_embedding"), io.String.Input("output_filename", default="merged_2_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -551,9 +593,9 @@ class EmbeddingTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -561,10 +603,11 @@ class EmbeddingTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} model_names = {"model_a": model_a, "model_b": model_b}
@@ -594,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_checkpoint"), io.String.Input("output_filename", default="merged_3_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -613,9 +661,9 @@ class CheckpointThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -623,10 +671,11 @@ class CheckpointThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -654,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_model"), io.String.Input("output_filename", default="merged_3_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -673,9 +727,9 @@ class ModelThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -683,10 +737,11 @@ class ModelThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -714,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_textencoder"), io.String.Input("output_filename", default="merged_3_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -733,9 +793,9 @@ class TextEncoderThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -743,10 +803,11 @@ class TextEncoderThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -774,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_lora"), io.String.Input("output_filename", default="merged_3_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -793,9 +859,9 @@ class LoRAThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -803,10 +869,11 @@ class LoRAThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -834,13 +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("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("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("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.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_embedding"), io.String.Input("output_filename", default="merged_3_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
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("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"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
], ],
@@ -853,9 +925,9 @@ class EmbeddingThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, process_device: str, exclude_patterns: str, discard_patterns: 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') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -863,10 +935,11 @@ class EmbeddingThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_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, "output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32, "device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} 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 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()] THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
+22 -22
View File
@@ -4,12 +4,11 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open from safetensors import safe_open
from .utils import convert_pt_to_safetensors from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation 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, 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: if not patterns:
raise ValueError("No keys/patterns provided to prune.") 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 # 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] model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") 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 # Cleanup after operation
del pruned_tensors
cleanup_after_operation() cleanup_after_operation()
return output_path return output_path
+14 -14
View File
@@ -3,12 +3,11 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open from safetensors import safe_open
from .utils import convert_pt_to_safetensors from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation 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)) 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 # 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] model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") 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 # Cleanup after operation
del renamed_tensors
cleanup_after_operation() cleanup_after_operation()
return output_path return output_path
+1 -1
View File
@@ -1,7 +1,7 @@
[project] [project]
name = "ComfyUI-ModelUtils" name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files." description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.3.5" version = "0.4.1"
license = { file = "LICENSE" } license = { file = "LICENSE" }
classifiers = [ classifiers = [
+1 -1
View File
@@ -1,5 +1,5 @@
requests requests
Pillow Pillow
mutagen mutagen
unifiedefficientloader==0.4.4 unifiedefficientloader>=0.5.1
av av