refactor some repeated code

Author: a-r-r-o-w

src/diffusers/hooks/group_offloading.py

@@ -669,3 +669,10 @@ def _is_group_offload_enabled(module: torch.nn.Module) -> bool:
         if hasattr(submodule, "_diffusers_hook") and submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING) is not None:
             return True
     return False
+
+
+def _get_group_onload_device(module: torch.nn.Module) -> torch.device:
+    for submodule in module.modules():
+        if hasattr(submodule, "_diffusers_hook") and submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING) is not None:
+            return submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING).group.onload_device
+    raise ValueError("Group offloading is not enabled for the provided module.")

src/diffusers/models/modeling_utils.py

@@ -86,16 +86,17 @@
 
 
 def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
+    from ..hooks.group_offloading import _get_group_onload_device
+
     try:
-        if hasattr(parameter, "_diffusers_hook"):
-            for submodule in parameter.modules():
-                if not hasattr(submodule, "_diffusers_hook"):
-                    continue
-                registry = parameter._diffusers_hook
-                hook = registry.get_hook("group_offloading")
-                if hook is not None:
-                    return hook.group.onload_device
+        # Try to get the onload device from the group offloading hook
+        return _get_group_onload_device(parameter)
+    except ValueError:
+        pass
 
+    try:
+        # If the onload device is not available due to no group offloading hooks, try to get the device
+        # from the first parameter or buffer
         parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
         return next(parameters_and_buffers).device
     except StopIteration:

src/diffusers/pipelines/pipeline_utils.py

@@ -462,7 +462,7 @@ def module_is_offloaded(module):
                     f"The module '{module.__class__.__name__}' has been loaded in `bitsandbytes` 8bit and moving it to {device} via `.to()` is not supported. Module is still on {module.device}."
                 )
 
-            # Note: we also handle this as the ModelMixin level. The reason for doing it here too is that modeling
+            # Note: we also handle this at the ModelMixin level. The reason for doing it here too is that modeling
             # components can be from outside diffusers too, but still have group offloading enabled.
             if (
                 self._maybe_raise_error_if_group_offload_active(raise_error=False, module=module)
@@ -1035,19 +1035,18 @@ 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
+
         # When apply group offloading at the leaf_level, we're in the same situation as accelerate's sequential
         # offloading. We need to return the onload device of the group offloading hooks so that the intermediates
         # required for computation (latents, prompt embeddings, etc.) can be created on the correct device.
         for name, model in self.components.items():
             if not isinstance(model, torch.nn.Module):
                 continue
-            for submodule in model.modules():
-                if not hasattr(submodule, "_diffusers_hook"):
-                    continue
-                registry = submodule._diffusers_hook
-                hook = registry.get_hook("group_offloading")
-                if hook is not None:
-                    return hook.group.onload_device
+            try:
+                return _get_group_onload_device(model)
+            except ValueError:
+                pass
 
         for name, model in self.components.items():
             if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload: