[feature] implement TeaCache

[sayakpaul]

Currently, we support PAB and FasterCache. More details: https://huggingface.co/docs/diffusers/main/en/optimization/cache.

It'd be cool to support Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model as it's quite popular in the community. Code is here:
https://github.com/ali-vilab/TeaCache

Before proceeding with a PR, please propose a design first in this thread. We can take it from there.

[LawJarp-A]

Hey @sayakpaul
I want to take this up, i have a few questions tho

• any preferred tests/datasets to test this on?
• it's model-specific due to different timestep embedding patterns across architectures. should I proto with CogVideoX first?

[LawJarp-A]

Additionally this is my thought process (I want to refine it still) @sayakpaul

Instead of monkey-patching model forward methods (like original TeaCache), use hooks:
-> Create a TeaCacheHook that uses new_forward() to replace entire transformer block forward methods. The hook intercepts the block before any computation starts, extracts timestep embeddings, makes the skip/compute decision, and either returns cached residuals or runs the full computation.
-> Each supported model gets an extractor class that knows how to pull timestep embeddings from that model's specific architecture. CogVideoX extracts from time_embedding(), FLUX from time_text_embed(),

State management
-> Use diffusers existing BaseState system to track accumulated distance, previous embeddings, and cached residuals across inference steps (need to refine this a bit)

Add TeaCacheConfig to the existing cache configuration system in CacheMixin. This ensures it works seamlessly

Note: I am exploring the repository more as I work on figuring out how best to integrate this. I took some of time to test the current caching play around with it. I think this can be integrated well. I'm open to any suggestions and corrections

Thanks!

[yiyixuxu]

hi @LawJarp-A
Thanks for your interest!
feel free to open a PR and we can start working on it from there

[sayakpaul]

Yes, we can prototype with a model of choice. In case the original repository already provides results for a fairly popular model like Flux, let's prototype with that instead. If not, we can start with CogVideoX, as you mentioned.

[LawJarp-A]

Awesome, will prototype on flux first then. Will keep you guys updated on the PR
@sayakpaul @yiyixuxu
Thanks :)

[LawJarp-A]

Working on this PR: #12652
As of now, writing it for just one model (FLUX), after testing for it, I want to write it with better strcuture to suuport future additions, something like below

  src/diffusers/hooks/teacache/
  ├── __init__.py              
  ├── base.py                  # TeaCacheConfig, TeaCacheHook base classes
  ├── extractors.py            # All model-specific extractors
  ├── flux.py                  
  ├── cogvideox.py         
and so on...

[sayakpaul]

Hey,

Thanks for your willingness to contribute. Could we try to follow the hooks-based implementation for this (we use it for FasterCache)?

[LawJarp-A]

I am using hook based only now. Is there any issue in the PR?
(Not completed yet, still a draft PR)
(Above structure was something I was thinking out loud for the future expansion)

[sayakpaul]

Cool, we will check the PR first.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.