Allow safety_checker to be None when using CPU offload (huggingfa…

…ce#1078) Allow None safety_checker when using CPU offload.
yoonseokjin · Oct 31, 2022 · 07d0f28 · 07d0f28
1 parent 94cea75
commit 07d0f28
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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__(