Add VACE support for WAN 2.1 conditional video generation (#1885)

Signed-off-by: Binh Tang <tangbinhna@gmail.com>
Signed-off-by: Binh Tang <binht@netflix.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Binh Tang <binht@netflix.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

docs/models/supported_models.md

@@ -31,6 +31,7 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |

docs/user_guide/diffusion_features.md

@@ -130,6 +130,7 @@ The following tables show which models support each feature:
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |

examples/offline_inference/vace/vace_video_generation.md

@@ -0,0 +1,88 @@
+# VACE Video Generation
+
+[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
+
+| Model | Architecture | Model Weights (bf16) | HuggingFace |
+|-------|-------------|----------------------|-------------|
+| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
+| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
+
+## Text-to-Video (T2V)
+
+```bash
+python vace_video_generation.py \
+  --mode t2v \
+  --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
+  --height 480 --width 832 --num-frames 81 \
+  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
+  --output t2v_output.mp4
+```
+
+## Image-to-Video (I2V)
+
+First frame is kept, remaining frames are generated:
+
+```bash
+python vace_video_generation.py \
+  --mode i2v \
+  --image astronaut.jpg \
+  --prompt "An astronaut emerging from a cracked egg on the moon" \
+  --height 480 --width 832 --num-frames 81 \
+  --output i2v_output.mp4
+```
+
+## First-Last-Frame Interpolation (FLF2V)
+
+```bash
+python vace_video_generation.py \
+  --mode flf2v \
+  --image first_frame.jpg --last-image last_frame.jpg \
+  --prompt "A bird takes off from a branch and lands on another" \
+  --height 512 --width 512 --num-frames 81 \
+  --output flf2v_output.mp4
+```
+
+## Inpainting
+
+Center vertical stripe is masked and regenerated:
+
+```bash
+python vace_video_generation.py \
+  --mode inpaint \
+  --image scene.jpg \
+  --prompt "Shrek walks out of a building" \
+  --height 480 --width 832 --num-frames 81 \
+  --output inpaint_output.mp4
+```
+
+## Reference Image-guided (R2V)
+
+```bash
+python vace_video_generation.py \
+  --mode r2v \
+  --image reference.jpg \
+  --prompt "Camera slowly zooms out from the character" \
+  --height 480 --width 832 --num-frames 81 \
+  --output r2v_output.mp4
+```
+
+## Key Arguments
+
+- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
+- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
+- `--image`: Input image for I2V, inpainting, and R2V modes.
+- `--last-image`: Last frame image for FLF2V mode.
+- `--prompt`: Text description of desired video.
+- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
+- `--num-frames`: Number of frames (default 81).
+- `--guidance-scale`: CFG scale (default 5.0).
+- `--flow-shift`: Scheduler flow shift (default 5.0).
+- `--num-inference-steps`: Number of denoising steps (default 30).
+- `--fps`: Frames per second for the saved MP4 (default 16).
+- `--output`: Path to save the generated video.
+- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
+- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
+- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
+- `--tensor-parallel-size`: Tensor parallel size.
+
+> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.

examples/offline_inference/vace/vace_video_generation.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""VACE video generation example.
+
+VACE (Video All-in-one Creation Engine) supports multiple video tasks:
+  - T2V:        Text-to-Video
+  - I2V:        Image-to-Video (first frame conditioning)
+  - V2LF:       Video-to-Last-Frame
+  - FLF2V:      First-Last-Frame interpolation
+  - Inpainting:  Masked region generation
+  - R2V:        Reference image-guided generation
+
+Usage examples:
+  # T2V (text-to-video)
+  python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
+
+  # I2V (image-to-video, first frame kept)
+  python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
+
+  # FLF2V (first-last frame interpolation)
+  python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
+
+  # R2V (reference image guided)
+  python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="VACE video generation.")
+    parser.add_argument(
+        "--model",
+        default="Wan-AI/Wan2.1-VACE-14B-diffusers",
+        help="VACE model ID or local path.",
+    )
+    parser.add_argument(
+        "--mode",
+        default="t2v",
+        choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
+        help="Generation mode.",
+    )
+    parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
+    parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
+    parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
+    parser.add_argument("--height", type=int, default=480, help="Video height.")
+    parser.add_argument("--width", type=int, default=832, help="Video width.")
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
+    parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
+    parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
+    parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
+    parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
+    parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
+    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
+    parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
+    return parser.parse_args()
+
+
+def build_prompts(args):
+    """Build prompt dict with multi_modal_data based on mode."""
+    h, w, nf = args.height, args.width, args.num_frames
+
+    gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
+    mask_black = PIL.Image.new("L", (w, h), 0)
+    mask_white = PIL.Image.new("L", (w, h), 255)
+
+    prompt_data = {
+        "prompt": args.prompt,
+        "negative_prompt": args.negative_prompt,
+    }
+
+    if args.mode == "t2v":
+        return prompt_data
+
+    if args.mode == "r2v":
+        assert args.image, "--image required for R2V mode"
+        ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
+        return prompt_data
+
+    if args.mode == "i2v":
+        assert args.image, "--image required for I2V mode"
+        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {
+            "video": [img] + [gray] * (nf - 1),
+            "mask": [mask_black] + [mask_white] * (nf - 1),
+        }
+        return prompt_data
+
+    if args.mode == "v2lf":
+        assert args.image, "--image required for V2LF mode"
+        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {
+            "video": [gray] * (nf - 1) + [img],
+            "mask": [mask_white] * (nf - 1) + [mask_black],
+        }
+        return prompt_data
+
+    if args.mode == "flf2v":
+        assert args.image and args.last_image, "--image and --last-image required for FLF2V"
+        first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {
+            "video": [first] + [gray] * (nf - 2) + [last],
+            "mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
+        }
+        return prompt_data
+
+    if args.mode == "inpaint":
+        assert args.image, "--image required for inpaint mode"
+        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        d = 80
+        frames, masks = [], []
+        for _ in range(nf):
+            base = np.array(img).copy()
+            mask = PIL.Image.new("L", (w, h), 0)
+            stripe = PIL.Image.new("L", (2 * d, h), 255)
+            mask.paste(stripe, (w // 2 - d, 0))
+            base[np.array(mask) > 128] = 128
+            frames.append(PIL.Image.fromarray(base))
+            masks.append(mask)
+        prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
+        return prompt_data
+
+    raise ValueError(f"Unknown mode: {args.mode}")
+
+
+def main():
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        vae_use_tiling=args.vae_use_tiling,
+        flow_shift=args.flow_shift,
+        enforce_eager=args.enforce_eager,
+        parallel_config=parallel_config,
+    )
+
+    prompt_data = build_prompts(args)
+
+    print(f"\n{'=' * 60}")
+    print(f"VACE {args.mode.upper()} Generation")
+    print(f"  Model: {args.model}")
+    print(f"  Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
+    print(f"{'=' * 60}\n")
+
+    start = time.perf_counter()
+    outputs = omni.generate(
+        prompt_data,
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_frames=args.num_frames,
+            num_inference_steps=args.num_inference_steps,
+            guidance_scale=args.guidance_scale,
+            generator=generator,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+
+    video = outputs[0].images
+    if isinstance(video, list):
+        video = video[0]
+    if isinstance(video, torch.Tensor):
+        video = video.cpu().numpy()
+    if video.ndim == 5:
+        video = video[0]
+    print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    from diffusers.utils import export_to_video
+
+    if np.issubdtype(video.dtype, np.integer):
+        video = video.astype(np.float32) / 255.0
+    export_to_video(list(video), str(output_path), fps=args.fps)
+    print(f"Saved to {output_path}")
+
+    omni.close()
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/online_serving/test_wan_2_1_vace_expansion.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Comprehensive e2e tests of diffusion features for Wan2.1-VACE in online serving mode.
+
+Wan2.1-VACE supports: Cache-DiT, Ulysses-SP, Ring, CFG-Parallel, TP,
+VAE-Patch-Parallel, HSDP. TeaCache is NOT supported for this model, so
+Cache-DiT is used in place of TeaCache for single-card and CFG tests.
+
+Uses the 1.3B variant for faster CI testing.
+
+Coverage:
+  Single GPU:
+    - Cache-DiT + layerwise CPU offload
+  Two GPUs:
+    - Cache-DiT + Ulysses-SP = 2
+    - Cache-DiT + Ring = 2
+    - Cache-DiT + CFG-Parallel = 2
+    - Cache-DiT + TP = 2 + VAE-Patch-Parallel = 2
+    - Cache-DiT + HSDP = 2 + VAE-Patch-Parallel = 2
+"""
+
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+)
+from tests.utils import hardware_marks
+
+MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
+PROMPT = "A cat walking slowly across a sunlit garden path"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+def _get_vace_feature_cases():
+    return [
+        # Single GPU: Cache-DiT + layerwise CPU offload
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--enable-layerwise-offload",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="single_card_001",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        # 2 GPUs: Cache-DiT + Ulysses-SP = 2
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ulysses-degree",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_001",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        # 2 GPUs: Cache-DiT + Ring = 2
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--ring",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_002",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        # 2 GPUs: Cache-DiT + CFG-Parallel = 2
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--cfg-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_003",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        # 2 GPUs: Cache-DiT + TP = 2 + VAE-Patch-Parallel = 2
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--tensor-parallel-size",
+                    "2",
+                    "--vae-patch-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_004",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        # 2 GPUs: Cache-DiT + HSDP = 2 + VAE-Patch-Parallel = 2
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--hsdp-shard-size",
+                    "2",
+                    "--vae-patch-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_005",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_vace_feature_cases(),
+    indirect=True,
+)
+def test_wan_2_1_vace(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test VACE T2V generation with all supported diffusion acceleration features."""
+    openai_client.send_video_diffusion_request(
+        {
+            "model": MODEL,
+            "form_data": {
+                "prompt": PROMPT,
+                "height": 480,
+                "width": 320,
+                "num_frames": 5,
+                "fps": 8,
+                "num_inference_steps": 2,
+                "guidance_scale": 5.0,
+                "seed": 42,
+            },
+        }
+    )

vllm_omni/diffusion/models/wan2_2/__init__.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
 from .pipeline_wan2_2 import (
     Wan22Pipeline,
     create_transformer_from_config,
@@ -16,7 +19,13 @@ from .pipeline_wan2_2_ti2v import (
     get_wan22_ti2v_post_process_func,
     get_wan22_ti2v_pre_process_func,
 )
+from .pipeline_wan2_2_vace import (
+    Wan22VACEPipeline,
+    get_wan22_vace_post_process_func,
+    get_wan22_vace_pre_process_func,
+)
 from .wan2_2_transformer import WanTransformer3DModel
+from .wan2_2_vace_transformer import VaceWanTransformerBlock, WanVACETransformer3DModel
 
 __all__ = [
     "Wan22Pipeline",
@@ -31,5 +40,10 @@ __all__ = [
     "Wan22TI2VPipeline",
     "get_wan22_ti2v_post_process_func",
     "get_wan22_ti2v_pre_process_func",
+    "Wan22VACEPipeline",
+    "get_wan22_vace_post_process_func",
+    "get_wan22_vace_pre_process_func",
     "WanTransformer3DModel",
+    "VaceWanTransformerBlock",
+    "WanVACETransformer3DModel",
 ]

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -278,13 +278,13 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
         # Initialize transformers with correct config (weights loaded via load_weights)
         if load_transformer:
             transformer_config = load_transformer_config(model, "transformer", local_files_only)
-            self.transformer = create_transformer_from_config(transformer_config)
+            self.transformer = self._create_transformer(transformer_config)
         else:
             self.transformer = None
 
         if load_transformer_2:
             transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
-            self.transformer_2 = create_transformer_from_config(transformer_2_config)
+            self.transformer_2 = self._create_transformer(transformer_2_config)
         else:
             self.transformer_2 = None
 
@@ -316,6 +316,10 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
 
+    def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+        """Create a transformer from a config dict. Subclasses may override."""
+        return create_transformer_from_config(config)
+
     @property
     def guidance_scale(self):
         return self._guidance_scale

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -0,0 +1,645 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+VACE (Video Creation and Editing) Pipeline for WAN models.
+
+VACE is an all-in-one model for video creation and editing. The mode is
+determined by which inputs are provided (no explicit mode flag):
+
+- T2V: Text-to-Video (prompt only)
+- R2V: Reference-to-Video (prompt + reference_images)
+- V2V: Video-to-Video editing (prompt + video)
+- MV2V: Masked Video-to-Video / inpainting (prompt + video + mask)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import replace
+
+import PIL.Image
+import torch
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    Wan22Pipeline,
+    retrieve_latents,
+)
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    get_wan22_post_process_func as get_wan22_vace_post_process_func,  # noqa: F401
+)
+from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniTextPrompt
+from vllm_omni.platforms import current_omni_platform
+
+logger = init_logger(__name__)
+
+
+def create_vace_transformer_from_config(config: dict) -> WanVACETransformer3DModel:
+    """Create WanVACETransformer3DModel from config dict."""
+    kwargs = {}
+    if "patch_size" in config:
+        kwargs["patch_size"] = tuple(config["patch_size"])
+    if "num_attention_heads" in config:
+        kwargs["num_attention_heads"] = config["num_attention_heads"]
+    if "attention_head_dim" in config:
+        kwargs["attention_head_dim"] = config["attention_head_dim"]
+    if "in_channels" in config:
+        kwargs["in_channels"] = config["in_channels"]
+    if "out_channels" in config:
+        kwargs["out_channels"] = config["out_channels"]
+    if "text_dim" in config:
+        kwargs["text_dim"] = config["text_dim"]
+    if "freq_dim" in config:
+        kwargs["freq_dim"] = config["freq_dim"]
+    if "ffn_dim" in config:
+        kwargs["ffn_dim"] = config["ffn_dim"]
+    if "num_layers" in config:
+        kwargs["num_layers"] = config["num_layers"]
+    if "cross_attn_norm" in config:
+        kwargs["cross_attn_norm"] = config["cross_attn_norm"]
+    if "eps" in config:
+        kwargs["eps"] = config["eps"]
+    if "image_dim" in config:
+        kwargs["image_dim"] = config["image_dim"]
+    if "added_kv_proj_dim" in config:
+        kwargs["added_kv_proj_dim"] = config["added_kv_proj_dim"]
+    if "rope_max_seq_len" in config:
+        kwargs["rope_max_seq_len"] = config["rope_max_seq_len"]
+    if "pos_embed_seq_len" in config:
+        kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
+    if "vace_layers" in config:
+        kwargs["vace_layers"] = config["vace_layers"]
+    if "vace_in_channels" in config:
+        kwargs["vace_in_channels"] = config["vace_in_channels"]
+
+    return WanVACETransformer3DModel(**kwargs)
+
+
+def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig):
+    """Pre-process function for VACE: handle reference images, source videos, and masks."""
+    import numpy as np
+
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
+
+            if not multi_modal_data:
+                request.prompts[i] = prompt
+                continue
+
+            # Handle reference images for R2V
+            # "image" is the standard key from online serving (SupportImageInput convention)
+            # "reference_images" is the offline API key for backwards compatibility
+            ref_images = multi_modal_data.get("image") or multi_modal_data.get("reference_images")
+            if ref_images is not None:
+                if isinstance(ref_images, str):
+                    ref_images = [PIL.Image.open(ref_images).convert("RGB")]
+                elif isinstance(ref_images, PIL.Image.Image):
+                    ref_images = [ref_images]
+                elif isinstance(ref_images, list):
+                    ref_images = [
+                        PIL.Image.open(img).convert("RGB") if isinstance(img, str) else img for img in ref_images
+                    ]
+
+                # Calculate dimensions from first reference image if not provided
+                if request.sampling_params.height is None or request.sampling_params.width is None:
+                    first_img = ref_images[0]
+                    max_area = 480 * 832  # VACE default is 480p
+                    aspect_ratio = first_img.height / first_img.width
+                    mod_value = 16
+                    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+                    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+
+                    if request.sampling_params.height is None:
+                        request.sampling_params.height = height
+                    if request.sampling_params.width is None:
+                        request.sampling_params.width = width
+
+                prompt["additional_information"]["reference_images"] = ref_images
+
+            # Handle source video for V2V / MV2V
+            source_video = multi_modal_data.get("video")
+            if source_video is not None:
+                if isinstance(source_video, list) and len(source_video) > 0:
+                    if isinstance(source_video[0], str):
+                        source_video = [PIL.Image.open(f).convert("RGB") for f in source_video]
+                prompt["additional_information"]["source_video"] = source_video
+
+            # Handle mask for MV2V / inpainting
+            mask = multi_modal_data.get("mask")
+            if mask is not None:
+                if isinstance(mask, list) and len(mask) > 0:
+                    if isinstance(mask[0], str):
+                        mask = [PIL.Image.open(m).convert("L") for m in mask]
+                elif isinstance(mask, str):
+                    mask = [PIL.Image.open(mask).convert("L")]
+                elif isinstance(mask, PIL.Image.Image):
+                    mask = [mask]
+                prompt["additional_information"]["mask"] = mask
+
+            request.prompts[i] = prompt
+        return request
+
+    return pre_process_func
+
+
+class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
+    """VACE (Video Creation and Editing) Pipeline for Wan2.1.
+
+    Extends Wan22Pipeline with VACE-specific context creation and weight loading.
+    All VACE modes (T2V, R2V, V2V, MV2V) are handled by varying the inputs.
+    """
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        # VACE defaults to flow_shift=3.0 for 480p (base WAN T2V uses 5.0 for 720p)
+        if od_config.flow_shift is None:
+            od_config = replace(od_config, flow_shift=3.0)
+
+        super().__init__(od_config=od_config, prefix=prefix)
+
+    def _create_transformer(self, config: dict) -> WanVACETransformer3DModel:
+        """Build VACE transformer directly from config dict."""
+        return create_vace_transformer_from_config(config)
+
+    def check_inputs(
+        self,
+        prompt,
+        negative_prompt,
+        height,
+        width,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        video=None,
+        mask=None,
+        reference_images=None,
+    ):
+        super().check_inputs(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+
+        # VACE-specific: validate video/mask/reference_images consistency
+        if video is not None:
+            if mask is not None and len(video) != len(mask):
+                raise ValueError(
+                    f"Length of `video` ({len(video)}) and `mask` ({len(mask)}) do not match. "
+                    "Please make sure that they have the same length."
+                )
+            if reference_images is not None:
+                is_pil_image = isinstance(reference_images, PIL.Image.Image)
+                is_list_of_pil_images = isinstance(reference_images, list) and all(
+                    isinstance(img, PIL.Image.Image) for img in reference_images
+                )
+                if not (is_pil_image or is_list_of_pil_images):
+                    raise ValueError(
+                        "`reference_images` has to be of type `PIL.Image.Image` or `list` of `PIL.Image.Image`, "
+                        f"but is {type(reference_images)}"
+                    )
+        elif mask is not None:
+            raise ValueError("`mask` can only be passed if `video` is passed as well.")
+
+    def preprocess_conditions(
+        self,
+        video: list | torch.Tensor | None,
+        mask: list | torch.Tensor | None,
+        reference_images: list[PIL.Image.Image] | None,
+        height: int,
+        width: int,
+        num_frames: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor, list[list[torch.Tensor]]]:
+        """Preprocess video, mask, and reference images for VACE conditioning.
+
+        - If video is None, create zero tensor (T2V mode)
+        - If mask is None, create all-ones tensor (generate everything)
+        - Reference images are resized maintaining aspect ratio and center-padded
+
+        Returns:
+            (video, mask, reference_images_processed) tensors ready for VAE encoding.
+        """
+        from diffusers.video_processor import VideoProcessor
+
+        video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+        if video is None:
+            video = torch.zeros(1, 3, num_frames, height, width, dtype=dtype, device=device)
+            image_size = (height, width)
+        else:
+            base = self.vae_scale_factor_spatial * self.transformer_config.patch_size[1]
+            if isinstance(video, list):
+                video_height, video_width = video_processor.get_default_height_width(video[0])
+                # Downscale if video exceeds target area
+                if video_height * video_width > height * width:
+                    scale = min(width / video_width, height / video_height)
+                    video_height, video_width = int(video_height * scale), int(video_width * scale)
+                # Align to base
+                video_height = (video_height // base) * base
+                video_width = (video_width // base) * base
+                video = video_processor.preprocess_video(video, video_height, video_width)
+            image_size = (video.shape[-2], video.shape[-1])
+
+        if mask is None:
+            mask = torch.ones_like(video)
+        else:
+            if isinstance(mask, list):
+                mask = video_processor.preprocess_video(mask, image_size[0], image_size[1])
+            mask = torch.clamp((mask + 1) / 2, min=0, max=1)
+
+        video = video.to(dtype=dtype, device=device)
+        mask = mask.to(dtype=dtype, device=device)
+
+        # Preprocess reference images: resize with aspect ratio, center-pad on white canvas
+        ref_images_processed: list[list[torch.Tensor]] = []
+        if reference_images is not None and len(reference_images) > 0:
+            preprocessed = []
+            for image in reference_images:
+                img_tensor = video_processor.preprocess(image, None, None)
+                img_h, img_w = img_tensor.shape[-2:]
+                scale = min(image_size[0] / img_h, image_size[1] / img_w)
+                new_h, new_w = int(img_h * scale), int(img_w * scale)
+                resized = torch.nn.functional.interpolate(
+                    img_tensor, size=(new_h, new_w), mode="bilinear", align_corners=False
+                ).squeeze(0)
+                canvas = torch.ones(3, *image_size, device=device, dtype=dtype)
+                top = (image_size[0] - new_h) // 2
+                left = (image_size[1] - new_w) // 2
+                canvas[:, top : top + new_h, left : left + new_w] = resized
+                preprocessed.append(canvas)
+            ref_images_processed = [preprocessed]
+        else:
+            ref_images_processed = [[]]
+
+        return video, mask, ref_images_processed
+
+    def prepare_video_latents(
+        self,
+        video: torch.Tensor,
+        mask: torch.Tensor,
+        reference_images: list[list[torch.Tensor]],
+        generator: torch.Generator | None,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Encode video and reference images into VACE conditioning latents.
+
+        - Encodes inactive (video * (1-mask)) and reactive (video * mask) regions
+        - Reference images are encoded and prepended as extra temporal frames
+        """
+        vae_dtype = self.vae.dtype
+
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=torch.float32).view(
+            1, self.vae.config.z_dim, 1, 1, 1
+        )
+        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std, device=device, dtype=torch.float32).view(
+            1, self.vae.config.z_dim, 1, 1, 1
+        )
+
+        # Binarize mask
+        mask = torch.where(mask > 0.5, 1.0, 0.0).to(dtype=vae_dtype)
+
+        # Encode inactive and reactive regions separately
+        video = video.to(dtype=vae_dtype)
+        inactive = video * (1 - mask)
+        reactive = video * mask
+
+        with torch.no_grad():
+            inactive_latent = retrieve_latents(self.vae.encode(inactive), generator, sample_mode="argmax")
+            reactive_latent = retrieve_latents(self.vae.encode(reactive), generator, sample_mode="argmax")
+
+        inactive_latent = ((inactive_latent.float() - latents_mean) * latents_std).to(vae_dtype)
+        reactive_latent = ((reactive_latent.float() - latents_mean) * latents_std).to(vae_dtype)
+
+        # Concatenate inactive + reactive along channels -> [B, 2*z_dim, T, H, W]
+        latents = torch.cat([inactive_latent, reactive_latent], dim=1)
+
+        # Prepend reference image latents along temporal dimension
+        latent_list = []
+        for latent, ref_batch in zip(latents, reference_images):
+            for ref_image in ref_batch:
+                ref_image = ref_image.to(dtype=vae_dtype)
+                ref_image = ref_image[None, :, None, :, :]  # [1, C, 1, H, W]
+                with torch.no_grad():
+                    ref_latent = retrieve_latents(self.vae.encode(ref_image), generator, sample_mode="argmax")
+                ref_latent = ((ref_latent.float() - latents_mean) * latents_std).to(vae_dtype)
+                ref_latent = ref_latent.squeeze(0)  # [z_dim, 1, H, W]
+                # Double channels with zeros (inactive=ref, reactive=zeros)
+                ref_latent = torch.cat([ref_latent, torch.zeros_like(ref_latent)], dim=0)
+                # Prepend along temporal dimension
+                latent = torch.cat([ref_latent, latent], dim=1)
+            latent_list.append(latent)
+
+        return torch.stack(latent_list)
+
+    def prepare_masks(
+        self,
+        mask: torch.Tensor,
+        reference_images: list[list[torch.Tensor]],
+    ) -> torch.Tensor:
+        """Encode mask using spatial stride sampling and prepend reference padding.
+
+        - 8x8 spatial stride encoding -> 64 channels
+        - Zero-masks prepended for reference image frames
+        """
+        patch_size = self.transformer_config.patch_size if hasattr(self.transformer_config, "patch_size") else (1, 2, 2)
+        if isinstance(self.transformer_config, dict):
+            patch_size = self.transformer_config.get("patch_size", (1, 2, 2))
+        transformer_patch_size = patch_size[1] if isinstance(patch_size, list | tuple) else 2
+
+        mask_list = []
+        for mask_, ref_batch in zip(mask, reference_images):
+            num_channels, num_frames, height, width = mask_.shape
+            new_num_frames = (num_frames + self.vae_scale_factor_temporal - 1) // self.vae_scale_factor_temporal
+            new_height = height // (self.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size
+            new_width = width // (self.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size
+
+            m = mask_[0, :, :, :]  # [T, H, W]
+            m = m.view(num_frames, new_height, self.vae_scale_factor_spatial, new_width, self.vae_scale_factor_spatial)
+            m = m.permute(2, 4, 0, 1, 3).flatten(0, 1)  # [64, T, H', W']
+            m = torch.nn.functional.interpolate(
+                m.unsqueeze(0), size=(new_num_frames, new_height, new_width), mode="nearest-exact"
+            ).squeeze(0)
+
+            # Prepend zero-masks for reference image frames
+            num_ref = len(ref_batch)
+            if num_ref > 0:
+                mask_padding = torch.zeros_like(m[:, :num_ref, :, :])
+                m = torch.cat([mask_padding, m], dim=1)
+
+            mask_list.append(m)
+
+        return torch.stack(mask_list)
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | None = None,
+        negative_prompt: str | None = None,
+        height: int = 480,
+        width: int = 832,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 5.0,
+        frame_num: int = 81,
+        output_type: str | None = "np",
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        attention_kwargs: dict | None = None,
+        vace_context_scale: float | list[float] = 1.0,
+        **kwargs,
+    ) -> DiffusionOutput:
+        """Generate or edit video using VACE.
+
+        The mode is determined by which inputs are provided in the request:
+        - T2V: prompt only (no video/mask/reference_images)
+        - R2V: prompt + reference_images (in multi_modal_data)
+        - V2V: prompt + video (in multi_modal_data)
+        - MV2V: prompt + video + mask (in multi_modal_data)
+
+        Args:
+            req: Diffusion request containing prompt and optional multi-modal data.
+            prompt: Text prompt (overridden by req.prompts if provided).
+            negative_prompt: Negative prompt for CFG.
+            height: Output video height.
+            width: Output video width.
+            num_inference_steps: Number of denoising steps.
+            guidance_scale: CFG scale.
+            frame_num: Number of output frames.
+            output_type: Output format ("np", "pt", or "latent").
+            generator: Random generator for reproducibility.
+            prompt_embeds: Pre-computed prompt embeddings.
+            negative_prompt_embeds: Pre-computed negative prompt embeddings.
+            attention_kwargs: Additional kwargs for attention layers.
+            vace_context_scale: VACE conditioning strength.
+        """
+        # Get parameters from request or arguments
+        if len(req.prompts) > 1:
+            raise ValueError(
+                "This model only supports a single prompt, not a batched request. "
+                "Please pass in a single prompt object or string, or a single-item list."
+            )
+
+        reference_images = None
+        source_video = None
+        source_mask = None
+
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            if isinstance(first_prompt, str):
+                prompt = first_prompt
+            else:
+                prompt = first_prompt.get("prompt")
+                negative_prompt = negative_prompt or first_prompt.get("negative_prompt")
+                prompt_embeds = prompt_embeds if prompt_embeds is not None else first_prompt.get("prompt_embeds")
+                negative_prompt_embeds = (
+                    negative_prompt_embeds
+                    if negative_prompt_embeds is not None
+                    else first_prompt.get("negative_prompt_embeds")
+                )
+
+                additional_info = first_prompt.get("additional_information", {})
+                reference_images = additional_info.get("reference_images")
+                source_video = additional_info.get("source_video")
+                source_mask = additional_info.get("mask")
+
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Prompt or prompt_embeds is required for VACE generation.")
+
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_frames = req.sampling_params.num_frames or frame_num
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        generator = req.sampling_params.generator or generator
+
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+
+        # Ensure dimensions are compatible with VAE and patch size
+        mod_value = self.vae_scale_factor_spatial * 2  # 8 * 2 = 16
+        height = (height // mod_value) * mod_value
+        width = (width // mod_value) * mod_value
+
+        if num_frames % self.vae_scale_factor_temporal != 1:
+            num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
+        num_frames = max(num_frames, 1)
+
+        self.check_inputs(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            video=source_video,
+            mask=source_mask,
+            reference_images=reference_images,
+        )
+
+        device = self.device
+        self._guidance_scale = guidance_scale
+        dtype = self.transformer.dtype if self.transformer is not None else torch.bfloat16
+
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
+
+        # Encode prompts
+        if prompt_embeds is None:
+            if prompt is None:
+                raise ValueError("Either prompt or prompt_embeds must be provided.")
+            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                do_classifier_free_guidance=guidance_scale > 1.0,
+                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                max_sequence_length=req.sampling_params.max_sequence_length or 512,
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
+            if negative_prompt_embeds is not None:
+                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+            elif guidance_scale > 1.0:
+                _, negative_prompt_embeds = self.encode_prompt(
+                    prompt="",
+                    negative_prompt=None,
+                    do_classifier_free_guidance=True,
+                    device=device,
+                    dtype=dtype,
+                )
+
+        num_reference_images = 0
+        if self.transformer.vace_patch_embedding is not None:
+            video, mask, ref_images_processed = self.preprocess_conditions(
+                video=source_video,
+                mask=source_mask,
+                reference_images=reference_images,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                dtype=dtype,
+                device=device,
+            )
+
+            conditioning_latents = self.prepare_video_latents(video, mask, ref_images_processed, generator, device)
+            mask_encoded = self.prepare_masks(mask, ref_images_processed)
+
+            # Unified VACE context: [video_latents, mask] along channels -> [B, C, T, H, W]
+            vace_context = torch.cat([conditioning_latents, mask_encoded], dim=1)
+
+            num_reference_images = len(ref_images_processed[0]) if ref_images_processed else 0
+        else:
+            vace_context = None
+
+        # Prepare noise latents (extra frames for reference images)
+        num_channels_latents = self.transformer_config.in_channels
+        noise_num_frames = num_frames + num_reference_images * self.vae_scale_factor_temporal
+        latents = self.prepare_latents(
+            batch_size=prompt_embeds.shape[0],
+            num_channels_latents=num_channels_latents,
+            height=height,
+            width=width,
+            num_frames=noise_num_frames,
+            dtype=torch.float32,
+            device=device,
+            generator=generator,
+            latents=req.sampling_params.latents,
+        )
+
+        # Set up scheduler
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        self._num_timesteps = len(timesteps)
+
+        # Denoising loop
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for t in timesteps:
+                self._current_timestep = t
+                latent_model_input = latents.to(dtype)
+                timestep = t.expand(latents.shape[0])
+
+                do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None
+
+                positive_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep,
+                    "encoder_hidden_states": prompt_embeds,
+                    "attention_kwargs": attention_kwargs,
+                    "vace_context": vace_context,
+                    "vace_context_scale": vace_context_scale,
+                    "return_dict": False,
+                }
+                negative_kwargs = (
+                    {
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep,
+                        "encoder_hidden_states": negative_prompt_embeds,
+                        "attention_kwargs": attention_kwargs,
+                        "vace_context": vace_context,
+                        "vace_context_scale": vace_context_scale,
+                        "return_dict": False,
+                    }
+                    if do_true_cfg
+                    else None
+                )
+
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=guidance_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    cfg_normalize=False,
+                )
+
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+                pbar.update()
+
+        self._current_timestep = None
+
+        if current_omni_platform.is_available():
+            current_omni_platform.empty_cache()
+
+        # Trim reference frames from output before decoding
+        # (reference images were prepended as extra temporal frames)
+        if output_type != "latent" and num_reference_images > 0:
+            latents = latents[:, :, num_reference_images:]
+
+        if output_type == "latent":
+            output = latents
+        else:
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            output = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(output=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights using AutoWeightsLoader for vLLM integration."""
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""VACE variant of WanTransformer3DModel for conditional video generation."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard
+from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
+    Transformer2DModelOutput,
+    WanTransformer3DModel,
+    WanTransformerBlock,
+)
+
+
+class VaceWanTransformerBlock(WanTransformerBlock):
+    """VACE variant of WanTransformerBlock with proj_in/proj_out for skip connections."""
+
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        added_kv_proj_dim: int | None = None,
+        cross_attn_norm: bool = False,
+        block_id: int = 0,
+    ):
+        super().__init__(dim, ffn_dim, num_heads, eps, added_kv_proj_dim, cross_attn_norm)
+        self.proj_in = nn.Linear(dim, dim) if block_id == 0 else None
+        self.proj_out = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        control_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        hidden_states_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.proj_in is not None:
+            control_hidden_states = self.proj_in(control_hidden_states)
+            control_hidden_states = control_hidden_states + hidden_states
+
+        control_hidden_states = super().forward(
+            control_hidden_states,
+            encoder_hidden_states,
+            temb,
+            rotary_emb,
+            hidden_states_mask,
+        )
+
+        conditioning_states = self.proj_out(control_hidden_states)
+        return conditioning_states, control_hidden_states
+
+
+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`.
+
+    # Shard hidden_states before VACE blocks (replaces parent's blocks.0)
+    _sp_plan = {
+        **{k: v for k, v in WanTransformer3DModel._sp_plan.items() if k != "blocks.0"},
+        "_sp_shard_point": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+        },
+    }
+
+    def __init__(
+        self,
+        *,
+        vace_layers: list[int] | None = None,
+        vace_in_channels: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.vace_blocks = None
+        self.vace_patch_embedding = None
+        self.vace_layers = None
+        self.vace_layers_mapping = None
+
+        # SP shard point: Identity module that _sp_plan hooks into to shard
+        # hidden_states before VACE processing (instead of at blocks.0)
+        self._sp_shard_point = nn.Identity()
+
+        if vace_layers is not None:
+            inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+            self.vace_layers = list(vace_layers)
+            self.vace_layers_mapping = {layer_idx: vace_idx for vace_idx, layer_idx in enumerate(vace_layers)}
+
+            vace_in_channels = vace_in_channels or self.config.in_channels
+            self.vace_patch_embedding = nn.Conv3d(
+                vace_in_channels,
+                inner_dim,
+                kernel_size=self.config.patch_size,
+                stride=self.config.patch_size,
+            )
+            self.vace_blocks = nn.ModuleList(
+                [
+                    VaceWanTransformerBlock(
+                        inner_dim,
+                        self.config.ffn_dim,
+                        self.config.num_attention_heads,
+                        self.config.eps,
+                        self.config.added_kv_proj_dim,
+                        self.config.cross_attn_norm,
+                        block_id=i,
+                    )
+                    for i in range(len(vace_layers))
+                ]
+            )
+
+    def embed_vace_context(
+        self,
+        vace_context: torch.Tensor,
+        seq_len: int,
+        sp_size: int = 1,
+    ) -> torch.Tensor:
+        """Compute VACE patch embeddings, aligned and sharded for SP.
+
+        Args:
+            vace_context: Raw conditioning tensor [B, C, T, H, W].
+            seq_len: Target full (padded) sequence length to align to.
+            sp_size: Sequence parallel world size.
+        """
+        vace_embeds = self.vace_patch_embedding(vace_context)
+        vace_embeds = vace_embeds.flatten(2).transpose(1, 2)
+
+        # Align to target seq_len (may include SP padding)
+        if vace_embeds.size(1) < seq_len:
+            vace_embeds = F.pad(vace_embeds, (0, 0, 0, seq_len - vace_embeds.size(1)))
+
+        if sp_size > 1:
+            vace_embeds = sp_shard(vace_embeds, dim=1)
+        return vace_embeds
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_image: torch.Tensor | None = None,
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        vace_context: torch.Tensor | None = None,
+        vace_context_scale: float | list[float] = 1.0,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        batch_size, _, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.config.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Compute RoPE embeddings (sharded by _sp_plan via split_output=True)
+        rotary_emb = self.rope(hidden_states)
+
+        # Patch embedding and flatten to sequence
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        if timestep.ndim == 2:
+            ts_seq_len = timestep.shape[1]
+            timestep = timestep.flatten()
+        else:
+            ts_seq_len = None
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
+        )
+        timestep_proj = self.timestep_proj_prepare(timestep_proj, ts_seq_len)
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        # Shard hidden_states via _sp_plan hook (before VACE, not at blocks.0)
+        hidden_states = self._sp_shard_point(hidden_states)
+
+        # SP state and attention mask for padding
+        hidden_states_mask = None
+        ctx = get_forward_context()
+        parallel_config = ctx.omni_diffusion_config.parallel_config
+        sp_size = parallel_config.sequence_parallel_size if parallel_config is not None else 1
+        if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
+            padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
+            hidden_states_mask = torch.ones(
+                batch_size,
+                padded_seq_len,
+                dtype=torch.bool,
+                device=hidden_states.device,
+            )
+            hidden_states_mask[:, ctx.sp_original_seq_len :] = False
+
+        # VACE: embed context and run conditioning blocks
+        vace_hints = None
+        if vace_context is not None and self.vace_blocks is not None:
+            full_seq_len = hidden_states.shape[1] * sp_size
+            control_hidden_states = self.embed_vace_context(vace_context.to(hidden_states.dtype), full_seq_len, sp_size)
+            vace_hints = []
+            for block in self.vace_blocks:
+                conditioning_states, control_hidden_states = block(
+                    hidden_states,
+                    encoder_hidden_states,
+                    control_hidden_states,
+                    timestep_proj,
+                    rotary_emb,
+                    hidden_states_mask,
+                )
+                vace_hints.append(conditioning_states)
+
+        # Normalize scale to per-layer list
+        if vace_hints is not None and isinstance(vace_context_scale, (int, float)):
+            vace_context_scale = [vace_context_scale] * len(vace_hints)
+
+        # Transformer blocks with VACE hint application
+        for i, block in enumerate(self.blocks):
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask)
+            if vace_hints is not None and self.vace_layers_mapping is not None and i in self.vace_layers_mapping:
+                vace_idx = self.vace_layers_mapping[i]
+                hidden_states = hidden_states + vace_hints[vace_idx] * vace_context_scale[vace_idx]
+
+        # Output norm, projection & unpatchify
+        shift, scale = self.output_scale_shift_prepare(temb)
+        shift = shift.to(hidden_states.device)
+        scale = scale.to(hidden_states.device)
+        if shift.ndim == 2:
+            shift = shift.unsqueeze(1)
+            scale = scale.unsqueeze(1)
+
+        hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1
+        )
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

vllm_omni/diffusion/registry.py

@@ -57,6 +57,11 @@ _DIFFUSION_MODELS = {
         "pipeline_wan2_2",
         "Wan22Pipeline",
     ),
+    "WanVACEPipeline": (
+        "wan2_2",
+        "pipeline_wan2_2_vace",
+        "Wan22VACEPipeline",
+    ),
     "LTX2Pipeline": (
         "ltx2",
         "pipeline_ltx2",
@@ -316,6 +321,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "ZImagePipeline": "get_post_process_func",
     "OvisImagePipeline": "get_ovis_image_post_process_func",
     "WanPipeline": "get_wan22_post_process_func",
+    "WanVACEPipeline": "get_wan22_vace_post_process_func",
     "LTX2Pipeline": "get_ltx2_post_process_func",
     "LTX2ImageToVideoPipeline": "get_ltx2_post_process_func",
     "StableAudioPipeline": "get_stable_audio_post_process_func",
@@ -346,6 +352,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",
     "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",
     "WanPipeline": "get_wan22_pre_process_func",
+    "WanVACEPipeline": "get_wan22_vace_pre_process_func",
     "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func",
     "OmniGen2Pipeline": "get_omnigen2_pre_process_func",
     "HeliosPipeline": "get_helios_pre_process_func",