Fix syntax error in quantization configuration

docs/source/en/optimization/cache.md

@@ -111,57 +111,3 @@ 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.

src/diffusers/__init__.py

@@ -168,14 +168,12 @@ 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",
         ]
@@ -934,14 +932,12 @@ 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,7 +23,6 @@ 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,13 +23,7 @@ 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",
-    "visual_transformer_blocks",
-)
+_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks", "layers")
 _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",)
 _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "layers")
 

src/diffusers/hooks/_helpers.py

@@ -26,7 +26,6 @@ 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
@@ -170,7 +169,7 @@ def _register_attention_processors_metadata():
 
 
 def _register_transformer_blocks_metadata():
-    from ..models.attention import BasicTransformerBlock, JointTransformerBlock
+    from ..models.attention import BasicTransformerBlock
     from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
     from ..models.transformers.transformer_bria import BriaTransformerBlock
     from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
@@ -185,7 +184,6 @@ 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
@@ -333,24 +331,6 @@ 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

@@ -1,468 +0,0 @@
-# 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,12 +68,10 @@ 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,
         )
@@ -87,8 +85,6 @@ 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):
@@ -103,13 +99,11 @@ 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
 
@@ -124,9 +118,6 @@ 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/schedulers/scheduling_cosine_dpmsolver_multistep.py

@@ -545,9 +545,7 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.

src/diffusers/schedulers/scheduling_deis_multistep.py

@@ -867,9 +867,7 @@ class DEISMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.

src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

@@ -245,26 +245,13 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
     ):
         if self.config.use_beta_sigmas and not is_scipy_available():
             raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
-        if (
-            sum(
-                [
-                    self.config.use_beta_sigmas,
-                    self.config.use_exponential_sigmas,
-                    self.config.use_karras_sigmas,
-                ]
-            )
-            > 1
-        ):
+        if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
             raise ValueError(
                 "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
             )
         if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
             deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
-            deprecate(
-                "algorithm_types dpmsolver and sde-dpmsolver",
-                "1.0.0",
-                deprecation_message,
-            )
+            deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message)
 
         if trained_betas is not None:
             self.betas = torch.tensor(trained_betas, dtype=torch.float32)
@@ -272,15 +259,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
             self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
-            self.betas = (
-                torch.linspace(
-                    beta_start**0.5,
-                    beta_end**0.5,
-                    num_train_timesteps,
-                    dtype=torch.float32,
-                )
-                ** 2
-            )
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
         elif beta_schedule == "squaredcos_cap_v2":
             # Glide cosine schedule
             self.betas = betas_for_alpha_bar(num_train_timesteps)
@@ -308,12 +287,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         self.init_noise_sigma = 1.0
 
         # settings for DPM-Solver
-        if algorithm_type not in [
-            "dpmsolver",
-            "dpmsolver++",
-            "sde-dpmsolver",
-            "sde-dpmsolver++",
-        ]:
+        if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
             if algorithm_type == "deis":
                 self.register_to_config(algorithm_type="dpmsolver++")
             else:
@@ -750,7 +724,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output: torch.Tensor,
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -764,7 +738,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output (`torch.Tensor`):
                 The direct output from the learned diffusion model.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
 
         Returns:
@@ -848,7 +822,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output: torch.Tensor,
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         noise: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
@@ -858,10 +832,8 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output (`torch.Tensor`):
                 The direct output from the learned diffusion model.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
-            noise (`torch.Tensor`, *optional*):
-                The noise tensor.
 
         Returns:
             `torch.Tensor`:
@@ -888,10 +860,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
                 "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
             )
 
-        sigma_t, sigma_s = (
-            self.sigmas[self.step_index + 1],
-            self.sigmas[self.step_index],
-        )
+        sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
         lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
@@ -922,7 +891,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output_list: List[torch.Tensor],
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         noise: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
@@ -932,7 +901,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output_list (`List[torch.Tensor]`):
                 The direct outputs from learned diffusion model at current and latter timesteps.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
 
         Returns:
@@ -1045,7 +1014,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output_list: List[torch.Tensor],
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         noise: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
@@ -1055,10 +1024,8 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output_list (`List[torch.Tensor]`):
                 The direct outputs from learned diffusion model at current and latter timesteps.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by diffusion process.
-            noise (`torch.Tensor`, *optional*):
-                The noise tensor.
 
         Returns:
             `torch.Tensor`:
@@ -1139,9 +1106,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         return x_t
 
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.
@@ -1251,10 +1216,7 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         sample = sample.to(torch.float32)
         if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None:
             noise = randn_tensor(
-                model_output.shape,
-                generator=generator,
-                device=model_output.device,
-                dtype=torch.float32,
+                model_output.shape, generator=generator, device=model_output.device, dtype=torch.float32
             )
         elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
             noise = variance_noise.to(device=model_output.device, dtype=torch.float32)

src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py

@@ -141,10 +141,6 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         use_beta_sigmas (`bool`, *optional*, defaults to `False`):
             Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta
             Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information.
-        use_flow_sigmas (`bool`, *optional*, defaults to `False`):
-            Whether to use flow sigmas for step sizes in the noise schedule during the sampling process.
-        flow_shift (`float`, *optional*, defaults to 1.0):
-            The flow shift factor. Valid only when `use_flow_sigmas=True`.
         lambda_min_clipped (`float`, defaults to `-inf`):
             Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
             cosine (`squaredcos_cap_v2`) noise schedule.
@@ -167,15 +163,15 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         num_train_timesteps: int = 1000,
         beta_start: float = 0.0001,
         beta_end: float = 0.02,
-        beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
+        beta_schedule: str = "linear",
         trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
         solver_order: int = 2,
-        prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon",
+        prediction_type: str = "epsilon",
         thresholding: bool = False,
         dynamic_thresholding_ratio: float = 0.995,
         sample_max_value: float = 1.0,
-        algorithm_type: Literal["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"] = "dpmsolver++",
-        solver_type: Literal["midpoint", "heun"] = "midpoint",
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
         lower_order_final: bool = True,
         euler_at_final: bool = False,
         use_karras_sigmas: Optional[bool] = False,
@@ -184,32 +180,19 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         use_flow_sigmas: Optional[bool] = False,
         flow_shift: Optional[float] = 1.0,
         lambda_min_clipped: float = -float("inf"),
-        variance_type: Optional[Literal["learned", "learned_range"]] = None,
-        timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
+        variance_type: Optional[str] = None,
+        timestep_spacing: str = "linspace",
         steps_offset: int = 0,
     ):
         if self.config.use_beta_sigmas and not is_scipy_available():
             raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
-        if (
-            sum(
-                [
-                    self.config.use_beta_sigmas,
-                    self.config.use_exponential_sigmas,
-                    self.config.use_karras_sigmas,
-                ]
-            )
-            > 1
-        ):
+        if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
             raise ValueError(
                 "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
             )
         if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
             deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
-            deprecate(
-                "algorithm_types dpmsolver and sde-dpmsolver",
-                "1.0.0",
-                deprecation_message,
-            )
+            deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message)
 
         if trained_betas is not None:
             self.betas = torch.tensor(trained_betas, dtype=torch.float32)
@@ -217,15 +200,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
             self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
-            self.betas = (
-                torch.linspace(
-                    beta_start**0.5,
-                    beta_end**0.5,
-                    num_train_timesteps,
-                    dtype=torch.float32,
-                )
-                ** 2
-            )
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
         elif beta_schedule == "squaredcos_cap_v2":
             # Glide cosine schedule
             self.betas = betas_for_alpha_bar(num_train_timesteps)
@@ -244,12 +219,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         self.init_noise_sigma = 1.0
 
         # settings for DPM-Solver
-        if algorithm_type not in [
-            "dpmsolver",
-            "dpmsolver++",
-            "sde-dpmsolver",
-            "sde-dpmsolver++",
-        ]:
+        if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
             if algorithm_type == "deis":
                 self.register_to_config(algorithm_type="dpmsolver++")
             else:
@@ -280,11 +250,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         """
         return self._step_index
 
-    def set_timesteps(
-        self,
-        num_inference_steps: Optional[int] = None,
-        device: Optional[Union[str, torch.device]] = None,
-    ):
+    def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -416,7 +382,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         return sample
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t
-    def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
+    def _sigma_to_t(self, sigma, log_sigmas):
         """
         Convert sigma values to corresponding timestep values through interpolation.
 
