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

[Save Pretrained] Remove dead code lines that can accidentally remove pytorch files (#2038) 

correct safetensors
This commit is contained in:
Patrick von Platen 2023-01-19 10:11:27 +01:00 committed by GitHub
parent 007c914c70
commit f354dd9e2f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 0 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -272,15 +272,6 @@ class ModelMixin(torch.nn.Module):
weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
# Clean the folder from a previous save
for filename in os.listdir(save_directory):
full_filename = os.path.join(save_directory, filename)
# If we have a shard file that is not going to be replaced, we delete it, but only from the main process
# in distributed settings to avoid race conditions.
weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "")
if filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and is_main_process:
os.remove(full_filename)
# Save the model # Save the model
save_function(state_dict, os.path.join(save_directory, weights_name)) save_function(state_dict, os.path.join(save_directory, weights_name))