diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e0cc17aa1..bba4e2513 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -32,6 +32,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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | diff --git a/examples/offline_inference/speech_to_video/README.md b/examples/offline_inference/speech_to_video/README.md new file mode 100644 index 000000000..28a64fb55 --- /dev/null +++ b/examples/offline_inference/speech_to_video/README.md @@ -0,0 +1,135 @@ +# Speech-To-Video + +This example demonstrates how to generate talking-head videos from a reference image and an audio clip using the Wan2.2 Speech-to-Video (S2V) pipeline with vLLM-Omni's offline inference API. + +The S2V pipeline generates multiple video clips autoregressively — the last frames of each clip become the motion context for the next — producing a seamless video that spans the full audio duration. + +## Prerequisites + +The S2V pipeline requires `librosa` for audio loading: + +```bash +pip install librosa decord +``` + +## Local CLI Usage + +### Singing Example (480p) + +```bash +wget -O "Five Hundred Miles.png" "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png" +wget -O "Five Hundred Miles.MP3" "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3" +python speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 448 --width 832 \ + --num-inference-steps 5 \ + --output s2v_singing_480p.mp4 +``` + +## Key Arguments + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--model` | *(required)* | Path to Wan2.2 S2V model (local path or HuggingFace ID) | +| `--image` | *(required)* | Path to reference image (face/portrait) | +| `--audio` | *(required)* | Path to audio file (wav/mp3) | +| `--prompt` | `"A person speaking naturally"` | Text prompt describing the scene | +| `--negative-prompt` | *(S2V default)* | Negative prompt (built-in Chinese quality filter if not set) | +| `--height` / `--width` | *auto* | Output resolution (auto-calculated from image aspect ratio if not set, must be divisible by 64) | +| `--num-frames` | `80` | Number of frames per clip (should be divisible by 4) | +| `--num-inference-steps` | `40` | Number of denoising steps | +| `--guidance-scale` | `4.5` | Classifier-free guidance scale | +| `--flow-shift` | `3.0` | Scheduler flow shift | +| `--seed` | `42` | Random seed for reproducibility | +| `--fps` | `16` | Frames per second for the saved MP4 | +| `--init-first-frame` | `False` | Use reference image as the first frame | +| `--output` | `s2v_output.mp4` | Output video file path | +| `--vae-use-slicing` | `False` | Enable VAE slicing for memory optimization | +| `--vae-use-tiling` | `False` | Enable VAE tiling for memory optimization | +| `--tensor-parallel-size` | `1` | Number of GPUs for tensor parallelism (TP) inside the DiT | +| `--cfg-parallel-size` | `1` | Number of GPUs for CFG parallelism (1 or 2) | +| `--vae-patch-parallel-size` | `1` | Number of GPUs for VAE patch parallelism | +| `--enable-cpu-offload` | `False` | Enable CPU offloading | +| `--enable-layerwise-offload` | `False` | Enable layerwise offloading on DiT | + +> If you encounter OOM errors, try `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage, or `--enable-cpu-offload` for more aggressive memory saving. + +## Performance Tips + +The dominant cost in S2V is self-attention across the full token sequence (~80K at 720p), which scales quadratically with resolution. Strategies to speed up generation: + +| Strategy | Speedup | How | +|---|---|---| +| **Lower resolution** (480p) | ~3.5x | `--height 448 --width 832` — reduces tokens from 80K to 33K | +| **Tensor parallelism** (2 GPUs) | ~1.4x | `--tensor-parallel-size 2` — splits attention/GEMM across GPUs | +| **Fewer steps** | Linear | `--num-inference-steps 5` — trades quality for speed | +| **Combine all above** | ~5x+ | 480p + TP=2 + 5 steps | + +## Profiling + +To profile the S2V pipeline and identify performance bottlenecks: + +### Enable Built-in Profiler + +Use `--enable-diffusion-pipeline-profiler` to display stage durations: + +```bash +python speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 448 --width 832 \ + --num-inference-steps 5 \ + --enable-diffusion-pipeline-profiler \ + --output s2v_output.mp4 +``` + +### Torch Profiler (Device-Level Profiling) + +For detailed GPU/XPU kernel-level profiling with Chrome trace visualization: + +```bash +python speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 448 --width 832 \ + --num-inference-steps 1 \ + --profile-dir /tmp/s2v_profile \ + --profile-record-shapes \ + --profile-with-memory \ + ----profile-with-stack \ + --output s2v_profiled.mp4 +``` + +**Profiler Arguments:** + +| Parameter | Description | +|-----------|-------------| +| `--profile-dir` | Enable torch profiler and save traces to this directory | +| `--profile-record-shapes` | Record tensor shapes (increases trace size) | +| `--profile-with-stack` | Record stack traces (increases overhead) | +| `--profile-with-memory` | Profile memory usage | +| `--profile-with-flops` | Estimate FLOPs for operations | + +**Output Files:** +- `trace_rank*.json.gz` - Chrome trace (open in `chrome://tracing`) +- `ops_rank*.xlsx` - Excel workbook with operation statistics + +> **Note:** Use `--num-inference-steps 1` for faster profiling. The trace captures one full denoising step with all transformer blocks and attention operations. + +## How It Works + +The S2V pipeline processes in the following stages: + +1. **Pre-process**: The reference image is resized/cropped to fit the target resolution. The audio path is validated. +2. **Text encoding**: The prompt is encoded using the UMT5 text encoder. +3. **Audio encoding**: The audio file is loaded via `librosa`, processed through Wav2Vec2, and bucketed to align with video frame-rate. The audio length determines the number of clips (`num_repeat`). +4. **Reference image encoding**: The reference image is VAE-encoded into latent tokens. +5. **Multi-clip denoising**: For each clip, random noise is generated and denoised through the `WanModel_S2V` transformer with audio injection. The last frames of each decoded clip become the motion context for the next. +6. **Post-process**: All clips are concatenated and converted to pixel-space video frames. diff --git a/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py new file mode 100644 index 000000000..97ce63700 --- /dev/null +++ b/examples/offline_inference/speech_to_video/speech_to_video.py @@ -0,0 +1,402 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Speech-to-Video generation example using Wan2.2 S2V. + +Generates talking-head videos from a reference image and an audio clip +using the Wan2.2 S2V pipeline with multi-clip autoregressive generation. + +Usage: + python speech_to_video.py \ + --model /path/to/Wan2.2-S2V-14B \ + --image reference.jpg \ + --audio speech.wav \ + --prompt "A person speaking naturally" +""" + +import argparse +import os +import time +from pathlib import Path + +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.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)." + ) + parser.add_argument( + "--model", + required=True, + help="Path to Wan2.2 S2V model (local path or HuggingFace ID).", + ) + parser.add_argument("--image", required=True, help="Path to reference image.") + parser.add_argument("--audio", required=True, help="Path to audio file (wav/mp3).") + parser.add_argument( + "--prompt", + default="A person speaking naturally", + help="Text prompt describing the scene.", + ) + parser.add_argument( + "--negative-prompt", + default=None, + help="Negative prompt (uses S2V default if not set).", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed.") + parser.add_argument( + "--guidance-scale", + type=float, + default=4.5, + help="CFG scale (default: 4.5).", + ) + parser.add_argument( + "--height", + type=int, + default=None, + help="Video height (auto-calculated from reference image if not set).", + ) + parser.add_argument( + "--width", + type=int, + default=None, + help="Video width (auto-calculated from reference image if not set).", + ) + parser.add_argument( + "--num-frames", + type=int, + default=80, + help="Frames per clip (should be divisible by 4, default: 80).", + ) + parser.add_argument( + "--num-inference-steps", + type=int, + default=40, + help="Number of denoising steps (default: 40).", + ) + parser.add_argument( + "--flow-shift", + type=float, + default=3.0, + help="Scheduler flow shift (default: 3.0).", + ) + parser.add_argument( + "--output", + type=str, + default="s2v_output.mp4", + help="Path to save the output video.", + ) + parser.add_argument( + "--fps", + type=int, + default=16, + help="Frames per second for the output video (default: 16).", + ) + parser.add_argument( + "--init-first-frame", + action="store_true", + help="Use the reference image as the first frame of the video.", + ) + parser.add_argument( + "--vae-use-slicing", + action="store_true", + help="Enable VAE slicing for memory optimization.", + ) + parser.add_argument( + "--vae-use-tiling", + action="store_true", + help="Enable VAE tiling for memory optimization.", + ) + parser.add_argument( + "--enable-cpu-offload", + action="store_true", + help="Enable CPU offloading for diffusion models.", + ) + parser.add_argument( + "--enable-layerwise-offload", + action="store_true", + help="Enable layerwise (blockwise) offloading on DiT modules.", + ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", + ) + parser.add_argument( + "--cfg-parallel-size", + type=int, + default=1, + choices=[1, 2], + help="Number of GPUs used for classifier free guidance parallel size.", + ) + parser.add_argument( + "--vae-patch-parallel-size", + type=int, + default=1, + help="Number of GPUs used for VAE patch/tile parallelism (decode).", + ) + parser.add_argument( + "--ring-degree", + type=int, + default=1, + help="Number of GPUs used for ring sequence parallelism.", + ) + parser.add_argument( + "--ulysses-degree", + type=int, + default=1, + help="Number of GPUs used for ulysses sequence parallelism.", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable torch.compile and force eager execution.", + ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) + parser.add_argument( + "--cache-backend", + type=str, + default=None, + choices=["cache_dit"], + help="Cache backend for acceleration. Default: None.", + ) + parser.add_argument( + "--enable-cache-dit-summary", + action="store_true", + help="Enable cache-dit summary logging after diffusion forward passes.", + ) + parser.add_argument( + "--profile-dir", + type=str, + default=None, + help="Enable torch profiler and save traces to this directory.", + ) + parser.add_argument( + "--profile-record-shapes", + action="store_true", + help="Record tensor shapes in profiler (increases trace size).", + ) + parser.add_argument( + "--profile-with-stack", + action="store_true", + help="Record stack traces in profiler (increases overhead).", + ) + parser.add_argument( + "--profile-with-memory", + action="store_true", + help="Profile memory usage.", + ) + parser.add_argument( + "--profile-with-flops", + action="store_true", + help="Estimate FLOPs for operations.", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + + # Load reference image + import PIL.Image + + image = PIL.Image.open(args.image).convert("RGB") + + # Cache-dit config + cache_config = None + if args.cache_backend == "cache_dit": + cache_config = { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "max_cached_steps": 20, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": False, + "taylorseer_order": 1, + "scm_steps_mask_policy": None, + "scm_steps_policy": "dynamic", + } + + parallel_config = DiffusionParallelConfig( + tensor_parallel_size=args.tensor_parallel_size, + cfg_parallel_size=args.cfg_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, + ring_degree=args.ring_degree, + ulysses_degree=args.ulysses_degree, + ) + + # Check if profiling is requested via CLI args or environment variable + profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR") + profiler_enabled = bool(profile_dir) + profiler_config = None + if profiler_enabled: + from vllm.config import ProfilerConfig + + profiler_config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=profile_dir, + torch_profiler_record_shapes=args.profile_record_shapes, + torch_profiler_with_stack=args.profile_with_stack, + torch_profiler_with_memory=args.profile_with_memory, + torch_profiler_with_flops=args.profile_with_flops, + ) + + omni = Omni( + model=args.model, + model_class_name="WanS2VPipeline", + flow_shift=args.flow_shift, + vae_use_slicing=args.vae_use_slicing, + vae_use_tiling=args.vae_use_tiling, + enable_cpu_offload=args.enable_cpu_offload, + enable_layerwise_offload=args.enable_layerwise_offload, + parallel_config=parallel_config, + enforce_eager=args.enforce_eager, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + cache_backend=args.cache_backend, + cache_config=cache_config, + enable_cache_dit_summary=args.enable_cache_dit_summary, + profiler_config=profiler_config, + ) + + # Print generation configuration + print(f"\n{'=' * 60}") + print("Generation Configuration:") + print(f" Model: {args.model}") + print(f" Reference image: {args.image}") + print(f" Audio: {args.audio}") + print(f" Inference steps: {args.num_inference_steps}") + print(f" Frames per clip: {args.num_frames}") + print(f" Guidance scale: {args.guidance_scale}") + print(f" Flow shift: {args.flow_shift}") + print(f" Init first frame: {args.init_first_frame}") + if args.height and args.width: + print(f" Video size: {args.width}x{args.height}") + else: + print(" Video size: auto (from reference image)") + print(f"{'=' * 60}\n") + + # Start profiling if enabled + if profiler_enabled: + print("[Profiler] Starting profiling...") + omni.start_profile() + + generation_start = time.perf_counter() + + result = omni.generate( + { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "multi_modal_data": { + "image": image, + "audio": args.audio, + "init_first_frame": args.init_first_frame, + }, + }, + OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_frames=args.num_frames, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + seed=args.seed, + generator=generator, + ), + ) + + generation_end = time.perf_counter() + generation_time = generation_end - generation_start + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") + + # Stop profiling if enabled + if profiler_enabled: + print("\n[Profiler] Stopping profiler and collecting results...") + profile_results = omni.stop_profile() + if profile_results and isinstance(profile_results, list): + print("\n" + "=" * 60) + print("PROFILING RESULTS:") + for i, result in enumerate(profile_results): + print(f"\nStage {i}:") + if result: + print(f" {result}") + print("=" * 60) + else: + print("[Profiler] No valid profiling data returned.") + + # Extract output from result + output = OmniRequestOutput.unwrap_result(result) + + if not output.images: + raise ValueError("No video frames found in OmniRequestOutput.") + + # Extract audio from multimodal_output (set by pipeline post-processor) + mm = output.multimodal_output or {} + audio_waveform = mm.get("audio") + output_fps = float(mm.get("fps", args.fps)) + output_sr = int(mm.get("audio_sample_rate", 16000)) + + if audio_waveform is not None: + print(f"Audio waveform: shape={audio_waveform.shape}, sr={output_sr}") + else: + print("Warning: no audio waveform in pipeline output") + + # Normalize frames to (T, H, W, C) uint8 numpy array + import numpy as np + + def _flatten_to_array(data): + """Unwrap to a single (T, H, W, C) numpy array.""" + if isinstance(data, np.ndarray): + if data.ndim == 5: + return data[0] # (B, T, H, W, C) → (T, H, W, C) + if data.ndim == 4: + return data # already (T, H, W, C) + if isinstance(data, list) and data: + first_elem = data[0] + if isinstance(first_elem, np.ndarray) and first_elem.ndim >= 4: + return _flatten_to_array(first_elem) + if isinstance(first_elem, np.ndarray) and first_elem.ndim == 3: + return np.stack(data) # list of (H, W, C) → (T, H, W, C) + return data + + video_frames = _flatten_to_array(output.images) + # postprocess_video returns float32 in [0,1]; mux_video_audio_bytes needs uint8 + if isinstance(video_frames, np.ndarray) and video_frames.dtype != np.uint8: + video_frames = (np.clip(video_frames, 0, 1) * 255).astype(np.uint8) + num_frames = video_frames.shape[0] if isinstance(video_frames, np.ndarray) else len(video_frames) + + # Save output video with audio + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes + + video_bytes = mux_video_audio_bytes( + video_frames, + audio_waveform, + fps=output_fps, + audio_sample_rate=output_sr, + ) + with open(output_path, "wb") as f: + f.write(video_bytes) + + print(f"Saved generated video to {output_path}") + print(f"Video has {num_frames} frames at {output_fps} fps ({num_frames / output_fps:.1f}s)") + if audio_waveform is not None: + print(f"Audio: {len(audio_waveform) / output_sr:.1f}s at {output_sr} Hz") + + +if __name__ == "__main__": + main() diff --git a/recipes/Wan-AI/Wan2.2-S2V.md b/recipes/Wan-AI/Wan2.2-S2V.md new file mode 100644 index 000000000..34470ed0a --- /dev/null +++ b/recipes/Wan-AI/Wan2.2-S2V.md @@ -0,0 +1,206 @@ +# Wan2.2 Speech To Video + +## Summary + +- Vendor: Wan-AI +- Model: `Wan-AI/Wan2.2-S2V-14B` +- Task: Speech-to-video generation (talking-head / singing from reference image + audio) +- Mode: Offline inference +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to generate talking-head or singing videos from a +reference image and audio clip using the Wan2.2 S2V 14B model with vLLM-Omni. +Two configurations are provided: + +1. **Multi-GPU with tensor parallelism** — splits the DiT across GPUs for faster + per-step latency. +2. **Multi-GPU with tensor parallelism + CPU offload** — enables generation on + GPUs with limited VRAM by offloading idle model components to CPU. + +## References + +- Upstream model card: +- Example reference assets: + +## Hardware Support + +## CUDA + +### 2× NVIDIA A100/H100 (80 GB) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver: NVIDIA driver with CUDA 12.x +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Prerequisites + +None + + +#### Command + +```bash +wget -O "Five Hundred Miles.png" \ + "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png" +wget -O "Five Hundred Miles.MP3" \ + "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3" +``` + +**TP=2 (no CPU offload, requires ~60 GB VRAM per GPU):** + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +python examples/offline_inference/speech_to_video/speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 720 --width 1280 --num-frames 81 \ + --num-inference-steps 40 \ + --fps 16 \ + --tensor-parallel-size 2 \ + --vae-use-slicing --vae-use-tiling \ + --output s2v_720p_tp2.mp4 +``` + +**TP=1 with CPU offload (reduces VRAM usage, adds ~12s overhead):** + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +python examples/offline_inference/speech_to_video/speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 720 --width 1280 --num-frames 81 \ + --num-inference-steps 40 \ + --fps 16 \ + --tensor-parallel-size 1 \ + --enable-cpu-offload \ + --vae-use-slicing --vae-use-tiling \ + --output s2v_720p_tp1_offload.mp4 +``` + +**TP=1 * cfg=2 with CPU offload:** + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +python examples/offline_inference/speech_to_video/speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 720 --width 1280 --num-frames 81 \ + --num-inference-steps 40 \ + --fps 16 \ + --tensor-parallel-size 1 \ + --enable-cpu-offload \ + --cfg-parallel-size 2 \ + --vae-use-slicing --vae-use-tiling \ + --output s2v_720p_tp1_cfg2_offload.mp4 +``` + +## XPU + +### 4x or 8x Intel Arc B70 (32 GB) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- torch: 2.11.0+xpu +- triton-xpu: 3.7.0 +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Prerequisites + +None + + +#### Command + +```bash +wget -O "Five Hundred Miles.png" \ + "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png" +wget -O "Five Hundred Miles.MP3" \ + "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3" +``` + +**TP=4:** + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +python examples/offline_inference/speech_to_video/speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 720 --width 1280 --num-frames 81 \ + --num-inference-steps 40 \ + --fps 16 \ + --tensor-parallel-size 4 \ + --vae-use-slicing --vae-use-tiling \ + --output s2v_720p_tp4.mp4 +``` + +**TP=4 cfg=2:** + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +python examples/offline_inference/speech_to_video/speech_to_video.py \ + --model Wan-AI/Wan2.2-S2V-14B \ + --image "Five Hundred Miles.png" \ + --audio "Five Hundred Miles.MP3" \ + --prompt "A person singing" \ + --height 720 --width 1280 --num-frames 81 \ + --num-inference-steps 40 \ + --fps 16 \ + --tensor-parallel-size 4 \ + --cfg-parallel-size 2 \ + --vae-use-slicing --vae-use-tiling \ + --output s2v_720p_tp4_cfg2_offload.mp4 +``` + + +#### Verification + +The output video should show a talking/singing person matching the reference +image with lip movements synchronized to the audio. Check: + +1. Video file is generated at the specified output path. +2. Audio is muxed into the output MP4. +3. Lip sync quality is reasonable for the given inference steps. + +#### Notes + +- **Key flags:** + - `--tensor-parallel-size ` — splits the DiT across N GPUs for TP. + - `--cfg-parallel-size ` - splits the conditioned / unconditioned pred to 2 GPUs run parallel. + - `--enable-cpu-offload` — enables model-level CPU offloading (transformer + and text_encoder alternate on GPU). Adds ~12s overhead but reduces VRAM. + - `--vae-use-slicing --vae-use-tiling` — reduces VAE peak memory by + processing frames in slices and spatial tiles. + - `--num-frames 81` — generates ~5s of video at 16 fps. + - `--num-inference-steps` — quality/speed tradeoff. 4 steps for fast preview, + 40 steps for production quality. +- **Performance tips:** + - Use `--enforce-eager` to skip torch.compile if you encounter recompilation + issues or want deterministic profiling. + - For profiling, add `--enable-diffusion-pipeline-profiler` for stage-level + timing, or `--profiler-config '{"profiler":"torch","torch_profiler_dir":"./perf","torch_profiler_record_shapes":true,"torch_profiler_with_stack":false}'` + for detailed op-level torch profiler traces. + - The first clip in multi-clip generation includes warmup overhead; subsequent + clips run at steady-state speed. +- **Known limitations:** + - `VLLM_WORKER_MULTIPROC_METHOD=spawn` is required for multi-GPU offline + inference to avoid CUDA context issues with forked processes. + - CPU offload adds ~12s latency for model-level shuttling between CPU and GPU. + - S2V self-attention is ~8.5× more expensive than T2V due to audio-visual + conditioning fused into self-attn blocks. diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 18f627bc8..9491b759a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -761,6 +761,11 @@ class OmniDiffusionConfig: self.model_class_name = "NextStep11Pipeline" self.tf_model_config = TransformerConfig() self.update_multimodal_support() + elif model_type == "s2v": + if self.model_class_name is None: + self.model_class_name = "WanS2VPipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() elif architectures and len(architectures) == 1: self.model_class_name = architectures[0] else: diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py index ecbfdd472..2159a1184 100644 --- a/vllm_omni/diffusion/models/wan2_2/__init__.py +++ b/vllm_omni/diffusion/models/wan2_2/__init__.py @@ -17,6 +17,11 @@ from .pipeline_wan2_2_i2v import ( get_wan22_i2v_post_process_func, get_wan22_i2v_pre_process_func, ) +from .pipeline_wan2_2_s2v import ( + Wan22S2VPipeline, + get_wan22_s2v_post_process_func, + get_wan22_s2v_pre_process_func, +) from .pipeline_wan2_2_ti2v import ( Wan22TI2VPipeline, get_wan22_ti2v_post_process_func, @@ -42,6 +47,9 @@ __all__ = [ "WanI2VDMD2Pipeline", "get_wan22_i2v_post_process_func", "get_wan22_i2v_pre_process_func", + "Wan22S2VPipeline", + "get_wan22_s2v_post_process_func", + "get_wan22_s2v_pre_process_func", "Wan22TI2VPipeline", "get_wan22_ti2v_post_process_func", "get_wan22_ti2v_pre_process_func", diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py new file mode 100644 index 000000000..0817d6f48 --- /dev/null +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -0,0 +1,1357 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Wan2.2 Speech-to-Video (S2V) Pipeline for vLLM-Omni. + +Migrated from Wan2.2/wan/speech2video.py (WanS2V class). +""" + +from __future__ import annotations + +import logging +import math +import os +from collections.abc import Iterable +from typing import Any + +import numpy as np +import PIL.Image +import torch +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from torchvision import transforms +from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel, Wav2Vec2ForCTC, Wav2Vec2Processor +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + load_transformer_config, + retrieve_latents, +) +from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( + create_s2v_transformer_from_config, +) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniTextPrompt +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + +# Default negative prompt from Wan2.2 S2V config (shared_config + s2v_14B) +_S2V_DEFAULT_NEG_PROMPT = ( + "画面模糊，最差质量，画面模糊，细节模糊不清，情绪激动剧烈，手快速抖动，字幕，" + "丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，" + "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" +) + +# --------------------------------------------------------------------------- +# Audio utility helpers (ported from Wan2.2/wan/modules/s2v/audio_encoder.py) +# --------------------------------------------------------------------------- + +_AUDIO_VIDEO_RATE = 30 + + +def _linear_interpolation( + features: torch.Tensor, + input_fps: float, + output_fps: float, + output_len: int | None = None, +) -> torch.Tensor: + """Resample audio features from *input_fps* to *output_fps*. + + Args: + features: Shape ``[1, T, D]``. + input_fps: Source frame-rate of *features*. + output_fps: Target frame-rate. + output_len: If given, force output length; otherwise derived from + ``features.shape[1] / input_fps * output_fps``. + + Returns: + Resampled tensor with shape ``[1, output_len, D]``. + """ + features = features.transpose(1, 2) # [1, D, T] + if output_len is None: + seq_len = features.shape[2] / float(input_fps) + output_len = int(seq_len * output_fps) + output_features = torch.nn.functional.interpolate(features, size=output_len, align_corners=True, mode="linear") + return output_features.transpose(1, 2) # [1, output_len, D] + + +def _get_sample_indices( + original_fps: float, + total_frames: int, + target_fps: float, + num_sample: int, + fixed_start: int | None = None, +) -> np.ndarray: + """Compute frame indices to sample *num_sample* frames at *target_fps*.""" + required_duration = num_sample / target_fps + required_origin_frames = int(np.ceil(required_duration * original_fps)) + if required_duration > total_frames / original_fps: + raise ValueError("required_duration must be less than video length") + + if fixed_start is not None and fixed_start >= 0: + start_frame = fixed_start + else: + max_start = total_frames - required_origin_frames + if max_start < 0: + raise ValueError("video length is too short") + start_frame = np.random.randint(0, max_start + 1) + start_time = start_frame / original_fps + + end_time = start_time + required_duration + time_points = np.linspace(start_time, end_time, num_sample, endpoint=False) + + frame_indices = np.round(np.array(time_points) * original_fps).astype(int) + frame_indices = np.clip(frame_indices, 0, total_frames - 1) + return frame_indices + + +def _get_audio_embed_bucket_fps( + audio_embed: torch.Tensor, + fps: int = 16, + batch_frames: int = 80, + m: int = 0, +) -> tuple[torch.Tensor, int, int]: + """Bucket audio embeddings to align with video frames at *fps*. + + Args: + audio_embed: Shape ``[num_layers, audio_frame_num, audio_dim]``. + fps: Video frame-rate. + batch_frames: Number of video frames per clip (``infer_frames``). + m: Context window radius for each audio sample. + + Returns: + ``(audio_embed_bucket, num_repeat, target_video_frames)`` where + *num_repeat* is the number of clips needed and *target_video_frames* + is the exact number of video frames to match audio duration. + """ + num_layers, audio_frame_num, audio_dim = audio_embed.shape + return_all_layers = num_layers > 1 + + scale = _AUDIO_VIDEO_RATE / fps + # Calculate exact video frames needed to match audio duration + target_video_frames = int(audio_frame_num / scale) + # Number of clips needed (use ceil to cover all audio) + min_batch_num = math.ceil(target_video_frames / batch_frames) + # Ensure at least 1 clip + min_batch_num = max(1, min_batch_num) + bucket_num = min_batch_num * batch_frames + + pad_audio_num = math.ceil(min_batch_num * batch_frames / fps * _AUDIO_VIDEO_RATE) - audio_frame_num + batch_idx = _get_sample_indices( + original_fps=_AUDIO_VIDEO_RATE, + total_frames=audio_frame_num + pad_audio_num, + target_fps=fps, + num_sample=bucket_num, + fixed_start=0, + ) + audio_sample_stride = int(_AUDIO_VIDEO_RATE / fps) + + batch_audio_eb = [] + for bi in batch_idx: + if bi < audio_frame_num: + chosen_idx = list( + range(bi - m * audio_sample_stride, bi + (m + 1) * audio_sample_stride, audio_sample_stride) + ) + chosen_idx = [max(0, c) for c in chosen_idx] + chosen_idx = [min(audio_frame_num - 1, c) for c in chosen_idx] + + if return_all_layers: + frame_audio_embed = audio_embed[:, chosen_idx].flatten(start_dim=-2, end_dim=-1) + else: + frame_audio_embed = audio_embed[0][chosen_idx].flatten() + else: + if return_all_layers: + frame_audio_embed = torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device) + else: + frame_audio_embed = torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) + batch_audio_eb.append(frame_audio_embed) + + batch_audio_eb = torch.stack(batch_audio_eb, dim=0) + return batch_audio_eb, min_batch_num, target_video_frames + + +# --------------------------------------------------------------------------- +# Size calculation helper (ported from WanS2V.get_size_less_than_area) +# --------------------------------------------------------------------------- + + +def _get_size_less_than_area( + height: int, + width: int, + target_area: int = 1024 * 704, + divisor: int = 64, +) -> tuple[int, int]: + """Compute (H, W) that fits within *target_area*, padded to *divisor*.""" + if height * width <= target_area: + max_upper_area = target_area + min_scale = 0.1 + max_scale = 1.0 + else: + max_upper_area = target_area + d = divisor - 1 + b = d * (height + width) + a = height * width + c = d**2 - max_upper_area + min_scale = (-b + math.sqrt(b**2 - 2 * a * c)) / (2 * a) + max_scale = math.sqrt(max_upper_area / (height * width)) + + for i in range(100): + scale = max_scale - (max_scale - min_scale) * i / 100 + new_height, new_width = int(height * scale), int(width * scale) + pad_height = (divisor - new_height % divisor) % divisor + pad_width = (divisor - new_width % divisor) % divisor + padded_height, padded_width = new_height + pad_height, new_width + pad_width + if padded_height * padded_width <= max_upper_area: + return padded_height, padded_width + + # Fallback + aspect_ratio = width / height + target_width = int((target_area * aspect_ratio) ** 0.5 // divisor * divisor) + target_height = int((target_area / aspect_ratio) ** 0.5 // divisor * divisor) + if target_width >= width or target_height >= height: + target_width = int(width // divisor * divisor) + target_height = int(height // divisor * divisor) + return target_height, target_width + + +# --------------------------------------------------------------------------- +# Pre / Post process functions (registered via registry.py) +# --------------------------------------------------------------------------- + + +def get_wan22_s2v_post_process_func( + od_config: OmniDiffusionConfig, +): + from diffusers.video_processor import VideoProcessor + + video_processor = VideoProcessor(vae_scale_factor=8) + + def post_process_func( + output, + output_type: str = "np", + ): + # output is (video_tensor, audio_waveform_np, audio_sample_rate) + if isinstance(output, tuple) and len(output) == 3: + video, audio_waveform, audio_sr = output + else: + video = output + audio_waveform = None + audio_sr = None + + if output_type == "latent": + return video + + processed_video = video_processor.postprocess_video(video, output_type=output_type) + + if audio_waveform is not None: + return { + "video": processed_video, + "audio": audio_waveform, + "audio_sample_rate": audio_sr, + "fps": 16, + } + return processed_video + + return post_process_func + + +def get_wan22_s2v_pre_process_func( + od_config: OmniDiffusionConfig, +): + """Pre-process function for S2V: load ref image, compute target size. + + Expects ``multi_modal_data`` to contain: + - ``"image"``: reference image (PIL.Image or file path) + - ``"audio"``: audio file path (str) + + Optionally: + - ``"pose_video"``: pose conditioning video path (str) + - ``"init_first_frame"``: bool, use ref image as first frame + """ + + 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"] = {} + + # -- Reference image -- + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if raw_image is None: + raise ValueError( + "No reference image provided. S2V requires a reference image. " + 'Set `"multi_modal_data": {"image": , ...}`' + ) + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image).convert("RGB") + elif isinstance(raw_image, PIL.Image.Image): + image = raw_image + else: + raise TypeError(f"Unsupported image type {type(raw_image)}") + + # -- Audio -- + raw_audio = multi_modal_data.get("audio", None) if multi_modal_data is not None else None + if raw_audio is None: + raise ValueError( + "No audio provided. S2V requires an audio file path. " + 'Set `"multi_modal_data": {"audio": "", ...}`' + ) + + # -- Compute target size -- + max_area = 720 * 1280 + if request.sampling_params.height is not None and request.sampling_params.width is not None: + height, width = request.sampling_params.height, request.sampling_params.width + else: + ref_h, ref_w = image.height, image.width + height, width = _get_size_less_than_area(ref_h, ref_w, target_area=max_area) + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + # Resize + center-crop reference image to target size + resize_op = transforms.Resize(min(height, width)) + crop_op = transforms.CenterCrop((height, width)) + ref_pil = crop_op(resize_op(image)) + + prompt["multi_modal_data"]["image"] = ref_pil + prompt["additional_information"]["audio_path"] = raw_audio + prompt["additional_information"]["pose_video"] = ( + multi_modal_data.get("pose_video", None) if multi_modal_data is not None else None + ) + prompt["additional_information"]["init_first_frame"] = ( + multi_modal_data.get("init_first_frame", False) if multi_modal_data is not None else False + ) + request.prompts[i] = prompt + + return request + + return pre_process_func + + +# --------------------------------------------------------------------------- +# Model format detection +# --------------------------------------------------------------------------- + + +def _is_diffusers_format(model_path: str) -> bool: + """Check if a model directory uses diffusers subfolder layout.""" + return os.path.isdir(os.path.join(model_path, "tokenizer")) and os.path.isdir(os.path.join(model_path, "vae")) + + +def _resolve_model_path(model: str) -> str: + """Resolve a HF repo ID or local path to a local directory. + + For local paths, returns as-is. For HF repo IDs, downloads needed files + and returns the snapshot directory. + """ + if os.path.isdir(model): + return model + from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + + # S2V models use a non-standard repo layout with files at root level. + # Download only the files needed for inference: + # - Transformer weights (sharded safetensors + index) + # - VAE weights + # - T5 encoder weights + tokenizer config + # - Audio encoder (wav2vec2) weights + config + # - Model config files + allow_patterns = [ + "*.safetensors", # Transformer weights (sharded) + "*.safetensors.index.json", # Weight shard index + "*.pth", # VAE and T5 weights + "*.json", # Config files (config.json, configuration.json) + "google/**", # T5 tokenizer/config directory + "wav2vec2-large-xlsr-53-english/**", # Audio encoder directory + ] + return download_weights_from_hf_specific(model, None, allow_patterns) + + +# --------------------------------------------------------------------------- +# Wan T5 → UMT5EncoderModel weight conversion +# --------------------------------------------------------------------------- + +# UMT5Config matching the Wan2.2 T5 encoder architecture (umt5-xxl variant). +_WAN_UMT5_CONFIG = UMT5Config( + vocab_size=256384, + d_model=4096, + d_kv=64, + d_ff=10240, + num_heads=64, + num_layers=24, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + dense_act_fn="gelu_new", + is_gated_act=True, + is_encoder_decoder=False, +) + + +def _convert_wan_t5_state_dict(wan_sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Convert a Wan2.2 T5 encoder state dict to HuggingFace UMT5 format. + + The Wan checkpoint uses flat ``blocks.{i}.attn.*`` naming whereas + HuggingFace ``UMT5EncoderModel`` expects ``encoder.block.{i}.layer.*`` + prefixed keys. All tensor shapes are identical; only key names change. + """ + hf_sd: dict[str, torch.Tensor] = {} + + # Embeddings (Wan has 1 copy; UMT5 ties shared + encoder.embed_tokens) + hf_sd["shared.weight"] = wan_sd["token_embedding.weight"] + hf_sd["encoder.embed_tokens.weight"] = wan_sd["token_embedding.weight"] + + # Final layer norm + hf_sd["encoder.final_layer_norm.weight"] = wan_sd["norm.weight"] + + # Per-block weights + num_layers = _WAN_UMT5_CONFIG.num_layers + for i in range(num_layers): + src = f"blocks.{i}" + dst = f"encoder.block.{i}" + + # Self-attention + for proj in ("q", "k", "v", "o"): + hf_sd[f"{dst}.layer.0.SelfAttention.{proj}.weight"] = wan_sd[f"{src}.attn.{proj}.weight"] + hf_sd[f"{dst}.layer.0.SelfAttention.relative_attention_bias.weight"] = wan_sd[ + f"{src}.pos_embedding.embedding.weight" + ] + hf_sd[f"{dst}.layer.0.layer_norm.weight"] = wan_sd[f"{src}.norm1.weight"] + + # Gated feed-forward (gate.0 → wi_0, fc1 → wi_1, fc2 → wo) + hf_sd[f"{dst}.layer.1.DenseReluDense.wi_0.weight"] = wan_sd[f"{src}.ffn.gate.0.weight"] + hf_sd[f"{dst}.layer.1.DenseReluDense.wi_1.weight"] = wan_sd[f"{src}.ffn.fc1.weight"] + hf_sd[f"{dst}.layer.1.DenseReluDense.wo.weight"] = wan_sd[f"{src}.ffn.fc2.weight"] + hf_sd[f"{dst}.layer.1.layer_norm.weight"] = wan_sd[f"{src}.norm2.weight"] + + return hf_sd + + +def _load_wan_t5_as_umt5( + model: UMT5EncoderModel, + checkpoint_path: str, + dtype: torch.dtype = torch.bfloat16, +) -> UMT5EncoderModel: + """Load a Wan2.2 T5 ``.pth`` checkpoint as a ``UMT5EncoderModel``.""" + logger.info("Loading Wan T5 checkpoint: %s", checkpoint_path) + wan_sd = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + hf_sd = _convert_wan_t5_state_dict(wan_sd) + + model.load_state_dict(hf_sd, assign=True) + model = model.to(dtype=dtype).eval() + + return model + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- + + +class Wan22S2VPipeline( + nn.Module, SupportImageInput, SupportAudioInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin +): + """ + Wan2.2 Speech-to-Video Pipeline. + + Migrated from ``Wan2.2/wan/speech2video.py`` (``WanS2V``). + + Key differences from I2V: + - Single transformer (``WanModel_S2V``), no MoE boundary switching. + - Audio conditioning via wav2vec2 features injected into transformer. + - Motion frame autoregressive chaining across multiple clips. + - Optional pose video conditioning. + - Reference image encoded as separate ``ref_latents`` tokens (not + channel-concatenated like I2V). + """ + + # Default config values from Wan2.2/wan/configs/wan_s2v_14B.py + _DEFAULT_MOTION_FRAMES = 73 + _DEFAULT_INFER_FRAMES = 80 + _DEFAULT_FPS = 16 + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + + self.device = get_local_device() + dtype = getattr(od_config, "dtype", torch.bfloat16) + + model = od_config.model + model_path = _resolve_model_path(model) + is_diffusers = _is_diffusers_format(model_path) + + if is_diffusers: + self._init_diffusers_format(model_path, dtype) + else: + self._init_original_format(model_path, dtype) + + # -- Scheduler -- + flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0 # S2V default + self.scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, + shift=flow_shift, + prediction_type="flow_prediction", + ) + + # -- VAE scale factors -- + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if hasattr(self.vae, "config") else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if hasattr(self.vae, "config") else 8 + + # -- Resolution divisor for padding (VAE spatial scale × transformer patch size) -- + # S2V uses patch_size=(1,2,2) in patch_embedding (3D conv with stride 2 in spatial dims) + # Total divisor = vae_scale_factor_spatial × patch_size_spatial = 8 × 2 = 16 + self.resolution_divisor = self.vae_scale_factor_spatial * 2 # 2 from patch_size spatial stride + + # -- S2V specific config -- + # These mirror Wan2.2/wan/configs/wan_s2v_14B.py + self.motion_frames = getattr(od_config, "motion_frames", self._DEFAULT_MOTION_FRAMES) + self.drop_first_motion = getattr(od_config, "drop_first_motion", True) + self.fps = getattr(od_config, "fps", self._DEFAULT_FPS) or self._DEFAULT_FPS + self.audio_sample_m = 0 + + self._guidance_scale = None + self._num_timesteps = None + self._current_timestep = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + # ------------------------------------------------------------------ + # Format-specific initialization + # ------------------------------------------------------------------ + + def _init_diffusers_format(self, model_path: str, dtype: torch.dtype) -> None: + """Initialize from a diffusers-format model with subfolders.""" + # -- Weight sources for transformer (loaded via load_weights) -- + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + + # -- Text encoder -- + self.tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") + self.text_encoder = UMT5EncoderModel.from_pretrained( + model_path, subfolder="text_encoder", torch_dtype=dtype + ).to(self.device) + + # -- VAE -- + self.vae = DistributedAutoencoderKLWan.from_pretrained(model_path, subfolder="vae", torch_dtype=dtype).to( + self.device + ) + + # -- Audio encoder (wav2vec2) -- + wav2vec_path = os.path.join(model_path, "wav2vec2-large-xlsr-53-english") + if not os.path.isdir(wav2vec_path): + wav2vec_path = "facebook/wav2vec2-large-xlsr-53" + self.audio_processor = Wav2Vec2Processor.from_pretrained(wav2vec_path) + self.audio_model = Wav2Vec2ForCTC.from_pretrained(wav2vec_path).to(self.device) + self.audio_model.eval() + + # -- Transformer (S2V variant, native vllm-omni ops) -- + transformer_config = load_transformer_config(model_path, "transformer", local_files_only=True) + self.transformer = create_s2v_transformer_from_config(transformer_config) + + def _init_original_format(self, model_path: str, dtype: torch.dtype) -> None: + """Initialize from the original Wan2.2 checkpoint layout. + + Expected structure:: + + model_path/ + ├── config.json # transformer config + ├── diffusion_pytorch_model-*.safetensors # transformer weights + ├── Wan2.1_VAE.pth # VAE weights + ├── models_t5_umt5-xxl-enc-bf16.pth # T5 encoder weights + ├── google/umt5-xxl/ # tokenizer + └── wav2vec2-large-xlsr-53-english/ # audio encoder + """ + # -- Weight sources: transformer weights live at root -- + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder=None, + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + + # -- Text encoder (original Wan2.2 T5, converted to UMT5EncoderModel) -- + tokenizer_path = os.path.join(model_path, "google", "umt5-xxl") + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + + t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth") + self.text_encoder = UMT5EncoderModel(_WAN_UMT5_CONFIG) + self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype).to(self.device) + + # -- VAE (original Wan2.1 VAE, loaded via diffusers from_single_file) -- + vae_pth = os.path.join(model_path, "Wan2.1_VAE.pth") + self.vae = DistributedAutoencoderKLWan.from_single_file(vae_pth, torch_dtype=dtype) + self.vae.init_distributed() + self.vae = self.vae.to(self.device) + + # -- Audio encoder (wav2vec2) -- + wav2vec_path = os.path.join(model_path, "wav2vec2-large-xlsr-53-english") + if not os.path.isdir(wav2vec_path): + wav2vec_path = "facebook/wav2vec2-large-xlsr-53" + self.audio_processor = Wav2Vec2Processor.from_pretrained(wav2vec_path) + self.audio_model = Wav2Vec2ForCTC.from_pretrained(wav2vec_path).to(self.device) + self.audio_model.eval() + + # -- Transformer (S2V variant, native vllm-omni ops) -- + transformer_config = load_transformer_config(model_path, ".", local_files_only=True) + self.transformer = create_s2v_transformer_from_config(transformer_config) + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale is not None and self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + # ------------------------------------------------------------------ + # VAE latent normalization + # ------------------------------------------------------------------ + + def _normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Apply (latent - mean) * (1/std) normalization. + + The original Wan2_1_VAE does this internally, but diffusers' + AutoencoderKLWan stores the values in config without applying them. + """ + if not hasattr(self.vae, "config") or not hasattr(self.vae.config, "latents_mean"): + return latents + mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + inv_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 + ) + return (latents - mean) * inv_std + + def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Reverse normalization before VAE decode: latent / (1/std) + mean.""" + if not hasattr(self.vae, "config") or not hasattr(self.vae.config, "latents_mean"): + return latents + mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + inv_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 + ) + return latents / inv_std + mean + + # ------------------------------------------------------------------ + # Encoding helpers + # ------------------------------------------------------------------ + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 512, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Encode text prompts using T5 text encoder. + + Identical to Wan22I2VPipeline.encode_prompt. + """ + device = device or self.device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt_clean = [self._prompt_clean(p) for p in prompt] + batch_size = len(prompt_clean) + + text_inputs = self.tokenizer( + prompt_clean, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], + dim=0, + ) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + neg_text_inputs = self.tokenizer( + [self._prompt_clean(p) for p in negative_prompt], + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask + seq_lens_neg = mask_neg.gt(0).sum(dim=1).long() + negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] + negative_prompt_embeds = torch.stack( + [ + torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + for u in negative_prompt_embeds + ], + dim=0, + ) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + @staticmethod + def _prompt_clean(text: str) -> str: + return " ".join(text.strip().split()) + + def encode_audio( + self, + audio_path: str | np.ndarray, + infer_frames: int, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, int, int]: + """Extract wav2vec2 audio features and bucket them to video frame-rate. + + Ported from ``WanS2V.encode_audio`` and ``AudioEncoder``. + + Args: + audio_path: Path to audio file, or raw numpy audio array (16 kHz). + infer_frames: Number of video frames per clip. + + Returns: + ``(audio_embed_bucket, num_repeat, target_video_frames)`` — audio + embeddings aligned to video frames with shape + ``[1, num_layers, C_a, T_total]``, the number of clips needed, and + the exact number of video frames to match audio duration. + """ + device = device or self.device + dtype = dtype or self.transformer.dtype + + # Extract wav2vec2 features + if isinstance(audio_path, np.ndarray): + audio_input = audio_path.astype(np.float32) + sample_rate = 16000 + else: + audio_input, sample_rate = load_audio(audio_path, sr=16000, mono=True) + input_values = self.audio_processor(audio_input, sampling_rate=sample_rate, return_tensors="pt").input_values + # Match dtype of audio model weights (may be bfloat16 in bundled checkpoints) + audio_model_dtype = next(self.audio_model.parameters()).dtype + res = self.audio_model( + input_values.to(device=self.audio_model.device, dtype=audio_model_dtype), + output_hidden_states=True, + ) + feat = torch.cat(res.hidden_states) # [num_layers, T, D] + feat = _linear_interpolation(feat, input_fps=50, output_fps=_AUDIO_VIDEO_RATE) + + # Bucket to video fps + audio_embed_bucket, num_repeat, target_video_frames = _get_audio_embed_bucket_fps( + feat.to(torch.float32), + fps=self.fps, + batch_frames=infer_frames, + m=self.audio_sample_m, + ) + audio_embed_bucket = audio_embed_bucket.to(device, dtype) + audio_embed_bucket = audio_embed_bucket.unsqueeze(0) # [1, T_bucket, ...] + + # Permute to match model expectation: [B, num_layers, C_a, T_a] + if len(audio_embed_bucket.shape) == 3: + audio_embed_bucket = audio_embed_bucket.permute(0, 2, 1) + elif len(audio_embed_bucket.shape) == 4: + audio_embed_bucket = audio_embed_bucket.permute(0, 2, 3, 1) + + return audio_embed_bucket, num_repeat, target_video_frames + + def encode_ref_image( + self, + image: PIL.Image.Image, + height: int, + width: int, + device: torch.device | None = None, + ) -> torch.Tensor: + """VAE-encode the reference image into latent space. + + Ported from ``WanS2V.generate`` lines 492-499. + + Args: + image: Reference PIL image (already resized/cropped). + height: Target height. + width: Target width. + + Returns: + ``ref_latents`` with shape ``[1, C, 1, H_lat, W_lat]``. + """ + device = device or self.device + + tensor_trans = transforms.ToTensor() + ref_pixel_values = tensor_trans(image) # [C, H, W] + ref_pixel_values = ref_pixel_values.unsqueeze(1).unsqueeze(0) * 2 - 1.0 # [1, C, 1, H, W] + ref_pixel_values = ref_pixel_values.to(dtype=self.vae.dtype, device=device) + ref_latents = retrieve_latents(self.vae.encode(ref_pixel_values), sample_mode="argmax") + return self._normalize_latents(ref_latents) + + def prepare_motion_latents( + self, + motion_pixels: torch.Tensor, + device: torch.device | None = None, + ) -> torch.Tensor: + """VAE-encode motion frame pixels into latent space. + + Args: + motion_pixels: Shape ``[1, C, T_motion, H, W]``, pixel values in + ``[-1, 1]``. + + Returns: + ``motion_latents`` with shape ``[1, C_lat, T_lat, H_lat, W_lat]``. + """ + device = device or self.device + motion_pixels = motion_pixels.to(dtype=self.vae.dtype, device=device) + motion_latents = retrieve_latents(self.vae.encode(motion_pixels), sample_mode="argmax") + return self._normalize_latents(motion_latents) + + def prepare_latents( + self, + infer_frames: int, + height: int, + width: int, + motion_frames: int, + lat_motion_frames: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Generate random noise latents for one S2V clip. + + Ported from ``WanS2V.generate`` lines 544-556. + + The target latent shape accounts for both the denoised portion and + excludes the motion prefix (which is prepended separately at decode + time). + + Returns: + Noise tensor with shape ``[C, T_lat, H_lat, W_lat]`` (no batch + dim, matching the original ``WanModel_S2V`` convention). + """ + lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames + lat_h = height // self.vae_scale_factor_spatial + lat_w = width // self.vae_scale_factor_spatial + shape = (16, lat_target_frames, lat_h, lat_w) + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def check_inputs( + self, + prompt: str | None, + image: PIL.Image.Image | None, + audio_path: str | None, + height: int, + width: int, + prompt_embeds: torch.Tensor | None = None, + ): + if prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if image is None: + raise ValueError("Reference image is required for S2V generation.") + if audio_path is None: + raise ValueError("Audio path is required for S2V generation.") + if height % self.resolution_divisor != 0 or width % self.resolution_divisor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.resolution_divisor}, got {height} and {width}." + ) + + # ------------------------------------------------------------------ + # Noise prediction + # ------------------------------------------------------------------ + + def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: + """Forward pass through the S2V transformer to predict noise. + + The S2V model returns a list of tensors (one per batch element). + We take the first element since S2V processes one sample at a time. + """ + if current_model is None: + current_model = self.transformer + # WanModel_S2V's norm layers compute in float32; autocast ensures + # the float32→bfloat16 casts happen automatically (matching the + # original Wan2.2 inference path). + param_dtype = next(current_model.parameters()).dtype + with torch.amp.autocast(self.device.type, dtype=param_dtype): + result = current_model(**kwargs) + # WanModel_S2V.forward returns a list of tensors + if isinstance(result, list): + return result[0] + return result[0] if isinstance(result, tuple) else result + + def diffuse( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + guidance_scale: float, + clip_generator: torch.Generator, + dtype: torch.dtype, + device: torch.device, + max_seq_len: int, + cond_latents: torch.Tensor, + input_motion_latents: torch.Tensor, + ref_latents: torch.Tensor, + motion_frames: list[int], + drop_first_motion: bool, + positive_audio_emb: torch.Tensor, + negative_audio_emb: torch.Tensor | None, + ) -> torch.Tensor: + """Denoising diffusion loop for one S2V clip. + + Args: + latents: Initial noise tensor [C, T, H, W] + timesteps: Denoising timesteps + prompt_embeds: Text embeddings [1, seq_len, dim] + negative_prompt_embeds: Negative text embeddings (optional) + guidance_scale: CFG scale + clip_generator: Random generator for this clip + dtype: Data type for computation + device: Device for computation + max_seq_len: Maximum sequence length for transformer + cond_latents: Pose condition latents [1, 16, T, H, W] + input_motion_latents: Motion latents from previous clip + ref_latents: Reference image latents + motion_frames: Motion frame counts [pixel_frames, latent_frames] + drop_first_motion: Whether to drop first motion frames (first clip only) + positive_audio_emb: Precomputed audio embeddings for positive prompt + negative_audio_emb: Precomputed audio embeddings for negative prompt (optional) + + Returns: + Denoised latents [C, T, H, W] + """ + do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None + + with self.progress_bar(total=len(timesteps)) as pbar: + for t in timesteps: + self._current_timestep = t + + latent_model_input = [latents.to(device)] + timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device) + + # -- Positive (conditional) prediction kwargs -- + positive_kwargs = { + "x": latent_model_input, + "t": timestep, + "context": prompt_embeds[0:1], + "seq_len": max_seq_len, + "cond_states": cond_latents, + "motion_latents": input_motion_latents, + "ref_latents": ref_latents, + "motion_frames": motion_frames, + "drop_motion_frames": drop_first_motion, + "audio_emb": positive_audio_emb, + } + + if do_true_cfg: + negative_kwargs = { + "x": latent_model_input, + "t": timestep, + "context": negative_prompt_embeds[0:1], + "seq_len": max_seq_len, + "cond_states": cond_latents, + "motion_latents": input_motion_latents, + "ref_latents": ref_latents, + "motion_frames": motion_frames, + "drop_motion_frames": drop_first_motion, + "audio_emb": negative_audio_emb, + } + else: + negative_kwargs = 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, + ) + + # Scheduler step + latents = self.scheduler.step( + noise_pred.unsqueeze(0), + t, + latents.unsqueeze(0), + return_dict=False, + generator=clip_generator, + )[0].squeeze(0) + + pbar.update() + + self._current_timestep = None + return latents + + # ------------------------------------------------------------------ + # Main forward + # ------------------------------------------------------------------ + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | None = None, + negative_prompt: str | None = None, + image: PIL.Image.Image | None = None, + audio_path: str | None = None, + height: int = 704, + width: int = 1024, + infer_frames: int | None = None, + num_inference_steps: int = 40, + guidance_scale: float = 4.5, + num_repeat: int | None = None, + pose_video: str | None = None, + init_first_frame: bool = False, + output_type: str | None = "np", + generator: torch.Generator | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + **kwargs, + ) -> DiffusionOutput: + """Run S2V generation — may produce multiple autoregressive clips. + + This method mirrors ``WanS2V.generate()``, reorganized into the + vLLM-Omni pipeline structure. + """ + # ---- Extract params from request ---- + if len(req.prompts) > 1: + raise ValueError("S2V only supports a single prompt per request.") + + if len(req.prompts) == 1: + prompt_data = req.prompts[0] + prompt = prompt_data if isinstance(prompt_data, str) else prompt_data.get("prompt") + negative_prompt = None if isinstance(prompt_data, str) else prompt_data.get("negative_prompt") + + multi_modal_data = prompt_data.get("multi_modal_data", {}) if not isinstance(prompt_data, str) else {} + additional_info = prompt_data.get("additional_information", {}) if not isinstance(prompt_data, str) else {} + + if image is None: + raw_image = multi_modal_data.get("image", None) + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image).convert("RGB") + elif isinstance(raw_image, PIL.Image.Image): + image = raw_image + + if audio_path is None: + audio_path = additional_info.get("audio_path") + if audio_path is None: + audio_path = multi_modal_data.get("audio") + + if pose_video is None: + pose_video = additional_info.get("pose_video") + + init_first_frame = additional_info.get("init_first_frame", init_first_frame) + + if negative_prompt is None: + negative_prompt = _S2V_DEFAULT_NEG_PROMPT + + height = req.sampling_params.height or height + width = req.sampling_params.width or width + # num_frames defaults to 1 (image-model sentinel); S2V needs many frames. + req_num_frames = req.sampling_params.num_frames + infer_frames = ( + infer_frames + or (req_num_frames if req_num_frames and req_num_frames > 1 else None) + or self._DEFAULT_INFER_FRAMES + ) + num_steps = req.sampling_params.num_inference_steps or num_inference_steps + + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + + self._guidance_scale = guidance_scale + + # ---- Validate ---- + self.check_inputs(prompt, image, audio_path, height, width, prompt_embeds) + + device = self.device + dtype = self.transformer.dtype if hasattr(self.transformer, "dtype") else torch.bfloat16 + + if generator is None: + generator = req.sampling_params.generator + seed = req.sampling_params.seed + if generator is None and seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + if seed is None: + seed = 0 + + # ---- 1. Text encoding ---- + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=1, + max_sequence_length=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) + + # ---- 2. Audio encoding ---- + audio_emb, audio_num_repeat, target_video_frames = self.encode_audio( + audio_path, infer_frames=infer_frames, device=device, dtype=dtype + ) + if num_repeat is None or num_repeat > audio_num_repeat: + num_repeat = audio_num_repeat + + # Load raw audio waveform for muxing into output video + if isinstance(audio_path, np.ndarray): + raw_audio_waveform = audio_path.astype(np.float32) + raw_audio_sr = 16000 + else: + raw_audio_waveform, raw_audio_sr = load_audio(audio_path, sr=None, mono=True) + + # ---- 3. Reference image encoding ---- + ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype) + + # ---- 4. Initial motion latents (zeros) ---- + motion_frames = self.motion_frames + motion_pixels = torch.zeros( + [1, 3, motion_frames, height, width], + dtype=dtype, + device=device, + ) + drop_first_motion = self.drop_first_motion + if init_first_frame: + drop_first_motion = False + # Place reference image in the last 6 frames of motion + tensor_trans = transforms.ToTensor() + ref_tensor = tensor_trans(image).to(device=device, dtype=dtype) * 2 - 1.0 # [C, H, W] + motion_pixels[:, :, -6:] = ref_tensor.unsqueeze(0).unsqueeze(2).expand(-1, -1, 6, -1, -1) + + motion_latents = self.prepare_motion_latents(motion_pixels, device=device).to(dtype=dtype) + + lat_motion_frames = (motion_frames + 3) // 4 + + # ---- 5. Pose conditioning (optional) ---- + # Pose conditioning is encoded via VAE ahead of the denoising loop. + # When no pose video is provided, zero-valued cond_latents are used. + # Full pose video reading/encoding (WanS2V.load_pose_cond) would be + # ported here in a future iteration when pose support is needed. + # For now, we support the default zero-condition path. + + # ---- 6. Multi-clip autoregressive denoising ---- + clips: list[torch.Tensor] = [] + # Keep a pixel-space buffer of the trailing motion_frames for the + # autoregressive connection between clips. + videos_last_frames = torch.zeros([1, 3, motion_frames, height, width], dtype=dtype, device=device) + + for r in range(num_repeat): + # Per-clip seed + clip_seed = seed + r + clip_generator = torch.Generator(device=device).manual_seed(clip_seed) + + # -- Noise -- + latents = self.prepare_latents( + infer_frames=infer_frames, + height=height, + width=width, + motion_frames=motion_frames, + lat_motion_frames=lat_motion_frames, + dtype=dtype, + device=device, + generator=clip_generator, + ) + + # -- Scheduler -- + self.scheduler.set_timesteps(num_steps, device=device) + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + # -- Slice audio for this clip -- + left_idx = r * infer_frames + right_idx = left_idx + infer_frames + audio_input = audio_emb[..., left_idx:right_idx] + + # -- Pose condition for this clip -- + # Default: zero condition (no pose driving) + # Shape: [B=1, C=16, T, H, W] — passed to transformer which iterates over batch dim + lat_h = height // self.vae_scale_factor_spatial + lat_w = width // self.vae_scale_factor_spatial + lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames + cond_latents = torch.zeros([1, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device) + + # -- Clone motion latents for this clip -- + input_motion_latents = motion_latents.clone() + + # Max sequence length for the transformer + max_seq_len = int(np.prod([lat_target_frames, lat_h, lat_w]) // 4) + + # -- Precompute audio embeddings once per clip -- + mf = [motion_frames, lat_motion_frames] + # When CPU offload is active the transformer lives on CPU until + # its __call__ hook fires. encode_audio() bypasses __call__, + # so move the audio sub-module to GPU explicitly. + _audio_enc = getattr(self.transformer, "casual_audio_encoder", None) + _audio_on_cpu = _audio_enc is not None and next(_audio_enc.parameters()).device.type == "cpu" + if _audio_on_cpu: + _audio_enc.to(device) + param_dtype = next(self.transformer.parameters()).dtype + with torch.amp.autocast(device.type, dtype=param_dtype): + positive_audio_emb = self.transformer.encode_audio(audio_input, mf) + do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None + if do_true_cfg: + with torch.amp.autocast(device.type, dtype=param_dtype): + negative_audio_emb = self.transformer.encode_audio(0.0 * audio_input, mf) + else: + negative_audio_emb = None + if _audio_on_cpu: + _audio_enc.to("cpu") + + # -- Denoising loop -- + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + guidance_scale=guidance_scale, + clip_generator=clip_generator, + dtype=dtype, + device=device, + max_seq_len=max_seq_len, + cond_latents=cond_latents, + input_motion_latents=input_motion_latents, + ref_latents=ref_latents, + motion_frames=mf, + drop_first_motion=drop_first_motion and r == 0, + positive_audio_emb=positive_audio_emb, + negative_audio_emb=negative_audio_emb, + ) + + # ---- Decode this clip ---- + latents_for_decode = latents.unsqueeze(0) # [1, C, T, H, W] + if not (drop_first_motion and r == 0): + decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2) + else: + decode_latents = torch.cat([ref_latents, latents_for_decode], dim=2) + + decode_latents = self._denormalize_latents(decode_latents) + decode_latents = decode_latents.to(self.vae.dtype) + clip_video = self.vae.decode(decode_latents, return_dict=False)[0] # [1, C, T, H, W] + + # Handle VAE patch parallel: only rank 0 gets result, broadcast to all ranks + # This is needed for S2V's autoregressive loop where all ranks need the decoded frames + if clip_video.numel() == 0: + # Non-rank0 received empty tensor from patch parallel decode + # Use the same broadcast mechanism as the VAE patch parallel code + import torch.distributed as dist + + total_frames = decode_latents.shape[2] + + # Create buffer for broadcast + clip_video = torch.empty( + (1, 3, total_frames, height, width), + device=decode_latents.device, + dtype=decode_latents.dtype, + ) + + # Get the VAE's patch parallel group (same one used in decode) + vae_pp_group = getattr(self.vae, "_vae_pp_group", None) + if vae_pp_group is not None: + # Broadcast using the same group as VAE patch parallel + dist.broadcast(clip_video, src=0, group=vae_pp_group) + + # Trim to the infer_frames of interest + clip_video = clip_video[:, :, -infer_frames:] + if drop_first_motion and r == 0: + # Drop the first 3 frames (artifact from ref_latents prepend) + clip_video = clip_video[:, :, 3:] + + # ---- Update motion for next clip (autoregressive) ---- + overlap_frames_num = min(motion_frames, clip_video.shape[2]) + videos_last_frames = torch.cat( + [videos_last_frames[:, :, overlap_frames_num:], clip_video[:, :, -overlap_frames_num:]], + dim=2, + ) + videos_last_frames = videos_last_frames.to(dtype=dtype, device=device) + # Only prepare motion latents if there's another clip coming + if r < num_repeat - 1: + motion_latents = self.prepare_motion_latents(videos_last_frames, device=device).to(dtype=dtype) + + clips.append(clip_video.cpu()) + + # Free VRAM between clips + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + # ---- Concatenate all clips ---- + output = torch.cat(clips, dim=2) # [1, C, T_total, H, W] + + return DiffusionOutput( + output=(output, raw_audio_waveform, raw_audio_sr), + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + 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_s2v_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py new file mode 100644 index 000000000..cc0bb78c1 --- /dev/null +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py @@ -0,0 +1,1887 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Wan2.2 Speech-to-Video (S2V) Transformer using vllm-omni ops. + +Audio and motion modules use standard nn.Linear (no TP needed). +""" + +import math +from collections.abc import Iterable +from copy import deepcopy + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.attention import AdaLayerNorm +from diffusers.models.normalization import FP32LayerNorm +from einops import rearrange, repeat +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.platforms import current_omni_platform + +from .wan2_2_transformer import DistributedRMSNorm, WanFeedForward + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# S2V-specific utility functions +# --------------------------------------------------------------------------- + + +def sinusoidal_embedding_1d(dim, position): + if dim % 2 != 0: + raise ValueError(f"dim ({dim}) must be even") + half = dim // 2 + position = position.type(torch.float64) + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@torch.amp.autocast(current_omni_platform.device_type, enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + if dim % 2 != 0: + raise ValueError(f"dim ({dim}) must be even") + freqs = torch.outer( + torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)) + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +def rope_precompute(x, grid_sizes, freqs, start=None): + """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.""" + b, s, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2 + + if isinstance(freqs, list): + trainable_freqs = freqs[1] + freqs = freqs[0] + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + output = torch.view_as_complex(x.detach().reshape(b, s, n, -1, 2).to(torch.float64)) + seq_bucket = [0] + if not isinstance(grid_sizes, list): + grid_sizes = [grid_sizes] + for g in grid_sizes: + if not isinstance(g, list): + g = [torch.zeros_like(g), g] + batch_size = g[0].shape[0] + for i in range(batch_size): + if start is None: + f_o, h_o, w_o = g[0][i] + else: + f_o, h_o, w_o = start[i] + + f, h, w = g[1][i] + t_f, t_h, t_w = g[2][i] + seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o + seq_len = int(seq_f * seq_h * seq_w) + if seq_len > 0: + if t_f > 0: + if f_o >= 0: + f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist() + else: + f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist() + h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist() + w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist() + + assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0 + freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() + freqs_0 = freqs_0.view(seq_f, 1, 1, -1) + + freqs_i = torch.cat( + [ + freqs_0.expand(seq_f, seq_h, seq_w, -1), + freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1), + freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + elif t_f < 0: + freqs_i = trainable_freqs.unsqueeze(1) + output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i + seq_bucket.append(seq_bucket[-1] + seq_len) + return output + + +@torch.amp.autocast(current_omni_platform.device_type, enabled=False) +def rope_apply(x, grid_sizes, freqs, start=None): + """Apply RoPE using grid-based frequency computation (for motioner/init).""" + n, c = x.size(2), x.size(3) // 2 + input_dtype = x.dtype + + if isinstance(freqs, list): + trainable_freqs = freqs[1] + freqs = freqs[0] + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + output = x.clone() + seq_bucket = [0] + if not isinstance(grid_sizes, list): + grid_sizes = [grid_sizes] + for g in grid_sizes: + if not isinstance(g, list): + g = [torch.zeros_like(g), g] + batch_size = g[0].shape[0] + for i in range(batch_size): + if start is None: + f_o, h_o, w_o = g[0][i] + else: + f_o, h_o, w_o = start[i] + + f, h, w = g[1][i] + t_f, t_h, t_w = g[2][i] + seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o + seq_len = int(seq_f * seq_h * seq_w) + if seq_len > 0: + if t_f > 0: + if f_o >= 0: + f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist() + else: + f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist() + h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist() + w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist() + + assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0 + freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() + freqs_0 = freqs_0.view(seq_f, 1, 1, -1) + + freqs_i = torch.cat( + [ + freqs_0.expand(seq_f, seq_h, seq_w, -1), + freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1), + freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + elif t_f < 0: + freqs_i = trainable_freqs.unsqueeze(1) + x_i = torch.view_as_complex( + x[i, seq_bucket[-1] : seq_bucket[-1] + seq_len].to(torch.float64).reshape(seq_len, n, -1, 2) + ) + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = x_i + seq_bucket.append(seq_bucket[-1] + seq_len) + return output.to(input_dtype) + + +@torch.compiler.disable +@torch.amp.autocast(current_omni_platform.device_type, enabled=False) +def rope_apply_s2v(x, grid_sizes, freqs, start=None): + """Apply RoPE using precomputed complex freqs (for S2V main transformer). + + Under TP, x has fewer heads than freqs (which is precomputed with all heads). + Since RoPE frequencies are identical across heads, we slice freqs to match. + """ + n = x.size(2) + input_dtype = x.dtype + output = [] + for i, _ in enumerate(x): + s = x.size(1) + x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2)) + freqs_i = freqs[i, :s, :n] + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, s:]]) + output.append(x_i) + return torch.stack(output).to(input_dtype) + + +@torch.amp.autocast(current_omni_platform.device_type, enabled=False) +def rope_apply_usp(x, grid_sizes, freqs): + """Apply RoPE for Ulysses sequence parallel (context parallel mode).""" + n = x.size(2) + input_dtype = x.dtype + output = [] + for i, _ in enumerate(x): + s = x.size(1) + x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2)) + freqs_i = freqs[i, :, :n] + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, s:]]) + output.append(x_i) + return torch.stack(output).to(input_dtype) + + +def zero_module(module): + for p in module.parameters(): + p.detach().zero_() + return module + + +def torch_dfs(model, parent_name="root"): + module_names, modules = [], [] + current_name = parent_name if parent_name else "root" + module_names.append(current_name) + modules.append(model) + for name, child in model.named_children(): + if parent_name: + child_name = f"{parent_name}.{name}" + else: + child_name = name + child_modules, child_names = torch_dfs(child, child_name) + module_names += child_names + modules += child_modules + return modules, module_names + + +# --------------------------------------------------------------------------- +# TP-enabled main transformer building blocks +# --------------------------------------------------------------------------- + + +class WanS2VSelfAttention(nn.Module): + """S2V self-attention using vllm-omni ops (QKVParallelLinear + Attention).""" + + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.inner_dim = num_heads * self.head_dim + + # Fused QKV projection + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + bias=True, + quant_config=None, + ) + + self.num_heads = self.to_qkv.num_heads + self.num_kv_heads = self.to_qkv.num_kv_heads + self.tp_inner_dim = self.num_heads * self.head_dim + + # QK normalization (TP-aware) + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + + # Output projection + self.to_out = RowParallelLinear( + self.inner_dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + ) + + # Unified attention layer + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + num_kv_heads=self.num_kv_heads, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + ) + + def forward(self, x, seq_lens, grid_sizes, freqs): + b, s = x.shape[:2] + + # Fused QKV + qkv, _ = self.to_qkv(x) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + # QK normalization + query = self.norm_q(query) + key = self.norm_k(key) + + # Reshape to [B, S, N, D] + query = query.view(b, s, self.num_heads, self.head_dim) + key = key.view(b, s, self.num_kv_heads, self.head_dim) + value = value.view(b, s, self.num_kv_heads, self.head_dim) + + # Apply S2V complex-valued RoPE + query = rope_apply_s2v(query, grid_sizes, freqs) + key = rope_apply_s2v(key, grid_sizes, freqs) + + # Reshape back after rope (rope_apply_s2v preserves [B, S, N, D]) + query = query.view(b, s, self.num_heads, self.head_dim) + key = key.view(b, s, self.num_kv_heads, self.head_dim) + + # Attention — rope_apply_s2v preserves input dtype (bf16), so + # query/key/value are all bf16 here, matching FlashAttention requirements. + hidden_states = self.attn(query, key, value) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(x) + + # Output + hidden_states = self.to_out(hidden_states) + return hidden_states + + +class WanS2VCrossAttention(nn.Module): + """S2V cross-attention using vllm-omni ops (ColumnParallelLinear + Attention).""" + + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.inner_dim = num_heads * self.head_dim + + # Separate Q/K/V projections for cross-attention + self.to_q = ColumnParallelLinear(dim, self.inner_dim, bias=True, gather_output=False, return_bias=False) + self.to_k = ColumnParallelLinear(dim, self.inner_dim, bias=True, gather_output=False, return_bias=False) + self.to_v = ColumnParallelLinear(dim, self.inner_dim, bias=True, gather_output=False, return_bias=False) + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads = num_heads // tp_size + self.tp_inner_dim = self.num_heads * self.head_dim + + # QK normalization (TP-aware) + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + + # Output projection + self.to_out = RowParallelLinear(self.inner_dim, dim, bias=True, input_is_parallel=True, return_bias=False) + + # Unified attention layer + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + num_kv_heads=self.num_heads, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + skip_sequence_parallel=True, + ) + + def forward(self, x, context, context_lens=None): + b, n, d = x.size(0), self.num_heads, self.head_dim + + q = self.norm_q(self.to_q(x)).view(b, -1, n, d) + k = self.norm_k(self.to_k(context)).view(b, -1, n, d) + v = self.to_v(context).view(b, -1, n, d) + + x = self.attn(q, k, v) + x = x.flatten(2, 3) + x = x.type_as(q) + + x = self.to_out(x) + return x + + +class WanS2VTransformerBlock(nn.Module): + """S2V transformer block with segment-wise modulation. + + Key S2V feature: `e` is a tuple [e0_tensor, seg_idx]. The sequence is split + into 2 segments at seg_idx (noisy tokens vs ref/motion tokens), and different + modulation is applied to each segment. + """ + + def __init__(self, dim, ffn_dim, num_heads, window_size=(-1, -1), qk_norm=True, cross_attn_norm=False, eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + + # Self-attention (TP-enabled) + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.self_attn = WanS2VSelfAttention(dim, num_heads, window_size, qk_norm, eps) + + # Cross-attention (TP-enabled) + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanS2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) + + # FFN (TP-enabled, reused from wan2_2_transformer.py) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim) + + # 6-way scale-shift modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens): + seg_boundary = e[1] + seg_idx = [0, min(max(0, seg_boundary), x.size(1)), x.size(1)] + e = e[0] + + # Modulation parameters — self.modulation will be cast to match e's dtype + e = (self.modulation.unsqueeze(2).to(e.dtype) + e).chunk(6, dim=1) + e = [element.squeeze(1) for element in e] + + # Self-attention with per-segment modulation + norm_x = self.norm1(x).type_as(x) + parts = [] + for i in range(2): + parts.append(norm_x[:, seg_idx[i] : seg_idx[i + 1]] * (1 + e[1][:, i : i + 1]) + e[0][:, i : i + 1]) + norm_x = torch.cat(parts, dim=1) + + y = self.self_attn(norm_x, seq_lens, grid_sizes, freqs) + y_parts = [] + for i in range(2): + y_parts.append(y[:, seg_idx[i] : seg_idx[i + 1]] * e[2][:, i : i + 1]) + x = x + torch.cat(y_parts, dim=1) + + # Cross-attention + x = x + self.cross_attn(self.norm3(x).type_as(x), context, context_lens) + + # FFN with per-segment modulation + norm2_x = self.norm2(x).type_as(x) + parts = [] + for i in range(2): + parts.append(norm2_x[:, seg_idx[i] : seg_idx[i + 1]] * (1 + e[4][:, i : i + 1]) + e[3][:, i : i + 1]) + norm2_x = torch.cat(parts, dim=1) + y = self.ffn(norm2_x) + y_parts = [] + for i in range(2): + y_parts.append(y[:, seg_idx[i] : seg_idx[i + 1]] * e[5][:, i : i + 1]) + x = x + torch.cat(y_parts, dim=1) + return x + + +class WanS2VHead(nn.Module): + """S2V output head with modulation.""" + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + + out_dim_full = math.prod(patch_size) * out_dim + self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim_full) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + e = (self.modulation.to(e.dtype) + e.unsqueeze(1)).chunk(2, dim=1) + x = self.norm(x).type_as(x) * (1 + e[1]) + e[0] + x = self.head(x) + return x + + +# --------------------------------------------------------------------------- +# Audio modules (no TP, keep as nn.Linear) +# --------------------------------------------------------------------------- + + +class CausalConv1d(nn.Module): + def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", **kwargs): + super().__init__() + self.pad_mode = pad_mode + self.time_causal_padding = (kernel_size - 1, 0) + self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + + +class MotionEncoder_tc(nn.Module): + def __init__(self, in_dim, hidden_dim, num_heads=4, need_global=True, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.num_heads = num_heads + self.need_global = need_global + + self.conv1_local = CausalConv1d(in_dim, hidden_dim // 4 * num_heads, 3, stride=1) + if need_global: + self.conv1_global = CausalConv1d(in_dim, hidden_dim // 4, 3, stride=1) + self.act = nn.SiLU() + self.conv2 = CausalConv1d(hidden_dim // 4, hidden_dim // 2, 3, stride=2) + self.conv3 = CausalConv1d(hidden_dim // 2, hidden_dim, 3, stride=2) + + if need_global: + self.final_linear = nn.Linear(hidden_dim, hidden_dim, **factory_kwargs) + + self.norm1 = nn.LayerNorm(hidden_dim // 4, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.norm2 = nn.LayerNorm(hidden_dim // 2, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.norm3 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim)) + + def forward(self, x): + x = rearrange(x, "b t c -> b c t") + x_ori = x.clone() + b, c, t = x.shape + x = self.conv1_local(x) + x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads) + x = self.norm1(x) + x = self.act(x) + x = rearrange(x, "b t c -> b c t") + x = self.conv2(x) + x = rearrange(x, "b c t -> b t c") + x = self.norm2(x) + x = self.act(x) + x = rearrange(x, "b t c -> b c t") + x = self.conv3(x) + x = rearrange(x, "b c t -> b t c") + x = self.norm3(x) + x = self.act(x) + x = rearrange(x, "(b n) t c -> b t n c", b=b) + padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1) + x = torch.cat([x, padding], dim=-2) + x_local = x.clone() + + if not self.need_global: + return x_local + + x = self.conv1_global(x_ori) + x = rearrange(x, "b c t -> b t c") + x = self.norm1(x) + x = self.act(x) + x = rearrange(x, "b t c -> b c t") + x = self.conv2(x) + x = rearrange(x, "b c t -> b t c") + x = self.norm2(x) + x = self.act(x) + x = rearrange(x, "b t c -> b c t") + x = self.conv3(x) + x = rearrange(x, "b c t -> b t c") + x = self.norm3(x) + x = self.act(x) + x = self.final_linear(x) + x = rearrange(x, "(b n) t c -> b t n c", b=b) + return x, x_local + + +class CausalAudioEncoder(nn.Module): + def __init__(self, dim=5120, num_layers=25, out_dim=2048, video_rate=8, num_token=4, need_global=False): + super().__init__() + self.encoder = MotionEncoder_tc(in_dim=dim, hidden_dim=out_dim, num_heads=num_token, need_global=need_global) + weight = torch.ones((1, num_layers, 1, 1)) * 0.01 + self.weights = torch.nn.Parameter(weight) + self.act = torch.nn.SiLU() + + def forward(self, features): + # weights param is float32; cast to features dtype for the weighted sum + weights = self.act(self.weights).to(features.dtype) + weights_sum = weights.sum(dim=1, keepdims=True) + weighted_feat = ((features * weights) / weights_sum).sum(dim=1) + weighted_feat = weighted_feat.permute(0, 2, 1) + res = self.encoder(weighted_feat) + return res + + +class AudioCrossAttention(nn.Module): + """Audio cross-attention using nn.Linear (no TP needed).""" + + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + self.attn = Attention( + num_heads=num_heads, + head_size=self.head_dim, + num_kv_heads=num_heads, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + ) + + def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: list[int] | None = None) -> torch.Tensor: + b, n, d = x.size(0), self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, -1, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + x = self.attn(q, k, v) + x = x.flatten(2, 3) + x = self.o(x) + return x + + +class AudioInjector_WAN(nn.Module): + def __init__( + self, + all_modules, + all_modules_names, + dim=2048, + num_heads=32, + inject_layer=None, + root_net=None, + enable_adain=False, + adain_dim=2048, + need_adain_ont=False, + ): + super().__init__() + if inject_layer is None: + inject_layer = [0, 27] + + self.injected_block_id = {} + audio_injector_id = 0 + for mod_name, mod in zip(all_modules_names, all_modules): + if isinstance(mod, WanS2VTransformerBlock): + for inject_id in inject_layer: + if f"transformer_blocks.{inject_id}" in mod_name: + self.injected_block_id[inject_id] = audio_injector_id + audio_injector_id += 1 + + self.injector = nn.ModuleList( + [AudioCrossAttention(dim=dim, num_heads=num_heads, qk_norm=True) for _ in range(audio_injector_id)] + ) + self.injector_pre_norm_feat = nn.ModuleList( + [nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) for _ in range(audio_injector_id)] + ) + self.injector_pre_norm_vec = nn.ModuleList( + [nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) for _ in range(audio_injector_id)] + ) + if enable_adain: + self.injector_adain_layers = nn.ModuleList( + [ + AdaLayerNorm(output_dim=dim * 2, embedding_dim=adain_dim, chunk_dim=1) + for _ in range(audio_injector_id) + ] + ) + if need_adain_ont: + self.injector_adain_output_layers = nn.ModuleList( + [nn.Linear(dim, dim) for _ in range(audio_injector_id)] + ) + + +# --------------------------------------------------------------------------- +# Motion modules (no TP, keep as nn.Linear) +# --------------------------------------------------------------------------- + + +class SimpleSelfAttention(nn.Module): + """Simple self-attention for motioner (no TP, uses nn.Linear).""" + + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + self.attn = Attention( + num_heads=num_heads, + head_size=self.head_dim, + num_kv_heads=num_heads, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + ) + + def forward( + self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + ) -> torch.Tensor: + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + + q = rope_apply(q, grid_sizes, freqs) + k = rope_apply(k, grid_sizes, freqs) + + # rope_apply preserves [B, S, N, D] shape and input dtype + q = q.view(b, s, n, d) + k = k.view(b, s, n, d) + + x = self.attn(q, k, v) + x = x.flatten(2, 3) + x = self.o(x) + return x + + +class SwinSelfAttention(SimpleSelfAttention): + def forward( + self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + ) -> torch.Tensor: + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + if b != 1: + raise ValueError( + f"SwinSelfAttention only supports batch_size=1, got batch_size={b}. " + "Batched inference requires refactoring the frame reference and rearrange logic." + ) + + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + + q = rope_apply(q, grid_sizes, freqs) + k = rope_apply(k, grid_sizes, freqs) + q = q.view(b, s, n, d) + k = k.view(b, s, n, d) + T, H, W = grid_sizes[0].tolist() + + q = rearrange(q, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + k = rearrange(k, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + v = rearrange(v, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + + # Skip ref_q (last frame query) - unused in current implementation + q = q[:-1] + + ref_k = repeat(k[-1:], "1 s n d -> t s n d", t=k.shape[0] - 1) + k = k[:-1] + k = torch.cat([k[:1], k, k[-1:]]) + k = torch.cat([k[1:-1], k[2:], k[:-2], ref_k], dim=1) + + ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=v.shape[0] - 1) + v = v[:-1] + v = torch.cat([v[:1], v, v[-1:]]) + v = torch.cat([v[1:-1], v[2:], v[:-2], ref_v], dim=1) + + out = self.attn(q, k, v) + out = torch.cat([out, ref_v[:1]], axis=0) + out = rearrange(out, "(b t) (h w) n d -> b (t h w) n d", t=T, h=H, w=W) + x = out.flatten(2, 3) + x = self.o(x) + return x + + +class CausalSelfAttention(SimpleSelfAttention): + def forward( + self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + ) -> torch.Tensor: + shifting = 3 + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + if b != 1: + raise ValueError( + f"CausalSelfAttention only supports batch_size=1, got batch_size={b}. " + "Batched inference requires refactoring the causal masking and frame reference logic." + ) + + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + + T, H, W = grid_sizes[0].tolist() + + q = rearrange(q, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + k = rearrange(k, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + v = rearrange(v, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + + # Skip ref_q (last frame query) - unused in current causal implementation + q = q[:-1] + + grid_sizes_q = torch.tensor([[1, H, W]] * q.shape[0], dtype=torch.long) + start = [[shifting, 0, 0]] * q.shape[0] + q = rope_apply(q, grid_sizes_q, freqs, start=start) + q = q.view(q.shape[0], -1, n, d) + + ref_k = k[-1:] + grid_sizes_ref = torch.tensor([[1, H, W]], dtype=torch.long) + start_ref = [[shifting + 10, 0, 0]] + ref_k = rope_apply(ref_k, grid_sizes_ref, freqs, start_ref) + ref_k = ref_k.view(1, -1, n, d) + ref_k = repeat(ref_k, "1 s n d -> t s n d", t=k.shape[0] - 1) + + k = k[:-1] + k = torch.cat([*([k[:1]] * shifting), k]) + cat_k = [] + for i in range(shifting): + cat_k.append(k[i : i - shifting]) + cat_k.append(k[shifting:]) + k = torch.cat(cat_k, dim=1) + + grid_sizes_k = torch.tensor([[shifting + 1, H, W]] * (q.shape[0]), dtype=torch.long) + k = rope_apply(k, grid_sizes_k, freqs) + k = k.view(k.shape[0], -1, n, d) + k = torch.cat([k, ref_k], dim=1) + + ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=q.shape[0]) + v = v[:-1] + v = torch.cat([*([v[:1]] * shifting), v]) + cat_v = [] + for i in range(shifting): + cat_v.append(v[i : i - shifting]) + cat_v.append(v[shifting:]) + v = torch.cat(cat_v, dim=1) + v = torch.cat([v, ref_v], dim=1) + + outs = [] + for i in range(q.shape[0]): + out = self.attn(q[i : i + 1], k[i : i + 1], v[i : i + 1]) + outs.append(out) + out = torch.cat(outs, dim=0) + out = torch.cat([out, ref_v[:1]], axis=0) + out = rearrange(out, "(b t) (h w) n d -> b (t h w) n d", t=T, h=H, w=W) + x = out.flatten(2, 3) + x = self.o(x) + return x + + +class MotionerAttentionBlock(nn.Module): + def __init__( + self, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + self_attn_block="SelfAttention", + ): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + + self.norm1 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) + if self_attn_block == "SelfAttention": + self.self_attn = SimpleSelfAttention(dim, num_heads, window_size, qk_norm, eps) + elif self_attn_block == "SwinSelfAttention": + self.self_attn = SwinSelfAttention(dim, num_heads, window_size, qk_norm, eps) + elif self_attn_block == "CausalSelfAttention": + self.self_attn = CausalSelfAttention(dim, num_heads, window_size, qk_norm, eps) + + self.norm2 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) + + def forward( + self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + ) -> torch.Tensor: + y = self.self_attn(self.norm1(x).type_as(x), seq_lens, grid_sizes, freqs) + x = x + y + y = self.ffn(self.norm2(x).type_as(x)) + x = x + y + return x + + +class MotionerHead(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + out_dim = math.prod(patch_size) * out_dim + self.norm = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) + self.head = nn.Linear(dim, out_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.head(self.norm(x)) + return x + + +class MotionerTransformers(nn.Module): + def __init__( + self, + patch_size=(1, 2, 2), + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + self_attn_block="SelfAttention", + motion_token_num=1024, + enable_tsm=False, + motion_stride=4, + expand_ratio=2, + trainable_token_pos_emb=False, + ): + super().__init__() + self.patch_size = patch_size + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + + self.enable_tsm = enable_tsm + self.motion_stride = motion_stride + self.expand_ratio = expand_ratio + self.sample_c = patch_size[0] + + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.blocks = nn.ModuleList( + [ + MotionerAttentionBlock( + dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, self_attn_block=self_attn_block + ) + for _ in range(num_layers) + ] + ) + + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") + if (dim // num_heads) % 2 != 0: + raise ValueError(f"dim // num_heads ({dim // num_heads}) must be even") + d = dim // num_heads + self.freqs = torch.cat( + [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], + dim=1, + ) + + self.gradient_checkpointing = False + self.motion_side_len = int(math.sqrt(motion_token_num)) + if self.motion_side_len**2 != motion_token_num: + raise ValueError(f"motion_token_num ({motion_token_num}) must be a perfect square") + self.token = nn.Parameter(torch.zeros(1, motion_token_num, dim).contiguous()) + + self.trainable_token_pos_emb = trainable_token_pos_emb + if trainable_token_pos_emb: + x = torch.zeros([1, motion_token_num, num_heads, d]) + x[..., ::2] = 1 + gride_sizes = [ + [ + torch.tensor([0, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([1, self.motion_side_len, self.motion_side_len]).unsqueeze(0).repeat(1, 1), + torch.tensor([1, self.motion_side_len, self.motion_side_len]).unsqueeze(0).repeat(1, 1), + ] + ] + token_freqs = rope_apply(x, gride_sizes, self.freqs) + token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2) + token_freqs = token_freqs * 0.01 + self.token_freqs = torch.nn.Parameter(token_freqs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x[0].shape[1] contains motion frames count, but not used directly here + device = self.patch_embedding.weight.device + freqs = self.freqs + if freqs.device != device: + freqs = freqs.to(device) + + if self.trainable_token_pos_emb: + with torch.amp.autocast(current_omni_platform.device_type, dtype=torch.float64): + token_freqs = self.token_freqs.to(torch.float64) + token_freqs = token_freqs / token_freqs.norm(dim=-1, keepdim=True) + freqs = [freqs, torch.view_as_complex(token_freqs)] + + if self.enable_tsm: + sample_idx = [ + _sample_indices(u.shape[1], stride=self.motion_stride, expand_ratio=self.expand_ratio, c=self.sample_c) + for u in x + ] + x = [torch.flip(torch.flip(u, [1])[:, idx], [1]) for idx, u in zip(sample_idx, x)] + + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + + seq_f, seq_h, seq_w = x[0].shape[-3:] + batch_size = len(x) + if not self.enable_tsm: + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + grid_sizes = [[torch.zeros_like(grid_sizes), grid_sizes, grid_sizes]] + seq_f = 0 + else: + grid_sizes = [] + for idx in sample_idx[0][::-1][:: self.sample_c]: + tsm_frame_grid_sizes = [ + [ + torch.tensor([idx, 0, 0]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([idx + 1, seq_h, seq_w]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([1, seq_h, seq_w]).unsqueeze(0).repeat(batch_size, 1), + ] + ] + grid_sizes += tsm_frame_grid_sizes + seq_f = sample_idx[0][-1] + 1 + + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + x = torch.cat([u for u in x]) + batch_size = len(x) + + token_grid_sizes = [ + [ + torch.tensor([seq_f, 0, 0]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([seq_f + 1, self.motion_side_len, self.motion_side_len]) + .unsqueeze(0) + .repeat(batch_size, 1), + torch.tensor([1 if not self.trainable_token_pos_emb else -1, seq_h, seq_w]) + .unsqueeze(0) + .repeat(batch_size, 1), + ] + ] + + grid_sizes = grid_sizes + token_grid_sizes + token_len = self.token.shape[1] + token = self.token.clone().repeat(x.shape[0], 1, 1).contiguous() + seq_lens = seq_lens + torch.tensor([t.size(0) for t in token], dtype=torch.long) + x = torch.cat([x, token], dim=1) + + kwargs = dict(seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=freqs) + for block in self.blocks: + x = block(x, **kwargs) + + out = x[:, -token_len:] + return out + + def unpatchify(self, x, grid_sizes): + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + +class FramePackMotioner(nn.Module): + def __init__(self, inner_dim=1024, num_heads=16, zip_frame_buckets=None, drop_mode="drop", *args, **kwargs): + super().__init__(*args, **kwargs) + if zip_frame_buckets is None: + zip_frame_buckets = [1, 2, 16] + + self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)) + self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4)) + self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8)) + self.zip_frame_buckets = torch.tensor(zip_frame_buckets, dtype=torch.long) + + self.inner_dim = inner_dim + self.num_heads = num_heads + + if inner_dim % num_heads != 0: + raise ValueError(f"inner_dim ({inner_dim}) must be divisible by num_heads ({num_heads})") + if (inner_dim // num_heads) % 2 != 0: + raise ValueError(f"inner_dim // num_heads ({inner_dim // num_heads}) must be even") + d = inner_dim // num_heads + self.freqs = torch.cat( + [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], + dim=1, + ) + self.drop_mode = drop_mode + + def forward(self, motion_latents: torch.Tensor, add_last_motion: int = 2) -> torch.Tensor: + mot = [] + mot_remb = [] + for m in motion_latents: + lat_height, lat_width = m.shape[2], m.shape[3] + pad_lat = torch.zeros(16, self.zip_frame_buckets.sum(), lat_height, lat_width).to( + device=m.device, dtype=m.dtype + ) + overlap_frame = min(pad_lat.shape[1], m.shape[1]) + if overlap_frame > 0: + pad_lat[:, -overlap_frame:] = m[:, -overlap_frame:] + + if add_last_motion < 2 and self.drop_mode != "drop": + zero_end_frame = self.zip_frame_buckets[: len(self.zip_frame_buckets) - add_last_motion - 1].sum() + pad_lat[:, -zero_end_frame:] = 0 + + pad_lat = pad_lat.unsqueeze(0) + clean_latents_4x, clean_latents_2x, clean_latents_post = pad_lat[ + :, :, -self.zip_frame_buckets.sum() :, :, : + ].split(list(self.zip_frame_buckets)[::-1], dim=2) + + clean_latents_post = self.proj(clean_latents_post).flatten(2).transpose(1, 2) + clean_latents_2x = self.proj_2x(clean_latents_2x).flatten(2).transpose(1, 2) + clean_latents_4x = self.proj_4x(clean_latents_4x).flatten(2).transpose(1, 2) + + if add_last_motion < 2 and self.drop_mode == "drop": + clean_latents_post = clean_latents_post[:, :0] if add_last_motion < 2 else clean_latents_post + clean_latents_2x = clean_latents_2x[:, :0] if add_last_motion < 1 else clean_latents_2x + + motion_lat = torch.cat([clean_latents_post, clean_latents_2x, clean_latents_4x], dim=1) + + # Rope + start_time_id = -(self.zip_frame_buckets[:1].sum()) + end_time_id = start_time_id + self.zip_frame_buckets[0] + grid_sizes = ( + [] + if add_last_motion < 2 and self.drop_mode == "drop" + else [ + [ + torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([end_time_id, lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.zip_frame_buckets[0], lat_height // 2, lat_width // 2]) + .unsqueeze(0) + .repeat(1, 1), + ] + ] + ) + + start_time_id = -(self.zip_frame_buckets[:2].sum()) + end_time_id = start_time_id + self.zip_frame_buckets[1] // 2 + grid_sizes_2x = ( + [] + if add_last_motion < 1 and self.drop_mode == "drop" + else [ + [ + torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([end_time_id, lat_height // 4, lat_width // 4]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.zip_frame_buckets[1], lat_height // 2, lat_width // 2]) + .unsqueeze(0) + .repeat(1, 1), + ] + ] + ) + + start_time_id = -(self.zip_frame_buckets[:3].sum()) + end_time_id = start_time_id + self.zip_frame_buckets[2] // 4 + grid_sizes_4x = [ + [ + torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([end_time_id, lat_height // 8, lat_width // 8]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.zip_frame_buckets[2], lat_height // 2, lat_width // 2]) + .unsqueeze(0) + .repeat(1, 1), + ] + ] + + grid_sizes = grid_sizes + grid_sizes_2x + grid_sizes_4x + + motion_rope_emb = rope_precompute( + motion_lat.detach().view(1, motion_lat.shape[1], self.num_heads, self.inner_dim // self.num_heads), + grid_sizes, + self.freqs, + start=None, + ) + + mot.append(motion_lat) + mot_remb.append(motion_rope_emb) + return mot, mot_remb + + +def _sample_indices(num_frames, stride, expand_ratio, c): + indices = [] + current_start = 0 + while current_start < num_frames: + bucket_width = int(stride * (expand_ratio ** (len(indices) / stride))) + interval = int(bucket_width / stride * c) + current_end = min(num_frames, current_start + bucket_width) + bucket_samples = [] + for i in range(current_end - 1, current_start - 1, -interval): + for near in range(c): + bucket_samples.append(i - near) + indices += bucket_samples[::-1] + current_start += bucket_width + return indices + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + + +class WanS2VTransformer3DModel(nn.Module): + """ + Wan2.2 S2V Transformer using vllm-omni ops for the main transformer blocks. + + This replaces WanModel_S2V from wan_modules, using: + - QKVParallelLinear for self-attention Q/K/V (fused) + - ColumnParallelLinear for cross-attention Q/K/V + - RowParallelLinear for output projections + - DistributedRMSNorm for QK normalization (TP-aware) + - WanFeedForward (ColumnParallelGELU + RowParallelLinear) + - FP32LayerNorm for layer normalization + - Attention layer for attention computation + """ + + _repeated_blocks = ["WanS2VTransformerBlock"] + _layerwise_offload_blocks_attrs = ["blocks"] + + @staticmethod + def _is_transformer_block(name, mod): + """Match transformer blocks for HSDP sharding.""" + return "blocks" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + + def __init__( + self, + cond_dim=0, + audio_dim=5120, + num_audio_token=4, + enable_adain=False, + adain_mode="attn_norm", + audio_inject_layers=None, + zero_init=False, + zero_timestep=False, + enable_motioner=True, + add_last_motion=True, + enable_tsm=False, + trainable_token_pos_emb=False, + motion_token_num=1024, + enable_framepack=False, + framepack_drop_mode="drop", + model_type="s2v", + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + *args, + **kwargs, + ): + super().__init__() + + if audio_inject_layers is None: + audio_inject_layers = [0, 4, 8, 12, 16, 20, 24, 27] + + assert model_type == "s2v" + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # Embeddings + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim)) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # Transformer blocks (TP-enabled) + self.blocks = nn.ModuleList( + [ + WanS2VTransformerBlock(dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ] + ) + + # Output head + self.head = WanS2VHead(dim, out_dim, patch_size, eps) + + # RoPE base frequencies (complex-valued) + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") + if (dim // num_heads) % 2 != 0: + raise ValueError(f"dim // num_heads ({dim // num_heads}) must be even") + d = dim // num_heads + self.freqs = torch.cat( + [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], + dim=1, + ) + + self.use_context_parallel = False + + # Condition encoder + if cond_dim > 0: + self.cond_encoder = nn.Conv3d(cond_dim, dim, kernel_size=patch_size, stride=patch_size) + + # Audio modules + self.enable_adain = enable_adain + self.casual_audio_encoder = CausalAudioEncoder( + dim=audio_dim, out_dim=dim, num_token=num_audio_token, need_global=enable_adain + ) + + all_modules, all_modules_names = torch_dfs(self.blocks, parent_name="root.transformer_blocks") + self.audio_injector = AudioInjector_WAN( + all_modules, + all_modules_names, + dim=dim, + num_heads=num_heads, + inject_layer=audio_inject_layers, + root_net=self, + enable_adain=enable_adain, + adain_dim=dim, + need_adain_ont=adain_mode != "attn_norm", + ) + self.adain_mode = adain_mode + + self.trainable_cond_mask = nn.Embedding(3, dim) + + if zero_init: + self._zero_init_weights() + + self.zero_timestep = zero_timestep + self.add_last_motion = add_last_motion + + # Motion modules + if enable_motioner and enable_framepack: + raise ValueError("enable_motioner and enable_framepack are mutually exclusive") + + self.enable_motioner = enable_motioner + if enable_motioner: + motioner_dim = 2048 + self.motioner = MotionerTransformers( + patch_size=(2, 4, 4), + dim=motioner_dim, + ffn_dim=motioner_dim, + freq_dim=256, + out_dim=16, + num_heads=16, + num_layers=13, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + motion_token_num=motion_token_num, + enable_tsm=enable_tsm, + motion_stride=4, + expand_ratio=2, + trainable_token_pos_emb=trainable_token_pos_emb, + ) + self.zip_motion_out = nn.Sequential( + FP32LayerNorm(motioner_dim, elementwise_affine=False), zero_module(nn.Linear(motioner_dim, dim)) + ) + + self.trainable_token_pos_emb = trainable_token_pos_emb + if trainable_token_pos_emb: + d = dim // num_heads + x = torch.zeros([1, motion_token_num, num_heads, d]) + x[..., ::2] = 1 + gride_sizes = [ + [ + torch.tensor([0, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([1, self.motioner.motion_side_len, self.motioner.motion_side_len]) + .unsqueeze(0) + .repeat(1, 1), + torch.tensor([1, self.motioner.motion_side_len, self.motioner.motion_side_len]) + .unsqueeze(0) + .repeat(1, 1), + ] + ] + token_freqs = rope_apply(x, gride_sizes, self.freqs) + token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2) + token_freqs = token_freqs * 0.01 + self.token_freqs = nn.Parameter(token_freqs) + + self.enable_framepack = enable_framepack + if enable_framepack: + self.frame_packer = FramePackMotioner( + inner_dim=dim, num_heads=num_heads, zip_frame_buckets=[1, 2, 16], drop_mode=framepack_drop_mode + ) + + def _zero_init_weights(self): + with torch.no_grad(): + self.trainable_cond_mask = zero_module(self.trainable_cond_mask) + if hasattr(self, "cond_encoder"): + self.cond_encoder = zero_module(self.cond_encoder) + for i in range(len(self.audio_injector.injector)): + self.audio_injector.injector[i].o = zero_module(self.audio_injector.injector[i].o) + if self.enable_adain: + self.audio_injector.injector_adain_layers[i].linear = zero_module( + self.audio_injector.injector_adain_layers[i].linear + ) + + # ------------------------------------------------------------------ + # Motion processing + # ------------------------------------------------------------------ + + def process_motion(self, motion_latents, drop_motion_frames=False): + if drop_motion_frames or motion_latents[0].shape[1] == 0: + return [], [] + self.lat_motion_frames = motion_latents[0].shape[1] + mot = [self.patch_embedding(m.unsqueeze(0)) for m in motion_latents] + batch_size = len(mot) + mot_remb = [] + flatten_mot = [] + for bs in range(batch_size): + height, width = mot[bs].shape[3], mot[bs].shape[4] + flat_mot = mot[bs].flatten(2).transpose(1, 2).contiguous() + motion_grid_sizes = [ + [ + torch.tensor([-self.lat_motion_frames, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([0, height, width]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.lat_motion_frames, height, width]).unsqueeze(0).repeat(1, 1), + ] + ] + motion_rope_emb = rope_precompute( + flat_mot.detach().view(1, flat_mot.shape[1], self.num_heads, self.dim // self.num_heads), + motion_grid_sizes, + self.freqs, + start=None, + ) + mot_remb.append(motion_rope_emb) + flatten_mot.append(flat_mot) + return flatten_mot, mot_remb + + def process_motion_frame_pack(self, motion_latents, drop_motion_frames=False, add_last_motion=2): + flatten_mot, mot_remb = self.frame_packer(motion_latents, add_last_motion) + if drop_motion_frames: + return [m[:, :0] for m in flatten_mot], [m[:, :0] for m in mot_remb] + else: + return flatten_mot, mot_remb + + def process_motion_transformer_motioner(self, motion_latents, drop_motion_frames=False, add_last_motion=True): + batch_size = len(motion_latents) + height = motion_latents[0].shape[2] // self.patch_size[1] + width = motion_latents[0].shape[3] // self.patch_size[2] + + freqs = self.freqs + device = self.patch_embedding.weight.device + if freqs.device != device: + freqs = freqs.to(device) + if self.trainable_token_pos_emb: + with torch.amp.autocast(current_omni_platform.device_type, dtype=torch.float64): + token_freqs = self.token_freqs.to(torch.float64) + token_freqs = token_freqs / token_freqs.norm(dim=-1, keepdim=True) + freqs = [freqs, torch.view_as_complex(token_freqs)] + + if not drop_motion_frames and add_last_motion: + last_motion_latent = [u[:, -1:] for u in motion_latents] + last_mot = [self.patch_embedding(m.unsqueeze(0)) for m in last_motion_latent] + last_mot = [m.flatten(2).transpose(1, 2) for m in last_mot] + last_mot = torch.cat(last_mot) + gride_sizes = [ + [ + torch.tensor([-1, 0, 0]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([0, height, width]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([1, height, width]).unsqueeze(0).repeat(batch_size, 1), + ] + ] + else: + last_mot = torch.zeros( + [batch_size, 0, self.dim], device=motion_latents[0].device, dtype=motion_latents[0].dtype + ) + gride_sizes = [] + + zip_motion = self.motioner(motion_latents) + zip_motion = self.zip_motion_out(zip_motion) + if drop_motion_frames: + zip_motion = zip_motion * 0.0 + zip_motion_grid_sizes = [ + [ + torch.tensor([-1, 0, 0]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([0, self.motioner.motion_side_len, self.motioner.motion_side_len]) + .unsqueeze(0) + .repeat(batch_size, 1), + torch.tensor([1 if not self.trainable_token_pos_emb else -1, height, width]) + .unsqueeze(0) + .repeat(batch_size, 1), + ] + ] + + mot = torch.cat([last_mot, zip_motion], dim=1) + gride_sizes = gride_sizes + zip_motion_grid_sizes + + motion_rope_emb = rope_precompute( + mot.detach().view(batch_size, mot.shape[1], self.num_heads, self.dim // self.num_heads), + gride_sizes, + freqs, + start=None, + ) + return [m.unsqueeze(0) for m in mot], [r.unsqueeze(0) for r in motion_rope_emb] + + def inject_motion( + self, x, seq_lens, rope_embs, mask_input, motion_latents, drop_motion_frames=False, add_last_motion=True + ): + if self.enable_motioner: + mot, mot_remb = self.process_motion_transformer_motioner( + motion_latents, drop_motion_frames=drop_motion_frames, add_last_motion=add_last_motion + ) + elif self.enable_framepack: + mot, mot_remb = self.process_motion_frame_pack( + motion_latents, drop_motion_frames=drop_motion_frames, add_last_motion=add_last_motion + ) + else: + mot, mot_remb = self.process_motion(motion_latents, drop_motion_frames=drop_motion_frames) + + if len(mot) > 0: + x = [torch.cat([u, m], dim=1) for u, m in zip(x, mot)] + seq_lens = seq_lens + torch.tensor([r.size(1) for r in mot], dtype=torch.long) + rope_embs = [torch.cat([u, m], dim=1) for u, m in zip(rope_embs, mot_remb)] + mask_input = [ + torch.cat([m, 2 * torch.ones([1, u.shape[1] - m.shape[1]], device=m.device, dtype=m.dtype)], dim=1) + for m, u in zip(mask_input, x) + ] + return x, seq_lens, rope_embs, mask_input + + def after_transformer_block(self, block_idx, hidden_states): + if block_idx in self.audio_injector.injected_block_id.keys(): + audio_attn_id = self.audio_injector.injected_block_id[block_idx] + audio_emb = self.merged_audio_emb + num_frames = audio_emb.shape[1] + + input_hidden_states = hidden_states[:, : self.original_seq_len].clone() + input_hidden_states = rearrange(input_hidden_states, "b (t n) c -> (b t) n c", t=num_frames) + + if self.enable_adain and self.adain_mode == "attn_norm": + audio_emb_global = self.audio_emb_global + audio_emb_global = rearrange(audio_emb_global, "b t n c -> (b t) n c") + adain_hidden_states = self.audio_injector.injector_adain_layers[audio_attn_id]( + input_hidden_states, temb=audio_emb_global[:, 0] + ) + attn_hidden_states = adain_hidden_states + else: + attn_hidden_states = self.audio_injector.injector_pre_norm_feat[audio_attn_id](input_hidden_states) + + audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c", t=num_frames) + attn_audio_emb = audio_emb + residual_out = self.audio_injector.injector[audio_attn_id]( + x=attn_hidden_states, + context=attn_audio_emb, + context_lens=torch.ones(attn_hidden_states.shape[0], dtype=torch.long, device=attn_hidden_states.device) + * attn_audio_emb.shape[1], + ) + residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames) + hidden_states[:, : self.original_seq_len] = hidden_states[:, : self.original_seq_len] + residual_out + + return hidden_states + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def encode_audio(self, audio_input: torch.Tensor, motion_frames: list[int]) -> dict[str, torch.Tensor]: + """Precompute audio embeddings from raw audio features. + + Call this once per clip before the denoising loop, then pass the + result to forward() via the ``audio_emb`` kwarg to avoid redundant + encoding at every denoising step. + + Returns: + dict with 'audio_emb' (and optionally 'audio_emb_global' when + enable_adain is True). + """ + audio_input = torch.cat( + [audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input], + dim=-1, + ) + audio_emb_res = self.casual_audio_encoder(audio_input) + result = {} + if self.enable_adain: + audio_emb_global, audio_emb = audio_emb_res + result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :].clone() + else: + audio_emb = audio_emb_res + result["audio_emb"] = audio_emb[:, motion_frames[1] :, :] + return result + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + context: torch.Tensor, + seq_len: int, + ref_latents: torch.Tensor, + motion_latents: torch.Tensor, + cond_states: torch.Tensor, + audio_input: torch.Tensor | None = None, + motion_frames: list[int] = [17, 5], + add_last_motion: int = 2, + drop_motion_frames: bool = False, + audio_emb: dict[str, torch.Tensor] | None = None, + *extra_args: object, + **extra_kwargs: object, + ) -> torch.Tensor: + """Forward pass for S2V transformer. + + Args: + x: Noisy latents [B, C, T, H, W] + t: Timesteps [B] + context: Text embeddings [B, seq_len, dim] + seq_len: Sequence length + ref_latents: Reference image latents [B, C, 1, H, W] + motion_latents: Motion context latents [B, C, T_motion, H, W] + cond_states: Pose condition latents [B, C, T, H, W] — typically zero tensor + for no pose control. Iterates over batch dimension (dim 0). + audio_input: Raw audio features (optional if audio_emb provided) + motion_frames: [num_frames, num_latent_frames] for audio alignment + add_last_motion: Number of last motion frames to repeat + drop_motion_frames: Whether to drop first motion frame + audio_emb: Precomputed audio embeddings from encode_audio() + + Returns: + Predicted noise [B, C, T, H, W] + """ + add_last_motion = self.add_last_motion * add_last_motion + + # Audio encoding — use precomputed embeddings if available + if audio_emb is not None: + self.merged_audio_emb = audio_emb["audio_emb"] + if self.enable_adain: + self.audio_emb_global = audio_emb["audio_emb_global"] + else: + audio_input = torch.cat( + [audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input], + dim=-1, + ) + audio_emb_res = self.casual_audio_encoder(audio_input) + if self.enable_adain: + audio_emb_global, audio_emb_local = audio_emb_res + self.audio_emb_global = audio_emb_global[:, motion_frames[1] :].clone() + else: + audio_emb_local = audio_emb_res + self.merged_audio_emb = audio_emb_local[:, motion_frames[1] :, :] + + # Patch embedding + condition + # NOTE: x and cond_states are 5D tensors [B, C, T, H, W]. Iterating over them + # yields 4D tensors [C, T, H, W], which we unsqueeze(0) to restore batch dim. + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + cond = [self.cond_encoder(c.unsqueeze(0)) for c in cond_states] + x = [x_ + pose for x_, pose in zip(x, cond)] + + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + + original_grid_sizes = deepcopy(grid_sizes) + grid_sizes = [[torch.zeros_like(grid_sizes), grid_sizes, grid_sizes]] + + # Reference image tokens + self.lat_motion_frames = motion_latents[0].shape[1] + ref = [self.patch_embedding(r.unsqueeze(0)) for r in ref_latents] + batch_size = len(ref) + height, width = ref[0].shape[3], ref[0].shape[4] + ref_grid_sizes = [ + [ + torch.tensor([30, 0, 0]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([31, height, width]).unsqueeze(0).repeat(batch_size, 1), + torch.tensor([1, height, width]).unsqueeze(0).repeat(batch_size, 1), + ] + ] + + ref = [r.flatten(2).transpose(1, 2) for r in ref] + self.original_seq_len = seq_lens[0].item() + + seq_lens = seq_lens + torch.tensor([r.size(1) for r in ref], dtype=torch.long) + grid_sizes = grid_sizes + ref_grid_sizes + x = [torch.cat([u, r], dim=1) for u, r in zip(x, ref)] + + # Masks: 0=noisy, 1=ref, 2=motion + mask_input = [torch.zeros([1, u.shape[1]], dtype=torch.long, device=x[0].device) for u in x] + for i in range(len(mask_input)): + mask_input[i][:, self.original_seq_len :] = 1 + + # Precompute RoPE + x = torch.cat(x) + b, s, n, d = x.size(0), x.size(1), self.num_heads, self.dim // self.num_heads + self.pre_compute_freqs = rope_precompute(x.detach().view(b, s, n, d), grid_sizes, self.freqs, start=None) + + x = [u.unsqueeze(0) for u in x] + self.pre_compute_freqs = [u.unsqueeze(0) for u in self.pre_compute_freqs] + + # Inject motion tokens + x, seq_lens, self.pre_compute_freqs, mask_input = self.inject_motion( + x, + seq_lens, + self.pre_compute_freqs, + mask_input, + motion_latents, + drop_motion_frames=drop_motion_frames, + add_last_motion=add_last_motion, + ) + + x = torch.cat(x, dim=0) + self.pre_compute_freqs = torch.cat(self.pre_compute_freqs, dim=0) + mask_input = torch.cat(mask_input, dim=0) + + x = x + self.trainable_cond_mask(mask_input).to(x.dtype) + + # Time embeddings — sinusoidal_embedding_1d uses float64 internally + # for precision, then we cast to the model dtype (bf16) since the + # downstream scale/shift modulation doesn't need float32. + if self.zero_timestep: + t = torch.cat([t, torch.zeros([1], dtype=t.dtype, device=t.device)]) + e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).to(x.dtype)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + + if self.zero_timestep: + e = e[:-1] + zero_e0 = e0[-1:] + e0 = e0[:-1] + e0 = torch.cat([e0.unsqueeze(2), zero_e0.unsqueeze(2).repeat(e0.size(0), 1, 1, 1)], dim=2) + e0 = [e0, self.original_seq_len] + else: + e0 = e0.unsqueeze(2).repeat(1, 1, 2, 1) + e0 = [e0, 0] + + # Context embedding + context_lens = None + context = self.text_embedding( + torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) + ) + + # Transformer blocks + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.pre_compute_freqs, + context=context, + context_lens=context_lens, + ) + + for idx, block in enumerate(self.blocks): + x = block(x, **kwargs) + x = self.after_transformer_block(idx, x) + + # Output + x = x[:, : self.original_seq_len] + x = self.head(x, e) + x = self.unpatchify(x, original_grid_sizes) + return [u.float() for u in x] + + def unpatchify(self, x, grid_sizes): + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights with QKV fusion for main transformer blocks.""" + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + # QKV fusion mapping (only for main transformer blocks) + stacked_params_mapping = [ + (".self_attn.to_qkv", ".self_attn.q", "q"), + (".self_attn.to_qkv", ".self_attn.k", "k"), + (".self_attn.to_qkv", ".self_attn.v", "v"), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + original_name = name + is_main_block = name.startswith("blocks.") + + if is_main_block: + # Try QKV fusion for self-attention + fused = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + lookup_name = name.replace(weight_name, param_name) + if lookup_name in params_dict: + param = params_dict[lookup_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + fused = True + break + + if fused: + loaded_params.add(original_name) + loaded_params.add(lookup_name) + continue + + # Name remapping for main blocks + name = name.replace(".self_attn.o.", ".self_attn.to_out.") + name = name.replace(".cross_attn.q.", ".cross_attn.to_q.") + name = name.replace(".cross_attn.k.", ".cross_attn.to_k.") + name = name.replace(".cross_attn.v.", ".cross_attn.to_v.") + name = name.replace(".cross_attn.o.", ".cross_attn.to_out.") + name = name.replace(".ffn.0.", ".ffn.net_0.proj.") + name = name.replace(".ffn.2.", ".ffn.net_2.") + + if name not in params_dict: + logger.warning(f"Skipping weight {original_name} -> {name}") + continue + + param = params_dict[name] + + # TP-shard norm weights for main blocks + if ( + tp_size > 1 + and is_main_block + and any( + norm_name in name + for norm_name in [ + ".self_attn.norm_q.", + ".self_attn.norm_k.", + ".cross_attn.norm_q.", + ".cross_attn.norm_k.", + ] + ) + ): + shard_size = loaded_weight.shape[0] // tp_size + loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + loaded_params.add(original_name) + loaded_params.add(name) + + return loaded_params + + +# --------------------------------------------------------------------------- +# Factory function +# --------------------------------------------------------------------------- + + +def create_s2v_transformer_from_config(config: dict) -> WanS2VTransformer3DModel: + """Create WanS2VTransformer3DModel from a config dict (diffusers format).""" + kwargs = {} + + # Map config keys to constructor parameters + key_mapping = { + "cond_dim": "cond_dim", + "audio_dim": "audio_dim", + "num_audio_token": "num_audio_token", + "enable_adain": "enable_adain", + "adain_mode": "adain_mode", + "audio_inject_layers": "audio_inject_layers", + "zero_init": "zero_init", + "zero_timestep": "zero_timestep", + "enable_motioner": "enable_motioner", + "add_last_motion": "add_last_motion", + "enable_tsm": "enable_tsm", + "trainable_token_pos_emb": "trainable_token_pos_emb", + "motion_token_num": "motion_token_num", + "enable_framepack": "enable_framepack", + "framepack_drop_mode": "framepack_drop_mode", + "model_type": "model_type", + "patch_size": "patch_size", + "text_len": "text_len", + "in_dim": "in_dim", + "dim": "dim", + "ffn_dim": "ffn_dim", + "freq_dim": "freq_dim", + "text_dim": "text_dim", + "out_dim": "out_dim", + "num_heads": "num_heads", + "num_layers": "num_layers", + "window_size": "window_size", + "qk_norm": "qk_norm", + "cross_attn_norm": "cross_attn_norm", + "eps": "eps", + } + + for config_key, param_name in key_mapping.items(): + if config_key in config: + value = config[config_key] + if config_key in ("patch_size", "window_size") and isinstance(value, list): + value = tuple(value) + kwargs[param_name] = value + + return WanS2VTransformer3DModel(**kwargs) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 44beb4bc1..753e62293 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -115,6 +115,11 @@ _DIFFUSION_MODELS = { "pipeline_wan2_2_i2v", "Wan22I2VPipeline", ), + "WanS2VPipeline": ( + "wan2_2", + "pipeline_wan2_2_s2v", + "Wan22S2VPipeline", + ), "WanT2VDMD2Pipeline": ( "wan2_2", "pipeline_wan2_2", @@ -434,6 +439,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "StableAudioPipeline": "get_stable_audio_post_process_func", "AudioXPipeline": "get_audiox_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", + "WanS2VPipeline": "get_wan22_s2v_post_process_func", "WanT2VDMD2Pipeline": "get_wan22_post_process_func", "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func", "LongCatImagePipeline": "get_longcat_image_post_process_func", @@ -469,6 +475,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "WanPipeline": "get_wan22_pre_process_func", "WanVACEPipeline": "get_wan22_vace_pre_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func", + "WanS2VPipeline": "get_wan22_s2v_pre_process_func", "WanT2VDMD2Pipeline": "get_wan22_pre_process_func", "WanI2VDMD2Pipeline": "get_wan22_i2v_pre_process_func", "OmniGen2Pipeline": "get_omnigen2_pre_process_func", diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index c73a56fab..b4d308ebf 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -307,6 +307,72 @@ class OmniRequestOutput: """Check if this is a pipeline stage output.""" return self.stage_id is not None and self.request_output is not None + def unwrap(self) -> "OmniRequestOutput": + """Unwrap nested OmniRequestOutput to get the innermost result. + + This helper handles the common pattern where pipeline outputs may wrap + other OmniRequestOutput instances. It recursively unwraps until it reaches + the final output with actual content (images, text, etc.). + + Returns: + The innermost OmniRequestOutput containing the actual generation results. + + Example: + ```python + result = omni.generate(...) + output = OmniRequestOutput.unwrap_result(result) + if output.images: + # Access images directly + video_frames = output.images + ``` + """ + current = self + # Unwrap nested pipeline outputs + while current.is_pipeline_output and current.request_output is not None: + if isinstance(current.request_output, OmniRequestOutput): + current = current.request_output + else: + break + return current + + @staticmethod + def unwrap_result(result: Any) -> "OmniRequestOutput": + """Unwrap result from omni.generate() to get the final OmniRequestOutput. + + This static helper handles the full unwrapping pattern including: + 1. Extracting from list if needed + 2. Type validation + 3. Recursive unwrapping of nested pipeline outputs + + Args: + result: The result from omni.generate() - may be a list or OmniRequestOutput + + Returns: + The innermost OmniRequestOutput with actual content + + Raises: + ValueError: If result is not an OmniRequestOutput or list containing one + + Example: + ```python + result = omni.generate(...) + output = OmniRequestOutput.unwrap_result(result) + # output is guaranteed to be the final OmniRequestOutput + ``` + """ + # Handle list wrapper + if isinstance(result, list): + if not result: + raise ValueError("Result list is empty") + result = result[0] + + # Validate type + if not isinstance(result, OmniRequestOutput): + raise ValueError(f"Expected OmniRequestOutput, got {type(result)}") + + # Unwrap nested outputs + return result.unwrap() + def to_dict(self) -> dict[str, Any]: """Convert to dictionary for JSON serialization.""" result = { diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index cc4218ae2..bd50163dc 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -74,3 +74,8 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.xpu.mem_get_info(device) return free + + @classmethod + def get_profiler_cls(cls) -> str: + """Return XPU-specific profiler that handles XPU events.""" + return "vllm_omni.platforms.xpu.profiler.XPUTorchProfilerWrapper" diff --git a/vllm_omni/platforms/xpu/profiler.py b/vllm_omni/platforms/xpu/profiler.py new file mode 100644 index 000000000..97bcca718 --- /dev/null +++ b/vllm_omni/platforms/xpu/profiler.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing_extensions import override +from vllm.logger import init_logger + +from vllm_omni.profiler.omni_torch_profiler import ( + OmniTorchProfilerWrapper, + TorchProfilerActivity, +) + +logger = init_logger(__name__) + + +class XPUTorchProfilerWrapper(OmniTorchProfilerWrapper): + """XPU-specific profiler wrapper. + + Uses torch.profiler with XPU ProfilerActivity to capture Intel GPU events. + XPU events are exposed via torch.profiler.ProfilerActivity.XPU and work + with the standard torch.profiler interface. + """ + + @override + def _get_default_activities(self) -> list[TorchProfilerActivity]: + """Default to CPU + XPU profiling for Intel GPUs.""" + return ["CPU", "XPU"] + + @override + def _on_stop_hook(self) -> None: + """XPU-specific stop hook that handles XPU event aggregation. + + XPU profiler follows CUDA-like patterns but uses xpu_time_total metrics. + """ + rank = self.local_rank + sort_key = "self_xpu_time_total" + + excel_sheets: dict[str, list[dict]] = {} + + # 1) Summary op table + summary_events = self.profiler.key_averages() + excel_sheets["summary"] = self._event_list_to_rows(summary_events) + + # 2) Shape-grouped op table + if self.profiler_config.torch_profiler_record_shapes: + try: + shape_events = self.profiler.key_averages( + group_by_input_shape=True, + ) + excel_sheets["by_shape"] = self._event_list_to_rows(shape_events) + except Exception as e: + logger.warning( + "[Rank %s] Failed to export shape-grouped op table: %s", + rank, + e, + ) + + # 3) Stack-grouped op table + if self.profiler_config.torch_profiler_with_stack: + try: + stack_events = self.profiler.key_averages( + group_by_stack_n=8, + ) + excel_sheets["by_stack"] = self._event_list_to_rows(stack_events) + except Exception as e: + logger.warning( + "[Rank %s] Failed to export stack-grouped op table: %s", + rank, + e, + ) + + # 4) Export stack files + try: + cpu_stack_file = self._artifact_path("stacks_cpu", ".txt") + self.profiler.export_stacks( + cpu_stack_file, + metric="self_cpu_time_total", + ) + self._artifact_paths["stacks_cpu"] = cpu_stack_file + except Exception as e: + logger.warning("[Rank %s] export_stacks(cpu) failed: %s", rank, e) + + # XPU stack traces + try: + xpu_stack_file = self._artifact_path("stacks_xpu", ".txt") + self.profiler.export_stacks( + xpu_stack_file, + metric="self_xpu_time_total", + ) + self._artifact_paths["stacks_xpu"] = xpu_stack_file + except Exception as e: + logger.warning("[Rank %s] export_stacks(xpu) failed: %s", rank, e) + + try: + self._table_path = self._write_excel_artifact("ops", excel_sheets) + except Exception as e: + logger.warning("[Rank %s] Failed to export Excel workbook: %s", rank, e) + + # Print XPU time table for rank 0 + if rank == 0: + try: + table = self.profiler.key_averages().table(sort_by=sort_key, row_limit=50) + logger.info("XPU Profiler Results:\n%s", table) + except Exception as e: + logger.warning("[Rank %s] Failed to generate XPU time table: %s", rank, e) + + if self.dump_cpu_time_total and rank == 0: + logger.info(self.profiler.key_averages().table(sort_by="self_cpu_time_total", row_limit=50))