From 458f4023235f1d49ea10e47fb641a051b431e438 Mon Sep 17 00:00:00 2001 From: Binh Tang Date: Thu, 2 Apr 2026 00:42:50 -0700 Subject: [PATCH] Add VACE support for WAN 2.1 conditional video generation (#1885) Signed-off-by: Binh Tang Signed-off-by: Binh Tang Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Binh Tang Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- docs/models/supported_models.md | 1 + docs/user_guide/diffusion_features.md | 1 + .../vace/vace_video_generation.md | 88 +++ .../vace/vace_video_generation.py | 209 ++++++ .../test_wan_2_1_vace_expansion.py | 161 +++++ vllm_omni/diffusion/models/wan2_2/__init__.py | 14 + .../models/wan2_2/pipeline_wan2_2.py | 8 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 645 ++++++++++++++++++ .../models/wan2_2/wan2_2_vace_transformer.py | 254 +++++++ vllm_omni/diffusion/registry.py | 7 + 10 files changed, 1386 insertions(+), 2 deletions(-) create mode 100644 examples/offline_inference/vace/vace_video_generation.md create mode 100644 examples/offline_inference/vace/vace_video_generation.py create mode 100644 tests/e2e/online_serving/test_wan_2_1_vace_expansion.py create mode 100644 vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py create mode 100644 vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0706a6786..68024e18b 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | ✅︎ | | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 607d9af73..9cd407d37 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | diff --git a/examples/offline_inference/vace/vace_video_generation.md b/examples/offline_inference/vace/vace_video_generation.md new file mode 100644 index 000000000..bbaf99452 --- /dev/null +++ b/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. diff --git a/examples/offline_inference/vace/vace_video_generation.py b/examples/offline_inference/vace/vace_video_generation.py new file mode 100644 index 000000000..6ca0d74c5 --- /dev/null +++ b/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() diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py new file mode 100644 index 000000000..0de70afe8 --- /dev/null +++ b/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, + }, + } + ) diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py index c337f58a4..d418001d9 100644 --- a/vllm_omni/diffusion/models/wan2_2/__init__.py +++ b/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", ] diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index d7d8bad52..d2d2bb860 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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 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 new file mode 100644 index 000000000..ea5233631 --- /dev/null +++ b/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) 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 new file mode 100644 index 000000000..4f4217dab --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 994dac04a..dcd227237 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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",