[Model] Add HunyuanVideo-1.5 T2V and I2V support (#1516)
Signed-off-by: lishunyang <lishunyang12@163.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
15 files changed
+2443
-66
No files matched your search
@@ -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` |
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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`.
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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': <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)
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user