[Feature] [HunyuanImage3] Add TeaCache support for inference acceleration (#1927)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

vllm_omni/diffusion/cache/teacache/backend.py

@@ -20,6 +20,24 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig
 logger = init_logger(__name__)
 
 
+def enable_hunyuan_image3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """
+    Enable TeaCache for HunyuanImage3 model.
+
+    HunyuanImage3 uses a GPT-based architecture with KV cache, which is incompatible
+    with the standard hook-based TeaCache approach. Instead, we store the TeaCacheConfig
+    on the pipeline so the denoising loop can implement caching directly.
+    """
+    teacache_config = TeaCacheConfig(
+        transformer_type="HunyuanImage3Pipeline",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    pipeline._tea_cache_config = teacache_config
+
+    logger.info(f"TeaCache enabled for HunyuanImage3 with rel_l1_thresh={teacache_config.rel_l1_thresh}")
+
+
 def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """
     Enable TeaCache for Bagel model.
@@ -70,6 +88,7 @@ def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) ->
 CUSTOM_TEACACHE_ENABLERS = {
     "BagelPipeline": enable_bagel_teacache,
     "Flux2KleinPipeline": enable_flux2_klein_teacache,
+    "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
 }
 
 
@@ -157,6 +176,17 @@ class TeaCacheBackend(CacheBackend):
                                 Currently not used by TeaCache but accepted for interface consistency.
             verbose: Whether to log refresh operations (default: True)
         """
+        # HunyuanImage3: tea cache state is managed inside the denoising loop,
+        # so refresh is a no-op (state is re-initialized every __call__).
+        if (
+            hasattr(pipeline, "_tea_cache_config")
+            and isinstance(pipeline._tea_cache_config, TeaCacheConfig)
+            and pipeline.__class__.__name__ == "HunyuanImage3Pipeline"
+        ):
+            if verbose:
+                logger.debug(f"TeaCache state refreshed for HunyuanImage3 (num_inference_steps={num_inference_steps})")
+            return
+
         # Extract transformer from pipeline
         transformer = pipeline.transformer
 

vllm_omni/diffusion/cache/teacache/config.py

@@ -54,6 +54,16 @@ _MODEL_COEFFICIENTS = {
         -12.281286412689623,
         1.0733905006198015,
     ],
+    # HunyuanImage3 pipeline coefficients
+    # Calibrated via polyfit on 3920 data points (80 prompts × 49 steps).
+    # Maps time_embed rel_l1 (range ~0.12-0.54) to output rel_l1 (range ~0.01-0.53).
+    "HunyuanImage3Pipeline": [
+        1.04117826e02,
+        -1.26848482e02,
+        5.68168652e01,
+        -1.04182570e01,
+        6.78098549e-01,
+    ],
 }