fix(mag_cache): stop mutating MagCacheConfig; broadcast per-transformer step count via plain attribute

MagCacheConfig is no longer mutated at runtime. Instead, pipelines broadcast
each transformer's actual step count via a plain _mag_cache_expected_steps
attribute on the root module, which MagCacheHeadHook/MagCacheBlockHook read
live through _get_effective_num_steps/_get_effective_mag_ratios (falling back
to config.num_inference_steps when unset, so single-transformer pipelines are
unaffected).

Addresses both review comments on #14027:
- Sharing one MagCacheConfig across transformer/transformer_2 under an
  asymmetric step split can no longer corrupt it, since it's never mutated.
- Pipelines no longer import hook internals (update_mag_cache_num_steps is
  removed); they only set a plain attribute, keeping pipeline code decoupled
  from hook implementation details.

Author: lcheng321

src/diffusers/hooks/__init__.py

@@ -23,7 +23,7 @@
     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, update_mag_cache_num_steps
+    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/mag_cache.py

@@ -169,12 +169,25 @@ def reset(self):
         self.calibration_ratios = []
 
 
+def _get_effective_num_steps(config: MagCacheConfig, root_module: torch.nn.Module) -> int:
+    expected = getattr(root_module, "_mag_cache_expected_steps", None)
+    return expected if expected is not None else config.num_inference_steps
+
+
+def _get_effective_mag_ratios(config: MagCacheConfig, effective_num_steps: int) -> torch.Tensor:
+    if effective_num_steps == config.num_inference_steps:
+        return config.mag_ratios
+    original = getattr(config, "_original_mag_ratios", config.mag_ratios)
+    return nearest_interp(original, effective_num_steps)
+
+
 class MagCacheHeadHook(ModelHook):
     _is_stateful = True
 
-    def __init__(self, state_manager: StateManager, config: MagCacheConfig):
+    def __init__(self, state_manager: StateManager, config: MagCacheConfig, root_module: torch.nn.Module):
         self.state_manager = state_manager
         self.config = config
+        self.root_module = root_module
         self._metadata = None
 
     def initialize_hook(self, module):
@@ -200,13 +213,16 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             should_compute = True
         else:
             # MagCache Logic
+            effective_num_steps = _get_effective_num_steps(self.config, self.root_module)
+            effective_mag_ratios = _get_effective_mag_ratios(self.config, effective_num_steps)
+
             current_step = state.step_index
-            if current_step >= len(self.config.mag_ratios):
+            if current_step >= len(effective_mag_ratios):
                 current_scale = 1.0
             else:
-                current_scale = self.config.mag_ratios[current_step]
+                current_scale = effective_mag_ratios[current_step]
 
-            retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5)
+            retention_step = int(self.config.retention_ratio * effective_num_steps + 0.5)
 
             if current_step >= retention_step:
                 state.accumulated_ratio *= current_scale
@@ -286,11 +302,18 @@ def reset_state(self, module):
 
 
 class MagCacheBlockHook(ModelHook):
-    def __init__(self, state_manager: StateManager, is_tail: bool = False, config: MagCacheConfig = None):
+    def __init__(
+        self,
+        state_manager: StateManager,
+        is_tail: bool = False,
+        config: MagCacheConfig = None,
+        root_module: torch.nn.Module = None,
+    ):
         super().__init__()
         self.state_manager = state_manager
         self.is_tail = is_tail
         self.config = config
+        self.root_module = root_module
         self._metadata = None
 
     def initialize_hook(self, module):
@@ -379,7 +402,8 @@ def _perform_calibration_step(self, state: MagCacheState, current_residual: torc
 
     def _advance_step(self, state: MagCacheState):
         state.step_index += 1
-        if state.step_index >= self.config.num_inference_steps:
+        effective_num_steps = _get_effective_num_steps(self.config, self.root_module)
+        if state.step_index >= effective_num_steps:
             # End of inference loop
             if self.config.calibrate:
                 print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):")
@@ -426,38 +450,44 @@ def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None:
     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)
