[BugFix] Fix Wan2.2-VACE-Fun cache and serving lifecycle (#4667)

Signed-off-by: Ianniu <985683179yzy@gmail.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: Ianniu <985683179yzy@gmail.com>

docs/models/supported_models.md

@@ -39,6 +39,7 @@ th {
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |

recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md

@@ -0,0 +1,109 @@
+# Wan2.2 VACE-Fun Diffusers
+
+> All-in-one video creation and editing with Wan2.2 VACE-Fun A14B in diffusers format.
+
+## Summary
+
+- Vendor: Pyros13
+- Model: `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers`
+- Base model: `alibaba-pai/Wan2.2-VACE-Fun-A14B`
+- Task: VACE video generation and editing (T2V / I2V / R2V / V2V / inpainting)
+- Mode: Offline inference and OpenAI-compatible online serving
+- Maintainer: Community
+
+## When To Use This Recipe
+
+Use this recipe to run Wan2.2-VACE-Fun-A14B through the standard diffusers-style
+`Wan22VACEPipeline` path in vLLM-Omni. The checkpoint is a diffusers conversion of
+`alibaba-pai/Wan2.2-VACE-Fun-A14B`, with `model_index.json`, `transformer/`,
+`transformer_2/`, `text_encoder/`, `tokenizer/`, `vae/`, and `scheduler/`
+subfolders. Because the model already uses the expected diffusers layout, vLLM-Omni
+does not need any original-format key remapping or `.pth` component loaders.
+
+The model is a two-expert architecture. The high-noise `transformer` and low-noise
+`transformer_2` are switched by the Wan2.2 boundary during denoising.
+
+## References
+
+- Model card: <https://huggingface.co/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers>
+- Base model card: <https://huggingface.co/alibaba-pai/Wan2.2-VACE-Fun-A14B>
+- Related example: [`examples/offline_inference/vace/vace_video_generation.md`](../../examples/offline_inference/vace/vace_video_generation.md)
+- Related issue: [vllm-project/vllm-omni#4206](https://github.com/vllm-project/vllm-omni/issues/4206)
+- PR: [vllm-project/vllm-omni#4667](https://github.com/vllm-project/vllm-omni/pull/4667)
+
+## Offline Inference
+
+```bash
+python examples/offline_inference/vace/vace_video_generation.py \
+  --model Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers \
+  --mode i2v --image ./i2v_input.JPG \
+  --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard" \
+  --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \
+  --seed 42 --guidance-scale 5.0 \
+  --vae-use-tiling --enforce-eager \
+  --output vace_fun_i2v.mp4
+```
+
+See [`examples/offline_inference/vace/vace_video_generation.md`](../../examples/offline_inference/vace/vace_video_generation.md)
+for other VACE modes and their input flags.
+
+## Online Serving
+
+### Server
+
+```bash
+vllm serve Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers --omni \
+  --model-class-name Wan22VACEPipeline \
+  --vae-use-tiling \
+  --enforce-eager \
+  --port 8091
+```
+
+For multi-GPU sequence-parallel serving, add the same parallelism flags you would
+use offline, for example `--ulysses-degree 4`. The explicit `--model-class-name`
+is optional for this diffusers checkpoint, but keeping it in examples makes the
+selected pipeline obvious.
+
+### Client
+
+```bash
+no_proxy=127.0.0.1 \
+curl -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -F "prompt=Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard" \
+  -F "input_reference=@/absolute/path/to/i2v_input.JPG" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=81" \
+  -F "fps=16" \
+  -F "num_inference_steps=30" \
+  -F "guidance_scale=5.0" \
+  --output vace_fun_i2v_serve.mp4
+```
+
+For video-to-video or inpainting-style requests, upload the reference video with
+`input_reference=@/absolute/path/to/input.mp4` or use `video_reference` with a URL
+or JSON-safe data URI.
+
+## Cache-DiT / TaylorSeer Notes
+
+`Wan22VACEPipeline` uses the shared Wan2.2 Cache-DiT enabler for both the
+high-noise and low-noise transformers. The adapter intentionally wraps only the
+main denoising `blocks`. VACE `vace_blocks` form a conditioning branch: each step
+they combine the current latent with the VACE context to produce hints that are
+then injected into selected main blocks. Recomputing that branch preserves the
+control signal; caching the main backbone still gives the acceleration target.
+
+## Validation
+
+Validated on NVIDIA B300 with `mkt_part_000.mp4`, 736x1280, 61 frames, 30
+inference steps, `seed=1`, and `boundary_ratio=0.875`.
+
+The Pyros13 diffusers checkpoint produced a byte-identical MP4 to the
+`alibaba-pai/Wan2.2-VACE-Fun-A14B` original-format checkpoint under the same
+vLLM-Omni pipeline and inputs:
+
+- SHA256: `2e2db74a612a009b0d6ad6d739e039883f0de84ce5cdf2676324daf78fb05ef8`
+- Frame shape: 61 x 736 x 1280
+- MAE / RMSE / max pixel diff: 0 / 0 / 0
+- PSNR: infinite
+- Mean SSIM: 1.0

tests/diffusion/cache/test_cache_dit.py

@@ -42,6 +42,10 @@ SEPARATE_CFG_TRANSFORMERS = [
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
+def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
+    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_backend.enable_cache_for_wan22
+
+
 @pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
 def test_cache_dit_configs_have_separate_cfg(transformer_model):
     """Check models with separate CFG set has_separate_cfg=True in their configs."""
@@ -62,6 +66,34 @@ def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_blo
     assert adapter_kwargs["has_separate_cfg"] is True
 
 
+@pytest.mark.parametrize("has_transformer_2", [False, True])
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+def test_wan22_custom_enabler_passes_taylorseer_calibrator(
+    mock_cache_dit,
+    mock_block_adapter,
+    has_transformer_2,
+):
+    mock_pipeline = Mock()
+    mock_pipeline.transformer.blocks = [Mock()]
+    if has_transformer_2:
+        mock_pipeline.transformer_2.blocks = [Mock()]
+    else:
+        mock_pipeline.transformer_2 = None
+    cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1)
+
+    cd_backend.enable_cache_for_wan22(mock_pipeline, cache_config)
+
+    enable_cache_kwargs = mock_cache_dit.enable_cache.call_args.kwargs
+    calibrator_config = enable_cache_kwargs["calibrator_config"]
+    assert calibrator_config is not None
+    assert calibrator_config.taylorseer_order == 1
+
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    for modifier in adapter_kwargs["params_modifiers"]:
+        assert modifier._context_kwargs["calibrator_config"] is calibrator_config
+
+
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
 def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -14,6 +14,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
     create_vace_transformer_from_config,
     get_wan22_vace_pre_process_func,
 )
+from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -93,6 +94,10 @@ def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch
     }
 
 
+def test_vace_transformer_layerwise_offloads_conditioning_blocks() -> None:
+    assert WanVACETransformer3DModel._layerwise_offload_blocks_attrs == ["vace_blocks", "blocks"]
+
+
 def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None:
     pipeline = _make_vace_pipeline()
     mask = torch.ones(1, 3, 5, 16, 16)

tests/entrypoints/openai_api/test_video_server.py

@@ -11,6 +11,7 @@ import json
 import os
 import threading
 import time
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 
 import numpy as np
@@ -109,6 +110,14 @@ class BlockingVideoHandler:
             raise
 
 
+class FakeServerSocket:
+    def __init__(self):
+        self.closed = False
+
+    def close(self):
+        self.closed = True
+
+
 @pytest.fixture(autouse=True)
 def isolated_video_backends(tmp_path, monkeypatch):
     """Use isolated in-memory metadata and local storage for each test."""
@@ -121,6 +130,91 @@ def isolated_video_backends(tmp_path, monkeypatch):
     return store, tasks, storage
 
 
+@pytest.mark.asyncio
+async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch):
+    events: list[str] = []
+    serve_started = asyncio.Event()
+    http_shutdown = asyncio.Event()
+    engine_context_exited = asyncio.Event()
+    sock = FakeServerSocket()
+
+    class FakeEngine:
+        stage_configs = []
+
+        async def get_supported_tasks(self):
+            return ("generate",)
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*args, **kwargs):
+        del args, kwargs
+        events.append("engine_enter")
+        try:
+            yield FakeEngine()
+        finally:
+            events.append("engine_exit")
+            engine_context_exited.set()
+
+    async def fake_serve_http(*args, **kwargs):
+        del args, kwargs
+        events.append("serve_http")
+        serve_started.set()
+
+        async def wait_for_shutdown():
+            await http_shutdown.wait()
+            events.append("http_shutdown")
+
+        return asyncio.create_task(wait_for_shutdown())
+
+    async def fake_storage_start():
+        events.append("storage_start")
+
+    async def fake_get_vllm_config(engine_client):
+        del engine_client
+        return None
+
+    async def fake_init_app_state(engine_client, state, args):
+        del engine_client, state, args
+        events.append("init_app_state")
+
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "build_openai_app", lambda args, supported_tasks: FastAPI())
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+    monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
+    monkeypatch.setattr(api_server, "_get_vllm_config", fake_get_vllm_config)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_init_app_state)
+    monkeypatch.setattr(api_server, "get_uvicorn_log_config", lambda args: None)
+
+    args = SimpleNamespace(
+        tool_parser_plugin="",
+        reasoning_parser_plugin="",
+        reasoning_parser=None,
+        structured_outputs_config=SimpleNamespace(reasoning_parser=None),
+        enable_ssl_refresh=False,
+        host="127.0.0.1",
+        port=0,
+        uvicorn_log_level="info",
+        disable_uvicorn_access_log=True,
+        ssl_keyfile=None,
+        ssl_certfile=None,
+        ssl_ca_certs=None,
+        ssl_cert_reqs=None,
+        ssl_ciphers=None,
+        h11_max_incomplete_event_size=None,
+        h11_max_header_count=None,
+    )
+
+    worker_task = asyncio.create_task(api_server.omni_run_server_worker("127.0.0.1:0", sock, args))
+    await asyncio.wait_for(serve_started.wait(), timeout=2)
+
+    assert not engine_context_exited.is_set()
+
+    http_shutdown.set()
+    await asyncio.wait_for(worker_task, timeout=2)
+
+    assert sock.closed
+    assert events.index("http_shutdown") < events.index("engine_exit")
+
+
 @pytest.fixture
 def test_client():
     app = FastAPI()

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -248,20 +248,27 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
     """
     # Build DBCacheConfig with optional SCM support
     db_cache_config = _build_db_cache_config(cache_config)
+    calibrator_config = _resolve_calibrator_config(cache_config)
 
     if getattr(pipeline, "transformer_2", None) is None:
         logger.info("transformer_2 not found, enabling cache-dit for single transformer mode")
         cache_dit.enable_cache(
             BlockAdapter(
                 transformer=pipeline.transformer,
+                # For VACE, cache only the main denoising blocks. The
+                # conditioning branch (vace_blocks) has a different forward
+                # contract and produces per-step hints from the current latent
+                # plus vace_context; keeping it outside CacheDiT preserves the
+                # control signal while still accelerating the repeated backbone.
                 blocks=[pipeline.transformer.blocks],
                 forward_pattern=[ForwardPattern.Pattern_2],
                 params_modifiers=[
-                    ParamsModifier(cache_config=db_cache_config),
+                    ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
                 ],
                 has_separate_cfg=True,
             ),
             cache_config=db_cache_config,
+            calibrator_config=calibrator_config,
         )
         return build_cache_context_refresh(cache_config)
 
@@ -272,6 +279,9 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
                 pipeline.transformer_2,
             ],
             blocks=[
+                # See the single-transformer branch above: VACE conditioning
+                # blocks are intentionally recomputed each step and are not
+                # wrapped by CacheDiT's main-block Pattern_2 adapter.
                 pipeline.transformer.blocks,
                 pipeline.transformer_2.blocks,
             ],
@@ -286,17 +296,20 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
                         max_warmup_steps=cache_config.max_warmup_steps,
                         max_cached_steps=cache_config.max_cached_steps,
                     ),
+                    calibrator_config=calibrator_config,
                 ),
                 ParamsModifier(
                     cache_config=DBCacheConfig().reset(
                         max_warmup_steps=2,
                         max_cached_steps=20,
                     ),
+                    calibrator_config=calibrator_config,
                 ),
             ],
             has_separate_cfg=True,
         ),
         cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
     )
 
     refresh_trans_one = build_cache_context_refresh(cache_config)
@@ -965,6 +978,7 @@ CUSTOM_DIT_ENABLERS.update(
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
         "Wan22TI2VPipeline": enable_cache_for_wan22,
+        "Wan22VACEPipeline": enable_cache_for_wan22,
         "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
         "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
     }

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -198,6 +198,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         attention_kwargs: dict[str, object],
         vace_context: torch.Tensor | None,
         vace_context_scale: float,
+        boundary_timestep: float | None = None,
     ) -> torch.Tensor:
         if attention_kwargs is None:
             attention_kwargs = {}
@@ -206,6 +207,11 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                 self._current_timestep = t
                 set_forward_context_denoise_step_idx(step_idx)
 
+                if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None:
+                    current_model = self.transformer_2
+                else:
+                    current_model = self.transformer if self.transformer is not None else self.transformer_2
+
                 latent_model_input = latents.to(dtype)
                 timestep = t.expand(latents.shape[0])
 
@@ -219,6 +225,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                     "vace_context": vace_context,
                     "vace_context_scale": vace_context_scale,
                     "return_dict": False,
+                    "current_model": current_model,
                 }
                 negative_kwargs = (
                     {
@@ -229,6 +236,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                         "vace_context": vace_context,
                         "vace_context_scale": vace_context_scale,
                         "return_dict": False,
+                        "current_model": current_model,
                     }
                     if do_true_cfg
                     else None
@@ -598,7 +606,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
                 )
 
         num_reference_images = 0
-        if self.transformer.vace_patch_embedding is not None:
+        # Either expert may be the one loaded (boundary_ratio 0.0/1.0 skips one),
+        # so read it from whichever transformer exists.
+        active_transformer = self.transformer if self.transformer is not None else self.transformer_2
+        if active_transformer.vace_patch_embedding is not None:
             video, mask, ref_images_processed = self.preprocess_conditions(
                 video=source_video,
                 mask=source_mask,
@@ -640,6 +651,19 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
 
+        # MoE boundary: only the two-expert checkpoints (Wan2.2 A14B) have a
+        # transformer_2 to switch to. For single-expert VACE (e.g. Wan2.1) this
+        # stays None and diffuse() uses the same transformer for every step.
+        boundary_timestep = None
+        if self.transformer_2 is not None:
+            boundary_ratio = (
+                self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio
+            )
+            if boundary_ratio is None:
+                boundary_ratio = 0.875
+                logger.warning_once("boundary_ratio is required for Wan2.2 VACE generation. using default value 0.875")
+            boundary_timestep = boundary_ratio * self.scheduler.config.num_train_timesteps
+
         latents = self.diffuse(
             latents=latents,
             timesteps=timesteps,
@@ -650,6 +674,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
             attention_kwargs=attention_kwargs,
             vace_context=vace_context,
             vace_context_scale=vace_context_scale,
+            boundary_timestep=boundary_timestep,
         )
 
         self._current_timestep = None

vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py

@@ -81,9 +81,7 @@ class VaceWanTransformerBlock(WanTransformerBlock):
 class WanVACETransformer3DModel(WanTransformer3DModel):
     """VACE-extended WAN Transformer with conditioning blocks for video editing."""
 
-    # TODO: `vace_blocks` are not layerwise-offloaded yet. The current offloader only
-    # supports a single block group (`blocks`); extend it to support both
-    # `vace_blocks` and `blocks`.
+    _layerwise_offload_blocks_attrs = ["vace_blocks", "blocks"]
 
     # Shard hidden_states before VACE blocks (replaces parent's blocks.0)
     _sp_plan = {

vllm_omni/entrypoints/openai/api_server.py

@@ -574,15 +574,14 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
             **uvicorn_kwargs,
         )
 
-    # NB: Await server shutdown only after the backend context is exited
-    try:
-        await shutdown_task
-    finally:
-        state = getattr(app, "state", None)
-        serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None
-        if serving_speech is not None:
-            serving_speech.shutdown()
-        sock.close()
+        try:
+            await shutdown_task
+        finally:
+            state = getattr(app, "state", None)
+            serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None
+            if serving_speech is not None:
+                serving_speech.shutdown()
+            sock.close()
 
 
 @asynccontextmanager