style

* add magcache

* formatting

* add magcache support with calibration mode

* add imports

* improvements

* Apply style fixes

* fix kandinsky errors

* add tests and documentation

* Apply style fixes

* improvements

* Apply style fixes

* make fix-copies.

* minor fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

docs/source/en/optimization/cache.md

@@ -111,3 +111,57 @@ config = TaylorSeerCacheConfig(
 )
 pipe.transformer.enable_cache(config)
 ```
+
+## MagCache
+
+[MagCache](https://github.com/Zehong-Ma/MagCache) accelerates inference by skipping transformer blocks based on the magnitude of the residual update. It observes that the magnitude of updates (Output - Input) decays predictably over the diffusion process. By accumulating an "error budget" based on pre-computed magnitude ratios, it dynamically decides when to skip computation and reuse the previous residual.
+
+MagCache relies on **Magnitude Ratios** (`mag_ratios`), which describe this decay curve. These ratios are specific to the model checkpoint and scheduler.
+
+### Usage
+
+To use MagCache, you typically follow a two-step process: **Calibration** and **Inference**.
+
+1.  **Calibration**: Run inference once with `calibrate=True`. The hook will measure the residual magnitudes and print the calculated ratios to the console.
+2.  **Inference**: Pass these ratios to `MagCacheConfig` to enable acceleration.
+
+```python
+import torch
+from diffusers import FluxPipeline, MagCacheConfig
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16
+).to("cuda")
+
+# 1. Calibration Step
+# Run full inference to measure model behavior.
+calib_config = MagCacheConfig(calibrate=True, num_inference_steps=4)
+pipe.transformer.enable_cache(calib_config)
+
+# Run a prompt to trigger calibration
+pipe("A cat playing chess", num_inference_steps=4)
+# Logs will print something like: "MagCache Calibration Results: [1.0, 1.37, 0.97, 0.87]"
+
+# 2. Inference Step
+# Apply the specific ratios obtained from calibration for optimized speed.
+# Note: For Flux models, you can also import defaults: 
+# from diffusers.hooks.mag_cache import FLUX_MAG_RATIOS
+mag_config = MagCacheConfig(
+    mag_ratios=[1.0, 1.37, 0.97, 0.87],
+    num_inference_steps=4
+)
+
+pipe.transformer.enable_cache(mag_config) 
+
+image = pipe("A cat playing chess", num_inference_steps=4).images[0]
+```
+
+> [!NOTE]
+> `mag_ratios` represent the model's intrinsic magnitude decay curve. Ratios calibrated for a high number of steps (e.g., 50) can be reused for lower step counts (e.g., 20). The implementation uses interpolation to map the curve to the current number of inference steps.
+
+> [!TIP]
+> For pipelines that run Classifier-Free Guidance sequentially (like Kandinsky 5.0), the calibration log might print two arrays: one for the Conditional pass and one for the Unconditional pass. In most cases, you should use the first array (Conditional).
+
+> [!TIP]
+> For pipelines that run Classifier-Free Guidance in a **batched** manner (like SDXL or Flux), the `hidden_states` processed by the model contain both conditional and unconditional branches concatenated together. The calibration process automatically accounts for this, producing a single array of ratios that represents the joint behavior. You can use this resulting array directly without modification.

docs/source/en/quantization/torchao.md

@@ -66,7 +66,7 @@ from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConf
 from torchao.quantization import Int4WeightOnlyConfig
 
 pipeline_quant_config = PipelineQuantizationConfig(
-    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
+    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128)))}
 )
 pipeline = DiffusionPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",

src/diffusers/__init__.py

@@ -168,12 +168,14 @@ else:
             "FirstBlockCacheConfig",
             "HookRegistry",
             "LayerSkipConfig",
+            "MagCacheConfig",
             "PyramidAttentionBroadcastConfig",
             "SmoothedEnergyGuidanceConfig",
             "TaylorSeerCacheConfig",
             "apply_faster_cache",
             "apply_first_block_cache",
             "apply_layer_skip",
+            "apply_mag_cache",
             "apply_pyramid_attention_broadcast",
             "apply_taylorseer_cache",
         ]
@@ -932,12 +934,14 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             FirstBlockCacheConfig,
             HookRegistry,
             LayerSkipConfig,
+            MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             SmoothedEnergyGuidanceConfig,
             TaylorSeerCacheConfig,
             apply_faster_cache,
             apply_first_block_cache,
             apply_layer_skip,
+            apply_mag_cache,
             apply_pyramid_attention_broadcast,
             apply_taylorseer_cache,
         )

src/diffusers/hooks/__init__.py

@@ -23,6 +23,7 @@ if is_torch_available():
     from .hooks import HookRegistry, ModelHook
     from .layer_skip import LayerSkipConfig, apply_layer_skip
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
+    from .mag_cache import MagCacheConfig, apply_mag_cache
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
     from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache

src/diffusers/hooks/_common.py

@@ -23,7 +23,13 @@ from ..models.attention_processor import Attention, MochiAttention
 _ATTENTION_CLASSES = (Attention, MochiAttention, AttentionModuleMixin)
 _FEEDFORWARD_CLASSES = (FeedForward, LuminaFeedForward)
 
-_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks", "layers")
+_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = (
+    "blocks",
+    "transformer_blocks",
+    "single_transformer_blocks",
+    "layers",
+    "visual_transformer_blocks",
+)
 _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",)
 _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "layers")
 

src/diffusers/hooks/_helpers.py

@@ -26,6 +26,7 @@ class AttentionProcessorMetadata:
 class TransformerBlockMetadata:
     return_hidden_states_index: int = None
     return_encoder_hidden_states_index: int = None
+    hidden_states_argument_name: str = "hidden_states"
 
     _cls: Type = None
     _cached_parameter_indices: Dict[str, int] = None
@@ -169,7 +170,7 @@ def _register_attention_processors_metadata():
 
 
 def _register_transformer_blocks_metadata():
-    from ..models.attention import BasicTransformerBlock
+    from ..models.attention import BasicTransformerBlock, JointTransformerBlock
     from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
     from ..models.transformers.transformer_bria import BriaTransformerBlock
     from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
@@ -184,6 +185,7 @@ def _register_transformer_blocks_metadata():
         HunyuanImageSingleTransformerBlock,
         HunyuanImageTransformerBlock,
     )
+    from ..models.transformers.transformer_kandinsky import Kandinsky5TransformerDecoderBlock
     from ..models.transformers.transformer_ltx import LTXVideoTransformerBlock
     from ..models.transformers.transformer_mochi import MochiTransformerBlock
     from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock
@@ -331,6 +333,24 @@ def _register_transformer_blocks_metadata():
         ),
     )
 
+    TransformerBlockRegistry.register(
+        model_class=JointTransformerBlock,
+        metadata=TransformerBlockMetadata(
+            return_hidden_states_index=1,
+            return_encoder_hidden_states_index=0,
+        ),
+    )
+
+    # Kandinsky 5.0 (Kandinsky5TransformerDecoderBlock)
+    TransformerBlockRegistry.register(
+        model_class=Kandinsky5TransformerDecoderBlock,
+        metadata=TransformerBlockMetadata(
+            return_hidden_states_index=0,
+            return_encoder_hidden_states_index=None,
+            hidden_states_argument_name="visual_embed",
+        ),
+    )
+
 
 # fmt: off
 def _skip_attention___ret___hidden_states(self, *args, **kwargs):

src/diffusers/hooks/mag_cache.py

@@ -0,0 +1,468 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+
+from ..utils import get_logger
+from ..utils.torch_utils import unwrap_module
+from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS
+from ._helpers import TransformerBlockRegistry
+from .hooks import BaseState, HookRegistry, ModelHook, StateManager
+
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+_MAG_CACHE_LEADER_BLOCK_HOOK = "mag_cache_leader_block_hook"
+_MAG_CACHE_BLOCK_HOOK = "mag_cache_block_hook"
+
+# Default Mag Ratios for Flux models (Dev/Schnell) are provided for convenience.
+# Users must explicitly pass these to the config if using Flux.
+# Reference: https://github.com/Zehong-Ma/MagCache
+FLUX_MAG_RATIOS = torch.tensor(
+    [1.0]
+    + [
+        1.21094,
+        1.11719,
+        1.07812,
+        1.0625,
+        1.03906,
+        1.03125,
+        1.03906,
+        1.02344,
+        1.03125,
+        1.02344,
+        0.98047,
+        1.01562,
+        1.00781,
+        1.0,
+        1.00781,
+        1.0,
+        1.00781,
+        1.0,
+        1.0,
+        0.99609,
+        0.99609,
+        0.98047,
+        0.98828,
+        0.96484,
+        0.95703,
+        0.93359,
+        0.89062,
+    ]
+)
+
+
+def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor:
+    """
+    Interpolate the source array to the target length using nearest neighbor interpolation.
+    """
+    src_length = len(src_array)
+    if target_length == 1:
+        return src_array[-1:]
+
+    scale = (src_length - 1) / (target_length - 1)
+    grid = torch.arange(target_length, device=src_array.device, dtype=torch.float32)
+    mapped_indices = torch.round(grid * scale).long()
+    return src_array[mapped_indices]
+
+
+@dataclass
+class MagCacheConfig:
+    r"""
+    Configuration for [MagCache](https://github.com/Zehong-Ma/MagCache).
+
+    Args:
+        threshold (`float`, defaults to `0.06`):
+            The threshold for the accumulated error. If the accumulated error is below this threshold, the block
+            computation is skipped. A higher threshold allows for more aggressive skipping (faster) but may degrade
+            quality.
+        max_skip_steps (`int`, defaults to `3`):
+            The maximum number of consecutive steps that can be skipped (K in the paper).
+        retention_ratio (`float`, defaults to `0.2`):
+            The fraction of initial steps during which skipping is disabled to ensure stability. For example, if
+            `num_inference_steps` is 28 and `retention_ratio` is 0.2, the first 6 steps will never be skipped.
+        num_inference_steps (`int`, defaults to `28`):
+            The number of inference steps used in the pipeline. This is required to interpolate `mag_ratios` correctly.
+        mag_ratios (`torch.Tensor`, *optional*):
+            The pre-computed magnitude ratios for the model. These are checkpoint-dependent. If not provided, you must
+            set `calibrate=True` to calculate them for your specific model. For Flux models, you can use
+            `diffusers.hooks.mag_cache.FLUX_MAG_RATIOS`.
+        calibrate (`bool`, defaults to `False`):
+            If True, enables calibration mode. In this mode, no blocks are skipped. Instead, the hook calculates the
+            magnitude ratios for the current run and logs them at the end. Use this to obtain `mag_ratios` for new
+            models or schedulers.
+    """
+
+    threshold: float = 0.06
+    max_skip_steps: int = 3
+    retention_ratio: float = 0.2
+    num_inference_steps: int = 28
+    mag_ratios: Optional[Union[torch.Tensor, List[float]]] = None
+    calibrate: bool = False
+
+    def __post_init__(self):
+        # User MUST provide ratios OR enable calibration.
+        if self.mag_ratios is None and not self.calibrate:
+            raise ValueError(
+                " `mag_ratios` must be provided for MagCache inference because these ratios are model-dependent.\n"
+                "To get them for your model:\n"
+                "1. Initialize `MagCacheConfig(calibrate=True, ...)`\n"
+                "2. Run inference on your model once.\n"
+                "3. Copy the printed ratios array and pass it to `mag_ratios` in the config.\n"
+                "For Flux models, you can import `FLUX_MAG_RATIOS` from `diffusers.hooks.mag_cache`."
+            )
+
+        if not self.calibrate and self.mag_ratios is not None:
+            if not torch.is_tensor(self.mag_ratios):
+                self.mag_ratios = torch.tensor(self.mag_ratios)
+
+            if len(self.mag_ratios) != self.num_inference_steps:
+                logger.debug(
+                    f"Interpolating mag_ratios from length {len(self.mag_ratios)} to {self.num_inference_steps}"
+                )
+                self.mag_ratios = nearest_interp(self.mag_ratios, self.num_inference_steps)
+
+
+class MagCacheState(BaseState):
+    def __init__(self) -> None:
+        super().__init__()
+        # Cache for the residual (output - input) from the *previous* timestep
+        self.previous_residual: torch.Tensor = None
+
+        # State inputs/outputs for the current forward pass
+        self.head_block_input: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None
+        self.should_compute: bool = True
+
+        # MagCache accumulators
+        self.accumulated_ratio: float = 1.0
+        self.accumulated_err: float = 0.0
+        self.accumulated_steps: int = 0
+
+        # Current step counter (timestep index)
+        self.step_index: int = 0
+
+        # Calibration storage
+        self.calibration_ratios: List[float] = []
+
+    def reset(self):
+        self.previous_residual = None
+        self.should_compute = True
+        self.accumulated_ratio = 1.0
+        self.accumulated_err = 0.0
+        self.accumulated_steps = 0
+        self.step_index = 0
+        self.calibration_ratios = []
+
+
+class MagCacheHeadHook(ModelHook):
+    _is_stateful = True
+
+    def __init__(self, state_manager: StateManager, config: MagCacheConfig):
+        self.state_manager = state_manager
+        self.config = config
+        self._metadata = None
+
+    def initialize_hook(self, module):
+        unwrapped_module = unwrap_module(module)
+        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
+        return module
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            self.state_manager.set_context("inference")
+
+        arg_name = self._metadata.hidden_states_argument_name
+        hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
+        state: MagCacheState = self.state_manager.get_state()
+        state.head_block_input = hidden_states
+
+        should_compute = True
+
+        if self.config.calibrate:
+            # Never skip during calibration
+            should_compute = True
+        else:
+            # MagCache Logic
+            current_step = state.step_index
+            if current_step >= len(self.config.mag_ratios):
+                current_scale = 1.0
+            else:
+                current_scale = self.config.mag_ratios[current_step]
+
+            retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5)
+
+            if current_step >= retention_step:
+                state.accumulated_ratio *= current_scale
+                state.accumulated_steps += 1
+                state.accumulated_err += abs(1.0 - state.accumulated_ratio)
+
+                if (
+                    state.previous_residual is not None
+                    and state.accumulated_err <= self.config.threshold
+                    and state.accumulated_steps <= self.config.max_skip_steps
+                ):
+                    should_compute = False
+                else:
+                    state.accumulated_ratio = 1.0
+                    state.accumulated_steps = 0
+                    state.accumulated_err = 0.0
+
+        state.should_compute = should_compute
+
+        if not should_compute:
+            logger.debug(f"MagCache: Skipping step {state.step_index}")
+            # Apply MagCache: Output = Input + Previous Residual
+
+            output = hidden_states
+            res = state.previous_residual
+
+            if res.device != output.device:
+                res = res.to(output.device)
+
+            # Attempt to apply residual handling shape mismatches (e.g., text+image vs image only)
+            if res.shape == output.shape:
+                output = output + res
+            elif (
+                output.ndim == 3
+                and res.ndim == 3
+                and output.shape[0] == res.shape[0]
+                and output.shape[2] == res.shape[2]
+            ):
+                # Assuming concatenation where image part is at the end (standard in Flux/SD3)
+                diff = output.shape[1] - res.shape[1]
+                if diff > 0:
+                    output = output.clone()
+                    output[:, diff:, :] = output[:, diff:, :] + res
+                else:
+                    logger.warning(
+                        f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. "
+                        "Cannot apply residual safely. Returning input without residual."
+                    )
+            else:
+                logger.warning(
+                    f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. "
+                    "Cannot apply residual safely. Returning input without residual."
+                )
+
+            if self._metadata.return_encoder_hidden_states_index is not None:
+                original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                    "encoder_hidden_states", args, kwargs
+                )
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = output
+                ret_list[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states
+                return tuple(ret_list)
+            else:
+                return output
+
+        else:
+            # Compute original forward
+            output = self.fn_ref.original_forward(*args, **kwargs)
+            return output
+
+    def reset_state(self, module):
+        self.state_manager.reset()
+        return module
+
+
+class MagCacheBlockHook(ModelHook):
+    def __init__(self, state_manager: StateManager, is_tail: bool = False, config: MagCacheConfig = None):
+        super().__init__()
+        self.state_manager = state_manager
+        self.is_tail = is_tail
+        self.config = config
+        self._metadata = None
+
+    def initialize_hook(self, module):
+        unwrapped_module = unwrap_module(module)
+        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
+        return module
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            self.state_manager.set_context("inference")
+        state: MagCacheState = self.state_manager.get_state()
+
+        if not state.should_compute:
+            arg_name = self._metadata.hidden_states_argument_name
+            hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
+            if self.is_tail:
+                # Still need to advance step index even if we skip
+                self._advance_step(state)
+
+            if self._metadata.return_encoder_hidden_states_index is not None:
+                encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                    "encoder_hidden_states", args, kwargs
+                )
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = hidden_states
+                ret_list[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states
+                return tuple(ret_list)
+
+            return hidden_states
+
+        output = self.fn_ref.original_forward(*args, **kwargs)
+
+        if self.is_tail:
+            # Calculate residual for next steps
+            if isinstance(output, tuple):
+                out_hidden = output[self._metadata.return_hidden_states_index]
+            else:
+                out_hidden = output
+
+            in_hidden = state.head_block_input
+
+            if in_hidden is None:
+                return output
+
+            # Determine residual
+            if out_hidden.shape == in_hidden.shape:
+                residual = out_hidden - in_hidden
+            elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]:
+                diff = in_hidden.shape[1] - out_hidden.shape[1]
+                if diff == 0:
+                    residual = out_hidden - in_hidden
+                else:
+                    residual = out_hidden - in_hidden  # Fallback to matching tail
+            else:
+                # Fallback for completely mismatched shapes
+                residual = out_hidden
+
+            if self.config.calibrate:
+                self._perform_calibration_step(state, residual)
+
+            state.previous_residual = residual
+            self._advance_step(state)
+
+        return output
+
+    def _perform_calibration_step(self, state: MagCacheState, current_residual: torch.Tensor):
+        if state.previous_residual is None:
+            # First step has no previous residual to compare against.
+            # log 1.0 as a neutral starting point.
+            ratio = 1.0
+        else:
+            # MagCache Calibration Formula: mean(norm(curr) / norm(prev))
+            # 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)
+
+            # Avoid division by zero
+            ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()
+
+        state.calibration_ratios.append(ratio)
+
+    def _advance_step(self, state: MagCacheState):
+        state.step_index += 1
+        if state.step_index >= self.config.num_inference_steps:
+            # End of inference loop
+            if self.config.calibrate:
+                print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):")
+                print(f"{state.calibration_ratios}\n")
+                logger.info(f"MagCache Calibration Results: {state.calibration_ratios}")
+
+            # Reset state
+            state.step_index = 0
+            state.accumulated_ratio = 1.0
+            state.accumulated_steps = 0
+            state.accumulated_err = 0.0
+            state.previous_residual = None
+            state.calibration_ratios = []
+
+
+def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None:
+    """
+    Applies MagCache to a given module (typically a Transformer).
+
+    Args:
+        module (`torch.nn.Module`):
+            The module to apply MagCache to.
+        config (`MagCacheConfig`):
+            The configuration for MagCache.
+    """
+    # Initialize registry on the root module so the Pipeline can set context.
+    HookRegistry.check_if_exists_or_initialize(module)
+
+    state_manager = StateManager(MagCacheState, (), {})
+    remaining_blocks = []
+
+    for name, submodule in module.named_children():
+        if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList):
+            continue
+        for index, block in enumerate(submodule):
+            remaining_blocks.append((f"{name}.{index}", block))
+
+    if not remaining_blocks:
+        logger.warning("MagCache: No transformer blocks found to apply hooks.")
+        return
+
+    # Handle single-block models
+    if len(remaining_blocks) == 1:
+        name, block = remaining_blocks[0]
+        logger.info(f"MagCache: Applying Head+Tail Hooks to single block '{name}'")
+        _apply_mag_cache_block_hook(block, state_manager, config, is_tail=True)
+        _apply_mag_cache_head_hook(block, state_manager, config)
+        return
+
+    head_block_name, head_block = remaining_blocks.pop(0)
+    tail_block_name, tail_block = remaining_blocks.pop(-1)
+
+    logger.info(f"MagCache: Applying Head Hook to {head_block_name}")
+    _apply_mag_cache_head_hook(head_block, state_manager, config)
+
+    for name, block in remaining_blocks:
+        _apply_mag_cache_block_hook(block, state_manager, config)
+
+    logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}")
+    _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True)
+
+
+def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None:
+    registry = HookRegistry.check_if_exists_or_initialize(block)
+
+    # Automatically remove existing hook to allow re-application (e.g. switching modes)
+    if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None:
+        registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)
+
+    hook = MagCacheHeadHook(state_manager, config)
+    registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK)
+
+
+def _apply_mag_cache_block_hook(
+    block: torch.nn.Module,
+    state_manager: StateManager,
+    config: MagCacheConfig,
+    is_tail: bool = False,
+) -> None:
+    registry = HookRegistry.check_if_exists_or_initialize(block)
+
+    # Automatically remove existing hook to allow re-application
+    if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None:
+        registry.remove_hook(_MAG_CACHE_BLOCK_HOOK)
+
+    hook = MagCacheBlockHook(state_manager, is_tail, config)
+    registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK)

src/diffusers/models/cache_utils.py

@@ -68,10 +68,12 @@ class CacheMixin:
         from ..hooks import (
             FasterCacheConfig,
             FirstBlockCacheConfig,
+            MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
             apply_faster_cache,
             apply_first_block_cache,
+            apply_mag_cache,
             apply_pyramid_attention_broadcast,
             apply_taylorseer_cache,
         )
@@ -85,6 +87,8 @@ class CacheMixin:
             apply_faster_cache(self, config)
         elif isinstance(config, FirstBlockCacheConfig):
             apply_first_block_cache(self, config)
+        elif isinstance(config, MagCacheConfig):
+            apply_mag_cache(self, config)
         elif isinstance(config, PyramidAttentionBroadcastConfig):
             apply_pyramid_attention_broadcast(self, config)
         elif isinstance(config, TaylorSeerCacheConfig):
@@ -99,11 +103,13 @@ class CacheMixin:
             FasterCacheConfig,
             FirstBlockCacheConfig,
             HookRegistry,
+            MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
         )
         from ..hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK
         from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK
+        from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK
         from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK
         from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK
 
@@ -118,6 +124,9 @@ class CacheMixin:
         elif isinstance(self._cache_config, FirstBlockCacheConfig):
             registry.remove_hook(_FBC_LEADER_BLOCK_HOOK, recurse=True)
             registry.remove_hook(_FBC_BLOCK_HOOK, recurse=True)
+        elif isinstance(self._cache_config, MagCacheConfig):
+            registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK, recurse=True)
+            registry.remove_hook(_MAG_CACHE_BLOCK_HOOK, recurse=True)
         elif isinstance(self._cache_config, PyramidAttentionBroadcastConfig):
             registry.remove_hook(_PYRAMID_ATTENTION_BROADCAST_HOOK, recurse=True)
         elif isinstance(self._cache_config, TaylorSeerCacheConfig):

src/diffusers/utils/dummy_pt_objects.py

@@ -227,6 +227,21 @@ class LayerSkipConfig(metaclass=DummyObject):
         requires_backends(cls, ["torch"])
 
 
+class MagCacheConfig(metaclass=DummyObject):
+    _backends = ["torch"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["torch"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+
 class PyramidAttentionBroadcastConfig(metaclass=DummyObject):
     _backends = ["torch"]
 
@@ -284,6 +299,10 @@ def apply_layer_skip(*args, **kwargs):
     requires_backends(apply_layer_skip, ["torch"])
 
 
+def apply_mag_cache(*args, **kwargs):
+    requires_backends(apply_mag_cache, ["torch"])
+
+
 def apply_pyramid_attention_broadcast(*args, **kwargs):
     requires_backends(apply_pyramid_attention_broadcast, ["torch"])
 

tests/hooks/test_mag_cache.py

@@ -0,0 +1,244 @@
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import torch
+
+from diffusers import MagCacheConfig, apply_mag_cache
+from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
+from diffusers.models import ModelMixin
+from diffusers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class DummyBlock(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+        # Output is double input
+        # This ensures Residual = 2*Input - Input = Input
+        return hidden_states * 2.0
+
+
+class DummyTransformer(ModelMixin):
+    def __init__(self):
+        super().__init__()
+        self.transformer_blocks = torch.nn.ModuleList([DummyBlock(), DummyBlock()])
+
+    def forward(self, hidden_states, encoder_hidden_states=None):
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states)
+        return hidden_states
+
+
+class TupleOutputBlock(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+        # Returns a tuple
+        return hidden_states * 2.0, encoder_hidden_states
+
+
+class TupleTransformer(ModelMixin):
+    def __init__(self):
+        super().__init__()
+        self.transformer_blocks = torch.nn.ModuleList([TupleOutputBlock()])
+
+    def forward(self, hidden_states, encoder_hidden_states=None):
+        for block in self.transformer_blocks:
+            # Emulate Flux-like behavior
+            output = block(hidden_states, encoder_hidden_states=encoder_hidden_states)
+            hidden_states = output[0]
+            encoder_hidden_states = output[1]
+        return hidden_states, encoder_hidden_states
+
+
+class MagCacheTests(unittest.TestCase):
+    def setUp(self):
+        # Register standard dummy block
+        TransformerBlockRegistry.register(
+            DummyBlock,
+            TransformerBlockMetadata(return_hidden_states_index=None, return_encoder_hidden_states_index=None),
+        )
+        # Register tuple block (Flux style)
+        TransformerBlockRegistry.register(
+            TupleOutputBlock,
+            TransformerBlockMetadata(return_hidden_states_index=0, return_encoder_hidden_states_index=1),
+        )
+
+    def _set_context(self, model, context_name):
+        """Helper to set context on all hooks in the model."""
+        for module in model.modules():
+            if hasattr(module, "_diffusers_hook"):
+                module._diffusers_hook._set_context(context_name)
+
+    def _get_calibration_data(self, model):
+        for module in model.modules():
+            if hasattr(module, "_diffusers_hook"):
+                hook = module._diffusers_hook.get_hook("mag_cache_block_hook")
+                if hook:
+                    return hook.state_manager.get_state().calibration_ratios
+        return []
+
+    def test_mag_cache_validation(self):
+        """Test that missing mag_ratios raises ValueError."""
+        with self.assertRaises(ValueError):
+            MagCacheConfig(num_inference_steps=10, calibrate=False)
+
+    def test_mag_cache_skipping_logic(self):
+        """
+        Tests that MagCache correctly calculates residuals and skips blocks when conditions are met.
+        """
+        model = DummyTransformer()
+
+        # Dummy ratios: [1.0, 1.0] implies 0 accumulated error if we skip
+        ratios = np.array([1.0, 1.0])
+
+        config = MagCacheConfig(
+            threshold=100.0,
+            num_inference_steps=2,
+            retention_ratio=0.0,  # Enable immediate skipping
+            max_skip_steps=5,
+            mag_ratios=ratios,
+        )
+
+        apply_mag_cache(model, config)
+        self._set_context(model, "test_context")
+
+        # Step 0: Input 10.0 -> Output 40.0 (2 blocks * 2x each)
+        # HeadInput=10. Output=40. Residual=30.
+        input_t0 = torch.tensor([[[10.0]]])
+        output_t0 = model(input_t0)
+        self.assertTrue(torch.allclose(output_t0, torch.tensor([[[40.0]]])), "Step 0 failed")
+
+        # Step 1: Input 11.0.
+        # If Skipped: Output = Input(11) + Residual(30) = 41.0
+        # If Computed: Output = 11 * 4 = 44.0
+        input_t1 = torch.tensor([[[11.0]]])
+        output_t1 = model(input_t1)
+
+        self.assertTrue(
+            torch.allclose(output_t1, torch.tensor([[[41.0]]])), f"Expected Skip (41.0), got {output_t1.item()}"
+        )
+
+    def test_mag_cache_retention(self):
+        """Test that retention_ratio prevents skipping even if error is low."""
+        model = DummyTransformer()
+        # Ratios that imply 0 error, so it *would* skip if retention allowed it
+        ratios = np.array([1.0, 1.0])
+
+        config = MagCacheConfig(
+            threshold=100.0,
+            num_inference_steps=2,
+            retention_ratio=1.0,  # Force retention for ALL steps
+            mag_ratios=ratios,
+        )
+
+        apply_mag_cache(model, config)
+        self._set_context(model, "test_context")
+
+        # Step 0
+        model(torch.tensor([[[10.0]]]))
+
+        # Step 1: Should COMPUTE (44.0) not SKIP (41.0) because of retention
+        input_t1 = torch.tensor([[[11.0]]])
+        output_t1 = model(input_t1)
+
+        self.assertTrue(
+            torch.allclose(output_t1, torch.tensor([[[44.0]]])),
+            f"Expected Compute (44.0) due to retention, got {output_t1.item()}",
+        )
+
+    def test_mag_cache_tuple_outputs(self):
+        """Test compatibility with models returning (hidden, encoder_hidden) like Flux."""
+        model = TupleTransformer()
+        ratios = np.array([1.0, 1.0])
+
+        config = MagCacheConfig(threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=ratios)
+
+        apply_mag_cache(model, config)
+        self._set_context(model, "test_context")
+
+        # Step 0: Compute. Input 10.0 -> Output 20.0 (1 block * 2x)
+        # Residual = 10.0
+        input_t0 = torch.tensor([[[10.0]]])
+        enc_t0 = torch.tensor([[[1.0]]])
+        out_0, _ = model(input_t0, encoder_hidden_states=enc_t0)
+        self.assertTrue(torch.allclose(out_0, torch.tensor([[[20.0]]])))
+
+        # Step 1: Skip. Input 11.0.
+        # Skipped Output = 11 + 10 = 21.0
+        input_t1 = torch.tensor([[[11.0]]])
+        out_1, _ = model(input_t1, encoder_hidden_states=enc_t0)
+
+        self.assertTrue(
+            torch.allclose(out_1, torch.tensor([[[21.0]]])), f"Tuple skip failed. Expected 21.0, got {out_1.item()}"
+        )
+
+    def test_mag_cache_reset(self):
+        """Test that state resets correctly after num_inference_steps."""
+        model = DummyTransformer()
+        config = MagCacheConfig(
+            threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=np.array([1.0, 1.0])
+        )
+        apply_mag_cache(model, config)
+        self._set_context(model, "test_context")
+
+        input_t = torch.ones(1, 1, 1)
+
+        model(input_t)  # Step 0
+        model(input_t)  # Step 1 (Skipped)
+
+        # Step 2 (Reset -> Step 0) -> Should Compute
+        # Input 2.0 -> Output 8.0
+        input_t2 = torch.tensor([[[2.0]]])
+        output_t2 = model(input_t2)
+
+        self.assertTrue(torch.allclose(output_t2, torch.tensor([[[8.0]]])), "State did not reset correctly")
+
+    def test_mag_cache_calibration(self):
+        """Test that calibration mode records ratios."""
+        model = DummyTransformer()
+        config = MagCacheConfig(num_inference_steps=2, calibrate=True)
+        apply_mag_cache(model, config)
+        self._set_context(model, "test_context")
+
+        # Step 0
+        # HeadInput = 10. Output = 40. Residual = 30.
+        # Ratio 0 is placeholder 1.0
+        model(torch.tensor([[[10.0]]]))
+
+        # Check intermediate state
+        ratios = self._get_calibration_data(model)
+        self.assertEqual(len(ratios), 1)
+        self.assertEqual(ratios[0], 1.0)
+
+        # Step 1
+        # HeadInput = 10. Output = 40. Residual = 30.
+        # PrevResidual = 30. CurrResidual = 30.
+        # Ratio = 30/30 = 1.0
+        model(torch.tensor([[[10.0]]]))
+
+        # Verify it computes fully (no skip)
+        # If it skipped, output would be 41.0. It should be 40.0
+        # Actually in test setup, input is same (10.0) so output 40.0.
+        # Let's ensure list is empty after reset (end of step 1)
+        ratios_after = self._get_calibration_data(model)
+        self.assertEqual(ratios_after, [])

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -2,6 +2,7 @@ import gc
 import tempfile
 from typing import Callable, Union
 
+import numpy as np
 import pytest
 import torch
 
@@ -37,6 +38,9 @@ class ModularPipelineTesterMixin:
     optional_params = frozenset(["num_inference_steps", "num_images_per_prompt", "latents", "output_type"])
     # this is modular specific: generator needs to be a intermediate input because it's mutable
     intermediate_params = frozenset(["generator"])
+    # Output type for the pipeline (e.g., "images" for image pipelines, "videos" for video pipelines)
+    # Subclasses can override this to change the expected output type
+    output_type = "images"
 
     def get_generator(self, seed=0):
         generator = torch.Generator("cpu").manual_seed(seed)
@@ -117,6 +121,30 @@ class ModularPipelineTesterMixin:
         pipeline.set_progress_bar_config(disable=None)
         return pipeline
 
+    def _convert_output_to_tensor(self, output):
+        if isinstance(output, torch.Tensor):
+            return output
+        elif isinstance(output, list):
+            # For video outputs (list of numpy arrays)
+            if len(output) > 0 and isinstance(output[0], np.ndarray):
+                return torch.from_numpy(output[0])
+            # For batched video outputs
+            return torch.stack([torch.from_numpy(item) for item in output])
+        elif isinstance(output, np.ndarray):
+            return torch.from_numpy(output)
+        else:
+            raise TypeError(f"Unsupported output type: {type(output)}")
+
+    def _get_batch_size_from_output(self, output):
+        if isinstance(output, torch.Tensor):
+            return output.shape[0]
+        elif isinstance(output, list):
+            return len(output)
+        elif isinstance(output, np.ndarray):
+            return output.shape[0]
+        else:
+            raise TypeError(f"Unsupported output type: {type(output)}")
+
     def test_pipeline_call_signature(self):
         pipe = self.get_pipeline()
         input_parameters = pipe.blocks.input_names
@@ -163,7 +191,7 @@ class ModularPipelineTesterMixin:
 
         logger.setLevel(level=diffusers.logging.WARNING)
         for batch_size, batched_input in zip(batch_sizes, batched_inputs):
-            output = pipe(**batched_input, output="images")
+            output = pipe(**batched_input, output=self.output_type)
             assert len(output) == batch_size, "Output is different from expected batch size"
 
     def test_inference_batch_single_identical(
@@ -197,12 +225,27 @@ class ModularPipelineTesterMixin:
         if "batch_size" in inputs:
             batched_inputs["batch_size"] = batch_size
 
-        output = pipe(**inputs, output="images")
-        output_batch = pipe(**batched_inputs, output="images")
+        output = pipe(**inputs, output=self.output_type)
+        output_batch = pipe(**batched_inputs, output=self.output_type)
 
-        assert output_batch.shape[0] == batch_size
+        assert self._get_batch_size_from_output(output_batch) == batch_size
 
-        max_diff = torch.abs(output_batch[0] - output[0]).max()
+        # Convert outputs to tensors for comparison
+        if isinstance(output, list) and isinstance(output_batch, list):
+            # Both are lists - compare first elements
+            if isinstance(output[0], np.ndarray):
+                output_tensor = torch.from_numpy(output[0])
+                output_batch_tensor = torch.from_numpy(output_batch[0])
+            else:
+                output_tensor = output[0]
+                output_batch_tensor = output_batch[0]
+        else:
+            output_tensor = self._convert_output_to_tensor(output)
+            output_batch_tensor = self._convert_output_to_tensor(output_batch)
+            if output_batch_tensor.shape[0] == batch_size and output_tensor.shape[0] == 1:
+                output_batch_tensor = output_batch_tensor[0:1]
+
+        max_diff = torch.abs(output_batch_tensor - output_tensor).max()
         assert max_diff < expected_max_diff, "Batch inference results different from single inference results"
 
     @require_accelerator
@@ -217,19 +260,31 @@ class ModularPipelineTesterMixin:
         # Reset generator in case it is used inside dummy inputs
         if "generator" in inputs:
             inputs["generator"] = self.get_generator(0)
-        output = pipe(**inputs, output="images")
+        output = pipe(**inputs, output=self.output_type)
 
         fp16_inputs = self.get_dummy_inputs()
         # Reset generator in case it is used inside dummy inputs
         if "generator" in fp16_inputs:
             fp16_inputs["generator"] = self.get_generator(0)
-        output_fp16 = pipe_fp16(**fp16_inputs, output="images")
+        output_fp16 = pipe_fp16(**fp16_inputs, output=self.output_type)
 
-        output = output.cpu()
-        output_fp16 = output_fp16.cpu()
+        # Convert outputs to tensors for comparison
+        output_tensor = self._convert_output_to_tensor(output).float().cpu()
+        output_fp16_tensor = self._convert_output_to_tensor(output_fp16).float().cpu()
 
-        max_diff = numpy_cosine_similarity_distance(output.flatten(), output_fp16.flatten())
-        assert max_diff < expected_max_diff, "FP16 inference is different from FP32 inference"
+        # Check for NaNs in outputs (can happen with tiny models in FP16)
+        if torch.isnan(output_tensor).any() or torch.isnan(output_fp16_tensor).any():
+            pytest.skip("FP16 inference produces NaN values - this is a known issue with tiny models")
+
+        max_diff = numpy_cosine_similarity_distance(
+            output_tensor.flatten().numpy(), output_fp16_tensor.flatten().numpy()
+        )
+
+        # Check if cosine similarity is NaN (which can happen if vectors are zero or very small)
+        if torch.isnan(torch.tensor(max_diff)):
+            pytest.skip("Cosine similarity is NaN - outputs may be too small for reliable comparison")
+
+        assert max_diff < expected_max_diff, f"FP16 inference is different from FP32 inference (max_diff: {max_diff})"
 
     @require_accelerator
     def test_to_device(self):
@@ -251,15 +306,17 @@ class ModularPipelineTesterMixin:
     def test_inference_is_not_nan_cpu(self):
         pipe = self.get_pipeline().to("cpu")
 
-        output = pipe(**self.get_dummy_inputs(), output="images")
-        assert torch.isnan(output).sum() == 0, "CPU Inference returns NaN"
+        output = pipe(**self.get_dummy_inputs(), output=self.output_type)
+        output_tensor = self._convert_output_to_tensor(output)
+        assert torch.isnan(output_tensor).sum() == 0, "CPU Inference returns NaN"
 
     @require_accelerator
     def test_inference_is_not_nan(self):
         pipe = self.get_pipeline().to(torch_device)
 
-        output = pipe(**self.get_dummy_inputs(), output="images")
-        assert torch.isnan(output).sum() == 0, "Accelerator Inference returns NaN"
+        output = pipe(**self.get_dummy_inputs(), output=self.output_type)
+        output_tensor = self._convert_output_to_tensor(output)
+        assert torch.isnan(output_tensor).sum() == 0, "Accelerator Inference returns NaN"
 
     def test_num_images_per_prompt(self):
         pipe = self.get_pipeline().to(torch_device)
@@ -278,9 +335,9 @@ class ModularPipelineTesterMixin:
                     if key in self.batch_params:
                         inputs[key] = batch_size * [inputs[key]]
 
-                images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output="images")
+                images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output=self.output_type)
 
-                assert images.shape[0] == batch_size * num_images_per_prompt
+                assert self._get_batch_size_from_output(images) == batch_size * num_images_per_prompt
 
     @require_accelerator
     def test_components_auto_cpu_offload_inference_consistent(self):
@@ -293,9 +350,10 @@ class ModularPipelineTesterMixin:
         image_slices = []
         for pipe in [base_pipe, offload_pipe]:
             inputs = self.get_dummy_inputs()
-            image = pipe(**inputs, output="images")
+            image = pipe(**inputs, output=self.output_type)
 
-            image_slices.append(image[0, -3:, -3:, -1].flatten())
+            image_tensor = self._convert_output_to_tensor(image)
+            image_slices.append(image_tensor[0, -3:, -3:, -1].flatten())
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
@@ -315,9 +373,10 @@ class ModularPipelineTesterMixin:
         image_slices = []
         for pipe in pipes:
             inputs = self.get_dummy_inputs()
-            image = pipe(**inputs, output="images")
+            image = pipe(**inputs, output=self.output_type)
 
-            image_slices.append(image[0, -3:, -3:, -1].flatten())
+            image_tensor = self._convert_output_to_tensor(image)
+            image_slices.append(image_tensor[0, -3:, -3:, -1].flatten())
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
@@ -331,13 +390,13 @@ class ModularGuiderTesterMixin:
         pipe.update_components(guider=guider)
 
         inputs = self.get_dummy_inputs()
-        out_no_cfg = pipe(**inputs, output="images")
+        out_no_cfg = pipe(**inputs, output=self.output_type)
 
         # forward pass with CFG applied
         guider = ClassifierFreeGuidance(guidance_scale=7.5)
         pipe.update_components(guider=guider)
         inputs = self.get_dummy_inputs()
-        out_cfg = pipe(**inputs, output="images")
+        out_cfg = pipe(**inputs, output=self.output_type)
 
         assert out_cfg.shape == out_no_cfg.shape
         max_diff = torch.abs(out_cfg - out_no_cfg).max()

tests/modular_pipelines/wan/test_modular_pipeline_wan.py

@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from diffusers.modular_pipelines import WanAutoBlocks, WanModularPipeline
+
+from ..test_modular_pipelines_common import ModularPipelineTesterMixin
+
+
+class TestWanModularPipelineFast(ModularPipelineTesterMixin):
+    pipeline_class = WanModularPipeline
+    pipeline_blocks_class = WanAutoBlocks
+    pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-modular-pipe"
+
+    params = frozenset(["prompt", "height", "width", "num_frames"])
+    batch_params = frozenset(["prompt"])
+    optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
+    output_type = "videos"
+
+    def get_dummy_inputs(self, seed=0):
+        generator = self.get_generator(seed)
+        inputs = {
+            "prompt": "A painting of a squirrel eating a burger",
+            "generator": generator,
+            "num_inference_steps": 2,
+            "height": 16,
+            "width": 16,
+            "num_frames": 9,
+            "max_sequence_length": 16,
+        }
+        return inputs
+
+    @pytest.mark.skip(reason="num_videos_per_prompt")
+    def test_num_images_per_prompt(self):
+        pass

tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py

@@ -0,0 +1,44 @@
+# coding=utf-8
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from diffusers.modular_pipelines import ZImageAutoBlocks, ZImageModularPipeline
+
+from ..test_modular_pipelines_common import ModularPipelineTesterMixin
+
+
+class TestZImageModularPipelineFast(ModularPipelineTesterMixin):
+    pipeline_class = ZImageModularPipeline
+    pipeline_blocks_class = ZImageAutoBlocks
+    pretrained_model_name_or_path = "hf-internal-testing/tiny-zimage-modular-pipe"
+
+    params = frozenset(["prompt", "height", "width"])
+    batch_params = frozenset(["prompt"])
+
+    def get_dummy_inputs(self, seed=0):
+        generator = self.get_generator(seed)
+        inputs = {
+            "prompt": "A painting of a squirrel eating a burger",
+            "generator": generator,
+            "num_inference_steps": 2,
+            "height": 32,
+            "width": 32,
+            "max_sequence_length": 16,
+            "output_type": "pt",
+        }
+        return inputs
+
+    def test_inference_batch_single_identical(self):
+        super().test_inference_batch_single_identical(expected_max_diff=5e-3)

tests/pipelines/flux/test_pipeline_flux.py

@@ -27,6 +27,7 @@ from ..test_pipelines_common import (
     FasterCacheTesterMixin,
     FirstBlockCacheTesterMixin,
     FluxIPAdapterTesterMixin,
+    MagCacheTesterMixin,
     PipelineTesterMixin,
     PyramidAttentionBroadcastTesterMixin,
     TaylorSeerCacheTesterMixin,
@@ -41,6 +42,7 @@ class FluxPipelineFastTests(
     FasterCacheTesterMixin,
     FirstBlockCacheTesterMixin,
     TaylorSeerCacheTesterMixin,
+    MagCacheTesterMixin,
     unittest.TestCase,
 ):
     pipeline_class = FluxPipeline

tests/pipelines/test_pipelines_common.py

@@ -35,6 +35,7 @@ from diffusers import (
 from diffusers.hooks import apply_group_offloading
 from diffusers.hooks.faster_cache import FasterCacheBlockHook, FasterCacheDenoiserHook
 from diffusers.hooks.first_block_cache import FirstBlockCacheConfig
+from diffusers.hooks.mag_cache import MagCacheConfig
 from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook
 from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig
 from diffusers.image_processor import VaeImageProcessor
@@ -2976,6 +2977,59 @@ class TaylorSeerCacheTesterMixin:
         )
 
 
+class MagCacheTesterMixin:
+    mag_cache_config = MagCacheConfig(
+        threshold=0.06,
+        max_skip_steps=3,
+        retention_ratio=0.2,
+        num_inference_steps=50,
+        mag_ratios=torch.ones(50),
+    )
+
+    def test_mag_cache_inference(self, expected_atol: float = 0.1):
+        device = "cpu"
+
+        def create_pipe():
+            torch.manual_seed(0)
+            num_layers = 2
+            components = self.get_dummy_components(num_layers=num_layers)
+            pipe = self.pipeline_class(**components)
+            pipe = pipe.to(device)
+            pipe.set_progress_bar_config(disable=None)
+            return pipe
+
+        def run_forward(pipe):
+            torch.manual_seed(0)
+            inputs = self.get_dummy_inputs(device)
+            # Match the config steps
+            inputs["num_inference_steps"] = 50
+            return pipe(**inputs)[0]
+
+        # 1. Run inference without MagCache (Baseline)
+        pipe = create_pipe()
+        output = run_forward(pipe).flatten()
+        original_image_slice = np.concatenate((output[:8], output[-8:]))
+
+        # 2. Run inference with MagCache ENABLED
+        pipe = create_pipe()
+        pipe.transformer.enable_cache(self.mag_cache_config)
+        output = run_forward(pipe).flatten()
+        image_slice_enabled = np.concatenate((output[:8], output[-8:]))
+
+        # 3. Run inference with MagCache DISABLED
+        pipe.transformer.disable_cache()
+        output = run_forward(pipe).flatten()
+        image_slice_disabled = np.concatenate((output[:8], output[-8:]))
+
+        assert np.allclose(original_image_slice, image_slice_enabled, atol=expected_atol), (
+            "MagCache outputs should not differ too much from baseline."
+        )
+
+        assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-4), (
+            "Outputs after disabling cache should match original inference exactly."
+        )
+
+
 # Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used.
 # This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a
 # reference image.