Compare commits

..
Author SHA1 Message Date
silveroxides 892dd733d3 Add node to extract text encoder LoRA 2026-05-03 18:30:45 +02:00
10 changed files with 675 additions and 1268 deletions
+4 -4
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, LoRAMultiMergeDAREEnhanced from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE
from .nodes.metakeys import ( from .nodes.metakeys import (
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys, ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys CheckpointMetaKeys, EmbeddingMetaKeys
@@ -22,7 +22,7 @@ from .nodes.merger import (
) )
from .nodes.lora_extract_svd import ( from .nodes.lora_extract_svd import (
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius LoRAExtractKnee, LoRAExtractFrobenius, LoRATextEncoderExtractFixed
) )
from .nodes.lora_resize import ( from .nodes.lora_resize import (
LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFixed, LoRAResizeRatio,
@@ -60,12 +60,12 @@ class ModelUtilsExtension(ComfyExtension):
CheckpointThreeMerger, EmbeddingThreeMerger, CheckpointThreeMerger, EmbeddingThreeMerger,
# LoRA Extraction # LoRA Extraction
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius, LoRAExtractKnee, LoRAExtractFrobenius, LoRATextEncoderExtractFixed,
# LoRA Resize # LoRA Resize
LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative, LoRAResizeFrobenius, LoRAResizeCumulative,
# LoRA Multi-Merge # LoRA Multi-Merge
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced, LoRAMultiMerge, LoRAMultiMergeDARE,
# LoRA Merge To Model # LoRA Merge To Model
LoRAMergeToModel, LoRAMergeToModel,
# Downloaders # Downloaders
+187 -119
View File
@@ -4,7 +4,6 @@ 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
@@ -13,12 +12,13 @@ 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, IncrementalSafetensorsWriter from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
@@ -399,22 +399,14 @@ def _extract_chunked_layer(
# Pattern Matching # Pattern Matching
# ============================================================================= # =============================================================================
def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list: def _compile_patterns(pattern_string: str) -> list:
"""Compile patterns from whitespace-separated string. """Compile regex 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 not p: if 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:
@@ -422,15 +414,8 @@ def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
return patterns return patterns
def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool: def _matches_any_pattern(key: str, patterns: list) -> 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)
@@ -446,7 +431,6 @@ 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,
@@ -457,10 +441,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,
glob_skip_patterns: bool = False, is_text_encoder: bool = False,
) -> None: ) -> dict[str, torch.Tensor]:
""" """
Extract LoRA from difference between two models, writing incrementally to disk. Extract LoRA from difference between two models.
Args: Args:
model_a_path: Finetuned model path model_a_path: Finetuned model path
@@ -470,7 +454,6 @@ 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
@@ -478,16 +461,17 @@ 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). is_text_encoder: Whether extracting from text encoder models
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, glob_mode=glob_skip_patterns) skip_patterns = _compile_patterns(skip_patterns_str)
# 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)
@@ -505,9 +489,9 @@ def extract_lora_from_files(
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key): def _process_layer(key):
lora_name = _format_lora_key(key) lora_name = _format_lora_key(key, is_text_encoder=is_text_encoder)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): if _matches_any_pattern(key, skip_patterns):
return "skipped", None return "skipped", None
# Load tensors with pinned memory for CUDA # Load tensors with pinned memory for CUDA
@@ -588,8 +572,9 @@ 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_B.weight"] = lora_up.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_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
del weight_diff del weight_diff
return "chunked", layer_results return "chunked", layer_results
@@ -603,31 +588,27 @@ 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_A.weight"] = lora_down.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_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype)
status = "extracted" status = "extracted"
del weight_diff del weight_diff
return status, layer_results return status, layer_results
writer = IncrementalSafetensorsWriter(output_path) for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
writer.__enter__() status, layer_sd = _process_layer(key)
try: stats[status] += 1
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"): if layer_sd:
status, layer_sd = _process_layer(key) output_sd.update(layer_sd)
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")
@@ -637,12 +618,22 @@ 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)
return 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)
print(f"[LoRA Extract] Saved to {output_path}")
return output_path
# ============================================================================= # =============================================================================
@@ -658,49 +649,98 @@ def _get_model_inputs():
] ]
def _format_lora_key(key: str) -> str: def _get_clip_model_inputs():
return [
io.Combo.Input("model_a", options=folder_paths.get_filename_list("text_encoders"),
tooltip="Finetuned CLIP/Text Encoder (A - B = LoRA)"),
io.Combo.Input("model_b", options=folder_paths.get_filename_list("text_encoders"),
tooltip="Base CLIP/Text Encoder (A - B = LoRA)"),
]
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, is_text_encoder: bool = False) -> str:
""" """
Format the key for LoRA saving. Format the key for LoRA saving.
Always standardizes to 'diffusion_model.' prefix for maximum compatibility. Standardizes to 'diffusion_model.' or 'transformer.' prefix.
If is_text_encoder is True, use 'text_encoders.transformer.'
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 (ldm/sgm format: model.diffusion_model.*) if is_text_encoder:
if key.startswith("model.diffusion_model."): return f"text_encoders.transformer.{key}"
inner_key = key[22:] # strip len("model.diffusion_model.")
# Handle ComfyUI Checkpoint wrapper
if key.startswith("model.diffusion_model.") or key.startswith("net."):
# Check if it's a Diffusers-format transformer inside a checkpoint
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # len("model.diffusion_model.")
else:
inner_key = key[4:] # len("net.")
if inner_key.startswith("transformer_blocks") or inner_key.startswith("single_transformer_blocks"):
return f"transformer.{inner_key}"
# Standard checkpoint or Original Flux
return f"diffusion_model.{inner_key}" return f"diffusion_model.{inner_key}"
# Handle net.* wrapper # Handle direct Diffusers keys
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"diffusion_model.{key}" return f"transformer.{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 already-prefixed keys # Handle known Unet blocks (Diffusers format UNet)
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 # Default fallback to diffusion_model
return f"diffusion_model.{key}" return f"diffusion_model.{key}"
@@ -719,10 +759,7 @@ 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="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"), tooltip="Regex patterns for layers to skip"),
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."),
] ]
@@ -753,20 +790,19 @@ 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, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: save_dtype, device, 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)
extract_lora_from_files( output_sd = 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, output_path, linear_dim, conv_dim, device, save_dtype, 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, glob_skip_patterns lazy_load, force_clear_cache
) )
return io.NodeOutput(output_path) return io.NodeOutput(_save_lora(output_sd, output_filename))
class LoRAExtractRatio(io.ComfyNode): class LoRAExtractRatio(io.ComfyNode):
@@ -798,21 +834,19 @@ 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, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: save_dtype, device, 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)
extract_lora_from_files( output_sd = 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, output_path, linear_max_rank, conv_max_rank, device, save_dtype, 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(output_path) return io.NodeOutput(_save_lora(output_sd, output_filename))
class LoRAExtractQuantile(io.ComfyNode): class LoRAExtractQuantile(io.ComfyNode):
@@ -844,21 +878,19 @@ 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, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: save_dtype, device, 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)
extract_lora_from_files( output_sd = 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, output_path, linear_max_rank, conv_max_rank, device, save_dtype, 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(output_path) return io.NodeOutput(_save_lora(output_sd, output_filename))
class LoRAExtractKnee(io.ComfyNode): class LoRAExtractKnee(io.ComfyNode):
@@ -888,21 +920,19 @@ 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, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: save_dtype, device, 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)
extract_lora_from_files( output_sd = 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, output_path, linear_max_rank, conv_max_rank, device, save_dtype, 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(output_path) return io.NodeOutput(_save_lora(output_sd, output_filename))
class LoRAExtractFrobenius(io.ComfyNode): class LoRAExtractFrobenius(io.ComfyNode):
@@ -934,18 +964,56 @@ 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, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: save_dtype, device, 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)
extract_lora_from_files( output_sd = 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, output_path, linear_max_rank, conv_max_rank, device, save_dtype, 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(output_path) return io.NodeOutput(_save_lora(output_sd, output_filename))
class LoRATextEncoderExtractFixed(io.ComfyNode):
"""Extract LoRA from Text Encoder with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRATextEncoderExtractFixed",
display_name="LoRA Extract (Text Encoder, Fixed Rank)",
category="ModelUtils/LoRA",
description="Extract LoRA from Text Encoder/CLIP models with specified fixed rank.",
inputs=[
*_get_clip_model_inputs(),
io.Int.Input("linear_dim", default=32, min=1, max=16384,
tooltip="Rank for linear layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, 0,
device, save_dtype, linear_dim, 0,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, is_text_encoder=True
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
+219 -663
View File
File diff suppressed because it is too large Load Diff
+166 -172
View File
@@ -12,10 +12,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 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, IncrementalSafetensorsWriter from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from typing import Optional, Dict, Tuple, List
@@ -418,7 +419,66 @@ 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}")
# Build metadata and output path before loop (all values known at this point) output_sd = {}
fro_list = []
pbar = comfy.utils.ProgressBar(len(pairs))
with torch.no_grad():
for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"):
if "down" not in block_keys or "up" not in block_keys:
pbar.update(1)
continue
lora_down = handler.get_tensor(block_keys["down"])
lora_up = handler.get_tensor(block_keys["up"])
is_conv = len(lora_down.shape) == 4
# Transfer to GPU with pinned memory if available
if device == 'cuda':
lora_down = transfer_to_gpu_pinned(lora_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(lora_up, device, torch.float32)
# Merge and re-extract
if is_conv:
weight = _merge_conv(lora_down, lora_up, device)
result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
else:
weight = _merge_linear(lora_down, lora_up, device)
result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
del weight, lora_down, lora_up
fro_list.append(result['fro_retained'])
# Store using same format suffixes as input
down_suffix = format_info["down_suffix"]
up_suffix = format_info["up_suffix"]
alpha_suffix = format_info["alpha_suffix"]
# Standardize prefix using heuristics if possible
new_block_name = _format_lora_key(block_name)
output_sd[f"{new_block_name}{down_suffix}"] = result["lora_down"].to(save_dtype)
output_sd[f"{new_block_name}{up_suffix}"] = result["lora_up"].to(save_dtype)
if alpha_suffix:
output_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype)
del result
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
if verbose and fro_list:
import numpy as np
avg_fro = np.mean(fro_list)
std_fro = np.std(fro_list)
print(f"[LoRA Resize] Average Frobenius retention: {avg_fro:.1%} ± {std_fro:.3f}")
# Update metadata
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"
@@ -428,78 +488,14 @@ 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")
fro_list = [] save_file(output_sd, output_path, metadata)
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}") print(f"[LoRA Resize] Saved to {output_path}")
return output_path return output_path
finally: finally:
@@ -719,7 +715,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.", "net."] BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."]
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."
@@ -765,11 +761,7 @@ 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}"
# Build output path before loop output_sd = {}
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))
@@ -777,116 +769,118 @@ 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...")
writer = IncrementalSafetensorsWriter(output_path, metadata=base_metadata) with torch.no_grad():
writer.__enter__() for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"):
try: # Check skip patterns
with torch.no_grad(): if _matches_any_pattern(base_key, skip_patterns):
for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): stats["skipped"] += 1
# Check skip patterns pbar.update(1)
if _matches_any_pattern(base_key, skip_patterns): continue
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:
# Standard LoRA format
if "down" not in block_keys or "up" not in block_keys:
continue
cpu_down = info["handler"].get_tensor(block_keys["down"])
cpu_up = info["handler"].get_tensor(block_keys["up"])
if device == 'cuda':
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
else:
lora_down = cpu_down.to(device=device, dtype=torch.float32)
lora_up = cpu_up.to(device=device, dtype=torch.float32)
del cpu_down, cpu_up
# Get alpha
if "alpha" in block_keys:
alpha_tensor = info["handler"].get_tensor(block_keys["alpha"])
layer_alpha = float(alpha_tensor.item())
else:
layer_alpha = float(info["network_dim"])
layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0
effective_scale = layer_scale * info["weight"]
# Compute delta
is_conv = len(lora_down.shape) == 4
if is_conv:
in_rank, in_size, kernel_size, k_ = lora_down.shape
out_size, out_rank, _, _ = lora_up.shape
delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
delta = delta.reshape(out_size, in_size, kernel_size, kernel_size)
else:
delta = lora_up @ lora_down
del lora_down, lora_up
# Apply delta to base weight
base_weight = base_weight + effective_scale * delta
del delta
# Write merged weight
writer.write(base_key, base_weight.to(save_dtype).cpu().contiguous())
del base_weight
stats["merged"] += 1
else: else:
# No LoRA contribution, write as-is base_weight = cpu_base.to(device=device, dtype=torch.float32)
writer.write(base_key, cpu_base.to(save_dtype).contiguous()) del cpu_base
del cpu_base
stats["copied"] += 1 # 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:
# Standard LoRA format
if "down" not in block_keys or "up" not in block_keys:
continue
cpu_down = info["handler"].get_tensor(block_keys["down"])
cpu_up = info["handler"].get_tensor(block_keys["up"])
if device == 'cuda':
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
else:
lora_down = cpu_down.to(device=device, dtype=torch.float32)
lora_up = cpu_up.to(device=device, dtype=torch.float32)
del cpu_down, cpu_up
# Get alpha
if "alpha" in block_keys:
alpha_tensor = info["handler"].get_tensor(block_keys["alpha"])
layer_alpha = float(alpha_tensor.item())
else:
layer_alpha = float(info["network_dim"])
layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0
effective_scale = layer_scale * info["weight"]
# Compute delta
is_conv = len(lora_down.shape) == 4
if is_conv:
in_rank, in_size, kernel_size, k_ = lora_down.shape
out_size, out_rank, _, _ = lora_up.shape
delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
delta = delta.reshape(out_size, in_size, kernel_size, kernel_size)
else:
delta = lora_up @ lora_down
del lora_down, lora_up
# Apply delta to base weight
base_weight = base_weight + effective_scale * delta
del delta
# Store merged weight
output_sd[base_key] = base_weight.to(save_dtype).cpu().contiguous()
del base_weight
stats["merged"] += 1
else: else:
# Non-weight tensor (bias, norm, etc.), write as-is # No LoRA contribution, copy as-is
writer.write(base_key, cpu_base.to(save_dtype).contiguous()) output_sd[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:
+60 -133
View File
@@ -1,4 +1,3 @@
import fnmatch
import os import os
import re import re
import torch import torch
@@ -6,6 +5,7 @@ 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,38 +25,26 @@ 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, glob_mode=False): def _compile_patterns(pattern_string):
"""Compiles whitespace-separated patterns into a list. """Compiles whitespace-separated regex patterns into a list of compiled patterns.
In regex mode (default): returns compiled re.Pattern objects. Returns empty list if pattern_string is empty or whitespace-only.
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():
if not pattern: try:
continue patterns.append(re.compile(pattern))
if glob_mode: except re.error as e:
patterns.append(pattern) raise ValueError(f"Invalid regex pattern '{pattern}': {e}")
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, glob_mode=False): def _matches_any_pattern(key, patterns):
"""Returns True if key matches any pattern. """Returns True if key matches any of the compiled regex patterns (substring search)."""
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
@@ -112,9 +100,8 @@ 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
glob_mode = recipe_params.get('glob_patterns', False) exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''))
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''), glob_mode=glob_mode) discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''))
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)
@@ -175,7 +162,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, glob_mode=glob_mode): if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1 discarded_keys += 1
continue continue
@@ -188,9 +175,9 @@ 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, glob_mode=glob_mode): if _matches_any_pattern(key, exclude_patterns):
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
batch_bytes += t.numel() * t.element_size() batch_bytes += t.numel() * t.element_size()
batch_key_count += 1 batch_key_count += 1
@@ -210,7 +197,7 @@ class MergerLogic:
result = recipe.merge() result = recipe.merge()
except MissingTensorError as e: except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR: if mismatch_mode == MissingTensorBehavior.ERROR:
_flush_batch() _flush_batch()
writer.__exit__(None, None, None) writer.__exit__(None, None, None)
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}") raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
@@ -225,7 +212,7 @@ class MergerLogic:
if isinstance(result, dict): if isinstance(result, dict):
for r_key, r_tensor in result.items(): for r_key, r_tensor in result.items():
t = r_tensor.detach().to(save_torch_dtype).cpu() t = r_tensor.detach().to(save_torch_dtype).cpu()
batch_buffer[r_key] = t batch_buffer[r_key] = t
batch_bytes += t.numel() * t.element_size() batch_bytes += t.numel() * t.element_size()
else: else:
@@ -239,11 +226,11 @@ class MergerLogic:
result = result[slices] result = result[slices]
t = result.detach().to(save_torch_dtype).cpu() t = result.detach().to(save_torch_dtype).cpu()
batch_buffer[key] = t batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size() batch_bytes += t.numel() * t.element_size()
batch_key_count += 1 batch_key_count += 1
# Clean up references to allow GC immediately. # Clean up references to allow GC immediately.
@@ -309,18 +296,13 @@ 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=2.0, 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("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"),
], ],
@@ -333,9 +315,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -343,11 +325,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -374,18 +355,13 @@ 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=2.0, 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("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"),
], ],
@@ -398,9 +374,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -408,11 +384,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -439,18 +414,13 @@ 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=2.0, 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("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"),
], ],
@@ -463,9 +433,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -473,11 +443,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -504,18 +473,13 @@ 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=2.0, 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("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"),
], ],
@@ -528,9 +492,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -538,11 +502,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -569,18 +532,13 @@ 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=2.0, 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("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"),
], ],
@@ -593,9 +551,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -603,11 +561,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -637,18 +594,13 @@ 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=2.0, 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("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"),
], ],
@@ -661,9 +613,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -671,11 +623,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -703,18 +654,13 @@ 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=2.0, 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("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"),
], ],
@@ -727,9 +673,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -737,11 +683,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -769,18 +714,13 @@ 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=2.0, 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("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"),
], ],
@@ -793,9 +733,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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,11 +743,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -835,18 +774,13 @@ 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=2.0, 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("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"),
], ],
@@ -859,9 +793,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -869,11 +803,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
@@ -901,18 +834,13 @@ 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=2.0, 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("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"),
], ],
@@ -925,9 +853,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, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, gamma: float, delta: 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,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput: 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)
@@ -935,11 +863,10 @@ 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, "epsilon": epsilon, "zeta": zeta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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}
+1 -139
View File
@@ -793,143 +793,5 @@ class DARETIESMode(CalcMode):
return op return op
class EnhancedPowerUpOp(Operation): TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
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,11 +4,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 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, IncrementalSafetensorsWriter from unifiedefficientloader import MemoryEfficientSafeOpen
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,
@@ -39,33 +40,32 @@ 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")
# Stream tensors, filter on the fly, write immediately save_file(pruned_tensors, output_path, metadata)
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,11 +3,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 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, IncrementalSafetensorsWriter from unifiedefficientloader import MemoryEfficientSafeOpen
@@ -42,25 +43,24 @@ 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")
# Stream tensors, rename on the fly, write immediately save_file(renamed_tensors, output_path, metadata)
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.4.1" version = "0.3.5"
license = { file = "LICENSE" } license = { file = "LICENSE" }
classifiers = [ classifiers = [
+1 -1
View File
@@ -1,5 +1,5 @@
requests requests
Pillow Pillow
mutagen mutagen
unifiedefficientloader>=0.5.1 unifiedefficientloader==0.4.4
av av