[hooks] Support Qwen Image 2.1 in MagCache

src/diffusers/hooks/_helpers.py

@@ -196,6 +196,7 @@ def _register_transformer_blocks_metadata():
         MotifVideoTransformerBlock,
     )
     from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock
+    from ..models.transformers.transformer_qwenimage21 import QwenImage21TransformerBlock
     from ..models.transformers.transformer_wan import WanTransformerBlock
     from ..models.transformers.transformer_z_image import ZImageTransformerBlock
 
@@ -358,6 +359,15 @@ def _register_transformer_blocks_metadata():
         ),
     )
 
+    # QwenImage2.1
+    TransformerBlockRegistry.register(
+        model_class=QwenImage21TransformerBlock,
+        metadata=TransformerBlockMetadata(
+            return_hidden_states_index=0,
+            return_encoder_hidden_states_index=None,
+        ),
+    )
+
     # HunyuanImage2.1
     TransformerBlockRegistry.register(
         model_class=HunyuanImageTransformerBlock,

src/diffusers/hooks/mag_cache.py

@@ -249,6 +249,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 if diff > 0:
                     output = output.clone()
                     output[:, diff:, :] = output[:, diff:, :] + res
+                elif diff < 0:
+                    # The residual was computed on a longer sequence (e.g. the QwenImage2.1 KV-cache prefill step
+                    # covers prefix + target tokens, later steps only the target); the target tokens sit at the end.
+                    output = output + res[:, -output.shape[1] :, :]
                 else:
                     logger.warning(
                         f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. "
@@ -371,6 +375,12 @@ def _perform_calibration_step(self, state: MagCacheState, current_residual: torc
             # norm(dim=-1) gives magnitude of each token vector
             curr_norm = torch.linalg.norm(current_residual.float(), dim=-1)
             prev_norm = torch.linalg.norm(state.previous_residual.float(), dim=-1)
+            if curr_norm.ndim == 2 and prev_norm.ndim == 2 and curr_norm.shape[1] != prev_norm.shape[1]:
+                # Sequence length changed between steps (e.g. KV-cache prefill vs. cached steps); compare the
+                # trailing tokens both steps share.
+                num_tokens = min(curr_norm.shape[1], prev_norm.shape[1])
+                curr_norm = curr_norm[:, -num_tokens:]
+                prev_norm = prev_norm[:, -num_tokens:]
 
             # Avoid division by zero
             ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()

tests/models/transformers/test_models_transformer_qwenimage21.py

@@ -25,6 +25,7 @@
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
+    MagCacheTesterMixin,
     MemoryTesterMixin,
     ModelTesterMixin,
     SingleFileTesterMixin,
@@ -364,6 +365,10 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode):
         assert not torch.allclose(with_kv_cache[3], uncached[3], atol=1e-5)
 
 
+class TestQwenImage21TransformerMagCache(QwenImage21TransformerTesterConfig, MagCacheTesterMixin):
+    """MagCache tests for QwenImage 2.1 Transformer."""
+
+
 class TestQwenImage21TransformerSingleFile(QwenImage21TransformerTesterConfig, SingleFileTesterMixin):
     @property
     def ckpt_path(self):

tests/pipelines/qwenimage21/test_qwenimage21.py

@@ -38,6 +38,7 @@
     BasePipelineTesterConfig,
     LoraMemoryTesterMixin,
     LoraTesterMixin,
+    MagCacheTesterMixin,
     MemoryTesterMixin,
     PipelineTesterMixin,
 )
@@ -347,3 +348,7 @@ def test_simple_inference_with_text_denoiser_lora_and_scale(self, base_pipe_outp
 
 class TestQwenImage21PipelineLoRAMemory(QwenImage21PipelineTesterConfig, LoraMemoryTesterMixin):
     """LoRA x memory-optimization tests for the Qwen-Image 2.1 pipeline."""
+
+
+class TestQwenImage21PipelineMagCache(QwenImage21PipelineTesterConfig, MagCacheTesterMixin):
+    """MagCache tests for the Qwen-Image 2.1 pipeline."""