Skip to content

Commit bf7b0bc

Browse files
authored
Allow safety_checker to be None when using CPU offload (#1078)
Allow None safety_checker when using CPU offload.
1 parent e4d264e commit bf7b0bc

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,8 @@ def enable_sequential_cpu_offload(self):
133133
device = torch.device("cuda")
134134

135135
for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
136-
cpu_offload(cpu_offloaded_model, device)
136+
if cpu_offloaded_model is not None:
137+
cpu_offload(cpu_offloaded_model, device)
137138

138139
@torch.no_grad()
139140
def __call__(

0 commit comments

Comments
 (0)