[BugFix] Fix Wan2.2-VACE-Fun cache and serving lifecycle

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 @@
 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 @@
     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 __init__(self, **kwargs) -> None:
     }
 
 
+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 os
 import threading
 import time
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 
 import numpy as np
@@ -109,6 +110,14 @@ async def generate_video_bytes(
             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 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheCo
         "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,
     }