Skip to content

Allow Flux2 to use MagCache #14351

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Ednaordinary wants to merge 5 commits into
base: main
Choose a base branch
from
+27 −1
Open

Allow Flux2 to use MagCache #14351

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
febae44
fix hooks
Ednaordinary Aug 1, 2026
17ea853
add tail overrides
Ednaordinary Aug 1, 2026
290efa3
fix
Ednaordinary Aug 1, 2026
f803437
remove debug
Ednaordinary Aug 1, 2026
3ca8570
make style, quality
Ednaordinary Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions src/diffusers/hooks/_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ def _register_attention_processors_metadata():
from ..models.attention_processor import AttnProcessor2_0
from ..models.transformers.transformer_cogview4 import CogView4AttnProcessor
from ..models.transformers.transformer_flux import FluxAttnProcessor
from ..models.transformers.transformer_flux2 import Flux2AttnProcessor
from ..models.transformers.transformer_hunyuanimage import HunyuanImageAttnProcessor
from ..models.transformers.transformer_qwenimage import QwenDoubleStreamAttnProcessor2_0
from ..models.transformers.transformer_wan import WanAttnProcessor2_0
Expand Down Expand Up @@ -144,6 +145,12 @@ def _register_attention_processors_metadata():
metadata=AttentionProcessorMetadata(skip_processor_output_fn=_skip_proc_output_fn_Attention_FluxAttnProcessor),
)

# Flux2AttnProcessor
AttentionProcessorRegistry.register(
model_class=Flux2AttnProcessor,
metadata=AttentionProcessorMetadata(skip_processor_output_fn=_skip_proc_output_fn_Attention_FluxAttnProcessor),
)

# QwenDoubleStreamAttnProcessor2
AttentionProcessorRegistry.register(
model_class=QwenDoubleStreamAttnProcessor2_0,
Expand Down Expand Up @@ -175,6 +182,7 @@ def _register_transformer_blocks_metadata():
from ..models.transformers.transformer_bria import BriaTransformerBlock
from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
from ..models.transformers.transformer_flux2 import Flux2SingleTransformerBlock, Flux2TransformerBlock
from ..models.transformers.transformer_hunyuan_video import (
HunyuanVideoSingleTransformerBlock,
HunyuanVideoTokenReplaceSingleTransformerBlock,
Expand Down Expand Up @@ -246,6 +254,22 @@ def _register_transformer_blocks_metadata():
),
)

# Flux2
TransformerBlockRegistry.register(
model_class=Flux2TransformerBlock,
metadata=TransformerBlockMetadata(
return_hidden_states_index=1,
return_encoder_hidden_states_index=0,
),
)
TransformerBlockRegistry.register(
model_class=Flux2SingleTransformerBlock,
metadata=TransformerBlockMetadata(
return_hidden_states_index=0,
return_encoder_hidden_states_index=None,
),
)

# HunyuanVideo
TransformerBlockRegistry.register(
model_class=HunyuanVideoTransformerBlock,
Expand Down
4 changes: 3 additions & 1 deletion src/diffusers/hooks/mag_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,8 +346,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
diff = in_hidden.shape[1] - out_hidden.shape[1]
if diff == 0:
residual = out_hidden - in_hidden
elif diff > 0:
residual = out_hidden - in_hidden[:, diff:] # Fallback to matching tail
else:
residual = out_hidden - in_hidden # Fallback to matching tail
residual = out_hidden[:, -diff:] - in_hidden # Fallback to matching tail
Comment on lines +349 to +352

@sayakpaul sayakpaul Aug 1, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you explain the changes?

@Ednaordinary Ednaordinary Aug 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

else:
# Fallback for completely mismatched shapes
residual = out_hidden
Expand Down
Loading