Allow `safety_checker` to be `None` when using CPU offload (#1078)

Allow None safety_checker when using CPU offload.
2022-10-31 15:03:33 +01:00 · 2022-10-31 15:03:33 +01:00 · bf7b0bc25b
parent e4d264e4eb
commit bf7b0bc25b
1 changed files with 2 additions and 1 deletions
--- 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__(