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).

[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.