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

docs/user_guide/diffusion_features.md

@@ -108,7 +108,7 @@ The following tables show which models support each feature:
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅  | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |

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)),
+            )

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1069,6 +1069,130 @@ def extract_flux2_context(
     )
 
 
+def extract_flux_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    pooled_projections: torch.Tensor,
+    timestep: torch.LongTensor,
+    img_ids: torch.Tensor,
+    txt_ids: torch.Tensor,
+    guidance: torch.Tensor | None = None,
+    joint_attention_kwargs: dict[str, Any] | None = None,
+    **kwargs: Any,
+) -> CacheContext:
+    """
+    Extract cache context for FluxTransformer2DModel.
+
+    This mirrors the standard FLUX.1 transformer forward path while exposing
+    the first modulated image stream tensor as TeaCache's similarity signal.
+
+    Args:
+        module: FluxTransformer2DModel instance
+        hidden_states: Input image hidden states tensor
+        encoder_hidden_states: Text encoder outputs
+        pooled_projections: Pooled text conditioning
+        timestep: Current diffusion timestep
+        img_ids: Image position IDs for RoPE
+        txt_ids: Text position IDs for RoPE
+        guidance: Optional guidance scale for guidance-distilled variants
+        joint_attention_kwargs: Additional attention kwargs
+        **kwargs: Additional keyword arguments
+
+    Returns:
+        CacheContext with all information needed for generic caching
+    """
+    from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+    if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
+        raise ValueError("Module must have transformer_blocks")
+
+    # ============================================================================
+    # PREPROCESSING (Flux1-specific)
+    # ============================================================================
+    hidden_states = module.x_embedder(hidden_states)
+    timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
+
+    if guidance is not None:
+        guidance = guidance.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
+
+    temb = (
+        module.time_text_embed(timestep, pooled_projections)
+        if guidance is None
+        else module.time_text_embed(timestep, guidance, pooled_projections)
+    )
+    encoder_hidden_states = module.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    if current_omni_platform.is_npu():
+        freqs_cos, freqs_sin = module.pos_embed(ids.cpu())
+        image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu())
+    else:
+        image_rotary_emb = module.pos_embed(ids)
+
+    # ============================================================================
+    # EXTRACT MODULATED INPUT (for cache decision)
+    # ============================================================================
+    first_block = module.transformer_blocks[0]
+    modulated_input, *_ = first_block.norm1(hidden_states, emb=temb)
+
+    # ============================================================================
+    # DEFINE TRANSFORMER EXECUTION (Flux1-specific)
+    # ============================================================================
+    def run_transformer_blocks() -> tuple[torch.Tensor, torch.Tensor]:
+        h = hidden_states
+        e = encoder_hidden_states
+
+        for block in module.transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        for block in module.single_transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        return (h, e)
+
+    return_dict = kwargs.get("return_dict", True)
+
+    # ============================================================================
+    # DEFINE POSTPROCESSING
+    # ============================================================================
+    def postprocess(h: torch.Tensor) -> Any:
+        h = module.norm_out(h, temb)
+        output = module.proj_out(h)
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    # ============================================================================
+    # RETURN CONTEXT
+    # ============================================================================
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+
+
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -1083,8 +1207,8 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "StableAudioDiTModel": extract_stable_audio_context,
     "Flux2Transformer2DModel": extract_flux2_context,
     "LongCatImageTransformer2DModel": extract_longcat_context,
+    "FluxTransformer2DModel": extract_flux_context,
     # Future models:
-    # "FluxTransformer2DModel": extract_flux_context,
     # "CogVideoXTransformer3DModel": extract_cogvideox_context,
 }