ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-13 22:49:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

Save entire pixart model again

Browse Source
This commit is contained in:
Jaret Burkett
2024-07-07 07:56:48 -06:00
parent cab8a1c7b8
commit 045e4a6e15
5 changed files with 177 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
toolkit/stable_diffusion_model.py
View File

@@ -1775,16 +1775,16 @@ class StableDiffusion:
# saving in diffusers format
if not output_file.endswith('.safetensors'):
# diffusers
if self.is_pixart:
self.unet.save_pretrained(
save_directory=output_file,
safe_serialization=True,
)
else:
self.pipeline.save_pretrained(
save_directory=output_file,
safe_serialization=True,
)
# if self.is_pixart:
# self.unet.save_pretrained(
# save_directory=output_file,
# safe_serialization=True,
# )
# else:
self.pipeline.save_pretrained(
save_directory=output_file,
safe_serialization=True,
)
# save out meta config
meta_path = os.path.join(output_file, 'aitk_meta.yaml')
with open(meta_path, 'w') as f: