[Bugfix] Fix cache-dit for single-transformer Wan2.2 models(eg. Wan2.2-TI2V-5B) (#1392)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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 |
 

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