hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-04 11:12:35 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Embed model merge metadata in .safetensors file

Browse Source
This commit is contained in:
space-nuko
2023-04-02 17:41:55 -05:00
parent 22bcc7be42
commit d132481058
3 changed files with 55 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
modules/extras.py
View File

@@ -1,6 +1,7 @@
import os
import re
import shutil
import json
import torch
@@ -71,7 +72,7 @@ def to_half(tensor, enable):
return tensor
def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights):
def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata):
shared.state.begin()
shared.state.job = 'model-merge'
@@ -241,13 +242,52 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
shared.state.textinfo = "Saving"
print(f"Saving to {output_modelname}...")
metadata = {"format": "pt", "models": {}, "merge_recipe": None}
if save_metadata:
merge_recipe = {
"primary_model_hash": primary_model_info.sha256,
"secondary_model_hash": secondary_model_info.sha256 if secondary_model_info else None,
"tertiary_model_hash": tertiary_model_info.sha256 if tertiary_model_info else None,
"interp_method": interp_method,
"multiplier": multiplier,
"save_as_half": save_as_half,
"custom_name": custom_name,
"config_source": config_source,
"bake_in_vae": bake_in_vae,
"discard_weights": discard_weights,
"is_inpainting": result_is_inpainting_model,
"is_instruct_pix2pix": result_is_instruct_pix2pix_model
}
metadata["merge_recipe"] = json.dumps(merge_recipe)
def add_model_metadata(checkpoint_info):
metadata["models"][checkpoint_info.sha256] = {
"name": checkpoint_info.name,
"legacy_hash": checkpoint_info.hash,
"merge_recipe": checkpoint_info.metadata.get("merge_recipe", None)
}
metadata["models"].update(checkpoint_info.metadata.get("models", {}))
add_model_metadata(primary_model_info)
if secondary_model_info:
add_model_metadata(secondary_model_info)
if tertiary_model_info:
add_model_metadata(tertiary_model_info)
metadata["models"] = json.dumps(metadata["models"])
_, extension = os.path.splitext(output_modelname)
if extension.lower() == ".safetensors":
safetensors.torch.save_file(theta_0, output_modelname, metadata={"format": "pt"})
safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
else:
torch.save(theta_0, output_modelname)
sd_models.list_models()
created_model = next((ckpt for ckpt in sd_models.checkpoints_list.values() if ckpt.name == filename), None)
if created_model:
created_model.calculate_shorthash()
create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)