[RFC] Clean and refactor TeaCache registry (#416)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

vllm_omni/diffusion/cache/teacache/backend.py

@@ -47,20 +47,20 @@ class TeaCacheBackend(CacheBackend):
         and applies the TeaCache hook to the transformer.
 
         Args:
-            pipeline: Diffusion pipeline instance. Extracts transformer and model_type:
+            pipeline: Diffusion pipeline instance. Extracts transformer and transformer_type:
                      - transformer: pipeline.transformer
-                     - model_type: pipeline.__class__.__name__
+                     - transformer_type: pipeline.transformer.__class__.__name__
         """
-        # Extract transformer and model_type from pipeline
+        # Extract transformer and transformer_type from pipeline
         transformer = pipeline.transformer
-        model_type = pipeline.__class__.__name__
+        transformer_type = transformer.__class__.__name__
 
-        # Create TeaCacheConfig from DiffusionCacheConfig with model_type
+        # Create TeaCacheConfig from DiffusionCacheConfig with transformer_type
         # Access parameters via attribute access: config.rel_l1_thresh
         # rel_l1_thresh already has a default value of 0.2 in DiffusionCacheConfig
         try:
             teacache_config = TeaCacheConfig(
-                model_type=model_type,
+                transformer_type=transformer_type,
                 rel_l1_thresh=self.config.rel_l1_thresh,
                 coefficients=self.config.coefficients,
             )
@@ -69,7 +69,7 @@ class TeaCacheBackend(CacheBackend):
             raise ValueError(
                 f"Invalid TeaCache configuration: {e}. "
                 f"Expected keys: rel_l1_thresh, coefficients (optional). "
-                f"model_type is automatically extracted from pipeline.__class__.__name__."
+                f"transformer_type is automatically extracted from pipeline.transformer.__class__.__name__."
             )
 
         # Apply hook to transformer
@@ -80,7 +80,7 @@ class TeaCacheBackend(CacheBackend):
 
         logger.info(
             f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
-            f"model_type={teacache_config.model_type}"
+            f"transformer_class={teacache_config.transformer_type}"
         )
 
     def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:

vllm_omni/diffusion/cache/teacache/config.py

@@ -7,17 +7,18 @@ from dataclasses import dataclass
 # These coefficients account for model-specific characteristics in how embeddings change
 # Source: TeaCache paper and ComfyUI-TeaCache empirical tuning
 _MODEL_COEFFICIENTS = {
-    # FLUX model coefficients from TeaCache paper
-    "FluxPipeline": [
+    # FLUX transformer coefficients from TeaCache paper
+    "FluxTransformer2DModel": [
         4.98651651e02,
         -2.83781631e02,
         5.58554382e01,
         -3.82021401e00,
         2.64230861e-01,
     ],
-    # Qwen-Image model coefficients from ComfyUI-TeaCache
+    # Qwen-Image transformer coefficients from ComfyUI-TeaCache
     # Tuned specifically for Qwen's dual-stream transformer architecture
-    "QwenImagePipeline": [
+    # Used for all Qwen-Image Family pipelines, in general
+    "QwenImageTransformer2DModel": [
         -4.50000000e02,
         2.80000000e02,
         -4.50000000e01,
@@ -43,15 +44,15 @@ class TeaCacheConfig:
             - 0.4: ~1.8x speedup with slight quality loss
             - 0.6: ~2.0x speedup with noticeable quality loss
         coefficients: Polynomial coefficients for rescaling L1 distance. If None, uses
-            model-specific defaults based on model_type.
-        model_type: Pipeline class name (e.g., "QwenImagePipeline", "FluxPipeline").
-            Must match OmniDiffusionConfig.model_class_name for proper extractor lookup.
-            Defaults to "QwenImagePipeline".
+            model-specific defaults based on transformer_type.
+        transformer_type: Transformer class name (e.g., "QwenImageTransformer2DModel").
+            Auto-detected from pipeline.transformer.__class__.__name__ in backend.
+            Defaults to "QwenImageTransformer2DModel".
     """
 
     rel_l1_thresh: float = 0.2
     coefficients: list[float] | None = None
-    model_type: str = "QwenImagePipeline"
+    transformer_type: str = "QwenImageTransformer2DModel"
 
     def __post_init__(self) -> None:
         """Validate and set default coefficients."""
@@ -59,8 +60,13 @@ class TeaCacheConfig:
             raise ValueError(f"rel_l1_thresh must be positive, got {self.rel_l1_thresh}")
 
         if self.coefficients is None:
-            # Use model-specific coefficients, fallback to FluxPipeline if not found
-            self.coefficients = _MODEL_COEFFICIENTS.get(self.model_type, _MODEL_COEFFICIENTS["FluxPipeline"])
+            # Use model-specific coefficients, explicitly check if the type exists or not
+            if self.transformer_type not in _MODEL_COEFFICIENTS:
+                raise KeyError(
+                    f"Cannot find coefficients for {self.transformer_type}. "
+                    f"Supported: {list(_MODEL_COEFFICIENTS.keys())}"
+                )
+            self.coefficients = _MODEL_COEFFICIENTS[self.transformer_type]
 
         if len(self.coefficients) != 5:
             raise ValueError(f"coefficients must contain exactly 5 elements, got {len(self.coefficients)}")

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -151,6 +151,7 @@ def extract_qwen_context(
     img_shapes: torch.Tensor,
     txt_seq_lens: torch.Tensor,
     guidance: torch.Tensor | None = None,
+    additional_t_cond: torch.Tensor | None = None,
     attention_kwargs: dict[str, Any] | None = None,
     **kwargs: Any,
 ) -> CacheContext:
@@ -170,6 +171,7 @@ def extract_qwen_context(
         img_shapes: Image shapes for position embedding
         txt_seq_lens: Text sequence lengths
         guidance: Optional guidance scale for CFG
+        additional_t_cond: Optional additional timestep conditioning
         attention_kwargs: Additional attention arguments
         **kwargs: Additional keyword arguments ignored by this extractor
 
@@ -193,9 +195,9 @@ def extract_qwen_context(
         guidance = guidance.to(hidden_states.dtype) * 1000
 
     temb = (
-        module.time_text_embed(timestep, hidden_states)
+        module.time_text_embed(timestep, hidden_states, additional_t_cond)
         if guidance is None
-        else module.time_text_embed(timestep, guidance, hidden_states)
+        else module.time_text_embed(timestep, guidance, hidden_states, additional_t_cond)
     )
 
     image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
@@ -254,22 +256,20 @@ def extract_qwen_context(
 
 
 # Registry for model-specific extractors
-# Key: Pipeline class name (from OmniDiffusionConfig.model_class_name)
+# Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
 #
-# Note: Use the exact pipeline class name as specified in OmniDiffusionConfig.model_class_name
-# to ensure consistent lookup without substring matching.
+# Note: Use the transformer class name as specified in pipelines as TeaCache hooks operate
+# on the transformer module and multiple pipelines can share the same transformer.
 EXTRACTOR_REGISTRY: dict[str, Callable] = {
-    "QwenImagePipeline": extract_qwen_context,
-    "QwenImageEditPipeline": extract_qwen_context,
-    "QwenImageEditPlusPipeline": extract_qwen_context,
+    "QwenImageTransformer2DModel": extract_qwen_context,
     # Future models:
-    # "FluxPipeline": extract_flux_context,
-    # "CogVideoXPipeline": extract_cogvideox_context,
+    # "FluxTransformer2DModel": extract_flux_context,
+    # "CogVideoXTransformer3DModel": extract_cogvideox_context,
 }
 
 
-def register_extractor(model_identifier: str, extractor_fn: Callable) -> None:
+def register_extractor(transformer_cls_name: str, extractor_fn: Callable) -> None:
     """
     Register a new extractor function for a model type.
 
@@ -277,7 +277,7 @@ def register_extractor(model_identifier: str, extractor_fn: Callable) -> None:
     the core TeaCache code.
 
     Args:
-        model_identifier: Model type identifier (class name or type string)
+        transformer_cls_name: Transformer model type identifier (class name or type string)
         extractor_fn: Function with signature (module, *args, **kwargs) -> CacheContext
 
     Example:
@@ -299,19 +299,19 @@ def register_extractor(model_identifier: str, extractor_fn: Callable) -> None:
         ...     return CacheContext(modulated, hidden_states, None, temb, run_blocks, postprocess)
         >>> register_extractor("FluxTransformer2DModel", extract_flux_context)
     """
-    EXTRACTOR_REGISTRY[model_identifier] = extractor_fn
+    EXTRACTOR_REGISTRY[transformer_cls_name] = extractor_fn
 
 
-def get_extractor(model_type: str) -> Callable:
+def get_extractor(transformer_cls_name: str) -> Callable:
     """
-    Get extractor function for given model type.
+    Get extractor function for given transformer class.
 
-    This function looks up the extractor based on the exact model_type string,
-    which should match the pipeline class name from OmniDiffusionConfig.model_class_name.
+    This function looks up the extractor based on the exact transformer_cls_name string,
+    which should match the transformer type in the pipeline (i.e., pipeline.transformer.__class__.__name__).
 
     Args:
-        model_type: Model type string (e.g., "QwenImagePipeline", "FluxPipeline")
-                   Must exactly match a key in EXTRACTOR_REGISTRY.
+        transformer_cls_name: Transformer class name (e.g., "QwenImageTransformer2DModel")
+                                Must exactly match a key in EXTRACTOR_REGISTRY.
 
     Returns:
         Extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -320,18 +320,18 @@ def get_extractor(model_type: str) -> Callable:
         ValueError: If model type not found in registry
 
     Example:
-        >>> # Get extractor for Qwen model
-        >>> extractor = get_extractor("QwenImagePipeline")
+        >>> # Get extractor for QwenImageTransformer2DModel
+        >>> extractor = get_extractor("QwenImageTransformer2DModel")
         >>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...)
     """
     # Direct lookup - no substring matching
-    if model_type in EXTRACTOR_REGISTRY:
-        return EXTRACTOR_REGISTRY[model_type]
+    if transformer_cls_name in EXTRACTOR_REGISTRY:
+        return EXTRACTOR_REGISTRY[transformer_cls_name]
 
     # No match found
     available_types = list(EXTRACTOR_REGISTRY.keys())
     raise ValueError(
-        f"Unknown model type: '{model_type}'. "
+        f"Unknown model type: '{transformer_cls_name}'. "
         f"Available types: {available_types}\n"
         f"To add support for a new model, use register_extractor() or add to EXTRACTOR_REGISTRY."
     )

vllm_omni/diffusion/cache/teacache/hook.py

@@ -62,7 +62,7 @@ class TeaCacheHook(ModelHook):
 
     def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
         """
-        Initialize hook with extractor from config model type.
+        Initialize hook with extractor from config transformer model type.
 
         Args:
             module: The module to initialize the hook for.
@@ -70,9 +70,9 @@ class TeaCacheHook(ModelHook):
         Returns:
             The initialized module.
         """
-        # Get extractor function based on model_type from config
-        # model_type should be the pipeline class name (e.g., "QwenImagePipeline")
-        self.extractor_fn = get_extractor(self.config.model_type)
+        # Get extractor function based on transformer_type from config
+        # transformer_type is the transformer class name (e.g., "QwenImageTransformer2DModel")
+        self.extractor_fn = get_extractor(self.config.transformer_type)
 
         # Set default context
         self.state_manager.set_context("teacache")
@@ -244,9 +244,13 @@ def apply_teacache_hook(module: torch.nn.Module, config: TeaCacheConfig) -> None
         config: TeaCacheConfig specifying caching parameters
 
     Example:
-        >>> config = TeaCacheConfig(rel_l1_thresh=0.2, model_type="Qwen")
+        >>> config = TeaCacheConfig(
+        ...     rel_l1_thresh=0.2,
+        ...     transformer_type="QwenImageTransformer2DModel"
+        ... )
         >>> apply_teacache_hook(transformer, config)
-        >>> # Model now uses TeaCache automatically, no code changes needed!
+        >>> # Transformer bound to the pipeline now uses TeaCache automatically,
+        ... # no code changes needed!
     """
     registry = HookRegistry.get_or_create(module)
     hook = TeaCacheHook(config)