[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>

docs/user_guide/diffusion/teacache.md

@@ -125,6 +125,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
+| `HunyuanImage3Pipeline` | HunyuanImage3 | `tencent/HunyuanImage-3.0-Instruct` |
 
 ### VideoGen
 

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,
+    ],
 }
 
 

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py

@@ -2667,6 +2667,15 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
 
+        # ---- TeaCache initialization ----
+        tea_cache_config = getattr(self.model, "_tea_cache_config", None)
+        if tea_cache_config is not None:
+            tc_rescale = np.poly1d(tea_cache_config.coefficients)
+            tc_acc_dist = 0.0
+            tc_prev_mod = None
+            tc_prev_pred = None
+            tc_cnt = 0
+
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 # expand the latents if we are doing classifier free guidance
@@ -2675,17 +2684,42 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
 
                 t_expand = t.repeat(latent_model_input.shape[0])
 
-                model_inputs = self.model.prepare_inputs_for_generation(
-                    input_ids,
-                    images=latent_model_input,
-                    timestep=t_expand,
-                    **model_kwargs,
-                )
+                # ---- TeaCache: decide whether to compute or reuse ----
+                should_compute = True
+                if tea_cache_config is not None:
+                    with torch.no_grad():
+                        # Use timestep embedding as the modulated input for cache decision
+                        cur_mod = self.model.time_embed(t.unsqueeze(0))
+                    if tc_cnt > 0 and tc_prev_mod is not None:
+                        rel_dist = (
+                            ((cur_mod - tc_prev_mod).abs().mean() / (tc_prev_mod.abs().mean() + 1e-8)).cpu().item()
+                        )
+                        rescaled = float(tc_rescale(rel_dist))
+                        tc_acc_dist += abs(rescaled)
+                        if tc_acc_dist < tea_cache_config.rel_l1_thresh:
+                            should_compute = False
+                        else:
+                            tc_acc_dist = 0.0
+                    tc_prev_mod = cur_mod.detach()
 
-                with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-                    model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
-                    pred = model_output["diffusion_prediction"]
-                pred = pred.to(dtype=torch.float32)
+                if should_compute:
+                    model_inputs = self.model.prepare_inputs_for_generation(
+                        input_ids,
+                        images=latent_model_input,
+                        timestep=t_expand,
+                        **model_kwargs,
+                    )
+
+                    with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+                        model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
+                        pred = model_output["diffusion_prediction"]
+                    pred = pred.to(dtype=torch.float32)
+
+                    if tea_cache_config is not None:
+                        tc_prev_pred = pred.clone()
+                else:
+                    # TeaCache fast path: reuse previous prediction
+                    pred = tc_prev_pred
 
                 # perform guidance
                 if self.do_classifier_free_guidance:
@@ -2695,7 +2729,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                 # compute the previous noisy sample x_t -> x_t-1
                 latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]
 
-                if i != len(timesteps) - 1:
+                if i != len(timesteps) - 1 and should_compute:
                     model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                         model_output,
                         model_kwargs,
@@ -2709,6 +2743,9 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     model_kwargs["query_lens"] = query_lens
                     model_kwargs["seq_lens"] = seq_lens
 
+                if tea_cache_config is not None:
+                    tc_cnt += 1
+
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
                     for k in callback_on_step_end_tensor_inputs: