Remove mm_prefix_lm patch because vllm==0.18.0 already support (#2047)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.18.0"
+          VLLM_VERSION: "v0.17.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: true

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.18.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -11,10 +11,29 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-# Use bash -c so that $(python3 -c ...) is expanded inside the container.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+RUN uv pip uninstall --system -y vllm || true
 
+# Install vLLM from precompiled wheel at the selected commit.
+# Must use direct URL because the wheel has a PEP 440 local version identifier
+# (e.g. +g0a0a1a198) which pip/uv refuse to install from a PEP 503 package index.
+ENV VLLM_PRECOMPILED_WHEEL_COMMIT=89138b21cc246ae944c741d5c399c148e2b770ab
+RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
+    html=urllib.request.urlopen('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/vllm/').read().decode(); \
+    m=re.search(r'>(\S+x86_64\.whl)<',html); \
+    print('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/'+m.group(1).replace('+','%2B'))") && \
+    echo "Installing vLLM from: ${VLLM_WHEEL_URL}" && \
+    uv pip install --system --force-reinstall "${VLLM_WHEEL_URL}"
+
+RUN uv pip install --system ".[dev]"
+
+RUN uv pip install --system --upgrade \
+        "flashinfer-cubin==0.6.6" \
+        "nvidia-cublas-cu12==12.9.1.4" \
+        "numpy==2.2.6"
+
+RUN uv pip install --system --upgrade \
+    "flashinfer-jit-cache==0.6.6" \
+    --index-url https://flashinfer.ai/whl/cu129
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 
 ENTRYPOINT []

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
+ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docker/Dockerfile.xpu

@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.18.0
+ARG VLLM_VERSION=v0.17.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/getting_started/installation/gpu.md

@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -34,13 +34,13 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.18.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 ```
 
-The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.18.0
+git checkout v0.17.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.18.0 \
+    vllm/vllm-omni:v0.16.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
-The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The release 0.17.0 of vLLM requires ROCm 7.0 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.18.0
+git checkout v0.17.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.18.0 \
+  vllm/vllm-omni-rocm:v0.16.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.18.0
+  vllm/vllm-omni-rocm:v0.16.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.18.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/models/supported_models.md

@@ -52,8 +52,6 @@ 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` |
 
 

docs/user_guide/diffusion/quantization/fp8.md

@@ -63,7 +63,6 @@ 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
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,16 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
 
 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:
-    # Wan I2V-A14B (MoE)
+    # 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"
 
@@ -25,11 +24,6 @@ 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
@@ -49,11 +43,11 @@ from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
+    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
+        help="Diffusers Wan2.2 I2V model ID or local path.",
     )
     parser.add_argument(
         "--model-class-name",

examples/offline_inference/text_to_video/text_to_video.md

@@ -1,19 +1,10 @@
 # Text-To-Video
 
-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 |
+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.
 
 ## 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." \
@@ -45,100 +36,22 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
-### HunyuanVideo-1.5 (480p)
+Key arguments:
 
-```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. 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.
+- `--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).
 - `--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.
 
-### 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`.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/text_to_video/text_to_video.py

@@ -15,92 +15,56 @@ 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 from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
-    )
+    parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        help="Diffusers model ID or local path. "
-        "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
-        "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
+        help="Diffusers Wan2.2 model ID or local path.",
     )
     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 (Wan2.2 only).")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    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("--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(
         "--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=None, help="Sampling steps. Default: model-specific."
-    )
+    parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
     parser.add_argument(
         "--boundary-ratio",
         type=float,
-        default=None,
-        help="(Wan2.2) Boundary split ratio for low/high DiT. Default 0.875.",
+        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).",
     )
     parser.add_argument(
-        "--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).",
+        "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument(
         "--cache-backend",
         type=str,
         default=None,
         choices=["cache_dit"],
-        help="Cache backend for acceleration (Wan2.2). Default: None.",
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
+            "Default: None (no cache acceleration)."
+        ),
     )
     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=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("--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(
         "--vae-use-slicing",
         action="store_true",
@@ -173,41 +137,32 @@ 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()
-
-    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)
+    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).
     cache_config = None
     if args.cache_backend == "cache_dit":
         cache_config = {
-            "Fn_compute_blocks": 1,
-            "Bn_compute_blocks": 0,
-            "max_warmup_steps": 4,
+            # 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)
             "max_cached_steps": 20,
-            "residual_diff_threshold": 0.24,
-            "max_continuous_cached_steps": 3,
-            "enable_taylorseer": False,
-            "taylorseer_order": 1,
-            "scm_steps_mask_policy": None,
-            "scm_steps_policy": "dynamic",
+            "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"
         }
-
     # Configure parallel settings
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
@@ -221,11 +176,14 @@ def main():
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
-    omni_kwargs = dict(
+    omni = Omni(
         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,
@@ -233,18 +191,6 @@ 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...")
@@ -264,25 +210,22 @@ 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_dict,
-        OmniDiffusionSamplingParams(**sampling_kwargs),
+        {
+            "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,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start

examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh

@@ -1,68 +0,0 @@
-#!/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}"

examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh

@@ -1,31 +0,0 @@
-#!/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

examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh

@@ -1,60 +0,0 @@
-#!/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}"

examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh

@@ -1,31 +0,0 @@
-#!/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

tests/e2e/online_serving/test_hunyuan_video_15_expansion.py

@@ -1,98 +0,0 @@
-# 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)

vllm_omni/diffusion/models/hunyuan_video/__init__.py

@@ -1,25 +0,0 @@
-# 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",
-]

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -1,791 +0,0 @@
-# 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

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -1,513 +0,0 @@
-# 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)

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py

@@ -1,594 +0,0 @@
-# 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)

vllm_omni/diffusion/registry.py

@@ -142,16 +142,6 @@ _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",
-    ),
 }
 
 
@@ -326,8 +316,6 @@ _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 = {
@@ -344,7 +332,6 @@ _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",
 }
 
 

vllm_omni/patch.py

@@ -1,8 +1,6 @@
 import sys
-from functools import cached_property as _cached_property
 
 from aenum import extend_enum
-from vllm.config import ModelConfig as _ModelConfig
 from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
 from vllm.model_executor.layers.rotary_embedding import (
     MRotaryEmbedding as _OriginalMRotaryEmbedding,
@@ -45,20 +43,6 @@ except ImportError:
     # GlmImageTextConfig not available, skip patching
     pass
 
-
-# Patch ModelConfig.is_mm_prefix_lm to include Bagel (bidirectional attention
-# for multimodal prefix positions, same as Gemma3/Molmo2/PaliGemma).
-_orig_is_mm_prefix_lm = _ModelConfig.__dict__["is_mm_prefix_lm"].func
-
-
-@_cached_property
-def _patched_is_mm_prefix_lm(self) -> bool:
-    return _orig_is_mm_prefix_lm(self) or getattr(self.hf_config, "model_type", None) == "bagel"
-
-
-_patched_is_mm_prefix_lm.__set_name__(_ModelConfig, "is_mm_prefix_lm")
-_ModelConfig.is_mm_prefix_lm = _patched_is_mm_prefix_lm
-
 # Extend RequestStatus enum with omni-specific statuses
 if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
     # The value - 1 is intentionally chosen to ensure it is treated