Implement TeaCache

src/diffusers/hooks/__init__.py

@@ -20,10 +20,11 @@
     from .faster_cache import FasterCacheConfig, apply_faster_cache
     from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache
     from .group_offloading import apply_group_offloading
-    from .hooks import HookRegistry, ModelHook
+    from .hooks import HookRegistry, ModelHook, StateManager
     from .layer_skip import LayerSkipConfig, apply_layer_skip
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
     from .mag_cache import MagCacheConfig, apply_mag_cache
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
     from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache
+    from .teacache import TeaCacheConfig, apply_teacache

src/diffusers/hooks/hooks.py

@@ -40,11 +40,14 @@ def __init__(self, state_cls: BaseState, init_args=None, init_kwargs=None):
         self._current_context = None
 
     def get_state(self):
-        if self._current_context is None:
-            raise ValueError("No context is set. Please set a context before retrieving the state.")
-        if self._current_context not in self._state_cache.keys():
-            self._state_cache[self._current_context] = self._state_cls(*self._init_args, **self._init_kwargs)
-        return self._state_cache[self._current_context]
+        context = self._current_context
+        if context is None:
+            # Fallback to default context for backward compatibility with
+            # pipelines that don't call cache_context()
+            context = "_default"
+        if context not in self._state_cache:
+            self._state_cache[context] = self._state_cls(*self._init_args, **self._init_kwargs)
+        return self._state_cache[context]
 
     def set_context(self, name: str) -> None:
         self._current_context = name