iframe-proxy

LawJarp-A · 2025-11-13T14:54:08Z

What does this PR do?

What is TeaCache?

TeaCache (Timestep Embedding Aware Cache) is a training-free caching technique that speeds up diffusion model inference by 1.5x-2.6x by reusing transformer block computations when consecutive timestep embeddings are similar.

Architecture & Design

TeaCache uses a ModelHook to intercept transformer forward passes without modifying model code. The algorithm:

Extracts modulated input from first transformer block (after norm1 + timestep embedding)
Computes relative L1 distance vs previous timestep
Applies model-specific polynomial rescaling: c[0]*x^4 + c[1]*x^3 + c[2]*x^2 + c[3]*x + c[4]
Accumulates rescaled distance across timesteps
If accumulated < threshold → Reuses cached residual (FAST)
If accumulated >= threshold → Full transformer pass (SLOW, update cache)

Key Design Features:

Hook-based: Integrates with HookRegistry and CacheMixin for lifecycle management
State Isolation: StateManager with context-aware state for CFG conditional/unconditional branches
Model Auto-Detection: Detects model type from class name and config path (specific variants checked first)
Boundary Guarantee: First and last timesteps always computed fully for quality
Specialized Strategies: Dual residual caching (CogVideoX), per-sequence-length caching (Lumina2)

Supported Models

All models support automatic coefficient detection based on model class name and config path. Custom coefficients can also be provided via TeaCacheConfig.

Benchmark Results (FLUX.1-dev)

Threshold	Time	Speedup
Baseline	9.26s	1.00x
0.2	6.85s	1.35x
0.4	5.24s	1.77x
0.6	4.64s	2.00x
0.8	4.18s	2.22x

Benchmark Results (Lumina2)

Threshold	Time	Speedup
Baseline	3.45s	1.00x
0.2	3.07s	1.12x
0.4	2.27s	1.52x
0.6	1.84s	1.88x

Benchmark Results (CogVideoX-2b)

Threshold	Time	Speedup
Baseline	26.27s	1.00x
0.3	23.97s	1.10x
0.5	22.57s	1.16x
0.7	19.31s	1.36x
0.9	17.38s	1.51x

Benchmark Results (Mochi)

Threshold	Time	Speedup
Baseline	7.71s	1.00x
0.05	6.27s	1.23x
0.06	6.03s	1.28x
0.08	5.73s	1.35x
0.10	5.41s	1.42x

Test Hardware: NVIDIA h100
Framework: Diffusers with TeaCache hooks
All tests: Same seed (42) for reproducibility

Usage

from diffusers import FluxPipeline
from diffusers.hooks import TeaCacheConfig

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
pipe.to("cuda")

# Enable TeaCache (1.75x speedup with 0.4 threshold)
config = TeaCacheConfig(rel_l1_thresh=0.4)
pipe.transformer.enable_cache(config)

image = pipe("A dragon on a crystal mountain", num_inference_steps=20).images[0]

pipe.transformer.disable_cache()

Configuration Options

The TeaCacheConfig supports the following parameters:

rel_l1_thresh (float, default=0.2): Threshold for accumulated relative L1 distance. Recommended values: 0.25 for ~1.5x speedup, 0.4 for ~1.8x, 0.6 for ~2.0x. Mochi models require lower thresholds (0.06-0.09).
coefficients (List[float], optional): Polynomial coefficients for rescaling L1 distance. Auto-detected based on model type if not provided.
num_inference_steps (int, optional): Total inference steps. Ensures first/last timesteps are always computed. Auto-detected if not provided.
num_inference_steps_callback (Callable[[], int], optional): Callback returning total inference steps. Alternative to num_inference_steps.
current_timestep_callback (Callable[[], int], optional): Callback returning current timestep. Used for debugging/statistics.

Files Changed

src/diffusers/hooks/teacache.py - Core implementation with model-specific forward functions
src/diffusers/models/cache_utils.py - CacheMixin integration
src/diffusers/hooks/__init__.py - Export TeaCacheConfig and apply_teacache
tests/hooks/test_teacache.py - Comprehensive unit tests

Fixes # (issue)
#12589
#12635

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

@sayakpaul @yiyixuxu @DN6

…nal_forward

LawJarp-A · 2025-11-13T17:29:06Z

Work done

Implement teacache for FLUX architecture using hooks (only flux for now)
add logging
add compatible tests

