Implement TeaCache

src/diffusers/hooks/__init__.py

@@ -20,9 +20,10 @@
     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 .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
@@ -133,7 +136,7 @@ def reset_state(self, module: torch.nn.Module):
         return module
 
     def _set_context(self, module: torch.nn.Module, name: str) -> None:
-        # Iterate over all attributes of the hook to see if any of them have the type `StateManager`. If so, call `set_context` on them.
+        """Set context on all StateManager attributes of this hook."""
         for attr_name in dir(self):
             attr = getattr(self, attr_name)
             if isinstance(attr, StateManager):
@@ -142,22 +145,12 @@ def _set_context(self, module: torch.nn.Module, name: str) -> None:
 
 
 class HookFunctionReference:
-    def __init__(self) -> None:
-        """A container class that maintains mutable references to forward pass functions in a hook chain.
-
-        Its mutable nature allows the hook system to modify the execution chain dynamically without rebuilding the
-        entire forward pass structure.
+    """Mutable container for forward pass function references in a hook chain.
 
-        Attributes:
-            pre_forward: A callable that processes inputs before the main forward pass.
-            post_forward: A callable that processes outputs after the main forward pass.
-            forward: The current forward function in the hook chain.
-            original_forward: The original forward function, stored when a hook provides a custom new_forward.
+    Enables dynamic modification of the execution chain without rebuilding.
+    """
 
-        The class enables hook removal by allowing updates to the forward chain through reference modification rather
-        than requiring reconstruction of the entire chain. When a hook is removed, only the relevant references need to
-        be updated, preserving the execution order of the remaining hooks.
-        """
+    def __init__(self) -> None:
         self.pre_forward = None
         self.post_forward = None
         self.forward = None