Add Merge Files Natively

modules/extras.py

@@ -8,7 +8,8 @@ import torch
 import tqdm
 import gradio as gr
 import safetensors.torch
-from sd_meh.merge import merge_models
+from modules.merging.merge import merge_models
+from modules.merging.utils import TRIPLE_METHODS
 
 from modules import shared, images, sd_models, sd_vae, sd_models_config
 
@@ -232,6 +233,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         metadata["sd_merge_models"] = json.dumps(metadata["sd_merge_models"])
 
     _, extension = os.path.splitext(output_modelname)
+    theta_0 = theta_0.to_dict()
     if extension.lower() == ".safetensors":
         safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
     else:
@@ -249,11 +251,29 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
 
 def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
-    shared.state.begin('model-merge')
+    shared.state.begin('merge')
+
+    def fail(message):
+        shared.state.textinfo = message
+        shared.state.end()
+        return [*[gr.update() for _ in range(4)], message]
+
     kwargs["models"] = {
         "model_a": sd_models.checkpoints_list[kwargs.get("primary_model_name", None)].filename,
         "model_b": sd_models.checkpoints_list[kwargs.get("secondary_model_name", None)].filename,
     }
+
+    if kwargs.get("primary_model_name", None) in [None, 'None']:
+        return fail("Failed: Merging requires a primary model.")
+    primary_model_info = sd_models.checkpoints_list[kwargs.get("primary_model_name", None)]
+    if kwargs.get("secondary_model_name", None) in [None, 'None']:
+        return fail("Failed: Merging requires a secondary model.")
+    secondary_model_info = sd_models.checkpoints_list[kwargs.get("secondary_model_name", None)]
+    if kwargs.get("tertiary_model_name", None) in [None, 'None'] and kwargs.get("merge_mode", None) in TRIPLE_METHODS:
+        return fail(f"Failed: Interpolation method ({kwargs.get('merge_mode', None)}) requires a tertiary model.")
+    tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if kwargs.get("merge_mode",
+                                                                                        None) in TRIPLE_METHODS else None
+
     del kwargs["primary_model_name"]
     del kwargs["secondary_model_name"]
     if kwargs.get("tertiary_model_name", None) is not None:
@@ -261,12 +281,14 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         del kwargs["tertiary_model_name"]
 
     try:
-        alpha = [float(x) for x in [kwargs["alpha_base"]] + kwargs["alpha_in_blocks"].split(",") + [kwargs["alpha_mid_block"]] + kwargs["alpha_out_blocks"].split(",")]
-        assert len(alpha) == 26 or len(alpha) == 20, "Alpha Block Weights are wrong length (26 or 20 for SDXL) falling back"
+        alpha = [float(x) for x in
+                 [kwargs["alpha_base"]] + kwargs["alpha_in_blocks"].split(",") + [kwargs["alpha_mid_block"]] + kwargs[
+                     "alpha_out_blocks"].split(",")]
+        assert len(alpha) == 26 or len(
+            alpha) == 20, "Alpha Block Weights are wrong length (26 or 20 for SDXL) falling back"
         kwargs["alpha"] = alpha
     except Exception as e:
         kwargs["alpha"] = kwargs.get("alpha_preset", kwargs["alpha"])
-        print(e)
     finally:
         kwargs.pop("alpha_base", None)
         kwargs.pop("alpha_in_blocks", None)
@@ -275,12 +297,14 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         kwargs.pop("alpha_preset", None)
     if kwargs.get("beta", False):
         try:
-            beta = [float(x) for x in [kwargs["beta_base"]] + kwargs["beta_in_blocks"].split(",") + [kwargs["beta_mid_block"]] + kwargs["beta_out_blocks"].split(",")]
-            assert len(beta) == 26 or len(beta) == 20, "Beta Block Weights are wrong length (26 or 20 for SDXL) falling back"
+            beta = [float(x) for x in
+                    [kwargs["beta_base"]] + kwargs["beta_in_blocks"].split(",") + [kwargs["beta_mid_block"]] + kwargs[
+                        "beta_out_blocks"].split(",")]
+            assert len(beta) == 26 or len(
+                beta) == 20, "Beta Block Weights are wrong length (26 or 20 for SDXL) falling back"
             kwargs["beta"] = beta
         except Exception as e:
             kwargs["beta"] = kwargs.get("beta_preset", kwargs["beta"])
-            print(e)
         finally:
             kwargs.pop("beta_base", None)
             kwargs.pop("beta_in_blocks", None)
@@ -288,31 +312,29 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             kwargs.pop("beta_out_blocks", None)
             kwargs.pop("beta_preset", None)
 
-    return [*[gr.update() for _ in range(4)], f"{kwargs}"]
+    try:
+        theta_0 = merge_models(**kwargs)
+    except Exception as e:
+        return fail(f"{e}")
 
-    def fail(message):
-        shared.state.textinfo = message
-        shared.state.end()
-        return [*[gr.update() for _ in range(4)], message]
-
-    theta_0 = merge_models(**kwargs)
     ckpt_dir = shared.opts.ckpt_dir or sd_models.model_path
-    filename = custom_name
+    filename = kwargs.get("custom_name", "Unamed_Merge")
     filename += "." + kwargs.get("checkpoint_format", None)
     output_modelname = os.path.join(ckpt_dir, filename)
     shared.state.textinfo = "Saving"
     metadata = None
-    if save_metadata:
+    if kwargs.get("save_metadata", False):
         metadata = {"format": "pt", "sd_merge_models": {}}
         merge_recipe = {
             "type": "SDNext",  # indicate this model was merged with webui's built-in merger
             "primary_model_hash": primary_model_info.sha256,
             "secondary_model_hash": secondary_model_info.sha256 if secondary_model_info else None,
             "tertiary_model_hash": tertiary_model_info.sha256 if tertiary_model_info else None,
-            "interp_method": interp_method,
-            "multiplier": multiplier,
-            "save_as_half": save_as_half,
-            "custom_name": custom_name,
+            "merge_mode": kwargs.get('merge_mode', None),
+            "alpha": kwargs.get('alpha', None),
+            "beta": kwargs.get('beta', None),
+            "precision": kwargs.get('precision', None),
+            "custom_name": kwargs.get("custom_name", "Unamed_Merge"),
         }
         metadata["sd_merge_recipe"] = json.dumps(merge_recipe)
 
@@ -333,6 +355,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         metadata["sd_merge_models"] = json.dumps(metadata["sd_merge_models"])
 
     _, extension = os.path.splitext(output_modelname)
+    theta_0 = theta_0.to_dict()
     if extension.lower() == ".safetensors":
         safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
     else:
@@ -342,7 +365,6 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     created_model = next((ckpt for ckpt in sd_models.checkpoints_list.values() if ckpt.name == filename), None)
     if created_model:
         created_model.calculate_shorthash()
-    create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)
     shared.log.info(f"Model merge saved: {output_modelname}.")
     shared.state.textinfo = "Checkpoint saved"
     shared.state.end()

modules/merging/merge.py

