From 39a37597b5be531322fcfb7f2a565d4fbaf1deb4 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 3 Nov 2023 14:55:57 -0500
Subject: [PATCH 01/25] Further UI Work

---
 modules/extras.py    | 106 +++++++++++++++++++++++++++++++
 modules/ui_common.py |   4 +-
 modules/ui_models.py | 147 ++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 254 insertions(+), 3 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index a84b9b28e..0d5fc2c76 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -8,6 +8,7 @@ import torch
 import tqdm
 import gradio as gr
 import safetensors.torch
+from sd_meh.merge import merge_models
 
 from modules import shared, images, sd_models, sd_vae, sd_models_config
 
@@ -239,6 +240,111 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     shared.state.end()
     return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname]
 
+
+def run_MEHmodelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, merge_mode,base_alpha,
+                       base_beta, weights_alpha,
+                       weights_beta, precision, custom_name, checkpoint_format, save_metadata, preset, weights_clip, prune,
+                       re_basin, re_basin_iterations, device):  # pylint: disable=unused-argument
+    shared.state.begin('model-merge')
+    models = {
+            "model_a": sd_models.checkpoints_list[primary_model_name].filename,
+            "model_b": sd_models.checkpoints_list[secondary_model_name].filename,
+            }
+    if tertiary_model_name is not None:
+        models |= {"model_c": sd_models.checkpoints_list[tertiary_model_name].filename}
+    work_device = device
+    threads = 1
+    preset = preset if preset != "None" else None
+    block_weights_preset_alpha = block_weights_preset_beta = block_weights_preset_alpha_b = block_weights_preset_beta_b = preset if preset is not None else None
+    presets_alpha_lambda = presets_beta_lambda = None
+    logging_level = "INFO"
+
+
+    def fail(message):
+        shared.state.textinfo = message
+        shared.state.end()
+        return [*[gr.update() for _ in range(4)], message]
+
+
+
+    theta_0 = main(
+            model_a,
+            model_b,
+            model_c,
+            merge_mode,
+            weights_clip,
+            precision,
+            str(weights_alpha),
+            base_alpha,
+            str(weights_beta),
+            base_beta,
+            re_basin,
+            re_basin_iterations,
+            device,
+            work_device,
+            prune,
+            block_weights_preset_alpha,
+            block_weights_preset_beta,
+            threads,
+            block_weights_preset_alpha_b,
+            block_weights_preset_beta_b,
+            presets_alpha_lambda,
+            presets_beta_lambda,
+            logging_level,
+            )
+    ckpt_dir = shared.opts.ckpt_dir or sd_models.model_path
+    filename = custom_name
+    filename += "." + checkpoint_format
+    output_modelname = os.path.join(ckpt_dir, filename)
+    shared.state.textinfo = "Saving"
+    metadata = None
+    if save_metadata:
+        metadata = {"format": "pt", "sd_merge_models": {}}
+        merge_recipe = {
+            "type": "webui",  # 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,
+        }
+        metadata["sd_merge_recipe"] = json.dumps(merge_recipe)
+
+        def add_model_metadata(checkpoint_info):
+            checkpoint_info.calculate_shorthash()
+            metadata["sd_merge_models"][checkpoint_info.sha256] = {
+                "name": checkpoint_info.name,
+                "legacy_hash": checkpoint_info.hash,
+                "sd_merge_recipe": checkpoint_info.metadata.get("sd_merge_recipe", None)
+            }
+            metadata["sd_merge_models"].update(checkpoint_info.metadata.get("sd_merge_models", {}))
+
+        add_model_metadata(primary_model_info)
+        if secondary_model_info:
+            add_model_metadata(secondary_model_info)
+        if tertiary_model_info:
+            add_model_metadata(tertiary_model_info)
+        metadata["sd_merge_models"] = json.dumps(metadata["sd_merge_models"])
+
+    _, extension = os.path.splitext(output_modelname)
+    if extension.lower() == ".safetensors":
+        safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
+    else:
+        torch.save(theta_0, output_modelname)
+
+    sd_models.list_models()
+    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()
+    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],
+            "Checkpoint saved to " + output_modelname]
+
 def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_name, unet_conv, text_encoder_conv, vae_conv, others_conv, fix_clip):
 
     # position_ids in clip is int64. model_ema.num_updates is int32
diff --git a/modules/ui_common.py b/modules/ui_common.py
index c5d1fe62e..8b4733b73 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -231,7 +231,7 @@ def create_output_panel(tabname):
             return result_gallery, generation_info, html_info, html_info_formatted, html_log
 
 
-def create_refresh_button(refresh_component, refresh_method, refreshed_args, elem_id):
+def create_refresh_button(refresh_component, refresh_method, refreshed_args, elem_id, visible: bool = True):
 
     def refresh():
         refresh_method()
@@ -241,7 +241,7 @@ def create_refresh_button(refresh_component, refresh_method, refreshed_args, ele
         return gr.update(**(args or {}))
 
     from modules.ui_components import ToolButton
-    refresh_button = ToolButton(value=symbols.refresh, elem_id=elem_id)
+    refresh_button = ToolButton(value=symbols.refresh, elem_id=elem_id, visible=visible)
     refresh_button.click(fn=refresh, inputs=[], outputs=[refresh_component])
     return refresh_button
 
diff --git a/modules/ui_models.py b/modules/ui_models.py
index e84300963..e5115662c 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -3,12 +3,16 @@ import json
 from datetime import datetime
 import gradio as gr
 from modules import sd_models, sd_vae, extras
-from modules.ui_components import FormRow, ToolButton
+from modules.ui_components import FormRow, ToolButton, InputAccordion
 from modules.ui_common import create_refresh_button
 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
+
 
 
 search_metadata_civit = None
@@ -136,6 +140,147 @@ def create_ui():
                         models_outcome,
                     ]
                 )
+            with gr.Tab(label="MEH 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")
+                        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)
+                        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)
+                        with InputAccordion(False, label="Block Merge", elem_id=f"block_merge") as block_accordion:
+                            with FormRow():
+                                alpha_label = gr.Markdown("# Alpha")
+                            with FormRow():
+                                preset = gr.Dropdown(choices=["None"]+list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="Block Weight Preset", multiselect=True, max_choices=2)
+                                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():
+                                base = gr.Textbox(value=None, label="Base", scale=1)
+                                in_blocks = gr.Textbox(value=None, label="In Blocks", scale=10)
+                                mid_block = gr.Textbox(value=None, label="Mid Block", scale=1)
+                                out_blocks = gr.Textbox(value=None, label="Out Block", scale=10)
+                            with FormRow():
+                                beta_label = gr.Markdown("# Beta", visible=False)
+                            with FormRow():
+                                beta_preset = gr.Dropdown(choices=["None"]+list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="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 FormRow():
+                                beta_base = gr.Textbox(value=None, label="Base", scale=1, interactive=True, visible=False)
+                                beta_in_blocks = gr.Textbox(value=None, label="In Blocks", interactive=True, scale=10, visible=False)
+                                beta_mid_block = gr.Textbox(value=None, label="Mid Block", interactive=True, scale=1, visible=False)
+                                beta_out_blocks = gr.Textbox(value=None, label="Out Block", interactive=True, scale=10, visible=False)
+                        with FormRow():
+                            weights_clip = gr.Checkbox(label="Weights Clip")
+                            prune = gr.Checkbox(label="Prune")
+                            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)
+                        with FormRow():
+                            checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors", 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", 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(*args):
+                    try:
+                        results = extras.run_MEHmodelmerger(*args)
+                    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 results
+
+                def tertiary(mode):
+                    if mode in TRIPLE_METHODS:
+                        return [gr.update(visible=True) for _ in range(2)]
+                    else:
+                        return [gr.update(visible=False) for _ in range(2)]
+
+                def beta_visibility(mode):
+                    if mode in BETA_METHODS:
+                        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):
+                    for i, p in enumerate(presets):
+                        presets[i] = BLOCK_WEIGHTS_PRESETS[p]
+                    if len(presets) == 2:
+                        preset = interpolate(presets, ratio)
+                    else:
+                        preset = presets[0]
+                    preset = [str(x) for x in preset]
+                    preset = [preset[0],",".join(preset[1:13]),preset[13],",".join(preset[14:])]
+                    print(preset)
+                    return [gr.update(value=x) for x in preset]
+
+                preset.change(fn=preset_visiblility, inputs=preset, outputs=preset_lambda)
+                beta_preset.change(fn=preset_visiblility, inputs=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=[preset, preset_lambda], outputs=[base,in_blocks,mid_block,out_blocks])
+                MEHmodelmerger_merge.click(
+                    fn=wrap_gradio_gpu_call(MEHmodelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
+                    _js='modelmerger',
+                    inputs=[
+                        dummy_component,
+                        primary_model_name,
+                        secondary_model_name,
+                        tertiary_model_name,
+                        merge_mode,
+                        alpha,
+                        beta,
+                        precision,
+                        custom_name,
+                        checkpoint_format,
+                        save_metadata,
+                        preset,
+                        weights_clip,
+                        prune,
+                        re_basin,
+                        re_basin_iterations,
+                        device
+                    ],
+                    outputs=[
+                        primary_model_name,
+                        secondary_model_name,
+                        tertiary_model_name,
+                        dummy_component,
+                        models_outcome,
+                    ]
+                )
 
             with gr.Tab(label="Validate"):
                 model_headers = ['name', 'type', 'filename', 'hash', 'added', 'size', 'metadata']

From 685f3929676ddf82505f6a05a362d112d6ef0e75 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sun, 5 Nov 2023 07:10:05 -0600
Subject: [PATCH 02/25] Added Tabs

---
 modules/extras.py    | 117 +++++++++++++++++++------------------
 modules/ui_models.py | 134 ++++++++++++++++++++++++++++++-------------
 2 files changed, 155 insertions(+), 96 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 0d5fc2c76..ae7b9af8a 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -12,7 +12,6 @@ from sd_meh.merge import merge_models
 
 from modules import shared, images, sd_models, sd_vae, sd_models_config
 
-
 checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
 
 
@@ -32,6 +31,7 @@ def create_config(ckpt_result, config_source, a, b, c):
     def config(x):
         res = sd_models_config.find_checkpoint_config_near_filename(x) if x else None
         return res if res != shared.sd_default_config else None
+
     if config_source == 0:
         cfg = config(a) or config(b) or config(c)
     elif config_source == 1:
@@ -54,7 +54,9 @@ def to_half(tensor, enable):
     return tensor
 
 
-def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata): # pylint: disable=unused-argument
+def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier,
+                    save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights,
+                    save_metadata):  # pylint: disable=unused-argument
     shared.state.begin('merge')
     save_as_half = save_as_half == 0
 
@@ -148,14 +150,18 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
             # have another 4 channels for unmasked picture's latent space, plus one channel for mask, for a total of 9
             if a.shape != b.shape and a.shape[0:1] + a.shape[2:] == b.shape[0:1] + b.shape[2:]:
                 if a.shape[1] == 4 and b.shape[1] == 9:
-                    raise RuntimeError("When merging inpainting model with a normal one, A must be the inpainting model.")
+                    raise RuntimeError(
+                        "When merging inpainting model with a normal one, A must be the inpainting model.")
                 if a.shape[1] == 4 and b.shape[1] == 8:
-                    raise RuntimeError("When merging instruct-pix2pix model with a normal one, A must be the instruct-pix2pix model.")
-                if a.shape[1] == 8 and b.shape[1] == 4:#If we have an Instruct-Pix2Pix model...
-                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)#Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.
+                    raise RuntimeError(
+                        "When merging instruct-pix2pix model with a normal one, A must be the instruct-pix2pix model.")
+                if a.shape[1] == 8 and b.shape[1] == 4:  # If we have an Instruct-Pix2Pix model...
+                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b,
+                                                             multiplier)  # Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.
                     result_is_instruct_pix2pix_model = True
                 else:
-                    assert a.shape[1] == 9 and b.shape[1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}"
+                    assert a.shape[1] == 9 and b.shape[
+                        1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}"
                     theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)
                     result_is_inpainting_model = True
             else:
@@ -193,7 +199,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     if save_metadata:
         metadata = {"format": "pt", "sd_merge_models": {}}
         merge_recipe = {
-            "type": "webui", # indicate this model was merged with webui's built-in merger
+            "type": "webui",  # 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,
@@ -238,70 +244,68 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     shared.log.info(f"Model merge saved: {output_modelname}.")
     shared.state.textinfo = "Checkpoint saved"
     shared.state.end()
-    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname]
+    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],
+            "Checkpoint saved to " + output_modelname]
 
 
-def run_MEHmodelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, merge_mode,base_alpha,
-                       base_beta, weights_alpha,
-                       weights_beta, precision, custom_name, checkpoint_format, save_metadata, preset, weights_clip, prune,
-                       re_basin, re_basin_iterations, device):  # pylint: disable=unused-argument
+def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     shared.state.begin('model-merge')
-    models = {
-            "model_a": sd_models.checkpoints_list[primary_model_name].filename,
-            "model_b": sd_models.checkpoints_list[secondary_model_name].filename,
-            }
-    if tertiary_model_name is not None:
-        models |= {"model_c": sd_models.checkpoints_list[tertiary_model_name].filename}
-    work_device = device
-    threads = 1
-    preset = preset if preset != "None" else None
-    block_weights_preset_alpha = block_weights_preset_beta = block_weights_preset_alpha_b = block_weights_preset_beta_b = preset if preset is not None else None
-    presets_alpha_lambda = presets_beta_lambda = None
-    logging_level = "INFO"
+    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,
+    }
+    del kwargs["primary_model_name"]
+    del kwargs["secondary_model_name"]
+    if kwargs.get("tertiary_model_name", None) is not None:
+        kwargs["models"] |= {"model_c": sd_models.checkpoints_list[kwargs.get("tertiary_model_name", None)].filename}
+        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"
+        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)
+        kwargs.pop("alpha_mid_block", None)
+        kwargs.pop("alpha_out_blocks", None)
+        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"
+            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)
+            kwargs.pop("beta_mid_block", None)
+            kwargs.pop("beta_out_blocks", None)
+            kwargs.pop("beta_preset", None)
+
+    return [*[gr.update() for _ in range(4)], f"{kwargs}"]
 
     def fail(message):
         shared.state.textinfo = message
         shared.state.end()
         return [*[gr.update() for _ in range(4)], message]
 
-
-
-    theta_0 = main(
-            model_a,
-            model_b,
-            model_c,
-            merge_mode,
-            weights_clip,
-            precision,
-            str(weights_alpha),
-            base_alpha,
-            str(weights_beta),
-            base_beta,
-            re_basin,
-            re_basin_iterations,
-            device,
-            work_device,
-            prune,
-            block_weights_preset_alpha,
-            block_weights_preset_beta,
-            threads,
-            block_weights_preset_alpha_b,
-            block_weights_preset_beta_b,
-            presets_alpha_lambda,
-            presets_beta_lambda,
-            logging_level,
-            )
+    theta_0 = merge_models(**kwargs)
     ckpt_dir = shared.opts.ckpt_dir or sd_models.model_path
     filename = custom_name
-    filename += "." + checkpoint_format
+    filename += "." + kwargs.get("checkpoint_format", None)
     output_modelname = os.path.join(ckpt_dir, filename)
     shared.state.textinfo = "Saving"
     metadata = None
     if save_metadata:
         metadata = {"format": "pt", "sd_merge_models": {}}
         merge_recipe = {
-            "type": "webui",  # indicate this model was merged with webui's built-in merger
+            "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,
@@ -345,8 +349,9 @@ def run_MEHmodelmerger(id_task, primary_model_name, secondary_model_name, tertia
     return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],
             "Checkpoint saved to " + output_modelname]
 
-def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_name, unet_conv, text_encoder_conv, vae_conv, others_conv, fix_clip):
 
+def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_name, unet_conv, text_encoder_conv,
+                     vae_conv, others_conv, fix_clip):
     # position_ids in clip is int64. model_ema.num_updates is int32
     dtypes_to_fp16 = {torch.float32, torch.float64, torch.bfloat16}
     dtypes_to_bf16 = {torch.float32, torch.float64, torch.float16}
@@ -384,7 +389,6 @@ def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_nam
         state_dict = m["state_dict"] if "state_dict" in m else m
         return state_dict
 
-
     def fix_model(model, fix_clip=False):
         # code from model-toolkit
         nai_keys = {
@@ -446,6 +450,7 @@ def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_nam
             ok[wk] = t
         elif conv_t == "delete":
             return
+
     shared.log.info("Model convert: running")
     if conv_type == "ema-only":
         for k in tqdm.tqdm(state_dict):
diff --git a/modules/ui_models.py b/modules/ui_models.py
index e5115662c..abcb9bce1 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -1,5 +1,6 @@
 import os
 import json
+import inspect
 from datetime import datetime
 import gradio as gr
 from modules import sd_models, sd_vae, extras
@@ -140,9 +141,9 @@ def create_ui():
                         models_outcome,
                     ]
                 )
-            with gr.Tab(label="MEH Merge"):
+            with gr.Tab(label="Advanced Merge"):
                 def sd_model_choices():
-                                return ['None'] + sd_models.checkpoint_tiles()
+                    return ['None'] + sd_models.checkpoint_tiles()
                 with gr.Row(equal_height=False):
                     with gr.Column(variant='compact'):
                         with FormRow():
