From 6838533998f1d8ef4d006e5429995118d9ba9e64 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 21 Mar 2026 21:26:39 +0800 Subject: [PATCH] [Model] Add HunyuanVideo-1.5 T2V and I2V support (#1516) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: lishunyang Signed-off-by: 汪志鹏 Co-authored-by: 汪志鹏 Co-authored-by: Hongsheng Liu --- docs/models/supported_models.md | 2 + docs/user_guide/diffusion/quantization/fp8.md | 1 + .../image_to_video/image_to_video.py | 14 +- .../text_to_video/text_to_video.md | 107 ++- .../text_to_video/text_to_video.py | 161 ++-- .../run_curl_hunyuan_video_15.sh | 68 ++ .../run_server_hunyuan_video_15.sh | 31 + .../run_curl_hunyuan_video_15.sh | 60 ++ .../run_server_hunyuan_video_15.sh | 31 + .../test_hunyuan_video_15_expansion.py | 98 +++ .../models/hunyuan_video/__init__.py | 25 + .../hunyuan_video_15_transformer.py | 791 ++++++++++++++++++ .../pipeline_hunyuan_video_1_5.py | 513 ++++++++++++ .../pipeline_hunyuan_video_1_5_i2v.py | 594 +++++++++++++ vllm_omni/diffusion/registry.py | 13 + 15 files changed, 2443 insertions(+), 66 deletions(-) create mode 100644 examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh create mode 100644 examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh create mode 100644 examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh create mode 100644 examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh create mode 100644 tests/e2e/online_serving/test_hunyuan_video_15_expansion.py create mode 100644 vllm_omni/diffusion/models/hunyuan_video/__init__.py create mode 100644 vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py create mode 100644 vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py create mode 100644 vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a3b1dbe76..b4a15f9de 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -52,6 +52,8 @@ th { |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | |`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` | +| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | +| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | |`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | diff --git a/docs/user_guide/diffusion/quantization/fp8.md b/docs/user_guide/diffusion/quantization/fp8.md index 95a62b05d..f7887ab76 100644 --- a/docs/user_guide/diffusion/quantization/fp8.md +++ b/docs/user_guide/diffusion/quantization/fp8.md @@ -63,6 +63,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None | | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` | | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None | +| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None | ## Combining with Other Features diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index e312c6e72..04e056691 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -2,15 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2. +Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5. Supports: - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B) - LTX2 image-to-video pipeline +- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning Usage: - # I2V-A14B (MoE) + # Wan I2V-A14B (MoE) python image_to_video.py --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \ --image input.jpg --prompt "A cat playing with yarn" @@ -24,6 +25,11 @@ Usage: --image input.jpg --prompt "A cinematic dolly shot of a boat" \ --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \ --frame-rate 24 --fps 24 --output ltx2_i2v.mp4 + + # HunyuanVideo-1.5 I2V (480p) + python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \ + --image input.jpg --prompt "A cat playing with yarn" \ + --flow-shift 5.0 --guidance-scale 6.0 """ import argparse @@ -43,11 +49,11 @@ from vllm_omni.platforms import current_omni_platform def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.") + parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).") parser.add_argument( "--model", default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", - help="Diffusers Wan2.2 I2V model ID or local path.", + help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).", ) parser.add_argument( "--model-class-name", diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index cdac3d39a..f852e980a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -1,10 +1,19 @@ # Text-To-Video -The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used -for `Lightricks/LTX-2` to generate video+audio. +A unified script for text-to-video generation. Supports multiple models with model-aware defaults. + +## Supported Models + +| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | +|---|---|---|---|---|---| +| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | +| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | ## Local CLI Usage +### Wan2.2 (default) + ```bash python text_to_video.py \ --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ @@ -36,22 +45,100 @@ python text_to_video.py \ --output ltx2_out.mp4 ``` -Key arguments: +### HunyuanVideo-1.5 (480p) +```bash +python text_to_video.py \ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ + --prompt "A cat walks through a sunlit garden, flowers swaying gently in the breeze." \ + --height 480 \ + --width 832 \ + --num-frames 121 \ + --guidance-scale 6.0 \ + --flow-shift 5.0 \ + --num-inference-steps 50 \ + --fps 24 \ + --output hunyuan_video_15_output.mp4 +``` + +### HunyuanVideo-1.5 (720p) + +```bash +python text_to_video.py \ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --height 720 \ + --width 1280 \ + --num-frames 121 \ + --guidance-scale 6.0 \ + --flow-shift 9.0 \ + --num-inference-steps 50 \ + --fps 24 \ + --output hunyuan_720p.mp4 +``` + +### HunyuanVideo-1.5 with FP8 Quantization + +```bash +python text_to_video.py \ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ + --prompt "A dog running across a field of golden wheat." \ + --quantization fp8 \ + --height 480 --width 832 --num-frames 121 \ + --guidance-scale 6.0 --flow-shift 5.0 \ + --output hunyuan_fp8.mp4 +``` + +Quick test (smaller resolution, fewer frames): + +```bash +python text_to_video.py \ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --height 320 --width 576 --num-frames 17 --num-inference-steps 30 \ + --flow-shift 5.0 \ + --output quick_test.mp4 +``` + +## Key Arguments + +### Common + +- `--model`: Diffusers model ID or local path. - `--prompt`: text description (string). -- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8). -- `--num-frames`: Number of frames (Wan default is 81). -- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high). -- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string). -- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality). -- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video). +- `--height/--width`: output resolution. Default depends on model. +- `--num-frames`: number of frames. Default depends on model. +- `--guidance-scale`: CFG scale. Default depends on model. +- `--num-inference-steps`: sampling steps. Default depends on model. +- `--fps`: frames per second for the saved MP4. - `--output`: path to save the generated video. - `--vae-use-slicing`: enable VAE slicing for memory optimization. - `--vae-use-tiling`: enable VAE tiling for memory optimization. - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel). - `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2). - `--enable-cpu-offload`: enable CPU offloading for diffusion models. +- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). - `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio). +- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). +- `--flow-shift`: scheduler flow_shift parameter. -> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. +### Wan2.2-specific + +- `--negative-prompt`: artifacts to suppress. +- `--guidance-scale-high`: separate CFG scale for high-noise stage. +- `--boundary-ratio`: boundary split for low/high DiT (default 0.875). +- `--flow-shift`: scheduler flow_shift (5.0 for 720p, 12.0 for 480p). +- `--cache-backend`: `cache_dit` for acceleration. + +### HunyuanVideo-1.5 Optimal Configs + +| Variant | flow_shift | guidance_scale | steps | +|---------|-----------|----------------|-------| +| 480p T2V | 5.0 | 6.0 | 50 | +| 720p T2V | 9.0 | 6.0 | 50 | +| 480p I2V | 5.0 | 6.0 | 50 | +| 720p I2V | 7.0 | 6.0 | 50 | +| CFG-distilled | (same) | 1.0 | 50 | + +> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`. diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index e0493d4e6..a3aa818d2 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -15,56 +15,92 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform +_MODEL_PRESETS = { + "wan": { + "height": 720, + "width": 1280, + "num_frames": 81, + "num_inference_steps": 40, + "guidance_scale": 4.0, + "fps": 24, + "output": "wan22_output.mp4", + }, + "hunyuan": { + "height": 480, + "width": 832, + "num_frames": 121, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "fps": 24, + "output": "hunyuan_video_15_output.mp4", + }, +} + + +def _detect_preset(model: str) -> dict: + model_lower = model.lower() + if "hunyuan" in model_lower: + return _MODEL_PRESETS["hunyuan"] + return _MODEL_PRESETS["wan"] + def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.") + parser = argparse.ArgumentParser( + description="Generate a video from a text prompt. " + "Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models." + ) parser.add_argument( "--model", default="Wan-AI/Wan2.2-T2V-A14B-Diffusers", - help="Diffusers Wan2.2 model ID or local path.", + help="Diffusers model ID or local path. " + "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, " + "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v", ) parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.") - parser.add_argument("--negative-prompt", default="", help="Negative prompt.") + parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).") parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).") - parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.") - parser.add_argument("--height", type=int, default=720, help="Video height.") - parser.add_argument("--width", type=int, default=1280, help="Video width.") - parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).") + parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.") + parser.add_argument( + "--guidance-scale-high", type=float, default=None, help="Separate CFG for high-noise stage (Wan2.2 only)." + ) + parser.add_argument("--height", type=int, default=None, help="Video height. Default: model-specific.") + parser.add_argument("--width", type=int, default=None, help="Video width. Default: model-specific.") + parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.") parser.add_argument( "--frame-rate", type=float, default=None, help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.", ) - parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.") + parser.add_argument( + "--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific." + ) parser.add_argument( "--boundary-ratio", type=float, - default=0.875, - help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).", + default=None, + help="(Wan2.2) Boundary split ratio for low/high DiT. Default 0.875.", ) parser.add_argument( - "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)." + "--flow-shift", + type=float, + default=None, + help="Scheduler flow_shift. Wan2.2: 5.0(720p)/12.0(480p). HunyuanVideo-1.5: 5.0(480p)/9.0(720p).", ) parser.add_argument( "--cache-backend", type=str, default=None, choices=["cache_dit"], - help=( - "Cache backend to use for acceleration. " - "Options: 'cache_dit' (DBCache + SCM + TaylorSeer). " - "Default: None (no cache acceleration)." - ), + help="Cache backend for acceleration (Wan2.2). 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("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).") - parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.") + parser.add_argument("--output", type=str, default=None, help="Output path (mp4). Default: model-specific.") + parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.") parser.add_argument( "--vae-use-slicing", action="store_true", @@ -137,32 +173,41 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--quantization", + type=str, + default=None, + choices=["fp8", "gguf"], + help="Quantization method for the transformer (fp8 for online FP8 quantization).", + ) return parser.parse_args() def main(): args = parse_args() - generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps) - # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment). + preset = _detect_preset(args.model) + for key, default_val in preset.items(): + if getattr(args, key.replace("-", "_"), None) is None: + setattr(args, key.replace("-", "_"), default_val) + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + # Cache-dit config (Wan2.2 only) cache_config = None if args.cache_backend == "cache_dit": cache_config = { - # DBCache parameters [cache-dit only] - "Fn_compute_blocks": 1, # Optimized for single-transformer models - "Bn_compute_blocks": 0, # Number of backward compute blocks - "max_warmup_steps": 4, # Maximum warmup steps (works for few-step models) + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, "max_cached_steps": 20, - "residual_diff_threshold": 0.24, # Higher threshold for more aggressive caching - "max_continuous_cached_steps": 3, # Limit to prevent precision degradation - # TaylorSeer parameters [cache-dit only] - "enable_taylorseer": False, # Disabled by default (not suitable for few-step models) - "taylorseer_order": 1, # TaylorSeer polynomial order - # SCM (Step Computation Masking) parameters [cache-dit only] - "scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra" - "scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static" + "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", } + # Configure parallel settings parallel_config = DiffusionParallelConfig( ulysses_degree=args.ulysses_degree, @@ -176,14 +221,11 @@ def main(): # Check if profiling is requested via environment variable profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) - omni = Omni( + omni_kwargs = dict( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, - boundary_ratio=args.boundary_ratio, - flow_shift=args.flow_shift, - enable_cache_dit_summary=args.enable_cache_dit_summary, enable_cpu_offload=args.enable_cpu_offload, parallel_config=parallel_config, enforce_eager=args.enforce_eager, @@ -191,6 +233,18 @@ def main(): cache_config=cache_config, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) + if args.boundary_ratio is not None: + omni_kwargs["boundary_ratio"] = args.boundary_ratio + if args.flow_shift is not None: + omni_kwargs["flow_shift"] = args.flow_shift + if args.quantization is not None: + omni_kwargs["quantization"] = args.quantization + if args.cache_backend is not None: + omni_kwargs["cache_backend"] = args.cache_backend + omni_kwargs["cache_config"] = cache_config + omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary + + omni = Omni(**omni_kwargs) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -210,22 +264,25 @@ def main(): print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") + prompt_dict = {"prompt": args.prompt} + if args.negative_prompt: + prompt_dict["negative_prompt"] = args.negative_prompt + + sampling_kwargs = dict( + height=args.height, + width=args.width, + generator=generator, + guidance_scale=args.guidance_scale, + num_inference_steps=args.num_inference_steps, + num_frames=args.num_frames, + ) + if args.guidance_scale_high is not None: + sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high + generation_start = time.perf_counter() frames = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - }, - OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - generator=generator, - guidance_scale=args.guidance_scale, - guidance_scale_2=args.guidance_scale_high, - num_inference_steps=args.num_inference_steps, - num_frames=args.num_frames, - frame_rate=frame_rate, - ), + prompt_dict, + OmniDiffusionSamplingParams(**sampling_kwargs), ) generation_end = time.perf_counter() generation_time = generation_end - generation_start diff --git a/examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh b/examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh new file mode 100644 index 000000000..23a5fcf07 --- /dev/null +++ b/examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# HunyuanVideo-1.5 image-to-video curl example using the async video job API. + +set -euo pipefail + +INPUT_IMAGE="${INPUT_IMAGE:-test_input.jpg}" +BASE_URL="${BASE_URL:-http://localhost:8099}" +OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_i2v.mp4}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +if [ ! -f "$INPUT_IMAGE" ]; then + echo "Input image not found: $INPUT_IMAGE" + echo "Provide an image via INPUT_IMAGE env var." + exit 1 +fi + +create_response=$( + curl -sS -X POST "${BASE_URL}/v1/videos" \ + -H "Accept: application/json" \ + -F "prompt=The camera follows the puppy as it runs forward on the grass, its four legs alternating steps, its tail held high and wagging side to side." \ + -F "input_reference=@${INPUT_IMAGE}" \ + -F "size=832x480" \ + -F "num_frames=33" \ + -F "fps=24" \ + -F "num_inference_steps=30" \ + -F "guidance_scale=6.0" \ + -F "flow_shift=5.0" \ + -F "seed=42" +) + +video_id="$(echo "${create_response}" | jq -r '.id')" +if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then + echo "Failed to create video job:" + echo "${create_response}" | jq . + exit 1 +fi + +echo "Created video job ${video_id}" +echo "${create_response}" | jq . + +while true; do + status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")" + status="$(echo "${status_response}" | jq -r '.status')" + + case "${status}" in + queued|in_progress) + echo "Video job ${video_id} status: ${status}" + sleep "${POLL_INTERVAL}" + ;; + completed) + echo "${status_response}" | jq . + break + ;; + failed) + echo "Video generation failed:" + echo "${status_response}" | jq . + exit 1 + ;; + *) + echo "Unexpected status response:" + echo "${status_response}" | jq . + exit 1 + ;; + esac +done + +curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}" +echo "Saved video to ${OUTPUT_PATH}" diff --git a/examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh b/examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh new file mode 100644 index 000000000..50a78372f --- /dev/null +++ b/examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# HunyuanVideo-1.5 image-to-video online serving startup script +# +# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB +# 720p: needs FP8 + VAE tiling, ~35 GB VRAM + +MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v}" +PORT="${PORT:-8099}" +FLOW_SHIFT="${FLOW_SHIFT:-5.0}" +QUANTIZATION="${QUANTIZATION:-}" +CACHE_BACKEND="${CACHE_BACKEND:-none}" + +echo "Starting HunyuanVideo-1.5 I2V server..." +echo "Model: $MODEL" +echo "Port: $PORT" +echo "Flow shift: $FLOW_SHIFT" +echo "Quantization: ${QUANTIZATION:-none}" +echo "Cache backend: $CACHE_BACKEND" + +EXTRA_FLAGS="" +if [ -n "$QUANTIZATION" ]; then + EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION" +fi +if [ "$CACHE_BACKEND" != "none" ]; then + EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND" +fi + +vllm serve "$MODEL" --omni \ + --port "$PORT" \ + --flow-shift "$FLOW_SHIFT" \ + $EXTRA_FLAGS diff --git a/examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh b/examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh new file mode 100644 index 000000000..e6a655704 --- /dev/null +++ b/examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# HunyuanVideo-1.5 text-to-video curl example using the async video job API. + +set -euo pipefail + +BASE_URL="${BASE_URL:-http://localhost:8098}" +OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_t2v.mp4}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +create_response=$( + curl -sS -X POST "${BASE_URL}/v1/videos" \ + -H "Accept: application/json" \ + -F "prompt=A little girl wearing a straw hat runs through a summer meadow full of wildflowers. A wide shot is used, with the camera panning right to follow her." \ + -F "size=832x480" \ + -F "num_frames=33" \ + -F "fps=24" \ + -F "num_inference_steps=30" \ + -F "guidance_scale=6.0" \ + -F "flow_shift=5.0" \ + -F "seed=42" +) + +video_id="$(echo "${create_response}" | jq -r '.id')" +if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then + echo "Failed to create video job:" + echo "${create_response}" | jq . + exit 1 +fi + +echo "Created video job ${video_id}" +echo "${create_response}" | jq . + +while true; do + status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")" + status="$(echo "${status_response}" | jq -r '.status')" + + case "${status}" in + queued|in_progress) + echo "Video job ${video_id} status: ${status}" + sleep "${POLL_INTERVAL}" + ;; + completed) + echo "${status_response}" | jq . + break + ;; + failed) + echo "Video generation failed:" + echo "${status_response}" | jq . + exit 1 + ;; + *) + echo "Unexpected status response:" + echo "${status_response}" | jq . + exit 1 + ;; + esac +done + +curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}" +echo "Saved video to ${OUTPUT_PATH}" diff --git a/examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh b/examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh new file mode 100644 index 000000000..5f6af72d7 --- /dev/null +++ b/examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# HunyuanVideo-1.5 text-to-video online serving startup script +# +# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB +# 720p: needs FP8 + VAE tiling, ~35 GB VRAM + +MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v}" +PORT="${PORT:-8098}" +FLOW_SHIFT="${FLOW_SHIFT:-5.0}" +QUANTIZATION="${QUANTIZATION:-}" +CACHE_BACKEND="${CACHE_BACKEND:-none}" + +echo "Starting HunyuanVideo-1.5 T2V server..." +echo "Model: $MODEL" +echo "Port: $PORT" +echo "Flow shift: $FLOW_SHIFT" +echo "Quantization: ${QUANTIZATION:-none}" +echo "Cache backend: $CACHE_BACKEND" + +EXTRA_FLAGS="" +if [ -n "$QUANTIZATION" ]; then + EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION" +fi +if [ "$CACHE_BACKEND" != "none" ]; then + EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND" +fi + +vllm serve "$MODEL" --omni \ + --port "$PORT" \ + --flow-shift "$FLOW_SHIFT" \ + $EXTRA_FLAGS diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py new file mode 100644 index 000000000..71c7e5a50 --- /dev/null +++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests of common diffusion feature combinations in online serving mode +for HunyuanVideo-1.5-T2V (480p). + +Coverage (H100, since model cannot fit L4): +- TeaCache + Layerwise CPU offloading (1 GPU) +- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) +""" + +import pytest + +from tests.conftest import ( + OmniServer, + OmniServerParams, + OpenAIClientHandler, + dummy_messages_from_mix_data, +) +from tests.utils import hardware_marks + +PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion." +NEGATIVE_PROMPT = "low quality, blurry, distorted" + +MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) +PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) + + +def _get_diffusion_feature_cases(model: str): + """Return L4 diffusion feature cases for HunyuanVideo-1.5. + + Designed for 2x H100 environment per issue #1832. + """ + return [ + # (1 GPU) TeaCache + Layerwise CPU offloading + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "tea_cache", + "--enable-layerwise-offload", + ], + ), + id="single_card_teacache_layerwise", + marks=SINGLE_CARD_MARKS, + ), + # (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2 + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "cache_dit", + "--tensor-parallel-size", + "2", + "--vae-patch-parallel-size", + "2", + "--vae-use-tiling", + ], + ), + id="parallel_cachedit_tp2_vae2", + marks=PARALLEL_MARKS, + ), + ] + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_hunyuan_video_15_t2v( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +): + """L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100.""" + messages = dummy_messages_from_mix_data(content_text=PROMPT) + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 480, + "width": 640, + "num_frames": 5, + "num_inference_steps": 2, + "guidance_scale": 6.0, + "negative_prompt": NEGATIVE_PROMPT, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) diff --git a/vllm_omni/diffusion/models/hunyuan_video/__init__.py b/vllm_omni/diffusion/models/hunyuan_video/__init__.py new file mode 100644 index 000000000..7be2ed33c --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_video/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HunyuanVideo-1.5 diffusion model components (T2V and I2V).""" + +from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import ( + HunyuanVideo15Transformer3DModel, +) +from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import ( + HunyuanVideo15Pipeline, + get_hunyuan_video_15_post_process_func, +) +from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v import ( + HunyuanVideo15I2VPipeline, + get_hunyuan_video_15_i2v_post_process_func, + get_hunyuan_video_15_i2v_pre_process_func, +) + +__all__ = [ + "HunyuanVideo15Pipeline", + "HunyuanVideo15I2VPipeline", + "HunyuanVideo15Transformer3DModel", + "get_hunyuan_video_15_post_process_func", + "get_hunyuan_video_15_i2v_post_process_func", + "get_hunyuan_video_15_i2v_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py new file mode 100644 index 000000000..2f7318cef --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py @@ -0,0 +1,791 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from diffusers.models.embeddings import ( + CombinedTimestepTextProjEmbeddings, + TimestepEmbedding, + Timesteps, + get_1d_rotary_pos_embed, +) +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero +from torch import nn +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward + +logger = init_logger(__name__) + + +class HunyuanVideo15PatchEmbed(nn.Module): + def __init__( + self, + patch_size: int | tuple[int, int, int] = 16, + in_chans: int = 3, + embed_dim: int = 768, + ) -> None: + super().__init__() + patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.proj(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) # BCFHW -> BNC + return hidden_states + + +class HunyuanVideo15AdaNorm(nn.Module): + def __init__(self, in_features: int, out_features: int | None = None) -> None: + super().__init__() + out_features = out_features or 2 * in_features + self.linear = nn.Linear(in_features, out_features) + self.nonlinearity = nn.SiLU() + + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + temb = self.linear(self.nonlinearity(temb)) + gate_msa, gate_mlp = temb.chunk(2, dim=1) + gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) + return gate_msa, gate_mlp + + +class HunyuanVideo15TimeEmbedding(nn.Module): + def __init__(self, embedding_dim: int, use_meanflow: bool = False): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.use_meanflow = use_meanflow + self.time_proj_r = None + self.timestep_embedder_r = None + if use_meanflow: + self.time_proj_r = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder_r = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward( + self, + timestep: torch.Tensor, + timestep_r: torch.Tensor | None = None, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype)) + + if timestep_r is not None and self.timestep_embedder_r is not None: + timesteps_proj_r = self.time_proj_r(timestep_r) + timesteps_emb_r = self.timestep_embedder_r(timesteps_proj_r.to(dtype=timestep.dtype)) + timesteps_emb = timesteps_emb + timesteps_emb_r + + return timesteps_emb + + +class HunyuanVideo15RotaryPosEmbed(nn.Module): + def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None: + super().__init__() + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.rope_dim = rope_dim + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + rope_sizes = [num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size] + + axes_grids = [] + for i in range(len(rope_sizes)): + grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32) + axes_grids.append(grid) + grid = torch.meshgrid(*axes_grids, indexing="ij") + grid = torch.stack(grid, dim=0) # [3, T', H', W'] + + freqs = [] + for i in range(3): + # use_real=False returns complex [seq, dim/2]; we extract cos/sin manually + # to get half-dim tensors compatible with RotaryEmbedding (which doubles internally). + freq_cis = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=False) + freqs.append((freq_cis.real, freq_cis.imag)) + + freqs_cos = torch.cat([f[0] for f in freqs], dim=1).float() + freqs_sin = torch.cat([f[1] for f in freqs], dim=1).float() + return freqs_cos, freqs_sin + + +class HunyuanVideo15IndividualTokenRefinerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + + # Not TP + from diffusers.models.attention_processor import Attention as DiffusersAttention + + self.attn = DiffusersAttention( + query_dim=hidden_size, + cross_attention_dim=None, + heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + ) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + + from diffusers.models.attention import FeedForward as DiffusersFeedForward + + self.ff = DiffusersFeedForward( + hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate + ) + + self.norm_out = HunyuanVideo15AdaNorm(hidden_size, 2 * hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states) + + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=None, + attention_mask=attention_mask, + ) + + gate_msa, gate_mlp = self.norm_out(temb) + hidden_states = hidden_states + attn_output * gate_msa + + ff_output = self.ff(self.norm2(hidden_states)) + hidden_states = hidden_states + ff_output * gate_mlp + + return hidden_states + + +class HunyuanVideo15IndividualTokenRefiner(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + self.refiner_blocks = nn.ModuleList( + [ + HunyuanVideo15IndividualTokenRefinerBlock( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + self_attn_mask = None + if attention_mask is not None: + batch_size = attention_mask.shape[0] + seq_len = attention_mask.shape[1] + attention_mask = attention_mask.to(hidden_states.device).bool() + self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1) + self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + + for block in self.refiner_blocks: + hidden_states = block(hidden_states, temb, self_attn_mask) + + return hidden_states + + +class HunyuanVideo15TokenRefiner(nn.Module): + def __init__( + self, + in_channels: int, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + hidden_size = num_attention_heads * attention_head_dim + + self.time_text_embed = CombinedTimestepTextProjEmbeddings( + embedding_dim=hidden_size, pooled_projection_dim=in_channels + ) + self.proj_in = nn.Linear(in_channels, hidden_size, bias=True) + self.token_refiner = HunyuanVideo15IndividualTokenRefiner( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_layers=num_layers, + mlp_width_ratio=mlp_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + attention_mask: torch.LongTensor | None = None, + ) -> torch.Tensor: + if attention_mask is None: + pooled_projections = hidden_states.mean(dim=1) + else: + original_dtype = hidden_states.dtype + mask_float = attention_mask.float().unsqueeze(-1) + pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1) + pooled_projections = pooled_projections.to(original_dtype) + + temb = self.time_text_embed(timestep, pooled_projections) + hidden_states = self.proj_in(hidden_states) + hidden_states = self.token_refiner(hidden_states, temb, attention_mask) + + return hidden_states + + +class HunyuanVideo15ByT5TextProjection(nn.Module): + def __init__(self, in_features: int, hidden_size: int, out_features: int): + super().__init__() + self.norm = nn.LayerNorm(in_features) + self.linear_1 = nn.Linear(in_features, hidden_size) + self.linear_2 = nn.Linear(hidden_size, hidden_size) + self.linear_3 = nn.Linear(hidden_size, out_features) + self.act_fn = nn.GELU() + + def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(encoder_hidden_states) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_2(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_3(hidden_states) + return hidden_states + + +class HunyuanVideo15ImageProjection(nn.Module): + def __init__(self, in_channels: int, hidden_size: int): + super().__init__() + self.norm_in = nn.LayerNorm(in_channels) + self.linear_1 = nn.Linear(in_channels, in_channels) + self.act_fn = nn.GELU() + self.linear_2 = nn.Linear(in_channels, hidden_size) + self.norm_out = nn.LayerNorm(hidden_size) + + def forward(self, image_embeds: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm_in(image_embeds) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_2(hidden_states) + hidden_states = self.norm_out(hidden_states) + return hidden_states + + +class HunyuanVideo15Attention(nn.Module): + """Dual-stream joint attention with TP optimization. + + Key difference from FluxAttention: RoPE is applied **only** to the video + stream Q/K *before* concatenation with the encoder stream, whereas Flux + applies RoPE to the concatenated Q/K. + """ + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + bias: bool = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, + out_bias: bool = True, + eps: float = 1e-6, + out_dim: int | None = None, + ): + super().__init__() + + self.head_dim = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.out_dim = out_dim if out_dim is not None else query_dim + self.heads = out_dim // dim_head if out_dim is not None else heads + self.added_kv_proj_dim = added_kv_proj_dim + + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.to_qkv = QKVParallelLinear( + hidden_size=query_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + bias=bias, + ) + + self.to_out = nn.ModuleList( + [ + RowParallelLinear( + self.inner_dim, + self.out_dim, + bias=out_bias, + input_is_parallel=True, + return_bias=False, + ), + nn.Identity(), # placeholder for dropout (none used) + ] + ) + + if added_kv_proj_dim is not None: + self.norm_added_q = RMSNorm(dim_head, eps=eps) + self.norm_added_k = RMSNorm(dim_head, eps=eps) + + self.add_kv_proj = QKVParallelLinear( + hidden_size=self.added_kv_proj_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + bias=added_proj_bias, + ) + + self.to_add_out = RowParallelLinear( + self.inner_dim, + query_dim, + bias=out_bias, + input_is_parallel=True, + return_bias=False, + ) + + self.rope = RotaryEmbedding(is_neox_style=False) + self.attn = Attention( + num_heads=self.to_qkv.num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + num_kv_heads=self.to_qkv.num_kv_heads, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + qkv, _ = self.to_qkv(hidden_states) + q_size = self.to_qkv.num_heads * self.head_dim + kv_size = self.to_qkv.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + query = query.unflatten(-1, (self.to_qkv.num_heads, -1)) + key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) + value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) + + query = self.norm_q(query) + key = self.norm_k(key) + + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) + + if encoder_hidden_states is not None: + encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) + add_q_size = self.add_kv_proj.num_heads * self.head_dim + add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim + encoder_query, encoder_key, encoder_value = encoder_qkv.split( + [add_q_size, add_kv_size, add_kv_size], dim=-1 + ) + + encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1)) + encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1)) + encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1)) + + encoder_query = self.norm_added_q(encoder_query) + encoder_key = self.norm_added_k(encoder_key) + + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + + attn_metadata = None + if attention_mask is not None: + seq_len = query.shape[1] + # Pad mask to full sequence length (video + encoder tokens) + # Keep mask 2D (batch_size, seq_len) — each attention backend + # handles reshaping internally (flash_attn uses unpadding, + # SDPA expands to 4D via _maybe_reshape_attn_mask). + attention_mask = F.pad(attention_mask, (seq_len - attention_mask.shape[1], 0), value=True) + attention_mask = attention_mask.bool() + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + if encoder_hidden_states is not None: + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + [hidden_states.shape[1] - encoder_hidden_states.shape[1], encoder_hidden_states.shape[1]], dim=1 + ) + hidden_states = self.to_out[0](hidden_states) + encoder_hidden_states = self.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + else: + hidden_states = self.to_out[0](hidden_states) + return hidden_states + + +class HunyuanVideo15TransformerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + + self.attn = HunyuanVideo15Attention( + query_dim=hidden_size, + added_kv_proj_dim=hidden_size, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + bias=True, + eps=1e-6, + ) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio) + + self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( + encoder_hidden_states, emb=temb + ) + + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=freqs_cis, + ) + + hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1) + + norm_hidden_states = self.norm2(hidden_states) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + + ff_output = self.ff(norm_hidden_states) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + + hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output + encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + + return hidden_states, encoder_hidden_states + + +class HunyuanVideo15Transformer3DModel(nn.Module): + """ + HunyuanVideo-1.5 Transformer with TP-optimized dual-stream attention. + + Ported from diffusers ``HunyuanVideo15Transformer3DModel`` with vllm-omni + tensor-parallel layers for the 54 main transformer blocks. + """ + + _repeated_blocks = ["HunyuanVideo15TransformerBlock"] + _layerwise_offload_blocks_attr = "transformer_blocks" + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], + } + + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return "transformer_blocks" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + + def __init__( + self, + od_config: OmniDiffusionConfig, + in_channels: int = 65, + out_channels: int = 32, + num_attention_heads: int = 16, + attention_head_dim: int = 128, + num_layers: int = 54, + num_refiner_layers: int = 2, + mlp_ratio: float = 4.0, + patch_size: int = 1, + patch_size_t: int = 1, + qk_norm: str = "rms_norm", + text_embed_dim: int = 3584, + text_embed_2_dim: int = 1472, + image_embed_dim: int = 1152, + rope_theta: float = 256.0, + rope_axes_dim: tuple[int, ...] = (16, 56, 56), + target_size: int = 640, + task_type: str = "i2v", + use_meanflow: bool = False, + ): + super().__init__() + + # Allow tf_model_config to override num_layers + model_config = od_config.tf_model_config + num_layers = getattr(model_config, "num_layers", num_layers) + + self.parallel_config = od_config.parallel_config + self.in_channels = in_channels + self.out_channels = out_channels or in_channels + inner_dim = num_attention_heads * attention_head_dim + self.inner_dim = inner_dim + self.patch_size = patch_size + self.patch_size_t = patch_size_t + + self.x_embedder = HunyuanVideo15PatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim) + self.image_embedder = HunyuanVideo15ImageProjection(image_embed_dim, inner_dim) + + self.context_embedder = HunyuanVideo15TokenRefiner( + text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers + ) + self.context_embedder_2 = HunyuanVideo15ByT5TextProjection(text_embed_2_dim, 2048, inner_dim) + + self.time_embed = HunyuanVideo15TimeEmbedding(inner_dim, use_meanflow=use_meanflow) + self.cond_type_embed = nn.Embedding(3, inner_dim) + + self.rope = HunyuanVideo15RotaryPosEmbed(patch_size, patch_size_t, list(rope_axes_dim), rope_theta) + + self.transformer_blocks = nn.ModuleList( + [ + HunyuanVideo15TransformerBlock( + num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * self.out_channels) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + timestep_r: torch.LongTensor | None = None, + encoder_hidden_states_2: torch.Tensor | None = None, + encoder_attention_mask_2: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + image_embeds_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size_t, self.patch_size, self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + image_rotary_emb = self.rope(hidden_states) + + temb = self.time_embed(timestep, timestep_r=timestep_r) + + hidden_states = self.x_embedder(hidden_states) + + encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) + + encoder_hidden_states_cond_emb = self.cond_type_embed( + torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long) + ) + encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb + + encoder_hidden_states_2 = self.context_embedder_2(encoder_hidden_states_2) + + encoder_hidden_states_2_cond_emb = self.cond_type_embed( + torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long) + ) + encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb + + encoder_hidden_states_3 = self.image_embedder(image_embeds) + if image_embeds_mask is not None: + encoder_attention_mask_3 = image_embeds_mask + else: + # Fallback: detect T2V by checking if image_embeds are all zeros + is_t2v = torch.all(image_embeds == 0) + if is_t2v: + encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0 + encoder_attention_mask_3 = torch.zeros( + (batch_size, encoder_hidden_states_3.shape[1]), + dtype=encoder_attention_mask.dtype, + device=encoder_attention_mask.device, + ) + else: + encoder_attention_mask_3 = torch.ones( + (batch_size, encoder_hidden_states_3.shape[1]), + dtype=encoder_attention_mask.dtype, + device=encoder_attention_mask.device, + ) + encoder_hidden_states_3_cond_emb = self.cond_type_embed( + 2 * torch.ones_like(encoder_hidden_states_3[:, :, 0], dtype=torch.long) + ) + encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb + + # Token reordering: [valid_image, valid_byte5, valid_mllm, padding] + encoder_attention_mask = encoder_attention_mask.bool() + encoder_attention_mask_2 = encoder_attention_mask_2.bool() + encoder_attention_mask_3 = encoder_attention_mask_3.bool() + new_encoder_hidden_states = [] + new_encoder_attention_mask = [] + + for text, text_mask, text_2, text_mask_2, image, image_mask in zip( + encoder_hidden_states, + encoder_attention_mask, + encoder_hidden_states_2, + encoder_attention_mask_2, + encoder_hidden_states_3, + encoder_attention_mask_3, + ): + new_encoder_hidden_states.append( + torch.cat( + [ + image[image_mask], + text_2[text_mask_2], + text[text_mask], + image[~image_mask], + torch.zeros_like(text_2[~text_mask_2]), + torch.zeros_like(text[~text_mask]), + ], + dim=0, + ) + ) + new_encoder_attention_mask.append( + torch.cat( + [ + image_mask[image_mask], + text_mask_2[text_mask_2], + text_mask[text_mask], + image_mask[~image_mask], + text_mask_2[~text_mask_2], + text_mask[~text_mask], + ], + dim=0, + ) + ) + + encoder_hidden_states = torch.stack(new_encoder_hidden_states) + encoder_attention_mask = torch.stack(new_encoder_attention_mask) + + for block in self.transformer_blocks: + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + encoder_attention_mask, + image_rotary_emb, + ) + + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, -1, p_t, p_h, p_w + ) + hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) + hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (hidden_states,) + + return Transformer2DModelOutput(sample=hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + # self-attn + (".to_qkv", ".to_q", "q"), + (".to_qkv", ".to_k", "k"), + (".to_qkv", ".to_v", "v"), + # cross-attn + (".add_kv_proj", ".add_q_proj", "q"), + (".add_kv_proj", ".add_k_proj", "k"), + (".add_kv_proj", ".add_v_proj", "v"), + ] + + params_dict = dict(self.named_parameters()) + + # Load buffers for beta and eps (if any normalizations use them) + for name, buffer in self.named_buffers(): + if name.endswith(".beta") or name.endswith(".eps"): + params_dict[name] = buffer + + loaded_params: set[str] = set() + for name, loaded_weight in weights: + original_name = name + lookup_name = name + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in original_name: + continue + lookup_name = original_name.replace(weight_name, param_name) + if lookup_name not in params_dict: + break + param = params_dict[lookup_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Handle .to_out.0. -> .to_out. remapping + if lookup_name not in params_dict and ".to_out.0." in lookup_name: + lookup_name = lookup_name.replace(".to_out.0.", ".to_out.") + if lookup_name not in params_dict: + continue + param = params_dict[lookup_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(original_name) + loaded_params.add(lookup_name) + return loaded_params diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py new file mode 100644 index 000000000..643d471ff --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +import os +import re +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +from diffusers import AutoencoderKLHunyuanVideo15 +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from torch import nn +from transformers import AutoConfig, ByT5Tokenizer, Qwen2_5_VLTextModel, Qwen2Tokenizer +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +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.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel +from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + + +def retrieve_latents( + encoder_output: torch.Tensor, + generator: torch.Generator | None = None, + sample_mode: str = "sample", +) -> torch.Tensor: + """Extract latents from VAE encoder output.""" + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +def format_text_input(prompt: list[str], system_message: str) -> list[list[dict[str, str]]]: + return [ + [{"role": "system", "content": system_message}, {"role": "user", "content": p if p else " "}] for p in prompt + ] + + +def extract_glyph_texts(prompt: str) -> str | None: + pattern = r'"(.*?)"|' + r"\u201c(.*?)\u201d" + matches = re.findall(pattern, prompt) + result = [match[0] or match[1] for match in matches] + result = list(dict.fromkeys(result)) if len(result) > 1 else result + + if result: + return ". ".join([f'Text "{text}"' for text in result]) + ". " + return None + + +def get_hunyuan_video_15_post_process_func(od_config: OmniDiffusionConfig): + video_processor = VideoProcessor(vae_scale_factor=16) + + def post_process_func(video: torch.Tensor, output_type: str = "np"): + if output_type == "latent": + return video + return video_processor.postprocess_video(video, output_type=output_type) + + return post_process_func + + +class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin): + 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 + local_files_only = os.path.exists(model) + + self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.text_encoder = Qwen2_5_VLTextModel.from_pretrained( + model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only + ).to(self.device) + + self.tokenizer_2 = ByT5Tokenizer.from_pretrained( + model, subfolder="tokenizer_2", local_files_only=local_files_only + ) + t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only) + self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device) + + self.vae = AutoencoderKLHunyuanVideo15.from_pretrained( + model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only + ).to(self.device) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + + # Override the scheduler's shift if flow_shift is explicitly provided. + # Must set _shift directly since the property has no setter. + if od_config.flow_shift is not None: + self.scheduler._shift = od_config.flow_shift + + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel) + self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs) + + # Check if model uses meanflow (distilled variants) + self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False) + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="text_encoder_2", + revision=None, + prefix="text_encoder_2.", + fall_back_to_pt=True, + ), + ] + + self.vae_scale_factor_temporal = ( + self.vae.temporal_compression_ratio if hasattr(self.vae, "temporal_compression_ratio") else 4 + ) + self.vae_scale_factor_spatial = ( + self.vae.spatial_compression_ratio if hasattr(self.vae, "spatial_compression_ratio") else 16 + ) + self.num_channels_latents = self.vae.config.latent_channels if hasattr(self.vae, "config") else 32 + + # fmt: off + # NOTE: whitespace must match the diffusers reference exactly because + # prompt_template_encode_start_idx=108 depends on the tokenised length of + # this system message. The backslash-continuation includes leading spaces. + self.system_message = "You are a helpful assistant. Describe the video by detailing the following aspects: \ + 1. The main content and theme of the video. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ + 4. background environment, light, style and atmosphere. \ + 5. camera angles, movements, and transitions used in the video." + # fmt: on + self.prompt_template_encode_start_idx = 108 + self.tokenizer_max_length = 1000 + self.tokenizer_2_max_length = 256 + self.vision_num_semantic_tokens = 729 + self.vision_states_dim = 1152 + + self._guidance_scale = None + self._num_timesteps = None + self._current_timestep = None + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + def _get_mllm_prompt_embeds( + self, + prompt: list[str], + device: torch.device, + dtype: torch.dtype, + num_hidden_layers_to_skip: int = 2, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt_formatted = format_text_input(prompt, self.system_message) + + text_inputs = self.tokenizer.apply_chat_template( + prompt_formatted, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + padding="max_length", + max_length=self.tokenizer_max_length + self.prompt_template_encode_start_idx, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device=device) + prompt_attention_mask = text_inputs.attention_mask.to(device=device) + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, + ).hidden_states[-(num_hidden_layers_to_skip + 1)] + + crop_start = self.prompt_template_encode_start_idx + if crop_start is not None and crop_start > 0: + prompt_embeds = prompt_embeds[:, crop_start:] + prompt_attention_mask = prompt_attention_mask[:, crop_start:] + + return prompt_embeds.to(dtype=dtype), prompt_attention_mask + + def _get_byte5_prompt_embeds( + self, + prompt: list[str], + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt_embeds_list = [] + prompt_embeds_mask_list = [] + + for p in prompt: + glyph_text = extract_glyph_texts(p) + + if glyph_text is None: + glyph_text_embeds = torch.zeros( + (1, self.tokenizer_2_max_length, self.text_encoder_2.config.d_model), + device=device, + dtype=dtype, + ) + glyph_text_embeds_mask = torch.zeros((1, self.tokenizer_2_max_length), device=device, dtype=torch.int64) + else: + txt_tokens = self.tokenizer_2( + glyph_text, + padding="max_length", + max_length=self.tokenizer_2_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ).to(device) + + glyph_text_embeds = self.text_encoder_2( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask.float(), + )[0] + glyph_text_embeds = glyph_text_embeds.to(device=device, dtype=dtype) + glyph_text_embeds_mask = txt_tokens.attention_mask.to(device=device) + + prompt_embeds_list.append(glyph_text_embeds) + prompt_embeds_mask_list.append(glyph_text_embeds_mask) + + return torch.cat(prompt_embeds_list, dim=0), torch.cat(prompt_embeds_mask_list, dim=0) + + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device, + dtype: torch.dtype, + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = False, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + prompt = [prompt] if isinstance(prompt, str) else prompt + + prompt_embeds, prompt_embeds_mask = self._get_mllm_prompt_embeds(prompt, device, dtype) + prompt_embeds_2, prompt_embeds_mask_2 = self._get_byte5_prompt_embeds(prompt, device, dtype) + + # Cast masks to model dtype to match diffusers behaviour + prompt_embeds_mask = prompt_embeds_mask.to(dtype=dtype) + prompt_embeds_mask_2 = prompt_embeds_mask_2.to(dtype=dtype) + + negative_prompt_embeds = None + negative_prompt_embeds_mask = None + negative_prompt_embeds_2 = None + negative_prompt_embeds_mask_2 = None + + if do_classifier_free_guidance and negative_prompt is not None: + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds( + negative_prompt, device, dtype + ) + negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds( + negative_prompt, device, dtype + ) + negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype) + negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype) + elif do_classifier_free_guidance: + negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds([""], device, dtype) + negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds([""], device, dtype) + negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype) + negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype) + + return ( + prompt_embeds, + prompt_embeds_mask, + prompt_embeds_2, + prompt_embeds_mask_2, + negative_prompt_embeds, + negative_prompt_embeds_mask, + negative_prompt_embeds_2, + negative_prompt_embeds_mask_2, + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + num_frames: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + shape = ( + batch_size, + self.num_channels_latents, + (num_frames - 1) // self.vae_scale_factor_temporal + 1, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.") + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def prepare_cond_latents_and_mask( + self, latents: torch.Tensor, dtype: torch.dtype, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare zero condition latents and mask for T2V mode.""" + batch, channels, frames, height, width = latents.shape + cond_latents = torch.zeros(batch, channels, frames, height, width, dtype=dtype, device=device) + mask = torch.zeros(batch, 1, frames, height, width, dtype=dtype, device=device) + return cond_latents, mask + + def predict_noise(self, **kwargs: Any) -> torch.Tensor: + return self.transformer(**kwargs)[0] + + def forward( + self, + req: OmniDiffusionRequest, + num_inference_steps: int = 50, + guidance_scale: float = 6.0, + height: int = 480, + width: int = 832, + num_frames: int = 121, + output_type: str | None = "np", + generator: torch.Generator | list[torch.Generator] | None = None, + **kwargs, + ) -> DiffusionOutput: + if len(req.prompts) > 1: + raise ValueError("This model only supports a single prompt per request.") + if len(req.prompts) == 1: + prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") + negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") + else: + raise ValueError("Prompt is required for HunyuanVideo-1.5 generation.") + + height = req.sampling_params.height or height + width = req.sampling_params.width or width + num_frames_val = req.sampling_params.num_frames if req.sampling_params.num_frames else num_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 + + do_cfg = guidance_scale > 1.0 + + device = self.device + dtype = self.transformer.transformer_blocks[0].norm1.linear.weight.dtype + + if generator is None: + generator = req.sampling_params.generator + if generator is None and req.sampling_params.seed is not None: + generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + + ( + prompt_embeds, + prompt_embeds_mask, + prompt_embeds_2, + prompt_embeds_mask_2, + negative_prompt_embeds, + negative_prompt_embeds_mask, + negative_prompt_embeds_2, + negative_prompt_embeds_mask_2, + ) = self.encode_prompt( + prompt=prompt, + device=device, + dtype=dtype, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_cfg, + ) + + batch_size = prompt_embeds.shape[0] + + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + num_frames=num_frames_val, + dtype=dtype, + device=device, + generator=generator, + latents=req.sampling_params.latents, + ) + cond_latents, mask = self.prepare_cond_latents_and_mask(latents, dtype, device) + + # Image embeds (zeros for T2V, no mask — let the transformer detect T2V + # via the all-zeros check, matching the diffusers reference behaviour) + image_embeds = torch.zeros( + batch_size, + self.vision_num_semantic_tokens, + self.vision_states_dim, + dtype=dtype, + device=device, + ) + + # Timesteps — the scheduler handles flow_shift via its config (`shift` param). + # We just provide unshifted linear sigmas, matching the diffusers reference. + sigmas = np.linspace(1.0, 0.0, num_steps + 1)[:-1] + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + for i, t in enumerate(timesteps): + self._current_timestep = t + + latent_model_input = torch.cat([latents, cond_latents, mask], dim=1) + timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype) + + timestep_r = None + if self.use_meanflow: + if i == len(timesteps) - 1: + timestep_r = torch.tensor([0.0], device=device) + else: + timestep_r = timesteps[i + 1] + timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "timestep_r": timestep_r, + "encoder_hidden_states": prompt_embeds, + "encoder_attention_mask": prompt_embeds_mask, + "encoder_hidden_states_2": prompt_embeds_2, + "encoder_attention_mask_2": prompt_embeds_mask_2, + "image_embeds": image_embeds, + "return_dict": False, + } + + negative_kwargs = None + if do_cfg and negative_prompt_embeds is not None: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "timestep_r": timestep_r, + "encoder_hidden_states": negative_prompt_embeds, + "encoder_attention_mask": negative_prompt_embeds_mask, + "encoder_hidden_states_2": negative_prompt_embeds_2, + "encoder_attention_mask_2": negative_prompt_embeds_mask_2, + "image_embeds": image_embeds, + "return_dict": False, + } + + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_cfg and negative_kwargs is not None, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=req.sampling_params.cfg_normalize, + ) + + latents = self.scheduler_step_maybe_with_cfg( + noise_pred, + t, + latents, + do_true_cfg=do_cfg and negative_kwargs is not None, + ) + + self._current_timestep = None + + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + if output_type == "latent": + output = latents + else: + latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor + output = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput(output=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py new file mode 100644 index 000000000..d68c43125 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py @@ -0,0 +1,594 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +import os +from collections.abc import Iterable +from typing import Any, cast + +import numpy as np +import PIL.Image +import torch +from diffusers import AutoencoderKLHunyuanVideo15 +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from torch import nn +from transformers import ( + AutoConfig, + ByT5Tokenizer, + Qwen2_5_VLTextModel, + Qwen2Tokenizer, + SiglipImageProcessor, + SiglipVisionModel, +) +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +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.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel +from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import ( + extract_glyph_texts, + format_text_input, + get_hunyuan_video_15_post_process_func, + retrieve_latents, +) +from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + + +def get_hunyuan_video_15_i2v_post_process_func(od_config: OmniDiffusionConfig): + return get_hunyuan_video_15_post_process_func(od_config) + + +def get_hunyuan_video_15_i2v_pre_process_func(od_config: OmniDiffusionConfig): + """Pre-process function for I2V: load and resize image.""" + + max_area = 480 * 832 + divisor = 16 # Must be divisible by VAE spatial compression + + def pre_process_func(req: OmniDiffusionRequest) -> OmniDiffusionRequest: + if not req.prompts: + return req + + prompt_data = req.prompts[0] + if isinstance(prompt_data, str): + return req + + multi_modal_data = prompt_data.get("multi_modal_data", {}) + raw_image = multi_modal_data.get("image", None) + + if raw_image is None: + return req + + if isinstance(raw_image, list): + raw_image = raw_image[0] + + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image).convert("RGB") + elif isinstance(raw_image, PIL.Image.Image): + image = raw_image.convert("RGB") + else: + return req + + height = req.sampling_params.height + width = req.sampling_params.width + + if height is None or width is None: + w, h = image.size + aspect = w / h + target_h = int((max_area / aspect) ** 0.5) + target_w = int(target_h * aspect) + height = height or ((target_h + divisor - 1) // divisor * divisor) + width = width or ((target_w + divisor - 1) // divisor * divisor) + req.sampling_params.height = height + req.sampling_params.width = width + + return req + + return pre_process_func + + +class HunyuanVideo15I2VPipeline(nn.Module, CFGParallelMixin, SupportImageInput): + support_image_input = True + color_format = "RGB" + + 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 + local_files_only = os.path.exists(model) + + self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.text_encoder = Qwen2_5_VLTextModel.from_pretrained( + model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only + ).to(self.device) + + self.tokenizer_2 = ByT5Tokenizer.from_pretrained( + model, subfolder="tokenizer_2", local_files_only=local_files_only + ) + t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only) + self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device) + + self.image_encoder = SiglipVisionModel.from_pretrained( + model, subfolder="image_encoder", torch_dtype=dtype, local_files_only=local_files_only + ).to(self.device) + self.feature_extractor = SiglipImageProcessor.from_pretrained( + model, subfolder="feature_extractor", local_files_only=local_files_only + ) + + self.vae = AutoencoderKLHunyuanVideo15.from_pretrained( + model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only + ).to(self.device) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + + # Override the scheduler's shift if flow_shift is explicitly provided. + # Must set _shift directly since the property has no setter. + if od_config.flow_shift is not None: + self.scheduler._shift = od_config.flow_shift + + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel) + self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs) + + self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False) + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="text_encoder_2", + revision=None, + prefix="text_encoder_2.", + fall_back_to_pt=True, + ), + ] + + self.vae_scale_factor_temporal = ( + self.vae.temporal_compression_ratio if hasattr(self.vae, "temporal_compression_ratio") else 4 + ) + self.vae_scale_factor_spatial = ( + self.vae.spatial_compression_ratio if hasattr(self.vae, "spatial_compression_ratio") else 16 + ) + self.num_channels_latents = self.vae.config.latent_channels if hasattr(self.vae, "config") else 32 + + # fmt: off + # NOTE: whitespace must match the diffusers reference exactly because + # prompt_template_encode_start_idx=108 depends on the tokenised length of + # this system message. The backslash-continuation includes leading spaces. + self.system_message = "You are a helpful assistant. Describe the video by detailing the following aspects: \ + 1. The main content and theme of the video. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ + 4. background environment, light, style and atmosphere. \ + 5. camera angles, movements, and transitions used in the video." + # fmt: on + self.prompt_template_encode_start_idx = 108 + self.tokenizer_max_length = 1000 + self.tokenizer_2_max_length = 256 + self.vision_num_semantic_tokens = 729 + self.vision_states_dim = 1152 + + self._guidance_scale = None + self._num_timesteps = None + self._current_timestep = None + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + def _get_mllm_prompt_embeds( + self, + prompt: list[str], + device: torch.device, + dtype: torch.dtype, + num_hidden_layers_to_skip: int = 2, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt_formatted = format_text_input(prompt, self.system_message) + + text_inputs = self.tokenizer.apply_chat_template( + prompt_formatted, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + padding="max_length", + max_length=self.tokenizer_max_length + self.prompt_template_encode_start_idx, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device=device) + prompt_attention_mask = text_inputs.attention_mask.to(device=device) + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, + ).hidden_states[-(num_hidden_layers_to_skip + 1)] + + crop_start = self.prompt_template_encode_start_idx + if crop_start is not None and crop_start > 0: + prompt_embeds = prompt_embeds[:, crop_start:] + prompt_attention_mask = prompt_attention_mask[:, crop_start:] + + return prompt_embeds.to(dtype=dtype), prompt_attention_mask + + def _get_byte5_prompt_embeds( + self, + prompt: list[str], + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt_embeds_list = [] + prompt_embeds_mask_list = [] + + for p in prompt: + glyph_text = extract_glyph_texts(p) + + if glyph_text is None: + glyph_text_embeds = torch.zeros( + (1, self.tokenizer_2_max_length, self.text_encoder_2.config.d_model), + device=device, + dtype=dtype, + ) + glyph_text_embeds_mask = torch.zeros((1, self.tokenizer_2_max_length), device=device, dtype=torch.int64) + else: + txt_tokens = self.tokenizer_2( + glyph_text, + padding="max_length", + max_length=self.tokenizer_2_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ).to(device) + + glyph_text_embeds = self.text_encoder_2( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask.float(), + )[0] + glyph_text_embeds = glyph_text_embeds.to(device=device, dtype=dtype) + glyph_text_embeds_mask = txt_tokens.attention_mask.to(device=device) + + prompt_embeds_list.append(glyph_text_embeds) + prompt_embeds_mask_list.append(glyph_text_embeds_mask) + + return torch.cat(prompt_embeds_list, dim=0), torch.cat(prompt_embeds_mask_list, dim=0) + + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device, + dtype: torch.dtype, + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = False, + ) -> tuple: + prompt = [prompt] if isinstance(prompt, str) else prompt + + prompt_embeds, prompt_embeds_mask = self._get_mllm_prompt_embeds(prompt, device, dtype) + prompt_embeds_2, prompt_embeds_mask_2 = self._get_byte5_prompt_embeds(prompt, device, dtype) + + # Cast masks to model dtype to match diffusers behaviour + prompt_embeds_mask = prompt_embeds_mask.to(dtype=dtype) + prompt_embeds_mask_2 = prompt_embeds_mask_2.to(dtype=dtype) + + negative_prompt_embeds = None + negative_prompt_embeds_mask = None + negative_prompt_embeds_2 = None + negative_prompt_embeds_mask_2 = None + + if do_classifier_free_guidance and negative_prompt is not None: + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds( + negative_prompt, device, dtype + ) + negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds( + negative_prompt, device, dtype + ) + negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype) + negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype) + elif do_classifier_free_guidance: + negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds([""], device, dtype) + negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds([""], device, dtype) + negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype) + negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype) + + return ( + prompt_embeds, + prompt_embeds_mask, + prompt_embeds_2, + prompt_embeds_mask_2, + negative_prompt_embeds, + negative_prompt_embeds_mask, + negative_prompt_embeds_2, + negative_prompt_embeds_mask_2, + ) + + def _get_image_embeds( + self, + image: PIL.Image.Image, + device: torch.device, + ) -> torch.Tensor: + """Get semantic image embeddings from SigLIP.""" + image_encoder_dtype = next(self.image_encoder.parameters()).dtype + pixel_values = self.feature_extractor(images=image, do_resize=True, return_tensors="pt").pixel_values.to( + device=device, dtype=image_encoder_dtype + ) + image_enc_hidden_states = self.image_encoder(pixel_values).last_hidden_state + return image_enc_hidden_states + + def _get_image_latents( + self, + image: PIL.Image.Image, + height: int, + width: int, + device: torch.device, + ) -> torch.Tensor: + """Encode image to VAE latents for first-frame conditioning.""" + video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + image_tensor = video_processor.preprocess(image, height=height, width=width) + image_tensor = image_tensor.unsqueeze(2).to(device=device, dtype=self.vae.dtype) # [B, C, 1, H, W] + image_latents = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax") + image_latents = image_latents * self.vae.config.scaling_factor + return image_latents + + def prepare_cond_latents_and_mask( + self, + latents: torch.Tensor, + image: PIL.Image.Image, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare condition latents and mask for I2V: first frame = image, rest = zeros.""" + batch, channels, frames, lat_height, lat_width = latents.shape + + image_latents = self._get_image_latents(image, height, width, device) + + # Repeat image latent across all frames, then zero out non-first frames + latent_condition = image_latents.repeat(batch, 1, frames, 1, 1) + latent_condition[:, :, 1:, :, :] = 0 + latent_condition = latent_condition.to(device=device, dtype=dtype) + + # Mask: 1.0 for first frame (conditioned), 0.0 for rest (to denoise) + latent_mask = torch.zeros(batch, 1, frames, lat_height, lat_width, dtype=dtype, device=device) + latent_mask[:, :, 0, :, :] = 1.0 + + return latent_condition, latent_mask + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + num_frames: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_channels = self.num_channels_latents + shape = ( + batch_size, + num_channels, + (num_frames - 1) // self.vae_scale_factor_temporal + 1, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.") + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def predict_noise(self, **kwargs: Any) -> torch.Tensor: + return self.transformer(**kwargs)[0] + + def forward( + self, + req: OmniDiffusionRequest, + num_inference_steps: int = 50, + guidance_scale: float = 6.0, + height: int = 480, + width: int = 832, + num_frames: int = 121, + output_type: str | None = "np", + generator: torch.Generator | list[torch.Generator] | None = None, + **kwargs, + ) -> DiffusionOutput: + if len(req.prompts) > 1: + raise ValueError("This model only supports a single prompt per request.") + if len(req.prompts) == 1: + prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") + negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") + else: + raise ValueError("Prompt is required for HunyuanVideo-1.5 I2V generation.") + + multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else {} + raw_image = multi_modal_data.get("image", None) + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning("Received a list of images. Only the first image will be used.") + raw_image = raw_image[0] + + if raw_image is None: + raise ValueError( + "Image is required for HunyuanVideo-1.5 I2V generation. Pass it via 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.convert("RGB") + else: + image = cast(PIL.Image.Image, raw_image) + + height = req.sampling_params.height or height + width = req.sampling_params.width or width + num_frames_val = req.sampling_params.num_frames if req.sampling_params.num_frames else num_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 + + do_cfg = guidance_scale > 1.0 + + device = self.device + dtype = self.transformer.transformer_blocks[0].norm1.linear.weight.dtype + + if generator is None: + generator = req.sampling_params.generator + if generator is None and req.sampling_params.seed is not None: + generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + + ( + prompt_embeds, + prompt_embeds_mask, + prompt_embeds_2, + prompt_embeds_mask_2, + negative_prompt_embeds, + negative_prompt_embeds_mask, + negative_prompt_embeds_2, + negative_prompt_embeds_mask_2, + ) = self.encode_prompt( + prompt=prompt, + device=device, + dtype=dtype, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_cfg, + ) + + batch_size = prompt_embeds.shape[0] + + # Get image embeddings from SigLIP (no explicit mask — the transformer + # detects I2V via non-zero image_embeds, matching diffusers behaviour) + image_embeds = self._get_image_embeds(image, device).to(dtype=dtype) + + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + num_frames=num_frames_val, + dtype=dtype, + device=device, + generator=generator, + latents=req.sampling_params.latents, + ) + + cond_latents, mask = self.prepare_cond_latents_and_mask(latents, image, height, width, dtype, device) + + # Timesteps — the scheduler handles flow_shift via its config (`shift` param). + sigmas = np.linspace(1.0, 0.0, num_steps + 1)[:-1] + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + for i, t in enumerate(timesteps): + self._current_timestep = t + + latent_model_input = torch.cat([latents, cond_latents, mask], dim=1) + timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype) + + timestep_r = None + if self.use_meanflow: + if i == len(timesteps) - 1: + timestep_r = torch.tensor([0.0], device=device) + else: + timestep_r = timesteps[i + 1] + timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "timestep_r": timestep_r, + "encoder_hidden_states": prompt_embeds, + "encoder_attention_mask": prompt_embeds_mask, + "encoder_hidden_states_2": prompt_embeds_2, + "encoder_attention_mask_2": prompt_embeds_mask_2, + "image_embeds": image_embeds, + "return_dict": False, + } + + negative_kwargs = None + if do_cfg and negative_prompt_embeds is not None: + # For I2V CFG, negative still uses image embeds (only text is unconditional) + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "timestep_r": timestep_r, + "encoder_hidden_states": negative_prompt_embeds, + "encoder_attention_mask": negative_prompt_embeds_mask, + "encoder_hidden_states_2": negative_prompt_embeds_2, + "encoder_attention_mask_2": negative_prompt_embeds_mask_2, + "image_embeds": image_embeds, + "return_dict": False, + } + + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_cfg and negative_kwargs is not None, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=req.sampling_params.cfg_normalize, + ) + + latents = self.scheduler_step_maybe_with_cfg( + noise_pred, + t, + latents, + do_true_cfg=do_cfg and negative_kwargs is not None, + ) + + self._current_timestep = None + + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + if output_type == "latent": + output = latents + else: + latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor + output = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput(output=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 77e7885c0..d4a738b81 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -142,6 +142,16 @@ _DIFFUSION_MODELS = { "pipeline_dreamid_omni", "DreamIDOmniPipeline", ), + "HunyuanVideo15Pipeline": ( + "hunyuan_video", + "pipeline_hunyuan_video_1_5", + "HunyuanVideo15Pipeline", + ), + "HunyuanVideo15ImageToVideoPipeline": ( + "hunyuan_video", + "pipeline_hunyuan_video_1_5_i2v", + "HunyuanVideo15I2VPipeline", + ), } @@ -316,6 +326,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "HeliosPipeline": "get_helios_post_process_func", "HeliosPyramidPipeline": "get_helios_post_process_func", "Flux2Pipeline": "get_flux2_post_process_func", + "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func", + "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { @@ -332,6 +344,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "OmniGen2Pipeline": "get_omnigen2_pre_process_func", "HeliosPipeline": "get_helios_pre_process_func", "HeliosPyramidPipeline": "get_helios_pre_process_func", + "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func", }