reimash/ai-toolkit

1import argparse2import os3# add project root to sys path4import sys5 6from tqdm import tqdm7 8sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))9 10import torch11from diffusers.loaders import LoraLoaderMixin12from safetensors.torch import load_file13from collections import OrderedDict14import json15 16from toolkit.config_modules import ModelConfig17from toolkit.paths import KEYMAPS_ROOT18from toolkit.saving import convert_state_dict_to_ldm_with_mapping, get_ldm_state_dict_from_diffusers19from toolkit.stable_diffusion_model import StableDiffusion20 21# this was just used to match the vae keys to the diffusers keys22# you probably wont need this. Unless they change them.... again... again23# on second thought, you probably will24 25project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))26 27device = torch.device('cpu')28dtype = torch.float3229 30parser = argparse.ArgumentParser()31 32# require at lease one config file33parser.add_argument(34    'file_1',35    nargs='+',36    type=str,37    help='Path an LDM model'38)39 40parser.add_argument(41    '--is_xl',42    action='store_true',43    help='Is the model an XL model'44)45 46parser.add_argument(47    '--is_v2',48    action='store_true',49    help='Is the model a v2 model'50)51 52args = parser.parse_args()53 54find_matches = False55 56print("Loading model")57state_dict_file_1 = load_file(args.file_1[0])58state_dict_1_keys = list(state_dict_file_1.keys())59 60print("Loading model into diffusers format")61model_config = ModelConfig(62    name_or_path=args.file_1[0],63    is_xl=args.is_xl64)65sd = StableDiffusion(66    model_config=model_config,67    device=device,68)69sd.load_model()70 71# load our base72base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl_ldm_base.safetensors')73mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl.json')74 75print("Converting model back to LDM")76version_string = '1'77if args.is_v2:78    version_string = '2'79if args.is_xl:80    version_string = 'sdxl'81# convert the state dict82state_dict_file_2 = get_ldm_state_dict_from_diffusers(83    sd.state_dict(),84    version_string,85    device='cpu',86    dtype=dtype87)88 89# state_dict_file_2 = load_file(args.file_2[0])90 91state_dict_2_keys = list(state_dict_file_2.keys())92keys_in_both = []93 94keys_not_in_state_dict_2 = []95for key in state_dict_1_keys:96    if key not in state_dict_2_keys:97        keys_not_in_state_dict_2.append(key)98 99keys_not_in_state_dict_1 = []100for key in state_dict_2_keys:101    if key not in state_dict_1_keys:102        keys_not_in_state_dict_1.append(key)103 104keys_in_both = []105for key in state_dict_1_keys:106    if key in state_dict_2_keys:107        keys_in_both.append(key)108 109# sort them110keys_not_in_state_dict_2.sort()111keys_not_in_state_dict_1.sort()112keys_in_both.sort()113 114if len(keys_not_in_state_dict_2) == 0 and len(keys_not_in_state_dict_1) == 0:115    print("All keys match!")116    print("Checking values...")117    mismatch_keys = []118    loss = torch.nn.MSELoss()119    tolerance = 1e-6120    for key in tqdm(keys_in_both):121        if loss(state_dict_file_1[key], state_dict_file_2[key]) > tolerance:122            print(f"Values for key {key} don't match!")123            print(f"Loss: {loss(state_dict_file_1[key], state_dict_file_2[key])}")124            mismatch_keys.append(key)125 126    if len(mismatch_keys) == 0:127        print("All values match!")128    else:129        print("Some valued font match!")130        print(mismatch_keys)131        mismatched_path = os.path.join(project_root, 'config', 'mismatch.json')132        with open(mismatched_path, 'w') as f:133            f.write(json.dumps(mismatch_keys, indent=4))134        exit(0)135 136else:137    print("Keys don't match!, generating info...")138 139json_data = {140    "both": keys_in_both,141    "not_in_state_dict_2": keys_not_in_state_dict_2,142    "not_in_state_dict_1": keys_not_in_state_dict_1143}144json_data = json.dumps(json_data, indent=4)145 146remaining_diffusers_values = OrderedDict()147for key in keys_not_in_state_dict_1:148    remaining_diffusers_values[key] = state_dict_file_2[key]149 150# print(remaining_diffusers_values.keys())151 152remaining_ldm_values = OrderedDict()153for key in keys_not_in_state_dict_2:154    remaining_ldm_values[key] = state_dict_file_1[key]155 156# print(json_data)157 158 159json_save_path = os.path.join(project_root, 'config', 'keys.json')160json_matched_save_path = os.path.join(project_root, 'config', 'matched.json')161json_duped_save_path = os.path.join(project_root, 'config', 'duped.json')162state_dict_1_filename = os.path.basename(args.file_1[0])163# state_dict_2_filename = os.path.basename(args.file_2[0])164# save key names for each in own file165with open(os.path.join(project_root, 'config', f'{state_dict_1_filename}.json'), 'w') as f:166    f.write(json.dumps(state_dict_1_keys, indent=4))167 168with open(os.path.join(project_root, 'config', f'{state_dict_1_filename}_loop.json'), 'w') as f:169    f.write(json.dumps(state_dict_2_keys, indent=4))170 171with open(json_save_path, 'w') as f:172    f.write(json_data)173