Mirrors
/
diffusers
mirror of https://github.com/ShivamShrirao/diffusers.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[Safetensors] Make sure metadata is saved (#2506) 

* [Safetensors] Make sure metadata is saved

* make style
This commit is contained in:
Patrick von Platen 2023-02-27 17:24:40 +02:00 committed by GitHub
parent 7f43f65235
commit 0e975e5ff6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/diffusers/models/modeling_utils.py
View File

@ -291,9 +291,6 @@ class ModelMixin(torch.nn.Module):
logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
return
if save_function is None:
save_function = safetensors.torch.save_file if safe_serialization else torch.save
os.makedirs(save_directory, exist_ok=True)
model_to_save = self
@ -310,7 +307,12 @@ class ModelMixin(torch.nn.Module):
weights_name = _add_variant(weights_name, variant)
# Save the model
save_function(state_dict, os.path.join(save_directory, weights_name))
if safe_serialization:
safetensors.torch.save_file(
state_dict, os.path.join(save_directory, weights_name), metadata={"format": "pt"}
)
else:
torch.save(state_dict, os.path.join(save_directory, weights_name))
logger.info(f"Model weights saved in {os.path.join(save_directory, weights_name)}")