Skip to content

Commit 1ca0943

Browse files
authored
[core] DiffusionPipeline.device deduction for split-device pipelines (#14383)
* feat:redefine device deduction * test: unit test * review: remove all cpu check
1 parent f6e1c4d commit 1ca0943

2 files changed

Lines changed: 37 additions & 1 deletion

File tree

src/diffusers/pipelines/pipeline_utils.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -601,12 +601,20 @@ def module_is_offloaded(module):
601601
def device(self) -> torch.device:
602602
r"""
603603
Returns:
604-
`torch.device`: The torch device on which the pipeline is located.
604+
`torch.device`: The torch device on which the pipeline is located. When components are split across devices
605+
(for example, text encoders on CPU while the denoising backbone runs on an accelerator), the accelerator
606+
device is returned.
605607
"""
606608
module_names, _ = self._get_signature_keys(self)
607609
modules = [getattr(self, n, None) for n in module_names]
608610
modules = [m for m in modules if isinstance(m, torch.nn.Module)]
609611

612+
# Prefer a non-CPU, non-meta component so a split pipeline reports the accelerator it computes on,
613+
# rather than whichever component happens to sort first.
614+
for module in modules:
615+
if module.device.type not in ("cpu", "meta"):
616+
return module.device
617+
610618
for module in modules:
611619
return module.device
612620

tests/pipelines/test_pipelines.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1944,6 +1944,34 @@ def test_pipe_to(self):
19441944
assert sd1.device.type == device_type
19451945
assert sd2.device.type == device_type
19461946

1947+
@require_torch_accelerator
1948+
def test_pipe_device_split_across_devices(self):
1949+
unet = self.dummy_cond_unet()
1950+
scheduler = PNDMScheduler(skip_prk_steps=True)
1951+
vae = self.dummy_vae
1952+
bert = self.dummy_text_encoder
1953+
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
1954+
1955+
sd = StableDiffusionPipeline(
1956+
unet=unet,
1957+
scheduler=scheduler,
1958+
vae=vae,
1959+
text_encoder=bert,
1960+
tokenizer=tokenizer,
1961+
safety_checker=None,
1962+
feature_extractor=self.dummy_extractor,
1963+
)
1964+
1965+
device_type = torch.device(torch_device).type
1966+
1967+
# Text encoder stays on CPU while the denoising backbone runs on the accelerator. `text_encoder` sorts
1968+
# before `unet`/`vae`, so a first-component rule would report `cpu` here.
1969+
sd.unet.to(torch_device)
1970+
sd.vae.to(torch_device)
1971+
1972+
assert sd.text_encoder.device.type == "cpu"
1973+
assert sd.device.type == device_type
1974+
19471975
def test_pipe_same_device_id_offload(self):
19481976
unet = self.dummy_cond_unet()
19491977
scheduler = PNDMScheduler(skip_prk_steps=True)

0 commit comments

Comments
 (0)