diff --git a/nodes/merger.py b/nodes/merger.py index fd458ab..da4d09f 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -93,6 +93,7 @@ class MergerLogic: mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip') mismatch_mode = MissingTensorBehavior(mismatch_mode_str) recipe_params['mismatch_mode'] = mismatch_mode + alignment_mode = recipe_params.get('alignment_mode', 'pad/crop') # Compile filter patterns exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', '')) @@ -168,6 +169,15 @@ class MergerLogic: for r_key, r_tensor in result.items(): merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone() else: + # Ensure compatibility with Model A's architecture. + # If alignment_mode is 'pad/crop', we crop results that were padded. + # If alignment_mode is 'interpolate', resizing happened during operators. + if alignment_mode == 'pad/crop': + target_shape = recipe_params['_tensor_a_shape'] + if result.shape != target_shape: + slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape)) + result = result[slices] + merged_state_dict[key] = result.to(save_torch_dtype).cpu().clone() # Clean up tensor_a reference to allow GC @@ -219,6 +229,7 @@ class CheckpointTwoMerger(io.ComfyNode): io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("checkpoints")), 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), @@ -247,7 +258,7 @@ class CheckpointTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, - "mismatch_mode": mismatch_mode, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -273,6 +284,7 @@ class ModelTwoMerger(io.ComfyNode): io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("diffusion_models")), 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), @@ -301,7 +313,7 @@ class ModelTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, - "mismatch_mode": mismatch_mode, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -327,6 +339,7 @@ class TextEncoderTwoMerger(io.ComfyNode): io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("text_encoders")), 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), @@ -355,7 +368,7 @@ class TextEncoderTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, - "mismatch_mode": mismatch_mode, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -381,6 +394,7 @@ class LoRATwoMerger(io.ComfyNode): io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("loras")), 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), @@ -409,7 +423,7 @@ class LoRATwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, - "mismatch_mode": mismatch_mode, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -435,6 +449,7 @@ class EmbeddingTwoMerger(io.ComfyNode): io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("embeddings")), 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), @@ -463,7 +478,7 @@ class EmbeddingTwoMerger(io.ComfyNode): recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, - "mismatch_mode": mismatch_mode, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -492,6 +507,7 @@ class CheckpointThreeMerger(io.ComfyNode): io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("checkpoints")), 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), @@ -520,7 +536,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, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -547,6 +563,7 @@ class ModelThreeMerger(io.ComfyNode): io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("diffusion_models")), 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), @@ -575,7 +592,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, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -602,6 +619,7 @@ class TextEncoderThreeMerger(io.ComfyNode): io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("text_encoders")), 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), @@ -630,7 +648,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, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -657,6 +675,7 @@ class LoRAThreeMerger(io.ComfyNode): io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("loras")), 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), @@ -685,7 +704,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, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, @@ -712,6 +731,7 @@ class EmbeddingThreeMerger(io.ComfyNode): io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("embeddings")), 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), @@ -740,7 +760,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, + "mismatch_mode": mismatch_mode, "alignment_mode": recipe_params.get('alignment_mode', 'pad/crop'), "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "output_filename": output_filename, "save_dtype": save_dtype, "device": process_device, "dtype": torch.float32, diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py index 96554d3..51fc384 100644 --- a/nodes/merger_ops.py +++ b/nodes/merger_ops.py @@ -18,8 +18,8 @@ class MissingTensorError(Exception): # ################# # # UTILITIES # ################# # -def resize_tensors(tensor1, tensor2): - """Pads tensors to match shapes (supports up to 4D).""" +def resize_by_interpolation(tensor1, tensor2): + """Resizes tensor2 to match tensor1's shape using interpolation.""" if tensor1 is None or tensor2 is None: return tensor1, tensor2 if tensor1.shape == tensor2.shape: @@ -28,6 +28,43 @@ def resize_tensors(tensor1, tensor2): dims1 = len(tensor1.shape) dims2 = len(tensor2.shape) + if dims1 != dims2 or dims1 > 4 or dims1 == 0: + return tensor1, tensor2 + + orig_dtype = tensor2.dtype + t2_f = tensor2.float() + + # F.interpolate expects [B, C, ...] + # We treat the tensor as [1, 1, ...dims] + view_shape = [1, 1] + list(tensor2.shape) + t2_f = t2_f.view(view_shape) + + # Select mode based on dimensions + if dims1 == 1: + mode = 'linear' + elif dims1 == 2: + mode = 'bilinear' + else: + mode = 'trilinear' + + t2_res = F.interpolate(t2_f, size=list(tensor1.shape), mode=mode, align_corners=False) + + return tensor1, t2_res.view(tensor1.shape).to(orig_dtype) + + +def resize_tensors(tensor1, tensor2, mode='pad/crop'): + """Resizes tensors to match shapes using either padding or interpolation.""" + if tensor1 is None or tensor2 is None: + return tensor1, tensor2 + if tensor1.shape == tensor2.shape: + return tensor1, tensor2 + + if mode == 'interpolate': + return resize_by_interpolation(tensor1, tensor2) + + dims1 = len(tensor1.shape) + dims2 = len(tensor2.shape) + if dims1 != dims2 or dims1 > 4 or dims1 == 0: return tensor1, tensor2 @@ -64,6 +101,7 @@ class Operation: self.key = key self.sources = tuple(sources) self.merge_func = self.recurse + self.alignment_mode = 'pad/crop' self.alpha = None self.beta = None self.gamma = None @@ -137,7 +175,7 @@ class Add(Operation): return b if b is None: return a - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) return a + b @@ -145,17 +183,33 @@ class Sub(Operation): def oper(self, a, b) -> torch.Tensor: if a is None or b is None: return None # Cannot compute difference with missing operand - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) return a - b class SVD(Operation): - def __init__(self, key, alpha, beta, gamma, *sources): + def __init__(self, key, alpha, beta, gamma, *sources, target_shape=None): super().__init__(key, *sources) self.alpha, self.beta, self.gamma = int(alpha), int(beta), gamma + self.target_shape = target_shape + def oper(self, a, b) -> torch.Tensor: if a is None or b is None: return None - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) + + # For SVD extraction, we must ensure the diff matches the original Model A shape + # so the resulting LoRA weights are compatible with the target architecture. + if self.target_shape and a.shape != self.target_shape: + if self.alignment_mode == 'interpolate': + # Re-align both to match target_shape via interpolation + dummy = torch.empty(self.target_shape, device=a.device) + _, a = resize_by_interpolation(dummy, a) + _, b = resize_by_interpolation(dummy, b) + else: + slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(a.shape, self.target_shape)) + a = a[slices] + b = b[slices] + 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) @@ -185,9 +239,9 @@ class TrainDiff(Operation): def oper(self, a, b, c) -> torch.Tensor: if a is None or b is None or c is None: return None - a, b = resize_tensors(a, b) - a, c = resize_tensors(a, c) - b, c = resize_tensors(b, c) + a, b = resize_tensors(a, b, mode=self.alignment_mode) + a, c = resize_tensors(a, c, mode=self.alignment_mode) + b, c = resize_tensors(b, c, mode=self.alignment_mode) if torch.allclose(b.float(), c.float()): return torch.zeros_like(a) diff_bc = b.float() - c.float() @@ -202,9 +256,9 @@ class ExtractOp(Operation): if a is None or b is None: return None if base is not None: - base, a = resize_tensors(base, a) - base, b = resize_tensors(base, b) - a, b = resize_tensors(a, b) + base, a = resize_tensors(base, a, mode=self.alignment_mode) + base, b = resize_tensors(base, b, mode=self.alignment_mode) + a, b = resize_tensors(a, b, mode=self.alignment_mode) dtype = base.dtype if base is not None else a.dtype base_f = base.float() if base is not None else 0 @@ -226,7 +280,7 @@ class PowerUpOp(Operation): def oper(self, a, b): if a is None or b is None: return None - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) delta = b - a keep_prob = 1 - self.alpha if keep_prob <= 0: @@ -243,7 +297,7 @@ class InterpolateDifference(Operation): def oper(self, a, b): if a is None or b is None: return None - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) alpha, delta = max(self.alpha, 0.001), torch.abs(a - b) max_delta = torch.max(delta) if max_delta == 0: return a @@ -260,7 +314,7 @@ class ManualEnhancedInterpolateDifferenceOp(Operation): def oper(self, a, b): if a is None or b is None: return None - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) 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) @@ -280,7 +334,7 @@ class AutoEnhancedInterpolateDifferenceOp(Operation): def oper(self, a, b): if a is None or b is None: return None - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) delta = torch.abs(a - b) max_delta = torch.max(delta) diff = torch.nan_to_num((max_delta - delta) / max_delta) @@ -303,7 +357,7 @@ class WeightSumCutoffOp(Operation): def oper(self, a, b): if a is None or b is None: return None - a, b = resize_tensors(a, b) + a, b = resize_tensors(a, b, mode=self.alignment_mode) delta = torch.abs(a - b) diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta)) mean = torch.mean(diff, 0, True) @@ -354,6 +408,7 @@ class WeightSum(CalcMode): param_docs = {'alpha': 'Interpolation weight. 0.0 is 100% Model A, 1.0 is 100% Model B.'} def create_recipe(self, key, **kwargs): mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP) + alignment_mode = kwargs.get('alignment_mode', 'pad/crop') # Use pre-loaded tensor_a if available if '_tensor_a' in kwargs: @@ -369,7 +424,9 @@ class WeightSum(CalcMode): if kwargs['alpha'] <= 0: return a c = Multiply(key, 1 - kwargs['alpha'], a) d = Multiply(key, kwargs['alpha'], b) - return Add(key, c, d) + res = Add(key, c, d) + res.alignment_mode = alignment_mode + return res class AddDifference(CalcMode): @@ -379,6 +436,7 @@ class AddDifference(CalcMode): param_docs = {'alpha': 'Multiplier for the (B - C) difference.', 'beta': 'Smoothing (0=Off, 1=On).'} 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']) @@ -392,7 +450,9 @@ class AddDifference(CalcMode): diff = Sub(key, b, c) if kwargs['beta'] == 1: diff = Smooth(key, diff) diffm = Multiply(key, kwargs['alpha'], diff) - return Add(key, a, diffm) + res = Add(key, a, diffm) + res.alignment_mode = alignment_mode + return res class TrainDifference(CalcMode): @@ -402,6 +462,7 @@ class TrainDifference(CalcMode): param_docs = {'alpha': 'Multiplier for the calculated 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']) @@ -414,7 +475,9 @@ class TrainDifference(CalcMode): c = LoadTensor(key, kwargs['model_c'], **loader_args_bc) diff = TrainDiff(key, a, b, c) diffm = Multiply(key, kwargs['alpha'], diff) - return Add(key, a, diffm) + res = Add(key, a, diffm) + res.alignment_mode = alignment_mode + return res class InterpDifference(CalcMode): name = 'Comparative-Interpolation' description = 'Interpolates A and B based on relative tensor value differences.' @@ -422,6 +485,7 @@ class InterpDifference(CalcMode): param_docs = {'alpha': 'Curve strength.', 'beta': 'Style (0=Similarity, 1=Difference).', 'gamma': 'Mix (0=Random, 1=Linear).'} 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']) @@ -435,7 +499,9 @@ class InterpDifference(CalcMode): loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode) b = LoadTensor(key, kwargs['model_b'], **loader_args_b) - return InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b) + res = InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b) + res.alignment_mode = alignment_mode + return res class Extract(CalcMode): @@ -445,6 +511,7 @@ class Extract(CalcMode): 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): 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']) @@ -457,7 +524,9 @@ class Extract(CalcMode): c = LoadTensor(key, kwargs['model_c'], **loader_args_bc) extracted = ExtractOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'] * 15, a, b, c) multiplied = Multiply(key, kwargs['delta'], extracted) - return Add(key, a, multiplied) + res = Add(key, a, multiplied) + res.alignment_mode = alignment_mode + return res class AddDisimilarity(CalcMode): @@ -467,6 +536,7 @@ class AddDisimilarity(CalcMode): param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Multiplier for final features.', 'gamma': 'Similarity bias.'} 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']) @@ -479,7 +549,9 @@ class AddDisimilarity(CalcMode): c = LoadTensor(key, kwargs['model_c'], **loader_args_bc) extracted = Similarities(key, kwargs['alpha'], 1, kwargs['gamma'] * 15, b, c) multiplied = Multiply(key, kwargs['beta'], extracted) - return Add(key, a, multiplied) + res = Add(key, a, multiplied) + res.alignment_mode = alignment_mode + return res class PowerUp(CalcMode): @@ -489,6 +561,7 @@ class PowerUp(CalcMode): param_docs = {'alpha': 'Dropout rate (proportion of weights from B to *drop*).', 'beta': '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']) @@ -499,8 +572,11 @@ class PowerUp(CalcMode): loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode) b = LoadTensor(key, kwargs['model_b'], **loader_args_b) deltahat = PowerUpOp(key, kwargs['alpha'], kwargs['seed'], a, b) - res = Multiply(key, kwargs['beta'], deltahat) - return Add(key, a, res) + deltahat.alignment_mode = alignment_mode # Important! + multiplied = Multiply(key, kwargs['beta'], deltahat) + res = Add(key, a, multiplied) + res.alignment_mode = alignment_mode + return res class SVDMode(CalcMode): @@ -510,6 +586,7 @@ class SVDMode(CalcMode): param_docs = {'alpha': 'Rank for standard layers.', 'beta': 'Rank for 3x3 conv layers.', 'gamma': 'Clamp quantile for weights.'} 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']) @@ -519,7 +596,9 @@ class SVDMode(CalcMode): loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode) b = LoadTensor(key, kwargs['model_b'], **loader_args_b) - return SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b) + res = SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b, target_shape=kwargs.get('_tensor_a_shape')) + res.alignment_mode = alignment_mode + return res class ManEnhInterpDifference(CalcMode): @@ -529,6 +608,7 @@ class ManEnhInterpDifference(CalcMode): 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) + alignment_mode = kwargs.get('alignment_mode', 'pad/crop') if '_tensor_a' in kwargs: a = PreloadedTensor(key, kwargs['_tensor_a']) @@ -541,7 +621,9 @@ class ManEnhInterpDifference(CalcMode): 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) + res = ManualEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b) + res.alignment_mode = alignment_mode + return res class AutoEnhInterpDifference(CalcMode): @@ -551,6 +633,7 @@ class AutoEnhInterpDifference(CalcMode): 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) + alignment_mode = kwargs.get('alignment_mode', 'pad/crop') if '_tensor_a' in kwargs: a = PreloadedTensor(key, kwargs['_tensor_a']) @@ -563,7 +646,9 @@ class AutoEnhInterpDifference(CalcMode): 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) + res = AutoEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b) + res.alignment_mode = alignment_mode + return res class WeightSumCutoffMode(CalcMode): @@ -573,6 +658,7 @@ class WeightSumCutoffMode(CalcMode): 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) + alignment_mode = kwargs.get('alignment_mode', 'pad/crop') if '_tensor_a' in kwargs: a = PreloadedTensor(key, kwargs['_tensor_a']) @@ -582,7 +668,9 @@ class WeightSumCutoffMode(CalcMode): 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) + res = WeightSumCutoffOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b) + res.alignment_mode = alignment_mode + return res TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]