Skip to content

GitLab

Unverified Commit bf7b0bc2 authored by Pedro Cuenca's avatar Pedro Cuenca Committed by GitHub
Browse files

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

Allow None safety_checker when using CPU offload.
parent e4d264e4
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
View file @ bf7b0bc2
...@@ -133,7 +133,8 @@ class StableDiffusionPipeline(DiffusionPipeline): ...@@ -133,7 +133,8 @@ class StableDiffusionPipeline(DiffusionPipeline):
device = torch.device("cuda") device = torch.device("cuda")
for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]: 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() @torch.no_grad()
def __call__( def __call__(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment