resolve conflicts

Author: sayakpaul

.ai/references/testing.md

@@ -30,8 +30,8 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers
   - `MemoryTesterMixin` — CPU offload, group offload, layerwise casting.
   - Cache mixins — `PyramidAttentionBroadcastTesterMixin`, `FasterCacheTesterMixin`, `FirstBlockCacheTesterMixin`, `TaylorSeerCacheTesterMixin`, `MagCacheTesterMixin`. Guidance-distilled models override the cache config (e.g. `FASTER_CACHE_CONFIG = {... "is_guidance_distilled": True}`). Don't introduce caching related tests in the first iteration. These tests are added on a case-by-case basis.
   - In the first pass, just add tests related to `PipelineTesterMixin` and `MemoryTesterMixin`.
-- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; `enable_group_offload` keeps excluded components on the accelerator, so every other component stays covered — including the VAE, which the component-scoped `test_group_offloading_inference` leaves out. Block-level offloading is usually unaffected, hence the level in the name — a component that fails at both levels does need a skip.
-  - `torch.nn.MultiheadAttention` is the common instance: it passes `self.out_proj.weight` straight to `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`, so the hook on `out_proj` never fires. `SiglipVisionModel`'s attention pooling head wraps one — see `tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py`, whose `image_encoder` is excluded for this reason.
+- **Declare a component that can't be offloaded — don't hand-write a skip.** For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does); for a third-party component you can't annotate, such as a `transformers` encoder, name it in `group_offloading_leaf_level_exclude_modules` or `group_offloading_block_level_exclude_modules` on the config class. **Every `nn.Module` component is group offloaded unless the list for that level names it**, and the two levels fail on opposite hazards: leaf-level when compute reads a leaf's `.weight` instead of calling the leaf, block-level when a component re-enters submodules without going through the group leader's `forward` (VAE decode paths, hence the `vae` / `image_encoder` defaults). A component that fails at both goes in both, and a name matching no component fails the test as a typo. **Each level's test runs twice, with and without `use_stream`** — a subclass overriding one must re-declare `@MemoryTesterMixin._USE_STREAM`, or the override collapses to a single un-parametrized test that errors on the missing argument and reports as a green `xfail`.
+  - `torch.nn.MultiheadAttention` is the common leaf-level instance: it passes `self.out_proj.weight` straight to `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`, so the hook on `out_proj` never fires. `SiglipVisionModel`'s attention pooling head wraps one — see `tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py`, whose `image_encoder` is excluded for this reason.
   - `HunyuanDiTAttentionPool` (`src/diffusers/models/embeddings.py`) shows the same failure without an MHA module: a plain `nn.Module` that hands its `q_proj` / `k_proj` / `v_proj` / `c_proj` weights to `torch.nn.functional.multi_head_attention_forward`, so all four projections stay offloaded rather than just one. `HunyuanDiT2DModel` opts out of group offloading entirely with `_supports_group_offloading = False`.
   - Before adding a skip or an exclusion, confirm the failure still reproduces — several existing skips are stale, having outlived the upstream cause.
 - **A migration that surfaces a `src/` gap marks the test `xfail`, it does not patch the pipeline.** Give the marker a module-level name and a `reason` naming the exact gap (`PNDM_*` in `tests/pipelines/pndm/test_pndm.py` is the worked example), and prefer `strict=True` so the marker reports XPASS — and gets deleted — the day the pipeline is fixed. Use `strict=False` only when one mark covers a group whose members do not all fail. Marking a whole test class keeps the mixin's own marks (`@is_memory`, `@require_accelerator`) intact; overriding individual inherited tests drops the decorators they were declared with, so re-declare those too.

src/diffusers/models/transformers/transformer_prx.py

@@ -161,7 +161,10 @@ def __call__(
             ones_img = torch.ones((bs, l_img), dtype=torch.bool, device=device)
             attention_mask = attention_mask.to(device=device, dtype=torch.bool)
             joint_mask = torch.cat([attention_mask, ones_img], dim=-1)
-            attn_mask_tensor = joint_mask[:, None, None, :].expand(-1, attn.heads, l_img, -1)
+            # Every attention backend broadcasts the mask itself, so materialising
+            # [B, heads, L_img, L_all] only costs bandwidth: at batch 32, 1024 image tokens
+            # and 28 heads that is a 1.1 GiB mask read per block. Keep it broadcastable.
+            attn_mask_tensor = joint_mask[:, None, None, :]
 
         # Apply attention using dispatch_attention_fn for backend support
         # Reshape to match dispatch_attention_fn expectations: [B, L, H, D]

src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py

@@ -28,6 +28,7 @@
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
 from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer
 from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
+from ...models.modeling_utils import get_parameter_device
 from ...models.transformers.transformer_cosmos3 import (
     Cosmos3OmniTransformer,
 )
@@ -489,7 +490,7 @@ def _get_execution_device(self) -> torch.device:
                     return torch.device(execution_device)
 
             try:
-                return next(component.parameters()).device
+                return get_parameter_device(component)
             except StopIteration:
                 continue
 

tests/pipelines/audioldm2/test_audioldm2.py

@@ -469,11 +469,17 @@ class TestAudioLDM2PipelineMemory(AudioLDM2PipelineTesterConfig, MemoryTesterMix
     def test_sequential_cpu_offload_forward_pass(self):
         pass
 
-    @pytest.mark.skip(
-        "The pipeline encodes prompts through `text_encoder.get_text_features()` rather than the CLAP model's "
-        "`forward()`, so the top-level group-offloading hook never fires and the embedding weights stay offloaded."
+    # The pipeline encodes prompts through `text_encoder.get_text_features()` rather than the CLAP model's
+    # `forward()`. Block-level offloading gates the ungrouped weights on that `forward`, so they stay on the offload
+    # device; leaf level onloads each leaf on its own `forward` and is unaffected, so only block level is skipped.
+    CLAP_BLOCK_OFFLOAD_SKIP = pytest.mark.skip(
+        "`text_encoder.get_text_features()` bypasses the CLAP model's `forward()`, so block-level group offloading "
+        "never onloads the embedding weights."
     )
-    def test_group_offloading_inference(self):
+
+    @CLAP_BLOCK_OFFLOAD_SKIP
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
         pass
 
     @pytest.mark.skip("Not supported yet due to CLAPModel.")

tests/pipelines/ideogram4/test_pipeline_ideogram4.py

@@ -50,8 +50,10 @@ class Ideogram4PipelineTesterConfig(BasePipelineTesterConfig):
     required_input_params_in_call_signature = frozenset(["prompt", "height", "width", "guidance_scale"])
     batch_input_params = frozenset(["prompt"])
     output_shape = (3, 16, 16)
-    # `encode_prompt` drives the Qwen3-VL decoder layers directly instead of calling `text_encoder.forward`, so the
-    # offloading hooks would leave its inputs on the offload device. Keep the text encoder out of group offloading.
+    # `encode_prompt` calls the Qwen3-VL decoder layers directly instead of calling `text_encoder.forward`, and
+    # pins its inputs to `self.text_encoder.device`. Leaf-level hooks onload each leaf on its own forward while the
+    # module keeps reporting the offload device, so the inputs are left behind; block-level onloads the whole group
+    # up front and is unaffected, which is where the text encoder does get covered.
     group_offloading_leaf_level_exclude_modules = ["text_encoder"]
 
     def get_dummy_components(self, num_layers: int = 1):
@@ -285,7 +287,8 @@ class TestIdeogram4PipelineMemory(Ideogram4PipelineTesterConfig, MemoryTesterMix
     pins its inputs to `self.text_encoder.device` so they follow the weights under `enable_model_cpu_offload`
     (whose `CpuOffload` hook wraps the bypassed `forward` and so never fires). That pinning is wrong for every
     mechanism that hooks the submodules instead: they onload to the accelerator while the module still reports the
-    offload device, so the inputs are left behind. Hence the skips below.
+    offload device, so the inputs are left behind. Hence the skips below, and the text encoder's leaf-level group
+    offload exclusion on the config class.
     """
 
     _SUBMODULE_OFFLOAD_SKIP = (
@@ -307,17 +310,6 @@ def test_sequential_cpu_offload_forward_pass(self, base_pipe_output, expected_ma
     def test_sequential_offload_forward_pass_twice(self, expected_max_diff=2e-4):
         pass
 
-    @pytest.mark.skip(
-        reason=(
-            "Block-level group offloading cannot cover `text_encoder`: it leaves ungrouped leaves such as "
-            "`embed_tokens` to the root module's forward pre-hook, which never fires because `encode_prompt` "
-            "drives the decoder layers directly. Leaf-level offloading hooks those leaves individually and is "
-            "bit-exact here; only the block-level half of this test fails."
-        )
-    )
-    def test_group_offloading_inference(self):
-        pass
-
 
 class TestIdeogram4PipelineLoRA(Ideogram4PipelineTesterConfig, LoraTesterMixin):
     """LoRA tests for the Ideogram4 pipeline."""

tests/pipelines/kandinsky/test_kandinsky_prior.py

@@ -46,6 +46,16 @@
     "`clip_mean`, `clip_std`), which group offloading never onloads."
 )
 
+# A second, independent gap: the component-scoped test only offloads the denoiser under the names
+# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator.
+# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text
+# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline.
+COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = (
+    "The pipeline calls `PriorTransformer.post_process_latents()` after the denoising loop, which reads the "
+    "`clip_mean` / `clip_std` parameters held directly on the model. Group offloading onloads those only for the "
+    "duration of `forward`, so by then they are back on the offload device."
+)
+
 
 class KandinskyPriorPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = KandinskyPriorPipeline
@@ -205,6 +215,20 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=
 class TestKandinskyPriorPipelineMemory(KandinskyPriorPipelineTesterConfig, MemoryTesterMixin):
     """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky prior pipeline."""
 
+    @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
+        super().test_group_offloading_inference_block_level(
+            base_pipe_output, use_stream, expected_max_difference=expected_max_difference
+        )
+
+    @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_leaf_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
+        super().test_group_offloading_inference_leaf_level(
+            base_pipe_output, use_stream, expected_max_difference=expected_max_difference
+        )
+
     @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
     def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4):
         super().test_pipeline_level_group_offloading_inference(

tests/pipelines/kandinsky2_2/test_kandinsky_prior.py

@@ -48,6 +48,16 @@
     "`clip_mean`, `clip_std`), which group offloading never onloads."
 )
 
+# A second, independent gap: the component-scoped test only offloads the denoiser under the names
+# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator.
+# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text
+# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline.
+COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = (
+    "The pipeline calls `PriorTransformer.post_process_latents()` after the denoising loop, which reads the "
+    "`clip_mean` / `clip_std` parameters held directly on the model. Group offloading onloads those only for the "
+    "duration of `forward`, so by then they are back on the offload device."
+)
+
 
 class KandinskyV22PriorPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = KandinskyV22PriorPipeline
@@ -240,6 +250,20 @@ class TestKandinskyV22PriorPipelineMemory(KandinskyV22PriorPipelineTesterConfig,
     """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky 2.2 prior
     pipeline."""
 
+    @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
+        super().test_group_offloading_inference_block_level(
+            base_pipe_output, use_stream, expected_max_difference=expected_max_difference
+        )
+
+    @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_leaf_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
+        super().test_group_offloading_inference_leaf_level(
+            base_pipe_output, use_stream, expected_max_difference=expected_max_difference
+        )
+
     @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
     def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4):
         super().test_pipeline_level_group_offloading_inference(

tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py

@@ -54,6 +54,16 @@
     "`clip_mean`, `clip_std`), which group offloading never onloads."
 )
 
+# A second, independent gap: the component-scoped test only offloads the denoiser under the names
+# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator.
+# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text
+# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline.
+COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = (
+    "The pipeline calls `PriorTransformer.post_process_latents()` after the denoising loop, which reads the "
+    "`clip_mean` / `clip_std` parameters held directly on the model. Group offloading onloads those only for the "
+    "duration of `forward`, so by then they are back on the offload device."
+)
+
 
 class KandinskyV22PriorEmb2EmbPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = KandinskyV22PriorEmb2EmbPipeline
@@ -224,6 +234,20 @@ class TestKandinskyV22PriorEmb2EmbPipelineMemory(KandinskyV22PriorEmb2EmbPipelin
     """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky 2.2 prior
     emb2emb pipeline."""
 
+    @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
+        super().test_group_offloading_inference_block_level(
+            base_pipe_output, use_stream, expected_max_difference=expected_max_difference
+        )
+
+    @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
+    @MemoryTesterMixin._USE_STREAM
+    def test_group_offloading_inference_leaf_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4):
+        super().test_group_offloading_inference_leaf_level(
+            base_pipe_output, use_stream, expected_max_difference=expected_max_difference
+        )
+
     @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True)
     def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4):
         super().test_pipeline_level_group_offloading_inference(

tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py

@@ -47,6 +47,10 @@ class LTX2DFRPipelineTesterConfig(BasePipelineTesterConfig):
             "return_dict",
         ]
     )
+    # This config subclasses `BasePipelineTesterConfig` rather than `LTX2BaseTesterConfig`, so it has to restate the
+    # family's `audio_vae` exclusion: its decode-time convolutions run without the group leader's `forward` having
+    # onloaded the group.
+    group_offloading_block_level_exclude_modules = ["vae", "audio_vae"]
 
     def get_dummy_components(self):
         return get_dfr_dummy_components()

tests/pipelines/ltx2/testing_utils.py

@@ -21,7 +21,6 @@
 does not restate them.
 """
 
-import pytest
 import torch
 from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
 
@@ -283,6 +282,10 @@ class LTX2BaseTesterConfig(BasePipelineTesterConfig):
     # See `DEFAULT_UNSET_COMPONENTS`; each config narrows or widens this to the components its pipeline accepts.
     unset_components = DEFAULT_UNSET_COMPONENTS
 
+    # `audio_vae` fails at block level for the same reason the other VAEs do: its decode-time convolutions run
+    # without the group leader's `forward` having onloaded the group.
+    group_offloading_block_level_exclude_modules = ["vae", "audio_vae"]
+
     def get_dummy_components(self):
         return get_ltx2_dummy_components(unset_components=self.unset_components)
 
@@ -299,15 +302,6 @@ def get_pipeline_with_duration_head(self):
 class LTX2MemoryTesterMixin(MemoryTesterMixin):
     """`MemoryTesterMixin` for the LTX2 pipelines in this directory."""
 
-    # The shared helper only group-offloads `text_encoder` / `transformer` and moves `vae`, leaving LTX2's extra
-    # module components (`connectors`, `audio_vae`, `vocoder`) on the CPU while they receive accelerator tensors
-    # from the offloaded text encoder, so the forward pass mixes devices. This is pre-existing for the whole LTX2
-    # family rather than specific to any one pipeline. Pipeline-level offloading, which walks every component, is
-    # exercised by `test_pipeline_level_group_offloading_inference`.
-    @pytest.mark.skip("Using test_pipeline_level_group_offloading_inference instead")
-    def test_group_offloading_inference(self):
-        pass
-
 
 class LTX2LoraTesterMixin(LoraTesterMixin):
     """`LoraTesterMixin` for the LTX2 pipelines in this directory.