@@ -157,31 +158,35 @@ def create_ui():
                             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():
-                            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 InputAccordion(False, label="Block Merge", elem_id=f"block_merge") as block_accordion:
-                            with FormRow():
-                                alpha_label = gr.Markdown("# Alpha")
-                            with FormRow():
-                                preset = gr.Dropdown(choices=["None"]+list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="Block Weight Preset", multiselect=True, max_choices=2)
-                                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():
-                                base = gr.Textbox(value=None, label="Base", scale=1)
-                                in_blocks = gr.Textbox(value=None, label="In Blocks", scale=10)
-                                mid_block = gr.Textbox(value=None, label="Mid Block", scale=1)
-                                out_blocks = gr.Textbox(value=None, label="Out Block", scale=10)
-                            with FormRow():
-                                beta_label = gr.Markdown("# Beta", visible=False)
-                            with FormRow():
-                                beta_preset = gr.Dropdown(choices=["None"]+list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="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 FormRow():
-                                beta_base = gr.Textbox(value=None, label="Base", scale=1, interactive=True, visible=False)
-                                beta_in_blocks = gr.Textbox(value=None, label="In Blocks", interactive=True, scale=10, visible=False)
-                                beta_mid_block = gr.Textbox(value=None, label="Mid Block", interactive=True, scale=1, visible=False)
-                                beta_out_blocks = gr.Textbox(value=None, label="Out Block", interactive=True, scale=10, visible=False)
+                            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)
+                                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)
+                                    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_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)
                         with FormRow():
                             weights_clip = gr.Checkbox(label="Weights Clip")
                             prune = gr.Checkbox(label="Prune")
@@ -189,22 +194,59 @@ def create_ui():
                         with FormRow():
                             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", 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", label="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(*args):
+                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):
+                    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,[]]:
+                            del kwargs[key]
+
+                    # return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], f"{kwargs}"]
+
                     try:
-                        results = extras.run_MEHmodelmerger(*args)
+                        results = extras.run_MEHmodelmerger(dummy_component, **kwargs)
                     except Exception as e:
                         modules.errors.display(e, 'model merge')
                         sd_models.list_models()  # to remove the potentially missing models from the list
@@ -240,17 +282,17 @@ def create_ui():
                         preset = interpolate(presets, ratio)
                     else:
                         preset = presets[0]
-                    preset = [str(x) for x in preset]
-                    preset = [preset[0],",".join(preset[1:13]),preset[13],",".join(preset[14:])]
-                    print(preset)
-                    return [gr.update(value=x) for x in preset]
+                    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.change(fn=preset_visiblility, inputs=preset, outputs=preset_lambda)
-                beta_preset.change(fn=preset_visiblility, inputs=preset, outputs=beta_preset_lambda)
+                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=[preset, preset_lambda], outputs=[base,in_blocks,mid_block,out_blocks])
+                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',
@@ -262,16 +304,28 @@ def create_ui():
                         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,
-                        preset,
                         weights_clip,
                         prune,
                         re_basin,
                         re_basin_iterations,
