From 25586efd20b5e19353b80e478e62d07d4bf8b141 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Fri, 11 Sep 2026 19:22:00 +0700 Subject: [PATCH 1/2] [hooks] Support Flux2 in MagCache MagCache could not be enabled on `Flux2Transformer2DModel` (Flux2 and Flux2 Klein): 1. `Flux2TransformerBlock` and `Flux2SingleTransformerBlock` were not registered in `TransformerBlockRegistry`, so `apply_mag_cache` raised `ValueError: Model class ... not registered`. 2. Once registered, the tail hook raised a shape error. The head hook records the image-only input of `transformer_blocks[0]`, but the Flux2 model concatenates the text and image streams before the single-stream blocks and calls them with `encoder_hidden_states=None`, so the last single block returns the fused `[text; image]` sequence. The residual `tail_output - head_input` could not be formed. Register the two Flux2 blocks (the single block with `return_encoder_hidden_states_index=None`, matching how the model calls it) and add `_align_to_hidden_states`, which keeps only the trailing image tokens of the fused tensor when computing the residual. The cached residual is therefore image-only, which is what the reference MagCache4FLUX implementation does, and it matches the head input on skipped steps. This also replaces a dead "fallback" branch in the tail hook that subtracted mismatched shapes anyway. Tests: `MagCacheTesterMixin` for `Flux2Transformer2DModel`, `Flux2Pipeline` and `Flux2KleinPipeline`. FirstBlockCache has the same head/tail structure and still fails on Flux2; left for a follow-up. --- docs/source/en/optimization/cache.md | 2 +- src/diffusers/hooks/_helpers.py | 17 +++++++++++ src/diffusers/hooks/mag_cache.py | 28 ++++++++++++++----- .../test_models_transformer_flux2.py | 5 ++++ tests/pipelines/flux2/test_pipeline_flux2.py | 5 ++++ .../flux2/test_pipeline_flux2_klein.py | 5 ++++ 6 files changed, 54 insertions(+), 8 deletions(-) diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md index 079f073b73f0..200a601e4dcf 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -116,7 +116,7 @@ pipe.transformer.enable_cache(config) [MagCache](https://github.com/Zehong-Ma/MagCache) accelerates inference by skipping transformer blocks based on the magnitude of the residual update. It observes that the magnitude of updates (Output - Input) decays predictably over the diffusion process. By accumulating an "error budget" based on pre-computed magnitude ratios, it dynamically decides when to skip computation and reuse the previous residual. -MagCache relies on **Magnitude Ratios** (`mag_ratios`), which describe this decay curve. These ratios are specific to the model checkpoint and scheduler. +MagCache relies on **Magnitude Ratios** (`mag_ratios`), which describe this decay curve. These ratios are specific to the model checkpoint and scheduler. The bundled `FLUX_MAG_RATIOS` were measured on FLUX.1; other models, including Flux2 and Flux2 Klein, need their own calibration run. To use MagCache, you typically follow a two-step process: **Calibration** and **Inference**. diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 9cbe5bc8108f..0f29e97661c4 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -175,6 +175,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, @@ -246,6 +247,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, diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index e5f0aaebc01a..bdc96cd99ebb 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -80,6 +80,25 @@ def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor: return src_array[mapped_indices] +def _align_to_hidden_states(tensor: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + """ + Slice `tensor` down to the trailing `hidden_states.shape[1]` tokens when it is longer along the sequence axis. + + Models such as Flux2 concatenate the text and image streams before their single-stream blocks, so the tail block + returns `[text; image]` while the head block only saw the image stream. The image stream sits at the end of the + fused sequence in these layouts. + """ + if ( + tensor.ndim == 3 + and hidden_states.ndim == 3 + and tensor.shape[0] == hidden_states.shape[0] + and tensor.shape[2] == hidden_states.shape[2] + and tensor.shape[1] > hidden_states.shape[1] + ): + return tensor[:, -hidden_states.shape[1] :] + return tensor + + @dataclass class MagCacheConfig: r""" @@ -339,15 +358,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): if in_hidden is None: return output - # Determine residual + # Keep the image tokens only, so the residual matches the head input it is added to on skipped steps. + out_hidden = _align_to_hidden_states(out_hidden, in_hidden) if out_hidden.shape == in_hidden.shape: residual = out_hidden - in_hidden - elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]: - diff = in_hidden.shape[1] - out_hidden.shape[1] - if diff == 0: - residual = out_hidden - in_hidden - else: - residual = out_hidden - in_hidden # Fallback to matching tail else: # Fallback for completely mismatched shapes residual = out_hidden diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 3263ce68202c..75fe4a1701dd 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -38,6 +38,7 @@ GGUFTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, + MagCacheTesterMixin, MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, @@ -713,3 +714,7 @@ def pretrained_model_name_or_path(self): @property def pretrained_model_kwargs(self): return {"subfolder": "transformer"} + + +class TestFlux2TransformerMagCache(Flux2TransformerTesterConfig, MagCacheTesterMixin): + """MagCache tests for Flux2 Transformer.""" diff --git a/tests/pipelines/flux2/test_pipeline_flux2.py b/tests/pipelines/flux2/test_pipeline_flux2.py index ac72d843cd05..9c9bdb2468ef 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2.py +++ b/tests/pipelines/flux2/test_pipeline_flux2.py @@ -13,6 +13,7 @@ BasePipelineTesterConfig, LoraMemoryTesterMixin, LoraTesterMixin, + MagCacheTesterMixin, MemoryTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist, @@ -204,3 +205,7 @@ class TestFlux2PipelineLoRAMemory(Flux2PipelineTesterConfig, LoraMemoryTesterMix # See `TestFlux2PipelineLoRA`. denoiser_target_modules = {"transformer": ["to_qkv_mlp_proj", "to_k"]} + + +class TestFlux2PipelineMagCache(Flux2PipelineTesterConfig, MagCacheTesterMixin): + """MagCache tests for the Flux2 pipeline.""" diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index 0d7139b21e16..84cc08abb53d 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein.py @@ -23,6 +23,7 @@ ) from ..testing_utils import ( BasePipelineTesterConfig, + MagCacheTesterMixin, MemoryTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist, @@ -283,3 +284,7 @@ def test_flux2_klein_neuron_compile_128(self): assert image.shape == (1, 128, 128, 3) assert not np.isnan(image).any(), "Output contains NaN values" assert (image >= 0.0).all() and (image <= 1.0).all(), "Output pixel values outside [0, 1]" + + +class TestFlux2KleinPipelineMagCache(Flux2KleinPipelineTesterConfig, MagCacheTesterMixin): + """MagCache tests for the Flux2 Klein pipeline.""" From 97696e9f3336749c4540a078948a7e670d18ceba Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Fri, 11 Sep 2026 23:49:13 +0700 Subject: [PATCH 2/2] [hooks] Slice the fused tail output in the existing mismatch branch Implement the Flux2 residual fix inside the tail hook's sequence-length mismatch branch instead of a separate `_align_to_hidden_states` helper. The surrounding structure stays as it was: the `diff == 0` case still covers a batch-size mismatch via broadcasting, and only the mismatched-length case changes to subtract the trailing image tokens. --- src/diffusers/hooks/mag_cache.py | 29 ++++++++--------------------- 1 file changed, 8 insertions(+), 21 deletions(-) diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index bdc96cd99ebb..4ae1b5764edc 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -80,25 +80,6 @@ def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor: return src_array[mapped_indices] -def _align_to_hidden_states(tensor: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: - """ - Slice `tensor` down to the trailing `hidden_states.shape[1]` tokens when it is longer along the sequence axis. - - Models such as Flux2 concatenate the text and image streams before their single-stream blocks, so the tail block - returns `[text; image]` while the head block only saw the image stream. The image stream sits at the end of the - fused sequence in these layouts. - """ - if ( - tensor.ndim == 3 - and hidden_states.ndim == 3 - and tensor.shape[0] == hidden_states.shape[0] - and tensor.shape[2] == hidden_states.shape[2] - and tensor.shape[1] > hidden_states.shape[1] - ): - return tensor[:, -hidden_states.shape[1] :] - return tensor - - @dataclass class MagCacheConfig: r""" @@ -358,10 +339,16 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): if in_hidden is None: return output - # Keep the image tokens only, so the residual matches the head input it is added to on skipped steps. - out_hidden = _align_to_hidden_states(out_hidden, in_hidden) + # Determine residual if out_hidden.shape == in_hidden.shape: residual = out_hidden - in_hidden + elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]: + diff = in_hidden.shape[1] - out_hidden.shape[1] + if diff == 0: + residual = out_hidden - in_hidden + else: + # The tail returned the fused text+image sequence (e.g. Flux2); the image tokens sit at the end. + residual = out_hidden[:, -in_hidden.shape[1] :] - in_hidden else: # Fallback for completely mismatched shapes residual = out_hidden