diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index a683973df5d7..24fe0eabfa6f 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -601,12 +601,20 @@ def module_is_offloaded(module): def device(self) -> torch.device: r""" Returns: - `torch.device`: The torch device on which the pipeline is located. + `torch.device`: The torch device on which the pipeline is located. When components are split across devices + (for example, text encoders on CPU while the denoising backbone runs on an accelerator), the accelerator + device is returned. """ module_names, _ = self._get_signature_keys(self) modules = [getattr(self, n, None) for n in module_names] modules = [m for m in modules if isinstance(m, torch.nn.Module)] + # Prefer a non-CPU, non-meta component so a split pipeline reports the accelerator it computes on, + # rather than whichever component happens to sort first. + for module in modules: + if module.device.type not in ("cpu", "meta"): + return module.device + for module in modules: return module.device diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index d975ef8bd2d8..5feef16871ee 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -1944,6 +1944,38 @@ def test_pipe_to(self): assert sd1.device.type == device_type assert sd2.device.type == device_type + @require_torch_accelerator + def test_pipe_device_split_across_devices(self): + unet = self.dummy_cond_unet() + scheduler = PNDMScheduler(skip_prk_steps=True) + vae = self.dummy_vae + bert = self.dummy_text_encoder + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + + sd = StableDiffusionPipeline( + unet=unet, + scheduler=scheduler, + vae=vae, + text_encoder=bert, + tokenizer=tokenizer, + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + + device_type = torch.device(torch_device).type + + # Text encoder stays on CPU while the denoising backbone runs on the accelerator. `text_encoder` sorts + # before `unet`/`vae`, so a first-component rule would report `cpu` here. + sd.unet.to(torch_device) + sd.vae.to(torch_device) + + assert sd.text_encoder.device.type == "cpu" + assert sd.device.type == device_type + + # With every component on CPU there is no accelerator to prefer. + sd.to("cpu") + assert sd.device.type == "cpu" + def test_pipe_same_device_id_offload(self): unet = self.dummy_cond_unet() scheduler = PNDMScheduler(skip_prk_steps=True)