+        _apply_mag_cache_block_hook(block, state_manager, config, module, is_tail=True)
+        _apply_mag_cache_head_hook(block, state_manager, config, module)
         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)
+    _apply_mag_cache_head_hook(head_block, state_manager, config, module)
 
     for name, block in remaining_blocks:
-        _apply_mag_cache_block_hook(block, state_manager, config)
+        _apply_mag_cache_block_hook(block, state_manager, config, module)
 
     logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}")
-    _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True)
+    _apply_mag_cache_block_hook(tail_block, state_manager, config, module, is_tail=True)
 
 
-def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None:
+def _apply_mag_cache_head_hook(
+    block: torch.nn.Module,
+    state_manager: StateManager,
+    config: MagCacheConfig,
+    root_module: torch.nn.Module,
+) -> 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)
+    hook = MagCacheHeadHook(state_manager, config, root_module)
     registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK)
 
 
 def _apply_mag_cache_block_hook(
     block: torch.nn.Module,
     state_manager: StateManager,
     config: MagCacheConfig,
+    root_module: torch.nn.Module,
     is_tail: bool = False,
 ) -> None:
     registry = HookRegistry.check_if_exists_or_initialize(block)
@@ -466,17 +496,6 @@ def _apply_mag_cache_block_hook(
     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)
+    hook = MagCacheBlockHook(state_manager, is_tail, config, root_module)
     registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK)
 
-
-def update_mag_cache_num_steps(module: torch.nn.Module, num_steps: int) -> None:
-    config: MagCacheConfig = getattr(module, "_mag_cache_config", None)
-    if config is None:
-        return
-    original_ratios = getattr(config, "_original_mag_ratios", config.mag_ratios)
-    config.num_inference_steps = num_steps
-    if original_ratios is not None:
-        config.mag_ratios = nearest_interp(original_ratios, num_steps)
-
-

src/diffusers/pipelines/wan/pipeline_wan.py

@@ -587,13 +587,12 @@ def __call__(
             boundary_timestep = None
 
         if boundary_timestep is not None:
-            from ...hooks.mag_cache import update_mag_cache_num_steps
             n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep)
             n_steps_t2 = len(timesteps) - n_steps_t1
             if self.transformer is not None:
-                update_mag_cache_num_steps(self.transformer, n_steps_t1)
+                self.transformer._mag_cache_expected_steps = n_steps_t1
             if self.transformer_2 is not None:
-                update_mag_cache_num_steps(self.transformer_2, n_steps_t2)
+                self.transformer_2._mag_cache_expected_steps = n_steps_t2
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):

src/diffusers/pipelines/wan/pipeline_wan_i2v.py

@@ -742,13 +742,12 @@ def __call__(
             boundary_timestep = None
 
         if boundary_timestep is not None:
-            from ...hooks.mag_cache import update_mag_cache_num_steps
             n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep)
             n_steps_t2 = len(timesteps) - n_steps_t1
             if self.transformer is not None:
-                update_mag_cache_num_steps(self.transformer, n_steps_t1)
+                self.transformer._mag_cache_expected_steps = n_steps_t1
             if self.transformer_2 is not None:
-                update_mag_cache_num_steps(self.transformer_2, n_steps_t2)
+                self.transformer_2._mag_cache_expected_steps = n_steps_t2
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):

src/diffusers/pipelines/wan/pipeline_wan_vace.py

@@ -956,13 +956,12 @@ def __call__(
             boundary_timestep = None
 
         if boundary_timestep is not None:
-            from ...hooks.mag_cache import update_mag_cache_num_steps
             n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep)
             n_steps_t2 = len(timesteps) - n_steps_t1
             if self.transformer is not None:
-                update_mag_cache_num_steps(self.transformer, n_steps_t1)
+                self.transformer._mag_cache_expected_steps = n_steps_t1
             if self.transformer_2 is not None:
-                update_mag_cache_num_steps(self.transformer_2, n_steps_t2)
+                self.transformer_2._mag_cache_expected_steps = n_steps_t2
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):