Spaces:
Runtime error
Runtime error
| 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 | |