Add First Block Cache support for Flux 2

Author: guptaishaan

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=1,
+            return_encoder_hidden_states_index=0,
+        ),
+    )
+
     # HunyuanVideo
     TransformerBlockRegistry.register(
         model_class=HunyuanVideoTransformerBlock,

src/diffusers/models/transformers/transformer_flux2.py

@@ -837,18 +837,13 @@ def __init__(
     def forward(
         self,
         hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor | None,
+        encoder_hidden_states: torch.Tensor,
         temb_mod: torch.Tensor,
         image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         joint_attention_kwargs: dict[str, Any] | None = None,
-        split_hidden_states: bool = False,
-        text_seq_len: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already
-        # concatenated
-        if encoder_hidden_states is not None:
-            text_seq_len = encoder_hidden_states.shape[1]
-            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 
         mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0]
 
@@ -866,11 +861,8 @@ def forward(
         if hidden_states.dtype == torch.float16:
             hidden_states = hidden_states.clip(-65504, 65504)
 
-        if split_hidden_states:
-            encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
-            return encoder_hidden_states, hidden_states
-        else:
-            return hidden_states
+        encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
+        return encoder_hidden_states, hidden_states
 
 
 class Flux2TransformerBlock(nn.Module):
@@ -1374,12 +1366,9 @@ def forward(
                     joint_attention_kwargs=kv_attn_kwargs,
                 )
 
-        # Concatenate text and image streams for single-block inference
-        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
-
         # Blend single block modulation for extract mode: [txt_mod, ref_mod, img_mod]
         if kv_cache_mode == "extract" and num_ref_tokens > 0:
-            total_single_len = hidden_states.shape[1]
+            total_single_len = num_txt_tokens + hidden_states.shape[1]
             single_stream_mod = _blend_single_block_mods(
                 single_stream_mod, ref_single_mod, num_txt_tokens, num_ref_tokens, total_single_len
             )
@@ -1396,28 +1385,26 @@ def forward(
                 kv_attn_kwargs_single["kv_cache"] = kv_cache.get_single(index_block)
 
             if torch.is_grad_enabled() and self.gradient_checkpointing:
-                hidden_states = self._gradient_checkpointing_func(
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
                     block,
                     hidden_states,
-                    None,
+                    encoder_hidden_states,
                     single_stream_mod,
                     concat_rotary_emb,
                     kv_attn_kwargs_single,
                 )
             else:
-                hidden_states = block(
+                encoder_hidden_states, hidden_states = block(
                     hidden_states=hidden_states,
-                    encoder_hidden_states=None,
+                    encoder_hidden_states=encoder_hidden_states,
                     temb_mod=single_stream_mod,
                     image_rotary_emb=concat_rotary_emb,
                     joint_attention_kwargs=kv_attn_kwargs_single,
                 )
 
-        # Remove text tokens (and ref tokens in extract mode) from concatenated stream
+        # Remove ref tokens (extract mode only) from the image stream
         if kv_cache_mode == "extract" and num_ref_tokens > 0:
-            hidden_states = hidden_states[:, num_txt_tokens + num_ref_tokens :, ...]
-        else:
-            hidden_states = hidden_states[:, num_txt_tokens:, ...]
+            hidden_states = hidden_states[:, num_ref_tokens:, ...]
 
         # 7. Output layers
         hidden_states = self.norm_out(hidden_states, temb)

tests/models/transformers/test_models_transformer_flux2.py

@@ -34,6 +34,7 @@
     BaseModelTesterConfig,
     BitsAndBytesTesterMixin,
     ContextParallelTesterMixin,
+    FirstBlockCacheTesterMixin,
     GGUFCompileTesterMixin,
     GGUFTesterMixin,
     LoraHotSwappingForModelTesterMixin,
@@ -198,6 +199,10 @@ def test_tensor_parallel_neuron_inference(self):
         )
 
 
+class TestFlux2TransformerFBCCache(Flux2TransformerTesterConfig, FirstBlockCacheTesterMixin):
+    """FirstBlockCache tests for Flux2 Transformer."""
+
+
 class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin):
     """LoRA adapter tests for Flux2 Transformer."""
 

tests/pipelines/flux2/test_pipeline_flux2_klein.py

@@ -23,6 +23,7 @@
 )
 from ..testing_utils import (
     BasePipelineTesterConfig,
+    FirstBlockCacheTesterMixin,
     MemoryTesterMixin,
     PipelineTesterMixin,
     check_qkv_fused_layers_exist,
@@ -198,6 +199,10 @@ class TestFlux2KleinPipelineMemory(Flux2KleinPipelineTesterConfig, MemoryTesterM
     """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux2 Klein pipeline."""
 
 
+class TestFlux2KleinPipelineFirstBlockCache(Flux2KleinPipelineTesterConfig, FirstBlockCacheTesterMixin):
+    """First Block Cache tests for the Flux2 Klein pipeline."""
+
+
 @require_torch_neuron
 class TestFlux2KleinPipelineIntegration:
     ckpt_id = "black-forest-labs/FLUX.2-klein-4B"