From bf7b0bc25b4a10da78ae9dd19ae00ba840585d40 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Mon, 31 Oct 2022 15:03:33 +0100 Subject: [PATCH] Allow `safety_checker` to be `None` when using CPU offload (#1078) Allow None safety_checker when using CPU offload. --- .../pipelines/stable_diffusion/pipeline_stable_diffusion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 2a33074a24..098e8e6a85 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -133,7 +133,8 @@ class StableDiffusionPipeline(DiffusionPipeline): device = torch.device("cuda") for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]: - cpu_offload(cpu_offloaded_model, device) + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) @torch.no_grad() def __call__(