diff --git a/docs/merger_2_model_modes.md b/docs/merger_2_model_modes.md index 3863f94..7532caf 100644 --- a/docs/merger_2_model_modes.md +++ b/docs/merger_2_model_modes.md @@ -33,6 +33,30 @@ --- +## Power-Up (DARE+TIES) +> Combines DARE (Drop and Rescale) with TIES (Trim, Elect Sign) — the current best-practice algorithm for merging fine-tuned models, particularly LoRA-tuned ones. DARE sparsifies the task vector via random dropout and rescaling; TIES then removes residual low-magnitude noise and enforces sign consistency. The result is a cleaner, more targeted capability transfer from B to A with less interference. + +**Models Used:** A, B +**Parameters:** +- **Alpha:** DARE drop rate. Fraction of delta parameters randomly zeroed. Higher values produce a sparser, more targeted delta. Typical range: `0.5`–`0.9`. +- **Beta:** TIES trim quantile. Fraction of the smallest-magnitude delta parameters zeroed after DARE, removing residual noise. `0.0` disables trimming. Typical range: `0.1`–`0.3`. +- **Gamma:** Lambda scale. Final multiplier applied to the merged delta before adding to Model A. Equivalent to the task arithmetic scaling coefficient. Typical range: `0.5`–`1.5`. +- **Seed:** Random seed for the DARE dropout mask. Controls reproducibility. + +**Algorithm:** +1. Compute task vector `δ = B − A` +2. **DARE:** apply random binary mask with keep probability `1 − alpha`; rescale survivors by `1/(1−alpha)` +3. **TIES trim:** zero out parameters below the `beta`-quantile magnitude threshold +4. **TIES elect:** determine dominant sign per position; zero out disagreeing parameters +5. Return `A + gamma × δ_filtered` + +**Guidance:** +- Start with `alpha=0.9, beta=0.2, gamma=1.0` for LoRA-derived models +- Lower `alpha` (e.g. `0.5`) for full fine-tunes where more delta mass should be retained +- Increase `gamma` beyond `1.0` to amplify B's influence; decrease below `1.0` to soften it + +--- + ## Enhanced Man Interp > Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control. diff --git a/nodes/merger.py b/nodes/merger.py index 101fb49..bf88066 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -313,10 +313,10 @@ class CheckpointTwoMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.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.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"]), @@ -372,10 +372,10 @@ class ModelTwoMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.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.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"]), @@ -431,10 +431,10 @@ class TextEncoderTwoMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.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.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"]), @@ -490,10 +490,10 @@ class LoRATwoMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.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.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"]), @@ -549,10 +549,10 @@ class EmbeddingTwoMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.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.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"]), @@ -611,10 +611,10 @@ class CheckpointThreeMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.5, min=-2.0, max=3.0, step=0.01), - io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + 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.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"]), @@ -671,10 +671,10 @@ class ModelThreeMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.5, min=-2.0, max=3.0, step=0.01), - io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + 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.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"]), @@ -731,10 +731,10 @@ class TextEncoderThreeMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.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.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"]), @@ -791,10 +791,10 @@ class LoRAThreeMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.5, min=-2.0, max=3.0, step=0.01), - io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + 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.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"]), @@ -851,10 +851,10 @@ class EmbeddingThreeMerger(io.ComfyNode): io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), - 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.5, min=-2.0, max=3.0, step=0.01), - io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + 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.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"]), diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py index 3479da5..fe19be3 100644 --- a/nodes/merger_ops.py +++ b/nodes/merger_ops.py @@ -319,6 +319,53 @@ class PowerUpOp(Operation): return (m * delta) / keep_prob +class DARETIESOp(Operation): + """DARE + TIES combined operation on a single task vector (delta = b - a). + + Step 1 — DARE: randomly drop `drop_rate` fraction of delta parameters, + rescale survivors by 1/(1-drop_rate). + Step 2 — TIES trim: zero out parameters below the `trim_quantile` magnitude threshold. + Step 3 — TIES elect: determine dominant sign per parameter position. + Step 4 — TIES merge: zero out parameters that disagree with the dominant sign. + + For a two-model merge the TIES elect/merge step degenerates to a no-op (single vector, + sign is always self-consistent) but the trim step still reduces interference from + near-zero noisy deltas. The dominant sign filter is included for correctness and + forward compatibility with n-model extensions. + """ + def __init__(self, key, alpha, beta, gamma, seed, *sources): + super().__init__(key, *sources) + # alpha = drop_rate (DARE), beta = trim_quantile (TIES), gamma = lambda scale + 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 + a, b = resize_tensors(a, b, mode=self.alignment_mode) + delta = (b - a).float() + + # --- DARE: drop and rescale --- + keep_prob = 1.0 - self.alpha + if keep_prob <= 0: + return torch.zeros_like(a) + rng = torch.Generator(device=delta.device).manual_seed(self.seed) + dare_mask = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).float() + delta = (delta * dare_mask) / keep_prob + + # --- TIES trim: zero parameters below magnitude threshold --- + if self.beta > 0: + flat_abs = delta.abs().flatten() + k = max(1, int(flat_abs.numel() * self.beta)) + threshold = torch.kthvalue(flat_abs, k).values + delta = torch.where(delta.abs() < threshold, torch.zeros_like(delta), delta) + + # --- TIES elect: dominant sign (trivially self-consistent for single vector) --- + dominant_sign = torch.sign(delta) + delta = torch.where(torch.sign(delta) == dominant_sign, delta, torch.zeros_like(delta)) + + return (a + self.gamma * delta).to(a.dtype) + + class InterpolateDifference(Operation): def __init__(self, key, alpha, beta, gamma, seed, *sources): super().__init__(key, *sources) @@ -710,5 +757,41 @@ class WeightSumCutoffMode(CalcMode): return res -TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()] +class DARETIESMode(CalcMode): + name = 'Power-Up (DARE+TIES)' + description = 'Adds capabilities of B to A using DARE (Drop and Rescale) followed by TIES (Trim, Elect Sign). Best practice for LoRA-tuned model merging.' + models_used = ['A', 'B'] + param_docs = { + 'alpha': 'DARE drop rate — fraction of delta parameters randomly zeroed. Higher = more aggressive sparsification.', + 'beta': 'TIES trim quantile — fraction of smallest-magnitude delta parameters zeroed after DARE. Removes residual noise.', + 'gamma': 'Lambda — scale multiplier applied to the final merged delta before adding to Model A.', + 'seed': 'Random seed for the DARE dropout mask.', + } + + 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 = DARETIESOp( + key, + kwargs['alpha'], # drop_rate + kwargs['beta'], # trim_quantile + kwargs['gamma'], # lambda scale + kwargs['seed'], + a, b, + ) + op.alignment_mode = alignment_mode + return op + + +TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()] THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()] \ No newline at end of file