-                        device
+                        device,
+                        bake_in_vae,
                     ],
                     outputs=[
                         primary_model_name,

From 020d8ed1dc25f018348ce17f15b7457e6c7f0d73 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sun, 5 Nov 2023 19:45:15 -0600
Subject: [PATCH 03/25] Add Merge Files Natively

---
 modules/extras.py                |   66 +-
 modules/merging/merge.py         |  443 ++++++
 modules/merging/merge_methods.py |  211 +++
 modules/merging/model.py         |   56 +
 modules/merging/presets.py       |   62 +
 modules/merging/rebasin.py       | 2287 ++++++++++++++++++++++++++++++
 modules/merging/utils.py         |  115 ++
 modules/ui_models.py             |  373 +++--
 requirements.txt                 |    1 +
 9 files changed, 3461 insertions(+), 153 deletions(-)
 create mode 100644 modules/merging/merge.py
 create mode 100644 modules/merging/merge_methods.py
 create mode 100644 modules/merging/model.py
 create mode 100644 modules/merging/presets.py
 create mode 100644 modules/merging/rebasin.py
 create mode 100644 modules/merging/utils.py

diff --git a/modules/extras.py b/modules/extras.py
index ae7b9af8a..fe02f002a 100644
--- a/modules/extras.py
+++ b/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()
diff --git a/modules/merging/merge.py b/modules/merging/merge.py
new file mode 100644
index 000000000..e55433d86
--- /dev/null
+++ b/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")
diff --git a/modules/merging/merge_methods.py b/modules/merging/merge_methods.py
new file mode 100644
index 000000000..c10c459f0
--- /dev/null
+++ b/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
diff --git a/modules/merging/model.py b/modules/merging/model.py
new file mode 100644
index 000000000..d2d6470ec
--- /dev/null
+++ b/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
diff --git a/modules/merging/presets.py b/modules/merging/presets.py
new file mode 100644
index 000000000..40ec85010
--- /dev/null
+++ b/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],
+}
diff --git a/modules/merging/rebasin.py b/modules/merging/rebasin.py
new file mode 100644
index 000000000..a0ef80d49
--- /dev/null
+++ b/modules/merging/rebasin.py
@@ -0,0 +1,2287 @@
+# https://github.com/ogkalu2/Merge-Stable-Diffusion-models-without-distortion
+from collections import defaultdict
+from random import shuffle
+from typing import NamedTuple
+import torch
+from scipy.optimize import linear_sum_assignment
+
+SPECIAL_KEYS = [
+    "first_stage_model.decoder.norm_out.weight",
+    "first_stage_model.decoder.norm_out.bias",
+    "first_stage_model.encoder.norm_out.weight",
+    "first_stage_model.encoder.norm_out.bias",
+    "model.diffusion_model.out.0.weight",
+    "model.diffusion_model.out.0.bias",
+]
+
+
+class PermutationSpec(NamedTuple):
+    perm_to_axes: dict
+    axes_to_perm: dict
+
+
+def permutation_spec_from_axes_to_perm(axes_to_perm: dict) -> PermutationSpec:
+    perm_to_axes = defaultdict(list)
+    for wk, axis_perms in axes_to_perm.items():
+        for axis, perm in enumerate(axis_perms):
+            if perm is not None:
+                perm_to_axes[perm].append((wk, axis))
+    return PermutationSpec(perm_to_axes=dict(perm_to_axes), axes_to_perm=axes_to_perm)
+
+
+def sdunet_permutation_spec() -> PermutationSpec:
+    conv = lambda name, p_in, p_out: {
+        f"{name}.weight": (
+            p_out,
+            p_in,
+        ),
+        f"{name}.bias": (p_out,),
+    }
+    norm = lambda name, p: {f"{name}.weight": (p,), f"{name}.bias": (p,)}
+    dense = (
+        lambda name, p_in, p_out, bias=True: {
+            f"{name}.weight": (p_out, p_in),
+            f"{name}.bias": (p_out,),
+        }
+        if bias
+        else {f"{name}.weight": (p_out, p_in)}
+    )
+    skip = lambda name, p_in, p_out: {
+        f"{name}": (
+            p_out,
+            p_in,
+            None,
+            None,
+        )
+    }
+
+    # Unet Res blocks
+    easyblock = lambda name, p_in, p_out: {
+        **norm(f"{name}.in_layers.0", p_in),
+        **conv(f"{name}.in_layers.2", p_in, f"P_{name}_inner"),
+        **dense(
+            f"{name}.emb_layers.1", f"P_{name}_inner2", f"P_{name}_inner3", bias=True
+        ),
+        **norm(f"{name}.out_layers.0", f"P_{name}_inner4"),
+        **conv(f"{name}.out_layers.3", f"P_{name}_inner4", p_out),
+    }
+
+    # Text Encoder blocks
+    easyblock2 = lambda name, p: {
+        **norm(f"{name}.norm1", p),
+        **conv(f"{name}.conv1", p, f"P_{name}_inner"),
+        **norm(f"{name}.norm2", f"P_{name}_inner"),
+        **conv(f"{name}.conv2", f"P_{name}_inner", p),
+    }
+
+    # This is for blocks that use a residual connection, but change the number of channels via a Conv.
+    shortcutblock = lambda name, p_in, p_out: {
+        **norm(f"{name}.norm1", p_in),
+        **conv(f"{name}.conv1", p_in, f"P_{name}_inner"),
+        **norm(f"{name}.norm2", f"P_{name}_inner"),
+        **conv(f"{name}.conv2", f"P_{name}_inner", p_out),
+        **conv(f"{name}.nin_shortcut", p_in, p_out),
+        **norm(f"{name}.nin_shortcut", p_out),
+    }
+
+    return permutation_spec_from_axes_to_perm(
+        {
+            # Skipped Layers
+            **skip("betas", None, None),
+            **skip("alphas_cumprod", None, None),
+            **skip("alphas_cumprod_prev", None, None),
+            **skip("sqrt_alphas_cumprod", None, None),
+            **skip("sqrt_one_minus_alphas_cumprod", None, None),
+            **skip("log_one_minus_alphas_cumprods", None, None),
+            **skip("sqrt_recip_alphas_cumprod", None, None),
+            **skip("sqrt_recipm1_alphas_cumprod", None, None),
+            **skip("posterior_variance", None, None),
+            **skip("posterior_log_variance_clipped", None, None),
+            **skip("posterior_mean_coef1", None, None),
+            **skip("posterior_mean_coef2", None, None),
+            **skip("log_one_minus_alphas_cumprod", None, None),
+            **skip("model_ema.decay", None, None),
+            **skip("model_ema.num_updates", None, None),
+            # initial
+            **dense("model.diffusion_model.time_embed.0", None, "P_bg0", bias=True),
+            **dense("model.diffusion_model.time_embed.2", "P_bg0", "P_bg1", bias=True),
+            **conv("model.diffusion_model.input_blocks.0.0", "P_bg2", "P_bg3"),
+            # input blocks
+            **easyblock("model.diffusion_model.input_blocks.1.0", "P_bg4", "P_bg5"),
+            **norm("model.diffusion_model.input_blocks.1.1.norm", "P_bg6"),
+            **conv("model.diffusion_model.input_blocks.1.1.proj_in", "P_bg6", "P_bg7"),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_q",
+                "P_bg8",
+                "P_bg9",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_k",
+                "P_bg8",
+                "P_bg9",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_v",
+                "P_bg8",
+                "P_bg9",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg8",
+                "P_bg9",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg10",
+                "P_bg11",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.2",
+                "P_bg12",
+                "P_bg13",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_q",
+                "P_bg14",
+                "P_bg15",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k",
+                "P_bg16",
+                "P_bg17",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_v",
+                "P_bg16",
+                "P_bg17",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg18",
+                "P_bg19",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm1",
+                "P_bg19",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm2",
+                "P_bg19",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm3",
+                "P_bg19",
+            ),
+            **conv(
+                "model.diffusion_model.input_blocks.1.1.proj_out", "P_bg19", "P_bg20"
+            ),
+            **easyblock("model.diffusion_model.input_blocks.2.0", "P_bg21", "P_bg22"),
+            **norm("model.diffusion_model.input_blocks.2.1.norm", "P_bg23"),
+            **conv(
+                "model.diffusion_model.input_blocks.2.1.proj_in", "P_bg23", "P_bg24"
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_q",
+                "P_bg25",
+                "P_bg26",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_k",
+                "P_bg25",
+                "P_bg26",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_v",
+                "P_bg25",
+                "P_bg26",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg25",
+                "P_bg26",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg27",
+                "P_bg28",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.2",
+                "P_bg29",
+                "P_bg30",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_q",
+                "P_bg31",
+                "P_bg32",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k",
+                "P_bg33",
+                "P_bg34",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_v",
+                "P_bg33",
+                "P_bg34",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg35",
+                "P_bg36",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm1",
+                "P_bg36",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm2",
+                "P_bg36",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm3",
+                "P_bg36",
+            ),
+            **conv(
+                "model.diffusion_model.input_blocks.2.1.proj_out", "P_bg36", "P_bg37"
+            ),
+            **conv("model.diffusion_model.input_blocks.3.0.op", "P_bg38", "P_bg39"),
+            **easyblock("model.diffusion_model.input_blocks.4.0", "P_bg40", "P_bg41"),
+            **conv(
+                "model.diffusion_model.input_blocks.4.0.skip_connection",
+                "P_bg42",
+                "P_bg43",
+            ),
+            **norm("model.diffusion_model.input_blocks.4.1.norm", "P_bg44"),
+            **conv(
+                "model.diffusion_model.input_blocks.4.1.proj_in", "P_bg44", "P_bg45"
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_q",
+                "P_bg46",
+                "P_bg47",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_k",
+                "P_bg46",
+                "P_bg47",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_v",
+                "P_bg46",
+                "P_bg47",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg46",
+                "P_bg47",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg48",
+                "P_bg49",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.2",
+                "P_bg50",
+                "P_bg51",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_q",
+                "P_bg52",
+                "P_bg53",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_k",
+                "P_bg54",
+                "P_bg55",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_v",
+                "P_bg54",
+                "P_bg55",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg56",
+                "P_bg57",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm1",
+                "P_bg57",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm2",
+                "P_bg57",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm3",
+                "P_bg57",
+            ),
+            **conv(
+                "model.diffusion_model.input_blocks.4.1.proj_out", "P_bg57", "P_bg58"
+            ),
+            **easyblock("model.diffusion_model.input_blocks.5.0", "P_bg59", "P_bg60"),
+            **norm("model.diffusion_model.input_blocks.5.1.norm", "P_bg61"),
+            **conv(
+                "model.diffusion_model.input_blocks.5.1.proj_in", "P_bg61", "P_bg62"
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_q",
+                "P_bg63",
+                "P_bg64",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_k",
+                "P_bg63",
+                "P_bg64",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_v",
+                "P_bg63",
+                "P_bg64",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg63",
+                "P_bg64",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg65",
+                "P_bg66",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.2",
+                "P_bg67",
+                "P_bg68",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_q",
+                "P_bg69",
+                "P_bg70",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_k",
+                "P_bg71",
+                "P_bg72",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_v",
+                "P_bg71",
+                "P_bg72",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg73",
+                "P_bg74",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm1",
+                "P_bg74",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm2",
+                "P_bg74",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm3",
+                "P_bg74",
+            ),
+            **conv(
+                "model.diffusion_model.input_blocks.5.1.proj_out", "P_bg74", "P_bg75"
+            ),
+            **conv("model.diffusion_model.input_blocks.6.0.op", "P_bg76", "P_bg77"),
+            **easyblock("model.diffusion_model.input_blocks.7.0", "P_bg78", "P_bg79"),
+            **conv(
+                "model.diffusion_model.input_blocks.7.0.skip_connection",
+                "P_bg80",
+                "P_bg81",
+            ),
+            **norm("model.diffusion_model.input_blocks.7.1.norm", "P_bg82"),
+            **conv(
+                "model.diffusion_model.input_blocks.7.1.proj_in", "P_bg82", "P_bg83"
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_q",
+                "P_bg84",
+                "P_bg85",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_k",
+                "P_bg84",
+                "P_bg85",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_v",
+                "P_bg84",
+                "P_bg85",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg84",
+                "P_bg85",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg86",
+                "P_bg87",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.2",
+                "P_bg88",
+                "P_bg89",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_q",
+                "P_bg90",
+                "P_bg91",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_k",
+                "P_bg92",
+                "P_bg93",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_v",
+                "P_bg92",
+                "P_bg93",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg94",
+                "P_bg95",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm1",
+                "P_bg95",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm2",
+                "P_bg95",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm3",
+                "P_bg95",
+            ),
+            **conv(
+                "model.diffusion_model.input_blocks.7.1.proj_out", "P_bg95", "P_bg96"
+            ),
+            **easyblock("model.diffusion_model.input_blocks.8.0", "P_bg97", "P_bg98"),
+            **norm("model.diffusion_model.input_blocks.8.1.norm", "P_bg99"),
+            **conv(
+                "model.diffusion_model.input_blocks.8.1.proj_in", "P_bg99", "P_bg100"
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_q",
+                "P_bg101",
+                "P_bg102",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_k",
+                "P_bg101",
+                "P_bg102",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_v",
+                "P_bg101",
+                "P_bg102",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg101",
+                "P_bg102",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg103",
+                "P_bg104",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.2",
+                "P_bg105",
+                "P_bg106",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_q",
+                "P_bg107",
+                "P_bg108",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_k",
+                "P_bg109",
+                "P_bg110",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_v",
+                "P_bg109",
+                "P_bg110",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg111",
+                "P_bg112",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm1",
+                "P_bg112",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm2",
+                "P_bg112",
+            ),
+            **norm(
+                "model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm3",
+                "P_bg112",
+            ),
+            **conv(
+                "model.diffusion_model.input_blocks.8.1.proj_out", "P_bg112", "P_bg113"
+            ),
+            **conv("model.diffusion_model.input_blocks.9.0.op", "P_bg114", "P_bg115"),
+            **easyblock(
+                "model.diffusion_model.input_blocks.10.0", "P_bg115", "P_bg116"
+            ),
+            **easyblock(
+                "model.diffusion_model.input_blocks.11.0", "P_bg116", "P_bg117"
+            ),
+            # middle blocks
+            **easyblock("model.diffusion_model.middle_block.0", "P_bg117", "P_bg118"),
+            **norm("model.diffusion_model.middle_block.1.norm", "P_bg119"),
+            **conv(
+                "model.diffusion_model.middle_block.1.proj_in", "P_bg119", "P_bg120"
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_q",
+                "P_bg121",
+                "P_bg122",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_k",
+                "P_bg121",
+                "P_bg122",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_v",
+                "P_bg121",
+                "P_bg122",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg121",
+                "P_bg122",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg123",
+                "P_bg124",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.2",
+                "P_bg125",
+                "P_bg126",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_q",
+                "P_bg127",
+                "P_bg128",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_k",
+                "P_bg129",
+                "P_bg130",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_v",
+                "P_bg129",
+                "P_bg130",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg131",
+                "P_bg132",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.norm1",
+                "P_bg132",
+            ),
+            **norm(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.norm2",
+                "P_bg132",
+            ),
+            **norm(
+                "model.diffusion_model.middle_block.1.transformer_blocks.0.norm3",
+                "P_bg132",
+            ),
+            **conv(
+                "model.diffusion_model.middle_block.1.proj_out", "P_bg132", "P_bg133"
+            ),
+            **easyblock("model.diffusion_model.middle_block.2", "P_bg134", "P_bg135"),
+            # output blocks
+            **easyblock(
+                "model.diffusion_model.output_blocks.0.0", "P_bg136", "P_bg137"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.0.0.skip_connection",
+                "P_bg138",
+                "P_bg139",
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.1.0", "P_bg140", "P_bg141"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.1.0.skip_connection",
+                "P_bg142",
+                "P_bg143",
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.2.0", "P_bg144", "P_bg145"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.2.0.skip_connection",
+                "P_bg146",
+                "P_bg147",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.2.1.conv", "P_bg148", "P_bg149"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.3.0", "P_bg150", "P_bg151"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.3.0.skip_connection",
+                "P_bg152",
+                "P_bg153",
+            ),
+            **norm("model.diffusion_model.output_blocks.3.1.norm", "P_bg154"),
+            **conv(
+                "model.diffusion_model.output_blocks.3.1.proj_in", "P_bg154", "P_bg155"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_q",
+                "P_bg156",
+                "P_bg157",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_k",
+                "P_bg156",
+                "P_bg157",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_v",
+                "P_bg156",
+                "P_bg157",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg156",
+                "P_bg157",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg158",
+                "P_bg159",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.2",
+                "P_bg160",
+                "P_bg161",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_q",
+                "P_bg162",
+                "P_bg163",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_k",
+                "P_bg164",
+                "P_bg165",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_v",
+                "P_bg164",
+                "P_bg165",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg166",
+                "P_bg167",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm1",
+                "P_bg167",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm2",
+                "P_bg167",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm3",
+                "P_bg167",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.3.1.proj_out", "P_bg167", "P_bg168"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.4.0", "P_bg169", "P_bg170"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.4.0.skip_connection",
+                "P_bg171",
+                "P_bg172",
+            ),
+            **norm("model.diffusion_model.output_blocks.4.1.norm", "P_bg173"),
+            **conv(
+                "model.diffusion_model.output_blocks.4.1.proj_in", "P_bg173", "P_bg174"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_q",
+                "P_bg175",
+                "P_bg176",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_k",
+                "P_bg175",
+                "P_bg176",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_v",
+                "P_bg175",
+                "P_bg176",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg175",
+                "P_bg176",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg177",
+                "P_bg178",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.2",
+                "P_bg179",
+                "P_bg180",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_q",
+                "P_bg181",
+                "P_bg182",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_k",
+                "P_bg183",
+                "P_bg184",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_v",
+                "P_bg183",
+                "P_bg184",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg185",
+                "P_bg186",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm1",
+                "P_bg186",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm2",
+                "P_bg186",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm3",
+                "P_bg186",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.4.1.proj_out", "P_bg186", "P_bg187"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.5.0", "P_bg188", "P_bg189"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.5.0.skip_connection",
+                "P_bg190",
+                "P_bg191",
+            ),
+            **norm("model.diffusion_model.output_blocks.5.1.norm", "P_bg192"),
+            **conv(
+                "model.diffusion_model.output_blocks.5.1.proj_in", "P_bg192", "P_bg193"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_q",
+                "P_bg194",
+                "P_bg195",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_k",
+                "P_bg194",
+                "P_bg195",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_v",
+                "P_bg194",
+                "P_bg195",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg194",
+                "P_bg195",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg196",
+                "P_bg197",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.2",
+                "P_bg198",
+                "P_bg199",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_q",
+                "P_bg200",
+                "P_bg201",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_k",
+                "P_bg202",
+                "P_bg203",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_v",
+                "P_bg202",
+                "P_bg203",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg204",
+                "P_bg205",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm1",
+                "P_bg205",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm2",
+                "P_bg205",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm3",
+                "P_bg205",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.5.1.proj_out", "P_bg205", "P_bg206"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.5.2.conv", "P_bg206", "P_bg207"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.6.0", "P_bg208", "P_bg209"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.6.0.skip_connection",
+                "P_bg210",
+                "P_bg211",
+            ),
+            **norm("model.diffusion_model.output_blocks.6.1.norm", "P_bg212"),
+            **conv(
+                "model.diffusion_model.output_blocks.6.1.proj_in", "P_bg212", "P_bg213"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_q",
+                "P_bg214",
+                "P_bg215",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_k",
+                "P_bg214",
+                "P_bg215",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_v",
+                "P_bg214",
+                "P_bg215",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg214",
+                "P_bg215",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg216",
+                "P_bg217",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.2",
+                "P_bg218",
+                "P_bg219",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_q",
+                "P_bg220",
+                "P_bg221",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_k",
+                "P_bg222",
+                "P_bg223",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_v",
+                "P_bg222",
+                "P_bg223",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg224",
+                "P_bg225",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm1",
+                "P_bg225",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm2",
+                "P_bg225",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm3",
+                "P_bg225",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.6.1.proj_out", "P_bg225", "P_bg226"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.7.0", "P_bg227", "P_bg228"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.7.0.skip_connection",
+                "P_bg229",
+                "P_bg230",
+            ),
+            **norm("model.diffusion_model.output_blocks.7.1.norm", "P_bg231"),
+            **conv(
+                "model.diffusion_model.output_blocks.7.1.proj_in", "P_bg231", "P_bg232"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_q",
+                "P_bg233",
+                "P_bg234",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k",
+                "P_bg233",
+                "P_bg234",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_v",
+                "P_bg233",
+                "P_bg234",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg233",
+                "P_bg234",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg235",
+                "P_bg236",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.2",
+                "P_bg237",
+                "P_bg238",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_q",
+                "P_bg239",
+                "P_bg240",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_k",
+                "P_bg241",
+                "P_bg242",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_v",
+                "P_bg241",
+                "P_bg242",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg243",
+                "P_bg244",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm1",
+                "P_bg244",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm2",
+                "P_bg244",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm3",
+                "P_bg244",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.7.1.proj_out", "P_bg244", "P_bg245"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.8.0", "P_bg246", "P_bg247"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.8.0.skip_connection",
+                "P_bg248",
+                "P_bg249",
+            ),
+            **norm("model.diffusion_model.output_blocks.8.1.norm", "P_bg250"),
+            **conv(
+                "model.diffusion_model.output_blocks.8.1.proj_in", "P_bg250", "P_bg251"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_q",
+                "P_bg252",
+                "P_bg253",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_k",
+                "P_bg252",
+                "P_bg253",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_v",
+                "P_bg252",
+                "P_bg253",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg252",
+                "P_bg253",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg254",
+                "P_bg255",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.2",
+                "P_bg256",
+                "P_bg257",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_q",
+                "P_bg258",
+                "P_bg259",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_k",
+                "P_bg260",
+                "P_bg261",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_v",
+                "P_bg260",
+                "P_bg261",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg262",
+                "P_bg263",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm1",
+                "P_bg263",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm2",
+                "P_bg263",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm3",
+                "P_bg263",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.8.1.proj_out", "P_bg263", "P_bg264"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.8.2.conv", "P_bg265", "P_bg266"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.9.0", "P_bg267", "P_bg268"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.9.0.skip_connection",
+                "P_bg269",
+                "P_bg270",
+            ),
+            **norm("model.diffusion_model.output_blocks.9.1.norm", "P_bg271"),
+            **conv(
+                "model.diffusion_model.output_blocks.9.1.proj_in", "P_bg271", "P_bg272"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_q",
+                "P_bg273",
+                "P_bg274",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_k",
+                "P_bg273",
+                "P_bg274",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_v",
+                "P_bg273",
+                "P_bg274",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg273",
+                "P_bg274",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg275",
+                "P_bg276",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.2",
+                "P_bg277",
+                "P_bg278",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_q",
+                "P_bg279",
+                "P_bg280",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_k",
+                "P_bg281",
+                "P_bg282",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_v",
+                "P_bg281",
+                "P_bg282",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg283",
+                "P_bg284",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm1",
+                "P_bg284",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm2",
+                "P_bg284",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm3",
+                "P_bg284",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.9.1.proj_out", "P_bg284", "P_bg285"
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.10.0", "P_bg286", "P_bg287"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.10.0.skip_connection",
+                "P_bg288",
+                "P_bg289",
+            ),
+            **norm("model.diffusion_model.output_blocks.10.1.norm", "P_bg290"),
+            **conv(
+                "model.diffusion_model.output_blocks.10.1.proj_in", "P_bg290", "P_bg291"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_q",
+                "P_bg292",
+                "P_bg293",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_k",
+                "P_bg292",
+                "P_bg293",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_v",
+                "P_bg292",
+                "P_bg293",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg292",
+                "P_bg293",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.0.proj",
+                "P_b294",
+                "P_bg295",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.2",
+                "P_bg296",
+                "P_bg297",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_q",
+                "P_bg298",
+                "P_bg299",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_k",
+                "P_bg300",
+                "P_bg301",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_v",
+                "P_bg300",
+                "P_bg301",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg302",
+                "P_bg303",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm1",
+                "P_bg303",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm2",
+                "P_bg303",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm3",
+                "P_bg303",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.10.1.proj_out",
+                "P_bg303",
+                "P_bg304",
+            ),
+            **easyblock(
+                "model.diffusion_model.output_blocks.11.0", "P_bg305", "P_bg306"
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.11.0.skip_connection",
+                "P_bg307",
+                "P_bg308",
+            ),
+            **norm("model.diffusion_model.output_blocks.11.1.norm", "P_bg309"),
+            **conv(
+                "model.diffusion_model.output_blocks.11.1.proj_in", "P_bg309", "P_bg310"
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_q",
+                "P_bg311",
+                "P_bg312",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_k",
+                "P_bg311",
+                "P_bg312",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_v",
+                "P_bg311",
+                "P_bg312",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_out.0",
+                "P_bg311",
+                "P_bg312",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.0.proj",
+                "P_bg313",
+                "P_bg314",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.2",
+                "P_bg315",
+                "P_bg316",
+                bias=True,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_q",
+                "P_bg317",
+                "P_bg318",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_k",
+                "P_bg319",
+                "P_bg320",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_v",
+                "P_bg319",
+                "P_bg320",
+                bias=False,
+            ),
+            **dense(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_out.0",
+                "P_bg321",
+                "P_bg322",
+                bias=True,
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm1",
+                "P_bg322",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm2",
+                "P_bg322",
+            ),
+            **norm(
+                "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm3",
+                "P_bg322",
+            ),
+            **conv(
+                "model.diffusion_model.output_blocks.11.1.proj_out",
+                "P_bg322",
+                "P_bg323",
+            ),
+            **norm("model.diffusion_model.out.0", "P_bg324"),
+            **conv("model.diffusion_model.out.2", "P_bg325", "P_bg326"),
+            # Text Encoder
+            # encoder down
+            **conv("first_stage_model.encoder.conv_in", "P_bg327", "P_bg328"),
+            **easyblock2("first_stage_model.encoder.down.0.block.0", "P_bg328"),
+            **easyblock2("first_stage_model.encoder.down.0.block.1", "P_bg328"),
+            **conv(
+                "first_stage_model.encoder.down.0.downsample.conv", "P_bg328", "P_bg329"
+            ),
+            **shortcutblock(
+                "first_stage_model.encoder.down.1.block.0", "P_bg330", "P_bg331"
+            ),
+            **easyblock2("first_stage_model.encoder.down.1.block.1", "P_bg331"),
+            **conv(
+                "first_stage_model.encoder.down.1.downsample.conv", "P_bg331", "P_bg332"
+            ),
+            **shortcutblock(
+                "first_stage_model.encoder.down.2.block.0", "P_bg332", "P_bg333"
+            ),
+            **easyblock2("first_stage_model.encoder.down.2.block.1", "P_bg333"),
+            **conv(
+                "first_stage_model.encoder.down.2.downsample.conv", "P_bg333", "P_bg334"
+            ),
+            **easyblock2("first_stage_model.encoder.down.3.block.0", "P_bg334"),
+            **easyblock2("first_stage_model.encoder.down.3.block.1", "P_bg334"),
+            # encoder mid-block
+            **easyblock2("first_stage_model.encoder.mid.block_1", "P_bg334"),
+            **norm("first_stage_model.encoder.mid.attn_1.norm", "P_bg334"),
+            **conv("first_stage_model.encoder.mid.attn_1.q", "P_bg334", "P_bg335"),
+            **conv("first_stage_model.encoder.mid.attn_1.k", "P_bg334", "P_bg335"),
+            **conv("first_stage_model.encoder.mid.attn_1.v", "P_bg334", "P_bg335"),
+            **conv(
+                "first_stage_model.encoder.mid.attn_1.proj_out", "P_bg335", "P_bg336"
+            ),
+            **easyblock2("first_stage_model.encoder.mid.block_2", "P_bg336"),
+            **norm("first_stage_model.encoder.norm_out", "P_bg337"),
+            **conv("first_stage_model.encoder.conv_out", "P_bg338", "P_bg339"),
+            **conv("first_stage_model.decoder.conv_in", "P_bg340", "P_bg341"),
+            # decoder mid-block
+            **easyblock2("first_stage_model.decoder.mid.block_1", "P_bg342"),
+            **norm("first_stage_model.decoder.mid.attn_1.norm", "P_bg342"),
+            **conv("first_stage_model.decoder.mid.attn_1.q", "P_bg342", "P_bg343"),
+            **conv("first_stage_model.decoder.mid.attn_1.k", "P_bg342", "P_bg343"),
+            **conv("first_stage_model.decoder.mid.attn_1.v", "P_bg342", "P_bg343"),
+            **conv(
+                "first_stage_model.decoder.mid.attn_1.proj_out", "P_bg343", "P_bg344"
+            ),
+            **easyblock2("first_stage_model.decoder.mid.block_2", "P_bg345"),
+            # decoder up
+            **shortcutblock(
+                "first_stage_model.decoder.up.0.block.0", "P_bg346", "P_bg347"
+            ),
+            **easyblock2("first_stage_model.decoder.up.0.block.1", "P_bg348"),
+            **easyblock2("first_stage_model.decoder.up.0.block.2", "P_bg349"),
+            **shortcutblock(
+                "first_stage_model.decoder.up.1.block.0", "P_bg350", "P_bg351"
+            ),
+            **easyblock2("first_stage_model.decoder.up.1.block.1", "P_bg352"),
+            **easyblock2("first_stage_model.decoder.up.1.block.2", "P_bg353"),
+            **conv(
+                "first_stage_model.decoder.up.1.upsample.conv", "P_bg353", "P_bg354"
+            ),
+            **easyblock2("first_stage_model.decoder.up.2.block.0", "P_bg355"),
+            **easyblock2("first_stage_model.decoder.up.2.block.1", "P_bg355"),
+            **easyblock2("first_stage_model.decoder.up.2.block.2", "P_bg355"),
+            **conv(
+                "first_stage_model.decoder.up.2.upsample.conv", "P_bg355", "P_bg356"
+            ),
+            **easyblock2("first_stage_model.decoder.up.3.block.0", "P_bg356"),
+            **easyblock2("first_stage_model.decoder.up.3.block.1", "P_bg356"),
+            **easyblock2("first_stage_model.decoder.up.3.block.2", "P_bg356"),
+            **conv(
+                "first_stage_model.decoder.up.3.upsample.conv", "P_bg356", "P_bg357"
+            ),
+            **norm("first_stage_model.decoder.norm_out", "P_bg358"),
+            **conv("first_stage_model.decoder.conv_out", "P_bg359", "P_bg360"),
+            **conv("first_stage_model.quant_conv", "P_bg361", "P_bg362"),
+            **conv("first_stage_model.post_quant_conv", "P_bg363", "P_bg364"),
+            **skip(
+                "cond_stage_model.transformer.text_model.embeddings.position_ids",
+                None,
+                None,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.embeddings.token_embedding",
+                "P_bg365",
+                "P_bg366",
+                bias=False,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.embeddings.token_embedding",
+                None,
+                None,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.embeddings.position_embedding",
+                "P_bg367",
+                "P_bg368",
+                bias=False,
+            ),
+            # cond stage text encoder
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj",
+                "P_bg369",
+                "P_bg370",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj",
+                "P_bg369",
+                "P_bg370",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj",
+                "P_bg369",
+                "P_bg370",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj",
+                "P_bg369",
+                "P_bg370",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1",
+                "P_bg370",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1",
+                "P_bg370",
+                "P_bg371",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2",
+                "P_bg371",
+                "P_bg372",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2",
+                "P_bg372",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj",
+                "P_bg372",
+                "P_bg373",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj",
+                "P_bg372",
+                "P_bg373",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj",
+                "P_bg372",
+                "P_bg373",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj",
+                "P_bg372",
+                "P_bg373",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1",
+                "P_bg373",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1",
+                "P_bg373",
+                "P_bg374",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2",
+                "P_bg374",
+                "P_bg375",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2",
+                "P_bg375",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj",
+                "P_bg375",
+                "P_bg376",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj",
+                "P_bg375",
+                "P_bg376",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj",
+                "P_bg375",
+                "P_bg376",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj",
+                "P_bg375",
+                "P_bg376",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1",
+                "P_bg376",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1",
+                "P_bg376",
+                "P_bg377",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2",
+                "P_bg377",
+                "P_bg378",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2",
+                "P_bg378",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj",
+                "P_bg378",
+                "P_bg379",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj",
+                "P_bg378",
+                "P_bg379",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj",
+                "P_bg378",
+                "P_bg379",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj",
+                "P_bg378",
+                "P_bg379",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1",
+                "P_bg379",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1",
+                "P_bg379",
+                "P_bg380",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2",
+                "P_bg380",
+                "P_b381",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2",
+                "P_bg381",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj",
+                "P_bg381",
+                "P_bg382",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj",
+                "P_bg381",
+                "P_bg382",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj",
+                "P_bg381",
+                "P_bg382",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj",
+                "P_bg381",
+                "P_bg382",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1",
+                "P_bg382",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1",
+                "P_bg382",
+                "P_bg383",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2",
+                "P_bg383",
+                "P_bg384",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2",
+                "P_bg384",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj",
+                "P_bg384",
+                "P_bg385",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj",
+                "P_bg384",
+                "P_bg385",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj",
+                "P_bg384",
+                "P_bg385",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj",
+                "P_bg384",
+                "P_bg385",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1",
+                "P_bg385",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1",
+                "P_bg385",
+                "P_bg386",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2",
+                "P_bg386",
+                "P_bg387",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2",
+                "P_bg387",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj",
+                "P_bg387",
+                "P_bg388",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj",
+                "P_bg387",
+                "P_bg388",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj",
+                "P_bg387",
+                "P_bg388",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj",
+                "P_bg387",
+                "P_bg388",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1",
+                "P_bg389",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1",
+                "P_bg389",
+                "P_bg390",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2",
+                "P_bg390",
+                "P_bg391",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2",
+                "P_bg391",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj",
+                "P_bg391",
+                "P_bg392",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj",
+                "P_bg391",
+                "P_bg392",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj",
+                "P_bg391",
+                "P_bg392",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj",
+                "P_bg391",
+                "P_bg392",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1",
+                "P_bg392",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1",
+                "P_bg392",
+                "P_bg393",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2",
+                "P_bg393",
+                "P_bg394",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2",
+                "P_bg394",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj",
+                "P_bg394",
+                "P_bg395",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj",
+                "P_bg394",
+                "P_bg395",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj",
+                "P_bg394",
+                "P_bg395",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj",
+                "P_bg394",
+                "P_bg395",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1",
+                "P_bg395",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1",
+                "P_bg395",
+                "P_bg396",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2",
+                "P_bg396",
+                "P_bg397",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2",
+                "P_bg397",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj",
+                "P_bg397",
+                "P_bg398",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj",
+                "P_bg397",
+                "P_bg398",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj",
+                "P_bg397",
+                "P_bg398",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj",
+                "P_bg397",
+                "P_bg398",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1",
+                "P_bg398",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1",
+                "P_bg398",
+                "P_bg399",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2",
+                "P_bg400",
+                "P_bg401",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2",
+                "P_bg401",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj",
+                "P_bg401",
+                "P_bg402",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj",
+                "P_bg401",
+                "P_bg402",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj",
+                "P_bg401",
+                "P_bg402",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj",
+                "P_bg401",
+                "P_bg402",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1",
+                "P_bg402",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1",
+                "P_bg402",
+                "P_bg403",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2",
+                "P_bg403",
+                "P_bg404",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2",
+                "P_bg404",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj",
+                "P_bg404",
+                "P_bg405",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj",
+                "P_bg404",
+                "P_bg405",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj",
+                "P_bg404",
+                "P_bg405",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj",
+                "P_bg404",
+                "P_bg405",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1",
+                "P_bg405",
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1",
+                "P_bg405",
+                "P_bg406",
+                bias=True,
+            ),
+            **dense(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2",
+                "P_bg406",
+                "P_bg407",
+                bias=True,
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2",
+                "P_bg407",
+            ),
+            **norm(
+                "cond_stage_model.transformer.text_model.final_layer_norm", "P_bg407"
+            ),
+        }
+    )
+
+
+def get_permuted_param(ps: PermutationSpec, perm, k: str, params, except_axis=None):
+    """Get parameter `k` from `params`, with the permutations applied."""
+    w = params[k]
+    for axis, p in enumerate(ps.axes_to_perm[k]):
+        # Skip the axis we're trying to permute.
+        if axis == except_axis:
+            continue
+
+        # None indicates that there is no permutation relevant to that axis.
+        if p:
+            w = torch.index_select(w, axis, perm[p].int())
+
+    return w
+
+
+def apply_permutation(ps: PermutationSpec, perm, params):
+    """Apply a `perm` to `params`."""
+    return {k: get_permuted_param(ps, perm, k, params) for k in params.keys()}
+
+
+def update_model_a(ps: PermutationSpec, perm, model_a, new_alpha):
+    for k in model_a:
+        try:
+            perm_params = get_permuted_param(
+                ps, perm, k, model_a
+            )
+            model_a[k] = model_a[k] * (1 - new_alpha) + new_alpha * perm_params
+        except RuntimeError:  # dealing with pix2pix and inpainting models
+            continue
+    return model_a
+
+
+def inner_matching(
+    n,
+    ps,
+    p,
+    params_a,
+    params_b,
+    usefp16,
+    progress,
+    number,
+    linear_sum,
+    perm,
+    device,
+):
+    A = torch.zeros((n, n), dtype=torch.float16) if usefp16 else torch.zeros((n, n))
+    A = A.to(device)
+
+    for wk, axis in ps.perm_to_axes[p]:
+        w_a = params_a[wk]
+        w_b = get_permuted_param(ps, perm, wk, params_b, except_axis=axis)
+        w_a = torch.moveaxis(w_a, axis, 0).reshape((n, -1)).to(device)
+        w_b = torch.moveaxis(w_b, axis, 0).reshape((n, -1)).T.to(device)
+
+        if usefp16:
+            w_a = w_a.half().to(device)
+            w_b = w_b.half().to(device)
+
+        try:
+            A += torch.matmul(w_a, w_b)
+        except RuntimeError:
+            A += torch.matmul(torch.dequantize(w_a), torch.dequantize(w_b))
+
+    A = A.cpu()
+    ri, ci = linear_sum_assignment(A.detach().numpy(), maximize=True)
+    A = A.to(device)
+
+    assert (torch.tensor(ri) == torch.arange(len(ri))).all()
+
+    eye_tensor = torch.eye(n).to(device)
+
+    oldL = torch.vdot(
+        torch.flatten(A).float(), torch.flatten(eye_tensor[perm[p].long()])
+    )
+    newL = torch.vdot(torch.flatten(A).float(), torch.flatten(eye_tensor[ci, :]))
+
+    if usefp16:
+        oldL = oldL.half()
+        newL = newL.half()
+
+    if newL - oldL != 0:
+        linear_sum += abs((newL - oldL).item())
+        number += 1
+        logging.info(f" permutation {p}: {newL - oldL}")
+
+    progress = progress or newL > oldL + 1e-12
+
+    perm[p] = torch.Tensor(ci).to(device)
+
+    return linear_sum, number, perm, progress
+
+
+def weight_matching(
+    ps: PermutationSpec,
+    params_a,
+    params_b,
+    max_iter=1,
+    init_perm=None,
+    usefp16=False,
+    device="cpu",
+):
+    perm_sizes = {
+        p: params_a[axes[0][0]].shape[axes[0][1]]
+        for p, axes in ps.perm_to_axes.items()
+        if axes[0][0] in params_a.keys()
+    }
+    perm = {}
+    perm = (
+        {p: torch.arange(n).to(device) for p, n in perm_sizes.items()}
+        if init_perm is None
+        else init_perm
+    )
+
+    linear_sum = 0
+    number = 0
+
+    special_layers = ["P_bg324", "P_bg358", "P_bg337"]
+    for _ in range(max_iter):
+        progress = False
+        shuffle(special_layers)
+        for p in special_layers:
+            n = perm_sizes[p]
+
+            linear_sum, number, perm, progress = inner_matching(
+                n,
+                ps,
+                p,
+                params_a,
+                params_b,
+                usefp16,
+                progress,
+                number,
+                linear_sum,
+                perm,
+                device,
+            )
+        if not progress:
+            break
+
+    average = linear_sum / number if number > 0 else 0
+    return perm, average
diff --git a/modules/merging/utils.py b/modules/merging/utils.py
new file mode 100644
index 000000000..4cee23db1
--- /dev/null
+++ b/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
diff --git a/modules/ui_models.py b/modules/ui_models.py
index abcb9bce1..87c1c2898 100644
--- a/modules/ui_models.py
+++ b/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])
diff --git a/requirements.txt b/requirements.txt
index 3cd2b7fe9..68b614704 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -33,6 +33,7 @@ rich
 safetensors
 scipy
 tb_nightly
