[hooks] Support Qwen Image 2.1 in MagCache

[rootonchair]

What does this PR do?

1. QwenImage21TransformerBlock was not registered in TransformerBlockRegistry, so enable_cache(MagCacheConfig(...)) raised ValueError: Model class ... not registered. It is now registered as a single-stream block (return_hidden_states_index=0, no encoder hidden states).
2. With the pipeline's KV cache (use_kv_cache=True, the default), the blocks see prefix + target tokens on the prefill step and only the target tokens on later steps. If step 1 was skipped, the step-0 residual was longer than the input, so the head hook returned the input without any residual (the whole transformer was silently skipped). Calibration crashed on the same change in length. Both now use the trailing target tokens of the longer residual, which is the same slice the model's own kv_cache_mode="cached" path takes (joint_hidden_states[:, prefix_len:]). Flux2 and other models are unaffected, since their sequence length does not change between steps.

Usage

Calibrate once with the same num_inference_steps you will use for inference, then pass the printed ratios:

import torch
from diffusers import MagCacheConfig, QwenImage21Pipeline

pipe = QwenImage21Pipeline.from_pretrained(
    "Qwen/Qwen-Image-2.1",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# 1. Calibration Step
# Run full inference to measure model behavior.
calib_config = MagCacheConfig(calibrate=True, num_inference_steps=40)
pipe.transformer.enable_cache(calib_config)

# Run a prompt to trigger calibration
pipe("A capybara wearing a wizard hat, readingainting", num_inference_steps=40)
# Logs will print something like: "MagCache Calibration Results: [1.0, 1.0637, 0.9952, ..., 1.0519, 1.1294]"

# Disable the calibration hooks before enabling a new config
pipe.transformer.disable_cache()

# 2. Inference Step
# Apply the specific ratios obtained from calibration for optimized speed.
# Note: You can also skip calibration and use these precomputed ratios
# (calibrated on Qwen-Image-2.1, bf16, 1024x1024, 40 steps).
QWENIMAGE21_MAG_RATIOS = [
    1.0, 1.0637, 0.9952, 0.9936, 0.9907, 0.995, 0.9871, 0.9901, 0.992, 0.9923,
    0.9934, 0.9952, 0.9886, 0.9944, 0.994, 0.9.9947,
    0.9907, 0.9973, 0.9919, 0.9949, 0.9973, 0.993, 0.997, 0.997, 0.9963, 1.0009,
    0.9986, 0.9999, 1.0017, 1.0079, 1.0072, 1. 1.1294,
]
mag_config = MagCacheConfig(
    mag_ratios=QWENIMAGE21_MAG_RATIOS,
    num_inference_steps=40
)

pipe.transformer.enable_cache(mag_config)

image = pipe("A capybara wearing a wizard hat, reading a book by candlelight, oil painting",
num_inference_steps=40).images[0]

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul @DN6

[rootonchair]

Self-review notes

Reviewed against .ai/references/review-rules.md (plus code_style.md, testing.md, pitfalls.md) using the self-review skill. Diff: 4 files, +30/-0.

Blocking issues

None.

• QwenImage21TransformerBlock is registered as a single-stream block (return_hidden_states_index=0, no encoder hidden states), which matches its forward.
• The mag_cache.py slices take the trailing target tokens. That is the same layout the model's own kv_cache_mode="cached" path relies on (joint_hidden_states[:, prefix_len:]).
• Tests follow the mixin-only pattern from [hooks] Support Flux2 in MagCache #14756. The 7 MagCache tests pass.

Non-blocking issues

1. No unit test covers the mag_cache.py change (hooks/mag_cache.py, the diff < 0 branch in the head hook and the trailing-token alignment in _perform_calibration_step). The MagCache mixin tests still pass with the change reverted: the model mixin uses no KV cache, and the pipeline mixin (50 steps, retention_ratio=0.2) never skips step 1 and never calibrates. In [hooks] Support Flux2 in MagCache #14756, by contrast, the mixins do exercise its mag_cache.py change on every step.
   The change was instead checked on the real Qwen/Qwen-Image-2.1 checkpoint: calibration produced all 40 ratios, 23/40 steps were skipped, no "Dimension mismatch" warnings were logged, and the result was 37.8 dB PSNR against no cache. Happy to add a KV-cache test if wanted.
2. A stale MagCache state crashes the KV-cache pipeline. If MagCacheConfig.num_inference_steps doesn't match the pipeline call, or a run is interrupted, the state isn't reset. The next call can then skip the step-0 KV prefill, and step 1 raises RuntimeError: KV cache has not been populated yet. (reproduced with the tiny model). Other models degrade silently in the same situation. Open question: reset MagCache state at the start of each pipeline call, or raise a clearer error?
3. Possibly redundant ndim guards in the calibration alignment. Residuals are always (B, L, D), so curr_norm.ndim == 2 and prev_norm.ndim == 2 is always true. Kept for consistency with the surrounding hook code.

Dead code

None. Every added line runs:

Location	Status	Reason
hooks/_helpers.py registration	Used	Looked up by TransformerBlockRegistry.get when MagCache is enabled
hooks/mag_cache.py head-hook diff < 0 branch	Used	A skip right after the KV prefill step
hooks/mag_cache.py calibration alignment	Used	Calibration across the prefill � cached transition
Test classes in both test files	Used	Collected by pytest

Documentation impact

• docs/source/en/optimization/cache.md: Qwen-Image 2.1 could be added to the list of models that need their own calibration.
• docs/source/en/optimization/cache.md: this bug predates the PR. The MagCache example calls enable_cache twice without disable_cache(), which raises ValueError: Caching has already been enabled.
• Possible pitfalls.md entry: a model with a KV/prefix cache changes the hooked sequence length between steps, so caches keeping per-step state must align on the trailing tokens. Both TaylorSeer (Fix TaylorSeer cache crash when the feature shape changes between steps #14831) and MagCache hit this.

Verdict

READY. Open for reviewer input: item 1 (KV-cache test for the hook change), item 2 (state reset vs. clearer error), and whether the disable_cache() docs fix belongs in this PR.

[infinite-Joy]

The existing guard checks rank, batch size, and channel count. This new branch takes the last input length tokens of the residual. Does it make sense to check whether the positions represent the same tokens that it assumes the trailing target image corresponds to.

[github-actions]

Hi @rootonchair, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.