From 11b728dc2f9908d5db023af4ac0b8c6e67df1d98 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 26 Apr 2023 08:09:36 +0000 Subject: [PATCH 1/2] Add all files --- .../pipeline_stable_diffusion_upscale.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index 45b26de284af..88710641fca4 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -697,15 +697,11 @@ def __call__( # 10. Post-processing # make sure the VAE is in float32 mode, as it overflows in float16 self.vae.to(dtype=torch.float32) - image = self.decode_latents(latents.float()) - - # Offload last model to CPU - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.final_offload_hook.offload() # 11. Convert to PIL # has_nsfw_concept = False if output_type == "pil": + image = self.decode_latents(latents.float()) image, has_nsfw_concept, _ = self.run_safety_checker(image, device, prompt_embeds.dtype) image = self.numpy_to_pil(image) @@ -713,9 +709,19 @@ def __call__( # 11. Apply watermark if self.watermarker is not None: image = self.watermarker.apply_watermark(image) + elif output_type == "pt": + latents = 1 / self.vae.config.scaling_factor * latents.float() + image = self.vae.decode(latents).sample + has_nsfw_concept = None else: + image = self.decode_latents(latents.float()) has_nsfw_concept = None + # Offload last model to CPU + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.final_offload_hook.offload() + + if not return_dict: return (image, has_nsfw_concept) From 5ae6f9ebbb3276b62b40ed84b298d2aaf186d472 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 26 Apr 2023 08:10:02 +0000 Subject: [PATCH 2/2] update --- .../stable_diffusion/pipeline_stable_diffusion_upscale.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index 88710641fca4..14e5c4ab7cd1 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -721,7 +721,6 @@ def __call__( if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: self.final_offload_hook.offload() - if not return_dict: return (image, has_nsfw_concept)