diff --git a/scripts/extract_locon.py b/scripts/extract_locon.py index 19b4bcc1..9a80d06f 100644 --- a/scripts/extract_locon.py +++ b/scripts/extract_locon.py @@ -13,7 +13,7 @@ import argparse from toolkit.lycoris_utils import extract_diff from toolkit.config import get_config -from toolkit.metadata import create_meta +from toolkit.metadata import create_meta, prep_meta_for_safetensors from toolkit.kohya_model_util import load_models_from_stable_diffusion_checkpoint import torch @@ -141,8 +141,7 @@ def main(): os.makedirs(output_folder, exist_ok=True) # having issues with meta - save_file(state_dict, output_path) - # save_file(state_dict, output_path, {'meta': json.dumps(save_meta, indent=4)}) + save_file(state_dict, output_path, prep_meta_for_safetensors(save_meta)) print(f"Saved to {output_path}") diff --git a/toolkit/metadata.py b/toolkit/metadata.py index 44d6a2f0..32ba7d23 100644 --- a/toolkit/metadata.py +++ b/toolkit/metadata.py @@ -22,3 +22,11 @@ def create_meta(dict_list, name=None): if name is not None: meta_string = meta_string.replace("[name]", name) return json.loads(meta_string) + + +def prep_meta_for_safetensors(meta): + # safetensors can only be one level deep + for key, value in meta.items(): + if isinstance(value, dict): + meta[key] = json.dumps(value) + return meta