mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
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("re_basin_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 "alpha" in self.ratioDict:
|
|
current_bases["alpha"] = self.step_weights_and_bases(self.ratioDict["alpha"])
|
|
if "beta" in self.ratioDict:
|
|
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
|