[refactor] teacache refactor (#395)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

vllm_omni/diffusion/cache/base.py

@@ -23,6 +23,8 @@ All backends implement the same interface:
 from abc import ABC, abstractmethod
 from typing import Any
 
+import torch.nn as nn
+
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
 
@@ -98,3 +100,13 @@ class CacheBackend(ABC):
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(config={self.config})"
+
+
+class CachedTransformer(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.do_true_cfg = False
+
+    def __init_subclass__(cls, enable_separate_cfg: bool = True, **kwargs):
+        cls.enable_separate_cfg = enable_separate_cfg
+        super().__init_subclass__(**kwargs)

vllm_omni/diffusion/cache/teacache/backend.py

@@ -75,10 +75,6 @@ class TeaCacheBackend(CacheBackend):
         # Apply hook to transformer
         apply_teacache_hook(transformer, teacache_config)
 
-        # Store backend reference on pipeline for pipeline's internal use
-        # (e.g., checking if cache is enabled for cache_branch parameter)
-        pipeline._cache_backend = self
-
         # Mark as enabled
         self.enabled = True
 

vllm_omni/diffusion/cache/teacache/hook.py

@@ -58,6 +58,7 @@ class TeaCacheHook(ModelHook):
         self.rescale_func = np.poly1d(config.coefficients)
         self.state_manager = StateManager(TeaCacheState)
         self.extractor_fn = None
+        self._forward_cnt = 0
 
     def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
         """
@@ -112,7 +113,11 @@ class TeaCacheHook(ModelHook):
         # GENERIC CACHING LOGIC (works for all models)
         # ============================================================================
         # Set context based on CFG branch for separate state tracking
-        cache_branch = kwargs.get("cache_branch", "default")
+        if module.do_true_cfg and self._forward_cnt % 2 == 1:
+            cache_branch = "negative"
+        else:
+            cache_branch = "positive"
+
         context_name = f"teacache_{cache_branch}"
         self.state_manager.set_context(context_name)
         state = self.state_manager.get_state()
@@ -155,6 +160,7 @@ class TeaCacheHook(ModelHook):
         # Update state
         state.previous_modulated_input = ctx.modulated_input.detach()
         state.cnt += 1
+        self._forward_cnt += 1
 
         # ============================================================================
         # POSTPROCESSING (model-specific, via callable)
@@ -221,6 +227,7 @@ class TeaCacheHook(ModelHook):
             The module with reset state.
         """
         self.state_manager.reset()
+        self._forward_cnt = 0
         return module