VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix an issue with saving as diffusers sd1/2 model close #1033

Browse Source
This commit is contained in:
Kohya S
2024-01-04 21:43:50 +09:00
parent 716bad188b
commit aab265e431
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
library/model_util.py
View File

@@ -1245,8 +1245,13 @@ def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_mod
if vae is None: if vae is None:
vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae") vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
# original U-Net cannot be saved, so we need to convert it to the Diffusers version
# TODO this consumes a lot of memory
diffusers_unet = diffusers.UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet")
diffusers_unet.load_state_dict(unet.state_dict())
pipeline = StableDiffusionPipeline( pipeline = StableDiffusionPipeline(
unet=unet, unet=diffusers_unet,
text_encoder=text_encoder, text_encoder=text_encoder,
vae=vae, vae=vae,
scheduler=scheduler, scheduler=scheduler,