Allow Flux2 to use MagCache

[Ednaordinary]

What does this PR do?

MagCache on Flux2 used to error out, this should fix it by adding flux2 to hooks (maybe allowing it to be used by other caches? I haven't checked) and fixes a case in MagCache that a code comment signaled intent to handle but never did.

Tested on 3090 ti

Test script

from transformers import Mistral3ForConditionalGeneration
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel, MagCacheConfig, AutoencoderKLFlux2

torch.set_float32_matmul_precision('medium')

transformer = Flux2Transformer2DModel.from_pretrained("Disty0/FLUX.2-dev-SDNQ-uint4-svd-r32", subfolder="transformer", dtype=torch.bfloat16)
text_enc = Mistral3ForConditionalGeneration.from_pretrained("Disty0/FLUX.2-dev-SDNQ-uint4-svd-r32", subfolder="text_encoder", dtype=torch.bfloat16)
vae = AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-small-decoder", torch_dtype=torch.bfloat16)
model = Flux2Pipeline.from_pretrained("black-forest-labs/FLUX.2-dev", transformer=transformer, text_encoder=text_enc, vae=vae, dtype=torch.bfloat16)
model.enable_model_cpu_offload()
model.vae.enable_slicing()

# calib_config = MagCacheConfig(calibrate=True, num_inference_steps=12)
# model.transformer.enable_cache(calib_config)

mag_config = MagCacheConfig(
    threshold=0.50,
    retention_ratio=0.0,
    max_skip_steps=10,
    # 50 steps
    # mag_ratios=[1.0, 1.0119190216064453, 0.9999825358390808, 1.0202345848083496, 1.0423189401626587, 1.0240278244018555, 1.0042935609817505, 1.0045005083084106, 1.0095233917236328, 0.9842535257339478, 0.9970784187316895, 0.9907071590423584, 0.9893150329589844, 0.9965107440948486, 0.9854446649551392, 0.9892123937606812, 0.988654375076294, 0.9940516352653503, 0.9938445687294006, 0.9816863536834717, 1.0045714378356934, 0.9928746223449707, 0.9888252019882202, 0.9963063597679138, 0.9899376630783081, 0.9925428032875061, 0.9984951019287109, 0.9905719757080078, 0.9902525544166565, 0.9907461404800415, 0.9965947866439819, 0.9878503084182739, 0.9880349636077881, 0.9960471987724304, 0.9860436320304871, 0.986957311630249, 0.9947977662086487, 0.985883355140686, 0.9843263626098633, 0.9875669479370117, 0.9781431555747986, 0.985008716583252, 0.9868263006210327, 0.9784890413284302, 0.9644964933395386, 0.988265872001648, 0.9522878527641296, 0.9906107783317566, 0.9935293197631836, 1.0806248188018799],
    # 28 steps
    # mag_ratios = [1.0, 1.0117623805999756, 1.0202363729476929, 1.0434367656707764, 1.0281867980957031, 1.0123460292816162, 0.9819282293319702, 0.9912159442901611, 0.985506534576416, 0.9743061065673828, 0.9888518452644348, 0.9879252910614014, 0.9852353930473328, 0.9878948330879211, 0.9867265224456787, 0.9788063764572144, 0.9950540065765381, 0.9860763549804688, 0.985837459564209, 0.9816434979438782, 0.9878530502319336, 0.9709774255752563, 0.9692974090576172, 0.9750104546546936, 0.9658247232437134, 0.9481929540634155, 0.9442211389541626, 0.971260666847229],
    # 12 steps
    mag_ratios = [1.0, 1.0323116779327393, 1.072181224822998, 0.9847173690795898, 0.9602146148681641, 0.9699361324310303, 0.9745591282844543, 0.9675220251083374, 0.9641062021255493, 0.9606950283050537, 0.933021068572998, 0.8983047008514404],
    num_inference_steps=12,
)

model.transformer.enable_cache(mag_config) 

prompt = "A cat holding a sign that says hello world"
image = model(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=4,
    num_inference_steps=12,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(42)
).images[0]
image.save("flux2-dev-uint4-s12-g-0.50.png")

Threshold > Steps v	0	0.2	0.5
50	4 min	50 sec	49 sec
28	2 min 18 sec	36 sec	36 sec
12	1 min 14 sec	36 sec	19 sec

this cat has many limbs

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.

@AlanPonnachan @sayakpaul

[sayakpaul]

Can you explain the changes?

[Ednaordinary]

This branch of the if statement appears to have the intent to slice out the text encoding if it remains within the hidden state, but it was never been implemented (if you look on main, there is an if-else that does the exact same thing). On flux 2, this will error out because out_hidden ends up as size 4608 while in_hidden ends up as 4096 due to the single transformer block returning a fused hidden state (there is a single transformer block kwarg to break it into encoder_hidden_state and hidden_state, but I believe it would break other things and possibly the pipeline and it's cleaner to just deal with it here).

[sayakpaul]

Let's add a comment about it and also, a test to https://github.com/huggingface/diffusers/blob/main/tests/hooks/test_mag_cache.py.

Cc: @AlanPonnachan could you comment on this?

[AlanPonnachan]

yes, flux2 has this exception. It receives un-fused tensors, concatenates them internally, and returns a single fused tensor without splitting it back (https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py#L847-L873).
This pr solution will work for where text tokens are preprended (which is true for flux2). But if a model returns a fused tensor but appends the text tokens, this math will need adjusting??

yes, please add a comment and test case.

[Ednaordinary]

returns a fused tensor but appends the text tokens

Does the happen to be a model in diffusers a model that does this? That would be helpful as reference

[github-actions]

Hi @Ednaordinary, 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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Ednaordinary]

failing tests seem unrelated (timeouts)

[github-actions]

Hi @Ednaordinary, a follow-up on the reminder above: this PR still does not link an issue it fixes.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice. Adding a closing keyword (e.g. Fixes #1234) to the PR description, or a maintainer adding the no-issue-needed label, will prevent that.

[Ednaordinary]

I guess I can open an issue (is this the recommended course of action? If so, the bot should probably be more specific about that instead of implying it)

[sayakpaul]

Just a small comment and very sorry about the delay.