From b1ec55ff423a0aa737ae410dd2d4bbc4c98d1396 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 23 Jan 2026 18:41:35 +0100 Subject: [PATCH] improve and correct merger operations based on papers --- docs/merger_2_model_modes.md | 41 ++++++++- docs/merger_3_model_modes.md | 8 +- nodes/merger.py | 25 +++--- nodes/merger_ops.py | 161 +++++++++++++++++++++++++++++++++-- 4 files changed, 213 insertions(+), 22 deletions(-) diff --git a/docs/merger_2_model_modes.md b/docs/merger_2_model_modes.md index d1100df..3863f94 100644 --- a/docs/merger_2_model_modes.md +++ b/docs/merger_2_model_modes.md @@ -23,12 +23,47 @@ --- ## Power-Up (DARE) -> Adds the unique capabilities of Model B to Model A using the Drop and Rescale (DARE) technique, which often preserves the knowledge of the base model better than simple additions. +> Adds the unique capabilities of Model B to Model A using the Drop and Rescale (DARE) technique. This implementation handles shape mismatches between models by padding and uses a randomized dropout mask. **Models Used:** A, B **Parameters:** -- **Alpha:** The dropout rate. This is the proportion of unique weights from Model B that are *dropped* before merging. A higher value means less of B is merged. +- **Alpha:** The dropout rate ($p$). This is the proportion of delta parameters from Model B that are randomly set to zero. - **Beta:** A final multiplier for the rescaled difference before it's added to Model A. +- **Rescaling Logic:** Remaining weights are automatically rescaled by $1/(1-p)$ as per the DARE paper to approximate the original embeddings. + +--- + +## Enhanced Man Interp +> Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control. + +**Models Used:** A, B +**Parameters:** +- **Alpha:** Interpolation strength. +- **Beta:** Lower mean threshold for filtering differences. +- **Gamma:** Upper mean threshold for filtering differences. +- **Delta:** Smoothness factor (mix between randomized mask and powered differences). + +--- + +## Enhanced Auto Interp +> Automated version of the enhanced interpolation mode that dynamically calculates thresholds based on mean differences. + +**Models Used:** A, B +**Parameters:** +- **Alpha:** Interpolation strength. +- **Beta:** Threshold adjustment factor. +- **Gamma:** Smoothness factor. + +--- + +## Weight-Sum Cutoff +> A linear interpolation mode that only applies the merge to weights whose differences fall within a specific threshold range. + +**Models Used:** A, B +**Parameters:** +- **Alpha:** Interpolation weight (multiplier for the difference). +- **Beta:** Upper threshold for the difference cutoff. +- **Gamma:** Lower threshold for the difference cutoff. --- @@ -84,4 +119,4 @@ Layers matching any pattern will be **removed entirely** from the output. | `text_model` | All text encoder layers | | `\.norm` | All normalization layers | | `attn\.(q\|k\|v)` | Query, key, value attention weights | -| `block\.[0-9]\.` | Blocks 0-9 | \ No newline at end of file +| `block\.[0-9]\.` | Blocks 0-9 | diff --git a/docs/merger_3_model_modes.md b/docs/merger_3_model_modes.md index ba5020d..5390066 100644 --- a/docs/merger_3_model_modes.md +++ b/docs/merger_3_model_modes.md @@ -22,19 +22,19 @@ --- ## Extract-Features -> A powerful mode that identifies features present in both `(B - A)` and `(C - A)` and adds them to A. Allows for fine-grained control over combining aspects based on their similarity. +> A powerful mode that identifies features present in both `(B - A)` and `(C - A)` and adds them to A. It uses per-vector cosine similarity to decide how much of each feature to keep, allowing for fine-grained control over combining aspects. **Models Used:** A, B, C **Parameters:** - **Alpha:** Weights the merge between Model B (`0.0`) and Model C (`1.0`). - **Beta:** Controls the focus on similarity (`0.0`) versus dissimilarity (`1.0`). -- **Gamma:** A bias exponent for similarity. Higher values increase the bias. +- **Gamma:** A bias exponent for similarity calculation. - **Delta:** A final multiplier for the extracted features before they are added to Model A. --- ## Add-Dissimilarities -> Identifies features that are dissimilar between Model B and Model C and adds them to Model A. Useful for combining unique aspects of two different models. +> Identifies features that are dissimilar between Model B and Model C (relative to A) and adds them to Model A. Useful for combining unique aspects of two different models. **Models Used:** A, B, C **Parameters:** @@ -80,4 +80,4 @@ Layers matching any pattern will be **removed entirely** from the output. **Pattern format:** - Whitespace-separated regex patterns - Patterns use **substring matching** (not full match) -- Example: `text_model lora` matches any key containing "text_model" OR "lora" \ No newline at end of file +- Example: `text_model lora` matches any key containing "text_model" OR "lora" diff --git a/nodes/merger.py b/nodes/merger.py index da4d625..fd458ab 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -222,6 +222,7 @@ class CheckpointTwoMerger(io.ComfyNode): io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.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"]), @@ -238,7 +239,7 @@ class CheckpointTwoMerger(io.ComfyNode): @classmethod def execute(cls, execution_mode: str, model_a: str, model_b: str, calc_mode: str, mismatch_mode: str, alpha: float, beta: float, - gamma: float, seed: int, output_filename: str, save_dtype: str, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": @@ -247,7 +248,7 @@ class CheckpointTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "mismatch_mode": mismatch_mode, - "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, @@ -275,6 +276,7 @@ class ModelTwoMerger(io.ComfyNode): io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.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"]), @@ -291,7 +293,7 @@ class ModelTwoMerger(io.ComfyNode): @classmethod def execute(cls, execution_mode: str, model_a: str, model_b: str, calc_mode: str, mismatch_mode: str, alpha: float, beta: float, - gamma: float, seed: int, output_filename: str, save_dtype: str, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": @@ -300,7 +302,7 @@ class ModelTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "mismatch_mode": mismatch_mode, - "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, @@ -328,6 +330,7 @@ class TextEncoderTwoMerger(io.ComfyNode): io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.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"]), @@ -344,7 +347,7 @@ class TextEncoderTwoMerger(io.ComfyNode): @classmethod def execute(cls, execution_mode: str, model_a: str, model_b: str, calc_mode: str, mismatch_mode: str, alpha: float, beta: float, - gamma: float, seed: int, output_filename: str, save_dtype: str, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": @@ -353,7 +356,7 @@ class TextEncoderTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "mismatch_mode": mismatch_mode, - "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, @@ -381,6 +384,7 @@ class LoRATwoMerger(io.ComfyNode): io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.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"]), @@ -397,7 +401,7 @@ class LoRATwoMerger(io.ComfyNode): @classmethod def execute(cls, execution_mode: str, model_a: str, model_b: str, calc_mode: str, mismatch_mode: str, alpha: float, beta: float, - gamma: float, seed: int, output_filename: str, save_dtype: str, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": @@ -406,7 +410,7 @@ class LoRATwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "mismatch_mode": mismatch_mode, - "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, @@ -434,6 +438,7 @@ class EmbeddingTwoMerger(io.ComfyNode): io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.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"]), @@ -450,7 +455,7 @@ class EmbeddingTwoMerger(io.ComfyNode): @classmethod def execute(cls, execution_mode: str, model_a: str, model_b: str, calc_mode: str, mismatch_mode: str, alpha: float, beta: float, - gamma: float, seed: int, output_filename: str, save_dtype: str, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": @@ -459,7 +464,7 @@ class EmbeddingTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "mismatch_mode": mismatch_mode, - "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "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 02b5f79..76c59b0 100644 --- a/nodes/merger_ops.py +++ b/nodes/merger_ops.py @@ -15,6 +15,33 @@ class MissingTensorError(Exception): """Raised when a tensor is missing and behavior is ERROR.""" pass +# ################# # +# UTILITIES +# ################# # +def resize_tensors(tensor1, tensor2): + """Pads 1D or 2D tensors to match shapes.""" + if len(tensor1.shape) not in [1, 2]: + return tensor1, tensor2 + + # Pad along the last dimension (width) + if tensor1.shape[-1] < tensor2.shape[-1]: + padding_size = tensor2.shape[-1] - tensor1.shape[-1] + tensor1 = F.pad(tensor1, (0, padding_size, 0, 0)) + elif tensor2.shape[-1] < tensor1.shape[-1]: + padding_size = tensor1.shape[-1] - tensor2.shape[-1] + tensor2 = F.pad(tensor2, (0, padding_size, 0, 0)) + + # Pad along the first dimension (height) + if tensor1.shape[0] < tensor2.shape[0]: + padding_size = tensor2.shape[0] - tensor1.shape[0] + tensor1 = F.pad(tensor1, (0, 0, 0, padding_size)) + elif tensor2.shape[0] < tensor1.shape[0]: + padding_size = tensor1.shape[0] - tensor2.shape[0] + tensor2 = F.pad(tensor2, (0, 0, 0, padding_size)) + + return tensor1, tensor2 + + # ################# # # MERGE OPERATORS # ################# # @@ -153,11 +180,12 @@ class ExtractOp(Operation): def oper(self, base: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: if a is None or b is None: return None + dtype = base.dtype if base is not None else a.dtype base_f = base.float() if base is not None else 0 a_f, b_f = a.float() - base_f, b.float() - base_f - c = torch.cosine_similarity(a_f.flatten(), b_f.flatten(), 0).clamp(-1, 1) + c = torch.cosine_similarity(a_f, b_f, -1).clamp(-1, 1).unsqueeze(-1) d = ((c + 1) / 2)**self.gamma - return (torch.lerp(a_f, b_f, self.alpha) * torch.lerp(d, 1 - d, self.beta)).to(a.dtype) + return (torch.lerp(a_f, b_f, self.alpha) * torch.lerp(d, 1 - d, self.beta)).to(dtype) class Similarities(ExtractOp): @@ -172,10 +200,14 @@ class PowerUpOp(Operation): def oper(self, a, b): if a is None or b is None: return None + a, b = resize_tensors(a, b) delta = b - a + keep_prob = 1 - self.alpha + if keep_prob <= 0: + return torch.zeros_like(delta) rng = torch.Generator(device=delta.device).manual_seed(self.seed) - m = torch.bernoulli(torch.full_like(delta, 1 - self.alpha), generator=rng) - return (m * delta) / (1 - self.alpha) if (1 - self.alpha) > 0 else torch.zeros_like(delta) + m = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).to(delta.dtype) + return (m * delta) / keep_prob class InterpolateDifference(Operation): @@ -194,6 +226,62 @@ class InterpolateDifference(Operation): mask = torch.lerp(torch.bernoulli(torch.clamp(diff, 0, 1), generator=rng), diff, self.gamma) return a * (1 - mask) + b * mask +class ManualEnhancedInterpolateDifferenceOp(Operation): + def __init__(self, key, alpha, beta, gamma, delta, seed, *sources): + super().__init__(key, *sources) + self.alpha, self.beta, self.gamma, self.delta, self.seed = alpha, beta, gamma, delta, seed + + def oper(self, a, b): + if a is None or b is None: return None + delta = torch.abs(a - b) + diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta)) + mean_diff = torch.mean(diff, 0, keepdim=True) + mask = torch.logical_and(self.beta < mean_diff, mean_diff < self.gamma) + powered_diff = torch.nan_to_num(diff ** (1 / max(self.alpha, 0.001) - 1)) + masked_diff = powered_diff * mask.float() + rng = torch.Generator(device=a.device).manual_seed(self.seed) + random_mask = torch.bernoulli(torch.clamp(masked_diff, 0, 1), generator=rng) + interpolated_mask = torch.lerp(random_mask, masked_diff, self.delta) + return a * (1 - interpolated_mask) + b * interpolated_mask + + +class AutoEnhancedInterpolateDifferenceOp(Operation): + def __init__(self, key, alpha, beta, gamma, seed, *sources): + super().__init__(key, *sources) + self.alpha, self.beta, self.gamma, self.seed = alpha, beta, gamma, seed + + def oper(self, a, b): + if a is None or b is None: return None + delta = torch.abs(a - b) + max_delta = torch.max(delta) + diff = torch.nan_to_num((max_delta - delta) / max_delta) + mean_diff = torch.mean(diff) + lower_threshold = mean_diff * (1 - self.beta) + upper_threshold = mean_diff * (1 + self.beta) + mask = torch.logical_and(lower_threshold < diff, diff < upper_threshold) + powered_diff = torch.nan_to_num(diff ** (1 / max(self.alpha, 0.001) - 1)) + masked_diff = powered_diff * mask.float() + rng = torch.Generator(device=a.device).manual_seed(self.seed) + random_mask = torch.bernoulli(torch.clamp(masked_diff, 0, 1), generator=rng) + interpolated_mask = torch.lerp(random_mask, masked_diff, self.gamma) + return a * (1 - interpolated_mask) + b * interpolated_mask + + +class WeightSumCutoffOp(Operation): + def __init__(self, key, alpha, beta, gamma, *sources): + super().__init__(key, *sources) + self.alpha, self.beta, self.gamma = alpha, beta, gamma + + def oper(self, a, b): + if a is None or b is None: return None + delta = torch.abs(a - b) + diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta)) + mean = torch.mean(diff, 0, True) + mask = torch.logical_and(mean < self.beta, self.gamma < mean) + mul = self.alpha * mask + return a * (1 - mul) + b * mul + + # ################# # # CALCULATION MODES # ################# # @@ -404,5 +492,68 @@ class SVDMode(CalcMode): return SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b) -TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode()] +class ManEnhInterpDifference(CalcMode): + name = 'Enhanced Man Interp' + description = 'Enhanced interpolation between each pair of values from A and B depending on their difference relative to other values.' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Interpolation strength.', 'beta': 'Lower mean threshold.', 'gamma': 'Upper mean threshold.', 'delta': 'Smoothness factor.'} + def create_recipe(self, key, **kwargs): + mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP) + + 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) + + if key.startswith('cond_stage_model.transformer.text_model.embeddings'): + return a + + loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode) + b = LoadTensor(key, kwargs['model_b'], **loader_args_b) + return ManualEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b) + + +class AutoEnhInterpDifference(CalcMode): + name = 'Enhanced Auto Interp' + description = 'Interpolates between each pair of values from A and B depending on their difference relative to other values.' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Interpolation strength.', 'beta': 'Threshold adjustment factor.', 'gamma': 'Smoothness factor.'} + def create_recipe(self, key, **kwargs): + mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP) + + 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) + + if key.startswith('cond_stage_model.transformer.text_model.embeddings'): + return a + + loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode) + b = LoadTensor(key, kwargs['model_b'], **loader_args_b) + return AutoEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b) + + +class WeightSumCutoffMode(CalcMode): + name = 'Weight-Sum Cutoff' + description = 'Weight-sum with cutoff based on value differences.' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Interpolation weight.', 'beta': 'Upper threshold.', 'gamma': 'Lower threshold.'} + def create_recipe(self, key, **kwargs): + mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP) + + 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) + return WeightSumCutoffOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b) + + +TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()] THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()] \ No newline at end of file