From 79593701aac623e44d75187bd1b0c1e3daa5a5f5 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 19 Oct 2025 10:19:52 +0200 Subject: [PATCH] Add model merging nodes --- __init__.py | 35 +++-- docs/merger_2_model_modes.md | 42 ++++++ docs/merger_3_model_modes.md | 43 +++++++ nodes/__init__.py | 18 +-- nodes/merger.py | 165 ++++++++++++++++++++++++ nodes/merger_ops.py | 241 +++++++++++++++++++++++++++++++++++ nodes/merger_utils.py | 80 ++++++++++++ nodes/metakeys.py | 2 +- nodes/prunekeys.py | 2 +- nodes/renamekeys.py | 2 +- 10 files changed, 596 insertions(+), 34 deletions(-) create mode 100644 docs/merger_2_model_modes.md create mode 100644 docs/merger_3_model_modes.md create mode 100644 nodes/merger.py create mode 100644 nodes/merger_ops.py create mode 100644 nodes/merger_utils.py diff --git a/__init__.py b/__init__.py index 7c1e909..7265f2f 100644 --- a/__init__.py +++ b/__init__.py @@ -3,29 +3,26 @@ from .nodes import * NODE_CLASS_MAPPINGS = { "UNetMetaKeys": UNetMetaKeys, "CLIPMetaKeys": CLIPMetaKeys, - "LoRAMetaKeys": LoRAMetaKeys, - "CheckpointMetaKeys": CheckpointMetaKeys, - "UNetRenameKeys": UNetRenameKeys, - "CLIPRenameKeys": CLIPRenameKeys, - "LoRARenameKeys": LoRARenameKeys, - "CheckpointRenameKeys": CheckpointRenameKeys, - "UNetPruneKeys": UNetPruneKeys, - "CLIPPruneKeys": CLIPPruneKeys, - "LoRAPruneKeys": LoRAPruneKeys, "CheckpointPruneKeys": CheckpointPruneKeys, + "CheckpointTwoMerger": CheckpointTwoMerger, + "UNetTwoMerger": UNetTwoMerger, + "CLIPTwoMerger": CLIPTwoMerger, + "LoRATwoMerger": LoRATwoMerger, + "CheckpointThreeMerger": CheckpointThreeMerger, + "UNetThreeMerger": UNetThreeMerger, + "CLIPThreeMerger": CLIPThreeMerger, + "LoRAThreeMerger": LoRAThreeMerger, } NODE_DISPLAY_NAME_MAPPINGS = { "UNetMetaKeys": "Get UNet Meta Keys", - "CLIPMetaKeys": "Get CLIP Meta Keys", - "LoRAMetaKeys": "Get LoRA Meta Keys", - "CheckpointMetaKeys": "Get Checkpoint Meta Keys", - "UNetRenameKeys": "Rename UNet Keys", - "CLIPRenameKeys": "Rename CLIP Keys", - "LoRARenameKeys": "Rename LoRA Keys", - "CheckpointRenameKeys": "Rename Checkpoint Keys", - "UNetPruneKeys": "Prune UNet Keys", - "CLIPPruneKeys": "Prune CLIP Keys", - "LoRAPruneKeys": "Prune LoRA Keys", "CheckpointPruneKeys": "Prune Checkpoint Keys", + "CheckpointTwoMerger": "Merge Checkpoints (2 Models)", + "UNetTwoMerger": "Merge UNets (2 Models)", + "CLIPTwoMerger": "Merge CLIPs (2 Models)", + "LoRATwoMerger": "Merge LoRAs (2 Models)", + "CheckpointThreeMerger": "Merge Checkpoints (3 Models)", + "UNetThreeMerger": "Merge UNets (3 Models)", + "CLIPThreeMerger": "Merge CLIPs (3 Models)", + "LoRAThreeMerger": "Merge LoRAs (3 Models)", } \ No newline at end of file diff --git a/docs/merger_2_model_modes.md b/docs/merger_2_model_modes.md new file mode 100644 index 0000000..0a6d5f8 --- /dev/null +++ b/docs/merger_2_model_modes.md @@ -0,0 +1,42 @@ +# 2-Model Merger Documentation + +--- + +## Weight-Sum +> `A * (1 - α) + B * α`. A simple linear interpolation between two models. + +**Models Used:** A, B +**Parameters:** +- **Alpha:** Interpolation weight. `0.0` is 100% Model A, `1.0` is 100% Model B. + +--- + +## Comparative-Interpolation +> Interpolates between A and B based on the relative differences in their tensor values. Creates a unique blend by deciding how much of each model to use for each individual weight. + +**Models Used:** A, B +**Parameters:** +- **Alpha:** Controls the interpolation curve's strength and shape. +- **Beta:** Switches the interpolation style. `0.0` focuses on similarity, `1.0` focuses on difference. +- **Gamma:** Mixes between a randomized (binomial) interpolation at `0.0` and a smooth (linear) one at `1.0`. + +--- + +## 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. + +**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. +- **Beta:** A final multiplier for the rescaled difference before it's added to Model A. + +--- + +## SVD LoRA Extraction +> Extracts the difference between Model B and Model A and saves it as a new LoRA file. **This mode outputs a LoRA, not a full model.** The output file will be saved to your `loras` directory. + +**Models Used:** A (Tuned Model), B (Base Model) +**Parameters:** +- **Alpha:** The Rank (dimension) for the LoRA's standard layers. +- **Beta:** The Rank (dimension) specifically for 3x3 convolution layers. +- **Gamma:** The clamp quantile for weight values. `0.99` is a typical value to prevent outlier weights from dominating. \ No newline at end of file diff --git a/docs/merger_3_model_modes.md b/docs/merger_3_model_modes.md new file mode 100644 index 0000000..06b273e --- /dev/null +++ b/docs/merger_3_model_modes.md @@ -0,0 +1,43 @@ +# 3-Model Merger Documentation + +--- + +## Add-Difference +> `A + (B - C) * α`. This classic technique applies the "difference" between models B and C to model A. For example, it can be used to add a style (B) to a base model (A) while subtracting the style's original base (C). + +**Models Used:** A, B, C +**Parameters:** +- **Alpha:** A multiplier for the `(B - C)` difference vector before it's added to A. +- **Beta:** A smoothing toggle. `0` for Off, `1` for On. When On, applies a median and gaussian filter to the difference, which can result in a smoother merge. + +--- + +## Train-Difference +> A variation of Add-Difference from the SuperMerger extension that uses a more complex scaling factor based on the relative distances between the tensors of A, B, and C. + +**Models Used:** A, B, C +**Parameters:** +- **Alpha:** A multiplier for the final calculated "training" difference. + +--- + +## 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. + +**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. +- **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. + +**Models Used:** A, B, C +**Parameters:** +- **Alpha:** Weights the comparison between Model B (`0.0`) and Model C (`1.0`). +- **Beta:** A multiplier for the final extracted dissimilar features. +- **Gamma:** A bias exponent for the similarity calculation. \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py index deae698..c94a632 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -1,18 +1,12 @@ from .metakeys import * from .renamekeys import * from .prunekeys import * +from .merger import * __all__ = [ - "UNetMetaKeys", - "CLIPMetaKeys", - "LoRAMetaKeys", - "CheckpointMetaKeys", - "UNetRenameKeys", - "CLIPRenameKeys", - "LoRARenameKeys", - "CheckpointRenameKeys", - "UNetPruneKeys", - "CLIPPruneKeys", - "LoRAPruneKeys", - "CheckpointPruneKeys", + "UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", + "UNetRenameKeys", "CLIPRenameKeys", "LoRARenameKeys", "CheckpointRenameKeys", + "UNetPruneKeys", "CLIPPruneKeys", "LoRAPruneKeys", "CheckpointPruneKeys", + "CheckpointTwoMerger", "UNetTwoMerger", "CLIPTwoMerger", "LoRATwoMerger", + "CheckpointThreeMerger", "UNetThreeMerger", "CLIPThreeMerger", "LoRAThreeMerger", ] \ No newline at end of file diff --git a/nodes/merger.py b/nodes/merger.py new file mode 100644 index 0000000..61fb18c --- /dev/null +++ b/nodes/merger.py @@ -0,0 +1,165 @@ +import os +import torch +import folder_paths +import comfy.utils +from safetensors.torch import save_file +from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES +from .merger_utils import MemoryEfficientSafeOpen + +def load_documentation_from_file(filename): + """Loads documentation from a markdown file in the ../docs/ directory.""" + current_dir = os.path.dirname(os.path.realpath(__file__)) + docs_path = os.path.join(current_dir, '..', 'docs', filename) + try: + with open(docs_path, 'r', encoding='utf-8') as f: + return f.read() + except FileNotFoundError: + return f"# Documentation File Not Found\n\nPlease ensure `{filename}` exists in the `docs` directory of the custom node." + +# --- UI Dictionary Helper Functions --- +def _create_two_model_inputs(model_folder_name): + return {"required": { + "execution_mode": (["MERGE", "DOCUMENTATION ONLY"],), + "model_a": (["None"] + folder_paths.get_filename_list(model_folder_name),), + "model_b": (["None"] + folder_paths.get_filename_list(model_folder_name),), + "calc_mode": ([m.name for m in TWO_MODEL_MODES],), + "alpha": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), + "beta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), + "gamma": ("FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.001}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "output_filename": ("STRING", {"default": "merged_2_model"}), + "save_dtype": (["fp32", "fp16", "bf16"],), + "process_device": (["cuda", "cpu"],), + }} + +def _create_three_model_inputs(model_folder_name): + return {"required": { + "execution_mode": (["MERGE", "DOCUMENTATION ONLY"],), + "model_a": (["None"] + folder_paths.get_filename_list(model_folder_name),), + "model_b": (["None"] + folder_paths.get_filename_list(model_folder_name),), + "model_c": (["None"] + folder_paths.get_filename_list(model_folder_name),), + "calc_mode": ([m.name for m in THREE_MODEL_MODES],), + "alpha": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), + "beta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), + "gamma": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), + "delta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "output_filename": ("STRING", {"default": "merged_3_model"}), + "save_dtype": (["fp32", "fp16", "bf16"],), + "process_device": (["cuda", "cpu"],), + }} + +# --- Base Classes for Backend Logic ONLY --- +class MergerLogic: + def _execute_merge(self, model_names, calc_mode, all_modes, recipe_params, model_type): + calc_mode_class = next((m for m in all_modes if m.name == calc_mode), None) + if not calc_mode_class: raise ValueError(f"Calc mode '{calc_mode}' not found.") + primary_model_name = model_names.get('model_a') + if not primary_model_name or primary_model_name == "None": raise ValueError("Model A is required to run the merge.") + handlers = {} + for name in model_names.values(): + if name and name != "None": + path = folder_paths.get_full_path(model_type, name) + if not path: raise FileNotFoundError(f"Model '{name}' not found.") + handlers[name] = MemoryEfficientSafeOpen(path) + all_keys = handlers[primary_model_name].keys() + metadata = handlers[primary_model_name].header.get("__metadata__", {}) + merged_state_dict = {} + pbar = comfy.utils.ProgressBar(len(all_keys)) + save_dtype = recipe_params.pop('save_dtype') + save_torch_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(save_dtype) + recipe_params.update({"handlers": handlers}) # Add handlers, device and dtype are already in + for key in all_keys: + recipe = calc_mode_class.create_recipe(key=key, **recipe_params) + result = recipe.merge() + if isinstance(result, dict): + for r_key, r_tensor in result.items(): merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone() + else: merged_state_dict[key] = result.to(save_torch_dtype).cpu().clone() + pbar.update(1) + for handler in handlers.values(): handler.__exit__(None, None, None) + output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type + output_dir = folder_paths.get_folder_paths(output_folder)[0] + os.makedirs(output_dir, exist_ok=True) + output_filename = recipe_params.get("output_filename") + output_path = os.path.join(output_dir, f"{output_filename}.safetensors") + save_file(merged_state_dict, output_path, metadata=metadata) + return (f"{output_filename}.safetensors",) + +class BaseTwoMerger(MergerLogic): + FUNCTION = "merge"; CATEGORY = "ModelUtils/Merging" + RETURN_TYPES = ("STRING", "STRING",); RETURN_NAMES = ("output_filename", "documentation",) + + def merge(self, execution_mode, model_a, model_b, calc_mode, alpha, beta, gamma, seed, output_filename, save_dtype, process_device): + doc = load_documentation_from_file('merger_2_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return ("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32 + } + model_names = {"model_a": model_a, "model_b": model_b} + filename, = self._execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, self.MODEL_TYPE) + return (filename, doc) + +class BaseThreeMerger(MergerLogic): + FUNCTION = "merge"; CATEGORY = "ModelUtils/Merging" + RETURN_TYPES = ("STRING", "STRING",); RETURN_NAMES = ("output_filename", "documentation",) + + def merge(self, execution_mode, model_a, model_b, model_c, calc_mode, alpha, beta, gamma, delta, seed, output_filename, save_dtype, process_device): + doc = load_documentation_from_file('merger_3_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return ("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32 + } + model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} + filename, = self._execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, self.MODEL_TYPE) + return (filename, doc) + +# --- Concrete 2-Model Nodes --- +class CheckpointTwoMerger(BaseTwoMerger): + MODEL_TYPE = "checkpoints" + @classmethod + def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) +class UNetTwoMerger(BaseTwoMerger): + MODEL_TYPE = "diffusion_models" + @classmethod + def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) +class CLIPTwoMerger(BaseTwoMerger): + MODEL_TYPE = "text_encoders" + @classmethod + def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) +class LoRATwoMerger(BaseTwoMerger): + MODEL_TYPE = "loras" + @classmethod + def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) + +# --- Concrete 3-Model Nodes --- +class CheckpointThreeMerger(BaseThreeMerger): + MODEL_TYPE = "checkpoints" + @classmethod + def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) +class UNetThreeMerger(BaseThreeMerger): + MODEL_TYPE = "diffusion_models" + @classmethod + def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) +class CLIPThreeMerger(BaseThreeMerger): + MODEL_TYPE = "text_encoders" + @classmethod + def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) +class LoRAThreeMerger(BaseThreeMerger): + MODEL_TYPE = "loras" + @classmethod + def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) + +__all__ = [ + "CheckpointTwoMerger", "UNetTwoMerger", "CLIPTwoMerger", "LoRATwoMerger", + "CheckpointThreeMerger", "UNetThreeMerger", "CLIPThreeMerger", "LoRAThreeMerger" +] \ No newline at end of file diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py new file mode 100644 index 0000000..f8152f4 --- /dev/null +++ b/nodes/merger_ops.py @@ -0,0 +1,241 @@ +import torch +import scipy.ndimage +import torch.nn.functional as F + +# ################# # +# MERGE OPERATORS +# ################# # +class Operation: + def __init__(self, key, *sources): + self.key = key + self.sources = tuple(sources) + self.merge_func = self.recurse + self.alpha = None + self.beta = None + self.gamma = None + self.delta = None + self.seed = None + + def __eq__(self, other): + return (type(self), self.key, self.alpha, self.beta, self.gamma, self.delta, self.seed, self.sources) == \ + (type(other), other.key, other.alpha, other.beta, other.gamma, other.delta, other.seed, other.sources) + + def __hash__(self): + return hash((type(self), self.key, self.alpha, self.beta, self.gamma, self.delta, self.seed, self.sources)) + + def oper(self, *args) -> torch.Tensor: + raise NotImplementedError + + def recurse(self, operation): + source_tensors = [source_oper.merge() for source_oper in operation.sources] + return operation.oper(*source_tensors) + + def merge(self): + return self.merge_func(self) + +class LoadTensor(Operation): + def __init__(self, key, model_name, handlers, device, dtype): + super().__init__(key) + self.model_name = model_name + self.handlers = handlers + self.device = device + self.dtype = dtype + + def merge(self) -> torch.Tensor: + if self.model_name not in self.handlers: + raise ValueError(f"Model '{self.model_name}' is required for this mode but was not provided.") + return self.handlers[self.model_name].get_tensor(self.key).to(device=self.device, dtype=self.dtype) + +class Multiply(Operation): + def __init__(self, key, alpha, *sources): + super().__init__(key, *sources) + self.alpha = alpha + def oper(self, a) -> torch.Tensor: + return a * self.alpha +class Add(Operation): + def oper(self, a, b) -> torch.Tensor: + return a + b +class Sub(Operation): + def oper(self, a, b) -> torch.Tensor: + return a - b + +class SVD(Operation): + def __init__(self, key, alpha, beta, gamma, *sources): + super().__init__(key, *sources) + self.alpha, self.beta, self.gamma = int(alpha), int(beta), gamma + def oper(self, a, b) -> torch.Tensor: + diff, weights, conv2d = a - b, {}, len(a.size()) == 4 + kernel_size = None if not conv2d else a.size()[2:4] + conv2d_3x3 = conv2d and kernel_size != (1, 1) + rank = self.alpha if not conv2d_3x3 or self.beta is None else self.beta + rank = min(rank, a.size()[1], a.size()[0]) + matrix = diff.flatten(start_dim=1) if conv2d_3x3 else diff.squeeze() if conv2d else diff + U, S, Vh = torch.linalg.svd(matrix) + U, S, Vh = U[:, :rank], S[:rank], Vh[:rank, :] + U = U @ torch.diag(S) + dist = torch.cat([U.flatten(), Vh.flatten()]) + hi_val = torch.quantile(dist, self.gamma) + U, Vh = U.clamp(-hi_val, hi_val), Vh.clamp(-hi_val, hi_val) + if conv2d: + Vh = Vh.reshape(rank, a.size()[1], *kernel_size) if conv2d_3x3 else Vh.reshape(rank, a.size()[1], 1, 1) + U = U.reshape(a.size()[0], rank, 1, 1) + weights_sd = {f"{self.key}.lora_up.weight": U, f"{self.key}.lora_down.weight": Vh, f"{self.key}.alpha": torch.tensor(float(rank))} + return weights_sd + +class Smooth(Operation): + def oper(self, a) -> torch.Tensor: + filtered = scipy.ndimage.median_filter(a.detach().cpu().to(torch.float32).numpy(), size=3) + filtered = scipy.ndimage.gaussian_filter(filtered, sigma=1) + return torch.tensor(filtered, dtype=a.dtype, device=a.device) +class TrainDiff(Operation): + def oper(self, a, b, c) -> torch.Tensor: + if torch.allclose(b.float(), c.float()): return torch.zeros_like(a) + diff_bc = b.float() - c.float() + dist_bc, dist_ab = torch.abs(diff_bc), torch.abs(b.float() - a.float()) + scale = torch.where(dist_bc + dist_ab != 0, dist_ab / (dist_bc + dist_ab), 0.0) + return (torch.sign(diff_bc) * torch.abs(scale) * torch.abs(diff_bc)).to(a.dtype) * 1.8 +class Extract(Operation): + def __init__(self, key, alpha, beta, gamma, *args): + super().__init__(key, *args) + self.alpha, self.beta, self.gamma = alpha, beta, gamma + def oper(self, base: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + 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) + 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) +class Similarities(Extract): + def oper(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return super().oper(None, a, b) +class PowerUp(Operation): + def __init__(self, key, alpha, seed, *sources): + super().__init__(key, *sources) + self.alpha, self.seed = alpha, seed + def oper(self, a, b): + delta = b - a + 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) +class InterpolateDifference(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): + alpha, delta = max(self.alpha, 0.001), torch.abs(a - b) + max_delta = torch.max(delta) + if max_delta == 0: return a + diff = ((max_delta - delta) / max_delta if self.beta != 1 else delta / max_delta) ** (1 / alpha - 1) + diff = torch.nan_to_num(diff) + rng = torch.Generator(device=diff.device).manual_seed(self.seed) + mask = torch.lerp(torch.bernoulli(torch.clamp(diff, 0, 1), generator=rng), diff, self.gamma) + return a * (1 - mask) + b * mask + +# ################# # +# CALCULATION MODES +# ################# # +class CalcMode: + name = 'calcmode'; description = 'description'; models_used = []; param_docs = {} + def create_recipe(self, key, **kwargs) -> Operation: raise NotImplementedError + +class WeightSum(CalcMode): + name = 'Weight-Sum' + description = 'A * (1 - α) + B * α. Simple linear interpolation.' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Interpolation weight. 0.0 is 100% Model A, 1.0 is 100% Model B.'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + if kwargs['alpha'] >= 1: return b + if kwargs['alpha'] <= 0: return a + c = Multiply(key, 1 - kwargs['alpha'], a) + d = Multiply(key, kwargs['alpha'], b) + return Add(key, c, d) +class AddDifference(CalcMode): + name = 'Add-Difference' + description = 'A + (B - C) * α. Applies the difference between B and C to A.' + models_used = ['A', 'B', 'C'] + param_docs = {'alpha': 'Multiplier for the (B - C) difference.', 'beta': 'Smoothing (0=Off, 1=On).'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + c = LoadTensor(key, kwargs['model_c'], **loader_args) + diff = Sub(key, b, c) + if kwargs['beta'] == 1: diff = Smooth(key, diff) + diffm = Multiply(key, kwargs['alpha'], diff) + return Add(key, a, diffm) +class TrainDifference(CalcMode): + name = 'Train-Difference' + description = 'A + (B - C) * α, using tensor distances.' + models_used = ['A', 'B', 'C'] + param_docs = {'alpha': 'Multiplier for the calculated difference.'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + c = LoadTensor(key, kwargs['model_c'], **loader_args) + diff = TrainDiff(key, a, b, c) + diffm = Multiply(key, kwargs['alpha'], diff) + return Add(key, a, diffm) +class InterpDifference(CalcMode): + name = 'Comparative-Interpolation' + description = 'Interpolates A and B based on relative tensor value differences.' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Curve strength.', 'beta': 'Style (0=Similarity, 1=Difference).', 'gamma': 'Mix (0=Random, 1=Linear).'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + if key.startswith('cond_stage_model.transformer.text_model.embeddings'): return LoadTensor(key, kwargs['model_a'], **loader_args) + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + return InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b) +class Extract(CalcMode): + name = 'Extract-Features' + description = 'Adds features from B and C to A, based on their similarity.' + models_used = ['A', 'B', 'C'] + param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Similarity vs dissimilarity.', 'gamma': 'Similarity bias.', 'delta': 'Multiplier for final features.'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + c = LoadTensor(key, kwargs['model_c'], **loader_args) + extracted = Extract(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'] * 15, a, b, c) + multiplied = Multiply(key, kwargs['delta'], extracted) + return Add(key, a, multiplied) +class AddDisimilarity(CalcMode): + name = 'Add-Dissimilarities' + description = 'Adds dissimilar features between B and C to model A.' + models_used = ['A', 'B', 'C'] + param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Multiplier for final features.', 'gamma': 'Similarity bias.'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + c = LoadTensor(key, kwargs['model_c'], **loader_args) + extracted = Similarities(key, kwargs['alpha'], 1, kwargs['gamma'] * 15, b, c) + multiplied = Multiply(key, kwargs['beta'], extracted) + return Add(key, a, multiplied) +class PowerUp(CalcMode): + name = 'Power-Up (DARE)' + description = 'Adds capabilities of B to A using Drop and Rescale (DARE).' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Dropout rate (proportion of weights from B to *drop*).', 'beta': 'Multiplier for the final difference.'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + deltahat = PowerUp(key, kwargs['alpha'], kwargs['seed'], a, b) + res = Multiply(key, kwargs['beta'], deltahat) + return Add(key, a, res) +class SVD(CalcMode): + name = 'SVD LoRA Extraction' + description = 'Extracts the difference between B and A into a LoRA file.' + models_used = ['A', 'B'] + param_docs = {'alpha': 'Rank for standard layers.', 'beta': 'Rank for 3x3 conv layers.', 'gamma': 'Clamp quantile for weights.'} + def create_recipe(self, key, **kwargs): + loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']} + a = LoadTensor(key, kwargs['model_a'], **loader_args) + b = LoadTensor(key, kwargs['model_b'], **loader_args) + return SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b) + +TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVD()] +THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()] \ No newline at end of file diff --git a/nodes/merger_utils.py b/nodes/merger_utils.py new file mode 100644 index 0000000..e18f872 --- /dev/null +++ b/nodes/merger_utils.py @@ -0,0 +1,80 @@ +import json +import struct +import torch + +class MemoryEfficientSafeOpen: + def __init__(self, filename, device='cpu'): + self.filename = filename + self.device = device + self.header, self.header_size = self._read_header() + self.file = open(filename, "rb") + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.file.close() + + def keys(self): + return [k for k in self.header.keys() if k != "__metadata__"] + + def get_tensor(self, key): + if key not in self.header: + raise KeyError(f"Tensor '{key}' not found in the file") + + metadata = self.header[key] + offset_start, offset_end = metadata["data_offsets"] + + tensor_bytes = None + if offset_start != offset_end: + self.file.seek(self.header_size + 8 + offset_start) + tensor_bytes = self.file.read(offset_end - offset_start) + + return self._deserialize_tensor(tensor_bytes, metadata) + + def _read_header(self): + with open(self.filename, "rb") as f: + header_size = struct.unpack("