Skip to content

Commit adabc57

Browse files
committed
Fix enable_sequential_cpu_offload in test_kandinsky_combined
1 parent 1826a1e commit adabc57

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -411,16 +411,16 @@ def __init__(
411411
def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
412412
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
413413

414-
def enable_sequential_cpu_offload(self, gpu_id=0):
414+
def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
415415
r"""
416416
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
417417
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
418418
`torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
419419
Note that offloading happens on a submodule basis. Memory savings are higher than with
420420
`enable_model_cpu_offload`, but performance is lower.
421421
"""
422-
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
423-
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id)
422+
self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
423+
self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device)
424424

425425
def progress_bar(self, iterable=None, total=None):
426426
self.prior_pipe.progress_bar(iterable=iterable, total=total)

0 commit comments

Comments
 (0)