[Feature]: support Flux.1-dev tea_cache (#2774)

tests/diffusion/cache/test_teacache_extractors.py

@@ -13,6 +13,8 @@ used by TeaCache. Each model's extractor can be tested by:
 
 Currently implemented:
 - TestFlux2KleinExtractor: Flux2Klein model extractor
+- TestFlux2Extractor: Flux2 model extractor
+- TestFluxExtractor: Flux model extractor
 """
 
 from abc import ABC, abstractmethod
@@ -22,7 +24,12 @@ import pytest
 import torch
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context
+from vllm_omni.diffusion.cache.teacache.extractors import (
+    extract_flux2_context,
+    extract_flux2_klein_context,
+    extract_flux_context,
+)
+from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
@@ -234,7 +241,7 @@ class TestFlux2Extractor(BaseExtractorTest):
         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)
+        context = extract_flux2_context(flux2_module, **sample_inputs)
 
         batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
         inner_dim = flux2_module.inner_dim
@@ -277,3 +284,139 @@ class TestFlux2Extractor(BaseExtractorTest):
                 img_ids=torch.randint(0, 64, (1, 1024, 4)),
                 txt_ids=torch.randint(0, 64, (1, 512, 4)),
             )
+
+
+@pytest.mark.cpu
+class TestFluxExtractor(BaseExtractorTest):
+    """Test extract_flux_context function."""
+
+    @pytest.fixture(autouse=True)
+    def cpu_vllm_config(self):
+        """Force CPU custom-op dispatch for this test class."""
+        from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+        with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+            yield
+
+    @pytest.fixture(autouse=True)
+    def mock_flux_attention_backend(self):
+        """Use the SDPA backend so FLUX can be instantiated in CPU tests."""
+        from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+
+        with patch("vllm_omni.diffusion.attention.layer.get_attn_backend", return_value=SDPABackend):
+            yield
+
+    def get_extractor(self):
+        return extract_flux_context
+
+    @pytest.fixture
+    def flux_module(self):
+        """Create a minimal FluxTransformer2DModel for testing."""
+        return FluxTransformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=16,
+            joint_attention_dim=32,
+            pooled_projection_dim=16,
+            axes_dims_rope=(4, 4, 8),
+        )
+
+    @pytest.fixture
+    def flux_module_without_guidance(self):
+        """Create a minimal non-guidance-distilled FLUX transformer."""
+        return FluxTransformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=16,
+            joint_attention_dim=32,
+            pooled_projection_dim=16,
+            guidance_embeds=False,
+            axes_dims_rope=(4, 4, 8),
+        )
+
+    def get_module(self, flux_module):
+        return flux_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        """Create sample input tensors for Flux."""
+        batch_size = 1
+        img_seq_len = 16
+        txt_seq_len = 8
+        in_channels = 64  # Flux default in_channels
+        txt_dim = 32
+        pooled_dim = 16
+
+        return {
+            "hidden_states": torch.randn(batch_size, img_seq_len, in_channels),
+            "encoder_hidden_states": torch.randn(batch_size, txt_seq_len, txt_dim),
+            "pooled_projections": torch.randn(batch_size, pooled_dim),
+            "timestep": torch.tensor([500]),
+            "img_ids": torch.randint(0, 64, (batch_size, img_seq_len, 3)),
+            "txt_ids": torch.randint(0, 64, (batch_size, txt_seq_len, 3)),
+            "guidance": torch.tensor([3.5]),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    def test_modulated_input_shape(self, flux_module, sample_inputs):
+        """Test that modulated_input has the projected FLUX inner dimension."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+
+        batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
+        assert context.modulated_input.shape == (batch_size, img_seq_len, flux_module.inner_dim)
+
+    def test_run_transformer_blocks_callable(self, flux_module, sample_inputs):
+        """Test that run_transformer_blocks is callable."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        assert callable(context.run_transformer_blocks)
+
+    def test_postprocess_callable(self, flux_module, sample_inputs):
+        """Test that postprocess is callable."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        assert callable(context.postprocess)
+
+    def test_postprocess_output_shape(self, flux_module, sample_inputs):
+        """Test that postprocess projects back to the input channel width."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        output = context.postprocess(context.hidden_states)
+
+        assert output.sample.shape == sample_inputs["hidden_states"].shape
+
+    def test_postprocess_return_tuple_when_return_dict_false(self, flux_module, sample_inputs):
+        """Test that postprocess honors return_dict=False."""
+        context = extract_flux_context(flux_module, **sample_inputs, return_dict=False)
+        output = context.postprocess(context.hidden_states)
+
+        assert isinstance(output, tuple)
+        assert len(output) == 1
+        assert output[0].shape == sample_inputs["hidden_states"].shape
+
+    def test_without_guidance(self, flux_module_without_guidance, sample_inputs):
+        """Test context extraction works for FLUX variants without guidance embeddings."""
+        inputs = sample_inputs.copy()
+        inputs["guidance"] = None
+
+        context = extract_flux_context(flux_module_without_guidance, **inputs)
+
+        assert context is not None
+        assert context.temb is not None
+
+    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_flux_context(
+                invalid_module,
+                hidden_states=torch.randn(1, 16, 64),
+                encoder_hidden_states=torch.randn(1, 8, 32),
+                pooled_projections=torch.randn(1, 16),
+                timestep=torch.tensor([500]),
+                img_ids=torch.randint(0, 64, (1, 16, 3)),
+                txt_ids=torch.randint(0, 64, (1, 8, 3)),
+            )