From 283a62fd769468fe0721cd32eccabe31fca94113 Mon Sep 17 00:00:00 2001 From: Ding Zuhao Date: Sat, 28 Mar 2026 23:24:21 +0800 Subject: [PATCH] [Feature] [HunyuanImage3] Add TeaCache support for inference acceleration (#1927) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Ding Zuhao Signed-off-by: 汪志鹏 Co-authored-by: 汪志鹏 --- docs/user_guide/diffusion/teacache.md | 1 + vllm_omni/diffusion/cache/teacache/backend.py | 30 ++++++++++ vllm_omni/diffusion/cache/teacache/config.py | 10 ++++ .../hunyuan_image_3_transformer.py | 59 +++++++++++++++---- 4 files changed, 89 insertions(+), 11 deletions(-) diff --git a/docs/user_guide/diffusion/teacache.md b/docs/user_guide/diffusion/teacache.md index 40dafeb88..c90076e20 100644 --- a/docs/user_guide/diffusion/teacache.md +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 132dcce0d..a5087fe0c 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 8568ad660..96cf3f03e 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/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, + ], } diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py index 69b39974b..d18913723 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/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: