import os import re import torch from safetensors import safe_open from safetensors.torch import save_file import hashlib from io import BytesIO import safetensors.torch from typing import Callable, Union, Optional re_digits = re.compile(r"\d+") re_x_proj = re.compile(r"(.*)_([qkv]_proj)$") re_compiled = {} suffix_conversion = { "attentions": {}, "resnets": { "conv1": "in_layers_2", "conv2": "out_layers_3", "time_emb_proj": "emb_layers_1", "conv_shortcut": "skip_connection", } } def convert_diffusers_name_to_compvis(key, is_sd2): def match(match_list, regex_text): regex = re_compiled.get(regex_text) if regex is None: regex = re.compile(regex_text) re_compiled[regex_text] = regex r = re.match(regex, key) if not r: return False match_list.clear() match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()]) return True m = [] if match(m, r"lora_unet_conv_in(.*)"): return f'diffusion_model_input_blocks_0_0{m[0]}' if match(m, r"lora_unet_conv_out(.*)"): return f'diffusion_model_out_2{m[0]}' if match(m, r"lora_unet_time_embedding_linear_(\d+)(.*)"): return f"diffusion_model_time_embed_{m[0] * 2 - 2}{m[1]}" if match(m, r"lora_unet_down_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"): suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3]) return f"diffusion_model_input_blocks_{1 + m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}" if match(m, r"lora_unet_mid_block_(attentions|resnets)_(\d+)_(.+)"): suffix = suffix_conversion.get(m[0], {}).get(m[2], m[2]) return f"diffusion_model_middle_block_{1 if m[0] == 'attentions' else m[1] * 2}_{suffix}" if match(m, r"lora_unet_up_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"): suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3]) return f"diffusion_model_output_blocks_{m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}" if match(m, r"lora_unet_down_blocks_(\d+)_downsamplers_0_conv"): return f"diffusion_model_input_blocks_{3 + m[0] * 3}_0_op" if match(m, r"lora_unet_up_blocks_(\d+)_upsamplers_0_conv"): return f"diffusion_model_output_blocks_{2 + m[0] * 3}_{2 if m[0]>0 else 1}_conv" if match(m, r"lora_te_text_model_encoder_layers_(\d+)_(.+)"): if is_sd2: if 'mlp_fc1' in m[1]: return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}" elif 'mlp_fc2' in m[1]: return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}" else: return f"model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" return f"transformer_text_model_encoder_layers_{m[0]}_{m[1]}" if match(m, r"lora_te2_text_model_encoder_layers_(\d+)_(.+)"): if 'mlp_fc1' in m[1]: return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}" elif 'mlp_fc2' in m[1]: return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}" else: return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" return key def safetensors_hashes(tensors, metadata): """Precalculate the model hashes needed by sd-webui-additional-networks to save time on indexing the model later.""" # Because writing user metadata to the file can change the result of # sd_models.model_hash(), only retain the training metadata for purposes of # calculating the hash, as they are meant to be immutable metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")} bytes = safetensors.torch.save(tensors, metadata) b = BytesIO(bytes) model_hash = addnet_hash_safetensors(b) legacy_hash = addnet_hash_legacy(b) return model_hash, legacy_hash def addnet_hash_legacy(b): """Old model hash used by sd-webui-additional-networks for .safetensors format files""" m = hashlib.sha256() b.seek(0x100000) m.update(b.read(0x10000)) return m.hexdigest()[0:8] def addnet_hash_safetensors(b): """New model hash used by sd-webui-additional-networks for .safetensors format files""" hash_sha256 = hashlib.sha256() blksize = 1024 * 1024 b.seek(0) header = b.read(8) n = int.from_bytes(header, "little") offset = n + 8 b.seek(offset) for chunk in iter(lambda: b.read(blksize), b""): hash_sha256.update(chunk) return hash_sha256.hexdigest() def lbw_lora(input_, output, ratios): print("Apply LBW") assert isinstance(input_, str) assert isinstance(output, str) assert isinstance(ratios, str) assert os.path.exists(input_), f"{input_} is not exists" assert os.path.exists(output) == False, f"{output} aleady exists" LOAD_PATH = input_ SAVE_PATH = output RATIOS = [float(x) for x in ratios.split(",")] LAYERS = len(RATIOS) assert LAYERS in [17, 26] BLOCKID17 = [ "BASE", "IN01", "IN02", "IN04", "IN05", "IN07", "IN08", "M00", "OUT03", "OUT04", "OUT05", "OUT06", "OUT07", "OUT08", "OUT09", "OUT10", "OUT11"] BLOCKID26 = [ "BASE", "IN00", "IN01", "IN02", "IN03", "IN04", "IN05", "IN06", "IN07", "IN08", "IN09", "IN10", "IN11", "M00", "OUT00", "OUT01", "OUT02", "OUT03", "OUT04", "OUT05", "OUT06", "OUT07", "OUT08", "OUT09", "OUT10", "OUT11"] if LAYERS == 17: RATIO_OF_ = dict(zip(BLOCKID17, RATIOS)) if LAYERS == 26: RATIO_OF_ = dict(zip(BLOCKID26, RATIOS)) print(RATIO_OF_) PATTERNS = [ r"^transformer_text_model_(encoder)_layers_(\d+)_.*", r"^diffusion_model_(in)put_blocks_(\d+)_.*", r"^diffusion_model_(middle)_block_(\d+)_.*", r"^diffusion_model_(out)put_blocks_(\d+)_.*"] def replacement(match): g1 = str(match.group(1)) # encoder, in, middle, out g2 = int(match.group(2)) # number assert g1 in ["encoder", "in", "middle", "out"] assert isinstance(g2, int) if g1 == "encoder": return "BASE" if g1 == "middle": return "M00" return f"{str.upper(g1)}{g2:02}" def compvis_name_to_blockid(compvis_name): strings = compvis_name for pattern in PATTERNS: strings = re.sub(pattern, replacement, strings) if strings != compvis_name: break assert strings != compvis_name blockid = strings if LAYERS == 17: assert blockid in BLOCKID26, f"Incorrect layer {blockid}" assert blockid in BLOCKID17, f"{blockid} is not included in 17 layers. May be 26 layers?" if LAYERS == 26: assert blockid in BLOCKID26, f"Incorrect layer {blockid}" return blockid with safe_open(LOAD_PATH, framework="pt", device="cpu") as f: tensors = {} for key in f.keys(): tensors[key] = f.get_tensor(key) # key = diffusers_name compvis_name = convert_diffusers_name_to_compvis(key, is_sd2=False) blockid = compvis_name_to_blockid(compvis_name) if compvis_name.endswith("lora_up.weight"): tensors[key] *= RATIO_OF_[blockid] print(f"({blockid}) {compvis_name} " f"updated with factor {RATIO_OF_[blockid]}") save_file(tensors, SAVE_PATH) print("Done")