From 09d0546ad0489b84cd8930c2267d9097fa4f19e2 Mon Sep 17 00:00:00 2001 From: dblunk88 <39381389+dblunk88@users.noreply.github.com> Date: Wed, 16 Nov 2022 11:40:16 -0500 Subject: [PATCH] cpu offloading: mutli GPU support (#1143) mutli GPU support --- .../pipelines/stable_diffusion/pipeline_stable_diffusion.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 65922451f0..963d75c58b 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -178,7 +178,7 @@ class StableDiffusionPipeline(DiffusionPipeline): # set slice_size = `None` to disable `attention slicing` self.enable_attention_slicing(None) - def enable_sequential_cpu_offload(self): + def enable_sequential_cpu_offload(self, gpu_id=0): r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a @@ -189,7 +189,7 @@ class StableDiffusionPipeline(DiffusionPipeline): else: raise ImportError("Please install accelerate via `pip install accelerate`") - device = torch.device("cuda") + device = torch.device(f"cuda:{gpu_id}") for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]: if cpu_offloaded_model is not None: