From a7840a424ed0c02a3a10d1510186fc6da5cf3996 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 3 May 2026 18:12:21 +0200 Subject: [PATCH] Add ability to adjust paramaters on a ler layer basis using regex/substring for layer colon separated with parameter=float for specific parameters --- nodes/lora_merger.py | 58 ++++++++++++++++++++++++------ nodes/merger.py | 48 +++++++++++++++++++++---- nodes/merger_utils.py | 84 ++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 172 insertions(+), 18 deletions(-) diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 2d74c60..ca72196 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -11,6 +11,7 @@ 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, @@ -31,6 +32,7 @@ def merge_multi_loras( device: str, save_dtype: torch.dtype, output_filename: str, + custom_weights: str = "", verbose: bool = True, ) -> str: """ @@ -51,6 +53,8 @@ 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: @@ -93,6 +97,9 @@ 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 @@ -100,7 +107,14 @@ def merge_multi_loras( # First pass: load and determine max rank for idx in indices: - info = lora_infos[idx] + 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"] + 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] @@ -278,6 +292,7 @@ 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"), ], @@ -289,7 +304,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, save_dtype, device) -> io.NodeOutput: + merge_mode, output_filename, custom_weights, 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] @@ -315,7 +330,8 @@ class LoRAMultiMerge(io.ComfyNode): merge_mode=merge_mode, device=device, save_dtype=dtype, - output_filename=output_filename + output_filename=output_filename, + custom_weights=custom_weights ) return io.NodeOutput(path) @@ -356,6 +372,7 @@ 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"), ], @@ -367,7 +384,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, save_dtype, device) -> io.NodeOutput: + drop_rate, trim_quantile, seed, output_filename, custom_weights, 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] @@ -395,7 +412,8 @@ class LoRAMultiMergeDARE(io.ComfyNode): seed=seed, device=device, save_dtype=dtype, - output_filename=output_filename + output_filename=output_filename, + custom_weights=custom_weights ) return io.NodeOutput(path) @@ -409,6 +427,7 @@ def merge_multi_loras_dare( device: str, save_dtype: torch.dtype, output_filename: str, + custom_weights: str = "", verbose: bool = True, ) -> str: """ @@ -421,6 +440,9 @@ 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) @@ -460,12 +482,22 @@ 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] + 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"] + 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 @@ -494,18 +526,22 @@ 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 drop_rate > 0: + if current_drop_rate > 0: for i in range(len(padded)): - mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float() - padded[i] = (padded[i] * mask) / (1 - drop_rate) + mask = (torch.rand(padded[i].shape, generator=rng, device=device) > current_drop_rate).float() + padded[i] = (padded[i] * mask) / (1 - current_drop_rate) # TIES # 1. Trim - if trim_quantile > 0: + if current_trim_quantile > 0: for i in range(len(padded)): flat = padded[i].abs().flatten() - k = int(len(flat) * trim_quantile) + k = int(len(flat) * current_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]) diff --git a/nodes/merger.py b/nodes/merger.py index fd90a77..1e1acb9 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -7,6 +7,7 @@ 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 ) @@ -103,6 +104,9 @@ 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(): @@ -177,7 +181,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 @@ -193,11 +197,13 @@ class MergerLogic: recipe_params['_tensor_a_dtype'] = tensor_a.dtype try: - recipe = calc_mode_class.create_recipe(key=key, **recipe_params) + # 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) 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}") @@ -212,7 +218,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: @@ -226,11 +232,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. @@ -303,6 +309,7 @@ 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"), ], @@ -317,6 +324,7 @@ 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": @@ -329,6 +337,7 @@ 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} @@ -362,6 +371,7 @@ 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"), ], @@ -376,6 +386,7 @@ 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": @@ -388,6 +399,7 @@ 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} @@ -421,6 +433,7 @@ 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"), ], @@ -435,6 +448,7 @@ 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": @@ -447,6 +461,7 @@ 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} @@ -480,6 +495,7 @@ 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"), ], @@ -494,6 +510,7 @@ 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": @@ -506,6 +523,7 @@ 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} @@ -539,6 +557,7 @@ 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"), ], @@ -553,6 +572,7 @@ 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": @@ -565,6 +585,7 @@ 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} @@ -601,6 +622,7 @@ 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"), ], @@ -615,6 +637,7 @@ 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": @@ -627,6 +650,7 @@ 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} @@ -661,6 +685,7 @@ 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"), ], @@ -675,6 +700,7 @@ 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": @@ -687,6 +713,7 @@ 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} @@ -721,6 +748,7 @@ 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"), ], @@ -735,6 +763,7 @@ 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": @@ -747,6 +776,7 @@ 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} @@ -781,6 +811,7 @@ 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"), ], @@ -795,6 +826,7 @@ 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": @@ -807,6 +839,7 @@ 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} @@ -841,6 +874,7 @@ 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"), ], @@ -855,6 +889,7 @@ 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": @@ -867,6 +902,7 @@ 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} diff --git a/nodes/merger_utils.py b/nodes/merger_utils.py index 4725508..652bc1e 100644 --- a/nodes/merger_utils.py +++ b/nodes/merger_utils.py @@ -1,7 +1,8 @@ """ Utilities for model merging operations. """ -from typing import Dict, List +from typing import Dict, List, Tuple, Any +import re def get_union_keys(*handlers) -> List[str]: """Returns the union of all tensor keys across multiple model handlers.""" @@ -21,3 +22,84 @@ 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