diff --git a/__init__.py b/__init__.py index 8896830..c004ead 100644 --- a/__init__.py +++ b/__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 diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 2d74c60..980fddd 100644 --- a/nodes/lora_merger.py +++ b/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() + diff --git a/nodes/merger.py b/nodes/merger.py index fd90a77..69dd80c 100644 --- a/nodes/merger.py +++ b/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, diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py index 71504d6..816a43b 100644 --- a/nodes/merger_ops.py +++ b/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()] \ No newline at end of file