address review comments

Author: a-r-r-o-w

src/diffusers/hooks/group_offloading.py

@@ -343,7 +343,7 @@ def _apply_group_offloading_block_level(
         for i in range(0, len(submodule), num_blocks_per_group):
             current_modules = submodule[i : i + num_blocks_per_group]
             group = ModuleGroup(
-                modules=submodule[i : i + num_blocks_per_group],
+                modules=current_modules,
                 offload_device=offload_device,
                 onload_device=onload_device,
                 offload_leader=current_modules[-1],

src/diffusers/models/autoencoders/autoencoder_oobleck.py

@@ -317,6 +317,7 @@ class AutoencoderOobleck(ModelMixin, ConfigMixin):
     """
 
     _supports_gradient_checkpointing = False
+    _supports_group_offloading = False
 
     @register_to_config
     def __init__(

src/diffusers/models/autoencoders/consistency_decoder_vae.py

@@ -68,6 +68,8 @@ class ConsistencyDecoderVAE(ModelMixin, ConfigMixin):
         ```
     """
 
+    _supports_group_offloading = False
+
     @register_to_config
     def __init__(
         self,

src/diffusers/models/autoencoders/vq_model.py

@@ -72,6 +72,7 @@ class VQModel(ModelMixin, ConfigMixin):
     """
 
     _skip_layerwise_casting_patterns = ["quantize"]
+    _supports_group_offloading = False
 
     @register_to_config
     def __init__(

src/diffusers/models/modeling_utils.py

@@ -174,6 +174,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
     _no_split_modules = None
     _keep_in_fp32_modules = None
     _skip_layerwise_casting_patterns = None
+    _supports_group_offloading = True
 
     def __init__(self):
         super().__init__()

src/diffusers/models/transformers/dit_transformer_2d.py

@@ -66,6 +66,7 @@ class DiTTransformer2DModel(ModelMixin, ConfigMixin):
 
     _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
     _supports_gradient_checkpointing = True
+    _supports_group_offloading = False
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/hunyuan_transformer_2d.py

@@ -245,6 +245,7 @@ class HunyuanDiT2DModel(ModelMixin, ConfigMixin):
     """
 
     _skip_layerwise_casting_patterns = ["pos_embed", "norm", "pooler"]
+    _supports_group_offloading = False
 
     @register_to_config
     def __init__(

src/diffusers/pipelines/pipeline_utils.py

@@ -1020,25 +1020,19 @@ def _execution_device(self):
         [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from
         Accelerate's module hooks.
         """
-        diffusers_hook_device = None
+        # 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:
-                    diffusers_hook_device = hook.group.onload_device
-                    break
-
-            if diffusers_hook_device is not None:
-                break
-
-        if diffusers_hook_device is not None:
-            return diffusers_hook_device
+                    return hook.group.onload_device
 
         for name, model in self.components.items():
             if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload:

tests/models/autoencoders/test_models_autoencoder_oobleck.py

@@ -132,15 +132,6 @@ def test_layerwise_casting_inference(self):
     def test_layerwise_casting_memory(self):
         pass
 
-    @unittest.skip(
-        "The convolution layers of AutoencoderOobleck are wrapped with torch.nn.utils.weight_norm. This causes the hook's pre_forward to not "
-        "cast the module weights to the expected device (as required by forward pass). As a result, forward pass errors out. To fix:\n"
-        "1. Make sure `nn::Module::to(device)` works with `torch.nn.utils.weight_norm` wrapped convolution layer.\n"
-        "2. Unskip this test."
-    )
-    def test_group_offloading(self):
-        pass
-
 
 @slow
 class AutoencoderOobleckIntegrationTests(unittest.TestCase):

tests/models/autoencoders/test_models_consistency_decoder_vae.py

@@ -155,10 +155,6 @@ def test_enable_disable_slicing(self):
             "Without slicing outputs should match with the outputs when slicing is manually disabled.",
         )
 
-    @unittest.skip("Not quite sure why this test fails and unable to debug.")
-    def test_group_offloading(self):
-        pass
-
 
 @slow
 class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase):