Fix TeaCache test threshold and single-block step advancement on skip

- Use float('inf') rel_l1_thresh in TeaCacheConfigMixin so _test_cache_inference
  reliably skips on the second pass despite large randn perturbations
- Advance step_index from TeaCacheHeadHook when skipping on single-block models
  where the co-located tail block hook is bypassed

Author: cursoragent

src/diffusers/hooks/teacache.py

@@ -132,11 +132,13 @@ def __init__(
         config: TeaCacheConfig,
         extract_modulated_input: Callable,
         coefficients: List[float],
+        advance_step_on_skip: bool = False,
     ):
         self.state_manager = state_manager
         self.config = config
         self.extract_modulated_input = extract_modulated_input
         self.coefficients = coefficients
+        self.advance_step_on_skip = advance_step_on_skip
         self._metadata = None
 
     def initialize_hook(self, module):
@@ -180,6 +182,9 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         if not should_compute:
             logger.debug(f"TeaCache: Skipping step {state.step_index}")
 
+            if self.advance_step_on_skip:
+                self._advance_step(state)
+
             output = hidden_states
             res = state.previous_residual
 
@@ -230,6 +235,14 @@ def reset_state(self, module):
         self.state_manager.reset()
         return module
 
+    def _advance_step(self, state: TeaCacheState):
+        state.step_index += 1
+        if state.step_index >= self.config.num_inference_steps:
+            state.step_index = 0
+            state.accumulated_distance = 0.0
+            state.previous_residual = None
+            state.previous_modulated_input = None
+
 
 class TeaCacheBlockHook(ModelHook):
     def __init__(self, state_manager: StateManager, is_tail: bool = False, config: TeaCacheConfig = None):
@@ -350,7 +363,9 @@ def apply_teacache(module: torch.nn.Module, config: TeaCacheConfig) -> None:
         name, block = remaining_blocks[0]
         logger.info(f"TeaCache: Applying Head+Tail Hooks to single block '{name}'")
         _apply_teacache_block_hook(block, state_manager, config, is_tail=True)
-        _apply_teacache_head_hook(block, state_manager, config, extract_modulated_input, coefficients)
+        _apply_teacache_head_hook(
+            block, state_manager, config, extract_modulated_input, coefficients, advance_step_on_skip=True
+        )
         return
 
     head_block_name, head_block = remaining_blocks.pop(0)
@@ -372,13 +387,16 @@ def _apply_teacache_head_hook(
     config: TeaCacheConfig,
     extract_modulated_input: Callable,
     coefficients: List[float],
+    advance_step_on_skip: bool = False,
 ) -> None:
     registry = HookRegistry.check_if_exists_or_initialize(block)
 
     if registry.get_hook(_TEACACHE_LEADER_BLOCK_HOOK) is not None:
         registry.remove_hook(_TEACACHE_LEADER_BLOCK_HOOK)
 
-    hook = TeaCacheHeadHook(state_manager, config, extract_modulated_input, coefficients)
+    hook = TeaCacheHeadHook(
+        state_manager, config, extract_modulated_input, coefficients, advance_step_on_skip=advance_step_on_skip
+    )
     registry.register_hook(hook, _TEACACHE_LEADER_BLOCK_HOOK)
 
 

tests/models/testing_utils/cache.py

@@ -643,10 +643,11 @@ class TeaCacheConfigMixin:
     """
 
     # Default TeaCache config - can be overridden by subclasses.
-    # Uses num_inference_steps=4 so interior steps can be skipped during _test_cache_inference.
+    # Uses num_inference_steps=4 and an infinite rel_l1_thresh so the second
+    # inference step is always skipped, which is required by _test_cache_inference.
     TEA_CACHE_CONFIG = {
         "num_inference_steps": 4,
-        "rel_l1_thresh": 100.0,
+        "rel_l1_thresh": float("inf"),
     }
 
     def _get_cache_config(self):