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..4ae1b5764edc 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -347,7 +347,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): if diff == 0: residual = out_hidden - in_hidden else: - residual = out_hidden - in_hidden # Fallback to matching tail + # 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 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."""