@@ -453,7 +419,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         return t
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
-    def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def _sigma_to_alpha_sigma_t(self, sigma):
         """
         Convert sigma values to alpha_t and sigma_t values.
 
@@ -475,7 +441,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         return alpha_t, sigma_t
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
-    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
         """
         Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
         Models](https://huggingface.co/papers/2206.00364).
@@ -601,7 +567,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output: torch.Tensor,
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -615,7 +581,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output (`torch.Tensor`):
                 The direct output from the learned diffusion model.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
 
         Returns:
@@ -700,7 +666,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output: torch.Tensor,
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         noise: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
@@ -710,10 +676,8 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output (`torch.Tensor`):
                 The direct output from the learned diffusion model.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
-            noise (`torch.Tensor`, *optional*):
-                The noise tensor.
 
         Returns:
             `torch.Tensor`:
@@ -740,10 +704,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
                 "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
             )
 
-        sigma_t, sigma_s = (
-            self.sigmas[self.step_index + 1],
-            self.sigmas[self.step_index],
-        )
+        sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
         lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
@@ -775,7 +736,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output_list: List[torch.Tensor],
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         noise: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
@@ -785,7 +746,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output_list (`List[torch.Tensor]`):
                 The direct outputs from learned diffusion model at current and latter timesteps.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
 
         Returns:
@@ -899,7 +860,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         self,
         model_output_list: List[torch.Tensor],
         *args,
-        sample: Optional[torch.Tensor] = None,
+        sample: torch.Tensor = None,
         noise: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
@@ -909,10 +870,8 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         Args:
             model_output_list (`List[torch.Tensor]`):
                 The direct outputs from learned diffusion model at current and latter timesteps.
-            sample (`torch.Tensor`, *optional*):
+            sample (`torch.Tensor`):
                 A current instance of a sample created by diffusion process.
-            noise (`torch.Tensor`, *optional*):
-                The noise tensor.
 
         Returns:
             `torch.Tensor`:
@@ -992,7 +951,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
             )
         return x_t
 
-    def _init_step_index(self, timestep: Union[int, torch.Tensor]):
+    def _init_step_index(self, timestep):
         if isinstance(timestep, torch.Tensor):
             timestep = timestep.to(self.timesteps.device)
 
@@ -1016,7 +975,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         model_output: torch.Tensor,
         timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
-        generator: Optional[torch.Generator] = None,
+        generator=None,
         variance_noise: Optional[torch.Tensor] = None,
         return_dict: bool = True,
     ) -> Union[SchedulerOutput, Tuple]:
@@ -1068,10 +1027,7 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
 
         if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None:
             noise = randn_tensor(
-                model_output.shape,
-                generator=generator,
-                device=model_output.device,
-                dtype=model_output.dtype,
+                model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype
             )
         elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
             noise = variance_noise
@@ -1118,21 +1074,6 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         noise: torch.Tensor,
         timesteps: torch.IntTensor,
     ) -> torch.Tensor:
-        """
-        Add noise to the clean `original_samples` using the scheduler's equivalent function.
-
-        Args:
-            original_samples (`torch.Tensor`):
-                The original samples to add noise to.
-            noise (`torch.Tensor`):
-                The noise tensor.
-            timesteps (`torch.IntTensor`):
-                The timesteps at which to add noise.
-
-        Returns:
-            `torch.Tensor`:
-                The noisy samples.
-        """
         # Make sure sigmas and timesteps have the same device and dtype as original_samples
         sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
         if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
@@ -1162,5 +1103,5 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin):
         noisy_samples = alpha_t * original_samples + sigma_t * noise
         return noisy_samples
 
-    def __len__(self) -> int:
+    def __len__(self):
         return self.config.num_train_timesteps

src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py

@@ -1120,9 +1120,7 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.

src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py

@@ -662,9 +662,7 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.

src/diffusers/schedulers/scheduling_sasolver.py

@@ -1122,9 +1122,7 @@ class SASolverScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.

src/diffusers/schedulers/scheduling_unipc_multistep.py

@@ -1083,9 +1083,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
     def index_for_timestep(
-        self,
-        timestep: Union[int, torch.Tensor],
-        schedule_timesteps: Optional[torch.Tensor] = None,
+        self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
     ) -> int:
         """
         Find the index for a given timestep in the schedule.

src/diffusers/utils/dummy_pt_objects.py

@@ -227,21 +227,6 @@ 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"]
 
@@ -299,10 +284,6 @@ 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

@@ -1,244 +0,0 @@
-# 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/pipelines/flux/test_pipeline_flux.py

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

tests/pipelines/test_pipelines_common.py

@@ -35,7 +35,6 @@ 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
@@ -2977,59 +2976,6 @@ 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.