Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
892dd733d3 |
+2
-2
@@ -22,7 +22,7 @@ from .nodes.merger import (
|
||||
)
|
||||
from .nodes.lora_extract_svd import (
|
||||
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius
|
||||
LoRAExtractKnee, LoRAExtractFrobenius, LoRATextEncoderExtractFixed
|
||||
)
|
||||
from .nodes.lora_resize import (
|
||||
LoRAResizeFixed, LoRAResizeRatio,
|
||||
@@ -60,7 +60,7 @@ class ModelUtilsExtension(ComfyExtension):
|
||||
CheckpointThreeMerger, EmbeddingThreeMerger,
|
||||
# LoRA Extraction
|
||||
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius, LoRATextEncoderExtractFixed,
|
||||
# LoRA Resize
|
||||
LoRAResizeFixed, LoRAResizeRatio,
|
||||
LoRAResizeFrobenius, LoRAResizeCumulative,
|
||||
|
||||
@@ -441,6 +441,7 @@ def extract_lora_from_files(
|
||||
svd_niter: int = 2,
|
||||
lazy_load: bool = True,
|
||||
force_clear_cache: bool = True,
|
||||
is_text_encoder: bool = False,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""
|
||||
Extract LoRA from difference between two models.
|
||||
@@ -460,6 +461,7 @@ def extract_lora_from_files(
|
||||
skip_patterns_str: Patterns for layers to skip
|
||||
mismatch_mode: How to handle mismatches
|
||||
chunk_large_layers: Enable chunked extraction
|
||||
is_text_encoder: Whether extracting from text encoder models
|
||||
"""
|
||||
output_sd = {}
|
||||
|
||||
@@ -487,7 +489,7 @@ def extract_lora_from_files(
|
||||
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
|
||||
|
||||
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):
|
||||
return "skipped", None
|
||||
@@ -647,6 +649,15 @@ def _get_model_inputs():
|
||||
]
|
||||
|
||||
|
||||
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
|
||||
@@ -683,14 +694,18 @@ def _reconstruct_dots(key: str) -> str:
|
||||
return key
|
||||
|
||||
|
||||
def _format_lora_key(key: str) -> str:
|
||||
def _format_lora_key(key: str, is_text_encoder: bool = False) -> str:
|
||||
"""
|
||||
Format the key for LoRA saving.
|
||||
Standardizes to 'diffusion_model.' or 'transformer.' prefix.
|
||||
If is_text_encoder is True, use 'text_encoders.transformer.'
|
||||
"""
|
||||
if key.endswith(".weight"):
|
||||
key = key[:-7]
|
||||
|
||||
if is_text_encoder:
|
||||
return f"text_encoders.transformer.{key}"
|
||||
|
||||
# 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
|
||||
@@ -962,3 +977,43 @@ class LoRAExtractFrobenius(io.ComfyNode):
|
||||
)
|
||||
|
||||
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))
|
||||
|
||||
+11
-47
@@ -11,7 +11,6 @@ from comfy_api.latest import io
|
||||
from safetensors.torch import save_file
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
|
||||
from .merger_utils import parse_custom_weights, get_custom_parameters
|
||||
from .lora_resize import (
|
||||
detect_lora_format,
|
||||
extract_lora_pairs,
|
||||
@@ -32,7 +31,6 @@ def merge_multi_loras(
|
||||
device: str,
|
||||
save_dtype: torch.dtype,
|
||||
output_filename: str,
|
||||
custom_weights: str = "",
|
||||
verbose: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -53,8 +51,6 @@ def merge_multi_loras(
|
||||
|
||||
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
|
||||
|
||||
parsed_custom_weights = parse_custom_weights(custom_weights)
|
||||
|
||||
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
|
||||
|
||||
try:
|
||||
@@ -97,9 +93,6 @@ def merge_multi_loras(
|
||||
# 2. Merge each core layer
|
||||
with torch.no_grad():
|
||||
for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"):
|
||||
# Apply per-layer overrides (regex/substring matching)
|
||||
layer_overrides = get_custom_parameters(core, parsed_custom_weights, {})
|
||||
|
||||
downs = []
|
||||
ups = []
|
||||
is_conv = False
|
||||
@@ -107,14 +100,7 @@ def merge_multi_loras(
|
||||
|
||||
# First pass: load and determine max rank
|
||||
for idx in indices:
|
||||
info = lora_infos[idx].copy()
|
||||
# Apply overrides: w1, w2, w3... or 'weight'
|
||||
lora_key = f'w{idx+1}'
|
||||
if lora_key in layer_overrides:
|
||||
info["weight"] = layer_overrides[lora_key]
|
||||
elif "weight" in layer_overrides:
|
||||
info["weight"] = layer_overrides["weight"]
|
||||
|
||||
info = lora_infos[idx]
|
||||
handler = handlers[idx]
|
||||
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
|
||||
block_keys = info["pairs"][orig_block]
|
||||
@@ -292,7 +278,6 @@ class LoRAMultiMerge(io.ComfyNode):
|
||||
io.Combo.Input("merge_mode", options=["concatenate", "weighted_sum"], default="concatenate",
|
||||
tooltip="concatenate: safe, increases rank. weighted_sum: fixed rank (to max input rank), mathematically lossy but standard in Kohya."),
|
||||
io.String.Input("output_filename", default="merged_lora"),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights for each LoRA. Syntax: 'regex:w1,w2,w3...' or 'regex:w1=0.5,w2=1.0'"),
|
||||
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
||||
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
||||
],
|
||||
@@ -304,7 +289,7 @@ class LoRAMultiMerge(io.ComfyNode):
|
||||
def execute(cls, lora_count,
|
||||
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
|
||||
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
|
||||
merge_mode, output_filename, custom_weights, save_dtype, device) -> io.NodeOutput:
|
||||
merge_mode, output_filename, save_dtype, device) -> io.NodeOutput:
|
||||
|
||||
count = int(lora_count)
|
||||
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
|
||||
@@ -330,8 +315,7 @@ class LoRAMultiMerge(io.ComfyNode):
|
||||
merge_mode=merge_mode,
|
||||
device=device,
|
||||
save_dtype=dtype,
|
||||
output_filename=output_filename,
|
||||
custom_weights=custom_weights
|
||||
output_filename=output_filename
|
||||
)
|
||||
|
||||
return io.NodeOutput(path)
|
||||
@@ -372,7 +356,6 @@ class LoRAMultiMergeDARE(io.ComfyNode):
|
||||
io.Float.Input("trim_quantile", default=0.2, min=0.0, max=1.0, step=0.01, tooltip="TIES trim quantile (drops smallest values)"),
|
||||
io.Int.Input("seed", default=42, min=0, max=0xffffffffffffffff),
|
||||
io.String.Input("output_filename", default="merged_lora_dare"),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights for each LoRA. Syntax: 'regex:w1,w2,w3...' or 'regex:w1=0.5,w2=1.0'"),
|
||||
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
||||
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
||||
],
|
||||
@@ -384,7 +367,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
|
||||
def execute(cls, lora_count,
|
||||
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
|
||||
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
|
||||
drop_rate, trim_quantile, seed, output_filename, custom_weights, save_dtype, device) -> io.NodeOutput:
|
||||
drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput:
|
||||
|
||||
count = int(lora_count)
|
||||
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
|
||||
@@ -412,8 +395,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
|
||||
seed=seed,
|
||||
device=device,
|
||||
save_dtype=dtype,
|
||||
output_filename=output_filename,
|
||||
custom_weights=custom_weights
|
||||
output_filename=output_filename
|
||||
)
|
||||
|
||||
return io.NodeOutput(path)
|
||||
@@ -427,7 +409,6 @@ def merge_multi_loras_dare(
|
||||
device: str,
|
||||
save_dtype: torch.dtype,
|
||||
output_filename: str,
|
||||
custom_weights: str = "",
|
||||
verbose: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -440,9 +421,6 @@ def merge_multi_loras_dare(
|
||||
print(f"[LoRA Multi-Merge DARE] Drop rate: {drop_rate}, Trim quantile: {trim_quantile}")
|
||||
|
||||
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
|
||||
|
||||
parsed_custom_weights = parse_custom_weights(custom_weights)
|
||||
|
||||
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
|
||||
|
||||
rng = torch.Generator(device=device).manual_seed(seed)
|
||||
@@ -482,22 +460,12 @@ def merge_multi_loras_dare(
|
||||
# 2. Merge
|
||||
with torch.no_grad():
|
||||
for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"):
|
||||
# Apply per-layer overrides (regex/substring matching)
|
||||
layer_overrides = get_custom_parameters(core, parsed_custom_weights, {})
|
||||
|
||||
downs = []
|
||||
ups = []
|
||||
max_rank = 0
|
||||
|
||||
for idx in indices:
|
||||
info = lora_infos[idx].copy()
|
||||
# Apply overrides: w1, w2, w3... or 'weight'
|
||||
lora_key = f'w{idx+1}'
|
||||
if lora_key in layer_overrides:
|
||||
info["weight"] = layer_overrides[lora_key]
|
||||
elif "weight" in layer_overrides:
|
||||
info["weight"] = layer_overrides["weight"]
|
||||
|
||||
info = lora_infos[idx]
|
||||
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
|
||||
block_keys = info["pairs"][orig_block]
|
||||
if "down" not in block_keys or "up" not in block_keys: continue
|
||||
@@ -526,22 +494,18 @@ def merge_multi_loras_dare(
|
||||
t = torch.nn.functional.pad(t, tuple(padding))
|
||||
padded.append(t * w)
|
||||
|
||||
# Apply per-layer overrides for DARE/TIES parameters
|
||||
current_drop_rate = layer_overrides.get('alpha', layer_overrides.get('drop_rate', drop_rate))
|
||||
current_trim_quantile = layer_overrides.get('beta', layer_overrides.get('trim_quantile', trim_quantile))
|
||||
|
||||
# DARE
|
||||
if current_drop_rate > 0:
|
||||
if drop_rate > 0:
|
||||
for i in range(len(padded)):
|
||||
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > current_drop_rate).float()
|
||||
padded[i] = (padded[i] * mask) / (1 - current_drop_rate)
|
||||
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float()
|
||||
padded[i] = (padded[i] * mask) / (1 - drop_rate)
|
||||
|
||||
# TIES
|
||||
# 1. Trim
|
||||
if current_trim_quantile > 0:
|
||||
if trim_quantile > 0:
|
||||
for i in range(len(padded)):
|
||||
flat = padded[i].abs().flatten()
|
||||
k = int(len(flat) * current_trim_quantile)
|
||||
k = int(len(flat) * trim_quantile)
|
||||
if k > 0:
|
||||
threshold = torch.kthvalue(flat, k).values
|
||||
padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i])
|
||||
|
||||
+6
-42
@@ -7,7 +7,6 @@ from tqdm import tqdm
|
||||
from comfy_api.latest import io
|
||||
from safetensors.torch import save_file
|
||||
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
|
||||
from .merger_utils import parse_custom_weights, get_custom_parameters
|
||||
from .device_utils import (
|
||||
estimate_model_size, prepare_for_large_operation, cleanup_after_operation
|
||||
)
|
||||
@@ -104,9 +103,6 @@ class MergerLogic:
|
||||
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''))
|
||||
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''))
|
||||
|
||||
# Parse custom per-layer weights
|
||||
parsed_custom_weights = parse_custom_weights(recipe_params.get('custom_weights', ''))
|
||||
|
||||
# Pre-compute key differences for logging
|
||||
primary_keys = set(all_keys)
|
||||
for name, handler in handlers.items():
|
||||
@@ -181,7 +177,7 @@ class MergerLogic:
|
||||
# Check exclude patterns - use Model A only, no merge
|
||||
if _matches_any_pattern(key, exclude_patterns):
|
||||
t = tensor_a.detach().to(save_torch_dtype).cpu()
|
||||
|
||||
|
||||
batch_buffer[key] = t
|
||||
batch_bytes += t.numel() * t.element_size()
|
||||
batch_key_count += 1
|
||||
@@ -197,13 +193,11 @@ class MergerLogic:
|
||||
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
|
||||
|
||||
try:
|
||||
# Apply per-layer parameter overrides (regex/substring matching)
|
||||
layer_params = get_custom_parameters(key, parsed_custom_weights, recipe_params)
|
||||
recipe = calc_mode_class.create_recipe(key=key, **layer_params)
|
||||
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
|
||||
result = recipe.merge()
|
||||
except MissingTensorError as e:
|
||||
if mismatch_mode == MissingTensorBehavior.ERROR:
|
||||
|
||||
|
||||
_flush_batch()
|
||||
writer.__exit__(None, None, None)
|
||||
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
|
||||
@@ -218,7 +212,7 @@ class MergerLogic:
|
||||
if isinstance(result, dict):
|
||||
for r_key, r_tensor in result.items():
|
||||
t = r_tensor.detach().to(save_torch_dtype).cpu()
|
||||
|
||||
|
||||
batch_buffer[r_key] = t
|
||||
batch_bytes += t.numel() * t.element_size()
|
||||
else:
|
||||
@@ -232,11 +226,11 @@ class MergerLogic:
|
||||
result = result[slices]
|
||||
|
||||
t = result.detach().to(save_torch_dtype).cpu()
|
||||
|
||||
|
||||
batch_buffer[key] = t
|
||||
batch_bytes += t.numel() * t.element_size()
|
||||
|
||||
|
||||
|
||||
batch_key_count += 1
|
||||
|
||||
# Clean up references to allow GC immediately.
|
||||
@@ -309,7 +303,6 @@ class CheckpointTwoMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -324,7 +317,6 @@ class CheckpointTwoMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_2_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -337,7 +329,6 @@ class CheckpointTwoMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b}
|
||||
@@ -371,7 +362,6 @@ class ModelTwoMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -386,7 +376,6 @@ class ModelTwoMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_2_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -399,7 +388,6 @@ class ModelTwoMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b}
|
||||
@@ -433,7 +421,6 @@ class TextEncoderTwoMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -448,7 +435,6 @@ class TextEncoderTwoMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_2_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -461,7 +447,6 @@ class TextEncoderTwoMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b}
|
||||
@@ -495,7 +480,6 @@ class LoRATwoMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -510,7 +494,6 @@ class LoRATwoMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_2_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -523,7 +506,6 @@ class LoRATwoMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b}
|
||||
@@ -557,7 +539,6 @@ class EmbeddingTwoMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -572,7 +553,6 @@ class EmbeddingTwoMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_2_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -585,7 +565,6 @@ class EmbeddingTwoMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b}
|
||||
@@ -622,7 +601,6 @@ class CheckpointThreeMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -637,7 +615,6 @@ class CheckpointThreeMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_3_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -650,7 +627,6 @@ class CheckpointThreeMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
|
||||
@@ -685,7 +661,6 @@ class ModelThreeMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -700,7 +675,6 @@ class ModelThreeMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_3_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -713,7 +687,6 @@ class ModelThreeMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
|
||||
@@ -748,7 +721,6 @@ class TextEncoderThreeMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -763,7 +735,6 @@ class TextEncoderThreeMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_3_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -776,7 +747,6 @@ class TextEncoderThreeMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
|
||||
@@ -811,7 +781,6 @@ class LoRAThreeMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -826,7 +795,6 @@ class LoRAThreeMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_3_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -839,7 +807,6 @@ class LoRAThreeMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
|
||||
@@ -874,7 +841,6 @@ class EmbeddingThreeMerger(io.ComfyNode):
|
||||
io.Combo.Input("process_device", options=["cuda", "cpu"]),
|
||||
io.String.Input("exclude_patterns", default="", multiline=True),
|
||||
io.String.Input("discard_patterns", default="", multiline=True),
|
||||
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
|
||||
io.Boolean.Input("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"),
|
||||
],
|
||||
@@ -889,7 +855,6 @@ class EmbeddingThreeMerger(io.ComfyNode):
|
||||
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
|
||||
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
|
||||
process_device: str, exclude_patterns: str, discard_patterns: str,
|
||||
custom_weights: str,
|
||||
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
|
||||
doc = load_documentation_from_file('merger_3_model_modes.md')
|
||||
if execution_mode == "DOCUMENTATION ONLY":
|
||||
@@ -902,7 +867,6 @@ class EmbeddingThreeMerger(io.ComfyNode):
|
||||
"output_filename": output_filename, "save_dtype": save_dtype,
|
||||
"device": process_device, "dtype": torch.float32,
|
||||
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
|
||||
"custom_weights": custom_weights,
|
||||
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
|
||||
}
|
||||
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
|
||||
|
||||
+1
-83
@@ -1,8 +1,7 @@
|
||||
"""
|
||||
Utilities for model merging operations.
|
||||
"""
|
||||
from typing import Dict, List, Tuple, Any
|
||||
import re
|
||||
from typing import Dict, List
|
||||
|
||||
def get_union_keys(*handlers) -> List[str]:
|
||||
"""Returns the union of all tensor keys across multiple model handlers."""
|
||||
@@ -22,84 +21,3 @@ def get_intersection_keys(*handlers) -> List[str]:
|
||||
for handler in handlers[1:]:
|
||||
all_keys &= set(handler.keys())
|
||||
return sorted(all_keys)
|
||||
|
||||
|
||||
def parse_custom_weights(custom_weights_str: str) -> List[Tuple[re.Pattern, Dict[str, float]]]:
|
||||
"""Parses multiline string of custom weights.
|
||||
Format: pattern:weight OR pattern:param=value,param2=value
|
||||
Returns list of (compiled_regex, parameter_dict).
|
||||
"""
|
||||
if not custom_weights_str or not custom_weights_str.strip():
|
||||
return []
|
||||
|
||||
results = []
|
||||
lines = custom_weights_str.strip().split('\n')
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line or line.startswith('#'):
|
||||
continue
|
||||
|
||||
if ':' not in line:
|
||||
continue
|
||||
|
||||
pattern_str, weight_str = line.rsplit(':', 1)
|
||||
pattern_str = pattern_str.strip()
|
||||
weight_str = weight_str.strip()
|
||||
|
||||
try:
|
||||
pattern = re.compile(pattern_str)
|
||||
except re.error:
|
||||
# Fallback to literal substring if regex invalid
|
||||
pattern = re.compile(re.escape(pattern_str))
|
||||
|
||||
params = {}
|
||||
if '=' in weight_str:
|
||||
# param=value,param2=value
|
||||
parts = weight_str.split(',')
|
||||
for part in parts:
|
||||
if '=' in part:
|
||||
k, v = part.split('=', 1)
|
||||
k = k.strip()
|
||||
# Map common aliases
|
||||
if k in ['a', 'alpha']: k = 'alpha'
|
||||
elif k in ['b', 'beta']: k = 'beta'
|
||||
elif k in ['g', 'gamma']: k = 'gamma'
|
||||
elif k in ['d', 'delta']: k = 'delta'
|
||||
elif k in ['w', 'weight']: k = 'weight'
|
||||
elif re.match(r'^w[1-8]$', k): pass # Keep w1-w8 as is
|
||||
|
||||
try:
|
||||
params[k] = float(v.strip())
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
# single float or comma-separated list of floats (for multi-LoRA)
|
||||
if ',' in weight_str:
|
||||
parts = weight_str.split(',')
|
||||
for i, p in enumerate(parts):
|
||||
try:
|
||||
params[f'w{i+1}'] = float(p.strip())
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
val = float(weight_str)
|
||||
params['alpha'] = val # default for merger
|
||||
params['weight'] = val # default for lora_merger
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if params:
|
||||
results.append((pattern, params))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def get_custom_parameters(key: str, parsed_weights: List[Tuple[re.Pattern, Dict[str, float]]], base_params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Returns updated parameters if key matches any pattern."""
|
||||
new_params = base_params.copy()
|
||||
for pattern, custom_dict in parsed_weights:
|
||||
if pattern.search(key):
|
||||
new_params.update(custom_dict)
|
||||
break # First match wins
|
||||
return new_params
|
||||
|
||||
Reference in New Issue
Block a user