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

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

docs/user_guide/acceleration/teacache.md

@@ -111,6 +111,7 @@ th {
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
+| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 
 ### VideoGen
 

tests/diffusion/cache/test_cache_backends.py

@@ -138,6 +138,7 @@ class TestTeaCacheBackend:
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "QwenImagePipeline"
         mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "QwenImageTransformer2DModel"
         mock_pipeline.transformer = mock_transformer
 
         config = DiffusionCacheConfig(rel_l1_thresh=0.3)
@@ -154,6 +155,7 @@ class TestTeaCacheBackend:
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "QwenImagePipeline"
         mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "QwenImageTransformer2DModel"
         mock_pipeline.transformer = mock_transformer
 
         config = DiffusionCacheConfig(rel_l1_thresh=0.3, coefficients=[1.0, 0.5, 0.2, 0.1, 0.05])
@@ -169,6 +171,7 @@ class TestTeaCacheBackend:
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "QwenImagePipeline"
         mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "QwenImageTransformer2DModel"
         mock_pipeline.transformer = mock_transformer
 
         # Mock hook registry

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)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -230,9 +230,6 @@ class QwenImageEditPipeline(
             model, subfolder="processor", local_files_only=local_files_only
         )
 
-        # Initialize cache backend to None (will be set by worker if needed)
-        self._cache_backend = None
-
         self.stage = None
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -194,7 +194,6 @@ class QwenImageEditPlusPipeline(
         )
 
         self.stage = None
-        self._cache_backend = None
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
         self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -264,8 +264,6 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         )
         self.default_sample_size = 128
 
-        # VLLM-omni related, Initialize cache backend to None (will be set by worker if needed)
-        self._cache_backend = None
         self.stage = None
 
     def check_inputs(
@@ -570,42 +568,36 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
 
             # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
-            # Forward pass for positive prompt (or unconditional if no CFG)
-            # cache_branch is passed to hook for CFG-aware state management
-            transformer_kwargs = {
-                "hidden_states": latent_model_input,
-                "timestep": timestep / 1000,
-                "guidance": guidance,
-                "encoder_hidden_states_mask": prompt_embeds_mask,
-                "encoder_hidden_states": prompt_embeds,
-                "img_shapes": img_shapes,
-                "txt_seq_lens": txt_seq_lens,
-                "attention_kwargs": self.attention_kwargs,
-                "additional_t_cond": is_rgb,
-                "return_dict": False,
-            }
-            if self._cache_backend is not None:
-                transformer_kwargs["cache_branch"] = "positive"
 
-            noise_pred = self.transformer(**transformer_kwargs)[0]
+            self.transformer.do_true_cfg = do_true_cfg  # used in teacache hook
+            # Forward pass for positive prompt (or unconditional if no CFG)
+            noise_pred = self.transformer(
+                hidden_states=latent_model_input,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=prompt_embeds_mask,
+                encoder_hidden_states=prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                additional_t_cond=is_rgb,
+                return_dict=False,
+            )[0]
             noise_pred = noise_pred[:, : latents.size(1)]
 
             if do_true_cfg:
-                neg_transformer_kwargs = {
-                    "hidden_states": latent_model_input,
-                    "timestep": timestep / 1000,
-                    "guidance": guidance,
-                    "encoder_hidden_states_mask": negative_prompt_embeds_mask,
-                    "encoder_hidden_states": negative_prompt_embeds,
-                    "img_shapes": img_shapes,
-                    "txt_seq_lens": negative_txt_seq_lens,
-                    "attention_kwargs": self.attention_kwargs,
-                    "additional_t_cond": is_rgb,
-                    "return_dict": False,
-                }
-                if self._cache_backend is not None:
-                    neg_transformer_kwargs["cache_branch"] = "negative"
-                neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
+                neg_noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    img_shapes=img_shapes,
+                    txt_seq_lens=negative_txt_seq_lens,
+                    attention_kwargs=self.attention_kwargs,
+                    additional_t_cond=is_rgb,
+                    return_dict=False,
+                )[0]
 
                 neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
                 comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)