Add new 'enhanced' PowerUp/DARE-Ties merge modes

__init__.py

@@ -1,7 +1,7 @@
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, io
 
-from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE
+from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced
 from .nodes.metakeys import (
     ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
     CheckpointMetaKeys, EmbeddingMetaKeys
@@ -65,7 +65,7 @@ class ModelUtilsExtension(ComfyExtension):
             LoRAResizeFixed, LoRAResizeRatio,
             LoRAResizeFrobenius, LoRAResizeCumulative,
             # LoRA Multi-Merge
-            LoRAMultiMerge, LoRAMultiMergeDARE,
+            LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
             # LoRA Merge To Model
             LoRAMergeToModel,
             # Downloaders

nodes/lora_merger.py

@@ -566,3 +566,258 @@ def merge_multi_loras_dare(
     finally:
         for h in handlers: h.__exit__(None, None, None)
         cleanup_after_operation()
+
+
+class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
+    """Merge multiple LoRAs into a single LoRA file using Enhanced DARE-Ties method."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LoRAMultiMergeDAREEnhanced",
+            display_name="LoRA Multi-Merge (Enhanced DARE-Ties)",
+            category="ModelUtils/LoRA/Merge",
+            description="Merge 1-8 LoRAs into a single LoRA file using Enhanced DARE-Ties. Uses dynamic probability masking based on value magnitudes.",
+            inputs=[
+                io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2"),
+                # LoRA 1
+                io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"),
+                io.Float.Input("weight_1", default=1.0, min=-10.0, max=10.0, step=0.01),
+                # ...
+                io.Combo.Input("lora_2", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_2", default=1.0, min=-10.0, max=10.0, step=0.01),
+                io.Combo.Input("lora_3", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_3", default=1.0, min=-10.0, max=10.0, step=0.01),
+                io.Combo.Input("lora_4", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01),
+                io.Combo.Input("lora_5", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_5", default=1.0, min=-10.0, max=10.0, step=0.01),
+                io.Combo.Input("lora_6", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_6", default=1.0, min=-10.0, max=10.0, step=0.01),
+                io.Combo.Input("lora_7", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_7", default=1.0, min=-10.0, max=10.0, step=0.01),
+                io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
+                io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01),
+
+                io.Float.Input("mask_power", default=2.0, min=0.001, max=10.0, step=0.01, tooltip="Mask power (Curve. 2.0 = quadratic)"),
+                io.Float.Input("min_keep_prob", default=0.01, min=0.0, max=1.0, step=0.01, tooltip="Minimum keep probability (Floor to prevent explosion)"),
+                io.Float.Input("mask_smooth", default=0.0, min=0.0, max=1.0, step=0.01, tooltip="Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale)"),
+                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_enhanced"),
+                io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
+                io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
+            ],
+            outputs=[io.String.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    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,
+                mask_power, min_keep_prob, mask_smooth, 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]
+        weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8]
+
+        valid_paths = []
+        valid_weights = []
+
+        for i in range(count):
+            if names[i] and names[i] != "None":
+                path = folder_paths.get_full_path_or_raise("loras", names[i])
+                valid_paths.append(path)
+                valid_weights.append(weights[i])
+
+        if not valid_paths:
+            raise ValueError("No LoRAs selected for merging.")
+
+        dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
+
+        path = merge_multi_loras_dare_enhanced(
+            lora_paths=valid_paths,
+            lora_weights=valid_weights,
+            mask_power=mask_power,
+            min_keep_prob=min_keep_prob,
+            mask_smooth=mask_smooth,
+            trim_quantile=trim_quantile,
+            seed=seed,
+            device=device,
+            save_dtype=dtype,
+            output_filename=output_filename
+        )
+
+        return io.NodeOutput(path)
+
+
+def merge_multi_loras_dare_enhanced(
+    lora_paths: List[str],
+    lora_weights: List[float],
+    mask_power: float,
+    min_keep_prob: float,
+    mask_smooth: float,
+    trim_quantile: float,
+    seed: int,
+    device: str,
+    save_dtype: torch.dtype,
+    output_filename: str,
+    verbose: bool = True,
+) -> str:
+    """
+    Merge multiple LoRAs using Enhanced DARE-Ties method applied to weights.
+    Different ranks are handled via zero-padding to the max rank.
+    """
+    total_size_gb = sum(estimate_model_size(p) for p in lora_paths)
+    if verbose:
+        print(f"[LoRA Multi-Merge Enhanced DARE] Preparing memory for {total_size_gb:.2f}GB operation...")
+
+    prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
+    handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
+
+    rng = torch.Generator(device=device).manual_seed(seed)
+
+    try:
+        # 1. Analyze
+        layer_map = {}
+        lora_infos = []
+        for i, handler in enumerate(handlers):
+            keys = handler.keys()
+            fmt = detect_lora_format(keys)
+            pairs = extract_lora_pairs(keys, fmt)
+            rank, alpha = detect_lora_rank(handler, pairs)
+            info = {
+                "name": os.path.basename(lora_paths[i]),
+                "pairs": pairs,
+                "rank": rank,
+                "alpha": alpha,
+                "weight": lora_weights[i]
+            }
+            lora_infos.append(info)
+
+            if verbose:
+                print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}")
+
+            for block_name in pairs.keys():
+                core = _format_lora_key(block_name)
+                if core not in layer_map: layer_map[core] = []
+                layer_map[core].append(i)
+
+        if verbose:
+            print(f"[LoRA Multi-Merge Enhanced DARE] Total unique layers after resolving naming: {len(layer_map)}")
+
+        output_sd = {}
+        pbar = comfy.utils.ProgressBar(len(layer_map))
+
+        # 2. Merge
+        with torch.no_grad():
+            for core, indices in tqdm(layer_map.items(), desc="Merging layers (Enhanced DARE)", unit="layers"):
+                downs = []
+                ups = []
+                max_rank = 0
+
+                for idx in indices:
+                    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
+
+                    t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
+                    t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
+
+                    max_rank = max(max_rank, t_d.shape[0])
+                    downs.append((t_d, info["weight"]))
+                    ups.append((t_u, info["weight"]))
+
+                if not downs:
+                    pbar.update(1)
+                    continue
+
+                def process_ties_dare_enhanced(tensors, weights, dim_to_pad):
+                    padded = []
+                    for t, w in tensors:
+                        if t.shape[dim_to_pad] < max_rank:
+                            padding = [0] * (len(t.shape) * 2)
+                            rev_dim = len(t.shape) - 1 - dim_to_pad
+                            padding[rev_dim*2 + 1] = max_rank - t.shape[dim_to_pad]
+                            t = torch.nn.functional.pad(t, tuple(padding))
+                        padded.append(t * w)
+
+                    # Enhanced DARE
+                    for i in range(len(padded)):
+                        t_val = padded[i]
+                        abs_t = torch.abs(t_val)
+                        max_val = torch.max(abs_t)
+                        if max_val > 0:
+                            prob = torch.clamp((abs_t / max_val) ** max(mask_power, 0.001), min=min_keep_prob, max=1.0)
+                            prob = torch.nan_to_num(prob)
+
+                            random_mask = torch.bernoulli(prob, generator=rng)
+                            interpolated_mask = torch.lerp(random_mask, prob, mask_smooth)
+
+                            padded[i] = (t_val * interpolated_mask) / prob
+
+                    # TIES
+                    if trim_quantile > 0:
+                        for i in range(len(padded)):
+                            flat = padded[i].abs().flatten()
+                            k = max(1, 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])
+
+                    stacked = torch.stack(padded)
+                    signs = torch.sign(stacked)
+                    sum_signs = signs.sum(dim=0)
+                    dominant_sign = torch.sign(sum_signs)
+
+                    mask = (signs == dominant_sign) & (dominant_sign != 0)
+                    filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
+
+                    count = mask.sum(dim=0)
+                    result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
+                    return result
+
+                merged_down = process_ties_dare_enhanced(downs, [1.0]*len(downs), 0)
+                merged_up = process_ties_dare_enhanced(ups, [1.0]*len(ups), 1)
+
+                output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
+                output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
+                output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype)
+
+                del merged_down, merged_up
+                for d, _ in downs: del d
+                for u, _ in ups: del u
+                downs.clear()
+                ups.clear()
+                import gc
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+
+                pbar.update(1)
+
+        # Final Summary
+        if verbose:
+            matched_stats = {i: 0 for i in range(len(lora_paths))}
+            for indices in layer_map.values():
+                for idx in indices:
+                    matched_stats[idx] += 1
+
+            print(f"[LoRA Multi-Merge Enhanced DARE] --- Merge Summary ---")
+            for i, info in enumerate(lora_infos):
+                print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
+            print(f"[LoRA Multi-Merge Enhanced DARE] Output state dict has {len(output_sd)} tensors")
+
+        # Save
+        output_dir = folder_paths.get_folder_paths("loras")[0]
+        os.makedirs(output_dir, exist_ok=True)
+        output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
+        save_file(output_sd, output_path, {"ss_training_comment": "Merged via Enhanced DARE-Ties"})
+        return output_path
+
+    finally:
+        for h in handlers: h.__exit__(None, None, None)
+        cleanup_after_operation()
+

nodes/merger.py

@@ -177,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,7 +197,7 @@ class MergerLogic:
                     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 +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:
@@ -226,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.
@@ -296,7 +296,9 @@ class CheckpointTwoMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_2_checkpoint"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -315,7 +317,7 @@ class CheckpointTwoMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_2_model_modes.md')
@@ -325,7 +327,7 @@ class CheckpointTwoMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -355,7 +357,9 @@ class ModelTwoMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_2_model"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -374,7 +378,7 @@ class ModelTwoMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_2_model_modes.md')
@@ -384,7 +388,7 @@ class ModelTwoMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -414,7 +418,9 @@ class TextEncoderTwoMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_2_textencoder"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -433,7 +439,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_2_model_modes.md')
@@ -443,7 +449,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -473,7 +479,9 @@ class LoRATwoMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_2_lora"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -492,7 +500,7 @@ class LoRATwoMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_2_model_modes.md')
@@ -502,7 +510,7 @@ class LoRATwoMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -532,7 +540,9 @@ class EmbeddingTwoMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_2_embedding"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -551,7 +561,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_2_model_modes.md')
@@ -561,7 +571,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -594,7 +604,9 @@ class CheckpointThreeMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_3_checkpoint"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -613,7 +625,7 @@ class CheckpointThreeMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_3_model_modes.md')
@@ -623,7 +635,7 @@ class CheckpointThreeMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -654,7 +666,9 @@ class ModelThreeMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_3_model"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -673,7 +687,7 @@ class ModelThreeMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_3_model_modes.md')
@@ -683,7 +697,7 @@ class ModelThreeMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -714,7 +728,9 @@ class TextEncoderThreeMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_3_textencoder"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -733,7 +749,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_3_model_modes.md')
@@ -743,7 +759,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -774,7 +790,9 @@ class LoRAThreeMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_3_lora"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -793,7 +811,7 @@ class LoRAThreeMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_3_model_modes.md')
@@ -803,7 +821,7 @@ class LoRAThreeMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
@@ -834,7 +852,9 @@ class EmbeddingThreeMerger(io.ComfyNode):
                 io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
                 io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
-                io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
+                io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
                 io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                 io.String.Input("output_filename", default="merged_3_embedding"),
                 io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
@@ -853,7 +873,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
     @classmethod
     def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
                 calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
-                gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
+                gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
                 process_device: str, exclude_patterns: str, discard_patterns: str,
                 lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
         doc = load_documentation_from_file('merger_3_model_modes.md')
@@ -863,7 +883,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
         recipe_params = {
             "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
             "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
-            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
+            "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
             "output_filename": output_filename, "save_dtype": save_dtype,
             "device": process_device, "dtype": torch.float32,
             "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,

nodes/merger_ops.py

@@ -793,5 +793,143 @@ class DARETIESMode(CalcMode):
         return op
 
 
-TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
+class EnhancedPowerUpOp(Operation):
+    def __init__(self, key, alpha, beta, gamma, delta, seed, *sources):
+        super().__init__(key, *sources)
+        self.mask_power = alpha       # Mask power
+        self.min_keep_prob = beta     # Min keep prob
+        self.mask_smooth = gamma      # Smoothness (0 or 1)
+        self.delta = delta            # Final multiplier
+        self.seed = seed
+
+    def oper(self, a, b):
+        if a is None or b is None: return None
+        a, b = resize_tensors(a, b, mode=self.alignment_mode)
+        delta_val = (b - a).float()
+        abs_delta = torch.abs(delta_val)
+        max_delta = torch.max(abs_delta)
+        if max_delta == 0: return torch.zeros_like(delta_val)
+
+        prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
+
+        rng = torch.Generator(device=a.device).manual_seed(self.seed)
+        random_mask = torch.bernoulli(prob, generator=rng)
+        interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
+
+        rescaled_delta = (delta_val * interpolated_mask) / prob
+        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) / prob
+
+        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()]