+tensordict
 toml
 torchdiffeq
 voluptuous

From 515a546a91b0bcc36020189b16855482ede35680 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Mon, 6 Nov 2023 08:19:50 -0600
Subject: [PATCH 04/25] Fix rebasin

---
 modules/merging/utils.py | 29 ++++++++++++-----------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/modules/merging/utils.py b/modules/merging/utils.py
index 4cee23db1..7feff1393 100644
--- a/modules/merging/utils.py
+++ b/modules/merging/utils.py
@@ -30,7 +30,7 @@ class WeightClass:
                  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.SDXL = "model.diffusion_model.middle_block.1.transformer_blocks.9.norm3.weight" in model_a.keys()
         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
@@ -59,13 +59,12 @@ class WeightClass:
                 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)
+            current_bases["alpha"] = self.step_weights_and_bases(self.ratioDict["alpha"])
         if self.ratioDict.get("beta", None):
-            current_bases["beta"] = self.step_weights_and_bases(self.ratioDict["beta"], it)
+            current_bases["beta"] = self.step_weights_and_bases(self.ratioDict["beta"])
 
         weight_index = 0
         if "model" in key:
@@ -92,22 +91,18 @@ class WeightClass:
                     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()
-        }
+    def step_weights_and_bases(self, ratio):
+        if not self.re_basin:
+            return ratio
 
+        new_ratio = [
+            1 - (1 - (1 + self.it) * v / self.iterations) / (1 - self.it * v / self.iterations)
+            if self.it > 0
+            else v / self.iterations
+            for v in ratio
+        ]
         return new_ratio
 
     def set_it(self, it):

From 864fc898face90451faba46aa6f4e1fe16a6d4bf Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 10 Nov 2023 06:51:38 -0600
Subject: [PATCH 05/25] Bugfix

---
 modules/extras.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index fe02f002a..044fc74ba 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -271,8 +271,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     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
+    tertiary_model_info = sd_models.checkpoints_list[kwargs.get("tertiary_model_name", None)] if kwargs.get("merge_mode", None) in TRIPLE_METHODS else None
 
     del kwargs["primary_model_name"]
     del kwargs["secondary_model_name"]
@@ -287,7 +286,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         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:
+    except:
         kwargs["alpha"] = kwargs.get("alpha_preset", kwargs["alpha"])
     finally:
         kwargs.pop("alpha_base", None)
@@ -295,6 +294,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         kwargs.pop("alpha_mid_block", None)
         kwargs.pop("alpha_out_blocks", None)
         kwargs.pop("alpha_preset", None)
