From aab265e43143537356dc9146b41ff6da96fb96b6 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Thu, 4 Jan 2024 21:43:50 +0900 Subject: [PATCH] Fix an issue with saving as diffusers sd1/2 model close #1033 --- library/model_util.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/library/model_util.py b/library/model_util.py index 6102d0a1..1f40ce32 100644 --- a/library/model_util.py +++ b/library/model_util.py @@ -1245,8 +1245,13 @@ def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_mod if vae is None: 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( - unet=unet, + unet=diffusers_unet, text_encoder=text_encoder, vae=vae, scheduler=scheduler,