Waiting for feedback and review :)
cc: @dhruvrnaik @sayakpaul @yiyixuxu

LawJarp-A · 2025-11-17T11:40:23Z

Hi @sayakpaul @dhruvrnaik any updates?

sayakpaul · 2025-11-23T11:27:37Z

@LawJarp-A sorry about the delay on our end. @DN6 will review it soon.

HuggingFaceDocBuilderDev · 2025-11-23T11:35:31Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

DN6 · 2025-11-24T02:56:23Z

Hi @LawJarp-A I think we would need TeaCache to be implemented in a model agnostic way in order to merge the PR. The First Block Cache implementation is a good reference for this.

LawJarp-A · 2025-11-24T03:40:06Z

Hi @LawJarp-A I think we would need TeaCache to be implemented in a model agnostic way in order to merge the PR. The First Block Cache implementation is a good reference for this.

Yep @DN6 , I agree, I wanted to first implement it just for a single model and get feedback on that before I work on Model agnostic full implementation. I'm sort of working on it, didn't push it yet. I'll take a look at First block cache for reference as well.
On the same note, lemme know if there is anything to add to the current implementation

LawJarp-A · 2025-11-26T08:53:31Z

@DN6 updated it in a more model agnostic way.
Requesting review and feedback

… quality checks

…th auto-detection

LawJarp-A · 2025-12-02T09:51:52Z

Added multi model support, testing it thoroughly though.

LawJarp-A · 2025-12-08T11:51:18Z

Hi @DN6 @sayakpaul
Two questions, I'm almost done testing, I'll update the PR with more descriptive results and changes. And do final cleanup/merging etc

Any tests I should write and anything I can refer to for the same?
Added support for other models, I'll add pictures comparison with speedup and threshold to the PR as well?

In the meantime any feedback would be appreciated

sayakpaul · 2025-12-08T12:20:26Z

Thanks @LawJarp-A!

Any tests I should write and anything I can refer to for the same?

You can refer to #12569 for testing

Added support for other models, I'll add pictures comparison with speedup and threshold to the PR as well?

Yes, I think that is informative for users.

sayakpaul

Some initial feedback. Most important question is it seems like we need to craft different logic based on different model? Can we not keep it model agnostic?

LawJarp-A · 2025-12-08T13:43:54Z

I am trying to think if ways we can avoid having the forward model for each model now. Initially that seemed like th ebe

Some initial feedback. Most important question is it seems like we need to craft different logic based on different model? Can we not keep it model agnostic?

t was fine when I wrote for flux, but lumina needed multi stage preprocessing.
I am trying to think how to , but keeping a generic forward might not work very well :/
Firstcache, FirstBlock all work block level, but TeaCache is more model level.
Defo open to ideas :)

sayakpaul · 2026-03-09T11:57:12Z

@LawJarp-A sounds good. Let us know whenever you're ready.

sayakpaul · 2026-03-16T08:49:32Z

@LawJarp-A let us know.

LawJarp-A · 2026-03-16T08:51:40Z

@LawJarp-A let us know.

Yessir. Will ping here

…on mismatch

LawJarp-A · 2026-03-25T17:21:32Z

Hi @sayakpaul, here’s a quick update addressing the remaining review comments.

Fixes in the latest commit:

Fixed infinite recursion in _is_peft_adapter. It now checks PeftAdapterMixin correctly using a lazy import.
Handled a Lumina2 edge case. _maybe_reset_state_for_new_inference now re-reads num_steps at cnt=0 and resets if it changes. This covers cases where cnt wraps without a normal reset. (In standard usage, maybe_free_model_hooks() already resets correctly.)
Fixed test assertion: updated message from “must be positive” to “must be non-negative” to match the actual error.
Updated apply_teacache docstring to recommend pipe.transformer.enable_cache(config) for pipeline usage.

Responses to open comments:

StateManager fallback (hooks.py:47): The default fallback is needed since some pipelines (like Lumina2) don’t use cache_context() in their loop. Throwing an error would break them. I can add a logger.debug if you prefer visibility.
Calibration utility: Custom coefficients can already be passed via TeaCacheConfig. A calibration step like [Feat] TaylorSeer Cache #12648 would be useful—I'll open a separate issue.
Configurable rescaling: The current 4th-degree polynomial comes from the paper and works across tested models. We can revisit this if needed.
.item() and torch.compile: All operations stay in tensor form except the final .item() for branching, which is required for Python control flow. This is the minimal graph break.
Copilot comments: All addressed—boundary checks, class name matching, zero threshold validation, and Lumina2 state reset.

sayakpaul · 2026-04-15T13:31:28Z

.item() and torch.compile: All operations stay in tensor form except the final .item() for branching, which is required for Python control flow. This is the minimal graph break.

But we can handle through pure tensors no? Problem with .item() is that it becomes data-dependent and that wouldn't be known during runtime.

Sorry for the delay.

LawJarp-A · 2026-05-04T11:15:26Z

But we can handle through pure tensors no? Problem with .item() is that it becomes data-dependent and that wouldn't be known during runtime.

would torch.cond address this? it keeps the predicae as a 0-dim bool tensor and traces both branches, so no graph break, but it needs state mutation moved out of the branches (gating the residual update with torch.where), which is a non-trivial refactor across the four per-model forwards. Happy to do it if that's the direction.

The lighter alternative is matching what first_block_cache.py:133 and mag_cache.py:300 do today, wrap the decision callsite with @torch.compiler.disable so the rest of the forward still compiles.
Happy to go either way; let me know which bar you want for TeaCache.
@sayakpaul

# Conflicts: # src/diffusers/hooks/__init__.py # src/diffusers/models/cache_utils.py

LawJarp-A · 2026-06-14T17:39:10Z

@sayakpaul @DN6 any more comments to address?

TeaCache was bundled inline in the model forward and threaded through every guidance branch of the pipeline denoising loop. Per reviewer guidance, drop it: diffusers already ships FirstBlockCache via CacheMixin and a TeaCache backend is in progress (huggingface#12652), so caching should live in the model framework, not be hand-rolled here. Remove the model-side enable_teacache / teacache_params / rescale-coefficient state and the cache short-circuit in the single-stream stage, the per-condition TeaCacheParams bookkeeping in the pipeline loop, and delete the now-unused utils/teacache_util.py. TeaCache defaulted to off, so the default inference path is unchanged. Addresses reviewer comments r3470572315 / r3470616028 / r3470619173 / r3470928656. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

LawJarp-A · 2026-09-06T10:30:32Z

@sayakpaul / @DN6 gentle ping on this one.

I also had some suggestions from claude, which i want to review and implement (full disclosure)

The skip/compute decision is data-dependent and has to become a Python bool for control flow. Mark the decision helpers with torch.compiler.disable so the graph break is contained there and the surrounding forward still compiles, matching first_block_cache.py, taylorseer_cache.py and mag_cache.py, which make the same decision the same way.

LawJarp-A · 2026-09-06T10:59:03Z

@sayakpaul rebased onto main, and answered the .item() question in code.

the decision helpers now carry @torch.compiler.disable, so the data-dependent branch is contained and the surrounding forward still compiles. same as first_block_cache.py:133, taylorseer_cache.py:225 and mag_cache.py:184/:300. checked with torch._dynamo.explain: teacache contributes exactly one break, at _should_compute. the model's other breaks are aten.nonzero from its own mask handling and are there with no cache applied. torch.cond would remove that last one but needs cache state threaded through both branches, better as a follow-up across all three hooks.

branch no longer conflicts, 27 teacache tests pass, ruff clean.

one still open: the StateManager fallback in hooks.py is only there because lumina2's pipeline doesn't wrap its loop in cache_context(). i'd rather add that to the lumina2 loop and revert hooks.py, if touching a pipeline file is fine.

@AtulJoshi1206 rebase is done, thanks for the nudge.

DN6 · 2026-09-17T10:48:41Z

Hi @LawJarp-A. Sorry for the delay on our end. I think one thing about the current design that in hindsight isn't sustainable is the custom forwards defined for each model. Ensuring there is no drift between cached forwards and the ones on main will be difficult, so I think it would be best to use the Head/Block hook approach, similar to MagCache and FirstBlockCache

diffusers/src/diffusers/hooks/mag_cache.py

Line 171 in ae2e4c7

class MagCacheHeadHook(ModelHook):

It would also be good to provide a way for users to calibrate the coefficients for a new model like Magcache.

fusheng-ji · 2026-10-11T15:47:28Z

Hi @LawJarp-A, a small thing I ran into while testing this with torch.compile, in case it's useful for the hook-based rewrite @DN6 suggested.

The logger.debug(...) calls in TeaCacheHook._maybe_reset_state_for_new_inference cause graph breaks. Dynamo doesn't trace logger calls ("Logger not supported for non-export cases ... config.ignore_logger_methods"). On a18b178 with a tiny Flux model, I see 4–5 of these breaks per call when the hook sets num_steps on the first step, resets after a completed run, or resets when num_steps changes.

Guarding the calls with if not torch.compiler.is_compiling(): removes them and keeps the debug logs in eager mode:
fusheng-ji@c11067a

The commit also adds a test that compiles the Flux forward through those three state transitions and asserts there are no logger graph breaks, using torch._dynamo.utils.counters["graph_break"]. It fails on a18b178 and passes with the change.

The code will probably change with the Head/Block hook approach, so I'm not opening a separate PR. It might be worth keeping in mind for any logging in the new hooks. Happy to adapt the test once the new version is up.

GiuseppeBellamacina

This adds a TeaCache hook that skips transformer blocks and reuses the cached residual when the modulated input changes little between steps. It covers Flux, Mochi, Lumina2 and CogVideoX, and also changes StateManager to fall back to a default context.

I read through the hook and the per-model forward paths. I have three points, two about when caching actually applies and one about LoRA, and they're in the inline comments below.

GiuseppeBellamacina · 2026-10-21T19:16:07Z

+    #     handle counting internally via the per-cache-key state
+    has_both_lengths = state.cond_seq_len != state.uncond_seq_len
+    is_cond_pass = cache_key == state.cond_seq_len
+    if has_both_lengths and is_cond_pass:


For Lumina2, state.cnt only increments when has_both_lengths and is_cond_pass. Without CFG, or when cond and uncond have the same padded length, cnt stays at 0. The first-step check would then always be true, and the cache would never be used. The comment at lines 852-855 says caching decisions still work in that case. Can you confirm that, or count steps independently of CFG, for example once per forward when only one length is seen?

GiuseppeBellamacina · 2026-10-21T19:16:07Z

+            if self.config.num_inference_steps is not None:
+                new_num_steps = self.config.num_inference_steps
+            elif self.config.num_inference_steps_callback is not None:
+                new_num_steps = self.config.num_inference_steps_callback()
+            elif hasattr(module, "num_steps"):
+                new_num_steps = module.num_steps


num_steps only comes from config.num_inference_steps, the callback, or a module.num_steps attribute. I couldn't find any transformer that defines that attribute, so with the default TeaCacheConfig(), num_steps stays 0. Then the reset at line 409 never fires and the last step is never forced to compute. A second pipeline call would reuse the previous run's residual and counters. The docstring on num_inference_steps (line 252) says "Auto-detected if not provided". Could this resolve num_steps from the pipeline or scheduler, or warn when it's unset?

GiuseppeBellamacina · 2026-10-21T19:16:07Z

+        ids = torch.cat((txt, img), dim=0)
+        image_rotary_emb = module.pos_embed(ids)
+
+        joint_attention_kwargs = kwargs.get("joint_attention_kwargs")


The Flux replacement forward passes joint_attention_kwargs straight to the blocks and never extracts or applies the LoRA scale, as the original forward does through its LoRA-scale handling. With a LoRA loaded and TeaCache enabled, I think the scale would be ignored. Could you apply the same scaling here, or tell me if I'm missing where it happens?

base implement teacahce flux; follow original impl

781c7e3

LawJarp-A mentioned this pull request Nov 13, 2025

[feature] implement TeaCache #12589

Open

update cache utils with teacache

549bf97

sayakpaul requested a review from DN6 November 13, 2025 16:49

change hook to inline transformer processing instead of calling origi…

29d4ffc

…nal_forward

sayakpaul mentioned this pull request Nov 14, 2025

The Diffusers MVP 🚀 #12635

Open

LawJarp-A added 4 commits November 14, 2025 06:43

add extensive docstring (auto gen); add repr

07c6718

add param validation and error messages

a7598a1

add basic logging

d9648e5

add compatible test

59cb890

LawJarp-A marked this pull request as ready for review November 14, 2025 08:23

Merge branch 'main' into teacache-flux

5227937

LawJarp-A and others added 2 commits November 26, 2025 12:30

Merge branch 'main' into teacache-flux

296aa8d

update it to make it model agnostic

44663de

LawJarp-A added 2 commits November 26, 2025 11:08

add TeaCache hook tests and ensure cache integration passes style and…

4d34020

… quality checks

Add multi-model TeaCache support for Mochi, Lumina2, and CogVideoX wi…

9dab52f

…th auto-detection

sayakpaul reviewed Dec 8, 2025

View reviewed changes

LawJarp-A and others added 2 commits December 8, 2025 13:44

simplify model extract; use logger

4cb71d6

Merge branch 'main' into teacache-flux

ad90044

Merge branch 'main' into teacache-flux

373a799

fix _is_peft_adapter recursion, lumina2 state reset, and test asserti…

f0cd223

…on mismatch

joangava approved these changes Mar 25, 2026

View reviewed changes

sayakpaul mentioned this pull request Apr 15, 2026

[feat]: General diffusers adapter backend to run diffusion models vllm-project/vllm-omni#2724

Merged

3 of 5 tasks

Merge remote-tracking branch 'upstream/main' into teacache-flux

2841390

# Conflicts: # src/diffusers/hooks/__init__.py # src/diffusers/models/cache_utils.py

github-actions Bot added models tests hooks size/L PR with diff > 200 LOC labels May 4, 2026

joangava approved these changes Jun 4, 2026

View reviewed changes

Merge branch 'main' into teacache-flux

375bda6

srlynch1 mentioned this pull request Jun 23, 2026

docs: TeaCache ideation for issue #12589 agentdevsl/diffusers#9

Merged

yiyixuxu mentioned this pull request Jun 24, 2026

Add Boogu-Image generation, editing, and turbo pipelines #14040

Open

8 of 9 tasks

LawJarp-A added 2 commits September 6, 2026 16:14

merge main

7e82c15

DN6 linked an issue Sep 7, 2026 that may be closed by this pull request

[feature] implement TeaCache #12589

Open

GiuseppeBellamacina reviewed Oct 21, 2026

View reviewed changes

Model	Coefficients	Status
FLUX	Auto-detected	Tested
FLUX-Kontext	Auto-detected	Ready
Mochi	Auto-detected	Ready
Lumina2	Auto-detected	Ready
CogVideoX (2b/5b/1.5-5B)	Auto-detected	Ready

Sunbelt Computer Software

PL/B Language Development and Support

Conversation

LawJarp-A commented Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

What is TeaCache?

Architecture & Design

Supported Models

Benchmark Results (FLUX.1-dev)

Benchmark Results (Lumina2)

Benchmark Results (CogVideoX-2b)

Benchmark Results (Mochi)

Usage

Configuration Options

Files Changed

Before submitting

Who can review?

Uh oh!

LawJarp-A commented Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

LawJarp-A commented Nov 17, 2025

Uh oh!

sayakpaul commented Nov 23, 2025

Uh oh!

HuggingFaceDocBuilderDev commented Nov 23, 2025

Uh oh!

DN6 commented Nov 24, 2025

Uh oh!

LawJarp-A commented Nov 24, 2025

Uh oh!

LawJarp-A commented Nov 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

LawJarp-A commented Dec 2, 2025

Uh oh!

LawJarp-A commented Dec 8, 2025

Uh oh!

sayakpaul commented Dec 8, 2025

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

LawJarp-A commented Dec 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sayakpaul commented Mar 9, 2026

Uh oh!

sayakpaul commented Mar 16, 2026

Uh oh!

LawJarp-A commented Mar 16, 2026

Uh oh!

LawJarp-A commented Mar 25, 2026

Uh oh!

sayakpaul commented Apr 15, 2026

Uh oh!

LawJarp-A commented May 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

LawJarp-A commented Jun 14, 2026

Uh oh!

LawJarp-A commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

LawJarp-A commented Sep 6, 2026

Uh oh!

DN6 commented Sep 17, 2026

Uh oh!

fusheng-ji commented Oct 11, 2026

Uh oh!

GiuseppeBellamacina left a comment

Choose a reason for hiding this comment

Uh oh!

LawJarp-A commented Nov 13, 2025 •

edited

Loading

LawJarp-A commented Nov 13, 2025 •

edited

Loading

LawJarp-A commented Nov 26, 2025 •

edited

Loading

LawJarp-A commented Dec 8, 2025 •

edited

Loading

LawJarp-A commented May 4, 2026 •

edited

Loading

LawJarp-A commented Sep 6, 2026 •

edited

Loading