AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-29 07:22:12 +03:00
Code Activity

Allow return pt x4 (#3236)

Browse Source 
* Add all files

* update
This commit is contained in:
Patrick von Platen
2023-04-26 10:34:34 +02:00
committed by GitHub
parent e51f19aee8
commit da2ce1a6b9
1 changed files with 11 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
View File

@@ -697,15 +697,11 @@ class StableDiffusionUpscalePipeline(DiffusionPipeline, TextualInversionLoaderMi
# 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,8 +709,17 @@ class StableDiffusionUpscalePipeline(DiffusionPipeline, TextualInversionLoaderMi
# 11. Apply watermark
if self.watermarker is not None:
image = self.watermarker.apply_watermark(image)
else:
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)