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

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,
 }