Unverified Commit 09d0546a authored by dblunk88's avatar dblunk88 Committed by GitHub
Browse files

cpu offloading: mutli GPU support (#1143)

mutli GPU support
parent 65d136e0
...@@ -178,7 +178,7 @@ class StableDiffusionPipeline(DiffusionPipeline): ...@@ -178,7 +178,7 @@ class StableDiffusionPipeline(DiffusionPipeline):
# set slice_size = `None` to disable `attention slicing` # set slice_size = `None` to disable `attention slicing`
self.enable_attention_slicing(None) self.enable_attention_slicing(None)
def enable_sequential_cpu_offload(self): def enable_sequential_cpu_offload(self, gpu_id=0):
r""" r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, 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 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): ...@@ -189,7 +189,7 @@ class StableDiffusionPipeline(DiffusionPipeline):
else: else:
raise ImportError("Please install accelerate via `pip install accelerate`") 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]: for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
if cpu_offloaded_model is not None: if cpu_offloaded_model is not None:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment