[Feature]: support Flux.2-dev tea_cache (#1871)

Co-authored-by: wuzhongjian <wuzhongjian@cmss.chinamobile.com>

tests/diffusion/cache/test_teacache_extractors.py

@@ -22,7 +22,7 @@ import pytest
 import torch
 
 from tests.utils import hardware_test
-from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_klein_context
+from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
@@ -174,3 +174,106 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
                 img_ids=torch.randint(0, 64, (1, 1024, 4)),
                 txt_ids=torch.randint(0, 64, (1, 512, 4)),
             )
+
+
+class TestFlux2Extractor(BaseExtractorTest):
+    """Test extract_flux2_context function."""
+
+    def get_extractor(self):
+        return extract_flux2_context
+
+    @pytest.fixture
+    def flux2_module(self):
+        """Create a minimal Flux2Transformer2DModel for testing."""
+        from vllm_omni.diffusion.models.flux2.flux2_transformer import Flux2Transformer2DModel
+
+        model = Flux2Transformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=48,
+            attention_head_dim=128,
+            joint_attention_dim=15360,
+        )
+        return model
+
+    def get_module(self, flux2_module):
+        return flux2_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        """Create sample input tensors for Flux2.
+
+        Note: hidden_states uses in_channels=128 (default for Flux2),
+        not inner_dim=6144. The x_embedder projects from 128 -> 6144.
+        encoder_hidden_states uses joint_attention_dim=15360 (model default),
+        which then gets projected to inner_dim=6144 by context_embedder.
+        """
+        batch_size = 1
+        img_seq_len = 1024
+        txt_seq_len = 512
+        in_channels = 128  # Model default in_channels
+        txt_dim = 15360  # Model default joint_attention_dim
+
+        return {
+            "hidden_states": torch.randn(batch_size, img_seq_len, in_channels),
+            "encoder_hidden_states": torch.randn(batch_size, txt_seq_len, txt_dim),
+            "timestep": torch.tensor([500]),
+            "img_ids": torch.randint(0, 64, (batch_size, img_seq_len, 4)),
+            "txt_ids": torch.randint(0, 64, (batch_size, txt_seq_len, 4)),
+            "guidance": torch.tensor([3.5]),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_modulated_input_shape(self, flux2_module, sample_inputs):
+        """Test that modulated_input has correct shape matching the model's inner_dim.
+
+        Note: After x_embedder projection, hidden_states are projected from
+        in_channels (128) to inner_dim (6144), so modulated_input should match
+        the projected shape, not the input shape.
+        """
+        context = extract_flux2_klein_context(flux2_module, **sample_inputs)
+
+        batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
+        inner_dim = flux2_module.inner_dim
+        assert context.modulated_input.shape == (batch_size, img_seq_len, inner_dim)
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_run_transformer_blocks_callable(self, flux2_module, sample_inputs):
+        """Test that run_transformer_blocks is callable."""
+        context = extract_flux2_context(flux2_module, **sample_inputs)
+        assert callable(context.run_transformer_blocks)
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_postprocess_callable(self, flux2_module, sample_inputs):
+        """Test that postprocess is callable."""
+        context = extract_flux2_context(flux2_module, **sample_inputs)
+        assert callable(context.postprocess)
+
+    def test_without_guidance(self, flux2_module, sample_inputs):
+        """Test context extraction works without guidance (no CFG)."""
+        inputs = sample_inputs.copy()
+        inputs["guidance"] = None
+
+        context = extract_flux2_context(flux2_module, **inputs)
+
+        assert context is not None
+        assert context.temb is not None
+
+    @pytest.mark.cpu
+    def test_invalid_module_raises_error(self):
+        """Test that invalid module without transformer_blocks raises ValueError."""
+        invalid_module = Mock()
+        invalid_module.transformer_blocks = []
+
+        with pytest.raises(ValueError, match="Module must have transformer_blocks"):
+            extract_flux2_context(
+                invalid_module,
+                hidden_states=torch.randn(1, 1024, 6144),
+                encoder_hidden_states=torch.randn(1, 512, 15360),
+                timestep=torch.tensor([500]),
+                img_ids=torch.randint(0, 64, (1, 1024, 4)),
+                txt_ids=torch.randint(0, 64, (1, 512, 4)),
+            )