AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

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

Browse Source 
Allow None safety_checker when using CPU offload.
This commit is contained in:
Pedro Cuenca
2022-10-31 15:03:33 +01:00
committed by GitHub
parent e4d264e4eb
commit bf7b0bc25b
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
View File

@@ -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__(