@@ -0,0 +1,443 @@
+import gc
+import os
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import safetensors.torch
+import torch
+from tqdm import tqdm
+
+from modules.merging import merge_methods
+from modules.merging.utils import WeightClass
+from modules.merging.model import SDModel
+from modules.merging.rebasin import (
+    apply_permutation,
+    sdunet_permutation_spec,
+    update_model_a,
+    weight_matching,
+)
+
+# logging.getLogger("sd_meh").addHandler(logging.NullHandler())
+
+MAX_TOKENS = 77
+
+
+KEY_POSITION_IDS = ".".join(
+    [
+        "cond_stage_model",
+        "transformer",
+        "text_model",
+        "embeddings",
+        "position_ids",
+    ]
+)
+
+
+NAI_KEYS = {
+    "cond_stage_model.transformer.embeddings.": "cond_stage_model.transformer.text_model.embeddings.",
+    "cond_stage_model.transformer.encoder.": "cond_stage_model.transformer.text_model.encoder.",
+    "cond_stage_model.transformer.final_layer_norm.": "cond_stage_model.transformer.text_model.final_layer_norm.",
+}
+
+
+def fix_clip(model: Dict) -> Dict:
+    if KEY_POSITION_IDS in model.keys():
+        model[KEY_POSITION_IDS] = torch.tensor(
+            [list(range(MAX_TOKENS))],
+            dtype=torch.int64,
+            device=model[KEY_POSITION_IDS].device,
+        )
+
+    return model
+
+
+def fix_key(model: Dict, key: str) -> Dict:
+    for nk in NAI_KEYS:
+        if key.startswith(nk):
+            model[key.replace(nk, NAI_KEYS[nk])] = model[key]
+            del model[key]
+
+    return model
+
+
+# https://github.com/j4ded/sdweb-merge-block-weighted-gui/blob/master/scripts/mbw/merge_block_weighted.py#L115
+def fix_model(model: Dict) -> Dict:
+    for k in model.keys():
+        model = fix_key(model, k)
+    return fix_clip(model)
+
+
+def load_sd_model(model: os.PathLike | str, device: str = "cpu") -> Dict:
+    if isinstance(model, str):
+        model = Path(model)
+
+    return SDModel(model, device).load_model()
+
+
+def prune_sd_model(model: Dict) -> Dict:
+    keys = list(model.keys())
+    for k in keys:
+        if (
+            not k.startswith("model.diffusion_model.")
+            and not k.startswith("first_stage_model.")
+            and not k.startswith("cond_stage_model.")
+        ):
+            del model[k]
+    return model
+
+
+def restore_sd_model(original_model: Dict, merged_model: Dict) -> Dict:
+    for k in original_model:
+        if k not in merged_model:
+            merged_model[k] = original_model[k]
+    return merged_model
+
+
+def log_vram(txt=""):
+    alloc = torch.cuda.memory_allocated(0)
+    # logging.debug(f"{txt} VRAM: {alloc*1e-9:5.3f}GB")
+
+
+def load_thetas(
+    models: Dict[str, os.PathLike | str],
+    prune: bool,
+    device: str,
+    precision: str,
+) -> Dict:
+    log_vram("before loading models")
+    if prune:
+        thetas = {k: prune_sd_model(load_sd_model(m, "cpu")) for k, m in models.items()}
+    else:
+        thetas = {k: load_sd_model(m, device) for k, m in models.items()}
+
+    if device == "cuda":
+        for model_key, model in thetas.items():
+            for key, block in model.items():
+                if precision == "fp16":
+                    thetas[model_key].update({key: block.to(device).half()})
+                else:
+                    thetas[model_key].update({key: block.to(device)})
+
+    log_vram("models loaded")
+    return thetas
+
+
+def merge_models(
+    models: Dict[str, os.PathLike | str],
+    merge_mode: str,
+    precision: str = "full",
+    weights_clip: bool = False,
+    re_basin: bool = False,
+    iterations: int = 1,
+    device: str = "cpu",
+    work_device: Optional[str] = None,
+    prune: bool = False,
+    threads: int = 1,
+    **kwargs,
+) -> Dict:
+    thetas = load_thetas(models, prune, device, precision)
+
+    # logging.info(f"start merging with {merge_mode} method")
+    weight_matcher = WeightClass(thetas["model_a"], **kwargs)
+    if re_basin:
+        merged = rebasin_merge(
+            thetas,
+            weight_matcher,
+            merge_mode,
+            precision=precision,
+            weights_clip=weights_clip,
+            iterations=iterations,
+            device=device,
+            work_device=work_device,
+            threads=threads,
+        )
+    else:
+        merged = simple_merge(
+            thetas,
+            weight_matcher,
+            merge_mode,
+            precision=precision,
+            weights_clip=weights_clip,
+            device=device,
+            work_device=work_device,
+            threads=threads,
+        )
+
+    return un_prune_model(merged, thetas, models, device, prune, precision)
+
+
+def un_prune_model(
+    merged: Dict,
+    thetas: Dict,
+    models: Dict,
+    device: str,
+    prune: bool,
+    precision: str,
+) -> Dict:
+    if prune:
+        # logging.info("Un-pruning merged model")
+        del thetas
+        gc.collect()
+        log_vram("remove thetas")
+        original_a = load_sd_model(models["model_a"], device)
+        for key in tqdm(original_a.keys(), desc="un-prune model a"):
+            if KEY_POSITION_IDS in key:
+                continue
+            if "model" in key and key not in merged.keys():
+                merged.update({key: original_a[key]})
+                if precision == "fp16":
+                    merged.update({key: merged[key].half()})
+        del original_a
+        gc.collect()
+        log_vram("remove original_a")
+        original_b = load_sd_model(models["model_b"], device)
+        for key in tqdm(original_b.keys(), desc="un-prune model b"):
+            if KEY_POSITION_IDS in key:
+                continue
+            if "model" in key and key not in merged.keys():
+                merged.update({key: original_b[key]})
+                if precision == "fp16":
+                    merged.update({key: merged[key].half()})
+        del original_b
+
+    return fix_model(merged)
+
+
+def simple_merge(
+    thetas: Dict[str, Dict],
+    weight_matcher: WeightClass,
+    merge_mode: str,
+    precision: str = "fp16",
+    weights_clip: bool = False,
+    device: str = "cpu",
+    work_device: Optional[str] = None,
+    threads: int = 1,
+) -> Dict:
+    futures = []
+    with tqdm(thetas["model_a"].keys(), desc="stage 1") as progress:
+        with ThreadPoolExecutor(max_workers=threads) as executor:
+            for key in thetas["model_a"].keys():
+                future = executor.submit(
+                    simple_merge_key,
+                    progress,
+                    key,
+                    thetas,
+                    weight_matcher,
+                    merge_mode,
+                    precision,
+                    weights_clip,
+                    device,
+                    work_device,
+                )
+                futures.append(future)
+
+        for res in futures:
+            res.result()
+
+    log_vram("after stage 1")
+
+    for key in tqdm(thetas["model_b"].keys(), desc="stage 2"):
+        if KEY_POSITION_IDS in key:
+            continue
+        if "model" in key and key not in thetas["model_a"].keys():
+            thetas["model_a"].update({key: thetas["model_b"][key]})
+            if precision == 16:
+                thetas["model_a"].update({key: thetas["model_a"][key].half()})
+
+    log_vram("after stage 2")
+
+    return fix_model(thetas["model_a"])
+
+
+def rebasin_merge(
+    thetas: Dict[str, os.PathLike | str],
+    weight_matcher: WeightClass,
+    merge_mode: str,
+    precision: str = "fp16",
+    weights_clip: bool = False,
+    iterations: int = 1,
+    device="cpu",
+    work_device=None,
+    threads: int = 1,
+):
+    # WARNING: not sure how this does when 3 models are involved...
+
+    model_a = thetas["model_a"].clone()
+    perm_spec = sdunet_permutation_spec()
+
+    print("Init rebasin iterations")
+    for it in range(iterations):
+        print(f"Rebasin iteration {it}")
+        log_vram(f"{it} iteration start")
+        weight_matcher.set_it(it)
+        log_vram("weights & bases, before simple merge")
+
+        # normal block merge we already know and love
+        thetas["model_a"] = simple_merge(
+            thetas,
+            weight_matcher,
+            merge_mode,
+            precision,
+            False,
+            device,
+            work_device,
+            threads,
+        )
+
+        log_vram("simple merge done")
+
+        # find permutations
+        perm_1, y = weight_matching(
+            perm_spec,
+            model_a,
+            thetas["model_a"],
+            max_iter=it,
+            init_perm=None,
+            usefp16=precision == 16,
+            device=device,
+        )
+
+        log_vram("weight matching #1 done")
+
+        thetas["model_a"] = apply_permutation(perm_spec, perm_1, thetas["model_a"])
+
+        log_vram("apply perm 1 done")
+
+        perm_2, z = weight_matching(
+            perm_spec,
+            thetas["model_b"],
+            thetas["model_a"],
+            max_iter=it,
+            init_perm=None,
+            usefp16=precision == 16,
+            device=device,
+        )
+
+        log_vram("weight matching #2 done")
+
+        new_alpha = torch.nn.functional.normalize(
+            torch.sigmoid(torch.Tensor([y, z])), p=1, dim=0
+        ).tolist()[0]
+        thetas["model_a"] = update_model_a(
+            perm_spec, perm_2, thetas["model_a"], new_alpha
+        )
+
+        log_vram("model a updated")
+
+    if weights_clip:
+        clip_thetas = thetas.copy()
+        clip_thetas["model_a"] = model_a
+        thetas["model_a"] = clip_weights(thetas, thetas["model_a"])
+
+    return thetas["model_a"]
+
+
+def simple_merge_key(progress, key, thetas, *args, **kwargs):
+    with merge_key_context(key, thetas, *args, **kwargs) as result:
+        if result is not None:
+            thetas["model_a"].update({key: result.detach().clone()})
+
+        progress.update()
+
+
+def merge_key(
+    key: str,
+    thetas: Dict,
+    weight_matcher: WeightClass,
+    merge_mode: str,
+    precision: int = 16,
+    weights_clip: bool = False,
+    device: str = "cpu",
+    work_device: Optional[str] = None,
+) -> Optional[Tuple[str, Dict]]:
+    if work_device is None:
+        work_device = device
+
+    if KEY_POSITION_IDS in key:
+        return
+
+    for theta in thetas.values():
+        if key not in theta.keys():
+            return
+
+        current_bases = weight_matcher(key)
+        try:
+            merge_method = getattr(merge_methods, merge_mode)
+        except AttributeError as e:
+            raise ValueError(f"{merge_mode} not implemented, aborting merge!") from e
+
+        merge_args = get_merge_method_args(current_bases, thetas, key, work_device)
+
+        # dealing with pix2pix and inpainting models
+        if (a_size := merge_args["a"].size()) != (b_size := merge_args["b"].size()):
+            if a_size[1] > b_size[1]:
+                merged_key = merge_args["a"]
+            else:
+                merged_key = merge_args["b"]
+        else:
+            merged_key = merge_method(**merge_args).to(device)
+
+        if weights_clip:
+            merged_key = clip_weights_key(thetas, merged_key, key)
+
+        if precision == 16:
+            merged_key = merged_key.half()
+
+        return merged_key
+
+
+def clip_weights(thetas, merged):
+    for k in thetas["model_a"].keys():
+        if k in thetas["model_b"].keys():
+            merged.update({k: clip_weights_key(thetas, merged[k], k)})
+    return merged
+
+
+def clip_weights_key(thetas, merged_weights, key):
+    t0 = thetas["model_a"][key]
+    t1 = thetas["model_b"][key]
+    maximums = torch.maximum(t0, t1)
+    minimums = torch.minimum(t0, t1)
+    return torch.minimum(torch.maximum(merged_weights, minimums), maximums)
+
+
+@contextmanager
+def merge_key_context(*args, **kwargs):
+    result = merge_key(*args, **kwargs)
+    try:
+        yield result
+    finally:
+        if result is not None:
+            del result
+
+
+def get_merge_method_args(
+    current_bases: Dict,
+    thetas: Dict,
+    key: str,
+    work_device: str,
+) -> Dict:
+    merge_method_args = {
+        "a": thetas["model_a"][key].to(work_device),
+        "b": thetas["model_b"][key].to(work_device),
+        **current_bases,
+    }
+
+    if "model_c" in thetas:
+        merge_method_args["c"] = thetas["model_c"][key].to(work_device)
+
+    return merge_method_args
+
+
+def save_model(model, output_file, file_format) -> None:
+    # logging.info(f"Saving {output_file}")
+    if file_format == "safetensors":
+        safetensors.torch.save_file(
+            model if type(model) == dict else model.to_dict(),
+            f"{output_file}.safetensors",
+            metadata={"format": "pt"},
+        )
+    else:
+        torch.save({"state_dict": model}, f"{output_file}.ckpt")

modules/merging/merge_methods.py

@@ -0,0 +1,211 @@
+import math
+from typing import Tuple
+
+import torch
+from torch import Tensor
+
+__all__ = [
+    "weighted_sum",
+    "weighted_subtraction",
+    "tensor_sum",
+    "add_difference",
+    "sum_twice",
+    "triple_sum",
+    "euclidean_add_difference",
+    "multiply_difference",
+    "top_k_tensor_sum",
+    "similarity_add_difference",
+    "distribution_crossover",
+    "ties_add_difference",
+]
+
+
+EPSILON = 1e-10  # Define a small constant EPSILON to prevent division by zero
+
+
+def weighted_sum(a: Tensor, b: Tensor, alpha: float, **kwargs) -> Tensor:
+    return (1 - alpha) * a + alpha * b
+
+
+def weighted_subtraction(
+    a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    # Adjust beta if both alpha and beta are 1.0 to avoid division by zero
+    if alpha == 1.0 and beta == 1.0:
+        beta -= EPSILON
+
+    return (a - alpha * beta * b) / (1 - alpha * beta)
+
+
+def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    if alpha + beta <= 1:
+        tt = a.clone()
+        talphas = int(a.shape[0] * beta)
+        talphae = int(a.shape[0] * (alpha + beta))
+        tt[talphas:talphae] = b[talphas:talphae].clone()
+    else:
+        talphas = int(a.shape[0] * (alpha + beta - 1))
+        talphae = int(a.shape[0] * beta)
+        tt = b.clone()
+        tt[talphas:talphae] = a[talphas:talphae].clone()
+    return tt
+
+
+def add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
+    return a + alpha * (b - c)
+
+
+def sum_twice(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    return (1 - beta) * ((1 - alpha) * a + alpha * b) + beta * c
+
+
+def triple_sum(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    return (1 - alpha - beta) * a + alpha * b + beta * c
+
+
+def euclidean_add_difference(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs
+) -> Tensor:
+    a_diff = a.float() - c.float()
+    b_diff = b.float() - c.float()
+    a_diff = torch.nan_to_num(a_diff / torch.linalg.norm(a_diff))
+    b_diff = torch.nan_to_num(b_diff / torch.linalg.norm(b_diff))
+
+    distance = (1 - alpha) * a_diff**2 + alpha * b_diff**2
+    distance = torch.sqrt(distance)
+    sum_diff = weighted_sum(a.float(), b.float(), alpha) - c.float()
+    distance = torch.copysign(distance, sum_diff)
+
+    target_norm = torch.linalg.norm(sum_diff)
+    return c + distance / torch.linalg.norm(distance) * target_norm
+
+
+def multiply_difference(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    diff_a = torch.pow(torch.abs(a.float() - c), (1 - alpha))
+    diff_b = torch.pow(torch.abs(b.float() - c), alpha)
+    difference = torch.copysign(diff_a * diff_b, weighted_sum(a, b, beta) - c)
+    return c + difference.to(c.dtype)
+
+
+def top_k_tensor_sum(
+    a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    a_flat = torch.flatten(a)
+    a_dist = torch.msort(a_flat)
+    b_indices = torch.argsort(torch.flatten(b), stable=True)
+    redist_indices = torch.argsort(b_indices)
+
+    start_i, end_i, region_is_inverted = ratio_to_region(alpha, beta, torch.numel(a))
+    start_top_k = kth_abs_value(a_dist, start_i)
+    end_top_k = kth_abs_value(a_dist, end_i)
+
+    indices_mask = (start_top_k < torch.abs(a_dist)) & (torch.abs(a_dist) <= end_top_k)
+    if region_is_inverted:
+        indices_mask = ~indices_mask
+    indices_mask = torch.gather(indices_mask.float(), 0, redist_indices)
+
+    a_redist = torch.gather(a_dist, 0, redist_indices)
+    a_redist = (1 - indices_mask) * a_flat + indices_mask * a_redist
+    return a_redist.reshape_as(a)
+
+
+def kth_abs_value(a: Tensor, k: int) -> Tensor:
+    if k <= 0:
+        return torch.tensor(-1, device=a.device)
+    else:
+        return torch.kthvalue(torch.abs(a.float()), k)[0]
+
+
+def ratio_to_region(width: float, offset: float, n: int) -> Tuple[int, int, bool]:
+    if width < 0:
+        offset += width
+        width = -width
+    width = min(width, 1)
+
+    if offset < 0:
+        offset = 1 + offset - int(offset)
+    offset = math.fmod(offset, 1.0)
+
+    if width + offset <= 1:
+        inverted = False
+        start = offset * n
+        end = (width + offset) * n
+    else:
+        inverted = True
+        start = (width + offset - 1) * n
+        end = offset * n
+
+    return round(start), round(end), inverted
+
+
+def similarity_add_difference(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    threshold = torch.maximum(torch.abs(a), torch.abs(b))
+    similarity = ((a * b / threshold**2) + 1) / 2
+    similarity = torch.nan_to_num(similarity * beta, nan=beta)
+
+    ab_diff = a + alpha * (b - c)
+    ab_sum = (1 - alpha / 2) * a + (alpha / 2) * b
+    return (1 - similarity) * ab_diff + similarity * ab_sum
+
+
+def distribution_crossover(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
+):
+    if a.shape == ():
+        return alpha * a + (1 - alpha) * b
+
+    c_indices = torch.argsort(torch.flatten(c))
+    a_dist = torch.gather(torch.flatten(a), 0, c_indices)
+    b_dist = torch.gather(torch.flatten(b), 0, c_indices)
+
+    a_dft = torch.fft.rfft(a_dist.float())
+    b_dft = torch.fft.rfft(b_dist.float())
+
+    dft_filter = torch.arange(0, torch.numel(a_dft), device=a_dft.device).float()
+    dft_filter /= torch.numel(a_dft)
+    if beta > EPSILON:
+        dft_filter = (dft_filter - alpha) / beta + 1 / 2
+        dft_filter = torch.clamp(dft_filter, 0.0, 1.0)
+    else:
+        dft_filter = (dft_filter >= alpha).float()
+
+    x_dft = (1 - dft_filter) * a_dft + dft_filter * b_dft
+    x_dist = torch.fft.irfft(x_dft, a_dist.shape[0])
+    x_values = torch.gather(x_dist, 0, torch.argsort(c_indices))
+    return x_values.reshape_as(a)
+
+
+def ties_add_difference(
+    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
+) -> Tensor:
+    deltas = []
+    signs = []
+    for m in [a, b]:
+        deltas.append(filter_top_k(m - c, beta))
+        signs.append(torch.sign(deltas[-1]))
+
+    signs = torch.stack(signs, dim=0)
+    final_sign = torch.sign(torch.sum(signs, dim=0))
+    delta_filters = (signs == final_sign).float()
+
+    res = torch.zeros_like(c, device=c.device)
+    for delta_filter, delta in zip(delta_filters, deltas):
+        res += delta_filter * delta
+
+    param_count = torch.sum(delta_filters, dim=0)
+    return c + alpha * torch.nan_to_num(res / param_count)
+
+
+def filter_top_k(a: Tensor, k: float):
+    k = max(int((1 - k) * torch.numel(a)), 1)
+    k_value, _ = torch.kthvalue(torch.abs(a.flatten()).float(), k)
+    top_k_filter = (torch.abs(a) >= k_value).float()
+    return a * top_k_filter

modules/merging/model.py

@@ -0,0 +1,56 @@
+# import logging
+import os
+from dataclasses import dataclass
+
+import safetensors
+import torch
+from tensordict import TensorDict
+
+# logging.getLogger("sd_meh").addHandler(logging.NullHandler())
+
+
+@dataclass
+class SDModel:
+    model_path: os.PathLike
+    device: str
+
+    def load_model(self):
+        # logging.info(f"Loading: {self.model_path}")
+        if self.model_path.suffix == ".safetensors":
+            ckpt = safetensors.torch.load_file(
+                self.model_path,
+                device=self.device,
+            )
+        else:
+            ckpt = torch.load(self.model_path, map_location=self.device)
+
+        return TensorDict.from_dict(get_state_dict_from_checkpoint(ckpt))
+
+
+# TODO: tidy up
+# from: stable-diffusion-webui/modules/sd_models.py
+def get_state_dict_from_checkpoint(pl_sd):
+    pl_sd = pl_sd.pop("state_dict", pl_sd)
+    pl_sd.pop("state_dict", None)
+    sd = {}
+    for k, v in pl_sd.items():
+        if new_key := transform_checkpoint_dict_key(k):
+            sd[new_key] = v
+
+    pl_sd.clear()
+    pl_sd.update(sd)
+    return pl_sd
+
+
+chckpoint_dict_replacements = {
+    "cond_stage_model.transformer.embeddings.": "cond_stage_model.transformer.text_model.embeddings.",
+    "cond_stage_model.transformer.encoder.": "cond_stage_model.transformer.text_model.encoder.",
+    "cond_stage_model.transformer.final_layer_norm.": "cond_stage_model.transformer.text_model.final_layer_norm.",
+}
+
+
+def transform_checkpoint_dict_key(k):
+    for text, replacement in chckpoint_dict_replacements.items():
+        if k.startswith(text):
+            k = replacement + k[len(text):]
+    return k

modules/merging/presets.py

@@ -0,0 +1,62 @@
+BLOCK_WEIGHTS_PRESETS = {
+    "GRAD_V": [0, 1, 0.9166666667, 0.8333333333, 0.75, 0.6666666667, 0.5833333333, 0.5, 0.4166666667, 0.3333333333, 0.25, 0.1666666667, 0.0833333333, 0, 0.0833333333, 0.1666666667, 0.25, 0.3333333333, 0.4166666667, 0.5, 0.5833333333, 0.6666666667, 0.75, 0.8333333333, 0.9166666667, 1.0],
+    "GRAD_A": [0, 0, 0.0833333333, 0.1666666667, 0.25, 0.3333333333, 0.4166666667, 0.5, 0.5833333333, 0.6666666667, 0.75, 0.8333333333, 0.9166666667, 1.0, 0.9166666667, 0.8333333333, 0.75, 0.6666666667, 0.5833333333, 0.5, 0.4166666667, 0.3333333333, 0.25, 0.1666666667, 0.0833333333, 0],
+    "FLAT_25": [0, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25],
+    "FLAT_75": [0, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75],
+    "WRAP08": [0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
+    "WRAP12": [0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
+    "WRAP14": [0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
+    "WRAP16": [0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
+    "MID12_50": [0, 0, 0, 0, 0, 0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0, 0, 0, 0, 0, 0],
+    "OUT07": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
+    "OUT12": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+    "OUT12_5": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.5, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+    "RING08_SOFT": [0, 0, 0, 0, 0, 0, 0.5, 1, 1, 1, 0.5, 0, 0, 0, 0, 0, 0.5, 1, 1, 1, 0.5, 0, 0, 0, 0, 0],
+    "RING08_5": [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
+    "RING10_5": [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
+    "RING10_3": [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
+    "SMOOTHSTEP": [0, 0, 0.00506365740740741, 0.0196759259259259, 0.04296875, 0.0740740740740741, 0.112123842592593, 0.15625, 0.205584490740741, 0.259259259259259, 0.31640625, 0.376157407407407, 0.437644675925926, 0.5, 0.562355324074074, 0.623842592592592, 0.68359375, 0.740740740740741, 0.794415509259259, 0.84375, 0.887876157407408, 0.925925925925926, 0.95703125, 0.980324074074074, 0.994936342592593, 1],
+    "REVERSE_SMOOTHSTEP": [0, 1, 0.994936342592593, 0.980324074074074, 0.95703125, 0.925925925925926, 0.887876157407407, 0.84375, 0.794415509259259, 0.740740740740741, 0.68359375, 0.623842592592593, 0.562355324074074, 0.5, 0.437644675925926, 0.376157407407408, 0.31640625, 0.259259259259259, 0.205584490740741, 0.15625, 0.112123842592592, 0.0740740740740742, 0.0429687499999996, 0.0196759259259258, 0.00506365740740744, 0],
+    "2SMOOTHSTEP": [0, 0, 0.0101273148148148, 0.0393518518518519, 0.0859375, 0.148148148148148, 0.224247685185185, 0.3125, 0.411168981481482, 0.518518518518519, 0.6328125, 0.752314814814815, 0.875289351851852, 1.0, 0.875289351851852, 0.752314814814815, 0.6328125, 0.518518518518519, 0.411168981481481, 0.3125, 0.224247685185184, 0.148148148148148, 0.0859375, 0.0393518518518512, 0.0101273148148153, 0],
+    "2R_SMOOTHSTEP": [0, 1, 0.989872685185185, 0.960648148148148, 0.9140625, 0.851851851851852, 0.775752314814815, 0.6875, 0.588831018518519, 0.481481481481481, 0.3671875, 0.247685185185185, 0.124710648148148, 0.0, 0.124710648148148, 0.247685185185185, 0.3671875, 0.481481481481481, 0.588831018518519, 0.6875, 0.775752314814816, 0.851851851851852, 0.9140625, 0.960648148148149, 0.989872685185185, 1],
+    "3SMOOTHSTEP": [0, 0, 0.0151909722222222, 0.0590277777777778, 0.12890625, 0.222222222222222, 0.336371527777778, 0.46875, 0.616753472222222, 0.777777777777778, 0.94921875, 0.871527777777778, 0.687065972222222, 0.5, 0.312934027777778, 0.128472222222222, 0.0507812500000004, 0.222222222222222, 0.383246527777778, 0.53125, 0.663628472222223, 0.777777777777778, 0.87109375, 0.940972222222222, 0.984809027777777, 1],
+    "3R_SMOOTHSTEP": [0, 1, 0.984809027777778, 0.940972222222222, 0.87109375, 0.777777777777778, 0.663628472222222, 0.53125, 0.383246527777778, 0.222222222222222, 0.05078125, 0.128472222222222, 0.312934027777778, 0.5, 0.687065972222222, 0.871527777777778, 0.94921875, 0.777777777777778, 0.616753472222222, 0.46875, 0.336371527777777, 0.222222222222222, 0.12890625, 0.0590277777777777, 0.0151909722222232, 0],
+    "4SMOOTHSTEP": [0, 0, 0.0202546296296296, 0.0787037037037037, 0.171875, 0.296296296296296, 0.44849537037037, 0.625, 0.822337962962963, 0.962962962962963, 0.734375, 0.49537037037037, 0.249421296296296, 0.0, 0.249421296296296, 0.495370370370371, 0.734375000000001, 0.962962962962963, 0.822337962962962, 0.625, 0.448495370370369, 0.296296296296297, 0.171875, 0.0787037037037024, 0.0202546296296307, 0],
+    "4R_SMOOTHSTEP": [0, 1, 0.97974537037037, 0.921296296296296, 0.828125, 0.703703703703704, 0.55150462962963, 0.375, 0.177662037037037, 0.0370370370370372, 0.265625, 0.50462962962963, 0.750578703703704, 1.0, 0.750578703703704, 0.504629629629629, 0.265624999999999, 0.0370370370370372, 0.177662037037038, 0.375, 0.551504629629631, 0.703703703703703, 0.828125, 0.921296296296298, 0.979745370370369, 1],
+    "HALF_SMOOTHSTEP": [0, 0, 0.0196759259259259, 0.0740740740740741, 0.15625, 0.259259259259259, 0.376157407407407, 0.5, 0.623842592592593, 0.740740740740741, 0.84375, 0.925925925925926, 0.980324074074074, 1.0, 0.980324074074074, 0.925925925925926, 0.84375, 0.740740740740741, 0.623842592592593, 0.5, 0.376157407407407, 0.259259259259259, 0.15625, 0.0740740740740741, 0.0196759259259259, 0],
+    "HALF_R_SMOOTHSTEP": [0, 1, 0.980324074074074, 0.925925925925926, 0.84375, 0.740740740740741, 0.623842592592593, 0.5, 0.376157407407407, 0.259259259259259, 0.15625, 0.0740740740740742, 0.0196759259259256, 0.0, 0.0196759259259256, 0.0740740740740742, 0.15625, 0.259259259259259, 0.376157407407407, 0.5, 0.623842592592593, 0.740740740740741, 0.84375, 0.925925925925926, 0.980324074074074, 1],
+    "ONE_THIRD_SMOOTHSTEP": [0, 0, 0.04296875, 0.15625, 0.31640625, 0.5, 0.68359375, 0.84375, 0.95703125, 1.0, 0.95703125, 0.84375, 0.68359375, 0.5, 0.31640625, 0.15625, 0.04296875, 0.0, 0.04296875, 0.15625, 0.31640625, 0.5, 0.68359375, 0.84375, 0.95703125, 1],
+    "ONE_THIRD_R_SMOOTHSTEP": [0, 1, 0.95703125, 0.84375, 0.68359375, 0.5, 0.31640625, 0.15625, 0.04296875, 0.0, 0.04296875, 0.15625, 0.31640625, 0.5, 0.68359375, 0.84375, 0.95703125, 1.0, 0.95703125, 0.84375, 0.68359375, 0.5, 0.31640625, 0.15625, 0.04296875, 0],
+    "ONE_FOURTH_SMOOTHSTEP": [0, 0, 0.0740740740740741, 0.259259259259259, 0.5, 0.740740740740741, 0.925925925925926, 1.0, 0.925925925925926, 0.740740740740741, 0.5, 0.259259259259259, 0.0740740740740741, 0.0, 0.0740740740740741, 0.259259259259259, 0.5, 0.740740740740741, 0.925925925925926, 1.0, 0.925925925925926, 0.740740740740741, 0.5, 0.259259259259259, 0.0740740740740741, 0],
+    "ONE_FOURTH_R_SMOOTHSTEP": [0, 1, 0.925925925925926, 0.740740740740741, 0.5, 0.259259259259259, 0.0740740740740742, 0.0, 0.0740740740740742, 0.259259259259259, 0.5, 0.740740740740741, 0.925925925925926, 1.0, 0.925925925925926, 0.740740740740741, 0.5, 0.259259259259259, 0.0740740740740742, 0.0, 0.0740740740740742, 0.259259259259259, 0.5, 0.740740740740741, 0.925925925925926, 1],
+    "COSINE": [0, 1, 0.995722430686905, 0.982962913144534, 0.961939766255643, 0.933012701892219, 0.896676670145617, 0.853553390593274, 0.80438071450436, 0.75, 0.691341716182545, 0.62940952255126, 0.565263096110026, 0.5, 0.434736903889974, 0.37059047744874, 0.308658283817455, 0.25, 0.195619285495639, 0.146446609406726, 0.103323329854382, 0.0669872981077805, 0.0380602337443566, 0.0170370868554658, 0.00427756931309475, 0],
+    "REVERSE_COSINE": [0, 0, 0.00427756931309475, 0.0170370868554659, 0.0380602337443566, 0.0669872981077808, 0.103323329854383, 0.146446609406726, 0.19561928549564, 0.25, 0.308658283817455, 0.37059047744874, 0.434736903889974, 0.5, 0.565263096110026, 0.62940952255126, 0.691341716182545, 0.75, 0.804380714504361, 0.853553390593274, 0.896676670145618, 0.933012701892219, 0.961939766255643, 0.982962913144534, 0.995722430686905, 1],
+    "CUBIC_HERMITE": [0, 0, 0.157576195987654, 0.28491512345679, 0.384765625, 0.459876543209877, 0.512996720679012, 0.546875, 0.564260223765432, 0.567901234567901, 0.560546875, 0.544945987654321, 0.523847415123457, 0.5, 0.476152584876543, 0.455054012345679, 0.439453125, 0.432098765432099, 0.435739776234568, 0.453125, 0.487003279320987, 0.540123456790124, 0.615234375, 0.71508487654321, 0.842423804012347, 1],
+    "REVERSE_CUBIC_HERMITE": [0, 1, 0.842423804012346, 0.71508487654321, 0.615234375, 0.540123456790123, 0.487003279320988, 0.453125, 0.435739776234568, 0.432098765432099, 0.439453125, 0.455054012345679, 0.476152584876543, 0.5, 0.523847415123457, 0.544945987654321, 0.560546875, 0.567901234567901, 0.564260223765432, 0.546875, 0.512996720679013, 0.459876543209876, 0.384765625, 0.28491512345679, 0.157576195987653, 0],
+    "ALL_A": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    "ALL_B": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+}
+
+
+SDXL_BLOCK_WEIGHTS_PRESETS = {
+    "SDXL_GRAD_V": [0, 1.0, 0.888889, 0.777778, 0.666667, 0.555556, 0.444444, 0.333333, 0.222222, 0.111111, 0.0, 0.111111, 0.222222, 0.333333, 0.444444, 0.555556, 0.666667, 0.777778, 0.888889, 1.0],
+    "SDXL_GRAD_A": [0, 0.0, 0.111111, 0.222222, 0.333333, 0.444444, 0.555556, 0.666667, 0.777778, 0.888889, 1.0, 0.888889, 0.777778, 0.666667, 0.555556, 0.444444, 0.333333, 0.222222, 0.111111, 0.0],
+    "SDXL_FLAT_25": [0, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25],
+    "SDXL_FLAT_75": [0, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75, 0.75],
+    "SDXL_WRAP08": [0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
+    "SDXL_WRAP12": [0, 1, 1, 1, 1, 1, 1, 0, 0, 0,  0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
+    "SDXL_WRAP14": [0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
+    "SDXL_OUT07": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
+    "SDXL_SMOOTHSTEP": [0, 0, 0.008916, 0.034294, 0.074074, 0.126200, 0.188615, 0.259259, 0.336077, 0.417010, 0.500000, 0.582990, 0.663923, 0.740741, 0.811385, 0.873800, 0.925926, 0.965706, 0.991084, 1],
+    "SDXL_REVERSE_SMOOTHSTEP": [0, 1, 0.991084, 0.965706, 0.925926, 0.873800, 0.811385, 0.740741, 0.663923, 0.582990, 0.500000, 0.417010, 0.336077, 0.259259, 0.188615, 0.126200, 0.074074, 0.034294, 0.008916, 0],
+    "SDXL_HALF_SMOOTHSTEP": [0, 0, 0.034294, 0.126200, 0.259259, 0.417010, 0.582990, 0.740741, 0.873800, 0.965706, 1, 0.965706, 0.873800, 0.740741, 0.582990, 0.417010, 0.259259, 0.126200, 0.034294, 0],
+    "SDXL_HALF_R_SMOOTHSTEP": [0, 1, 0.965706, 0.873800, 0.740741, 0.582990, 0.417010, 0.259259, 0.126200, 0.034294, 0, 0.034294, 0.126200, 0.259259, 0.417010, 0.582990, 0.740741, 0.873800, 0.965706, 1],
+    "SDXL_ONE_THIRD_SMOOTHSTEP": [0, 0, 0.074074, 0.259259, 0.500000, 0.740741, 0.925926, 1, 0.907407, 0.592593, 0, 0.592593, 0.907407, 1, 0.925926, 0.740741, 0.500000, 0.259259, 0.074074, 0],
+    "SDXL_ONE_THIRD_R_SMOOTHSTEP": [0, 1, 0.925926, 0.740741, 0.500000, 0.259259, 0.074074, 0, 0.092593, 0.407407, 1, 0.407407, 0.092593, 0, 0.074074, 0.259259, 0.500000, 0.740741, 0.925926, 1],
+    "SDXL_COSINE": [0, 1, 0.992404, 0.969846, 0.933013, 0.883022, 0.821394, 0.750000, 0.671010, 0.586824, 0.500000, 0.413176, 0.328990, 0.250000, 0.178606, 0.116978, 0.066987, 0.030154, 0.007596, 0],
+    "SDXL_REVERSE_COSINE": [0, 0, 0.007596, 0.030154, 0.066987, 0.116978, 0.178606, 0.250000, 0.328990, 0.413176, 0.500000, 0.586824, 0.671010, 0.750000, 0.821394, 0.883022, 0.933013, 0.969846, 0.992404, 1],
+    "SDXL_CUBIC_HERMITE": [0, 0, 0.268023, 0.461058, 0.588477, 0.659656, 0.683966, 0.670782, 0.629477, 0.569425, 0.500000, 0.430575, 0.370523, 0.329218, 0.316034, 0.340344, 0.411523, 0.538942, 0.731977, 1],
+    "SDXL_REVERSE_CUBIC_HERMITE": [0, 1, 0.731977, 0.538942, 0.411523, 0.340344, 0.316034, 0.329218, 0.370523, 0.430575, 0.500000, 0.569425, 0.629477, 0.670782, 0.683966, 0.659656, 0.588477, 0.461058, 0.268023, 0],
+    "SDXL_ALL_A": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    "SDXL_ALL_B": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+}

modules/merging/utils.py

@@ -0,0 +1,115 @@
+import inspect
+# import logging
+import re
+from modules.merging import merge_methods
+from modules.merging.presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
+
+BLOCK_WEIGHTS_PRESETS |= SDXL_BLOCK_WEIGHTS_PRESETS
+MERGE_METHODS = dict(inspect.getmembers(merge_methods, inspect.isfunction))
+BETA_METHODS = [
+    name
+    for name, fn in MERGE_METHODS.items()
+    if "beta" in inspect.getfullargspec(fn)[0]
+]
+TRIPLE_METHODS = [
+    name
+    for name, fn in MERGE_METHODS.items()
+    if "c" in inspect.getfullargspec(fn)[0]
+]
+
+
+def interpolate(values, interp_lambda):
+    interpolated = []
+    for i in range(len(values[0])):
+        interpolated.append((1 - interp_lambda) * values[0][i] + interp_lambda * values[1][i])
+    return interpolated
+
+
+class WeightClass:
+    def __init__(self,
+                 model_a,
+                 **kwargs,
+                 ):
+        self.SDXL = True if "model.diffusion_model.middle_block.1.transformer_blocks.9.norm3.weight" in model_a.keys() else False
+        self.NUM_INPUT_BLOCKS = 12 if not self.SDXL else 9
+        self.NUM_MID_BLOCK = 1
+        self.NUM_OUTPUT_BLOCKS = 12 if not self.SDXL else 9
+        self.NUM_TOTAL_BLOCKS = self.NUM_INPUT_BLOCKS + self.NUM_MID_BLOCK + self.NUM_OUTPUT_BLOCKS
+        self.iterations = kwargs.get("iterations", 1)
+        self.it = 0
+        self.re_basin = kwargs.get("re_basin", False)
+        self.ratioDict = {}
+        for key, value in kwargs.items():
+            if isinstance(value, list) or (key.lower() not in ["alpha", "beta"]):
+                self.ratioDict[key.lower()] = value
+            else:
+                self.ratioDict[key.lower()] = [value]
+
+        for key, value in self.ratioDict.items():
+            if key in ["alpha", "beta"]:
+                for i, v in enumerate(value):
+                    if isinstance(v, str) and v.upper() in BLOCK_WEIGHTS_PRESETS.keys():
+                        value[i] = BLOCK_WEIGHTS_PRESETS[v.upper()]
+                    else:
+                        value[i] = [float(x) for x in v.split(",")] if isinstance(v, str) else v
+                        if not isinstance(value[i], list):
+                            value[i] = [value[i]] * (self.NUM_TOTAL_BLOCKS + 1)
+                if len(value) > 1 and isinstance(value[0], list):
+                    self.ratioDict[key] = interpolate(value, self.ratioDict.get(key + "_lambda", 0))
+                else:
+                    self.ratioDict[key] = self.ratioDict[key][0]
+
+
+    def __call__(self, key, it=0):
+        current_bases = {}
+        if self.ratioDict.get("alpha", None):
+            current_bases["alpha"] = self.step_weights_and_bases(self.ratioDict["alpha"], it)
+        if self.ratioDict.get("beta", None):
+            current_bases["beta"] = self.step_weights_and_bases(self.ratioDict["beta"], it)
+
+        weight_index = 0
+        if "model" in key:
+
+            if "model.diffusion_model." in key:
+                weight_index = -1
+
+                re_inp = re.compile(r"\.input_blocks\.(\d+)\.")  # 12
+                re_mid = re.compile(r"\.middle_block\.(\d+)\.")  # 1
+                re_out = re.compile(r"\.output_blocks\.(\d+)\.")  # 12
+
+                if "time_embed" in key:
+                    weight_index = 0  # before input blocks
+                elif ".out." in key:
+                    weight_index = self.NUM_TOTAL_BLOCKS - 1  # after output blocks
+                elif m := re_inp.search(key):
+                    weight_index = int(m.groups()[0])
+                elif re_mid.search(key):
+                    weight_index = self.NUM_INPUT_BLOCKS
+                elif m := re_out.search(key):
+                    weight_index = self.NUM_INPUT_BLOCKS + self.NUM_MID_BLOCK + int(m.groups()[0])
+
+                if weight_index >= self.NUM_TOTAL_BLOCKS:
+                    raise ValueError(f"illegal block index {key}")
+
+        current_bases = {k: w[weight_index] for k, w in current_bases.items()}
+        if self.re_basin:
+            current_bases = self.step_weights_and_bases(current_bases,self.it)
+        return current_bases
+
+    def step_weights_and_bases(self,
+                               ratio,
+                               it: int = 0,
+                               ):
+        new_ratio = {
+            k:
+            1 - (1 - (1 + it) * v / self.iterations) / (1 - it * v / self.iterations)
+            if it > 0
+            else v / self.iterations
+            for k, v in ratio.items()
+        }
+
+        return new_ratio
+
+    def set_it(self, it):
+        self.it = it
+        return

modules/ui_models.py

@@ -10,11 +10,9 @@ from modules.call_queue import wrap_gradio_gpu_call
 from modules.shared import opts, log, req
 import modules.errors
 import modules.hashes
-from sd_meh import merge_methods
-from sd_meh.utils import BETA_METHODS, TRIPLE_METHODS, interpolate
-from sd_meh.presets import BLOCK_WEIGHTS_PRESETS
-
-
+from modules.merging import merge_methods
+from modules.merging.utils import BETA_METHODS, TRIPLE_METHODS, interpolate
+from modules.merging.presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
 
 search_metadata_civit = None
 
@@ -37,14 +35,17 @@ def create_ui():
             with gr.Tab(label="Convert"):
                 with gr.Row():
                     model_name = gr.Dropdown(sd_models.checkpoint_tiles(), label="Original model")
-                    create_refresh_button(model_name, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_tiles()}, "refresh_checkpoint_Z")
+                    create_refresh_button(model_name, sd_models.list_models,
+                                          lambda: {"choices": sd_models.checkpoint_tiles()}, "refresh_checkpoint_Z")
                 with gr.Row():
                     custom_name = gr.Textbox(label="New model name")
                 with gr.Row():
                     precision = gr.Radio(choices=["fp32", "fp16", "bf16"], value="fp16", label="Model precision")
-                    m_type = gr.Radio(choices=["disabled", "no-ema", "ema-only"], value="disabled", label="Model pruning methods")
+                    m_type = gr.Radio(choices=["disabled", "no-ema", "ema-only"], value="disabled",
+                                      label="Model pruning methods")
                 with gr.Row():
-                    checkpoint_formats = gr.CheckboxGroup(choices=["ckpt", "safetensors"], value=["safetensors"], label="Model Format")
+                    checkpoint_formats = gr.CheckboxGroup(choices=["ckpt", "safetensors"], value=["safetensors"],
+                                                          label="Model Format")
                 with gr.Row():
                     show_extra_options = gr.Checkbox(label="Show extra options", value=False)
                     fix_clip = gr.Checkbox(label="Fix clip", value=False)
@@ -81,24 +82,37 @@ def create_ui():
                         with FormRow():
                             def sd_model_choices():
                                 return ['None'] + sd_models.checkpoint_tiles()
+
                             primary_model_name = gr.Dropdown(sd_model_choices(), label="Primary model", value="None")
-                            create_refresh_button(primary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_A")
-                            secondary_model_name = gr.Dropdown(sd_model_choices(), label="Secondary model", value="None")
-                            create_refresh_button(secondary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_B")
+                            create_refresh_button(primary_model_name, sd_models.list_models,
+                                                  lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_A")
+                            secondary_model_name = gr.Dropdown(sd_model_choices(), label="Secondary model",
+                                                               value="None")
+                            create_refresh_button(secondary_model_name, sd_models.list_models,
+                                                  lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_B")
                             tertiary_model_name = gr.Dropdown(sd_model_choices(), label="Tertiary model", value="None")
-                            create_refresh_button(tertiary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_C")
+                            create_refresh_button(tertiary_model_name, sd_models.list_models,
+                                                  lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_C")
                         with FormRow():
-                            interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method")
-                            interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Interpolation ratio from Primary to Secondary', value=0.5)
+                            interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"],
+                                                     value="Weighted sum", label="Interpolation Method")
+                            interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05,
+                                                      label='Interpolation ratio from Primary to Secondary', value=0.5)
                         with FormRow():
-                            checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors", label="Model format")
+                            checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors",
+                                                         label="Model format")
                         with gr.Box():
-                            save_as_half = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision", type="index")
+                            save_as_half = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision",
+                                                    type="index")
                         with FormRow():
-                            config_source = gr.Radio(choices=["Primary", "Secondary", "Tertiary", "None"], value="Primary", label="Model configuration", type="index")
+                            config_source = gr.Radio(choices=["Primary", "Secondary", "Tertiary", "None"],
+                                                     value="Primary", label="Model configuration", type="index")
                         with FormRow():
-                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE")
-                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
+                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None",
+                                                      label="Bake in VAE")
+                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list,
+                                                  lambda: {"choices": ["None"] + list(sd_vae.vae_dict)},
+                                                  "modelmerger_refresh_bake_in_vae")
                         with FormRow():
                             discard_weights = gr.Textbox(value="", label="Discard weights with matching name")
                         with FormRow():
@@ -112,7 +126,8 @@ def create_ui():
                     except Exception as e:
                         modules.errors.display(e, 'model merge')
                         sd_models.list_models()  # to remove the potentially missing models from the list
-                        return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
+                        return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],
+                                f"Error merging checkpoints: {e}"]
                     return results
 
                 modelmerger_merge.click(
@@ -144,103 +159,135 @@ def create_ui():
             with gr.Tab(label="Advanced Merge"):
                 def sd_model_choices():
                     return ['None'] + sd_models.checkpoint_tiles()
+
                 with gr.Row(equal_height=False):
                     with gr.Column(variant='compact'):
                         with FormRow():
                             custom_name = gr.Textbox(label="New model name")
                         with FormRow():
-                            merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum", label="Interpolation Method")
+                            merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum",
+                                                     label="Interpolation Method")
                         with FormRow():
                             primary_model_name = gr.Dropdown(sd_model_choices(), label="Primary model", value="None")
-                            create_refresh_button(primary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_A")
-                            secondary_model_name = gr.Dropdown(sd_model_choices(), label="Secondary model", value="None")
-                            create_refresh_button(secondary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_B")
-                            tertiary_model_name = gr.Dropdown(sd_model_choices(), label="Tertiary model", value="None", visible=False)
-                            tertiary_refresh = create_refresh_button(tertiary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_C",visible=False)
+                            create_refresh_button(primary_model_name, sd_models.list_models,
+                                                  lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_A")
+                            secondary_model_name = gr.Dropdown(sd_model_choices(), label="Secondary model",
+                                                               value="None")
+                            create_refresh_button(secondary_model_name, sd_models.list_models,
+                                                  lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_B")
+                            tertiary_model_name = gr.Dropdown(sd_model_choices(), label="Tertiary model", value="None",
+                                                              visible=False)
+                            tertiary_refresh = create_refresh_button(tertiary_model_name, sd_models.list_models,
+                                                                     lambda: {"choices": sd_model_choices()},
+                                                                     "refresh_checkpoint_C", visible=False)
                         with FormRow():
                             with gr.Tabs() as tabs:
                                 with gr.TabItem(label="Simple Merge", id=0):
                                     with FormRow():
-                                        alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Alpha Ratio', value=0.5)
-                                        beta = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Beta Ratio', value=None, visible=False)
+                                        alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Alpha Ratio',
+                                                          value=0.5)
+                                        beta = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Beta Ratio',
+                                                         value=None, visible=False)
                                 with gr.TabItem(label="Preset Block Merge", id=1):
                                     with FormRow():
-                                        alpha_preset = gr.Dropdown(choices=["None"]+list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="ALPHA Block Weight Preset", multiselect=True, max_choices=2)
-                                        alpha_preset_lambda = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Preset Interpolation Ratio', value=None, visible=False)
-                                        apply_preset = ToolButton('⇩', visible=True)
+                                        sdxl = gr.Checkbox(label="SDXL")
                                     with FormRow():
-                                        beta_preset = gr.Dropdown(choices=["None"]+list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="BETA Block Weight Preset", multiselect=True, max_choices=2, interactive=True, visible=False)
-                                        beta_preset_lambda = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Preset Interpolation Ratio', value=None, interactive=True, visible=False)
-                                        beta_apply_preset = ToolButton('⇩', interactive=True, visible=False)
+                                        alpha_preset = gr.Dropdown(
+                                            choices=["None"] + list(BLOCK_WEIGHTS_PRESETS.keys()), value=None,
+                                            label="ALPHA Block Weight Preset", multiselect=True, max_choices=2)
+                                        alpha_preset_lambda = gr.Slider(minimum=0.0, maximum=1.0, step=0.05,
+                                                                        label='Preset Interpolation Ratio', value=None,
+                                                                        visible=False)
+                                        apply_preset = ToolButton('⇨', visible=True)
+                                    with FormRow():
+                                        beta_preset = gr.Dropdown(choices=["None"] + list(BLOCK_WEIGHTS_PRESETS.keys()),
+                                                                  value=None, label="BETA Block Weight Preset",
+                                                                  multiselect=True, max_choices=2, interactive=True,
+                                                                  visible=False)
+                                        beta_preset_lambda = gr.Slider(minimum=0.0, maximum=1.0, step=0.05,
+                                                                       label='Preset Interpolation Ratio', value=None,
+                                                                       interactive=True, visible=False)
+                                        beta_apply_preset = ToolButton('⇨', interactive=True, visible=False)
                                 with gr.TabItem(label="Manual Block Merge", id=2):
                                     with FormRow():
                                         alpha_label = gr.Markdown("# Alpha")
                                     with FormRow():
                                         alpha_base = gr.Textbox(value=None, label="Base", min_width=70, scale=1)
                                         alpha_in_blocks = gr.Textbox(value=None, label="In Blocks", scale=15)
-                                        alpha_mid_block = gr.Textbox(value=None, label="Mid Block", min_width=80, scale=1)
+                                        alpha_mid_block = gr.Textbox(value=None, label="Mid Block", min_width=80,
+                                                                     scale=1)
                                         alpha_out_blocks = gr.Textbox(value=None, label="Out Block", scale=15)
                                     with FormRow():
                                         beta_label = gr.Markdown("# Beta", visible=False)
                                     with FormRow():
-                                        beta_base = gr.Textbox(value=None, label="Base", min_width=70, scale=1, interactive=True, visible=False)
-                                        beta_in_blocks = gr.Textbox(value=None, label="In Blocks", interactive=True, scale=15, visible=False)
-                                        beta_mid_block = gr.Textbox(value=None, label="Mid Block", min_width=80, interactive=True, scale=1, visible=False)
-                                        beta_out_blocks = gr.Textbox(value=None, label="Out Block", interactive=True, scale=15, visible=False)
+                                        beta_base = gr.Textbox(value=None, label="Base", min_width=70, scale=1,
+                                                               interactive=True, visible=False)
+                                        beta_in_blocks = gr.Textbox(value=None, label="In Blocks", interactive=True,
+                                                                    scale=15, visible=False)
+                                        beta_mid_block = gr.Textbox(value=None, label="Mid Block", min_width=80,
+                                                                    interactive=True, scale=1, visible=False)
+                                        beta_out_blocks = gr.Textbox(value=None, label="Out Block", interactive=True,
+                                                                     scale=15, visible=False)
                         with FormRow():
                             weights_clip = gr.Checkbox(label="Weights Clip")
-                            prune = gr.Checkbox(label="Prune")
+                            prune = gr.Checkbox(label="Prune", value=True, visible=False)
                             re_basin = gr.Checkbox(label="ReBasin")
                         with FormRow():
-                            re_basin_iterations = gr.Slider(minimum=0, maximum=25, step=1, label='Number of ReBasin Iterations', value=None, visible=False)
+                            re_basin_iterations = gr.Slider(minimum=0, maximum=25, step=1,
+                                                            label='Number of ReBasin Iterations', value=None,
+                                                            visible=False)
                         with FormRow():
-                            checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors", visible=False, label="Model format")
+                            checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors",
+                                                         visible=False, label="Model format")
                         with FormRow():
                             precision = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision")
                         with FormRow():
                             device = gr.Radio(choices=["cpu", "cuda"], value="cpu", label="Device")
                         with FormRow():
-                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None",  interactive=True, label="Bake in VAE")
-                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
+                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None",
+                                                      interactive=True, label="Bake in VAE")
+                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list,
+                                                  lambda: {"choices": ["None"] + list(sd_vae.vae_dict)},
+                                                  "modelmerger_refresh_bake_in_vae")
                         with FormRow():
                             save_metadata = gr.Checkbox(value=True, label="Save metadata")
                         with gr.Row():
                             MEHmodelmerger_merge = gr.Button(value="Merge", variant='primary')
 
                 def MEHmodelmerger(dummy_component,
-                        primary_model_name,
-                        secondary_model_name,
-                        tertiary_model_name,
-                        merge_mode,
-                        alpha,
-                        beta,
-                        alpha_preset,
-                        alpha_preset_lambda,
-                        alpha_base,
-                        alpha_in_blocks,
-                        alpha_mid_block,
-                        alpha_out_blocks,
-                        beta_preset,
-                        beta_preset_lambda,
-                        beta_base,
-                        beta_in_blocks,
-                        beta_mid_block,
-                        beta_out_blocks,
-                        precision,
-                        custom_name,
-                        checkpoint_format,
-                        save_metadata,
-                        weights_clip,
-                        prune,
-                        re_basin,
-                        re_basin_iterations,
-                        device,
-                        bake_in_vae):
+                                   primary_model_name,
+                                   secondary_model_name,
+                                   tertiary_model_name,
+                                   merge_mode,
+                                   alpha,
+                                   beta,
+                                   alpha_preset,
+                                   alpha_preset_lambda,
+                                   alpha_base,
+                                   alpha_in_blocks,
+                                   alpha_mid_block,
+                                   alpha_out_blocks,
+                                   beta_preset,
+                                   beta_preset_lambda,
+                                   beta_base,
+                                   beta_in_blocks,
+                                   beta_mid_block,
+                                   beta_out_blocks,
+                                   precision,
+                                   custom_name,
+                                   checkpoint_format,
+                                   save_metadata,
+                                   weights_clip,
+                                   prune,
+                                   re_basin,
+                                   re_basin_iterations,
+                                   device,
+                                   bake_in_vae):
                     kwargs = {}
                     for x in inspect.getfullargspec(MEHmodelmerger)[0]:
                         kwargs[x] = locals()[x]
                     for key in list(kwargs.keys()):
-                        if kwargs[key] in [None,"None","",0,[]]:
+                        if kwargs[key] in [None, "None", "", 0, []]:
                             del kwargs[key]
 
                     # return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"{kwargs}"]
@@ -250,7 +297,8 @@ def create_ui():
                     except Exception as e:
                         modules.errors.display(e, 'model merge')
                         sd_models.list_models()  # to remove the potentially missing models from the list
-                        return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
+                        return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],
+                                f"Error merging checkpoints: {e}"]
                     return results
 
                 def tertiary(mode):
@@ -264,18 +312,20 @@ def create_ui():
                         return [gr.update(visible=True) for _ in range(9)]
                     else:
                         return [gr.update(visible=False) for _ in range(9)]
+
                 def show_iters(show):
                     if show:
                         return gr.Slider.update(value=5, visible=True)
                     else:
                         return gr.Slider.update(value=None, visible=False)
+
                 def preset_visiblility(x):
                     if len(x) == 2:
                         return gr.Slider.update(value=0.5, visible=True)
                     else:
                         return gr.Slider.update(value=None, visible=False)
 
-                def load_presets(presets,ratio):
+                def load_presets(presets, ratio):
                     for i, p in enumerate(presets):
                         presets[i] = BLOCK_WEIGHTS_PRESETS[p]
                     if len(presets) == 2:
@@ -283,16 +333,28 @@ def create_ui():
                     else:
                         preset = presets[0]
                     preset = ['%.3f' % x for x in preset]
-                    preset = [preset[0], ",".join(preset[1:13]),preset[13], ",".join(preset[14:])]
-                    return [gr.update(value=x) for x in preset]+[gr.update(selected=2)]
+                    preset = [preset[0], ",".join(preset[1:13]), preset[13], ",".join(preset[14:])]
+                    return [gr.update(value=x) for x in preset] + [gr.update(selected=2)]
 
+                def preset_choices(sdxl):
+                    if sdxl:
+                        return [gr.update(choices=["None"] + list(SDXL_BLOCK_WEIGHTS_PRESETS.keys())) for _ in range(2)]
+                    else:
+                        return [gr.update(choices=["None"] + list(BLOCK_WEIGHTS_PRESETS.keys())) for _ in range(2)]
+
+                sdxl.change(fn=preset_choices, inputs=sdxl, outputs=[alpha_preset, beta_preset])
                 alpha_preset.change(fn=preset_visiblility, inputs=alpha_preset, outputs=alpha_preset_lambda)
                 beta_preset.change(fn=preset_visiblility, inputs=alpha_preset, outputs=beta_preset_lambda)
                 merge_mode.input(fn=tertiary, inputs=merge_mode, outputs=[tertiary_model_name, tertiary_refresh])
-                merge_mode.input(fn=beta_visibility, inputs=merge_mode, outputs=[beta, alpha_label, beta_label, beta_apply_preset, beta_preset, beta_base, beta_in_blocks, beta_mid_block, beta_out_blocks])
-                re_basin.change(fn=show_iters, inputs=re_basin,outputs=re_basin_iterations)
-                apply_preset.click(fn=load_presets,inputs=[alpha_preset, alpha_preset_lambda], outputs=[alpha_base,alpha_in_blocks,alpha_mid_block,alpha_out_blocks,tabs])
-                beta_apply_preset.click(fn=load_presets,inputs=[beta_preset, beta_preset_lambda], outputs=[beta_base,beta_in_blocks,beta_mid_block,beta_out_blocks,tabs])
+                merge_mode.input(fn=beta_visibility, inputs=merge_mode,
+                                 outputs=[beta, alpha_label, beta_label, beta_apply_preset, beta_preset, beta_base,
+                                          beta_in_blocks, beta_mid_block, beta_out_blocks])
+                re_basin.change(fn=show_iters, inputs=re_basin, outputs=re_basin_iterations)
+                apply_preset.click(fn=load_presets, inputs=[alpha_preset, alpha_preset_lambda],
+                                   outputs=[alpha_base, alpha_in_blocks, alpha_mid_block, alpha_out_blocks, tabs])
+                beta_apply_preset.click(fn=load_presets, inputs=[beta_preset, beta_preset_lambda],
+                                        outputs=[beta_base, beta_in_blocks, beta_mid_block, beta_out_blocks, tabs])
+
                 MEHmodelmerger_merge.click(
                     fn=wrap_gradio_gpu_call(MEHmodelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
                     _js='modelmerger',
@@ -346,14 +408,14 @@ def create_ui():
                     model_checkhash_btn.click(fn=sd_models.update_model_hashes, inputs=[], outputs=[models_outcome])
                 with gr.Row():
                     model_table = gr.DataFrame(
-                        value = None,
-                        headers = model_headers,
-                        label = 'Model data',
-                        show_label = True,
-                        interactive = False,
-                        wrap = True,
-                        overflow_row_behaviour = 'paginate',
-                        max_rows = 50,
+                        value=None,
+                        headers=model_headers,
+                        label='Model data',
+                        show_label=True,
+                        interactive=False,
+                        wrap=True,
+                        overflow_row_behaviour='paginate',
+                        max_rows=50,
                     )
 
                 def list_models():
@@ -392,11 +454,15 @@ def create_ui():
                         # task='text-to-image',
                         library=['diffusers'],
                     )
-                    models = hf_api.list_models(filter=model_filter, full=True, limit=50, sort="downloads", direction=-1)
+                    models = hf_api.list_models(filter=model_filter, full=True, limit=50, sort="downloads",
+                                                direction=-1)
                     data.clear()
                     for model in models:
-                        tags = [t for t in model.tags if not t.startswith('diffusers') and not t.startswith('license') and not t.startswith('arxiv') and len(t) > 2]
-                        data.append([model.modelId, model.pipeline_tag, tags, model.downloads, model.lastModified, f'https://huggingface.co/{model.modelId}'])
+                        tags = [t for t in model.tags if
+                                not t.startswith('diffusers') and not t.startswith('license') and not t.startswith(
+                                    'arxiv') and len(t) > 2]
+                        data.append([model.modelId, model.pipeline_tag, tags, model.downloads, model.lastModified,
+                                     f'https://huggingface.co/{model.modelId}'])
                     return data
 
                 def hf_select(evt: gr.SelectData, data):
@@ -404,8 +470,9 @@ def create_ui():
 
                 def hf_download_model(hub_id: str, token, variant, revision, mirror, custom_pipeline):
                     from modules.modelloader import download_diffusers_model
-                    download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror, custom_pipeline=custom_pipeline)
-                    from modules.sd_models import list_models # pylint: disable=W0621
+                    download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token, variant=variant,
+                                             revision=revision, mirror=mirror, custom_pipeline=custom_pipeline)
+                    from modules.sd_models import list_models  # pylint: disable=W0621
                     list_models()
                     log.info(f'Diffuser model downloaded: model="{hub_id}"')
                     return f'Diffuser model downloaded: model="{hub_id}"'
@@ -413,20 +480,25 @@ def create_ui():
                 with gr.Column(scale=6):
                     gr.HTML('<h2>Search for models</h2>Select a model from the search results to download<br><br>')
                     with gr.Row():
-                        hf_search_text = gr.Textbox('', label = 'Search models', placeholder='search huggingface models')
+                        hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models')
                         hf_search_btn = ToolButton(value="🔍", label="Search")
                     with gr.Row():
                         with gr.Column(scale=2):
                             with gr.Row():
-                                hf_selected = gr.Textbox('', label = 'Select model', placeholder='select model from search results or enter model name manually')
+                                hf_selected = gr.Textbox('', label='Select model',
+                                                         placeholder='select model from search results or enter model name manually')
                         with gr.Column(scale=1):
                             with gr.Row():
-                                hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label = 'Specify model variant', placeholder='')
-                                hf_revision = gr.Textbox('', label = 'Specify model revision', placeholder='')
+                                hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label='Specify model variant',
+                                                        placeholder='')
+                                hf_revision = gr.Textbox('', label='Specify model revision', placeholder='')
                     with gr.Row():
-                        hf_token = gr.Textbox('', label = 'Huggingface token', placeholder='optional access token for private or gated models')
-                        hf_mirror = gr.Textbox('', label = 'Huggingface mirror', placeholder='optional mirror site for downloads')
-                        hf_custom_pipeline = gr.Textbox('', label = 'Custom pipeline', placeholder='optional pipeline for downloads')
+                        hf_token = gr.Textbox('', label='Huggingface token',
+                                              placeholder='optional access token for private or gated models')
+                        hf_mirror = gr.Textbox('', label='Huggingface mirror',
+                                               placeholder='optional mirror site for downloads')
+                        hf_custom_pipeline = gr.Textbox('', label='Custom pipeline',
+                                                        placeholder='optional pipeline for downloads')
                 with gr.Column(scale=1):
                     gr.HTML('<br>')
                     hf_download_model_btn = gr.Button(value="Download model", variant='primary')
@@ -434,12 +506,16 @@ def create_ui():
                 with gr.Row():
                     hf_headers = ['Name', 'Pipeline', 'Tags', 'Downloads', 'Updated', 'URL']
                     hf_types = ['str', 'str', 'str', 'number', 'date', 'markdown']
-                    hf_results = gr.DataFrame(None, label = 'Search results', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = hf_headers, datatype = hf_types, type='array')
+                    hf_results = gr.DataFrame(None, label='Search results', show_label=True, interactive=False,
+                                              wrap=True, overflow_row_behaviour='paginate', max_rows=10,
+                                              headers=hf_headers, datatype=hf_types, type='array')
 
                 hf_search_text.submit(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results])
                 hf_search_btn.click(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results])
                 hf_results.select(fn=hf_select, inputs=[hf_results], outputs=[hf_selected])
-                hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_custom_pipeline], outputs=[models_outcome])
+                hf_download_model_btn.click(fn=hf_download_model,
+                                            inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror,
+                                                    hf_custom_pipeline], outputs=[models_outcome])
 
             with gr.Tab(label="CivitAI"):
                 data = []
@@ -482,7 +558,8 @@ def create_ui():
                                 model['stats']['rating']
                             ])
                     res = f'Search result: name={name} tag={tag or "none"} type={model_type} models={len(data1)}'
