From d94108693aab1bdeaf89441c709e4dc4c383edd2 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 25 Jun 2026 11:20:58 +0800 Subject: [PATCH] [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 | 1 + .../Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md | 109 ++++++++++++++++++ tests/diffusion/cache/test_cache_dit.py | 32 +++++ .../models/wan2_2/test_wan22_vace_pipeline.py | 5 + .../openai_api/test_video_server.py | 94 +++++++++++++++ .../diffusion/cache/cache_dit_backend.py | 16 ++- .../models/wan2_2/pipeline_wan2_2_vace.py | 27 ++++- .../models/wan2_2/wan2_2_vace_transformer.py | 4 +- vllm_omni/entrypoints/openai/api_server.py | 17 ++- 9 files changed, 291 insertions(+), 14 deletions(-) create mode 100644 recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 86e526a8d..95f93f5c7 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | | | diff --git a/recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md b/recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md new file mode 100644 index 000000000..f98d78b62 --- /dev/null +++ b/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: +- Base model card: +- 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 diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 0a05aeb76..68545fe27 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/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): diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 9fa9b67c4..97a000df1 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/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) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index da97c4b94..c4ac824da 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/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() diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 5add71daa..c1f194ee8 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/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, } diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 5ba2c6c69..7f3458d5e 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/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 diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index 871f0df86..a13c04c99 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/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 = { diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index fbfd3aafe..2c6a3d87a 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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