Mirrors
/
diffusers
mirror of https://github.com/ShivamShrirao/diffusers.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

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
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats 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") 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__(