-                    return res, gr.update(visible=len(data1) > 0, value=data1 if len(data1) > 0 else []), gr.update(visible=False, value=None), gr.update(visible=False, value=None)
+                    return res, gr.update(visible=len(data1) > 0, value=data1 if len(data1) > 0 else []), gr.update(
+                        visible=False, value=None), gr.update(visible=False, value=None)
 
                 def civit_select1(evt: gr.SelectData, in_data):
                     model_id = in_data[evt.index[0]][0]
@@ -491,7 +568,8 @@ def create_ui():
                     for model in data:
                         if model['id'] == model_id:
                             for d in model['modelVersions']:
-                                if d.get('images') is not None and len(d['images']) > 0 and len(d['images'][0]['url']) > 0:
+                                if d.get('images') is not None and len(d['images']) > 0 and len(
+                                        d['images'][0]['url']) > 0:
                                     preview_img = d['images'][0]['url']
                                 data2.append([
                                     d['id'],
@@ -513,7 +591,8 @@ def create_ui():
                                 if variant['id'] == variant_id:
                                     for f in variant['files']:
                                         try:
-                                            if os.path.splitext(f['name'])[1].lower() in ['.safetensors', '.ckpt', '.pt', '.pth', '.bin']:
+                                            if os.path.splitext(f['name'])[1].lower() in ['.safetensors', '.ckpt',
+                                                                                          '.pt', '.pth', '.bin']:
                                                 data3.append([
                                                     f['name'],
                                                     round(f['sizeKB']),
@@ -529,7 +608,8 @@ def create_ui():
                     log.debug(f'CivitAI select: variant={in_data[evt.index[0]]}')
                     return in_data[evt.index[0]][3], in_data[evt.index[0]][0], gr.update(interactive=True)
 
-                def civit_download_model(model_url: str, model_name: str, model_path: str, model_type: str, image_url: str):
+                def civit_download_model(model_url: str, model_name: str, model_path: str, model_type: str,
+                                         image_url: str):
                     if model_url is None or len(model_url) == 0:
                         return 'No model selected'
                     try:
@@ -539,7 +619,7 @@ def create_ui():
                         res = f"CivitAI model downloaded error: model={model_url} {e}"
                         log.error(res)
                         return res
-                    from modules.sd_models import list_models # pylint: disable=W0621
+                    from modules.sd_models import list_models  # pylint: disable=W0621
                     list_models()
                     return res
 
@@ -556,12 +636,14 @@ def create_ui():
                             continue
                         for item in page.list_items():
                             meta = os.path.splitext(item['filename'])[0] + '.json'
-                            if ('card-no-preview.png' in item['preview'] or not os.path.isfile(meta)) and os.path.isfile(item['filename']):
+                            if ('card-no-preview.png' in item['preview'] or not os.path.isfile(
+                                    meta)) and os.path.isfile(item['filename']):
                                 sha = item.get('hash', None)
                                 found = False
                                 if sha is not None and len(sha) > 0:
                                     r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}')
-                                    log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}')
+                                    log.debug(
+                                        f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}')
                                     if r.status_code == 200:
                                         d = r.json()
                                         res.append(download_civit_meta(item['filename'], d['modelId']))
@@ -573,10 +655,12 @@ def create_ui():
                                                 if 'error' not in img_res:
                                                     found = True
                                                     break
-                                if not found and civit_previews_rehash and os.stat(item['filename']).st_size < (1024 * 1024 * 1024):
+                                if not found and civit_previews_rehash and os.stat(item['filename']).st_size < (
+                                        1024 * 1024 * 1024):
                                     sha = modules.hashes.calculate_sha256(item['filename'], quiet=True)[:10]
                                     r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}')
-                                    log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}')
+                                    log.debug(
+                                        f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}')
                                     if r.status_code == 200:
                                         d = r.json()
                                         res.append(download_civit_meta(item['filename'], d['modelId']))
@@ -591,11 +675,12 @@ def create_ui():
                     txt = '<br>'.join([r for r in res if len(r) > 0])
                     return txt
 
-                global search_metadata_civit # pylint: disable=global-statement
+                global search_metadata_civit  # pylint: disable=global-statement
                 search_metadata_civit = civit_search_metadata
 
                 with gr.Row():
-                    gr.HTML('<h2>Fetch information</h2>Fetches preview and metadata information for all models with missing information<br>Models with existing previews and information are not updated<br>')
+                    gr.HTML(
+                        '<h2>Fetch information</h2>Fetches preview and metadata information for all models with missing information<br>Models with existing previews and information are not updated<br>')
                 with gr.Row():
                     civit_previews_btn = gr.Button(value="Start", variant='primary')
                 with gr.Row():
@@ -605,11 +690,12 @@ def create_ui():
                     gr.HTML('<h2>Search for models</h2>')
                 with gr.Row():
                     with gr.Column(scale=1):
-                        civit_model_type = gr.Dropdown(label='Model type', choices=['SD 1.5', 'SD XL', 'LoRA', 'Other'], value='LoRA')
+                        civit_model_type = gr.Dropdown(label='Model type', choices=['SD 1.5', 'SD XL', 'LoRA', 'Other'],
+                                                       value='LoRA')
                     with gr.Column(scale=15):
                         with gr.Row():
