ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-28 10:11:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added method to get specific keys from model

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-15 14:51:04 -06:00
parent e96874241d
commit 55a5fcc7d9
2 changed files with 30 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
testing/compare_keys.py
View File

@@ -2,6 +2,7 @@ import argparse
import os
import torch
from diffusers.loaders import LoraLoaderMixin
from safetensors.torch import load_file
from collections import OrderedDict
import json
@@ -63,8 +64,8 @@ keys_in_both.sort()
json_data = {
"both": keys_in_both,
"state_dict_2": keys_not_in_state_dict_2,
"state_dict_1": keys_not_in_state_dict_1
"not_in_state_dict_2": keys_not_in_state_dict_2,
"not_in_state_dict_1": keys_not_in_state_dict_1
}
json_data = json.dumps(json_data, indent=4)
@@ -84,6 +85,15 @@ project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
json_save_path = os.path.join(project_root, 'config', 'keys.json')
json_matched_save_path = os.path.join(project_root, 'config', 'matched.json')
json_duped_save_path = os.path.join(project_root, 'config', 'duped.json')
state_dict_1_filename = os.path.basename(args.file_1[0])
state_dict_2_filename = os.path.basename(args.file_2[0])
# save key names for each in own file
with open(os.path.join(project_root, 'config', f'{state_dict_1_filename}.json'), 'w') as f:
f.write(json.dumps(state_dict_1_keys, indent=4))
with open(os.path.join(project_root, 'config', f'{state_dict_2_filename}.json'), 'w') as f:
f.write(json.dumps(state_dict_2_keys, indent=4))
with open(json_save_path, 'w') as f:
f.write(json_data)