diff --git a/html/locale_en.json b/html/locale_en.json
index 65f239e10..45074cb63 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,18 @@
{"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":"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 a84b9b28e..a0c048b4e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,5 +1,4 @@
import os
-import re
import html
import json
import shutil
@@ -8,11 +7,10 @@ import torch
import tqdm
import gradio as gr
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
-
-
-checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
+from modules import shared, images, sd_models, sd_vae, sd_models_config, devices
def run_pnginfo(image):
@@ -31,6 +29,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:
@@ -53,158 +52,116 @@ 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, **kwargs): # 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),
+ kwargs["models"] = {
+ "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,
}
- 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':
+
+ 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[primary_model_name]
- if theta_func2 and (not secondary_model_name or secondary_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[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)
+ 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.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.get_closet_checkpoint_match(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:
+ shared.log.warn(f"Merge: {e}")
+ kwargs["alpha"] = kwargs.get("alpha_preset", kwargs["alpha"])
+ 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:
+ shared.log.warn(f"Merge: {e}")
+ kwargs["beta"] = kwargs.get("beta_preset", kwargs["beta"])
+ 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)
+
+ if kwargs["device"] == "gpu":
+ kwargs["device"] = devices.device
+ elif kwargs["device"] == "shuffle":
+ kwargs["device"] = torch.device("cpu")
+ kwargs["work_device"] = devices.device
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)
+ 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"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():
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)
+ theta_0[theta_0_key] = to_half(vae_dict[key], kwargs.get("precision", "fp16") == "fp16")
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
+ filename = kwargs.get("custom_name", "Unnamed_Merge")
+ filename += "." + kwargs.get("checkpoint_format", None)
output_modelname = os.path.join(ckpt_dir, filename)
- shared.state.nextjob()
shared.state.textinfo = "Saving"
metadata = None
- if save_metadata:
+ if kwargs.get("save_metadata", False):
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,
- "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
+ "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)
@@ -225,22 +182,28 @@ 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)
+
+
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}.")
+ if kwargs["device"].type != "cpu":
+ devices.torch_gc(force=True)
+ 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)], "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_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}
@@ -278,7 +241,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 = {
@@ -340,6 +302,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/merging/merge.py b/modules/merging/merge.py
new file mode 100644
index 000000000..09d03685b
--- /dev/null
+++ b/modules/merging/merge.py
@@ -0,0 +1,418 @@
+import os
+from concurrent.futures import ThreadPoolExecutor
+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
+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_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
+
+
+KEY_POSITION_IDS = ".".join(
+ [
+ "cond_stage_model",
+ "transformer",
+ "text_model",
+ "embeddings",
+ "position_ids",
+ ]
+)
+
+
+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 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=""):
+ log.debug(f"{txt} VRAM: {modules.memstats.memory_stats()}")
+
+
+def load_thetas(
+ models: Dict[str, os.PathLike | str],
+ prune: bool,
+ device: torch.device,
+ precision: str,
+) -> Dict:
+ log_vram("before loading models")
+ if prune:
+ thetas = {k: prune_sd_model(TensorDict.from_dict(read_state_dict(m, "cpu"))) for k, m in models.items()}
+ else:
+ 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():
+ 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 = "fp16",
+ weights_clip: bool = False,
+ re_basin: bool = False,
+ device: torch.device = None,
+ work_device: torch.device = None,
+ prune: bool = False,
+ threads: int = 1,
+ **kwargs,
+) -> Dict:
+ iterations = kwargs.get("re_basin_iterations", 1)
+ thetas = load_thetas(models, prune, device, precision)
+
+ log.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: torch.device,
+ prune: bool,
+ precision: str,
+) -> Dict:
+ if prune:
+ log.info("Un-pruning merged model")
+ del thetas
+ devices.torch_gc(force=True)
+ log_vram("remove thetas")
+ original_a = TensorDict.from_dict(read_state_dict(models["model_a"], device))
+ 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 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)
+
+
+def simple_merge(
+ thetas: Dict[str, Dict],
+ weight_matcher: WeightClass,
+ merge_mode: str,
+ precision: str = "fp16",
+ weights_clip: bool = False,
+ device: torch.device = None,
+ work_device: torch.device = 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 == "fp16":
+ thetas["model_a"].update({key: thetas["model_a"][key].half()})
+
+ log_vram("after stage 2")
+
+ return fix_clip(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: torch.device = None,
+ work_device: torch.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()
+
+ for it in range(iterations):
+ log_vram(f"Rebasin iteration {it}")
+ weight_matcher.set_it(it)
+
+ # 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 == "fp16",
+ 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 == "fp16",
+ 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( # pylint: disable=inconsistent-return-statements
+ key: str,
+ thetas: Dict,
+ weight_matcher: WeightClass,
+ merge_mode: str,
+ precision: str = "fp16",
+ weights_clip: bool = False,
+ device: torch.device = None,
+ work_device: torch.device = 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 == "fp16":
+ 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: torch.device,
+) -> 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:
+ log.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..3f704c20f
--- /dev/null
+++ b/modules/merging/merge_methods.py
@@ -0,0 +1,247 @@
+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: # pylint: disable=unused-argument
+ """
+ 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: # pylint: disable=unused-argument
+ """
+ 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
+
+ return (a - alpha * beta * b) / (1 - alpha * beta)
+
+
+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
+ 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)
+ 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: # 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: # pylint: disable=unused-argument
+ """
+ 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: # pylint: disable=unused-argument
+ """
+ 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: # pylint: disable=unused-argument
+ """
+ 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))
+ 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: # pylint: disable=unused-argument
+ """
+ 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: # pylint: disable=unused-argument
+ """
+ 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)
+ 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: # pylint: disable=unused-argument
+ """
+ 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)
+
+ 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): # pylint: disable=unused-argument
+ """
+ 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
+
+ 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: # pylint: disable=unused-argument
+ """
+ An implementation of arXiv:2306.01708
+ """
+ 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/merge_presets.py b/modules/merging/merge_presets.py
new file mode 100644
index 000000000..40ec85010
--- /dev/null
+++ b/modules/merging/merge_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/merge_rebasin.py b/modules/merging/merge_rebasin.py
new file mode 100644
index 000000000..2606a52f9
--- /dev/null
+++ b/modules/merging/merge_rebasin.py
@@ -0,0 +1,2288 @@
+# 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
+from modules.shared import log
+
+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: { # 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,)} # pylint: disable=unnecessary-lambda-assignment
+ dense = (
+ 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: { # pylint: disable=unnecessary-lambda-assignment
+ f"{name}": (
+ p_out,
+ p_in,
+ None,
+ None,
+ )
+ }
+
+ # Unet Res blocks
+ 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(
+ 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: { # 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"),
+ **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: { # 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"),
+ **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
+ log.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/merge_utils.py b/modules/merging/merge_utils.py
new file mode 100644
index 000000000..5bdfee246
--- /dev/null
+++ b/modules/merging/merge_utils.py
@@ -0,0 +1,111 @@
+import inspect
+import re
+from modules.merging import merge_methods
+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)
+
+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 = "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
+ 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"])
+ if self.ratioDict.get("beta", None):
+ current_bases["beta"] = self.step_weights_and_bases(self.ratioDict["beta"])
+
+ 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()}
+ return current_bases
+
+ 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):
+ self.it = it
diff --git a/modules/ui_common.py b/modules/ui_common.py
index a2b8aa1ee..b1a5bf364 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -239,7 +239,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()
@@ -249,7 +249,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..9141dc9bb 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
@@ -9,7 +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 modules.merging import merge_methods
+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
@@ -32,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)
@@ -69,47 +75,217 @@ def create_ui():
)
with gr.Tab(label="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():
- def sd_model_choices():
- return ['None'] + sd_models.checkpoint_tiles()
+ 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", "
"))
+ 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")
- create_refresh_button(tertiary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "refresh_checkpoint_C")
+ 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():
- 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 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():
+ sdxl = gr.Checkbox(label="SDXL")
+ 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():
- 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")
+ weights_clip = gr.Checkbox(label="Weights Clip")
+ prune = gr.Checkbox(label="Prune", value=True, visible=False)
+ re_basin = gr.Checkbox(label="ReBasin")
with FormRow():
- config_source = gr.Radio(choices=["Primary", "Secondary", "Tertiary", "None"], value="Primary", label="Model configuration", type="index")
+ re_basin_iterations = gr.Slider(minimum=0, maximum=25, step=1,
+ label='Number of ReBasin Iterations', value=None,
+ visible=False)
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")
+ checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors",
+ visible=False, label="Model format")
with FormRow():
- discard_weights = gr.Textbox(value="", label="Discard weights with matching name")
+ precision = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision")
+ with FormRow():
+ 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="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")
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):
+ 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,
+ unload,
+ bake_in_vae):
+ kwargs = {}
+ 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_modelmerger(*args)
+ 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
- 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):
+ 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 show_help(mode):
+ doc = getattr(merge_methods, mode).__doc__.replace("\n", "
")
+ 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)
+ 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 = ['%.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)]
+
+ 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)]
+ 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)
+ 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])
+
modelmerger_merge.click(
fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
_js='modelmerger',
@@ -118,15 +294,32 @@ def create_ui():
primary_model_name,
secondary_model_name,
tertiary_model_name,
- interp_method,
- interp_amount,
- save_as_half,
+ 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,
- config_source,
- bake_in_vae,
- discard_weights,
save_metadata,
+ weights_clip,
+ prune,
+ re_basin,
+ re_basin_iterations,
+ device,
+ unload,
+ bake_in_vae,
],
outputs=[
primary_model_name,
@@ -147,14 +340,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():
@@ -193,11 +386,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):
@@ -205,8 +402,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}"'
@@ -214,20 +412,25 @@ def create_ui():
with gr.Column(scale=6):
gr.HTML('Search for models
Select a model from the search results to download
')
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('
')
hf_download_model_btn = gr.Button(value="Download model", variant='primary')
@@ -235,12 +438,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 = []
@@ -283,7 +490,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]
@@ -292,7 +500,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'],
@@ -314,7 +523,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']),
@@ -330,7 +540,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:
@@ -340,7 +551,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
@@ -357,12 +568,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']))
@@ -374,10 +587,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']))
@@ -392,11 +607,12 @@ def create_ui():
txt = '
'.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('Fetch information
Fetches preview and metadata information for all models with missing information
Models with existing previews and information are not updated
')
+ gr.HTML(
+ 'Fetch information
Fetches preview and metadata information for all models with missing information
Models with existing previews and information are not updated
')
with gr.Row():
civit_previews_btn = gr.Button(value="Start", variant='primary')
with gr.Row():
@@ -406,11 +622,12 @@ def create_ui():
gr.HTML('Search for models
')
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('')
@@ -418,39 +635,64 @@ def create_ui():
gr.HTML('Download model
')
with gr.Row():
civit_download_model_btn = gr.Button(value="Download", variant='primary')
- gr.HTML('Select a model, model version and and model variant from the search results to download or enter model URL manually
')
+ gr.HTML(
+ 'Select a model, model version and and model variant from the search results to download or enter model URL manually
')
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('Search results
')
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 04dd5360f..7e5eaefc0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -33,6 +33,7 @@ rich
safetensors
scipy
tb_nightly
+tensordict
toml
torchdiffeq
voluptuous