[Feat] support TeaCache for FLUX1.dev

[RuixiangMa]

Purpose

support TeaCache for FLUX1

Test Plan

Test Result

• 2* NVIDIA 4090(24G)
• TP=2 + enable_cpu_offload

curl -X POST http://localhost:8004/v1/images/generations   -H "Content-Type: application/json"   -d '{
    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
    "size": "1024x1024",
    "num_inference_steps": 50,
    "cfg_scale": 4.0,
    "guidance_scale": 4.0,
    "seed": 42
  }' | jq -r '.data[0].b64_json' | base64 -d > dragon.png

Metric	NO TeaCache	TeaCache (0.2)	TeaCache (0.4)	TeaCache (0.6)
image
Time	27.652 s/img	17.858 s/img	12.524 s/img	10.549 s/img

Refactored when PR #1234 is merged

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6102c32feb

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Preserve Flux forward return type in TeaCache postprocess

The Flux extractor ignores the caller's return_dict flag and always returns Transformer2DModelOutput, which changes the public contract of FluxTransformer2DModel.forward under TeaCache. In this repo, pipeline_flux.diffuse explicitly calls self.transformer(..., return_dict=False) and currently indexes [0], but other call sites (or downstream users) can legitimately rely on receiving a tuple when return_dict=False; this mismatch can cause behavioral regressions only when TeaCache is enabled.

Useful? React with 👍 / 👎.