From 1a6def3ddbbf560b815c03faa4a7193a0d030591 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Tue, 4 Apr 2023 08:52:55 -1000 Subject: [PATCH] fix post-processing (#2968) Co-authored-by: yiyixuxu --- .../alt_diffusion/pipeline_alt_diffusion_img2img.py | 13 +++++++------ .../pipeline_stable_diffusion_img2img.py | 13 +++++++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py index f9dfe3f38f..bb8116f2f5 100644 --- a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py +++ b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py @@ -734,14 +734,15 @@ class AltDiffusionImg2ImgPipeline(DiffusionPipeline, TextualInversionLoaderMixin image = latents has_nsfw_concept = None - image = self.decode_latents(latents) - - if self.safety_checker is not None: - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) else: - has_nsfw_concept = False + image = self.decode_latents(latents) - image = self.image_processor.postprocess(image, output_type=output_type) + if self.safety_checker is not None: + image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + else: + has_nsfw_concept = False + + image = self.image_processor.postprocess(image, output_type=output_type) # Offload last model to CPU if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index a91431f719..a0befdae73 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -742,14 +742,15 @@ class StableDiffusionImg2ImgPipeline(DiffusionPipeline, TextualInversionLoaderMi image = latents has_nsfw_concept = None - image = self.decode_latents(latents) - - if self.safety_checker is not None: - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) else: - has_nsfw_concept = False + image = self.decode_latents(latents) - image = self.image_processor.postprocess(image, output_type=output_type) + if self.safety_checker is not None: + image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + else: + has_nsfw_concept = False + + image = self.image_processor.postprocess(image, output_type=output_type) # Offload last model to CPU if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: