summaryrefslogtreecommitdiffstats
path: root/pipelines/stable_diffusion
diff options
context:
space:
mode:
Diffstat (limited to 'pipelines/stable_diffusion')
-rw-r--r--pipelines/stable_diffusion/vlpn_stable_diffusion.py20
1 files changed, 14 insertions, 6 deletions
diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index 13ea2ac..a0dff54 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -591,15 +591,23 @@ class VlpnStableDiffusion(DiffusionPipeline):
591 if callback is not None and i % callback_steps == 0: 591 if callback is not None and i % callback_steps == 0:
592 callback(i, t, latents) 592 callback(i, t, latents)
593 593
594 # 9. Post-processing
595 image = self.decode_latents(latents)
596
597 # 10. Run safety checker
598 has_nsfw_concept = None 594 has_nsfw_concept = None
599 595
600 # 11. Convert to PIL 596 if output_type == "latent":
601 if output_type == "pil": 597 image = latents
598 elif output_type == "pil":
599 # 9. Post-processing
600 image = self.decode_latents(latents)
601
602 # 10. Convert to PIL
602 image = self.numpy_to_pil(image) 603 image = self.numpy_to_pil(image)
604 else:
605 # 9. Post-processing
606 image = self.decode_latents(latents)
607
608 # Offload last model to CPU
609 if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
610 self.final_offload_hook.offload()
603 611
604 if not return_dict: 612 if not return_dict:
605 return (image, has_nsfw_concept) 613 return (image, has_nsfw_concept)