Spaces:
Runtime error
Runtime error
File size: 4,354 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
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
|