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

Allow None safety_checker when using CPU offload.

Allow `safety_checker` to be `None` when using CPU offload (#1078)
Allow None safety_checker when using CPU offload.
bf7b0bc2 · Pedro Cuenca · GitHub · e4d264e4 · bf7b0bc2
Unverified Commit bf7b0bc2 authored Oct 31, 2022 by Pedro Cuenca Committed by GitHub Oct 31, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py ...s/pipelines/stable_diffusion/pipeline_stable_diffusion.py +2 -1

No files found.
--- 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__(