+
     if kwargs.get("beta", False):
         try:
             beta = [float(x) for x in
@@ -303,7 +303,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             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:
+        except:
             kwargs["beta"] = kwargs.get("beta_preset", kwargs["beta"])
         finally:
             kwargs.pop("beta_base", None)
@@ -355,7 +355,11 @@ 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()
+    try:
+        theta_0 = theta_0.to_dict()
+    except:
+        pass
+
     if extension.lower() == ".safetensors":
         safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
     else:

From 02a32e71088401bb6b0629f674f13eb866e589b3 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 10 Nov 2023 07:02:17 -0600
Subject: [PATCH 06/25] Add Logging

---
 modules/merging/merge.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index e55433d86..76c5b036d 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -9,6 +9,7 @@ import safetensors.torch
 import torch
 from tqdm import tqdm
 
+from modules.shared import log
 from modules.merging import merge_methods
 from modules.merging.utils import WeightClass
 from modules.merging.model import SDModel
@@ -19,8 +20,6 @@ from modules.merging.rebasin import (
     weight_matching,
 )
 
-# logging.getLogger("sd_meh").addHandler(logging.NullHandler())
-
 MAX_TOKENS = 77
 
 
@@ -97,7 +96,7 @@ def restore_sd_model(original_model: Dict, merged_model: Dict) -> Dict:
 
 def log_vram(txt=""):
     alloc = torch.cuda.memory_allocated(0)
-    # logging.debug(f"{txt} VRAM: {alloc*1e-9:5.3f}GB")
+    log.info(f"{txt} VRAM: {alloc*1e-9:5.3f}GB")
 
 
 def load_thetas(
@@ -130,16 +129,16 @@ def merge_models(
     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:
+    iterations = kwargs.get("iterations", 1)
     thetas = load_thetas(models, prune, device, precision)
 
-    # logging.info(f"start merging with {merge_mode} method")
+    log.info(f"start merging with {merge_mode} method")
     weight_matcher = WeightClass(thetas["model_a"], **kwargs)
     if re_basin:
         merged = rebasin_merge(
@@ -177,7 +176,7 @@ def un_prune_model(
     precision: str,
 ) -> Dict:
     if prune:
-        # logging.info("Un-pruning merged model")
+        log.info("Un-pruning merged model")
         del thetas
         gc.collect()
         log_vram("remove thetas")
@@ -432,7 +431,7 @@ def get_merge_method_args(
 
 
 def save_model(model, output_file, file_format) -> None:
-    # logging.info(f"Saving {output_file}")
+    log.info(f"Saving {output_file}")
     if file_format == "safetensors":
         safetensors.torch.save_file(
             model if type(model) == dict else model.to_dict(),

From b65dfae35ad70e7989f4f1db2227c3fefca19d58 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 10 Nov 2023 07:21:24 -0600
Subject: [PATCH 07/25] Fix rebasin

---
 modules/extras.py        | 2 +-
 modules/merging/merge.py | 9 +++------
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 044fc74ba..b91010958 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -318,7 +318,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         return fail(f"{e}")
 
     ckpt_dir = shared.opts.ckpt_dir or sd_models.model_path
-    filename = kwargs.get("custom_name", "Unamed_Merge")
+    filename = kwargs.get("custom_name", "Unnamed_Merge")
     filename += "." + kwargs.get("checkpoint_format", None)
     output_modelname = os.path.join(ckpt_dir, filename)
     shared.state.textinfo = "Saving"
diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index 76c5b036d..6b00afe2b 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -135,7 +135,7 @@ def merge_models(
     threads: int = 1,
     **kwargs,
 ) -> Dict:
-    iterations = kwargs.get("iterations", 1)
+    iterations = kwargs.get("re_basin_iterations", 1)
     thetas = load_thetas(models, prune, device, precision)
 
     log.info(f"start merging with {merge_mode} method")
@@ -190,7 +190,7 @@ def un_prune_model(
                     merged.update({key: merged[key].half()})
         del original_a
         gc.collect()
-        log_vram("remove original_a")
+        # 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:
@@ -266,12 +266,9 @@ def rebasin_merge(
     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")
+        log_vram(f"Rebasin iteration {it}")
         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(

From 39acd0af662c15950f55c57640f062be94a72830 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 10 Nov 2023 07:38:37 -0600
Subject: [PATCH 08/25] Add VAE bake

---
 modules/extras.py          | 11 +++++++++++
 modules/merging/rebasin.py |  3 ++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/modules/extras.py b/modules/extras.py
index b91010958..d5293f2f6 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -317,6 +317,17 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     except Exception as e:
         return fail(f"{e}")
 
+    bake_in_vae_filename = sd_vae.vae_dict.get(kwargs.get("bake_in_vae", None), None)
+    if bake_in_vae_filename is not None:
+        shared.log.info(f"Model merge: baking in VAE: {bake_in_vae_filename}")
+        shared.state.textinfo = 'Baking in VAE'
+        vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename)
+        for key in vae_dict.keys():
+            theta_0_key = 'first_stage_model.' + key
+            if theta_0_key in theta_0:
+                theta_0[theta_0_key] = to_half(vae_dict[key], kwargs.get("precision", "fp16") == "fp16")
+        del vae_dict
+
     ckpt_dir = shared.opts.ckpt_dir or sd_models.model_path
     filename = kwargs.get("custom_name", "Unnamed_Merge")
     filename += "." + kwargs.get("checkpoint_format", None)
diff --git a/modules/merging/rebasin.py b/modules/merging/rebasin.py
index a0ef80d49..26aac9964 100644
--- a/modules/merging/rebasin.py
+++ b/modules/merging/rebasin.py
@@ -4,6 +4,7 @@ from random import shuffle
 from typing import NamedTuple
 import torch
 from scipy.optimize import linear_sum_assignment
+from modules.shared import log
 
 SPECIAL_KEYS = [
     "first_stage_model.decoder.norm_out.weight",
@@ -2227,7 +2228,7 @@ def inner_matching(
     if newL - oldL != 0:
         linear_sum += abs((newL - oldL).item())
         number += 1
-        logging.info(f" permutation {p}: {newL - oldL}")
+        log.info(f" permutation {p}: {newL - oldL}")
 
     progress = progress or newL > oldL + 1e-12
 

From db20da7b1a52d988b09c380a1f8b03cf6449513e Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Thu, 16 Nov 2023 19:01:03 -0600
Subject: [PATCH 09/25] Fixes

---
 modules/merging/merge.py                         | 10 +++++-----
 modules/merging/{model.py => merge_model.py}     |  0
 modules/merging/{presets.py => merge_presets.py} |  0
 modules/merging/{rebasin.py => merge_rebasin.py} |  0
 modules/merging/{utils.py => merge_utils.py}     |  7 +++++--
 modules/ui_models.py                             | 12 +++++-------
 6 files changed, 15 insertions(+), 14 deletions(-)
 rename modules/merging/{model.py => merge_model.py} (100%)
 rename modules/merging/{presets.py => merge_presets.py} (100%)
 rename modules/merging/{rebasin.py => merge_rebasin.py} (100%)
 rename modules/merging/{utils.py => merge_utils.py} (95%)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index 6b00afe2b..fe0097da9 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -9,11 +9,12 @@ import safetensors.torch
 import torch
 from tqdm import tqdm
 
+import modules.memstats
 from modules.shared import log
 from modules.merging import merge_methods
-from modules.merging.utils import WeightClass
-from modules.merging.model import SDModel
-from modules.merging.rebasin import (
+from modules.merging.merge_utils import WeightClass
+from modules.merging.merge_model import SDModel
+from modules.merging.merge_rebasin import (
     apply_permutation,
     sdunet_permutation_spec,
     update_model_a,
@@ -95,8 +96,7 @@ def restore_sd_model(original_model: Dict, merged_model: Dict) -> Dict:
 
 
 def log_vram(txt=""):
-    alloc = torch.cuda.memory_allocated(0)
-    log.info(f"{txt} VRAM: {alloc*1e-9:5.3f}GB")
+    log.info(f"{txt} VRAM: {modules.memstats.memory_stats()}")
 
 
 def load_thetas(
diff --git a/modules/merging/model.py b/modules/merging/merge_model.py
similarity index 100%
rename from modules/merging/model.py
rename to modules/merging/merge_model.py
diff --git a/modules/merging/presets.py b/modules/merging/merge_presets.py
similarity index 100%
rename from modules/merging/presets.py
rename to modules/merging/merge_presets.py
diff --git a/modules/merging/rebasin.py b/modules/merging/merge_rebasin.py
similarity index 100%
rename from modules/merging/rebasin.py
rename to modules/merging/merge_rebasin.py
diff --git a/modules/merging/utils.py b/modules/merging/merge_utils.py
similarity index 95%
rename from modules/merging/utils.py
rename to modules/merging/merge_utils.py
index 7feff1393..41b6f7442 100644
--- a/modules/merging/utils.py
+++ b/modules/merging/merge_utils.py
@@ -2,9 +2,12 @@ import inspect
 # import logging
 import re
 from modules.merging import merge_methods
-from modules.merging.presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
+from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
+
+ALL_PRESETS = {}
+ALL_PRESETS.update(BLOCK_WEIGHTS_PRESETS)
+ALL_PRESETS.update(SDXL_BLOCK_WEIGHTS_PRESETS)
 
-BLOCK_WEIGHTS_PRESETS |= SDXL_BLOCK_WEIGHTS_PRESETS
 MERGE_METHODS = dict(inspect.getmembers(merge_methods, inspect.isfunction))
 BETA_METHODS = [
     name
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 87c1c2898..a9c99e8e3 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -11,8 +11,8 @@ from modules.shared import opts, log, req
 import modules.errors
 import modules.hashes
 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
+from modules.merging.merge_utils import BETA_METHODS, TRIPLE_METHODS, interpolate
+from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
 
 search_metadata_civit = None
 
@@ -242,7 +242,7 @@ def create_ui():
                         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")
+                            work_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")
@@ -281,7 +281,7 @@ def create_ui():
                                    prune,
                                    re_basin,
                                    re_basin_iterations,
-                                   device,
+                                   work_device,
                                    bake_in_vae):
                     kwargs = {}
                     for x in inspect.getfullargspec(MEHmodelmerger)[0]:
@@ -290,8 +290,6 @@ def create_ui():
                         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}"]
-
                     try:
                         results = extras.run_MEHmodelmerger(dummy_component, **kwargs)
                     except Exception as e:
@@ -386,7 +384,7 @@ def create_ui():
                         prune,
                         re_basin,
                         re_basin_iterations,
-                        device,
+                        work_device,
                         bake_in_vae,
                     ],
                     outputs=[

From 26f6e717f67d7e89c4318c9e0c695c363c59a2e9 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Thu, 16 Nov 2023 19:17:13 -0600
Subject: [PATCH 10/25] Update extras.py

---
 modules/extras.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/extras.py b/modules/extras.py
index d5293f2f6..467946166 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -9,7 +9,7 @@ import tqdm
 import gradio as gr
 import safetensors.torch
 from modules.merging.merge import merge_models
-from modules.merging.utils import TRIPLE_METHODS
+from modules.merging.merge_utils import TRIPLE_METHODS
 
 from modules import shared, images, sd_models, sd_vae, sd_models_config
 

From 6c34317a576f66b4e3fa137c3d8cf25f3e2007f1 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Thu, 16 Nov 2023 19:31:36 -0600
Subject: [PATCH 11/25] fix precision

---
 modules/merging/merge.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index fe0097da9..ae68fba15 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -126,7 +126,7 @@ def load_thetas(
 def merge_models(
     models: Dict[str, os.PathLike | str],
     merge_mode: str,
-    precision: str = "full",
+    precision: str = "fp16",
     weights_clip: bool = False,
     re_basin: bool = False,
     device: str = "cpu",
@@ -242,7 +242,7 @@ def simple_merge(
             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:
+            if precision == "fp16":
                 thetas["model_a"].update({key: thetas["model_a"][key].half()})
 
     log_vram("after stage 2")
@@ -291,7 +291,7 @@ def rebasin_merge(
             thetas["model_a"],
             max_iter=it,
             init_perm=None,
-            usefp16=precision == 16,
+            usefp16=precision == "fp16",
             device=device,
         )
 
@@ -307,7 +307,7 @@ def rebasin_merge(
             thetas["model_a"],
             max_iter=it,
             init_perm=None,
-            usefp16=precision == 16,
+            usefp16=precision == "fp16",
             device=device,
         )
 
@@ -343,7 +343,7 @@ def merge_key(
     thetas: Dict,
     weight_matcher: WeightClass,
     merge_mode: str,
-    precision: int = 16,
+    precision: str = "fp16",
     weights_clip: bool = False,
     device: str = "cpu",
     work_device: Optional[str] = None,
@@ -378,7 +378,7 @@ def merge_key(
         if weights_clip:
             merged_key = clip_weights_key(thetas, merged_key, key)
 
-        if precision == 16:
+        if precision == "fp16":
             merged_key = merged_key.half()
 
         return merged_key

From 17d5b152246d7b42bc601a04e7e9d9dc3dcc5479 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Thu, 16 Nov 2023 21:00:27 -0600
Subject: [PATCH 12/25] Doc Strings for Merge Methods

---
 modules/merging/merge_methods.py | 90 ++++++++++++++++++++++----------
 1 file changed, 63 insertions(+), 27 deletions(-)

diff --git a/modules/merging/merge_methods.py b/modules/merging/merge_methods.py
index c10c459f0..a80a18106 100644
--- a/modules/merging/merge_methods.py
+++ b/modules/merging/merge_methods.py
@@ -24,12 +24,20 @@ EPSILON = 1e-10  # Define a small constant EPSILON to prevent division by zero
 
 
 def weighted_sum(a: Tensor, b: Tensor, alpha: float, **kwargs) -> Tensor:
+    """
+    Basic Merge:
+    alpha 0 returns Primary Model
+    alpha 1 returns Secondary Model
+    """
     return (1 - alpha) * a + alpha * b
 
 
-def weighted_subtraction(
-    a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs
-) -> Tensor:
+def weighted_subtraction(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    The inverse of a Weighted Sum Merge
+    Returns Primary Model when alpha*beta = 0
+    High values of alpha*beta are likely to break the merged model
+    """
     # 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
@@ -38,6 +46,12 @@ def weighted_subtraction(
 
 
 def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    Takes a slice of Secondary Model and pastes it into Primary Model
+    Alpha sets the width of the slice
+    Beta sets the start point of the slice
+    ie Alpha = 0.5 Beta = 0.25 is (ABBA) Alpha = 0.25 Beta = 0 is (BAAA)
+    """
     if alpha + beta <= 1:
         tt = a.clone()
         talphas = int(a.shape[0] * beta)
@@ -52,24 +66,37 @@ def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Ten
 
 
 def add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
+    """
+    Classic Add Difference Merge
+    """
     return a + alpha * (b - c)
 
 
-def sum_twice(
-    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
-) -> Tensor:
+def sum_twice(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    Stacked Basic Merge:
+    Equivalent to Merging Primary and Secondary @ alpha
+    Then merging the result with Tertiary @ beta
+    """
     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:
+def triple_sum(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    Weights Secondary and Tertiary at alpha and beta respectively
+    Fills in the rest with Primary
+    Expect odd results if alpha + beta > 1 as Primary will be merged with a negative ratio
+    """
     return (1 - alpha - beta) * a + alpha * b + beta * c
 
 
-def euclidean_add_difference(
-    a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs
-) -> Tensor:
+def euclidean_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
+    """
+    Subtract Primary and Secondary from Tertiary
+    Compare the remainders via Euclidean distance
+    Add to Tertiary
+    Note: Slow
+    """
     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))
@@ -84,18 +111,20 @@ def euclidean_add_difference(
     return c + distance / torch.linalg.norm(distance) * target_norm
 
 
-def multiply_difference(
-    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
-) -> Tensor:
+def multiply_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    Similar to Add Difference but with geometric mean instead of arithmatic mean
+    """
     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:
+def top_k_tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    Redistributes the largest weights of Secondary Model into Primary Model
+    """
     a_flat = torch.flatten(a)
     a_dist = torch.msort(a_flat)
     b_indices = torch.argsort(torch.flatten(b), stable=True)
@@ -144,9 +173,10 @@ def ratio_to_region(width: float, offset: float, n: int) -> Tuple[int, int, bool
     return round(start), round(end), inverted
 
 
-def similarity_add_difference(
-    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
-) -> Tensor:
+def similarity_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    Weighted Sum where A and B are similar and Add Difference where A and B are dissimilar
+    """
     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)
@@ -156,9 +186,14 @@ def similarity_add_difference(
     return (1 - similarity) * ab_diff + similarity * ab_sum
 
 
-def distribution_crossover(
-    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
-):
+def distribution_crossover(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs):
+    """
+    From the creator:
+    It's Primary high-passed + Secondary low-passed. Takes the fourrier transform of the weights of
+    Primary and Secondary when ordered with respect to Tertiary. Split the frequency domain
+    using a linear function. Alpha is the split frequency and Beta is the inclination of the line.
+    add everything under the line as the contribution of Primary and everything over the line as the contribution of Secondary
+    """
     if a.shape == ():
         return alpha * a + (1 - alpha) * b
 
@@ -183,9 +218,10 @@ def distribution_crossover(
     return x_values.reshape_as(a)
 
 
-def ties_add_difference(
-    a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs
-) -> Tensor:
+def ties_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+    """
+    An implementation of arXiv:2306.01708
+    """
     deltas = []
     signs = []
     for m in [a, b]:

From 9b80d5163099136c110ca3f820fddf6f1932a9f5 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Thu, 16 Nov 2023 21:21:52 -0600
Subject: [PATCH 13/25] Add Doc Strings to UI

---
 modules/merging/merge.py | 2 +-
 modules/ui_models.py     | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index ae68fba15..e0ed32145 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -96,7 +96,7 @@ def restore_sd_model(original_model: Dict, merged_model: Dict) -> Dict:
 
 
 def log_vram(txt=""):
-    log.info(f"{txt} VRAM: {modules.memstats.memory_stats()}")
+    log.debug(f"{txt} VRAM: {modules.memstats.memory_stats()}")
 
 
 def load_thetas(
diff --git a/modules/ui_models.py b/modules/ui_models.py
index a9c99e8e3..41d691e30 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -167,6 +167,8 @@ def create_ui():
                         with FormRow():
                             merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum",
                                                      label="Interpolation Method")
+                            merge_mode_help = ToolButton('❔', visible=True)
+                            merge_mode_docs = gr.Textbox(value=None, visible=False)
                         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,
@@ -323,6 +325,10 @@ def create_ui():
                     else:
                         return gr.Slider.update(value=None, visible=False)
 
+                def show_help(mode):
+                    doc = getattr(merge_methods, mode).__doc__
+                    return gr.update(value=doc, visible=True, interactive=False)
+
                 def load_presets(presets, ratio):
                     for i, p in enumerate(presets):
                         presets[i] = BLOCK_WEIGHTS_PRESETS[p]
@@ -339,7 +345,7 @@ def create_ui():
                         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)]
-
+                merge_mode_help.click(fn=show_help, inputs=merge_mode, outputs=merge_mode_docs)
                 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)

From 2515b832515d8b8995ec333f6ab6007178b6d8ad Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 17 Nov 2023 19:30:33 -0600
Subject: [PATCH 14/25] GUI Improvement

---
 html/locale_en.json  | 10 +++++++++-
 modules/ui_models.py | 18 +++++++++++-------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/html/locale_en.json b/html/locale_en.json
index 65f239e10..daf5785a9 100644
--- a/html/locale_en.json
+++ b/html/locale_en.json
@@ -14,6 +14,7 @@
   {"id":"","label":"⊗","localized":"","hint":"Clear prompt"},
   {"id":"","label":"🗁","localized":"","hint":"Show/hide extra networks"},
   {"id":"","label":"⇰","localized":"","hint":"Apply selected styles to current prompt"},
+  {"id":"","label":"⇨","localized":"","hint":"Apply preset to Manual Block Merge tab"},
   {"id":"","label":"⇩","localized":"","hint":"Save parameters from last generated image as style template"},
   {"id":"","label":"🕮","localized":"","hint":"Save parameters from last generated image as style template"},
   {"id":"","label":"⇕","localized":"","hint":"Sort by: Name asc/desc, Size largest/smallest, Time newest/oldest"},
@@ -255,7 +256,14 @@
   {"id":"","label":"specify model variant","localized":"","hint":""},
   {"id":"","label":"specify model revision","localized":"","hint":""},
   {"id":"","label":"huggingface token","localized":"","hint":""},
-  {"id":"","label":"huggingface mirror","localized":"","hint":""}
+  {"id":"","label":"huggingface mirror","localized":"","hint":""},
+  {"id":"","label":"Weights Clip","localized":"","hint":"Forced merged weights to be no heavier than the original model, preventing burn in and overly saturated models"},
+  {"id":"","label":"ReBasin","localized":"","hint":"Performs multiple merges with permutations in order to keep more features from both models"},
+  {"id":"","label":"Number of ReBasin Iterations","localized":"","hint":"Number of times to merge and permute the model before saving"},
+  {"id":"","label":"Base","localized":"","hint":"Text Encoder and a few unaligned keys (1 value)"},
+  {"id":"","label":"In Blocks","localized":"","hint":"Downsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"},
+  {"id":"","label":"Mid Block","localized":"","hint":"Central Block of the UNet (1 value)"},
+  {"id":"","label":"Out Block","localized":"","hint":"Upsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"}
 ],
 "train tabs": [
   {"id":"","label":"Preprocess","localized":"","hint":""},
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 41d691e30..16697df7c 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -167,8 +167,8 @@ def create_ui():
                         with FormRow():
                             merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum",
                                                      label="Interpolation Method")
-                            merge_mode_help = ToolButton('❔', visible=True)
-                            merge_mode_docs = gr.Textbox(value=None, visible=False)
+                        with FormRow():
+                            merge_mode_docs = gr.HTML(value=None, visible=False)
                         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,
@@ -244,7 +244,7 @@ def create_ui():
                         with FormRow():
                             precision = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision")
                         with FormRow():
-                            work_device = gr.Radio(choices=["cpu", "cuda"], value="cpu", label="Device")
+                            device = gr.Radio(choices=["cpu", "shuffle", "cuda"], value="cpu", label="Merge Device")
                         with FormRow():
                             bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None",
                                                       interactive=True, label="Bake in VAE")
@@ -283,7 +283,7 @@ def create_ui():
                                    prune,
                                    re_basin,
                                    re_basin_iterations,
-                                   work_device,
+                                   device,
                                    bake_in_vae):
                     kwargs = {}
                     for x in inspect.getfullargspec(MEHmodelmerger)[0]:
@@ -291,6 +291,9 @@ def create_ui():
                     for key in list(kwargs.keys()):
                         if kwargs[key] in [None, "None", "", 0, []]:
                             del kwargs[key]
+                    if kwargs["device"] == "shuffle":
+                        kwargs["device"] = "cpu"
+                        kwargs["work_device"] = "cuda"
 
                     try:
                         results = extras.run_MEHmodelmerger(dummy_component, **kwargs)
@@ -327,7 +330,7 @@ def create_ui():
 
                 def show_help(mode):
                     doc = getattr(merge_methods, mode).__doc__
-                    return gr.update(value=doc, visible=True, interactive=False)
+                    return gr.update(value=doc, visible=True)
 
                 def load_presets(presets, ratio):
                     for i, p in enumerate(presets):
@@ -345,7 +348,8 @@ def create_ui():
                         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)]
-                merge_mode_help.click(fn=show_help, inputs=merge_mode, outputs=merge_mode_docs)
+
+                merge_mode.change(fn=show_help, inputs=merge_mode, outputs=merge_mode_docs)
                 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)
@@ -390,7 +394,7 @@ def create_ui():
                         prune,
                         re_basin,
                         re_basin_iterations,
-                        work_device,
+                        device,
                         bake_in_vae,
                     ],
                     outputs=[

From abb4739590e117b937c4ca26a47e93ee4c742b68 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 17 Nov 2023 20:18:55 -0600
Subject: [PATCH 15/25] More UI Improvements

---
 html/locale_en.json  |  4 ++++
 modules/extras.py    |  3 +++
 modules/ui_models.py | 13 +++++++------
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/html/locale_en.json b/html/locale_en.json
index daf5785a9..e7289c456 100644
--- a/html/locale_en.json
+++ b/html/locale_en.json
@@ -260,10 +260,14 @@
   {"id":"","label":"Weights Clip","localized":"","hint":"Forced merged weights to be no heavier than the original model, preventing burn in and overly saturated models"},
   {"id":"","label":"ReBasin","localized":"","hint":"Performs multiple merges with permutations in order to keep more features from both models"},
   {"id":"","label":"Number of ReBasin Iterations","localized":"","hint":"Number of times to merge and permute the model before saving"},
+  {"id":"","label":"cpu","localized":"","hint":"Uses cpu and RAM only: slowest but least likely to OOM"},
+  {"id":"","label":"shuffle","localized":"","hint":"Loads full model in RAM and calculates on VRAM: Less speedup, suggested for SDXL merges"},
+  {"id":"","label":"cuda","localized":"","hint":"Loads models into VRAM automatically unloading current model: fastest option but unlikely to handle SDXL Models without OOM"},
   {"id":"","label":"Base","localized":"","hint":"Text Encoder and a few unaligned keys (1 value)"},
   {"id":"","label":"In Blocks","localized":"","hint":"Downsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"},
   {"id":"","label":"Mid Block","localized":"","hint":"Central Block of the UNet (1 value)"},
   {"id":"","label":"Out Block","localized":"","hint":"Upsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"}
+  {"id":"","label":"Preset Interpolation Ratio","localized":"","hint":"If two presets are selected, interpolate between them"}
 ],
 "train tabs": [
   {"id":"","label":"Preprocess","localized":"","hint":""},
diff --git a/modules/extras.py b/modules/extras.py
index 467946166..3dee5f6a2 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -311,6 +311,9 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             kwargs.pop("beta_mid_block", None)
             kwargs.pop("beta_out_blocks", None)
             kwargs.pop("beta_preset", None)
+            
+    if kwargs["device"] == "cuda":
+        sd_models.unload_model_weights()
 
     try:
         theta_0 = merge_models(**kwargs)
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 16697df7c..95fbe1522 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -160,6 +160,7 @@ def create_ui():
                 def sd_model_choices():
                     return ['None'] + sd_models.checkpoint_tiles()
 
+
                 with gr.Row(equal_height=False):
                     with gr.Column(variant='compact'):
                         with FormRow():
@@ -167,8 +168,7 @@ def create_ui():
                         with FormRow():
                             merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum",
                                                      label="Interpolation Method")
-                        with FormRow():
-                            merge_mode_docs = gr.HTML(value=None, visible=False)
+                            merge_mode_docs = gr.HTML(value=getattr(merge_methods, "weighted_sum").__doc__.replace("\n", "<br>"))
                         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,
@@ -322,15 +322,16 @@ def create_ui():
                     else:
                         return gr.Slider.update(value=None, visible=False)
 
+                def show_help(mode):
+                    doc = getattr(merge_methods, mode).__doc__.replace("\n", "<br>")
+                    return gr.update(value=doc, visible=True)
+
                 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 show_help(mode):
-                    doc = getattr(merge_methods, mode).__doc__
-                    return gr.update(value=doc, visible=True)
 
                 def load_presets(presets, ratio):
                     for i, p in enumerate(presets):
@@ -339,7 +340,7 @@ def create_ui():
                         preset = interpolate(presets, ratio)
                     else:
                         preset = presets[0]
-                    preset = ['%.3f' % x for x in preset]
+                    preset = ['%.3f' % x if int(x) != x else str(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)]
 

From cb8be0fc6e99f66c7b8c3d570a9fd4c835ea8cec Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 17 Nov 2023 20:35:06 -0600
Subject: [PATCH 16/25] Add torch.device

---
 modules/extras.py              |  2 +-
 modules/merging/merge.py       | 41 +++++++++++++++++-----------------
 modules/merging/merge_model.py |  6 ++---
 3 files changed, 23 insertions(+), 26 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 3dee5f6a2..c60568bf6 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -311,7 +311,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             kwargs.pop("beta_mid_block", None)
             kwargs.pop("beta_out_blocks", None)
             kwargs.pop("beta_preset", None)
-            
+
     if kwargs["device"] == "cuda":
         sd_models.unload_model_weights()
 
diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index e0ed32145..2330d50d5 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -1,4 +1,3 @@
-import gc
 import os
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
@@ -10,6 +9,7 @@ import torch
 from tqdm import tqdm
 
 import modules.memstats
+import modules.devices as devices
 from modules.shared import log
 from modules.merging import merge_methods
 from modules.merging.merge_utils import WeightClass
@@ -69,7 +69,7 @@ def fix_model(model: Dict) -> Dict:
     return fix_clip(model)
 
 
-def load_sd_model(model: os.PathLike | str, device: str = "cpu") -> Dict:
+def load_sd_model(model: os.PathLike | str, device: torch.device = None) -> Dict:
     if isinstance(model, str):
         model = Path(model)
 
@@ -102,7 +102,7 @@ def log_vram(txt=""):
 def load_thetas(
     models: Dict[str, os.PathLike | str],
     prune: bool,
-    device: str,
+    device: torch.device,
     precision: str,
 ) -> Dict:
     log_vram("before loading models")
@@ -111,13 +111,12 @@ def load_thetas(
     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)})
+    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
@@ -129,8 +128,8 @@ def merge_models(
     precision: str = "fp16",
     weights_clip: bool = False,
     re_basin: bool = False,
-    device: str = "cpu",
-    work_device: Optional[str] = None,
+    device: torch.device = None,
+    work_device: torch.device = None,
     prune: bool = False,
     threads: int = 1,
     **kwargs,
@@ -178,7 +177,7 @@ def un_prune_model(
     if prune:
         log.info("Un-pruning merged model")
         del thetas
-        gc.collect()
+        devices.torch_gc(force=True)
         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"):
@@ -189,7 +188,7 @@ def un_prune_model(
                 if precision == "fp16":
                     merged.update({key: merged[key].half()})
         del original_a
-        gc.collect()
+        devices.torch_gc(force=True)
         # 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"):
@@ -210,8 +209,8 @@ def simple_merge(
     merge_mode: str,
     precision: str = "fp16",
     weights_clip: bool = False,
-    device: str = "cpu",
-    work_device: Optional[str] = None,
+    device: torch.device = None,
+    work_device: torch.device = None,
     threads: int = 1,
 ) -> Dict:
     futures = []
@@ -257,8 +256,8 @@ def rebasin_merge(
     precision: str = "fp16",
     weights_clip: bool = False,
     iterations: int = 1,
-    device="cpu",
-    work_device=None,
+    device: torch.device = None,
+    work_device: torch.device = None,
     threads: int = 1,
 ):
     # WARNING: not sure how this does when 3 models are involved...
@@ -345,8 +344,8 @@ def merge_key(
     merge_mode: str,
     precision: str = "fp16",
     weights_clip: bool = False,
-    device: str = "cpu",
-    work_device: Optional[str] = None,
+    device: torch.device = None,
+    work_device: torch.device = None,
 ) -> Optional[Tuple[str, Dict]]:
     if work_device is None:
         work_device = device
@@ -413,7 +412,7 @@ def get_merge_method_args(
     current_bases: Dict,
     thetas: Dict,
     key: str,
-    work_device: str,
+    work_device: torch.device,
 ) -> Dict:
     merge_method_args = {
         "a": thetas["model_a"][key].to(work_device),
diff --git a/modules/merging/merge_model.py b/modules/merging/merge_model.py
index d2d6470ec..15e617808 100644
--- a/modules/merging/merge_model.py
+++ b/modules/merging/merge_model.py
@@ -1,4 +1,3 @@
-# import logging
 import os
 from dataclasses import dataclass
 
@@ -6,17 +5,16 @@ import safetensors
 import torch
 from tensordict import TensorDict
 
-# logging.getLogger("sd_meh").addHandler(logging.NullHandler())
 
 
 @dataclass
 class SDModel:
     model_path: os.PathLike
-    device: str
+    device: torch.device
 
     def load_model(self):
         # logging.info(f"Loading: {self.model_path}")
-        if self.model_path.suffix == ".safetensors":
+        if os.path.splitext(self.model_path) == ".safetensors":
             ckpt = safetensors.torch.load_file(
                 self.model_path,
                 device=self.device,

From 6aa859198f7fdea96e0e3d4c02fbd0087779edf3 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 17 Nov 2023 21:10:56 -0600
Subject: [PATCH 17/25] device fixes

---
 modules/extras.py    | 9 ++++++++-
 modules/ui_models.py | 5 -----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index c60568bf6..6f990cac8 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -11,7 +11,7 @@ import safetensors.torch
 from modules.merging.merge import merge_models
 from modules.merging.merge_utils import TRIPLE_METHODS
 
-from modules import shared, images, sd_models, sd_vae, sd_models_config
+from modules import shared, images, sd_models, sd_vae, sd_models_config, devices
 
 checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
 
@@ -313,7 +313,14 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             kwargs.pop("beta_preset", None)
 
     if kwargs["device"] == "cuda":
+        kwargs["device"] == devices.device
         sd_models.unload_model_weights()
+    elif kwargs["device"] == "shuffle":
+        kwargs["device"] = torch.device("cpu")
+        kwargs["work_device"] = devices.device
+    else:
+        kwargs["device"] = torch.device("cpu")
+
 
     try:
         theta_0 = merge_models(**kwargs)
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 95fbe1522..8f9746316 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -160,7 +160,6 @@ def create_ui():
                 def sd_model_choices():
                     return ['None'] + sd_models.checkpoint_tiles()
 
-
                 with gr.Row(equal_height=False):
                     with gr.Column(variant='compact'):
                         with FormRow():
@@ -291,9 +290,6 @@ def create_ui():
                     for key in list(kwargs.keys()):
                         if kwargs[key] in [None, "None", "", 0, []]:
                             del kwargs[key]
-                    if kwargs["device"] == "shuffle":
-                        kwargs["device"] = "cpu"
-                        kwargs["work_device"] = "cuda"
 
                     try:
                         results = extras.run_MEHmodelmerger(dummy_component, **kwargs)
@@ -332,7 +328,6 @@ def create_ui():
                     else:
                         return gr.Slider.update(value=None, visible=False)
 
-
                 def load_presets(presets, ratio):
                     for i, p in enumerate(presets):
                         presets[i] = BLOCK_WEIGHTS_PRESETS[p]

From c93d6c7e2669650ec3f90787daa25b72a753786b Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 17 Nov 2023 21:47:33 -0600
Subject: [PATCH 18/25] typos

---
 html/locale_en.json            | 2 +-
 modules/merging/merge_model.py | 2 +-
 modules/ui_models.py           | 1 -
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/html/locale_en.json b/html/locale_en.json
index e7289c456..45074cb63 100644
--- a/html/locale_en.json
+++ b/html/locale_en.json
@@ -266,7 +266,7 @@
   {"id":"","label":"Base","localized":"","hint":"Text Encoder and a few unaligned keys (1 value)"},
   {"id":"","label":"In Blocks","localized":"","hint":"Downsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"},
   {"id":"","label":"Mid Block","localized":"","hint":"Central Block of the UNet (1 value)"},
-  {"id":"","label":"Out Block","localized":"","hint":"Upsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"}
+  {"id":"","label":"Out Block","localized":"","hint":"Upsampling Blocks of the UNet (12 values for SD1.5, 9 values for SDXL)"},
   {"id":"","label":"Preset Interpolation Ratio","localized":"","hint":"If two presets are selected, interpolate between them"}
 ],
 "train tabs": [
diff --git a/modules/merging/merge_model.py b/modules/merging/merge_model.py
index 15e617808..37427c44b 100644
--- a/modules/merging/merge_model.py
+++ b/modules/merging/merge_model.py
@@ -14,7 +14,7 @@ class SDModel:
 
     def load_model(self):
         # logging.info(f"Loading: {self.model_path}")
-        if os.path.splitext(self.model_path) == ".safetensors":
+        if os.path.splitext(self.model_path)[1] == ".safetensors":
             ckpt = safetensors.torch.load_file(
                 self.model_path,
                 device=self.device,
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 8f9746316..e2c13a287 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -290,7 +290,6 @@ def create_ui():
                     for key in list(kwargs.keys()):
                         if kwargs[key] in [None, "None", "", 0, []]:
                             del kwargs[key]
-
                     try:
                         results = extras.run_MEHmodelmerger(dummy_component, **kwargs)
                     except Exception as e:

From cea7464553461faf3364574082db22e8464b89b9 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 18 Nov 2023 10:45:54 -0600
Subject: [PATCH 19/25] Use Internal read_state_dict

---
 modules/merging/merge.py | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index 2330d50d5..92446b2a8 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -1,7 +1,7 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from pathlib import Path
+# from pathlib import Path
 from typing import Dict, Optional, Tuple
 
 import safetensors.torch
@@ -11,9 +11,10 @@ from tqdm import tqdm
 import modules.memstats
 import modules.devices as devices
 from modules.shared import log
+from modules.sd_models import read_state_dict
 from modules.merging import merge_methods
 from modules.merging.merge_utils import WeightClass
-from modules.merging.merge_model import SDModel
+# from modules.merging.merge_model import SDModel
 from modules.merging.merge_rebasin import (
     apply_permutation,
     sdunet_permutation_spec,
@@ -69,11 +70,11 @@ def fix_model(model: Dict) -> Dict:
     return fix_clip(model)
 
 
-def load_sd_model(model: os.PathLike | str, device: torch.device = None) -> Dict:
-    if isinstance(model, str):
-        model = Path(model)
-
-    return SDModel(model, device).load_model()
+# def load_sd_model(model: os.PathLike | str, device: torch.device = None) -> Dict:
+#     if isinstance(model, str):
+#         model = Path(model)
+#
+#     return SDModel(model, device).load_model()
 
 
 def prune_sd_model(model: Dict) -> Dict:
@@ -107,9 +108,9 @@ def load_thetas(
 ) -> Dict:
     log_vram("before loading models")
     if prune:
-        thetas = {k: prune_sd_model(load_sd_model(m, "cpu")) for k, m in models.items()}
+        thetas = {k: prune_sd_model(read_state_dict(m, "cpu")) for k, m in models.items()}
     else:
-        thetas = {k: load_sd_model(m, device) for k, m in models.items()}
+        thetas = {k: read_state_dict(m, device) for k, m in models.items()}
 
     for model_key, model in thetas.items():
         for key, block in model.items():
@@ -170,7 +171,7 @@ def un_prune_model(
     merged: Dict,
     thetas: Dict,
     models: Dict,
-    device: str,
+    device: torch.device,
     prune: bool,
     precision: str,
 ) -> Dict:
@@ -179,7 +180,7 @@ def un_prune_model(
         del thetas
         devices.torch_gc(force=True)
         log_vram("remove thetas")
-        original_a = load_sd_model(models["model_a"], device)
+        original_a = read_state_dict(models["model_a"], device)
         for key in tqdm(original_a.keys(), desc="un-prune model a"):
             if KEY_POSITION_IDS in key:
                 continue
@@ -190,7 +191,7 @@ def un_prune_model(
         del original_a
         devices.torch_gc(force=True)
         # log_vram("remove original_a")
-        original_b = load_sd_model(models["model_b"], device)
+        original_b = read_state_dict(models["model_b"], device)
         for key in tqdm(original_b.keys(), desc="un-prune model b"):
             if KEY_POSITION_IDS in key:
                 continue

From d2d54af7aca1d4868b98ada8f2d516f0359f2432 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 18 Nov 2023 11:04:43 -0600
Subject: [PATCH 20/25] Cleanup and Credits

---
 modules/merging/merge.py       | 42 +++++---------------------
 modules/merging/merge_model.py | 54 ----------------------------------
 modules/merging/merge_utils.py |  1 -
 3 files changed, 7 insertions(+), 90 deletions(-)
 delete mode 100644 modules/merging/merge_model.py

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index 92446b2a8..a057ed015 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -1,26 +1,27 @@
 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
-
 import modules.memstats
 import modules.devices as devices
 from modules.shared import log
 from modules.sd_models import read_state_dict
 from modules.merging import merge_methods
 from modules.merging.merge_utils import WeightClass
-# from modules.merging.merge_model import SDModel
 from modules.merging.merge_rebasin import (
     apply_permutation,
     sdunet_permutation_spec,
     update_model_a,
     weight_matching,
 )
+##########################################################
+# Files in modules.merging are heavily modified
+# versions of sd-meh by @s1dxl used with his blessing
+# orginal code can be found @ https://github.com/s1dlx/meh
+##########################################################
 
 MAX_TOKENS = 77
 
@@ -36,12 +37,6 @@ KEY_POSITION_IDS = ".".join(
 )
 
 
-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():
@@ -54,29 +49,6 @@ def fix_clip(model: Dict) -> Dict:
     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: torch.device = None) -> 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:
@@ -201,7 +173,7 @@ def un_prune_model(
                     merged.update({key: merged[key].half()})
         del original_b
 
-    return fix_model(merged)
+    return fix_clip(merged)
 
 
 def simple_merge(
@@ -247,7 +219,7 @@ def simple_merge(
 
     log_vram("after stage 2")
 
-    return fix_model(thetas["model_a"])
+    return fix_clip(thetas["model_a"])
 
 
 def rebasin_merge(
diff --git a/modules/merging/merge_model.py b/modules/merging/merge_model.py
deleted file mode 100644
index 37427c44b..000000000
--- a/modules/merging/merge_model.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import os
-from dataclasses import dataclass
-
-import safetensors
-import torch
-from tensordict import TensorDict
-
-
-
-@dataclass
-class SDModel:
-    model_path: os.PathLike
-    device: torch.device
-
-    def load_model(self):
-        # logging.info(f"Loading: {self.model_path}")
-        if os.path.splitext(self.model_path)[1] == ".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
diff --git a/modules/merging/merge_utils.py b/modules/merging/merge_utils.py
index 41b6f7442..da0059967 100644
--- a/modules/merging/merge_utils.py
+++ b/modules/merging/merge_utils.py
@@ -1,5 +1,4 @@
 import inspect
-# import logging
 import re
 from modules.merging import merge_methods
 from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS

From 026790fe4c746b13ca2b0f41298f54f873abec12 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 18 Nov 2023 11:14:15 -0600
Subject: [PATCH 21/25] Cleanup and optimization

---
 modules/extras.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 6f990cac8..672db16e4 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -286,7 +286,8 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         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:
+    except AssertionError as e:
+        shared.log.info(e)
         kwargs["alpha"] = kwargs.get("alpha_preset", kwargs["alpha"])
     finally:
         kwargs.pop("alpha_base", None)
@@ -303,7 +304,8 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             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:
+        except AssertionError as e:
+            shared.log.info(e)
             kwargs["beta"] = kwargs.get("beta_preset", kwargs["beta"])
         finally:
             kwargs.pop("beta_base", None)
@@ -313,7 +315,7 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             kwargs.pop("beta_preset", None)
 
     if kwargs["device"] == "cuda":
-        kwargs["device"] == devices.device
+        kwargs["device"] = devices.device
         sd_models.unload_model_weights()
     elif kwargs["device"] == "shuffle":
         kwargs["device"] = torch.device("cpu")
@@ -321,7 +323,6 @@ def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     else:
         kwargs["device"] = torch.device("cpu")
 
-
     try:
         theta_0 = merge_models(**kwargs)
     except Exception as e:
@@ -390,6 +391,7 @@ 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()
+    devices.torch_gc(force=True)
     shared.log.info(f"Model merge saved: {output_modelname}.")
     shared.state.textinfo = "Checkpoint saved"
     shared.state.end()

From d9d3c2e7fce159514be487c2c2e370afa2d3309a Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 18 Nov 2023 11:33:20 -0600
Subject: [PATCH 22/25] Replace Original Merge Method Major Cleanup

---
 modules/extras.py        | 204 +--------------------------------------
 modules/merging/merge.py |  11 +--
 modules/ui_models.py     | 155 +++++++----------------------
 3 files changed, 46 insertions(+), 324 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 672db16e4..2b2852844 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,5 +1,4 @@
 import os
-import re
 import html
 import json
 import shutil
@@ -13,8 +12,6 @@ from modules.merging.merge_utils import TRIPLE_METHODS
 
 from modules import shared, images, sd_models, sd_vae, sd_models_config, devices
 
-checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
-
 
 def run_pnginfo(image):
     if image is None:
@@ -55,202 +52,7 @@ def to_half(tensor, enable):
     return tensor
 
 
-def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier,
-                    save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights,
-                    save_metadata):  # pylint: disable=unused-argument
-    shared.state.begin('merge')
-    save_as_half = save_as_half == 0
-
-    def fail(message):
-        shared.state.textinfo = message
-        shared.state.end()
-        return [*[gr.update() for _ in range(4)], message]
-
-    def weighted_sum(theta0, theta1, alpha):
-        return ((1 - alpha) * theta0) + (alpha * theta1)
-
-    def get_difference(theta1, theta2):
-        return theta1 - theta2
-
-    def add_difference(theta0, theta1_2_diff, alpha):
-        return theta0 + (alpha * theta1_2_diff)
-
-    def filename_weighted_sum():
-        a = primary_model_info.model_name
-        b = secondary_model_info.model_name
-        Ma = round(1 - multiplier, 2)
-        Mb = round(multiplier, 2)
-        return f"{Ma}({a}) + {Mb}({b})"
-
-    def filename_add_difference():
-        a = primary_model_info.model_name
-        b = secondary_model_info.model_name
-        c = tertiary_model_info.model_name
-        M = round(multiplier, 2)
-        return f"{a} + {M}({b} - {c})"
-
-    def filename_nothing():
-        return primary_model_info.model_name
-
-    theta_funcs = {
-        "Weighted sum": (filename_weighted_sum, None, weighted_sum),
-        "Add difference": (filename_add_difference, get_difference, add_difference),
-        "No interpolation": (filename_nothing, None, None),
-    }
-    filename_generator, theta_func1, theta_func2 = theta_funcs[interp_method]
-    shared.state.job_count = (1 if theta_func1 else 0) + (1 if theta_func2 else 0)
-    if not primary_model_name or primary_model_name == 'None':
-        return fail("Failed: Merging requires a primary model.")
-    primary_model_info = sd_models.checkpoints_list[primary_model_name]
-    if theta_func2 and (not secondary_model_name or secondary_model_name == 'None'):
-        return fail("Failed: Merging requires a secondary model.")
-    secondary_model_info = sd_models.checkpoints_list[secondary_model_name] if theta_func2 else None
-    if theta_func1 and (not tertiary_model_name or tertiary_model_name == 'None'):
-        return fail(f"Failed: Interpolation method ({interp_method}) requires a tertiary model.")
-    tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None
-    result_is_inpainting_model = False
-    result_is_instruct_pix2pix_model = False
-    if theta_func2:
-        shared.state.textinfo = "Loading B"
-        shared.log.info(f"Model merge loading secondary model: {secondary_model_info.filename}")
-        theta_1 = sd_models.read_state_dict(secondary_model_info.filename)
-    else:
-        theta_1 = None
-    if theta_func1:
-        shared.state.textinfo = "Loading C"
-        shared.log.info(f"Model merge loading tertiary model: {tertiary_model_info.filename}")
-        theta_2 = sd_models.read_state_dict(tertiary_model_info.filename)
-        shared.state.textinfo = 'Merging B and C'
-        shared.state.sampling_steps = len(theta_1.keys())
-        for key in tqdm.tqdm(theta_1.keys()):
-            if key in checkpoint_dict_skip_on_merge:
-                continue
-            if 'model' in key:
-                if key in theta_2:
-                    t2 = theta_2.get(key, torch.zeros_like(theta_1[key]))
-                    theta_1[key] = theta_func1(theta_1[key], t2)
-                else:
-                    theta_1[key] = torch.zeros_like(theta_1[key])
-            shared.state.sampling_step += 1
-        del theta_2
-        shared.state.nextjob()
-    shared.state.textinfo = f"Loading {primary_model_info.filename}..."
-    shared.log.info(f"Model merge loading primary model: {primary_model_info.filename}")
-    theta_0 = sd_models.read_state_dict(primary_model_info.filename)
-    shared.log.info("Model merge: running")
-    shared.state.textinfo = 'Merging A and B'
-    shared.state.sampling_steps = len(theta_0.keys())
-    for key in tqdm.tqdm(theta_0.keys()):
-        if theta_1 and 'model' in key and key in theta_1:
-            if key in checkpoint_dict_skip_on_merge:
-                continue
-            a = theta_0[key]
-            b = theta_1[key]
-            # this enables merging an inpainting model (A) with another one (B);
-            # where normal model would have 4 channels, for latenst space, inpainting model would
-            # have another 4 channels for unmasked picture's latent space, plus one channel for mask, for a total of 9
-            if a.shape != b.shape and a.shape[0:1] + a.shape[2:] == b.shape[0:1] + b.shape[2:]:
-                if a.shape[1] == 4 and b.shape[1] == 9:
-                    raise RuntimeError(
-                        "When merging inpainting model with a normal one, A must be the inpainting model.")
-                if a.shape[1] == 4 and b.shape[1] == 8:
-                    raise RuntimeError(
-                        "When merging instruct-pix2pix model with a normal one, A must be the instruct-pix2pix model.")
-                if a.shape[1] == 8 and b.shape[1] == 4:  # If we have an Instruct-Pix2Pix model...
-                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b,
-                                                             multiplier)  # Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.
-                    result_is_instruct_pix2pix_model = True
-                else:
-                    assert a.shape[1] == 9 and b.shape[
-                        1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}"
-                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)
-                    result_is_inpainting_model = True
-            else:
-                theta_0[key] = theta_func2(a, b, multiplier)
-            theta_0[key] = to_half(theta_0[key], save_as_half)
-        shared.state.sampling_step += 1
-    del theta_1
-    bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None)
-    if bake_in_vae_filename is not None:
-        shared.log.info(f"Model merge: baking in VAE: {bake_in_vae_filename}")
-        shared.state.textinfo = 'Baking in VAE'
-        vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename)
-        for key in vae_dict.keys():
-            theta_0_key = 'first_stage_model.' + key
-            if theta_0_key in theta_0:
-                theta_0[theta_0_key] = to_half(vae_dict[key], save_as_half)
-        del vae_dict
-    if save_as_half and not theta_func2:
-        for key in theta_0.keys():
-            theta_0[key] = to_half(theta_0[key], save_as_half)
-    if discard_weights:
-        regex = re.compile(discard_weights)
-        for key in list(theta_0):
-            if re.search(regex, key):
-                theta_0.pop(key, None)
-    ckpt_dir = shared.opts.ckpt_dir or sd_models.model_path
-    filename = filename_generator() if custom_name == '' else custom_name
-    filename += ".inpainting" if result_is_inpainting_model else ""
-    filename += ".instruct-pix2pix" if result_is_instruct_pix2pix_model else ""
-    filename += "." + checkpoint_format
-    output_modelname = os.path.join(ckpt_dir, filename)
-    shared.state.nextjob()
-    shared.state.textinfo = "Saving"
-    metadata = None
-    if save_metadata:
-        metadata = {"format": "pt", "sd_merge_models": {}}
-        merge_recipe = {
-            "type": "webui",  # 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,
-            "config_source": config_source,
-            "bake_in_vae": bake_in_vae,
-            "discard_weights": discard_weights,
-            "is_inpainting": result_is_inpainting_model,
-            "is_instruct_pix2pix": result_is_instruct_pix2pix_model
-        }
-        metadata["sd_merge_recipe"] = json.dumps(merge_recipe)
-
-        def add_model_metadata(checkpoint_info):
-            checkpoint_info.calculate_shorthash()
-            metadata["sd_merge_models"][checkpoint_info.sha256] = {
-                "name": checkpoint_info.name,
-                "legacy_hash": checkpoint_info.hash,
-                "sd_merge_recipe": checkpoint_info.metadata.get("sd_merge_recipe", None)
-            }
-            metadata["sd_merge_models"].update(checkpoint_info.metadata.get("sd_merge_models", {}))
-
-        add_model_metadata(primary_model_info)
-        if secondary_model_info:
-            add_model_metadata(secondary_model_info)
-        if tertiary_model_info:
-            add_model_metadata(tertiary_model_info)
-        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:
-        torch.save(theta_0, output_modelname)
-    sd_models.list_models()
-    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()
-    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],
-            "Checkpoint saved to " + output_modelname]
-
-
-def run_MEHmodelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
+def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     shared.state.begin('merge')
 
     def fail(message):
@@ -377,6 +179,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)
+
     try:
         theta_0 = theta_0.to_dict()
     except:
@@ -391,7 +194,8 @@ 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()
-    devices.torch_gc(force=True)
+    if kwargs["device"].type != "cpu":
+        devices.torch_gc(force=True)
     shared.log.info(f"Model merge saved: {output_modelname}.")
     shared.state.textinfo = "Checkpoint saved"
     shared.state.end()
diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index a057ed015..f0cb89ce5 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -4,6 +4,7 @@ from contextlib import contextmanager
 from typing import Dict, Optional, Tuple
 import safetensors.torch
 import torch
+from tensordict import TensorDict
 from tqdm import tqdm
 import modules.memstats
 import modules.devices as devices
@@ -37,7 +38,6 @@ KEY_POSITION_IDS = ".".join(
 )
 
 
-
 def fix_clip(model: Dict) -> Dict:
     if KEY_POSITION_IDS in model.keys():
         model[KEY_POSITION_IDS] = torch.tensor(
@@ -80,9 +80,9 @@ def load_thetas(
 ) -> Dict:
     log_vram("before loading models")
     if prune:
-        thetas = {k: prune_sd_model(read_state_dict(m, "cpu")) for k, m in models.items()}
+        thetas = {k: prune_sd_model(TensorDict.from_dict(read_state_dict(m, "cpu"))) for k, m in models.items()}
     else:
-        thetas = {k: read_state_dict(m, device) for k, m in models.items()}
+        thetas = {k: TensorDict.from_dict(read_state_dict(m, device)) for k, m in models.items()}
 
     for model_key, model in thetas.items():
         for key, block in model.items():
@@ -152,7 +152,7 @@ def un_prune_model(
         del thetas
         devices.torch_gc(force=True)
         log_vram("remove thetas")
-        original_a = read_state_dict(models["model_a"], device)
+        original_a = TensorDict.from_dict(read_state_dict(models["model_a"], device))
         for key in tqdm(original_a.keys(), desc="un-prune model a"):
             if KEY_POSITION_IDS in key:
                 continue
@@ -162,8 +162,7 @@ def un_prune_model(
                     merged.update({key: merged[key].half()})
         del original_a
         devices.torch_gc(force=True)
-        # log_vram("remove original_a")
-        original_b = read_state_dict(models["model_b"], device)
+        original_b = TensorDict.from_dict(read_state_dict(models["model_b"], device))
         for key in tqdm(original_b.keys(), desc="un-prune model b"):
             if KEY_POSITION_IDS in key:
                 continue
diff --git a/modules/ui_models.py b/modules/ui_models.py
index e2c13a287..509e828a6 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -4,7 +4,7 @@ import inspect
 from datetime import datetime
 import gradio as gr
 from modules import sd_models, sd_vae, extras
-from modules.ui_components import FormRow, ToolButton, InputAccordion
+from modules.ui_components import FormRow, ToolButton
 from modules.ui_common import create_refresh_button
 from modules.call_queue import wrap_gradio_gpu_call
 from modules.shared import opts, log, req
@@ -75,88 +75,6 @@ def create_ui():
                 )
 
             with gr.Tab(label="Merge"):
-                with gr.Row(equal_height=False):
-                    with gr.Column(variant='compact'):
-                        with FormRow():
-                            custom_name = gr.Textbox(label="New model name")
-                        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")
-                            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")
-                        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)
-                        with FormRow():
-                            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")
-                        with FormRow():
-                            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")
-                        with FormRow():
-                            discard_weights = gr.Textbox(value="", label="Discard weights with matching name")
-                        with FormRow():
-                            save_metadata = gr.Checkbox(value=True, label="Save metadata")
-                        with gr.Row():
-                            modelmerger_merge = gr.Button(value="Merge", variant='primary')
-
-                def modelmerger(*args):
-                    try:
-                        results = extras.run_modelmerger(*args)
-                    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 results
-
-                modelmerger_merge.click(
-                    fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
-                    _js='modelmerger',
-                    inputs=[
-                        dummy_component,
-                        primary_model_name,
-                        secondary_model_name,
-                        tertiary_model_name,
-                        interp_method,
-                        interp_amount,
-                        save_as_half,
-                        custom_name,
-                        checkpoint_format,
-                        config_source,
-                        bake_in_vae,
-                        discard_weights,
-                        save_metadata,
-                    ],
-                    outputs=[
-                        primary_model_name,
-                        secondary_model_name,
-                        tertiary_model_name,
-                        dummy_component,
-                        models_outcome,
-                    ]
-                )
-            with gr.Tab(label="Advanced Merge"):
                 def sd_model_choices():
                     return ['None'] + sd_models.checkpoint_tiles()
 
@@ -167,7 +85,8 @@ def create_ui():
                         with FormRow():
                             merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum",
                                                      label="Interpolation Method")
-                            merge_mode_docs = gr.HTML(value=getattr(merge_methods, "weighted_sum").__doc__.replace("\n", "<br>"))
+                            merge_mode_docs = gr.HTML(
+                                value=getattr(merge_methods, "weighted_sum").__doc__.replace("\n", "<br>"))
                         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,
@@ -253,45 +172,45 @@ def create_ui():
                         with FormRow():
                             save_metadata = gr.Checkbox(value=True, label="Save metadata")
                         with gr.Row():
-                            MEHmodelmerger_merge = gr.Button(value="Merge", variant='primary')
+                            modelmerger_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):
+                def modelmerger(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):
                     kwargs = {}
-                    for x in inspect.getfullargspec(MEHmodelmerger)[0]:
+                    for x in inspect.getfullargspec(modelmerger)[0]:
                         kwargs[x] = locals()[x]
                     for key in list(kwargs.keys()):
                         if kwargs[key] in [None, "None", "", 0, []]:
                             del kwargs[key]
                     try:
-                        results = extras.run_MEHmodelmerger(dummy_component, **kwargs)
+                        results = extras.run_modelmerger(dummy_component, **kwargs)
                     except Exception as e:
                         modules.errors.display(e, 'model merge')
                         sd_models.list_models()  # to remove the potentially missing models from the list
@@ -358,8 +277,8 @@ def create_ui():
                 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)]),
+                modelmerger_merge.click(
+                    fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
                     _js='modelmerger',
                     inputs=[
                         dummy_component,

From f2b66b8720f8b9f3005a7cb6d7c0c5d945ee3209 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 18 Nov 2023 11:45:54 -0600
Subject: [PATCH 23/25] Error handling

---
 modules/extras.py | 21 +++++++++------------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 2b2852844..cc57de2db 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -85,11 +85,10 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         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"
+        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 AssertionError as e:
-        shared.log.info(e)
+    except Exception as e:
+        shared.log.warn(f"Merge: {e}")
         kwargs["alpha"] = kwargs.get("alpha_preset", kwargs["alpha"])
     finally:
         kwargs.pop("alpha_base", None)
@@ -101,13 +100,11 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
     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"
+                    [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 AssertionError as e:
-            shared.log.info(e)
+        except Exception as e:
+            shared.log.warn(f"Merge: {e}")
             kwargs["beta"] = kwargs.get("beta_preset", kwargs["beta"])
         finally:
             kwargs.pop("beta_base", None)
@@ -132,7 +129,7 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
 
     bake_in_vae_filename = sd_vae.vae_dict.get(kwargs.get("bake_in_vae", None), None)
     if bake_in_vae_filename is not None:
-        shared.log.info(f"Model merge: baking in VAE: {bake_in_vae_filename}")
+        shared.log.info(f"Merge: baking in VAE: {bake_in_vae_filename}")
         shared.state.textinfo = 'Baking in VAE'
         vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename)
         for key in vae_dict.keys():
@@ -196,7 +193,7 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         created_model.calculate_shorthash()
     if kwargs["device"].type != "cpu":
         devices.torch_gc(force=True)
-    shared.log.info(f"Model merge saved: {output_modelname}.")
+    shared.log.info(f"Merge saved: {output_modelname}.")
     shared.state.textinfo = "Checkpoint saved"
     shared.state.end()
     return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)],

From c5371a5087bb5bf72901cb5f83d06cbe0eca4535 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 18 Nov 2023 12:05:15 -0600
Subject: [PATCH 24/25] pylint fixes

---
 modules/merging/merge.py         |  2 +-
 modules/merging/merge_methods.py | 24 ++++++++++++------------
 modules/merging/merge_rebasin.py | 14 +++++++-------
 modules/merging/merge_utils.py   |  1 -
 4 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index f0cb89ce5..289d069af 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -309,7 +309,7 @@ def simple_merge_key(progress, key, thetas, *args, **kwargs):
         progress.update()
 
 
-def merge_key(
+def merge_key(  # pylint: disable=inconsistent-return-statements
     key: str,
     thetas: Dict,
     weight_matcher: WeightClass,
diff --git a/modules/merging/merge_methods.py b/modules/merging/merge_methods.py
index a80a18106..3f704c20f 100644
--- a/modules/merging/merge_methods.py
+++ b/modules/merging/merge_methods.py
@@ -23,7 +23,7 @@ __all__ = [
 EPSILON = 1e-10  # Define a small constant EPSILON to prevent division by zero
 
 
-def weighted_sum(a: Tensor, b: Tensor, alpha: float, **kwargs) -> Tensor:
+def weighted_sum(a: Tensor, b: Tensor, alpha: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Basic Merge:
     alpha 0 returns Primary Model
@@ -32,7 +32,7 @@ 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:
+def weighted_subtraction(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     The inverse of a Weighted Sum Merge
     Returns Primary Model when alpha*beta = 0
@@ -45,7 +45,7 @@ def weighted_subtraction(a: Tensor, b: Tensor, alpha: float, beta: float, **kwar
     return (a - alpha * beta * b) / (1 - alpha * beta)
 
 
-def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Takes a slice of Secondary Model and pastes it into Primary Model
     Alpha sets the width of the slice
@@ -65,14 +65,14 @@ def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Ten
     return tt
 
 
-def add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
+def add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Classic Add Difference Merge
     """
     return a + alpha * (b - c)
 
 
-def sum_twice(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+def sum_twice(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Stacked Basic Merge:
     Equivalent to Merging Primary and Secondary @ alpha
@@ -81,7 +81,7 @@ def sum_twice(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwar
     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:
+def triple_sum(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Weights Secondary and Tertiary at alpha and beta respectively
     Fills in the rest with Primary
@@ -90,7 +90,7 @@ def triple_sum(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwa
     return (1 - alpha - beta) * a + alpha * b + beta * c
 
 
-def euclidean_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
+def euclidean_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Subtract Primary and Secondary from Tertiary
     Compare the remainders via Euclidean distance
@@ -111,7 +111,7 @@ def euclidean_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kw
     return c + distance / torch.linalg.norm(distance) * target_norm
 
 
-def multiply_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+def multiply_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Similar to Add Difference but with geometric mean instead of arithmatic mean
     """
@@ -121,7 +121,7 @@ def multiply_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: flo
     return c + difference.to(c.dtype)
 
 
-def top_k_tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+def top_k_tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Redistributes the largest weights of Secondary Model into Primary Model
     """
@@ -173,7 +173,7 @@ def ratio_to_region(width: float, offset: float, n: int) -> Tuple[int, int, bool
     return round(start), round(end), inverted
 
 
-def similarity_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+def similarity_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     Weighted Sum where A and B are similar and Add Difference where A and B are dissimilar
     """
@@ -186,7 +186,7 @@ def similarity_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, bet
     return (1 - similarity) * ab_diff + similarity * ab_sum
 
 
-def distribution_crossover(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs):
+def distribution_crossover(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs):  # pylint: disable=unused-argument
     """
     From the creator:
     It's Primary high-passed + Secondary low-passed. Takes the fourrier transform of the weights of
@@ -218,7 +218,7 @@ def distribution_crossover(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta:
     return x_values.reshape_as(a)
 
 
-def ties_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
+def ties_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:  # pylint: disable=unused-argument
     """
     An implementation of arXiv:2306.01708
     """
diff --git a/modules/merging/merge_rebasin.py b/modules/merging/merge_rebasin.py
index 26aac9964..2606a52f9 100644
--- a/modules/merging/merge_rebasin.py
+++ b/modules/merging/merge_rebasin.py
@@ -31,23 +31,23 @@ def permutation_spec_from_axes_to_perm(axes_to_perm: dict) -> PermutationSpec:
 
 
 def sdunet_permutation_spec() -> PermutationSpec:
-    conv = lambda name, p_in, p_out: {
+    conv = lambda name, p_in, p_out: {  # pylint: disable=unnecessary-lambda-assignment
         f"{name}.weight": (
             p_out,
             p_in,
         ),
         f"{name}.bias": (p_out,),
     }
-    norm = lambda name, p: {f"{name}.weight": (p,), f"{name}.bias": (p,)}
+    norm = lambda name, p: {f"{name}.weight": (p,), f"{name}.bias": (p,)}  # pylint: disable=unnecessary-lambda-assignment
     dense = (
-        lambda name, p_in, p_out, bias=True: {
+        lambda name, p_in, p_out, bias=True: {  # pylint: disable=unnecessary-lambda-assignment
             f"{name}.weight": (p_out, p_in),
             f"{name}.bias": (p_out,),
         }
         if bias
         else {f"{name}.weight": (p_out, p_in)}
     )
-    skip = lambda name, p_in, p_out: {
+    skip = lambda name, p_in, p_out: {  # pylint: disable=unnecessary-lambda-assignment
         f"{name}": (
             p_out,
             p_in,
@@ -57,7 +57,7 @@ def sdunet_permutation_spec() -> PermutationSpec:
     }
 
     # Unet Res blocks
-    easyblock = lambda name, p_in, p_out: {
+    easyblock = lambda name, p_in, p_out: {  # pylint: disable=unnecessary-lambda-assignment
         **norm(f"{name}.in_layers.0", p_in),
         **conv(f"{name}.in_layers.2", p_in, f"P_{name}_inner"),
         **dense(
@@ -68,7 +68,7 @@ def sdunet_permutation_spec() -> PermutationSpec:
     }
 
     # Text Encoder blocks
-    easyblock2 = lambda name, p: {
+    easyblock2 = lambda name, p: {  # pylint: disable=unnecessary-lambda-assignment
         **norm(f"{name}.norm1", p),
         **conv(f"{name}.conv1", p, f"P_{name}_inner"),
         **norm(f"{name}.norm2", f"P_{name}_inner"),
@@ -76,7 +76,7 @@ def sdunet_permutation_spec() -> PermutationSpec:
     }
 
     # This is for blocks that use a residual connection, but change the number of channels via a Conv.
-    shortcutblock = lambda name, p_in, p_out: {
+    shortcutblock = lambda name, p_in, p_out: {  # pylint: disable=unnecessary-lambda-assignment
         **norm(f"{name}.norm1", p_in),
         **conv(f"{name}.conv1", p_in, f"P_{name}_inner"),
         **norm(f"{name}.norm2", f"P_{name}_inner"),
diff --git a/modules/merging/merge_utils.py b/modules/merging/merge_utils.py
index da0059967..5bdfee246 100644
--- a/modules/merging/merge_utils.py
+++ b/modules/merging/merge_utils.py
@@ -109,4 +109,3 @@ class WeightClass:
 
     def set_it(self, it):
         self.it = it
-        return

From 581ab58bdd26420a9cad88fa84c91a9c2cc0bbad Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sun, 19 Nov 2023 12:03:13 -0600
Subject: [PATCH 25/25] Fix item 1,2,3,5,6

---
 modules/extras.py        | 26 ++++++++++++++------------
 modules/merging/merge.py | 14 +++++++++++---
 modules/ui_models.py     | 16 +++++++++++++---
 3 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index cc57de2db..a0c048b4e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -61,24 +61,24 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
         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,
+        "model_a": sd_models.get_closet_checkpoint_match(kwargs.get("primary_model_name", None)).filename,
+        "model_b": sd_models.get_closet_checkpoint_match(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)]
+    primary_model_info = sd_models.get_closet_checkpoint_match(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)]
+    secondary_model_info = sd_models.get_closet_checkpoint_match(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[kwargs.get("tertiary_model_name", None)] if kwargs.get("merge_mode", None) in TRIPLE_METHODS else None
+    tertiary_model_info = sd_models.get_closet_checkpoint_match(kwargs.get("tertiary_model_name", None)) 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:
-        kwargs["models"] |= {"model_c": sd_models.checkpoints_list[kwargs.get("tertiary_model_name", None)].filename}
+        kwargs["models"] |= {"model_c": sd_models.get_closet_checkpoint_match(kwargs.get("tertiary_model_name", None)).filename}
         del kwargs["tertiary_model_name"]
 
     try:
@@ -113,20 +113,26 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
             kwargs.pop("beta_out_blocks", None)
             kwargs.pop("beta_preset", None)
 
-    if kwargs["device"] == "cuda":
+    if kwargs["device"] == "gpu":
         kwargs["device"] = devices.device
-        sd_models.unload_model_weights()
     elif kwargs["device"] == "shuffle":
         kwargs["device"] = torch.device("cpu")
         kwargs["work_device"] = devices.device
     else:
         kwargs["device"] = torch.device("cpu")
+    if kwargs.pop("unload", False):
+        sd_models.unload_model_weights()
 
     try:
         theta_0 = merge_models(**kwargs)
     except Exception as e:
         return fail(f"{e}")
 
+    try:
+        theta_0 = theta_0.to_dict() #TensorDict -> Dict if necessary
+    except:
+        pass
+
     bake_in_vae_filename = sd_vae.vae_dict.get(kwargs.get("bake_in_vae", None), None)
     if bake_in_vae_filename is not None:
         shared.log.info(f"Merge: baking in VAE: {bake_in_vae_filename}")
@@ -177,10 +183,6 @@ def run_modelmerger(id_task, **kwargs):  # pylint: disable=unused-argument
 
     _, extension = os.path.splitext(output_modelname)
 
-    try:
-        theta_0 = theta_0.to_dict()
-    except:
-        pass
 
     if extension.lower() == ".safetensors":
         safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index 289d069af..09d03685b 100644
--- a/modules/merging/merge.py
+++ b/modules/merging/merge.py
@@ -54,7 +54,7 @@ def prune_sd_model(model: Dict) -> Dict:
     for k in keys:
         if (
             not k.startswith("model.diffusion_model.")
-            and not k.startswith("first_stage_model.")
+            # and not k.startswith("first_stage_model.")
             and not k.startswith("cond_stage_model.")
         ):
             del model[k]
@@ -153,23 +153,31 @@ def un_prune_model(
         devices.torch_gc(force=True)
         log_vram("remove thetas")
         original_a = TensorDict.from_dict(read_state_dict(models["model_a"], device))
-        for key in tqdm(original_a.keys(), desc="un-prune model a"):
+        unpruned = 0
+        for key in original_a.keys():
             if KEY_POSITION_IDS in key:
                 continue
             if "model" in key and key not in merged.keys():
                 merged.update({key: original_a[key]})
+                unpruned += 1
                 if precision == "fp16":
                     merged.update({key: merged[key].half()})
+        if unpruned != 0:
+            log.info(f"Merge: {unpruned} unmerged keys restored from Primary Model")
+            unpruned = 0
         del original_a
         devices.torch_gc(force=True)
         original_b = TensorDict.from_dict(read_state_dict(models["model_b"], device))
-        for key in tqdm(original_b.keys(), desc="un-prune model b"):
+        for key in original_b.keys():
             if KEY_POSITION_IDS in key:
                 continue
             if "model" in key and key not in merged.keys():
                 merged.update({key: original_b[key]})
+                unpruned += 1
                 if precision == "fp16":
                     merged.update({key: merged[key].half()})
+        if unpruned != 0:
+            log.info(f"Merge: {unpruned} unmerged keys restored from Secondary Model")
         del original_b
 
     return fix_clip(merged)
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 509e828a6..9141dc9bb 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -162,10 +162,11 @@ def create_ui():
                         with FormRow():
                             precision = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision")
                         with FormRow():
-                            device = gr.Radio(choices=["cpu", "shuffle", "cuda"], value="cpu", label="Merge Device")
+                            device = gr.Radio(choices=["cpu", "shuffle", "gpu"], value="cpu", label="Merge Device")
+                            unload = gr.Checkbox(label="Unload Current Model from VRAM", value=False, visible=False)
                         with FormRow():
                             bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None",
-                                                      interactive=True, label="Bake in VAE")
+                                                      interactive=True, label="Replace 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")
@@ -202,6 +203,7 @@ def create_ui():
                                 re_basin,
                                 re_basin_iterations,
                                 device,
+                                unload,
                                 bake_in_vae):
                     kwargs = {}
                     for x in inspect.getfullargspec(modelmerger)[0]:
@@ -240,6 +242,13 @@ def create_ui():
                     doc = getattr(merge_methods, mode).__doc__.replace("\n", "<br>")
                     return gr.update(value=doc, visible=True)
 
+                def show_unload(device):
+                    if device == "gpu":
+                        return gr.update(visble=True)
+                    else:
+                        return gr.update(visble=False)
+
+
                 def preset_visiblility(x):
                     if len(x) == 2:
                         return gr.Slider.update(value=0.5, visible=True)
@@ -262,7 +271,7 @@ def create_ui():
                         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)]
-
+                device.change(fn=show_unload, inputs=device, outputs=unload)
                 merge_mode.change(fn=show_help, inputs=merge_mode, outputs=merge_mode_docs)
                 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)
@@ -309,6 +318,7 @@ def create_ui():
                         re_basin,
                         re_basin_iterations,
                         device,
+                        unload,
                         bake_in_vae,
                     ],
                     outputs=[