fix(mag_cache): correct per-transformer step count for dual-transformer pipelines

[lcheng321]

Fixes #14025

Root cause

MagCacheConfig was applied to both transformer and transformer_2 with num_inference_steps set to the full pipeline step count. In dual transformer pipelines such as Wan 2.2, transformer only runs the high noise stage and transformer_2 only runs the low noise stage. Each of them executes a subset of the total steps, never the full count.

This causes three downstream problems:

• _advance_step: the check state.step_index >= config.num_inference_steps never fires for either transformer, since neither ever reaches the full step count. State never resets and calibration never completes.
• Retention window: int(retention_ratio * num_inference_steps) is computed against the full step count instead of the actual per transformer step count, so it protects more steps from skipping than intended.
• mag_ratios indexing: since mag_ratios is sized to the full step count, step_index no longer maps to the correct entry, so each transformer reads ratios that do not correspond to its own denoising timestep.

First attempt and why it changed

The initial fix added update_mag_cache_num_steps(module, actual_steps), which mutated config.num_inference_steps and config.mag_ratios in place on the attached MagCacheConfig.

Review feedback identified two problems with that approach:

1. If the same MagCacheConfig object is passed to both transformer and transformer_2, calling this function for one transformer corrupts the config for the other, since they share the same object.
2. The pipeline had to import a hook internal function directly, coupling pipeline code to hook implementation details.

Fix

MagCacheConfig is now never mutated after construction. Instead:

• Pipelines set a plain attribute on each transformer, transformer._mag_cache_expected_steps = n_steps, no hook internals imported.
• The hooks hold a reference to the root transformer module (passed in at apply_mag_cache time) and read _mag_cache_expected_steps off it live, at forward and reset time, falling back to config.num_inference_steps if the attribute was never set.
• mag_ratios are computed on demand from _original_mag_ratios at the effective step count, never written back to config.mag_ratios.

This means a shared config object across both transformers can no longer be corrupted, since nothing ever writes to it after construction.

Files changed:

File	Change
src/diffusers/hooks/mag_cache.py	Add _get_effective_num_steps and _get_effective_mag_ratios, both pure functions that read from a root module attribute and a config, no mutation. Hooks now take a root_module reference. Remove update_mag_cache_num_steps.
src/diffusers/hooks/__init__.py	Remove the export of the now deleted update_mag_cache_num_steps.
src/diffusers/pipelines/wan/pipeline_wan.py	Set transformer._mag_cache_expected_steps / transformer_2._mag_cache_expected_steps instead of calling into hook internals.
src/diffusers/pipelines/wan/pipeline_wan_i2v.py	Same as above.
src/diffusers/pipelines/wan/pipeline_wan_vace.py	Same as above.

All pipeline changes are no-ops when MagCache is not applied.

Evidence

Run locally, CPU only, no GPU or model weights required. Uses tiny WanTransformerBlock instances (dim=32) just to satisfy the hook registry.

1. The bug this PR fixes: shared config corruption under asymmetric split

Before this PR, the mutating approach corrupted a config object shared across two transformers. Reproduced directly against a plain MagCacheConfig, no pipeline needed:

from diffusers.hooks.mag_cache import MagCacheConfig, nearest_interp

shared = MagCacheConfig(mag_ratios=[1.0, 0.9, 0.8, 0.7], num_inference_steps=10)

# old mutating approach, simulated inline since update_mag_cache_num_steps is removed
shared.num_inference_steps = 7
shared.mag_ratios = nearest_interp(shared._original_mag_ratios, 7)
print("after setting for transformer:", shared.num_inference_steps)

shared.num_inference_steps = 3
shared.mag_ratios = nearest_interp(shared._original_mag_ratios, 3)
print("after setting for transformer_2:", shared.num_inference_steps)

print("what transformer now reads:", shared.num_inference_steps)

Output:
after setting for transformer: 7
after setting for transformer_2: 3
what transformer now reads: 3

transformer was supposed to run 7 steps but the shared config now reports 3, because transformer_2's update overwrote it. This is the exact corruption the reviewer flagged.

2. After the fix, the same shared config is never touched

from diffusers.hooks.mag_cache import MagCacheConfig, apply_mag_cache, _get_effective_num_steps
import torch.nn as nn
from diffusers.models.transformers.transformer_wan import WanTransformerBlock

def make_tiny():
    class T(nn.Module):
        def __init__(self):
            super().__init__()
            self.transformer_blocks = nn.ModuleList([
                WanTransformerBlock(dim=32, ffn_dim=64, num_heads=2,
                    qk_norm="rms_norm_across_heads", cross_attn_norm=True,
                    eps=1e-6, added_kv_proj_dim=None)
                for _ in range(3)
            ])
    return T()

shared = MagCacheConfig(mag_ratios=[1.0, 0.9, 0.8, 0.7], num_inference_steps=10)
t1, t2 = make_tiny(), make_tiny()
apply_mag_cache(t1, shared)
apply_mag_cache(t2, shared)

t1._mag_cache_expected_steps = 7
t2._mag_cache_expected_steps = 3

print("t1 effective steps:", _get_effective_num_steps(shared, t1))
print("t2 effective steps:", _get_effective_num_steps(shared, t2))
print("shared.num_inference_steps, unchanged:", shared.num_inference_steps)

Output:
t1 effective steps: 7
t2 effective steps: 3
shared.num_inference_steps, unchanged: 10

Each transformer sees its own correct step count and the shared config object is never written to.

3. Full unit test suite

$ python -m pytest tests/hooks/test_mag_cache_wan22_steps.py -v
19 passed in 0.30s

Covers, beyond the case above:

• Retention window computed against the effective per transformer step count, not the global one.
• mag_ratios re-interpolated per transformer without mutating config.mag_ratios.
• Asymmetric splits (3,2), (7,3), (1,9) on a single shared config object.
• Single transformer pipelines (Flux, SDXL, Wan 2.1) fall back to config.num_inference_steps unchanged when _mag_cache_expected_steps was never set, confirming no behavior change there.
• Hooks correctly hold a reference to the root module they were applied to.

4. Double interpolation is still avoided

Unchanged from the original fix, still verified by the test suite: mag_ratios are always re-computed from _original_mag_ratios, never from an already interpolated mag_ratios, so calling this multiple times at different step counts does not compound rounding error.

Scope note

This PR only fixes the step accounting and config mutation issues described above. The separately reported quality degradation on heavily step distilled models (50 step calibration collapsed into a 4 step schedule) is a distinct issue tracked in #14025 and not addressed here.

[lcheng321]

@sayakpaul I saw your name in the git blame history for mag_cache.py and noticed you are also a maintainer on this file, would you have time to review this PR. It fixes the per-transformer step accounting bug reported in #14025, where num_inference_steps was set to the full pipeline count instead of the actual per-transformer execution count. I verified the fix locally with a CPU only repro and 20 passing unit tests covering asymmetric splits and edge cases. Happy to answer any questions about the approach.

[sayakpaul]

Cc: @AlanPonnachan

[lcheng321]

@AlanPonnachan any chance you could take a look when free?

[AlanPonnachan]

Sorry for the late reply. I'll review it today

[AlanPonnachan]

It's recommended to keep MagCacheConfig fixed during pipeline execution rather than modifying it dynamically. If a user passes the same MagCacheConfig object to both transformers, or reuses the pipeline for a new request with a different step count, mutating config.num_inference_steps and config.mag_ratios in place will corrupt the configuration for subsequent runs?

[lcheng321]

Fixed in 6d19597.

Config was stored by reference and read live every step. Same config on both transformers under an asymmetric split meant the second update_mag_cache_num_steps call clobbered the first.

Now config is never mutated. Each transformer reads its own step count from a plain attribute on the root module (_mag_cache_expected_steps), falling back to config.num_inference_steps when unset. update_mag_cache_num_steps is gone, nothing left to mutate.

def test_shared_config_object_untouched_by_apply(self):
    shared = _make_config(self.RATIOS, self.TOTAL)
    t1, t2 = _make_tiny_transformer(), _make_tiny_transformer()
    apply_mag_cache(t1, shared)
    apply_mag_cache(t2, shared)
    assert shared.num_inference_steps == self.TOTAL
    assert len(shared.mag_ratios) == self.TOTAL

Also covered asymmetric splits (3,2), (7,3), (1,9) on a shared config. 19/19 tests pass.

[AlanPonnachan]

we should avoid tightly coupling the core pipelines to specific hook implementations (like importing update_mag_cache_num_steps directly). I think a better solution would be
can we just have the pipeline generically broadcast the expected step counts to the modules? For example:

n_steps_t1 = sum(1 for t in timesteps if t >= boundary_timestep)
if self.transformer is not None:
    self.transformer._expected_steps = n_steps_t1
if self.transformer_2 is not None:
    self.transformer_2._expected_steps = len(timesteps) - n_steps_t1

Then, MagCacheBlockHook can simply read getattr(module, "_expected_steps", default) during initialization. This keeps the pipeline cleanly decoupled from hook internals."

[lcheng321]

Fixed in 6d19597, close to your suggestion.

Pipeline just sets a plain attribute now:

if self.transformer is not None:
    self.transformer._mag_cache_expected_steps = n_steps_t1
if self.transformer_2 is not None:
    self.transformer_2._mag_cache_expected_steps = n_steps_t2

One tweak: reading it during initialize_hook doesn't work, that runs before the pipeline computes boundary_timestep, so the attribute isn't set yet. Hooks also live on individual blocks, not the root transformer. So hooks hold a root_module reference and read the attribute lazily at forward/reset time instead.