mirror of
https://github.com/vladmandic/automatic
synced 2026-08-29 08:31:00 +02:00
Merge pull request #2443 from AI-Casanova/Extended-Merging
Advanced Merging
This commit is contained in:
Vladimir Mandic
GitHub
+13
-1
@@ -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":""},
|
||||
|
||||
+98
-135
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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],
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
+321
-79
@@ -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", "<br>"))
|
||||
with FormRow():
|
||||
primary_model_name = gr.Dropdown(sd_model_choices(), label="Primary model", value="None")
|
||||
create_refresh_button(primary_model_name, sd_models.list_models, 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", "<br>")
|
||||
return gr.update(value=doc, visible=True)
|
||||
|
||||
def show_unload(device):
|
||||
if device == "gpu":
|
||||
return gr.update(visble=True)
|
||||
else:
|
||||
return gr.update(visble=False)
|
||||
|
||||
|
||||
def preset_visiblility(x):
|
||||
if len(x) == 2:
|
||||
return gr.Slider.update(value=0.5, visible=True)
|
||||
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('<h2>Search for models</h2>Select a model from the search results to download<br><br>')
|
||||
with gr.Row():
|
||||
hf_search_text = gr.Textbox('', label = 'Search models', placeholder='search huggingface models')
|
||||
hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models')
|
||||
hf_search_btn = ToolButton(value="🔍", label="Search")
|
||||
with gr.Row():
|
||||
with gr.Column(scale=2):
|
||||
with gr.Row():
|
||||
hf_selected = gr.Textbox('', label = 'Select model', placeholder='select model from search results or enter model name manually')
|
||||
hf_selected = gr.Textbox('', label='Select model',
|
||||
placeholder='select model from search results or enter model name manually')
|
||||
with gr.Column(scale=1):
|
||||
with gr.Row():
|
||||
hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label = 'Specify model variant', placeholder='')
|
||||
hf_revision = gr.Textbox('', label = 'Specify model revision', placeholder='')
|
||||
hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label='Specify model variant',
|
||||
placeholder='')
|
||||
hf_revision = gr.Textbox('', label='Specify model revision', placeholder='')
|
||||
with gr.Row():
|
||||
hf_token = gr.Textbox('', label = 'Huggingface token', placeholder='optional access token for private or gated models')
|
||||
hf_mirror = gr.Textbox('', label = 'Huggingface mirror', placeholder='optional mirror site for downloads')
|
||||
hf_custom_pipeline = gr.Textbox('', label = 'Custom pipeline', placeholder='optional pipeline for downloads')
|
||||
hf_token = gr.Textbox('', label='Huggingface token',
|
||||
placeholder='optional access token for private or gated models')
|
||||
hf_mirror = gr.Textbox('', label='Huggingface mirror',
|
||||
placeholder='optional mirror site for downloads')
|
||||
hf_custom_pipeline = gr.Textbox('', label='Custom pipeline',
|
||||
placeholder='optional pipeline for downloads')
|
||||
with gr.Column(scale=1):
|
||||
gr.HTML('<br>')
|
||||
hf_download_model_btn = gr.Button(value="Download model", variant='primary')
|
||||
@@ -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 = '<br>'.join([r for r in res if len(r) > 0])
|
||||
return txt
|
||||
|
||||
global search_metadata_civit # pylint: disable=global-statement
|
||||
global search_metadata_civit # pylint: disable=global-statement
|
||||
search_metadata_civit = civit_search_metadata
|
||||
|
||||
with gr.Row():
|
||||
gr.HTML('<h2>Fetch information</h2>Fetches preview and metadata information for all models with missing information<br>Models with existing previews and information are not updated<br>')
|
||||
gr.HTML(
|
||||
'<h2>Fetch information</h2>Fetches preview and metadata information for all models with missing information<br>Models with existing previews and information are not updated<br>')
|
||||
with gr.Row():
|
||||
civit_previews_btn = gr.Button(value="Start", variant='primary')
|
||||
with gr.Row():
|
||||
@@ -406,11 +622,12 @@ def create_ui():
|
||||
gr.HTML('<h2>Search for models</h2>')
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
civit_model_type = gr.Dropdown(label='Model type', choices=['SD 1.5', 'SD XL', 'LoRA', 'Other'], value='LoRA')
|
||||
civit_model_type = gr.Dropdown(label='Model type', choices=['SD 1.5', 'SD XL', 'LoRA', 'Other'],
|
||||
value='LoRA')
|
||||
with gr.Column(scale=15):
|
||||
with gr.Row():
|
||||
civit_search_text = gr.Textbox('', label = 'Search models', placeholder='keyword')
|
||||
civit_search_tag = gr.Textbox('', label = '', placeholder='tags')
|
||||
civit_search_text = gr.Textbox('', label='Search models', placeholder='keyword')
|
||||
civit_search_tag = gr.Textbox('', label='', placeholder='tags')
|
||||
civit_search_btn = ToolButton(value="🔍", label="Search", interactive=False)
|
||||
with gr.Row():
|
||||
civit_search_res = gr.HTML('')
|
||||
@@ -418,39 +635,64 @@ def create_ui():
|
||||
gr.HTML('<h2>Download model</h2>')
|
||||
with gr.Row():
|
||||
civit_download_model_btn = gr.Button(value="Download", variant='primary')
|
||||
gr.HTML('<span style="line-height: 2em">Select a model, model version and and model variant from the search results to download or enter model URL manually</span><br>')
|
||||
gr.HTML(
|
||||
'<span style="line-height: 2em">Select a model, model version and and model variant from the search results to download or enter model URL manually</span><br>')
|
||||
with gr.Row():
|
||||
civit_name = gr.Textbox('', label = 'Model name', placeholder='select model from search results', visible=True)
|
||||
civit_selected = gr.Textbox('', label = 'Model URL', placeholder='select model from search results', visible=True)
|
||||
civit_path = gr.Textbox('', label = 'Download path', placeholder='optional subfolder path where to save model', visible=True)
|
||||
civit_name = gr.Textbox('', label='Model name', placeholder='select model from search results',
|
||||
visible=True)
|
||||
civit_selected = gr.Textbox('', label='Model URL', placeholder='select model from search results',
|
||||
visible=True)
|
||||
civit_path = gr.Textbox('', label='Download path',
|
||||
placeholder='optional subfolder path where to save model', visible=True)
|
||||
with gr.Row():
|
||||
gr.HTML('<h2>Search results</h2>')
|
||||
with gr.Row():
|
||||
civit_headers1 = ['ID', 'Name', 'Tags', 'Downloads', 'Rating']
|
||||
civit_types1 = ['number', 'str', 'str', 'number', 'number']
|
||||
civit_results1 = gr.DataFrame(value = None, label = None, show_label = False, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers1, datatype = civit_types1, type='array', visible=False)
|
||||
civit_results1 = gr.DataFrame(value=None, label=None, show_label=False, interactive=False,
|
||||
wrap=True, overflow_row_behaviour='paginate', max_rows=10,
|
||||
headers=civit_headers1, datatype=civit_types1, type='array',
|
||||
visible=False)
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
civit_headers2 = ['ID', 'ModelID', 'Name', 'Base', 'Created', 'Preview']
|
||||
civit_types2 = ['number', 'number', 'str', 'str', 'date', 'str']
|
||||
civit_results2 = gr.DataFrame(value = None, label = 'Model versions', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers2, datatype = civit_types2, type='array', visible=False)
|
||||
civit_results2 = gr.DataFrame(value=None, label='Model versions', show_label=True,
|
||||
interactive=False, wrap=True, overflow_row_behaviour='paginate',
|
||||
max_rows=10, headers=civit_headers2, datatype=civit_types2,
|
||||
type='array', visible=False)
|
||||
with gr.Column():
|
||||
civit_headers3 = ['Name', 'Size', 'Metadata', 'URL']
|
||||
civit_types3 = ['str', 'number', 'str', 'str']
|
||||
civit_results3 = gr.DataFrame(value = None, label = 'Model variants', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers3, datatype = civit_types3, type='array', visible=False)
|
||||
civit_results3 = gr.DataFrame(value=None, label='Model variants', show_label=True,
|
||||
interactive=False, wrap=True, overflow_row_behaviour='paginate',
|
||||
max_rows=10, headers=civit_headers3, datatype=civit_types3,
|
||||
type='array', visible=False)
|
||||
|
||||
def is_visible(component):
|
||||
visible = len(component) > 0 if component is not None else False
|
||||
return gr.update(visible=visible)
|
||||
|
||||
civit_search_text.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
|
||||
civit_search_tag.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
|
||||
civit_search_btn.click(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
|
||||
civit_results1.select(fn=civit_select1, inputs=[civit_results1], outputs=[civit_results2, civit_results3, models_image])
|
||||
civit_search_text.submit(fn=civit_search_model,
|
||||
inputs=[civit_search_text, civit_search_tag, civit_model_type],
|
||||
outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
|
||||
civit_search_tag.submit(fn=civit_search_model,
|
||||
inputs=[civit_search_text, civit_search_tag, civit_model_type],
|
||||
outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
|
||||
civit_search_btn.click(fn=civit_search_model,
|
||||
inputs=[civit_search_text, civit_search_tag, civit_model_type],
|
||||
outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
|
||||
civit_results1.select(fn=civit_select1, inputs=[civit_results1],
|
||||
outputs=[civit_results2, civit_results3, models_image])
|
||||
civit_results2.select(fn=civit_select2, inputs=[civit_results2], outputs=[civit_results3])
|
||||
civit_results3.select(fn=civit_select3, inputs=[civit_results3], outputs=[civit_selected, civit_name, civit_search_btn])
|
||||
civit_results3.select(fn=civit_select3, inputs=[civit_results3],
|
||||
outputs=[civit_selected, civit_name, civit_search_btn])
|
||||
civit_results1.change(fn=is_visible, inputs=[civit_results1], outputs=[civit_results1])
|
||||
civit_results2.change(fn=is_visible, inputs=[civit_results2], outputs=[civit_results2])
|
||||
civit_results3.change(fn=is_visible, inputs=[civit_results3], outputs=[civit_results3])
|
||||
civit_download_model_btn.click(fn=civit_download_model, inputs=[civit_selected, civit_name, civit_path, civit_model_type, models_image], outputs=[models_outcome])
|
||||
civit_previews_btn.click(fn=civit_search_metadata, inputs=[civit_previews_rehash, civit_previews_rehash], outputs=[models_outcome])
|
||||
civit_download_model_btn.click(fn=civit_download_model,
|
||||
inputs=[civit_selected, civit_name, civit_path, civit_model_type,
|
||||
models_image], outputs=[models_outcome])
|
||||
civit_previews_btn.click(fn=civit_search_metadata,
|
||||
inputs=[civit_previews_rehash, civit_previews_rehash],
|
||||
outputs=[models_outcome])
|
||||
|
||||
@@ -33,6 +33,7 @@ rich
|
||||
safetensors
|
||||
scipy
|
||||
tb_nightly
|
||||
tensordict
|
||||
toml
|
||||
torchdiffeq
|
||||
voluptuous
|
||||
|
||||
Reference in New Issue
Block a user