From 07d0f28a3bda7e99f11539505c8beb172b5b62f2 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/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 2a33074a2473..098e8e6a85c9 100644 --- a/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -133,7 +133,8 @@ def enable_sequential_cpu_offload(self): 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__(