From d8b35035ec0be860d2927f4b0c49308dc1124791 Mon Sep 17 00:00:00 2001 From: Lancer Date: Sun, 22 Mar 2026 17:57:28 +0800 Subject: [PATCH] [Bugfix] Fix cache-dit for single-transformer Wan2.2 models(eg. Wan2.2-TI2V-5B) (#1392) Signed-off-by: Lancer Co-authored-by: Hongsheng Liu --- docs/design/feature/cache_dit.md | 8 ++-- .../diffusion/cache/cache_dit_backend.py | 43 +++++++++++++++++-- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/docs/design/feature/cache_dit.md b/docs/design/feature/cache_dit.md index 3112bb591..a53959956 100644 --- a/docs/design/feature/cache_dit.md +++ b/docs/design/feature/cache_dit.md @@ -113,13 +113,13 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No Some models require custom handling: -- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2) +- **Single or dual-transformer:** Models that may use one or two transformers (e.g., Wan2.2) - **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`) - **Special forward patterns:** Models with non-standard block execution patterns -### Example 1: Dual-Transformer Model (Wan2.2) +### Example 1: Single or Dual-Transformer Model (Wan2.2) -Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps. +Wan2.2 can use either a single transformer or two transformers (one for high-noise steps and one for low-noise steps). The implementation automatically detects the mode based on the presence of `transformer_2`. **Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations. @@ -267,7 +267,7 @@ Complete examples in the codebase: | Model | Path | Pattern | Notes | |-------|------|---------|-------| | **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic | -| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers | +| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Single or dual-transformer | Auto-detects mode based on transformer_2 presence | | **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer | | **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture | diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 32c0c8ba0..831d55f98 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -65,10 +65,10 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Wan2.2 dual-transformer architecture. + """Enable cache-dit for Wan2.2 single or dual-transformer architecture. - Wan2.2 uses two transformers (transformer and transformer_2) that need - to be enabled together using BlockAdapter. + Wan2.2 can use single or dual transformers (transformer and transformer_2) that need + to be enabled using BlockAdapter. Args: pipeline: The Wan2.2 pipeline instance. @@ -78,6 +78,43 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], A refresh function that can be called to update cache context with new num_inference_steps. """ + if getattr(pipeline, "transformer_2", None) is None: + logger.info("transformer_2 not found, enabling cache-dit for single transformer mode") + db_cache_config = _build_db_cache_config(cache_config) + cache_dit.enable_cache( + BlockAdapter( + transformer=pipeline.transformer, + blocks=[pipeline.transformer.blocks], + forward_pattern=[ForwardPattern.Pattern_2], + params_modifiers=[ + ParamsModifier(cache_config=db_cache_config), + ], + has_separate_cfg=True, + ), + cache_config=db_cache_config, + ) + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for single transformer.""" + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context( + pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose + ) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + cache_dit.enable_cache( BlockAdapter( transformer=[