diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index b4eef4c621b0..cdb5aff095a3 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2171,6 +2171,16 @@ def _execution_device(self): [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from Accelerate's module hooks. """ + from ..hooks.group_offloading import _get_group_onload_device + + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module): + continue + try: + return _get_group_onload_device(model) + except ValueError: + pass + for name, model in self.components.items(): if not isinstance(model, torch.nn.Module): continue diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 4880aee08ba1..70e30b753693 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -371,6 +371,33 @@ def test_components_auto_cpu_offload_inference_consistent(self): assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + @require_accelerator + def test_group_offloading_execution_device(self): + from diffusers.hooks import apply_group_offloading + + pipe = self.get_pipeline().to("cpu") + assert pipe._execution_device.type == "cpu" + + offloaded = None + for name, component in pipe.components.items(): + if not isinstance(component, torch.nn.Module): + continue + if not getattr(component, "_supports_group_offloading", True): + continue + apply_group_offloading( + component, + onload_device=torch.device(torch_device), + offload_device=torch.device("cpu"), + offload_type="leaf_level", + ) + offloaded = name + break + + if offloaded is None: + pytest.skip("No component supports group offloading.") + + assert pipe._execution_device.type == torch.device(torch_device).type + def test_save_from_pretrained(self, tmp_path): pipes = [] base_pipe = self.get_pipeline().to(torch_device)