-                            civit_search_text = gr.Textbox('', label = 'Search models', placeholder='keyword')
-                            civit_search_tag = gr.Textbox('', label = '', placeholder='tags')
+                            civit_search_text = gr.Textbox('', label='Search models', placeholder='keyword')
+                            civit_search_tag = gr.Textbox('', label='', placeholder='tags')
                             civit_search_btn = ToolButton(value="🔍", label="Search", interactive=False)
                         with gr.Row():
                             civit_search_res = gr.HTML('')
@@ -617,39 +703,64 @@ def create_ui():
                     gr.HTML('<h2>Download model</h2>')
                 with gr.Row():
                     civit_download_model_btn = gr.Button(value="Download", variant='primary')
-                    gr.HTML('<span style="line-height: 2em">Select a model, model version and and model variant from the search results to download or enter model URL manually</span><br>')
+                    gr.HTML(
+                        '<span style="line-height: 2em">Select a model, model version and and model variant from the search results to download or enter model URL manually</span><br>')
                 with gr.Row():
-                    civit_name = gr.Textbox('', label = 'Model name', placeholder='select model from search results', visible=True)
-                    civit_selected = gr.Textbox('', label = 'Model URL', placeholder='select model from search results', visible=True)
-                    civit_path = gr.Textbox('', label = 'Download path', placeholder='optional subfolder path where to save model', visible=True)
+                    civit_name = gr.Textbox('', label='Model name', placeholder='select model from search results',
+                                            visible=True)
+                    civit_selected = gr.Textbox('', label='Model URL', placeholder='select model from search results',
+                                                visible=True)
+                    civit_path = gr.Textbox('', label='Download path',
+                                            placeholder='optional subfolder path where to save model', visible=True)
                 with gr.Row():
                     gr.HTML('<h2>Search results</h2>')
                 with gr.Row():
                     civit_headers1 = ['ID', 'Name', 'Tags', 'Downloads', 'Rating']
                     civit_types1 = ['number', 'str', 'str', 'number', 'number']
-                    civit_results1 = gr.DataFrame(value = None, label = None, show_label = False, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers1, datatype = civit_types1, type='array', visible=False)
+                    civit_results1 = gr.DataFrame(value=None, label=None, show_label=False, interactive=False,
+                                                  wrap=True, overflow_row_behaviour='paginate', max_rows=10,
+                                                  headers=civit_headers1, datatype=civit_types1, type='array',
+                                                  visible=False)
                 with gr.Row():
                     with gr.Column():
                         civit_headers2 = ['ID', 'ModelID', 'Name', 'Base', 'Created', 'Preview']
                         civit_types2 = ['number', 'number', 'str', 'str', 'date', 'str']
-                        civit_results2 = gr.DataFrame(value = None, label = 'Model versions', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers2, datatype = civit_types2, type='array', visible=False)
+                        civit_results2 = gr.DataFrame(value=None, label='Model versions', show_label=True,
+                                                      interactive=False, wrap=True, overflow_row_behaviour='paginate',
+                                                      max_rows=10, headers=civit_headers2, datatype=civit_types2,
+                                                      type='array', visible=False)
                     with gr.Column():
                         civit_headers3 = ['Name', 'Size', 'Metadata', 'URL']
                         civit_types3 = ['str', 'number', 'str', 'str']
-                        civit_results3 = gr.DataFrame(value = None, label = 'Model variants', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers3, datatype = civit_types3, type='array', visible=False)
+                        civit_results3 = gr.DataFrame(value=None, label='Model variants', show_label=True,
+                                                      interactive=False, wrap=True, overflow_row_behaviour='paginate',
+                                                      max_rows=10, headers=civit_headers3, datatype=civit_types3,
+                                                      type='array', visible=False)
 
                 def is_visible(component):
                     visible = len(component) > 0 if component is not None else False
                     return gr.update(visible=visible)
 
-                civit_search_text.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
-                civit_search_tag.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
-                civit_search_btn.click(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
-                civit_results1.select(fn=civit_select1, inputs=[civit_results1], outputs=[civit_results2, civit_results3, models_image])
+                civit_search_text.submit(fn=civit_search_model,
+                                         inputs=[civit_search_text, civit_search_tag, civit_model_type],
+                                         outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
+                civit_search_tag.submit(fn=civit_search_model,
+                                        inputs=[civit_search_text, civit_search_tag, civit_model_type],
+                                        outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
+                civit_search_btn.click(fn=civit_search_model,
+                                       inputs=[civit_search_text, civit_search_tag, civit_model_type],
+                                       outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
+                civit_results1.select(fn=civit_select1, inputs=[civit_results1],
+                                      outputs=[civit_results2, civit_results3, models_image])
                 civit_results2.select(fn=civit_select2, inputs=[civit_results2], outputs=[civit_results3])
-                civit_results3.select(fn=civit_select3, inputs=[civit_results3], outputs=[civit_selected, civit_name, civit_search_btn])
+                civit_results3.select(fn=civit_select3, inputs=[civit_results3],
+                                      outputs=[civit_selected, civit_name, civit_search_btn])
                 civit_results1.change(fn=is_visible, inputs=[civit_results1], outputs=[civit_results1])
                 civit_results2.change(fn=is_visible, inputs=[civit_results2], outputs=[civit_results2])
                 civit_results3.change(fn=is_visible, inputs=[civit_results3], outputs=[civit_results3])
-                civit_download_model_btn.click(fn=civit_download_model, inputs=[civit_selected, civit_name, civit_path, civit_model_type, models_image], outputs=[models_outcome])
-                civit_previews_btn.click(fn=civit_search_metadata, inputs=[civit_previews_rehash, civit_previews_rehash], outputs=[models_outcome])
+                civit_download_model_btn.click(fn=civit_download_model,
+                                               inputs=[civit_selected, civit_name, civit_path, civit_model_type,
+                                                       models_image], outputs=[models_outcome])
+                civit_previews_btn.click(fn=civit_search_metadata,
+                                         inputs=[civit_previews_rehash, civit_previews_rehash],
+                                         outputs=[models_outcome])

requirements.txt

@@ -33,6 +33,7 @@ rich
 safetensors
 scipy
 tb_nightly
+tensordict
 toml
 torchdiffeq
 voluptuous