Merge remote-tracking branch 'origin/main' into dev/vllm-align

.buildkite/test-merge.yml

@@ -119,23 +119,23 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Tensor Parallelism Test"
-    timeout_in_minutes: 20
-    depends_on: upload-merge-pipeline
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+  # - label: "Diffusion Tensor Parallelism Test"
+  #   timeout_in_minutes: 20
+  #   depends_on: upload-merge-pipeline
+  #   commands:
+  #     - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
+  #   agents:
+  #     queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+  #   plugins:
+  #     - docker#v5.2.0:
+  #         image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #         always-pull: true
+  #         propagate-environment: true
+  #         shm-size: "8gb"
+  #         environment:
+  #           - "HF_HOME=/fsx/hf_cache"
+  #         volumes:
+  #           - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion GPU Worker Test"
     timeout_in_minutes: 20

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -6,6 +6,7 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
 Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
 """
 
+import asyncio
 import json
 from typing import Any
 
@@ -18,8 +19,8 @@ from .format import (
     audio_to_base64,
     base64_to_audio,
     base64_to_image_tensor,
-    base64_to_video,
     bytes_to_audio,
+    bytes_to_video,
     image_tensor_to_base64,
     image_tensor_to_png_bytes,
     video_to_base64,
@@ -31,10 +32,43 @@ from .types import AudioFormat
 logger = get_logger(__name__)
 
 
+async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            try:
+                return await response.json()
+            except aiohttp.ContentTypeError as e:
+                raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
+async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            return await response.read()
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
 class VLLMOmniClient:
-    def __init__(self, base_url: str, timeout: float | None = None):
+    def __init__(
+        self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
+    ):
         self.base_url = base_url
         self.timeout = aiohttp.ClientTimeout(total=timeout)
+        self.poll_interval = poll_interval
+        self.max_poll_duration = max_poll_duration
 
     async def generate_image(
         self,
@@ -262,33 +296,41 @@ class VLLMOmniClient:
                 content_type="image/png",
             )
 
-        url = self.base_url + "/videos"
         async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            # Start the video generation job
+            url = f"{self.base_url}/videos"
+            data = await url_json(session, url, "post", data=form)
+            if (job_id := data.get("id", None)) is None:
+                raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
+            if (job_status := data.get("status", None)) is None:
+                raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+
+            # Poll for video generation job completion
+            deadline = asyncio.get_running_loop().time() + self.max_poll_duration
+            url = f"{self.base_url}/videos/{job_id}"
+            while job_status not in {"completed", "failed"}:
+                await asyncio.sleep(self.poll_interval)
+
+                data = await url_json(session, url)
+                if (job_status := data.get("status", None)) is None:
+                    raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+                if asyncio.get_running_loop().time() >= deadline:
+                    raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
+
+            if job_status == "failed":
+                raise RuntimeError(f"Video job failed: {data}")
+
+            # Retrieve completed content
+            video_bytes = await url_bytes(session, f"{url}/content")
+
+            # Decode video and make a best effort at cleaning up server resources
             try:
-                async with session.post(url, data=form) as response:
-                    if not response.ok:
-                        error_text = await response.text()
-                        raise (ValueError if response.status < 500 else RuntimeError)(
-                            f"vLLM-Omni API returned status {response.status}: {error_text}"
-                        )
-
-                    try:
-                        data = await response.json()
-                    except aiohttp.ContentTypeError as e:
-                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
-            except aiohttp.ClientError as e:
-                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
-
-        if "data" not in data:
-            raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
-        if not data["data"]:
-            raise RuntimeError("API returned empty data array")
-        try:
-            base64_str = data["data"][0]["b64_json"]
-        except (KeyError, IndexError):
-            raise RuntimeError("API response missing 'b64_json' field in first data item")
-
-        return base64_to_video(base64_str)
+                return bytes_to_video(video_bytes)
+            finally:
+                try:
+                    await url_json(session, url, "delete")
+                except Exception as exc:
+                    logger.warning("Failed to clean up video job %s: %s", job_id, exc)
 
     async def generate_understanding_chat_completion(
         self,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py

@@ -143,15 +143,7 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
     return f"data:{mime_type};base64,{base64_str}"
 
 
-def base64_to_video(base64_str: str) -> VideoInput:
-    if base64_str.startswith("data:video"):
-        _, base64_str = base64_str.split(",", 1)
-
-    try:
-        video_bytes = base64.b64decode(base64_str)
-    except Exception as e:
-        raise ValueError(f"Invalid base64 string: {e}")
-
+def bytes_to_video(video_bytes: bytes) -> VideoInput:
     video_buffer = BytesIO(video_bytes)
 
     try:
@@ -208,6 +200,18 @@ def base64_to_video(base64_str: str) -> VideoInput:
     return InputImpl.VideoFromComponents(components)
 
 
+def base64_to_video(base64_str: str) -> VideoInput:
+    if base64_str.startswith("data:video"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        video_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_video(video_bytes)
+
+
 def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
     waveform = audio["waveform"][0]  # Shape: (C, T)
     sample_rate = audio["sample_rate"]

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -70,13 +70,17 @@ def add_sampling_parameters_to_stage(
 
     stages = pipeline_spec["stages"]
     if isinstance(sampling_param_list, dict):
+        sampling_param_list = sampling_param_list.__class__(sampling_param_list)
         sampling_param_list.update(params_to_add)
     elif sampling_param_list is None:
         sampling_param_list = params_to_add.copy()
     else:
         for i, stage in enumerate(stages):
             if stage == stage_type:
-                sampling_param_list[i].update(params_to_add)
+                stage_param = sampling_param_list[i]
+                stage_param = stage_param.__class__(stage_param)
+                stage_param.update(params_to_add)
+                sampling_param_list[i] = stage_param
 
     return sampling_param_list
 

docs/.nav.yml

@@ -28,6 +28,7 @@ nav:
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
+      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md

docs/models/supported_models.md

@@ -51,6 +51,7 @@ th {
 |`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
+|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
 
 
 ## List of Supported Models for NPU

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -39,6 +39,7 @@ The following table shows which models are currently supported by parallelism me
 | **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |        N/A        |
 | **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
 | **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        ✅        |
+| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |        N/A        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.

docs/user_guide/diffusion_acceleration.md

@@ -77,6 +77,7 @@ The following table shows which models are currently supported by each accelerat
 |-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
 
 ### Quantization
 

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -0,0 +1,58 @@
+# X-To-Video-Audio
+
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
+
+## Local CLI Usage
+### Download the Model locally
+Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
+
+```bash
+python download_dreamid_omni.py --output-dir ./dreamid_omni
+```
+After download, the model directory will look like this:
+
+```
+dreamid_omni/
+├── DreamID-Omni/
+│   ├── dreamid_omni.safetensors
+├── MMAudio/
+│   ├── ext_weights/
+│   │   ├── best_netG.pt
+│   │   ├── v1-16.pth
+├── Wan2.2-TI2V-5B/
+│   ├── google/*
+│   ├── models_t5_umt5-xxl-enc-bf16.pth
+│   ├── Wan2.2_VAE.pth
+│
+├── model_index.json # create by download_dreamid_omni.py
+```
+
+### Run the Inference
+```
+python x_to_video_audio.py \
+  --model /xx/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output dreamid_omni.mp4
+```
+In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
+The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+
+Key arguments:
+- `--prompt`: text description (string).
+- `--model`: path to the model local directory.
+- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--image-path`: path to the input image list.
+- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
+- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--num-inference-steps`: number of denoising steps (defaults 45).
+- `--video-negative-prompt`: negative prompt for video generation.
+- `--audio-negative-prompt`: negative prompt for audio generation.
+- `--enable-cpu-offload`: enable CPU offload (defaults False).

docs/user_guide/examples/online_serving/image_to_video.md

@@ -29,6 +29,35 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+Local file storage behavior can be controlled via the following environment variables:
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -38,7 +67,7 @@ The script allows overriding:
 bash run_curl_image_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -X POST http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=A bear playing with yarn, smooth motion" \
   -F "negative_prompt=low quality, blurry, static" \
@@ -52,7 +81,23 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "guidance_scale_2=1.0" \
   -F "boundary_ratio=0.875" \
   -F "flow_shift=12.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
 ```
 
 ## Request Format
@@ -66,6 +111,18 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "input_reference=@/path/to/qwen-bear.png"
 ```
 
+### Alternative JSON-Safe Reference Input
+
+Use `image_reference` when you want to pass a URL or JSON-safe image reference
+instead of uploading a file. Do not send `input_reference` and
+`image_reference` together.
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
+```
+
 ### Generation with Parameters
 
 ```bash
@@ -85,6 +142,63 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "seed=42"
 ```
 
+## Create Response Format
+
+`POST /v1/videos` returns a job record.
+
+```json
+{
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+  "prompt": "A bear playing with yarn, smooth motion",
+  "created_at": 1234567890
+}
+```
+
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
+
+```bash
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+```
+
 ## Example materials
 
 ??? abstract "run_curl_image_to_video.sh"

docs/user_guide/examples/online_serving/text_to_video.md

@@ -29,6 +29,34 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -38,7 +66,7 @@ The script allows overriding:
 bash run_curl_text_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -s http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   -F "width=832" \
@@ -50,7 +78,24 @@ curl -s http://localhost:8091/v1/videos \
   -F "guidance_scale=4.0" \
   -F "guidance_scale_2=4.0" \
   -F "boundary_ratio=0.875" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+  -F "flow_shift=5.0" \
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
 ```
 
 ## Request Format
@@ -85,8 +130,9 @@ curl -X POST http://localhost:8091/v1/videos \
 | Parameter             | Type   | Default | Description                                      |
 | --------------------- | ------ | ------- | ------------------------------------------------ |
 | `prompt`              | str    | -       | Text description of the desired video            |
+| `seconds`             | str    | None    | Clip duration in seconds                         |
+| `size`                | str    | None    | Output size in `WIDTHxHEIGHT` format             |
 | `negative_prompt`     | str    | None    | Negative prompt                                  |
-| `n`                   | int    | 1       | Number of videos to generate                     |
 | `width`               | int    | None    | Video width in pixels                            |
 | `height`              | int    | None    | Video height in pixels                           |
 | `num_frames`          | int    | None    | Number of frames to generate                     |
@@ -98,24 +144,62 @@ curl -X POST http://localhost:8091/v1/videos \
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
 | `lora`                | object | None    | LoRA configuration                               |
-| `extra_body`          | object | None    | Model-specific extra parameters                  |
 
-## Response Format
+## Create Response Format
+
+`POST /v1/videos` returns a job record.
 
 ```json
 {
-  "created": 1234567890,
-  "data": [
-    { "b64_json": "<base64-mp4>" }
-  ]
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+  "prompt": "A cinematic view of a futuristic city at sunset",
+  "created_at": 1234567890
 }
 ```
 
-## Extract Video
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
 
 ```bash
-# Extract base64 from response and decode to video
-cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
 ```
 
 ## Example materials

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -0,0 +1,103 @@
+import argparse
+import fcntl
+import os
+import site
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
+DEPENDENCY_BRANCH = "omni"
+CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
+LOCK_FILE = CACHE_DIR / ".install.lock"
+DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
+
+
+def download_dependency():
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+    with open(LOCK_FILE, "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        if not DEPENDENCY_DIR.exists():
+            print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
+            subprocess.run(
+                ["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
+                check=True,
+            )
+            print("Download finished.")
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+    # write .pth to site-packages
+    site_packages = Path(site.getsitepackages()[0])
+    pth_file = site_packages / "vllm_omni_dependency.pth"
+    pth_file.write_text(str(DEPENDENCY_DIR))
+    print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
+
+
+def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
+    """Download files from HF repo and log time + destination."""
+    if os.path.exists(local_dir):
+        print(f"Directory {local_dir} already exists. Skipping download.")
+        return
+    print(f"Starting download from {repo_id} into {local_dir}")
+    start_time = time.time()
+
+    snapshot_download(
+        repo_id=repo_id,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+        allow_patterns=allow_patterns,
+    )
+
+    elapsed = time.time() - start_time
+    print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
+
+
+def main(output_dir: str):
+    # Wan2.2
+    wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
+    timed_download(
+        repo_id="Wan-AI/Wan2.2-TI2V-5B",
+        local_dir=wan_dir,
+        allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
+    )
+
+    # MMAudio
+    mm_audio_dir = os.path.join(output_dir, "MMAudio")
+    timed_download(
+        repo_id="hkchengrex/MMAudio",
+        local_dir=mm_audio_dir,
+        allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
+    )
+
+    dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
+
+    timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
+
+    # Now we construct the config file
+    import json
+
+    data = {
+        "_class_name": "DreamIDOmniPipeline",
+        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
+    }
+
+    with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+    print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
+
+    # now we download the dependency code
+    download_dependency()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download models from Hugging Face")
+    parser.add_argument(
+        "--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
+    )
+    args = parser.parse_args()
+    main(args.output_dir)

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -0,0 +1,58 @@
+# X-To-Video-Audio
+
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
+
+## Local CLI Usage
+### Download the Model locally
+Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
+
+```bash
+python download_dreamid_omni.py --output-dir ./dreamid_omni
+```
+After download, the model directory will look like this:
+
+```
+dreamid_omni/
+├── DreamID-Omni/
+│   ├── dreamid_omni.safetensors
+├── MMAudio/
+│   ├── ext_weights/
+│   │   ├── best_netG.pt
+│   │   ├── v1-16.pth
+├── Wan2.2-TI2V-5B/
+│   ├── google/*
+│   ├── models_t5_umt5-xxl-enc-bf16.pth
+│   ├── Wan2.2_VAE.pth
+│
+├── model_index.json # create by download_dreamid_omni.py
+```
+
+### Run the Inference
+```
+python x_to_video_audio.py \
+  --model /xx/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output dreamid_omni.mp4
+```
+In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
+The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+
+Key arguments:
+- `--prompt`: text description (string).
+- `--model`: path to the model local directory.
+- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--image-path`: path to the input image list.
+- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
+- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--num-inference-steps`: number of denoising steps (defaults 45).
+- `--video-negative-prompt`: negative prompt for video generation.
+- `--audio-negative-prompt`: negative prompt for audio generation.
+- `--enable-cpu-offload`: enable CPU offload (defaults False).

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import re
+import time
+
+import librosa
+from PIL import Image
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
+AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
+    parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
+    parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
+    parser.add_argument("--prompt", default=None, help="Text prompt.")
+
+    parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
+    parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
+    parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
+
+    parser.add_argument("--height", type=int, default=704, help="Video height.")
+    parser.add_argument("--width", type=int, default=1280, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
+    parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
+    parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
+    parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2],
+        help="Number of GPUs used for classifier free guidance parallel size.",
+    )
+    parser.add_argument(
+        "--video-negative-prompt",
+        default="jitter, bad hands, blur, distortion",
+        help="Negative prompt for video.",
+    )
+    parser.add_argument(
+        "--audio-negative-prompt",
+        default="robotic, muffled, echo, distorted",
+        help="Negative prompt for audio.",
+    )
+    parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        default=False,
+        help="Enable CPU offloading for diffusion models.",
+    )
+    return parser.parse_args()
+
+
+def load_image_and_audio(image_paths, audio_paths):
+    image = []
+    audio = []
+
+    for path in image_paths:
+        with Image.open(path) as img:
+            img = img.convert("RGB")
+            image.append(img)
+
+    for path in audio_paths:
+        audio_array, sr = librosa.load(path, sr=16000)
+        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
+        audio.append(audio_array)
+    return image, audio
+
+
+def main() -> None:
+    args = parse_args()
+    if args.prompt is None and args.prompt_file is None:
+        raise ValueError("Either --prompt or --prompt-file must be provided.")
+
+    text_prompt = args.prompt
+    if args.prompt_file:
+        import json
+
+        with open(args.prompt_file) as f:
+            text_prompt = json.load(f)
+            text_prompt = re.sub(
+                r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(
+                r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
+            text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
+
+    image, audio = load_image_and_audio(args.image_path, args.audio_path)
+
+    prompt = {
+        "prompt": text_prompt,
+        "video_negative_prompt": args.video_negative_prompt,
+        "audio_negative_prompt": args.audio_negative_prompt,
+        "multi_modal_data": {"image": image, "audio": audio},
+    }
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "solver_name": args.solver_name,
+            "shift": args.shift,
+        },
+    )
+
+    parallel_config = DiffusionParallelConfig(
+        cfg_parallel_size=args.cfg_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        parallel_config=parallel_config,
+        model_type=args.model_type,
+        enable_cpu_offload=args.enable_cpu_offload,
+    )
+    start = time.perf_counter()
+    outputs = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
+
+    if not outputs:
+        raise RuntimeError("No output returned from DreamID-Omni.")
+    output = outputs[0].request_output
+    generated_video = output[0].images[0][0]
+    generated_audio = output[0].images[0][1]
+    try:
+        from dreamid_omni.utils.io_utils import save_video
+    except Exception as e:
+        raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}")
+    output_path = args.output
+    save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
+    print(f"Saved generated video to {output_path}")
+    print(f"Total time: {elapsed:.2f}s")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/image_to_video/README.md

@@ -26,6 +26,35 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+Local file storage behavior can be controlled via the following environment variables:
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -35,7 +64,7 @@ The script allows overriding:
 bash run_curl_image_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -X POST http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=A bear playing with yarn, smooth motion" \
   -F "negative_prompt=low quality, blurry, static" \
@@ -49,7 +78,23 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "guidance_scale_2=1.0" \
   -F "boundary_ratio=0.875" \
   -F "flow_shift=12.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
 ```
 
 ## Request Format
@@ -63,6 +108,18 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "input_reference=@/path/to/qwen-bear.png"
 ```
 
+### Alternative JSON-Safe Reference Input
+
+Use `image_reference` when you want to pass a URL or JSON-safe image reference
+instead of uploading a file. Do not send `input_reference` and
+`image_reference` together.
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
+```
+
 ### Generation with Parameters
 
 ```bash
@@ -81,3 +138,60 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "flow_shift=12.0" \
   -F "seed=42"
 ```
+
+## Create Response Format
+
+`POST /v1/videos` returns a job record, not inline base64 video data.
+
+```json
+{
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+  "prompt": "A bear playing with yarn, smooth motion",
+  "created_at": 1234567890
+}
+```
+
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
+
+```bash
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+```

examples/online_serving/image_to_video/run_curl_image_to_video.sh

@@ -1,33 +1,75 @@
 #!/bin/bash
-# Wan2.2 image-to-video curl example (OpenAI-style multipart)
+# Wan2.2 image-to-video curl example using the async video job API.
+
+set -euo pipefail
 
 INPUT_IMAGE="${INPUT_IMAGE:-../../offline_inference/image_to_video/qwen-bear.png}"
+BASE_URL="${BASE_URL:-http://localhost:8099}"
 OUTPUT_PATH="${OUTPUT_PATH:-wan22_i2v_output.mp4}"
 NEGATIVE_PROMPT="${NEGATIVE_PROMPT:-}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
 
 if [ ! -f "$INPUT_IMAGE" ]; then
     echo "Input image not found: $INPUT_IMAGE"
     exit 1
 fi
 
-NEGATIVE_PROMPT_FLAG=""
-if [ -n "$NEGATIVE_PROMPT" ]; then
-    NEGATIVE_PROMPT_FLAG="-F negative_prompt=${NEGATIVE_PROMPT}"
+create_cmd=(
+  curl -sS -X POST "${BASE_URL}/v1/videos"
+  -H "Accept: application/json"
+  -F "prompt=A bear playing with yarn, smooth motion"
+  -F "input_reference=@${INPUT_IMAGE}"
+  -F "seconds=2"
+  -F "size=832x480"
+  -F "fps=16"
+  -F "num_inference_steps=40"
+  -F "guidance_scale=1.0"
+  -F "guidance_scale_2=1.0"
+  -F "boundary_ratio=0.875"
+  -F "flow_shift=12.0"
+  -F "seed=42"
+)
+
+if [ -n "${NEGATIVE_PROMPT}" ]; then
+  create_cmd+=(-F "negative_prompt=${NEGATIVE_PROMPT}")
 fi
 
-curl -X POST http://localhost:8099/v1/videos \
-  -H "Accept: application/json" \
-  -F "prompt=A bear playing with yarn, smooth motion" \
-  $NEGATIVE_PROMPT_FLAG \
-  -F "input_reference=@${INPUT_IMAGE}" \
-  -F "seconds=2" \
-  -F "size=832x480" \
-  -F "fps=16" \
-  -F "num_inference_steps=40" \
-  -F "guidance_scale=1.0" \
-  -F "guidance_scale_2=1.0" \
-  -F "boundary_ratio=0.875" \
-  -F "flow_shift=12.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
+create_response="$("${create_cmd[@]}")"
+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/README.md

@@ -26,6 +26,35 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+Local file storage behavior can be controlled via the following environment variables:
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -35,7 +64,7 @@ The script allows overriding:
 bash run_curl_text_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -s http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   -F "width=832" \
@@ -47,7 +76,24 @@ curl -s http://localhost:8091/v1/videos \
   -F "guidance_scale=4.0" \
   -F "guidance_scale_2=4.0" \
   -F "boundary_ratio=0.875" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+  -F "flow_shift=5.0" \
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
 ```
 
 ## Request Format
@@ -82,8 +128,9 @@ curl -X POST http://localhost:8091/v1/videos \
 | Parameter             | Type   | Default | Description                                      |
 | --------------------- | ------ | ------- | ------------------------------------------------ |
 | `prompt`              | str    | -       | Text description of the desired video            |
+| `seconds`             | str    | None    | Clip duration in seconds                         |
+| `size`                | str    | None    | Output size in `WIDTHxHEIGHT` format             |
 | `negative_prompt`     | str    | None    | Negative prompt                                  |
-| `n`                   | int    | 1       | Number of videos to generate                     |
 | `width`               | int    | None    | Video width in pixels                            |
 | `height`              | int    | None    | Video height in pixels                           |
 | `num_frames`          | int    | None    | Number of frames to generate                     |
@@ -95,22 +142,60 @@ curl -X POST http://localhost:8091/v1/videos \
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
 | `lora`                | object | None    | LoRA configuration                               |
-| `extra_body`          | object | None    | Model-specific extra parameters                  |
 
-## Response Format
+## Create Response Format
+
+`POST /v1/videos` returns a job record, not inline base64 video data.
 
 ```json
 {
-  "created": 1234567890,
-  "data": [
-    { "b64_json": "<base64-mp4>" }
-  ]
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+  "prompt": "A cinematic view of a futuristic city at sunset",
+  "created_at": 1234567890
 }
 ```
 
-## Extract Video
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
 
 ```bash
-# Extract base64 from response and decode to video
-cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
 ```

examples/online_serving/text_to_video/run_curl_text_to_video.sh

@@ -1,20 +1,63 @@
 #!/bin/bash
-# Wan2.2 text-to-video curl example
+# Wan2.2 text-to-video curl example using the async video job API.
 
-OUTPUT_PATH="wan22_output.mp4"
+set -euo pipefail
 
-curl -X POST http://localhost:8098/v1/videos \
-  -H "Accept: application/json" \
-  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-  -F "seconds=2" \
-  -F "size=832x480" \
-  -F "negative_prompt=色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
-  -F "fps=16" \
-  -F "num_inference_steps=40" \
-  -F "guidance_scale=4.0" \
-  -F "guidance_scale_2=4.0" \
-  -F "boundary_ratio=0.875" \
-  -F "flow_shift=5.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
+BASE_URL="${BASE_URL:-http://localhost:8098}"
+OUTPUT_PATH="${OUTPUT_PATH:-wan22_output.mp4}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
 
+create_response=$(
+  curl -sS -X POST "${BASE_URL}/v1/videos" \
+    -H "Accept: application/json" \
+    -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+    -F "seconds=2" \
+    -F "size=832x480" \
+    -F "negative_prompt=色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+    -F "fps=16" \
+    -F "num_inference_steps=40" \
+    -F "guidance_scale=4.0" \
+    -F "guidance_scale_2=4.0" \
+    -F "boundary_ratio=0.875" \
+    -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}"

tests/comfyui/test_comfyui_integration.py

@@ -82,6 +82,15 @@ DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
     }
 )
 
+DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
+    {
+        "num_inference_steps": 30,
+        "guidance_scale": 6.0,
+        "true_cfg_scale": 1.5,
+    }
+)
+
+
 AR_LIST_SAMPLING_PARAMS = [
     AutoregressionSamplingParams(
         {
@@ -310,8 +319,8 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
             )
         elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE:
             assert len(received_sampling_params_list) == 1
-            expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy()
-            expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
+            expected = DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS.copy()
+            # expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
             _assert_sampling_param_values(
                 received_sampling_params_list[0],
                 {
@@ -690,7 +699,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
         pytest.param(
             SamplingCase(
                 kind=SamplingKind.VIDEO_DIFFUSION_SINGLE,
-                sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
+                sampling_params=DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS,
                 lora=LORA_PARAMS,
             ),
             id="single-diffusion-sampling-params",

tests/e2e/offline_inference/test_zimage_parallelism.py

@@ -17,12 +17,14 @@ import time
 from pathlib import Path
 
 import numpy as np
-import pytest
+
+# import pytest
 import torch
 from PIL import Image
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import DeviceMemoryMonitor, hardware_test
+# from tests.utils import DeviceMemoryMonitor, hardware_test
+from tests.utils import DeviceMemoryMonitor
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -155,69 +157,69 @@ def _run_zimage_generate(
         cleanup_dist_env_and_memory()
 
 
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_tensor_parallel_tp2(tmp_path: Path):
-    if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
-        pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
-    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
-        pytest.skip("Z-Image TP=2 requires >= 2 devices.")
+# @pytest.mark.advanced_model
+# @pytest.mark.diffusion
+# @pytest.mark.parallel
+# @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+# def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+#     if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
+#         pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
+#     if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+#         pytest.skip("Z-Image TP=2 requires >= 2 devices.")
 
-    enforce_eager = False
+#     enforce_eager = False
 
-    height = 512
-    width = 512
-    num_inference_steps = 2
-    seed = 42
+#     height = 512
+#     width = 512
+#     num_inference_steps = 2
+#     seed = 42
 
-    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
-        tp_size=1,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        seed=seed,
-        enforce_eager=enforce_eager,
-    )
-    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
-        tp_size=2,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        seed=seed,
-        enforce_eager=enforce_eager,
-    )
+#     tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
+#         tp_size=1,
+#         height=height,
+#         width=width,
+#         num_inference_steps=num_inference_steps,
+#         seed=seed,
+#         enforce_eager=enforce_eager,
+#     )
+#     tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
+#         tp_size=2,
+#         height=height,
+#         width=width,
+#         num_inference_steps=num_inference_steps,
+#         seed=seed,
+#         enforce_eager=enforce_eager,
+#     )
 
-    tp1_path = tmp_path / "zimage_tp1.png"
-    tp2_path = tmp_path / "zimage_tp2.png"
-    tp1_img.save(tp1_path)
-    tp2_img.save(tp2_path)
+#     tp1_path = tmp_path / "zimage_tp1.png"
+#     tp2_path = tmp_path / "zimage_tp2.png"
+#     tp1_img.save(tp1_path)
+#     tp2_img.save(tp2_path)
 
-    assert tp1_img.width == width and tp1_img.height == height
-    assert tp2_img.width == width and tp2_img.height == height
+#     assert tp1_img.width == width and tp1_img.height == height
+#     assert tp2_img.width == width and tp2_img.height == height
 
-    mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img)
-    mean_threshold = 3e-2
-    p99_threshold = 2.5e-1
-    print(
-        "Z-Image TP image diff stats (TP=1 vs TP=2): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
-        f"tp1_img={tp1_path}, tp2_img={tp2_path}"
-    )
-    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
-    )
+#     mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img)
+#     mean_threshold = 3e-2
+#     p99_threshold = 2.5e-1
+#     print(
+#         "Z-Image TP image diff stats (TP=1 vs TP=2): "
+#         f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+#         f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+#         f"tp1_img={tp1_path}, tp2_img={tp2_path}"
+#     )
+#     assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+#         f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+#         f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
+#     )
 
-    print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
-    assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
+#     print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
+#     assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
 
-    print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
-    assert tp2_peak_mem < tp1_peak_mem, (
-        f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
-    )
+#     print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+#     assert tp2_peak_mem < tp1_peak_mem, (
+#         f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
+#     )
 
 
 # @pytest.mark.integration

tests/e2e/online_serving/test_video_generation_api.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving tests for the async video generation API.
+
+These tests exercise the real `/v1/videos` job lifecycle against a running
+diffusion model, including create, retrieve, list, download, and delete flows.
+"""
+
+import os
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+import pytest
+import requests
+
+from tests.conftest import OmniServer
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+VIDEO_POLL_INTERVAL_S = 2.0
+VIDEO_TIMEOUT_S = 900.0
+
+
+@pytest.fixture(scope="module")
+def video_server():
+    """Start a real video generation server for async video API tests."""
+    with OmniServer(
+        MODEL,
+        [
+            "--num-gpus",
+            "1",
+            "--boundary-ratio",
+            "0.875",
+            "--flow-shift",
+            "5.0",
+            "--disable-log-stats",
+        ],
+    ) as server:
+        yield server
+
+
+def _video_api_url(server: OmniServer, suffix: str = "") -> str:
+    return f"http://{server.host}:{server.port}/v1/videos{suffix}"
+
+
+def _multipart_fields(payload: dict[str, Any]) -> list[tuple[str, tuple[None, str]]]:
+    return [(key, (None, str(value))) for key, value in payload.items() if value is not None]
+
+
+def _create_video_job(
+    server: OmniServer,
+    *,
+    prompt: str | None = None,
+    **overrides: Any,
+) -> requests.Response:
+    payload: dict[str, Any] = {
+        "prompt": prompt or f"video api e2e {uuid.uuid4().hex[:8]}",
+        "width": 640,
+        "height": 480,
+        "num_frames": 5,
+        "fps": 8,
+        "num_inference_steps": 2,
+        "guidance_scale": 1.0,
+        "guidance_scale_2": 1.0,
+        "boundary_ratio": 0.875,
+        "flow_shift": 5.0,
+        "seed": 42,
+    }
+    payload.update(overrides)
+    return requests.post(
+        _video_api_url(server),
+        files=_multipart_fields(payload),
+        timeout=VIDEO_TIMEOUT_S,
+    )
+
+
+def _retrieve_video_job(server: OmniServer, video_id: str) -> requests.Response:
+    return requests.get(_video_api_url(server, f"/{video_id}"), timeout=VIDEO_TIMEOUT_S)
+
+
+def _wait_for_video_status(server: OmniServer, video_id: str, expected_status: str) -> dict[str, Any]:
+    deadline = time.time() + VIDEO_TIMEOUT_S
+    last_payload: dict[str, Any] | None = None
+
+    while time.time() < deadline:
+        response = _retrieve_video_job(server, video_id)
+        assert response.status_code == 200, response.text
+        last_payload = response.json()
+        status = last_payload["status"]
+        if status == expected_status:
+            return last_payload
+        if status == "failed":
+            raise AssertionError(f"Video job {video_id} failed unexpectedly: {last_payload}")
+        time.sleep(VIDEO_POLL_INTERVAL_S)
+
+    raise AssertionError(
+        f"Timed out waiting for video job {video_id} to reach status={expected_status}. Last payload: {last_payload}"
+    )
+
+
+def _wait_for_video_missing(server: OmniServer, video_id: str) -> None:
+    deadline = time.time() + 30.0
+    last_status: int | None = None
+
+    while time.time() < deadline:
+        response = _retrieve_video_job(server, video_id)
+        last_status = response.status_code
+        if response.status_code == 404:
+            return
+        time.sleep(0.5)
+
+    raise AssertionError(f"Timed out waiting for video job {video_id} to disappear. Last status={last_status}")
+
+
+def _delete_video_job(server: OmniServer, video_id: str) -> requests.Response:
+    return requests.delete(_video_api_url(server, f"/{video_id}"), timeout=VIDEO_TIMEOUT_S)
+
+
+def _delete_video_job_with_retry(server: OmniServer, video_id: str) -> requests.Response:
+    deadline = time.time() + 30.0
+    last_response: requests.Response | None = None
+
+    while time.time() < deadline:
+        response = _delete_video_job(server, video_id)
+        last_response = response
+        if response.status_code != 409:
+            return response
+        time.sleep(1.0)
+
+    raise AssertionError(
+        f"Timed out waiting to delete video job {video_id}. "
+        f"Last response: {None if last_response is None else last_response.text}"
+    )
+
+
+def _best_effort_delete(server: OmniServer, video_id: str) -> None:
+    try:
+        response = _delete_video_job_with_retry(server, video_id)
+        if response.status_code not in (200, 404):
+            print(f"Cleanup delete for {video_id} returned {response.status_code}: {response.text}")
+    except Exception as exc:
+        print(f"Cleanup delete for {video_id} failed: {exc}")
+
+
+def _assert_mp4_payload(content: bytes) -> None:
+    assert len(content) > 32, f"Downloaded video payload is unexpectedly small: {len(content)} bytes"
+    assert content[4:8] == b"ftyp", "Downloaded payload does not look like an MP4 file."
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_create_list_retrieve_download_video(video_server: OmniServer, tmp_path: Path) -> None:
+    video_id: str | None = None
+    try:
+        create_response = _create_video_job(video_server, prompt="A small paper airplane gliding over a lake at dawn.")
+        assert create_response.status_code == 200, create_response.text
+
+        created = create_response.json()
+        video_id = created["id"]
+        assert created["object"] == "video"
+        assert created["status"] == "queued"
+        assert created["model"] == MODEL
+
+        completed = _wait_for_video_status(video_server, video_id, "completed")
+        assert completed["file_name"] is not None
+        assert completed["completed_at"] is not None
+        assert completed["progress"] == 100
+
+        retrieve_response = _retrieve_video_job(video_server, video_id)
+        assert retrieve_response.status_code == 200, retrieve_response.text
+        retrieved = retrieve_response.json()
+        assert retrieved["id"] == video_id
+        assert retrieved["status"] == "completed"
+
+        list_response = requests.get(_video_api_url(video_server), timeout=VIDEO_TIMEOUT_S)
+        assert list_response.status_code == 200, list_response.text
+        listed = list_response.json()
+        assert listed["object"] == "list"
+        assert any(item["id"] == video_id for item in listed["data"])
+
+        download_response = requests.get(
+            _video_api_url(video_server, f"/{video_id}/content"),
+            timeout=VIDEO_TIMEOUT_S,
+        )
+        assert download_response.status_code == 200, download_response.text
+        assert download_response.headers["content-type"].startswith("video/mp4")
+        assert completed["file_name"] in download_response.headers.get("content-disposition", "")
+        _assert_mp4_payload(download_response.content)
+
+        output_path = tmp_path / completed["file_name"]
+        output_path.write_bytes(download_response.content)
+        assert output_path.stat().st_size == len(download_response.content)
+    finally:
+        if video_id is not None:
+            _best_effort_delete(video_server, video_id)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_delete_video_while_queued_or_in_progress(video_server: OmniServer) -> None:
+    create_response = _create_video_job(
+        video_server,
+        prompt="A toy boat sailing through fog.",
+        num_frames=13,
+        num_inference_steps=4,
+    )
+    assert create_response.status_code == 200, create_response.text
+
+    created = create_response.json()
+    video_id = created["id"]
+    assert created["status"] == "queued"
+
+    # Delete immediately so the job is still queued or just entering execution.
+    delete_response = _delete_video_job_with_retry(video_server, video_id)
+    assert delete_response.status_code == 200, delete_response.text
+    deleted = delete_response.json()
+    assert deleted["id"] == video_id
+    assert deleted["deleted"] is True
+    assert deleted["object"] == "video.deleted"
+
+    _wait_for_video_missing(video_server, video_id)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_delete_video_after_processing(video_server: OmniServer) -> None:
+    create_response = _create_video_job(
+        video_server,
+        prompt="A lantern drifting above a quiet forest.",
+        seed=43,
+    )
+    assert create_response.status_code == 200, create_response.text
+
+    video_id = create_response.json()["id"]
+    completed = _wait_for_video_status(video_server, video_id, "completed")
+    assert completed["file_name"] is not None
+
+    delete_response = _delete_video_job_with_retry(video_server, video_id)
+    assert delete_response.status_code == 200, delete_response.text
+    deleted = delete_response.json()
+    assert deleted["id"] == video_id
+    assert deleted["deleted"] is True
+    assert deleted["object"] == "video.deleted"
+
+    _wait_for_video_missing(video_server, video_id)
+
+    download_response = requests.get(
+        _video_api_url(video_server, f"/{video_id}/content"),
+        timeout=VIDEO_TIMEOUT_S,
+    )
+    assert download_response.status_code == 404, download_response.text

tests/entrypoints/openai_api/test_video_server.py

@@ -4,19 +4,31 @@
 Unit tests for OpenAI-compatible video generation endpoints.
 """
 
+import asyncio
+import base64
 import io
+import json
+import os
+import threading
+import time
 from types import SimpleNamespace
 
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from PIL import Image
-from pydantic import ValidationError
 from pytest_mock import MockerFixture
 
+from vllm_omni.entrypoints.openai import api_server
 from vllm_omni.entrypoints.openai.api_server import router
-from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoResponseFormat
+from vllm_omni.entrypoints.openai.protocol.videos import (
+    VideoGenerationRequest,
+    VideoGenerationStatus,
+    VideoResponse,
+)
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
+from vllm_omni.entrypoints.openai.storage import LocalStorageManager
+from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -44,6 +56,38 @@ class FakeAsyncOmni:
         yield MockVideoResult(videos)
 
 
+class BlockingVideoHandler:
+    def __init__(self):
+        self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+        self.stage_configs = None
+        self.started = threading.Event()
+        self.cancelled = threading.Event()
+
+    def set_stage_configs_if_missing(self, stage_configs):
+        if self.stage_configs is None:
+            self.stage_configs = stage_configs
+
+    async def generate_videos(self, request, reference_id, *, reference_image=None):
+        self.started.set()
+        try:
+            await asyncio.Future()
+        except asyncio.CancelledError:
+            self.cancelled.set()
+            raise
+
+
+@pytest.fixture(autouse=True)
+def isolated_video_backends(tmp_path, monkeypatch):
+    """Use isolated in-memory metadata and local storage for each test."""
+    store: AsyncDictStore[VideoResponse] = AsyncDictStore()
+    tasks = TaskRegistry()
+    storage = LocalStorageManager(storage_path=str(tmp_path / "storage"))
+    monkeypatch.setattr(api_server, "VIDEO_STORE", store)
+    monkeypatch.setattr(api_server, "VIDEO_TASKS", tasks)
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", storage)
+    return store, tasks, storage
+
+
 @pytest.fixture
 def test_client():
     app = FastAPI()
@@ -52,7 +96,8 @@ def test_client():
         diffusion_engine=FakeAsyncOmni(),
         model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
     )
-    return TestClient(app)
+    with TestClient(app) as client:
+        yield client
 
 
 def _make_test_image_bytes(size=(64, 64)) -> bytes:
@@ -62,6 +107,34 @@ def _make_test_image_bytes(size=(64, 64)) -> bytes:
     return buf.getvalue()
 
 
+def _make_test_image_data_url(size=(64, 64)) -> str:
+    image_bytes = _make_test_image_bytes(size)
+    encoded = base64.b64encode(image_bytes).decode("utf-8")
+    return f"data:image/png;base64,{encoded}"
+
+
+def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0):
+    deadline = time.time() + timeout_s
+    last_payload = None
+    while time.time() < deadline:
+        response = client.get(f"/v1/videos/{video_id}")
+        assert response.status_code == 200
+        last_payload = response.json()
+        if last_payload["status"] == status:
+            return last_payload
+        time.sleep(0.02)
+    raise AssertionError(f"Timed out waiting for status={status}. Last payload: {last_payload}")
+
+
+def _wait_until(predicate, timeout_s: float = 2.0, interval_s: float = 0.02):
+    deadline = time.time() + timeout_s
+    while time.time() < deadline:
+        if predicate():
+            return
+        time.sleep(interval_s)
+    raise AssertionError("Timed out waiting for condition")
+
+
 def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
     fps_values = []
 
@@ -80,24 +153,22 @@ def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
             "size": "640x360",
             "seconds": "2",
             "fps": "12",
-            "n": "2",
         },
     )
 
     assert response.status_code == 200
-    data = response.json()
-    assert "data" in data and len(data["data"]) == 2
-    assert all(item["b64_json"] == "Zg==" for item in data["data"])
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
 
     engine = test_client.app.state.openai_serving_video._engine_client
     captured = engine.captured_sampling_params_list[0]
-    assert captured.num_outputs_per_prompt == 2
+    assert captured.num_outputs_per_prompt == 1
     assert captured.width == 640
     assert captured.height == 360
     assert captured.num_frames == 24
     assert captured.fps == 12
     assert captured.frame_rate == 12.0
-    assert fps_values == [12, 12]
+    assert fps_values == [12]
 
 
 def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
@@ -114,9 +185,8 @@ def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
     )
 
     assert response.status_code == 200
-    data = response.json()
-    assert "data" in data and len(data["data"]) == 1
-    assert data["data"][0]["b64_json"] == "Zg=="
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
 
     engine = test_client.app.state.openai_serving_video._engine_client
     prompt = engine.captured_prompt
@@ -145,6 +215,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
     )
 
     assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
 
     engine = test_client.app.state.openai_serving_video._engine_client
     prompt = engine.captured_prompt
@@ -153,6 +225,30 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
     assert input_image.size == (96, 64)
 
 
+def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="Zg==",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A fox running through snow.",
+            "image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    prompt = engine.captured_prompt
+    input_image = prompt["multi_modal_data"]["image"]
+    assert isinstance(input_image, Image.Image)
+    assert input_image.size == (40, 24)
+
+
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     fps_values = []
 
@@ -173,10 +269,13 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     )
 
     assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
     engine = test_client.app.state.openai_serving_video._engine_client
     captured = engine.captured_sampling_params_list[0]
     assert captured.num_frames == 72
     assert captured.fps == 24
+    assert captured.frame_rate == 24.0
     assert fps_values == [24]
 
 
@@ -194,16 +293,50 @@ def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
     )
 
     assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
     engine = test_client.app.state.openai_serving_video._engine_client
     captured = engine.captured_sampling_params_list[0]
     assert captured.width == 320
     assert captured.height == 240
 
 
+def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="Zg==",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "param pass",
+            "num_inference_steps": "30",
+            "guidance_scale": "6.5",
+            "guidance_scale_2": "8.0",
+            "true_cfg_scale": "4.0",
+            "boundary_ratio": "0.7",
+            "flow_shift": "0.25",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 30
+    assert captured.guidance_scale == 6.5
+    assert captured.guidance_scale_2 == 8.0
+    assert captured.true_cfg_scale == 4.0
+    assert captured.boundary_ratio == 0.7
+    assert captured.extra_args["flow_shift"] == 0.25
+
+
 def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture):
     audio_sample_rates = []
 
     def _fake_encode(video, fps, audio=None, audio_sample_rate=None):
+        del video, fps, audio
         audio_sample_rates.append(audio_sample_rate)
         return "Zg=="
 
@@ -233,38 +366,11 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
     )
 
     assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
     assert audio_sample_rates == [16000]
 
 
-def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
-    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
-    )
-    response = test_client.post(
-        "/v1/videos",
-        data={
-            "prompt": "param pass",
-            "num_inference_steps": "30",
-            "guidance_scale": "6.5",
-            "guidance_scale_2": "8.0",
-            "true_cfg_scale": "4.0",
-            "boundary_ratio": "0.7",
-            "flow_shift": "0.25",
-        },
-    )
-
-    assert response.status_code == 200
-    engine = test_client.app.state.openai_serving_video._engine_client
-    captured = engine.captured_sampling_params_list[0]
-    assert captured.num_inference_steps == 30
-    assert captured.guidance_scale == 6.5
-    assert captured.guidance_scale_2 == 8.0
-    assert captured.true_cfg_scale == 4.0
-    assert captured.boundary_ratio == 0.7
-    assert captured.extra_args["flow_shift"] == 0.25
-
-
 def test_missing_handler_returns_503():
     app = FastAPI()
     app.include_router(router)
@@ -287,29 +393,29 @@ def test_missing_prompt_returns_422(test_client):
     assert response.status_code == 422
 
 
-def test_invalid_size_format_raises_validation_error(test_client):
-    with pytest.raises(ValidationError):
-        test_client.post(
-            "/v1/videos",
-            data={"prompt": "bad size", "size": "invalid"},
-        )
-
-
-def test_invalid_size_parse_returns_500(test_client):
+def test_invalid_size_parse_returns_422(test_client):
     response = test_client.post(
         "/v1/videos",
         data={"prompt": "bad size", "size": "640x"},
     )
-    assert response.status_code == 500
-    assert "video generation failed" in response.json()["detail"].lower()
+    assert response.status_code == 422
+    body = response.json()
+    assert body["detail"][0]["loc"] == ["body", "size"]
+    assert body["detail"][0]["type"] == "string_pattern_mismatch"
+    assert body["detail"][0]["input"] == "640x"
 
 
-def test_invalid_response_format_raises_validation_error(test_client):
-    with pytest.raises(ValidationError):
-        test_client.post(
-            "/v1/videos",
-            data={"prompt": "bad format", "response_format": "url"},
-        )
+def test_rejects_input_reference_and_image_reference_together(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "bad refs",
+            "image_reference": '{"image_url": "https://example.com/cat.png"}',
+        },
+        files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
+    )
+    assert response.status_code == 400
+    assert "either input_reference or image_reference" in response.json()["detail"].lower()
 
 
 def test_invalid_seconds_returns_422(test_client):
@@ -320,20 +426,6 @@ def test_invalid_seconds_returns_422(test_client):
     assert response.status_code == 422
 
 
-def test_invalid_n_raises_validation_error(test_client):
-    with pytest.raises(ValidationError):
-        test_client.post(
-            "/v1/videos",
-            data={"prompt": "bad n", "n": "0"},
-        )
-
-    with pytest.raises(ValidationError):
-        test_client.post(
-            "/v1/videos",
-            data={"prompt": "bad n", "n": "5"},
-        )
-
-
 def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture):
     mocker.patch(
         "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
@@ -349,6 +441,8 @@ def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixtur
     )
 
     assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
     engine = test_client.app.state.openai_serving_video._engine_client
     captured_prompt = engine.captured_prompt
     captured_params = engine.captured_sampling_params_list[0]
@@ -364,19 +458,38 @@ def test_invalid_lora_returns_400(test_client):
             "lora": '{"name": "bad-lora"}',
         },
     )
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
+    assert failed["error"]["code"] == "HTTPException"
+    assert "lora object" in failed["error"]["message"].lower()
+
+
+def test_unsupported_image_reference_file_id_returns_400(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "unsupported ref",
+            "image_reference": '{"file_id": "file-123"}',
+        },
+    )
     assert response.status_code == 400
-    assert "lora" in response.json()["detail"].lower()
+    assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet."
+
+
+def test_invalid_uploaded_input_reference_returns_400(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "bad upload"},
+        files={"input_reference": ("input.png", b"not-an-image", "image/png")},
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image."
 
 
 def test_video_request_validation():
     req = VideoGenerationRequest(prompt="test")
     assert req.prompt == "test"
-    assert req.n == 1
-    assert req.response_format == VideoResponseFormat.B64_JSON
-
-    with pytest.raises(ValueError):
-        VideoGenerationRequest(prompt="test", response_format="url")
-
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", size="invalid")
 
@@ -384,4 +497,144 @@ def test_video_request_validation():
         VideoGenerationRequest(prompt="test", seconds="abc")
 
     with pytest.raises(ValueError):
-        VideoGenerationRequest(prompt="test", n=0)
+        VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
+
+
+def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="Zg==",
+    )
+    ids = []
+    for i in range(3):
+        create_resp = test_client.post("/v1/videos", data={"prompt": f"video-{i}"})
+        assert create_resp.status_code == 200
+        video_id = create_resp.json()["id"]
+        _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+        ids.append(video_id)
+
+    asyncio.run(api_server.VIDEO_STORE.update_fields(ids[0], {"created_at": 100}))
+    asyncio.run(api_server.VIDEO_STORE.update_fields(ids[1], {"created_at": 200}))
+    asyncio.run(api_server.VIDEO_STORE.update_fields(ids[2], {"created_at": 300}))
+
+    asc_resp = test_client.get("/v1/videos", params={"order": "asc"})
+    assert asc_resp.status_code == 200
+    asc_body = asc_resp.json()
+    asc_ids = [item["id"] for item in asc_body["data"]]
+    assert asc_ids == [ids[0], ids[1], ids[2]]
+    assert asc_body["object"] == "list"
+    assert asc_body["first_id"] == ids[0]
+    assert asc_body["last_id"] == ids[2]
+    assert asc_body["has_more"] is False
+
+    desc_resp = test_client.get("/v1/videos", params={"order": "desc", "limit": 2})
+    assert desc_resp.status_code == 200
+    desc_body = desc_resp.json()
+    desc_ids = [item["id"] for item in desc_body["data"]]
+    assert desc_ids == [ids[2], ids[1]]
+    assert desc_body["object"] == "list"
+    assert desc_body["first_id"] == ids[2]
+    assert desc_body["last_id"] == ids[1]
+    assert desc_body["has_more"] is True
+
+    after_resp = test_client.get("/v1/videos", params={"order": "asc", "after": ids[0]})
+    assert after_resp.status_code == 200
+    after_body = after_resp.json()
+    after_ids = [item["id"] for item in after_body["data"]]
+    assert after_ids == [ids[1], ids[2]]
+    assert after_body["object"] == "list"
+    assert after_body["first_id"] == ids[1]
+    assert after_body["last_id"] == ids[2]
+    assert after_body["has_more"] is False
+
+    zero_limit_resp = test_client.get("/v1/videos", params={"order": "asc", "limit": 0})
+    assert zero_limit_resp.status_code == 200
+    zero_limit_body = zero_limit_resp.json()
+    assert zero_limit_body["data"] == []
+    assert zero_limit_body["object"] == "list"
+    assert zero_limit_body["first_id"] is None
+    assert zero_limit_body["last_id"] is None
+    assert zero_limit_body["has_more"] is True
+
+    zero_limit_after_resp = test_client.get(
+        "/v1/videos",
+        params={"order": "asc", "after": ids[2], "limit": 0},
+    )
+    assert zero_limit_after_resp.status_code == 200
+    zero_limit_after_body = zero_limit_after_resp.json()
+    assert zero_limit_after_body["data"] == []
+    assert zero_limit_after_body["object"] == "list"
+    assert zero_limit_after_body["first_id"] is None
+    assert zero_limit_after_body["last_id"] is None
+    assert zero_limit_after_body["has_more"] is False
+
+
+def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="Zg==",
+    )
+    create_resp = test_client.post("/v1/videos", data={"prompt": "Delete this video"})
+    assert create_resp.status_code == 200
+    video_id = create_resp.json()["id"]
+
+    final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    file_name = final["file_name"]
+    assert file_name is not None
+    file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, file_name)
+    assert os.path.exists(file_path)
+
+    delete_resp = test_client.delete(f"/v1/videos/{video_id}")
+    assert delete_resp.status_code == 200
+    assert delete_resp.json()["id"] == video_id
+    assert delete_resp.json()["deleted"] is True
+    assert delete_resp.json()["object"] == "video.deleted"
+    assert not os.path.exists(file_path)
+
+
+def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client):
+    handler = BlockingVideoHandler()
+    test_client.app.state.openai_serving_video = handler
+
+    create_resp = test_client.post("/v1/videos", data={"prompt": "Cancel this video"})
+    assert create_resp.status_code == 200
+    video_id = create_resp.json()["id"]
+
+    assert handler.started.wait(timeout=2.0)
+
+    delete_resp = test_client.delete(f"/v1/videos/{video_id}")
+    assert delete_resp.status_code == 200
+    assert delete_resp.json()["id"] == video_id
+    assert delete_resp.json()["deleted"] is True
+    assert delete_resp.json()["object"] == "video.deleted"
+
+    assert handler.cancelled.wait(timeout=2.0)
+    _wait_until(lambda: asyncio.run(api_server.VIDEO_TASKS.get(video_id)) is None)
+    assert asyncio.run(api_server.VIDEO_STORE.get(video_id)) is None
+
+    retrieve_resp = test_client.get(f"/v1/videos/{video_id}")
+    assert retrieve_resp.status_code == 404
+
+
+def test_video_response_file_extension_is_robust():
+    response = VideoResponse(model="test-model", prompt="Make something beautiful")
+    assert response.file_extension == "mp4"
+
+    with_params = VideoResponse.model_construct(
+        model="test-model",
+        media_type="video/mp4; charset=binary",
+    )
+    assert with_params.file_extension == "mp4"
+
+    webm = VideoResponse.model_construct(
+        model="test-model",
+        media_type="video/webm",
+    )
+    assert webm.file_extension == "webm"
+
+    with pytest.raises(ValueError):
+        unknown = VideoResponse.model_construct(
+            model="test-model",
+            media_type="application/x-custom-video",
+        )
+        _ = unknown.file_extension

vllm_omni/benchmarks/metrics/metrics.py

@@ -228,8 +228,8 @@ def calculate_metrics(
                 good_completed += 1
 
     if completed == 0:
-        warnings.formatwarning = (
-            lambda msg, category, filename, lineno, line=None: f"{filename}:{lineno}: {category.__name__}: {msg}\n"
+        warnings.formatwarning = lambda msg, category, filename, lineno, line=None: (
+            f"{filename}:{lineno}: {category.__name__}: {msg}\n"
         )
         warnings.warn(
             "All requests failed. This is likely due to a misconfiguration on the benchmark arguments.",

vllm_omni/diffusion/data.py

@@ -322,6 +322,7 @@ class OmniDiffusionConfig:
 
     dtype: torch.dtype = torch.bfloat16
 
+    model_config: dict[str, Any] = field(default_factory=dict)
     tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
 
     # Attention

vllm_omni/diffusion/diffusion_engine.py

@@ -6,6 +6,7 @@ import time
 from collections.abc import Iterable
 from typing import Any
 
+import numpy as np
 import PIL.Image
 from vllm.logger import init_logger
 
@@ -30,6 +31,13 @@ def supports_image_input(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_image_input", False))
 
 
+def supports_audio_input(model_class_name: str) -> bool:
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is None:
+        return False
+    return bool(getattr(model_cls, "support_audio_input", False))
+
+
 def image_color_format(model_class_name: str) -> str:
     model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
     return getattr(model_cls, "color_format", "RGB")
@@ -369,7 +377,19 @@ class DiffusionEngine:
             dummy_image = PIL.Image.new(color_format, (width, height))
         else:
             dummy_image = None
-        prompt: OmniTextPrompt = {"prompt": "dummy run", "multi_modal_data": {"image": dummy_image}}
+
+        if supports_audio_input(self.od_config.model_class_name):
+            audio_sr = 16000
+            audio_duration_sec = 4
+            audio_array = np.random.randn(audio_sr * audio_duration_sec).astype(np.float32)
+            dummy_audio = audio_array[audio_sr * 1 : audio_sr * 3]
+        else:
+            dummy_audio = None
+
+        prompt: OmniTextPrompt = {
+            "prompt": "dummy run",
+            "multi_modal_data": {"image": dummy_image, "audio": dummy_audio},
+        }
         req = OmniDiffusionRequest(
             prompts=[prompt],
             sampling_params=OmniDiffusionSamplingParams(

vllm_omni/diffusion/envs.py

@@ -20,7 +20,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # used in distributed environment to determine the master address
     "MASTER_ADDR": lambda: os.getenv("MASTER_ADDR", ""),
     # used in distributed environment to manually set the communication port
-    "MASTER_PORT": lambda: (int(os.getenv("MASTER_PORT", "0")) if "MASTER_PORT" in os.environ else None),
+    "MASTER_PORT": lambda: int(os.getenv("MASTER_PORT", "0")) if "MASTER_PORT" in os.environ else None,
     # path to cudatoolkit home directory, under which should be bin, include,
     # and lib directories.
     "CUDA_HOME": lambda: os.environ.get("CUDA_HOME", None),

vllm_omni/diffusion/hooks/base.py

@@ -8,6 +8,7 @@ and modifying model forward passes without invasive changes to model code.
 
 from __future__ import annotations
 
+import inspect
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
@@ -166,7 +167,9 @@ class HookRegistry:
             # `_original_forward` attribute on the same module.
             if not hasattr(module, "_omni_original_forward"):
                 module._omni_original_forward = module.forward  # type: ignore[attr-defined]
-                module.forward = _WrappedForward(module)  # type: ignore[assignment]
+                wrapped = _WrappedForward(module)
+                wrapped.__signature__ = inspect.signature(module.forward)  # type: ignore[attr-defined]
+                module.forward = wrapped  # type: ignore[assignment]
 
         return registry
 

vllm_omni/diffusion/models/dreamid_omni/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .pipeline_dreamid_omni import DreamIDOmniPipeline
+
+__all__ = [
+    "DreamIDOmniPipeline",
+]

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -0,0 +1,335 @@
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.attention.layer import Attention
+
+try:
+    from dreamid_omni.modules.model import WanLayerNorm, WanRMSNorm
+except ImportError:
+    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
+
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply
+
+
+class FusionModel(nn.Module):
+    def __init__(self, video_config=None, audio_config=None):
+        super().__init__()
+        has_video = True
+        has_audio = True
+        if video_config is not None:
+            self.video_model = WanModel(**video_config)
+        else:
+            has_video = False
+            self.video_model = None
+            print("Warning: No video model is provided!")
+
+        if audio_config is not None:
+            self.audio_model = WanModel(**audio_config)
+        else:
+            has_audio = False
+            self.audio_model = None
+            print("Warning: No audio model is provided!")
+
+        if has_video and has_audio:
+            assert len(self.video_model.blocks) == len(self.audio_model.blocks)
+            self.num_blocks = len(self.video_model.blocks)
+
+            self.inject_cross_attention_kv_projections()
+        self.device = get_local_device()
+
+        self.num_heads = self.video_model.num_heads
+        self.head_dim = self.video_model.dim // self.video_model.num_heads
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            num_kv_heads=self.num_heads,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
+
+    def inject_cross_attention_kv_projections(self):
+        for vid_block in self.video_model.blocks:
+            vid_block.cross_attn.k_fusion = nn.Linear(vid_block.dim, vid_block.dim)
+            vid_block.cross_attn.v_fusion = nn.Linear(vid_block.dim, vid_block.dim)
+            vid_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
+            vid_block.cross_attn.norm_k_fusion = (
+                WanRMSNorm(vid_block.dim, eps=1e-6) if vid_block.qk_norm else nn.Identity()
+            )
+
+        for audio_block in self.audio_model.blocks:
+            audio_block.cross_attn.k_fusion = nn.Linear(audio_block.dim, audio_block.dim)
+            audio_block.cross_attn.v_fusion = nn.Linear(audio_block.dim, audio_block.dim)
+            audio_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
+            audio_block.cross_attn.norm_k_fusion = (
+                WanRMSNorm(audio_block.dim, eps=1e-6) if audio_block.qk_norm else nn.Identity()
+            )
+
+    def merge_kwargs(self, vid_kwargs, audio_kwargs):
+        """
+        keys in each kwarg:
+        e
+        seq_lens
+        grid_sizes
+        freqs
+        context
+        context_lens
+        """
+        merged_kwargs = {}
+        for key in vid_kwargs:
+            merged_kwargs[f"vid_{key}"] = vid_kwargs[key]
+        for key in audio_kwargs:
+            merged_kwargs[f"audio_{key}"] = audio_kwargs[key]
+        return merged_kwargs
+
+    def single_fusion_cross_attention_forward(
+        self,
+        cross_attn_block,
+        src_seq,
+        src_grid_sizes,
+        src_freqs,
+        target_seq,
+        target_seq_lens,
+        target_grid_sizes,
+        target_freqs,
+        context,
+        context_lens,
+        src_ref_lengths=None,
+        src_freqs_scaling=None,
+        target_ref_lengths=None,
+        target_freqs_scaling=None,
+    ):
+        b, n, d = src_seq.size(0), cross_attn_block.num_heads, cross_attn_block.head_dim
+        if hasattr(cross_attn_block, "k_img"):
+            ## means is i2v block
+            q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context)
+        else:
+            ## means is t2v block
+            q, k, v = cross_attn_block.qkv_fn(src_seq, context)
+            k_img = v_img = None
+
+        x = self.attn(q, k, v)
+
+        if k_img is not None:
+            img_x = self.attn(q, k_img, v_img)
+            x = x + img_x
+
+        # is_vid = src_grid_sizes.shape[1] > 1
+        # compute target attention
+        target_seq = cross_attn_block.pre_attn_norm_fusion(target_seq)
+        k_target = cross_attn_block.norm_k_fusion(cross_attn_block.k_fusion(target_seq)).view(b, -1, n, d)
+        v_target = cross_attn_block.v_fusion(target_seq).view(b, -1, n, d)
+
+        q = rope_apply(q, src_grid_sizes, src_freqs, ref_lengths=src_ref_lengths, freqs_scaling=src_freqs_scaling)
+        k_target = rope_apply(
+            k_target,
+            target_grid_sizes,
+            target_freqs,
+            ref_lengths=target_ref_lengths,
+            freqs_scaling=target_freqs_scaling,
+        )
+
+        target_x = self.attn(q, k_target, v_target)
+
+        x = x + target_x
+
+        x = x.flatten(2)  # [B, L/P, C]
+
+        x = cross_attn_block.o(x)
+        return x
+
+    def single_fusion_cross_attention_ffn_forward(
+        self,
+        attn_block,
+        src_seq,
+        src_grid_sizes,
+        src_freqs,
+        target_seq,
+        target_seq_lens,
+        target_grid_sizes,
+        target_freqs,
+        context,
+        context_lens,
+        src_e,
+        src_ref_lengths=None,
+        src_freqs_scaling=None,
+        target_ref_lengths=None,
+        target_freqs_scaling=None,
+    ):
+        src_seq = src_seq + self.single_fusion_cross_attention_forward(
+            attn_block.cross_attn,
+            attn_block.norm3(src_seq),
+            src_grid_sizes=src_grid_sizes,
+            src_freqs=src_freqs,
+            target_seq=target_seq,
+            target_seq_lens=target_seq_lens,
+            target_grid_sizes=target_grid_sizes,
+            target_freqs=target_freqs,
+            context=context,
+            context_lens=context_lens,
+            src_ref_lengths=src_ref_lengths,
+            src_freqs_scaling=src_freqs_scaling,
+            target_ref_lengths=target_ref_lengths,
+            target_freqs_scaling=target_freqs_scaling,
+        )
+        y = attn_block.ffn(attn_block.norm2(src_seq).bfloat16() * (1 + src_e[4].squeeze(2)) + src_e[3].squeeze(2))
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            src_seq = src_seq + y * src_e[5].squeeze(2)
+        return src_seq
+
+    def single_fusion_block_forward(
+        self,
+        vid_block,
+        audio_block,
+        vid,
+        audio,
+        vid_e,
+        vid_seq_lens,
+        vid_grid_sizes,
+        vid_freqs,
+        vid_context,
+        vid_context_lens,
+        vid_ref_lengths,
+        vid_freqs_scaling,
+        audio_e,
+        audio_seq_lens,
+        audio_grid_sizes,
+        audio_freqs,
+        audio_context,
+        audio_context_lens,
+        audio_ref_lengths,
+        audio_freqs_scaling,
+    ):
+        ## audio modulation
+        assert audio_e.dtype == torch.bfloat16
+        assert len(audio_e.shape) == 4 and audio_e.size(2) == 6 and audio_e.shape[1] == audio.shape[1], (
+            f"{audio_e.shape}, {audio.shape}"
+        )
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            audio_e = audio_block.modulation(audio_e).chunk(6, dim=2)
+        assert audio_e[0].dtype == torch.bfloat16
+
+        # audio self-attention
+        audio_y = audio_block.self_attn(
+            audio_block.norm1(audio).bfloat16() * (1 + audio_e[1].squeeze(2)) + audio_e[0].squeeze(2),
+            audio_seq_lens,
+            audio_grid_sizes,
+            audio_freqs,
+            ref_lengths=audio_ref_lengths,
+            freqs_scaling=audio_freqs_scaling,
+        )
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            audio = audio + audio_y * audio_e[2].squeeze(2)
+
+        ## video modulation
+        assert len(vid_e.shape) == 4 and vid_e.size(2) == 6 and vid_e.shape[1] == vid.shape[1], (
+            f"{vid_e.shape}, {vid.shape}"
+        )
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            vid_e = vid_block.modulation(vid_e).chunk(6, dim=2)
+
+        # video self-attention
+        vid_y = vid_block.self_attn(
+            vid_block.norm1(vid).bfloat16() * (1 + vid_e[1].squeeze(2)) + vid_e[0].squeeze(2),
+            vid_seq_lens,
+            vid_grid_sizes,
+            vid_freqs,
+            ref_lengths=vid_ref_lengths,
+        )
+
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            vid = vid + vid_y * vid_e[2].squeeze(2)
+
+        og_audio = audio
+
+        # audio cross-attention
+        audio = self.single_fusion_cross_attention_ffn_forward(
+            audio_block,
+            audio,
+            audio_grid_sizes,
+            audio_freqs,
+            vid,
+            vid_seq_lens,
+            vid_grid_sizes,
+            vid_freqs,
+            audio_context,
+            audio_context_lens,
+            audio_e,
+            src_ref_lengths=audio_ref_lengths,
+            src_freqs_scaling=audio_freqs_scaling,
+            target_ref_lengths=vid_ref_lengths,
+            target_freqs_scaling=None,
+        )
+
+        assert not torch.equal(og_audio, audio), "Audio should be changed after cross-attention!"
+
+        # video cross-attention
+        vid = self.single_fusion_cross_attention_ffn_forward(
+            vid_block,
+            vid,
+            vid_grid_sizes,
+            vid_freqs,
+            og_audio,
+            audio_seq_lens,
+            audio_grid_sizes,
+            audio_freqs,
+            vid_context,
+            vid_context_lens,
+            vid_e,
+            src_ref_lengths=vid_ref_lengths,
+            src_freqs_scaling=None,
+            target_ref_lengths=audio_ref_lengths,
+            target_freqs_scaling=audio_freqs_scaling,
+        )
+
+        return vid, audio
+
+    def forward(
+        self,
+        vid,
+        audio,
+        t,
+        vid_context,
+        audio_context,
+        vid_seq_len,
+        audio_seq_len,
+        ref_ip_lengths=None,
+        ref_audio_lengths=None,
+        slg_layer=False,
+        freqs_scaling=None,
+    ):
+        vid, vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs(
+            x=vid, t=t, context=vid_context, seq_len=vid_seq_len, ref_lengths=ref_ip_lengths
+        )
+
+        audio, audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs(
+            x=audio,
+            t=t,
+            context=audio_context,
+            seq_len=audio_seq_len,
+            ref_lengths=ref_audio_lengths,
+            freqs_scaling=freqs_scaling,
+        )
+
+        kwargs = self.merge_kwargs(vid_kwargs, audio_kwargs)
+
+        for i in range(self.num_blocks):
+            """
+            1 fusion block refers to 1 audio block with 1 video block.
+            """
+
+            vid_block = self.video_model.blocks[i]
+            audio_block = self.audio_model.blocks[i]
+
+            vid, audio = self.single_fusion_block_forward(
+                vid_block=vid_block, audio_block=audio_block, vid=vid, audio=audio, **kwargs
+            )
+
+        vid = self.video_model.post_transformer_block_out(vid, vid_kwargs["grid_sizes"], vid_e)
+        audio = self.audio_model.post_transformer_block_out(audio, audio_kwargs["grid_sizes"], audio_e)
+
+        return vid, audio
+
+    def set_rope_params(self):
+        self.video_model.set_rope_params()
+        self.audio_model.set_rope_params()

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -0,0 +1,508 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import math
+import os
+
+import torch
+import torch.distributed
+from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import Image, ImageOps
+from torch import nn
+from torchvision.transforms import Compose, Normalize
+from tqdm import tqdm
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+try:
+    from dreamid_omni.utils.divisible_crop import DivisibleCrop
+    from dreamid_omni.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
+    from dreamid_omni.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+    from dreamid_omni.utils.model_loading_utils import (
+        init_mmaudio_vae,
+        init_text_model,
+        init_wan_vae_2_2,
+        load_fusion_checkpoint,
+    )
+    from dreamid_omni.utils.rearrange import Rearrange
+    from dreamid_omni.utils.resize import NaResize
+except ImportError:
+    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
+from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
+
+logger = logging.getLogger(__name__)
+
+
+AUDIO_CONFIG = {
+    "patch_size": [1],
+    "model_type": "t2a",
+    "dim": 3072,
+    "ffn_dim": 14336,
+    "freq_dim": 256,
+    "num_heads": 24,
+    "num_layers": 30,
+    "in_dim": 20,
+    "out_dim": 20,
+    "text_len": 512,
+    "window_size": [-1, -1],
+    "qk_norm": True,
+    "cross_attn_norm": True,
+    "eps": 1e-6,
+    "temporal_rope_scaling_factor": 0.19676,
+}
+
+VIDEO_CONFIG = {
+    "patch_size": [1, 2, 2],
+    "model_type": "ti2v",
+    "dim": 3072,
+    "ffn_dim": 14336,
+    "freq_dim": 256,
+    "num_heads": 24,
+    "num_layers": 30,
+    "in_dim": 48,
+    "out_dim": 48,
+    "text_len": 512,
+    "window_size": [-1, -1],
+    "qk_norm": True,
+    "cross_attn_norm": True,
+    "eps": 1e-6,
+}
+
+
+class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput):
+    """DreamID-Omni pipeline for vLLM-Omni."""
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+
+        self.device = get_local_device()
+        model = od_config.model
+        self.target_dtype = od_config.dtype
+
+        # Init Models
+        ## Load VAEs
+        vae_model_video = init_wan_vae_2_2(model, rank=self.device)
+        vae_model_video.model.requires_grad_(False).eval()
+        vae_model_video.model = vae_model_video.model.bfloat16()
+        self.vae_model_video = vae_model_video
+        self.vae = self.vae_model_video
+
+        vae_model_audio = init_mmaudio_vae(model, rank=self.device)
+        vae_model_audio.requires_grad_(False).eval()
+        self.vae_model_audio = vae_model_audio.bfloat16()
+
+        # Load T5 text model
+        self.text_model = init_text_model(model, rank=self.device)
+        self.text_encoder = self.text_model.model
+
+        # Fusion model
+        ## load audio/video model config
+        Fusion_model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG)
+
+        checkpoint_path = self.od_config.model_config.get("fusion", None)
+        assert checkpoint_path is not None, "fusion checkpoint path is None"
+        load_fusion_checkpoint(Fusion_model, checkpoint_path=os.path.join(model, checkpoint_path))
+        self.model = Fusion_model
+        self.transformer = self.model
+
+        # Fixed attributes, non-configurable
+        self.audio_latent_channel = AUDIO_CONFIG.get("in_dim")
+        self.video_latent_channel = VIDEO_CONFIG.get("in_dim")
+        self.video_latent_length = 31
+        self.audio_latent_length = 157
+        self.target_area = 960 * 960
+
+        # Fixed attributes, CFG scales
+        self.video_cfg_scale = 3.0
+        self.video_ref_cfg_scale = 1.5
+        self.audio_cfg_scale = 4.0
+        self.audio_ref_cfg_scale = 2.0
+
+        # Schedulers will be set in forward
+        self.scheduler_video = None
+        self.scheduler_audio = None
+
+    def load_image_latent_ref_ip_video(
+        self,
+        images,
+        audios,
+        video_frame_height_width,
+    ):
+        # Load size.
+        patch_size = self.model.video_model.patch_size
+        vae_stride = [4, 16, 16]
+
+        # Load image and video.
+        ref_vae_latents = {
+            "image": [],
+            "audio": [],
+        }
+        video_h = video_frame_height_width[0]
+        video_w = video_frame_height_width[1]
+        ref_audio_lengths = []
+        for img in images:
+            # Calculate the required size to keep aspect ratio and fill the rest with padding.
+            img_ratio = img.width / img.height
+            target_ratio = video_w / video_h
+
+            if img_ratio > target_ratio:  # Image is wider than target
+                new_width = video_w
+                new_height = int(new_width / img_ratio)
+            else:  # Image is taller than target
+                new_height = video_h
+                new_width = int(new_height * img_ratio)
+
+            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+            # Create a new image with the target size and place the resized image in the center
+            delta_w = video_w - img.size[0]
+            delta_h = video_h - img.size[1]
+            padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+            new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
+
+            # Transform to tensor and normalize.
+            image_transform = Compose(
+                [
+                    NaResize(
+                        resolution=math.sqrt(
+                            video_frame_height_width[0] * video_frame_height_width[1]
+                        ),  # 256*448, 480*832
+                        mode="area",
+                        downsample_only=True,
+                    ),
+                    DivisibleCrop((vae_stride[1] * patch_size[1], vae_stride[2] * patch_size[2])),
+                    Normalize(0.5, 0.5),
+                    Rearrange("t c h w -> c t h w"),
+                ]
+            )
+            new_img = image_transform([new_img])
+            new_img = new_img.transpose(0, 1)
+            new_img = new_img.to(self.device)
+            new_img = new_img.to(self.target_dtype)
+
+            img_vae_latent = self.vae_model_video.wrapped_encode(new_img[:, :, None]).to(self.target_dtype).squeeze(0)
+            ref_vae_latents["image"].append(img_vae_latent)
+
+        for audio_array in audios:
+            audio_tensor = torch.from_numpy(audio_array).float().unsqueeze(0)
+            # print(audio_tensor.shape)#torch.Size([1, 81280])
+            audio_vae_latent = self.vae_model_audio.wrapped_encode(audio_tensor)
+            # print(audio_vae_latent.shape)#torch.Size([1, 20, length])
+            audio_length = audio_vae_latent.shape[2]
+            ref_audio_lengths.append(audio_length)
+            audio_vae_latent = audio_vae_latent.squeeze(0).transpose(0, 1)
+            ref_vae_latents["audio"].append(audio_vae_latent)
+
+        ref_vae_latents["image"] = torch.cat(ref_vae_latents["image"], dim=1)
+
+        ref_vae_latents["audio"] = torch.cat(ref_vae_latents["audio"], dim=0)
+
+        return ref_vae_latents, ref_audio_lengths
+
+    def load_weights(self, weights):
+        pass
+
+    def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0):
+        torch.manual_seed(4)
+
+        if solver_name == "unipc":
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
+            )
+            sample_scheduler.set_timesteps(sampling_steps, device=device, shift=shift)
+            timesteps = sample_scheduler.timesteps
+
+        elif solver_name == "dpm++":
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
+            )
+            sampling_sigmas = get_sampling_sigmas(sampling_steps, shift=shift)
+            timesteps, _ = retrieve_timesteps(sample_scheduler, device=device, sigmas=sampling_sigmas)
+
+        elif solver_name == "euler":
+            sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=shift)
+            timesteps, sampling_steps = retrieve_timesteps(
+                sample_scheduler,
+                sampling_steps,
+                device=device,
+            )
+
+        else:
+            raise NotImplementedError("Unsupported solver.")
+
+        return sample_scheduler, timesteps
+
+    def diffuse(
+        self,
+        video_noise: torch.Tensor,
+        audio_noise: torch.Tensor,
+        latents_ref_image: torch.Tensor,
+        latents_ref_audio: torch.Tensor,
+        timesteps_video: torch.Tensor,
+        timesteps_audio: torch.Tensor,
+        text_embeddings_video_pos: torch.Tensor,
+        text_embeddings_video_neg: torch.Tensor,
+        text_embeddings_audio_pos: torch.Tensor,
+        text_embeddings_audio_neg: torch.Tensor,
+        max_seq_len_video: int,
+        max_seq_len_audio: int,
+        freqs_scaling_tensor: torch.Tensor,
+        ref_ip_num: int,
+        ref_audio_length: int,
+        ref_audio_lengths: list,
+        scheduler_video,
+        scheduler_audio,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Diffusion loop with CFG parallel support for DreamID-Omni.
+        """
+        for i, (t_v, t_a) in tqdm(enumerate(zip(timesteps_video, timesteps_audio))):
+            timestep_input = torch.full((1,), t_v, device=self.device)
+
+            model_input_video = torch.cat([video_noise[:, :-ref_ip_num], latents_ref_image], dim=1)
+            model_input_video_neg = torch.cat(
+                [video_noise[:, :-ref_ip_num], torch.zeros_like(latents_ref_image)], dim=1
+            )
+
+            model_input_audio = torch.cat([audio_noise[:-ref_audio_length, :], latents_ref_audio], dim=0)
+
+            model_input_audio_neg = torch.cat(
+                [audio_noise[:-ref_audio_length, :], torch.zeros_like(latents_ref_audio)], dim=0
+            )
+            ref_ip_lengths = [ref_ip_num]
+
+            common_args = {
+                "vid_seq_len": max_seq_len_video,
+                "audio_seq_len": max_seq_len_audio,
+                "freqs_scaling": freqs_scaling_tensor,
+                "ref_ip_lengths": [ref_ip_lengths],
+                "ref_audio_lengths": [ref_audio_lengths],
+            }
+            pos_args = {
+                **common_args,
+                "audio_context": [text_embeddings_audio_pos],
+                "vid_context": [text_embeddings_video_pos],
+            }
+            neg_args = {
+                **common_args,
+                "audio_context": [text_embeddings_audio_neg],
+                "vid_context": [text_embeddings_video_neg],
+            }
+
+            if get_classifier_free_guidance_world_size() > 1:
+                # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
+                cfg_group = get_cfg_group()
+                cfg_rank = get_classifier_free_guidance_rank()
+
+                if cfg_rank == 0:
+                    pred_vid, pred_audio = self.model(
+                        vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args
+                    )
+                    pre_vid_ip_neg, _ = self.model(
+                        vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args
+                    )
+                    pred_vid_0 = pred_vid[0]
+                    pred_audio_0 = pred_audio[0]
+                    pre_vid_ip_0 = pre_vid_ip_neg[0]
+                    pred_refaudio_0 = torch.zeros_like(pred_audio_0)  # dummy tensor
+                else:
+                    pred_vid, pred_audio = self.model(
+                        vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args
+                    )
+                    _, pred_refaudio_neg = self.model(
+                        vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args
+                    )
+                    pred_vid_0 = pred_vid[0]
+                    pred_audio_0 = pred_audio[0]
+                    pre_vid_ip_0 = torch.zeros_like(pred_vid_0)  # dummy tensor
+                    pred_refaudio_0 = pred_refaudio_neg[0]
+
+                pred_vid_gathered = cfg_group.all_gather(pred_vid_0, separate_tensors=True)
+                pred_audio_gathered = cfg_group.all_gather(pred_audio_0, separate_tensors=True)
+                pre_vid_ip_gathered = cfg_group.all_gather(pre_vid_ip_0, separate_tensors=True)
+                pred_refaudio_gathered = cfg_group.all_gather(pred_refaudio_0, separate_tensors=True)
+
+                pred_vid_pos = [pred_vid_gathered[0]]
+                pred_vid_neg = [pred_vid_gathered[1]]
+                pred_audio_pos = [pred_audio_gathered[0]]
+                pred_audio_neg = [pred_audio_gathered[1]]
+                pre_vid_ip_neg = [pre_vid_ip_gathered[0]]
+                pred_refaudio_neg = [pred_refaudio_gathered[1]]
+            else:
+                pred_vid_pos, pred_audio_pos = self.model(
+                    vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args
+                )
+
+                pred_vid_neg, pred_audio_neg = self.model(
+                    vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args
+                )
+
+                pre_vid_ip_neg, _ = self.model(
+                    vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args
+                )
+
+                _, pred_refaudio_neg = self.model(
+                    vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args
+                )
+
+            pred_video_guided = (
+                pred_vid_neg[0]
+                + self.video_cfg_scale * (pred_vid_pos[0] - pred_vid_neg[0])
+                + self.video_ref_cfg_scale * (pred_vid_pos[0] - pre_vid_ip_neg[0])
+            )
+
+            pred_audio_guided = (
+                pred_audio_neg[0]
+                + self.audio_cfg_scale * (pred_audio_pos[0] - pred_audio_neg[0])
+                + self.audio_ref_cfg_scale * (pred_audio_pos[0] - pred_refaudio_neg[0])
+            )
+            video_noise = scheduler_video.step(
+                pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False
+            )[0].squeeze(0)
+            audio_noise = scheduler_audio.step(
+                pred_audio_guided.unsqueeze(0), t_a, audio_noise.unsqueeze(0), return_dict=False
+            )[0].squeeze(0)
+        return video_noise, audio_noise
+
+    def forward(
+        self,
+        request: OmniDiffusionRequest,
+        **kwargs,
+    ) -> DiffusionOutput:
+        """Main forward pass for DreamID-Omni pipeline for R2AV task."""
+        # Extract parameters from request
+        r_prompts = request.prompts[0]
+        if isinstance(r_prompts, str):
+            prompt = r_prompts
+            video_negative_prompt = ""
+            audio_negative_prompt = ""
+        else:
+            prompt = r_prompts.get("prompt", "")
+            video_negative_prompt = r_prompts.get("video_negative_prompt", "")
+            audio_negative_prompt = r_prompts.get("audio_negative_prompt", "")
+
+        multi_modal_data = r_prompts.get("multi_modal_data", {}) if not isinstance(r_prompts, str) else None
+        raw_images = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        raw_audios = multi_modal_data.get("audio", None) if multi_modal_data is not None else None
+        if raw_images is None or raw_audios is None:
+            raise ValueError("This model requires image and audio to run.")
+        if not isinstance(raw_images, list):
+            raw_images = [raw_images]
+        if not isinstance(raw_audios, list):
+            raw_audios = [raw_audios]
+        if len(raw_images) != len(raw_audios):
+            raise ValueError("'image' and 'audio' lists must have the same length.")
+
+        height = request.sampling_params.height
+        width = request.sampling_params.width
+        num_inference_steps = request.sampling_params.num_inference_steps
+        shift = request.sampling_params.extra_args.get("shift", 5.0)
+        solver_name = request.sampling_params.extra_args.get("solver_name", "unipc")
+        seed = request.sampling_params.seed if request.sampling_params.seed is not None else 42
+
+        # 1. Prepare reference latents
+        ref_vae_latents, ref_audio_lengths = self.load_image_latent_ref_ip_video(
+            images=raw_images,
+            audios=raw_audios,
+            video_frame_height_width=(height, width),
+        )
+
+        latents_ref_image = ref_vae_latents["image"]
+        latents_ref_audio = ref_vae_latents["audio"]
+        ref_ip_num = latents_ref_image.shape[1]
+        ref_audio_length = latents_ref_audio.shape[0]
+
+        # 2. scheduler
+        scheduler_video, timesteps_video = self.get_scheduler_time_steps(
+            sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
+        )
+        scheduler_audio, timesteps_audio = self.get_scheduler_time_steps(
+            sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
+        )
+
+        # 3. text embedding
+        text_embeddings = self.text_model([prompt, video_negative_prompt, audio_negative_prompt], device=self.device)
+        text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
+        text_embeddings_audio_pos = text_embeddings[0]
+        text_embeddings_video_pos = text_embeddings[0]
+        text_embeddings_video_neg = text_embeddings[1]
+        text_embeddings_audio_neg = text_embeddings[2]
+
+        video_latent_h, video_latent_w = height // 16, width // 16
+
+        video_noise_len = self.video_latent_length + ref_ip_num
+        audio_noise_len = self.audio_latent_length + ref_audio_length
+        freqs_scaling_tensor = torch.tensor(
+            self.video_latent_length / self.audio_latent_length, device=self.device, dtype=self.target_dtype
+        )
+
+        video_noise = torch.randn(
+            (self.video_latent_channel, video_noise_len, video_latent_h, video_latent_w),
+            device=self.device,
+            dtype=self.target_dtype,
+            generator=torch.Generator(device=self.device).manual_seed(seed),
+        )
+        audio_noise = torch.randn(
+            (audio_noise_len, self.audio_latent_channel),
+            device=self.device,
+            dtype=self.target_dtype,
+            generator=torch.Generator(device=self.device).manual_seed(seed),
+        )
+
+        _patch_size_h, _patch_size_w = self.model.video_model.patch_size[1], self.model.video_model.patch_size[2]
+
+        max_seq_len_video = (
+            video_noise.shape[1] * video_noise.shape[2] * video_noise.shape[3] // (_patch_size_h * _patch_size_w)
+        )
+
+        max_seq_len_audio = audio_noise_len
+
+        with torch.amp.autocast("cuda", enabled=self.target_dtype != torch.float32, dtype=self.target_dtype):
+            video_noise, audio_noise = self.diffuse(
+                video_noise,
+                audio_noise,
+                latents_ref_image,
+                latents_ref_audio,
+                timesteps_video,
+                timesteps_audio,
+                text_embeddings_video_pos,
+                text_embeddings_video_neg,
+                text_embeddings_audio_pos,
+                text_embeddings_audio_neg,
+                max_seq_len_video,
+                max_seq_len_audio,
+                freqs_scaling_tensor,
+                ref_ip_num,
+                ref_audio_length,
+                ref_audio_lengths,
+                scheduler_video,
+                scheduler_audio,
+            )
+
+        video_noise_for_decode = video_noise[:, :-ref_ip_num]
+        audio_noise_for_decode = audio_noise[:-ref_audio_length, :]
+
+        audio_latents_for_vae = audio_noise_for_decode.unsqueeze(0).transpose(1, 2)
+        generated_audio = self.vae_model_audio.wrapped_decode(audio_latents_for_vae).squeeze().cpu().float().numpy()
+
+        video_latents_for_vae = video_noise_for_decode.unsqueeze(0)
+        generated_video = self.vae_model_video.wrapped_decode(video_latents_for_vae).squeeze(0).cpu().float().numpy()
+
+        return DiffusionOutput(output=(generated_video, generated_audio))

vllm_omni/diffusion/models/dreamid_omni/wan2_2.py

@@ -0,0 +1,506 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from vllm_omni.diffusion.attention.layer import Attention
+
+try:
+    from dreamid_omni.modules.model import (
+        ChannelLastConv1d,
+        ConvMLP,
+        Head,
+        MLPProj,
+        ModulationAdd,
+        WanLayerNorm,
+        WanRMSNorm,
+        rope_apply,  # diff from  wan2.2
+        rope_params,
+        sinusoidal_embedding_1d,
+    )
+except ImportError:
+    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
+
+
+class WanSelfAttention(nn.Module):
+    """Optimized self-attention module using vLLM layers."""
+
+    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        # Unified attention layer
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            num_kv_heads=self.num_heads,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
+
+    # query, key, value function
+    def qkv_fn(self, x):
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+            ref_lengths(Tensor): Shape [B]
+        """
+        q, k, v = self.qkv_fn(x)
+
+        x = self.attn(
+            rope_apply(q, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
+            rope_apply(k, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
+            v,
+        )
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+    def qkv_fn(self, x, context):
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        return q, k, v
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        q, k, v = self.qkv_fn(x, context)
+
+        # compute attention
+        x = self.attn(q, k, v)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanI2VCrossAttention(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        raise RuntimeError("WanT2VCrossAttention is currently disabled and should not be used.")
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    "t2v_cross_attn": WanT2VCrossAttention,
+    "i2v_cross_attn": WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        cross_attn_type,
+        dim,
+        ffn_dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=False,
+        eps=1e-6,
+        additional_emb_length=None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        if cross_attn_type == "i2v_cross_attn":
+            assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn"
+            self.cross_attn = WanI2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, additional_emb_length)
+        else:
+            assert additional_emb_length is None, "additional_emb_length should be None for t2v_cross_attn"
+            self.cross_attn = WanT2VCrossAttention(
+                dim,
+                num_heads,
+                (-1, -1),
+                qk_norm,
+                eps,
+            )
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = ModulationAdd(dim, 6)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, L1, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        assert e.dtype == torch.bfloat16
+        assert len(e.shape) == 4 and e.size(2) == 6 and e.shape[1] == x.shape[1], f"{e.shape}, {x.shape}"
+        with amp.autocast("cuda", dtype=torch.bfloat16):
+            e = self.modulation(e).chunk(6, dim=2)
+        assert e[0].dtype == torch.bfloat16
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x).bfloat16() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), seq_lens, grid_sizes, freqs
+        )
+        with amp.autocast("cuda", dtype=torch.bfloat16):
+            x = x + y * e[2].squeeze(2)
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            y = self.ffn(self.norm2(x).bfloat16() * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
+            with amp.autocast("cuda", dtype=torch.bfloat16):
+                x = x + y * e[5].squeeze(2)
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+class WanModel(ModelMixin, ConfigMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video, text-to-audio.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type="t2v",
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=16,
+        dim=2048,
+        ffn_dim=8192,
+        freq_dim=256,
+        text_dim=4096,
+        additional_emb_dim=None,
+        additional_emb_length=None,
+        out_dim=16,
+        num_heads=16,
+        num_layers=32,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=True,
+        temporal_rope_scaling_factor=1.0,
+        eps=1e-6,
+    ):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in [
+            "t2v",
+            "i2v",
+            "t2a",
+            "tt2a",
+            "ti2v",
+        ]  ## tt2a means text transcript + text description to audio (to support both TTS and T2A
+        self.model_type = model_type
+        is_audio_type = "a" in self.model_type
+        is_video_type = "v" in self.model_type
+        assert is_audio_type ^ is_video_type, "Either audio or video model should be specified"
+        if is_audio_type:
+            ## audio model
+            assert len(patch_size) == 1 and patch_size[0] == 1, (
+                "Audio model should only accept 1 dimensional input, and we dont do patchify"
+            )
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        self.temporal_rope_scaling_factor = temporal_rope_scaling_factor
+        self.is_audio_type = is_audio_type
+        self.is_video_type = is_video_type
+        # embeddings
+        if is_audio_type:
+            ## hardcoded to MMAudio
+            self.patch_embedding = nn.Sequential(
+                ChannelLastConv1d(in_dim, dim, kernel_size=7, padding=3),
+                nn.SiLU(),
+                ConvMLP(dim, dim * 4, kernel_size=7, padding=3),
+            )
+        else:
+            self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
+
+        self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        ## so i2v and tt2a share the same cross attention while t2v and t2a share the same cross attention
+        cross_attn_type = "t2v_cross_attn" if model_type in ["t2v", "t2a", "ti2v"] else "i2v_cross_attn"
+
+        if cross_attn_type == "t2v_cross_attn":
+            assert additional_emb_dim is None and additional_emb_length is None, (
+                "additional_emb_length should be None for t2v and t2a model"
+            )
+        else:
+            assert additional_emb_dim is not None and additional_emb_length is not None, (
+                "additional_emb_length should be specified for i2v and tt2a model"
+            )
+
+        self.blocks = nn.ModuleList(
+            [
+                WanAttentionBlock(
+                    cross_attn_type,
+                    dim,
+                    ffn_dim,
+                    num_heads,
+                    window_size,
+                    qk_norm,
+                    cross_attn_norm,
+                    eps,
+                    additional_emb_length,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        self.set_rope_params()
+
+        if model_type in ["i2v", "tt2a"]:
+            self.img_emb = MLPProj(additional_emb_dim, dim)
+
+    def set_rope_params(self):
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        dim = self.dim
+        num_heads = self.num_heads
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+
+        if self.is_audio_type:
+            ## to be determined
+            # self.freqs = rope_params(1024, d, freqs_scaling=temporal_rope_scaling_factor)
+            self.freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=self.temporal_rope_scaling_factor)
+        else:
+            self.freqs = torch.cat(
+                [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
+                dim=1,
+            )
+
+    def prepare_transformer_block_kwargs(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        ref_lengths=None,
+        freqs_scaling=None,
+    ):
+        # params
+        ## need to change!
+        device = next(self.patch_embedding.parameters()).device
+
+        if self.is_audio_type and freqs_scaling is not None:
+            if isinstance(freqs_scaling, torch.Tensor):
+                scale_val = freqs_scaling.item()
+            else:
+                scale_val = freqs_scaling
+            d = self.dim // self.num_heads
+
+            current_freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=scale_val).to(device)
+        else:
+            current_freqs = self.freqs
+            if current_freqs.device != device:
+                current_freqs = current_freqs.to(device)
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]  ## x is list of [B L D] or [B C F H W]
+        if self.is_audio_type:
+            # [B, 1]
+            grid_sizes = torch.stack([torch.tensor(u.shape[1:2], dtype=torch.long) for u in x])
+        else:
+            # [B, 3]
+            grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+            x = [u.flatten(2).transpose(1, 2) for u in x]  # [B C F H W] -> [B (F H W) C] -> [B L C]
+
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len, f"Sequence length {seq_lens.max()} exceeds maximum {seq_len}."
+        x = torch.cat(
+            [torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]
+        )  # single [B, L, C]
+
+        # time embeddings
+        if t.dim() == 1:
+            t = t.unsqueeze(1).expand(t.size(0), seq_len)
+        with amp.autocast("cuda", dtype=torch.bfloat16):
+            bt = t.size(0)
+            t = t.flatten()
+            e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, seq_len)).bfloat16())
+            e0 = self.time_projection(e).unflatten(2, (6, self.dim))  # [1, 26784, 6, 3072] - B, seq_len, 6, dim
+            assert e.dtype == torch.bfloat16 and e0.dtype == torch.bfloat16
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
+        )
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=current_freqs,
+            context=context,
+            context_lens=context_lens,
+            ref_lengths=ref_lengths,
+            freqs_scaling=freqs_scaling,
+        )
+
+        return x, e, kwargs
+
+    def post_transformer_block_out(self, x, grid_sizes, e):
+        # head
+        x = self.head(x, e)
+        # unpatchify
+        if self.is_audio_type:
+            ## grid_sizes is [B 1] where 1 is L,
+            # converting grid_sizes from [B 1] -> [B]
+            grid_sizes = [gs[0] for gs in grid_sizes]
+            assert len(x) == len(grid_sizes)
+            x = [u[:gs] for u, gs in zip(x, grid_sizes)]
+        else:
+            ## grid_sizes is [B 3] where 3 is F H w
+            x = self.unpatchify(x, grid_sizes)
+
+        return [u.bfloat16() for u in x]
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError("WanModel model does not support forward pass.")
+
+    def unpatchify(self, x, grid_sizes) -> list[torch.Tensor]:
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            # v is [F H w] F * H * 80, 100, it was right padded by 20.
+            u = u[: math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum("fhwpqrc->cfphqwr", u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        # out is list of [C F H W]
+        return out

vllm_omni/diffusion/models/interface.py

@@ -13,6 +13,11 @@ class SupportImageInput(Protocol):
     color_format: ClassVar[str] = "RGB"  # Default color format
 
 
+@runtime_checkable
+class SupportAudioInput(Protocol):
+    support_audio_input: ClassVar[bool] = True
+
+
 @runtime_checkable
 class SupportAudioOutput(Protocol):
     support_audio_output: ClassVar[bool] = True

vllm_omni/diffusion/models/sd3/sd3_transformer.py

@@ -156,8 +156,12 @@ class SD3CrossAttention(nn.Module):
         else:
             self.to_out = None
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_added_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        if added_kv_proj_dim is not None and qk_norm:
+            self.norm_added_q = RMSNorm(head_dim, eps=eps)
+            self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        else:
+            self.norm_added_q = nn.Identity()
+            self.norm_added_k = nn.Identity()
 
         self.attn = Attention(
             num_heads=self.to_qkv.num_heads,
@@ -299,6 +303,8 @@ class SD3TransformerBlock(nn.Module):
                 dim=dim,
                 num_heads=num_attention_heads,
                 head_dim=attention_head_dim,
+                added_kv_proj_dim=None,
+                context_pre_only=True,
                 out_dim=dim,
                 qk_norm=True if qk_norm == "rms_norm" else False,
                 eps=1e-6,

vllm_omni/diffusion/registry.py

@@ -137,6 +137,11 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2",
         "Flux2Pipeline",
     ),
+    "DreamIDOmniPipeline": (
+        "dreamid_omni",
+        "pipeline_dreamid_omni",
+        "DreamIDOmniPipeline",
+    ),
 }
 
 
@@ -196,9 +201,9 @@ def initialize_model(
             od_config.vae_use_tiling = True
 
         # Configure VAE memory optimization settings from config
-        if hasattr(model.vae, "use_slicing"):
+        if hasattr(model, "vae") and hasattr(model.vae, "use_slicing"):
             model.vae.use_slicing = od_config.vae_use_slicing
-        if hasattr(model.vae, "use_tiling"):
+        if hasattr(model, "vae") and hasattr(model.vae, "use_tiling"):
             model.vae.use_tiling = od_config.vae_use_tiling
 
         if is_distributed_vae:

vllm_omni/entrypoints/async_omni.py

@@ -20,6 +20,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length, try_send_via_connector
 from vllm_omni.distributed.ray_utils.utils import try_close_ray
 from vllm_omni.engine.input_processor import OmniInputProcessor
+from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import OmniBase
 from vllm_omni.entrypoints.omni_stage import OmniStage
@@ -28,7 +29,7 @@ from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
 )
-from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
 
 # Internal imports (our code)
 from vllm_omni.lora.request import LoRARequest
@@ -125,6 +126,9 @@ class AsyncOmni(OmniBase):
         # Used to avoid race condition between output_handler and collective_rpc
         self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {}
 
+        # CFG companion → parent request ID mapping for output routing
+        self._companion_to_parent: dict[str, str] = {}
+
         super().__init__(model, **kwargs)
 
         # Register weak reference cleanup (called on garbage collection)
@@ -389,6 +393,19 @@ class AsyncOmni(OmniBase):
             req_state = ClientRequestState(request_id)
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
+
+            # Ensure modalities is in the prompt dict for CFG expansion
+            # (offline path includes it; online serving passes it separately)
+            if isinstance(prompt, dict) and output_modalities and "modalities" not in prompt:
+                prompt["modalities"] = output_modalities
+
+            # CFG companion tracking (prompt expansion + lifecycle management)
+            cfg = CfgCompanionTracker(
+                prompt_expand_func=getattr(self.stage_list[0], "prompt_expand_func", None),
+                stage0_sampling_params=sampling_params_list[0],
+            )
+            expanded_companions = cfg.expand_prompts({request_id: prompt})
+
             sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
             task = {
                 "request_id": request_id,
@@ -396,6 +413,18 @@ class AsyncOmni(OmniBase):
                 "sampling_params": sp0,
             }
             self.stage_list[0].submit(task)
+
+            # Submit CFG companion requests to stage-0
+            if cfg.is_active:
+                for companion_id, companion_prompt in expanded_companions:
+                    self._companion_to_parent[companion_id] = request_id
+                    companion_task = {
+                        "request_id": companion_id,
+                        "engine_inputs": companion_prompt,
+                        "sampling_params": cfg.stage0_sampling_params,
+                    }
+                    self.stage_list[0].submit(companion_task)
+
             metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
             _req_start_ts[request_id] = time.time()
             logger.info(
@@ -421,6 +450,7 @@ class AsyncOmni(OmniBase):
                     final_stage_id_for_e2e,
                     sampling_params_list,
                     prompt,
+                    cfg=cfg,
                 ):
                     yield output
 
@@ -440,6 +470,9 @@ class AsyncOmni(OmniBase):
                 logger.exception(f"[{self._name}] Request {request_id} Failed to finalized/build/log summary: {e}")
             finally:
                 self.request_states.pop(request_id, None)
+                if cfg.is_active:
+                    for cid in cfg.get_companion_request_ids(request_id).values():
+                        self._companion_to_parent.pop(cid, None)
         except (asyncio.CancelledError, GeneratorExit):
             await self.abort(request_id)
             logger.info("[AsyncOrchestrator] Request %s aborted.", request_id)
@@ -603,12 +636,29 @@ class AsyncOmni(OmniBase):
         final_stage_id_for_e2e: int,
         sampling_params_list: list[SamplingParams],
         prompt: Any,
+        cfg: CfgCompanionTracker | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
+            cfg_stage0 = stage_id == 0 and cfg is not None and cfg.is_active
             finished = False
-            while not finished:
+
+            while True:
+                if finished and (
+                    not cfg_stage0 or cfg.all_companions_done(request_id) or cfg.is_parent_failed(request_id)
+                ):
+                    break
+
                 result = await req_state.queue.get()
-                assert stage_id == req_state.stage_id
+
+                if cfg is not None and cfg.is_companion(result.get("request_id", "")):
+                    if cfg_stage0:
+                        rid = result.get("request_id")
+                        if "error" in result:
+                            cfg.on_companion_error(rid)
+                        else:
+                            cfg.on_companion_completed(rid)
+                    continue
+
                 engine_outputs, finished, output_to_yield = self._process_single_result(
                     result,
                     stage,
@@ -629,6 +679,16 @@ class AsyncOmni(OmniBase):
                     next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
                 sp_next: SamplingParams = sampling_params_list[next_stage_id]
 
+                if cfg is not None and cfg.is_active and not cfg.is_parent_failed(request_id):
+                    if isinstance(sp_next, OmniDiffusionSamplingParams):
+                        sp_next = copy.deepcopy(sp_next)
+                        sp_next.cfg_kv_request_ids = cfg.get_companion_request_ids(request_id)
+                        logger.info(
+                            "Attaching cfg_kv_request_ids=%s to request %s",
+                            sp_next.cfg_kv_request_ids,
+                            request_id,
+                        )
+
                 # Check if we have a connector for this edge
                 connector_key = (str(stage_id), str(next_stage_id))
                 connector = self.connectors.get(connector_key)
@@ -747,6 +807,7 @@ class AsyncOmni(OmniBase):
 
         stage_list = self.stage_list
         request_states = self.request_states
+        companion_to_parent = self._companion_to_parent
 
         async def output_handler():
             try:
@@ -773,6 +834,10 @@ class AsyncOmni(OmniBase):
                             continue
                         req_id = result.get("request_id")
                         req_state = request_states.get(req_id)
+                        if req_state is None:
+                            parent_id = companion_to_parent.get(req_id)
+                            if parent_id is not None:
+                                req_state = request_states.get(parent_id)
                         if req_state is None:
                             logger.debug(
                                 f"[{self._name}] Request may have been aborted; \

vllm_omni/entrypoints/omni_diffusion.py

@@ -61,18 +61,23 @@ class OmniDiffusion:
                 "model_index.json",
                 od_config.model,
             )
-            if config_dict is not None:
-                if od_config.model_class_name is None:
-                    od_config.model_class_name = config_dict.get("_class_name", None)
-                od_config.update_multimodal_support()
 
+            if config_dict is None:
+                raise FileNotFoundError("model_index.json not found")
+
+            if od_config.model_class_name is None:
+                od_config.model_class_name = config_dict.get("_class_name", None)
+            od_config.update_multimodal_support()
+
+            if od_config.model_class_name == "DreamIDOmniPipeline":
+                od_config.model_config = config_dict
+            else:
                 tf_config_dict = get_hf_file_to_dict(
                     "transformer/config.json",
                     od_config.model,
                 )
                 od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
-            else:
-                raise FileNotFoundError("model_index.json not found")
+
         except (AttributeError, OSError, ValueError, FileNotFoundError):
             cfg = get_hf_file_to_dict("config.json", od_config.model)
             if cfg is None:

vllm_omni/entrypoints/openai/api_server.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
 import base64
 import io
 import json
@@ -15,13 +16,14 @@ from argparse import Namespace
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from http import HTTPStatus
-from typing import Annotated, Any, cast
+from typing import Annotated, Any, Literal, cast
 
 import httpx
 import vllm.envs as envs
-from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile, WebSocket
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from PIL import Image
+from pydantic import BaseModel, Field
 from starlette.datastructures import State
 from starlette.routing import Route
 from vllm import SamplingParams
@@ -82,6 +84,7 @@ from vllm.tool_parsers import ToolParserManager
 from vllm.utils.system_utils import decorate_logs
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
     encode_image_base64,
     parse_size,
@@ -93,19 +96,52 @@ from vllm_omni.entrypoints.openai.protocol.images import (
     ImageGenerationResponse,
 )
 from vllm_omni.entrypoints.openai.protocol.videos import (
+    SecondStr,
+    SizeStr,
+    VideoDeleteResponse,
+    VideoError,
     VideoGenerationRequest,
-    VideoGenerationResponse,
+    VideoGenerationStatus,
+    VideoListResponse,
+    VideoResponse,
 )
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
-from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
+from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
+from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
+from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
+from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 
 logger = init_logger(__name__)
 router = APIRouter()
+profiler_router = APIRouter()
+
+
+def _should_enable_profiler_endpoints(args: Namespace) -> bool:
+    # Check upstream vLLM's profiler_config
+    profiler_config = getattr(args, "profiler_config", None)
+    if profiler_config is not None:
+        # profiler_config exists, check if profiler is set
+        profiler = getattr(profiler_config, "profiler", None)
+        if profiler is not None:
+            return True
+
+    # TODO: remove this env after refactoring torch profiler to CLI args
+    env_value = os.environ.get("VLLM_TORCH_PROFILER_DIR")
+    return env_value is not None
+
+
+class ProfileRequest(BaseModel):
+    """Request model for profiling endpoints."""
+
+    stages: list[int] | None = Field(
+        default=None,
+        description="List of stage IDs to profile. If None, profiles all stages.",
+    )
 
 
 def _remove_route_from_router(
@@ -235,6 +271,11 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         await omni_init_app_state(engine_client, app.state, args)
 
+        # Conditionally register profiler endpoints based on config or env var
+        if _should_enable_profiler_endpoints(args):
+            logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!")
+            app.include_router(profiler_router)
+
         vllm_config = await engine_client.get_vllm_config()
 
         # Check if pure diffusion mode (vllm_config will be None)
@@ -823,6 +864,20 @@ _remove_route_from_router(router, "/v1/audio/speech", {"POST"})
 @with_cancellation
 @load_aware_call
 async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request):
+    """Generate speech audio from text using the loaded TTS model.
+
+    Args:
+        request: Speech synthesis request in OpenAI-compatible format.
+        raw_request: Raw FastAPI request for accessing app state.
+
+    Returns:
+        The generated audio response, or an OpenAI-style error payload when
+        the request cannot be fulfilled.
+
+    Raises:
+        HTTPException: If the server does not support speech generation or the
+        synthesis request fails unexpectedly.
+    """
     handler = Omnispeech(raw_request)
     if handler is None:
         base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
@@ -854,13 +909,127 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
     },
 )
 async def list_voices(raw_request: Request):
-    """List available TTS voices/speakers from the loaded model."""
+    """List available TTS voices exposed by the loaded speech model.
+
+    Args:
+        raw_request: Raw FastAPI request for accessing app state.
+
+    Returns:
+        A JSON payload containing the sorted set of supported speaker names, or
+        an OpenAI-style error response when the current server configuration
+        does not support the Speech API.
+    """
     handler = Omnispeech(raw_request)
     if handler is None:
         return base(raw_request).create_error_response(message="The model does not support Speech API")
 
+    # Get all speakers (both model built-in and uploaded)
     speakers = sorted(handler.supported_speakers) if handler.supported_speakers else []
-    return JSONResponse(content={"voices": speakers})
+
+    # Get uploaded speakers details
+    uploaded_speakers = []
+    if hasattr(handler, "uploaded_speakers"):
+        for voice_name, info in handler.uploaded_speakers.items():
+            uploaded_speakers.append(
+                {
+                    "name": info.get("name", voice_name),
+                    "consent": info.get("consent", ""),
+                    "created_at": info.get("created_at", 0),
+                    "file_size": info.get("file_size", 0),
+                    "mime_type": info.get("mime_type", ""),
+                }
+            )
+
+    return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers})
+
+
+@router.post(
+    "/v1/audio/voices",
+    responses={
+        HTTPStatus.OK.value: {"model": dict},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def upload_voice(
+    raw_request: Request,
+    audio_sample: UploadFile = File(...),
+    consent: str = Form(...),
+    name: str = Form(...),
+):
+    """Upload a new voice sample for voice cloning.
+
+    Uploads an audio file that can be used as a reference for voice cloning
+    in Base task TTS requests. The voice can then be referenced by name
+    in subsequent TTS requests.
+
+    Args:
+        audio_sample: Audio file (max 10MB)
+        consent: Consent recording ID
+        name: Name for the new voice
+        raw_request: Raw FastAPI request
+
+    Returns:
+        JSON response with voice information
+    """
+    handler = Omnispeech(raw_request)
+    if handler is None:
+        return base(raw_request).create_error_response(message="The model does not support Speech API")
+
+    try:
+        # Upload the voice
+        result = await handler.upload_voice(audio_sample, consent, name)
+
+        return JSONResponse(content={"success": True, "voice": result})
+
+    except ValueError as e:
+        return base(raw_request).create_error_response(message=str(e))
+    except Exception as e:
+        logger.exception(f"Failed to upload voice: {e}")
+        return base(raw_request).create_error_response(message=f"Failed to upload voice: {str(e)}")
+
+
+@router.delete(
+    "/v1/audio/voices/{name}",
+    responses={
+        HTTPStatus.OK.value: {"model": dict},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def delete_voice(name: str, raw_request: Request):
+    """Delete an uploaded voice.
+
+    Deletes the voice sample and associated metadata. Also removes any
+    cached voice clone prompts for this voice.
+
+    Args:
+        name: Name of the voice to delete
+        raw_request: Raw FastAPI request
+
+    Returns:
+        JSON response indicating success or failure
+    """
+    handler = Omnispeech(raw_request)
+    if handler is None:
+        return base(raw_request).create_error_response(message="The model does not support Speech API")
+
+    try:
+        # Delete the voice
+        success = await handler.delete_voice(name)
+        if not success:
+            return JSONResponse(
+                content={"success": False, "error": f"Voice '{name}' not found"},
+                status_code=HTTPStatus.NOT_FOUND.value,
+            )
+
+        return JSONResponse(content={"success": True, "message": f"Voice '{name}' deleted successfully"})
+
+    except ValueError as e:
+        return base(raw_request).create_error_response(message=str(e))
+    except Exception as e:
+        logger.exception(f"Failed to delete voice '{name}': {e}")
+        return base(raw_request).create_error_response(message=f"Failed to delete voice: {str(e)}")
 
 
 @router.websocket("/v1/audio/speech/stream")
@@ -1397,7 +1566,7 @@ async def _generate_with_async_omni(
     return result
 
 
-def _update_if_not_none(object: any, key: str, val: any) -> None:
+def _update_if_not_none(object: Any, key: str, val: Any) -> None:
     if val is not None:
         setattr(object, key, val)
 
@@ -1503,7 +1672,7 @@ def _encode_image_base64_with_compression(
 
 def apply_stage_default_sampling_params(
     default_params_json: str | None,
-    sampling_params: any,
+    sampling_params: Any,
     stage_key: str,
 ) -> None:
     """
@@ -1523,29 +1692,16 @@ def apply_stage_default_sampling_params(
                     setattr(sampling_params, param_name, param_value)
 
 
-async def _run_video_generation(
-    request: VideoGenerationRequest,
-    raw_request: Request,
-    *,
-    input_reference_bytes: bytes | None = None,
-) -> VideoGenerationResponse:
-    handler = Omnivideo(raw_request)
-    if handler is None:
-        raise HTTPException(
-            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
-            detail="Video generation handler not initialized.",
-        )
-    logger.info("Video generation handler: %s", type(handler).__name__)
-    try:
-        return await handler.generate_videos(request, raw_request, input_reference_bytes=input_reference_bytes)
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.exception("Video generation failed: %s", e)
-        raise HTTPException(
-            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
-            detail=f"Video generation failed: {str(e)}",
-        )
+def _resolve_video_runtime_context(raw_request: Request) -> tuple[str | None, list[Any] | None]:
+    app_model_name = None
+    serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
+    if serving_models and getattr(serving_models, "base_model_paths", None):
+        base_paths = serving_models.base_model_paths
+        if base_paths:
+            app_model_name = base_paths[0].name
+
+    app_stage_configs = getattr(raw_request.app.state, "stage_configs", None)
+    return app_model_name, app_stage_configs
 
 
 def _parse_form_json(value: str | None) -> Any:
@@ -1560,10 +1716,97 @@ def _parse_form_json(value: str | None) -> Any:
         ) from exc
 
 
+def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse:
+    resp = VideoResponse(
+        model=model_name,
+        status=VideoGenerationStatus.QUEUED,
+        size=req.size,
+        prompt=req.prompt,
+    )
+    resp.seconds = str(req.seconds or resp.seconds)
+    return resp
+
+
+async def decode_and_save_video_output(output: Any, file_name: str) -> str:
+    if not output.b64_json:
+        raise RuntimeError(f"Video output for {file_name} did not include b64_json content.")
+
+    try:
+        video_bytes = base64.b64decode(output.b64_json)
+    except Exception as decode_exc:
+        raise RuntimeError(f"Failed to decode generated video payload for {file_name}") from decode_exc
+
+    return await STORAGE_MANAGER.save(video_bytes, file_name)
+
+
+def _cleanup_video(video_id: str, output_path: str | None):
+    try:
+        if output_path is not None:
+            os.remove(output_path)
+    except OSError:
+        logger.warning("Failed to cleanup partial video file '%s' for id=%s", output_path, video_id)
+
+
+async def _run_video_generation_job(
+    handler: OmniOpenAIServingVideo,
+    request: VideoGenerationRequest,
+    video_id: str,
+    reference_image: ReferenceImage | None = None,
+) -> None:
+    job = await VIDEO_STORE.get(video_id)
+    if job is None:
+        logger.warning("Video job %s missing before generation task started; skipping", video_id)
+        return
+
+    await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS})
+    started_at = time.perf_counter()
+    output_path = None
+    try:
+        response = await handler.generate_videos(request, video_id, reference_image=reference_image)
+        if not response.data:
+            raise RuntimeError("Video generation completed but returned no outputs.")
+
+        if (video_count := len(response.data)) > 1:
+            logger.warning("Video request %s generated %s outputs but we only expected one.", video_id, video_count)
+
+        file_name = f"{video_id}.{job.file_extension}"
+        output_path = await decode_and_save_video_output(response.data[0], file_name)
+        logger.info("Video request %s persisted %s output file.", video_id, output_path)
+
+        await VIDEO_STORE.update_fields(
+            video_id,
+            {
+                "status": VideoGenerationStatus.COMPLETED,
+                "progress": 100,
+                "file_name": file_name,
+                "completed_at": int(time.time()),
+                "inference_time_s": time.perf_counter() - started_at,
+            },
+        )
+    except Exception as exc:
+        logger.exception("Video generation failed for id=%s", video_id)
+
+        _cleanup_video(video_id, output_path)
+        # TODO: It would be better to have a finite collection of errors to return rather than the exception name
+        await VIDEO_STORE.update_fields(
+            video_id,
+            {
+                "status": VideoGenerationStatus.FAILED,
+                "completed_at": int(time.time()),
+                "error": VideoError(code=type(exc).__name__, message=str(exc)),
+                "inference_time_s": time.perf_counter() - started_at,
+            },
+        )
+    except asyncio.CancelledError:
+        _cleanup_video(video_id, output_path)
+        await VIDEO_STORE.pop(video_id)
+        raise
+
+
 @router.post(
     "/v1/videos",
     responses={
-        HTTPStatus.OK.value: {"model": VideoGenerationResponse},
+        HTTPStatus.OK.value: {"model": VideoResponse},
         HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
         HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
         HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
@@ -1573,11 +1816,10 @@ async def create_video(
     raw_request: Request,
     prompt: str = Form(...),
     input_reference: UploadFile | None = File(default=None),
+    image_reference: str | None = Form(default=None),
     model: str | None = Form(default=None),
-    n: int | None = Form(default=None),
-    seconds: int | None = Form(default=None),
-    size: str | None = Form(default=None),
-    response_format: str | None = Form(default=None),
+    seconds: SecondStr | None = Form(default=None),
+    size: SizeStr | None = Form(default=None),
     user: str | None = Form(default=None),
     width: int | None = Form(default=None),
     height: int | None = Form(default=None),
@@ -1592,19 +1834,60 @@ async def create_video(
     seed: int | None = Form(default=None),
     negative_prompt: str | None = Form(default=None),
     lora: str | None = Form(default=None),
-) -> VideoGenerationResponse:
-    """OpenAI-style video create endpoint (multipart form-data)."""
+) -> VideoResponse:
+    """Create an asynchronous video generation job.
+
+    This OpenAI-style endpoint accepts multipart form-data, validates the
+    request payload, persists a queued job record, and starts generation in the
+    background. The response contains metadata for polling job status rather
+    than the generated video bytes.
+
+    Args:
+        raw_request: Raw FastAPI request for accessing app state.
+        prompt: Text prompt describing the requested video.
+        input_reference: Optional uploaded reference image file.
+        image_reference: Optional JSON-encoded reference image descriptor.
+        model: Optional model name supplied by the client.
+        seconds: Optional target duration string accepted by the video API.
+        size: Optional output size string such as ``1280x720``.
+        user: Optional user identifier forwarded in the stored request.
+        width: Optional explicit output width override.
+        height: Optional explicit output height override.
+        num_frames: Optional explicit frame count override.
+        fps: Optional explicit frame rate override.
+        num_inference_steps: Optional inference step override.
+        guidance_scale: Optional primary guidance scale override.
+        guidance_scale_2: Optional secondary guidance scale override.
+        boundary_ratio: Optional boundary ratio override.
+        flow_shift: Optional flow shift override.
+        true_cfg_scale: Optional true CFG scale override.
+        seed: Optional random seed override.
+        negative_prompt: Optional negative prompt.
+        lora: Optional JSON-encoded per-request LoRA configuration.
+
+    Returns:
+        A queued ``VideoResponse`` that includes the generated job identifier
+        and initial metadata for later retrieval.
+
+    Raises:
+        HTTPException: If the request is invalid, the video handler is
+        unavailable, or job initialization fails.
+    """
     input_reference_bytes = await input_reference.read() if input_reference is not None else None
+    parsed_image_reference = _parse_form_json(image_reference)
+    if parsed_image_reference is not None and input_reference_bytes is not None:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="Provide either input_reference or image_reference, not both.",
+        )
 
     request_data: dict[str, Any] = {
         "prompt": prompt,
         "model": model,
         "seconds": seconds,
         "size": size,
+        "image_reference": parsed_image_reference,
         "user": user,
-        "response_format": response_format,
-        "input_reference": None,
-        "n": n,
         "width": width,
         "height": height,
         "num_frames": num_frames,
@@ -1619,6 +1902,235 @@ async def create_video(
         "negative_prompt": negative_prompt,
         "lora": _parse_form_json(lora),
     }
+
     request_data = {k: v for k, v in request_data.items() if v is not None}
     request = VideoGenerationRequest(**request_data)
-    return await _run_video_generation(request, raw_request, input_reference_bytes=input_reference_bytes)
+
+    handler = Omnivideo(raw_request)
+    if handler is None:
+        raise HTTPException(
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+            detail="Video generation handler not initialized.",
+        )
+    logger.info("Video generation handler: %s", type(handler).__name__)
+    try:
+        app_model_name, app_stage_configs = _resolve_video_runtime_context(raw_request)
+        effective_model_name = handler.model_name or app_model_name or request.model or "unknown"
+        if request.model is not None and effective_model_name is not None and request.model != effective_model_name:
+            logger.warning(
+                "Model mismatch: request specifies '%s' but server is running '%s'. Using server model.",
+                request.model,
+                effective_model_name,
+            )
+        handler.set_stage_configs_if_missing(app_stage_configs)
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.exception("Video generation failed: %s", e)
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+            detail=f"Video generation failed: {str(e)}",
+        )
+    ref = video_response_from_request(effective_model_name, request)
+
+    try:
+        image_data = await decode_input_reference(request.image_reference, input_reference_bytes)
+    except InvalidInputReferenceError as exc:
+        raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc
+
+    reference_image = ReferenceImage(data=image_data) if image_data is not None else image_data
+    await VIDEO_STORE.upsert(ref.id, ref)
+    task = asyncio.create_task(_run_video_generation_job(handler, request, ref.id, reference_image))
+    await VIDEO_TASKS.upsert(ref.id, task)
+    return ref
+
+
+@router.get("/v1/videos", response_model=VideoListResponse)
+async def list_videos(
+    after: str | None = None,
+    limit: int | None = Query(None, ge=0, le=100),
+    order: Annotated[Literal["asc", "desc"], Query()] = "desc",
+):
+    """List stored video generation jobs.
+
+    Args:
+        after: Optional cursor indicating the last seen video ID.
+        limit: Optional maximum number of jobs to return.
+        order: Sort order for the returned jobs by creation time.
+
+    Returns:
+        A ``VideoListResponse`` containing paginated job metadata and cursor
+        information.
+    """
+    jobs = await VIDEO_STORE.list_values()
+    jobs.sort(key=lambda j: j.created_at, reverse=order == "desc")
+
+    if after is not None:
+        idx = next((i for i, job in enumerate(jobs) if job.id == after), None)
+        jobs = [] if idx is None else jobs[idx + 1 :]
+
+    has_more = False
+    if limit is not None:
+        has_more = len(jobs) > limit
+        jobs = jobs[:limit]
+
+    first_id, last_id = None, None
+    if len(jobs) > 0:
+        first_id = jobs[0].id
+        last_id = jobs[-1].id
+
+    return VideoListResponse(data=jobs, has_more=has_more, first_id=first_id, last_id=last_id)
+
+
+@router.get("/v1/videos/{video_id}")
+async def retrieve_video(video_id: str) -> VideoResponse:
+    """Retrieve metadata for a previously created video job.
+
+    Args:
+        video_id: Identifier returned by ``POST /v1/videos``.
+
+    Returns:
+        The stored ``VideoResponse`` for the requested job.
+
+    Raises:
+        HTTPException: If the video job does not exist.
+    """
+    job = await VIDEO_STORE.get(video_id)
+    if job is None:
+        raise HTTPException(status_code=404, detail="Video not found")
+    return job
+
+
+@router.delete("/v1/videos/{video_id}")
+async def delete_video(video_id: str) -> VideoDeleteResponse:
+    """Delete a stored video job and any generated output.
+
+    If the job is still queued or running, this endpoint first attempts to
+    cancel the in-flight generation task before removing the stored metadata.
+
+    Args:
+        video_id: Identifier of the video job to delete.
+
+    Returns:
+        A ``VideoDeleteResponse`` confirming the job was removed.
+
+    Raises:
+        HTTPException: If the video job does not exist, cancellation is still
+        in progress, or output is not yet ready for a completed job.
+    """
+    job = await VIDEO_STORE.get(video_id)
+    if job is None:
+        raise HTTPException(status_code=404, detail="Video not found")
+
+    if job.status in (VideoGenerationStatus.QUEUED, VideoGenerationStatus.IN_PROGRESS):
+        task = await VIDEO_TASKS.get(video_id)
+        if task is not None:
+            task.cancel()
+            try:
+                await asyncio.wait_for(task, timeout=2.0)
+            except asyncio.TimeoutError:
+                raise HTTPException(status_code=409, detail="Cancellation in progress. Please try again later.")
+            except asyncio.CancelledError:
+                pass
+
+            await VIDEO_STORE.pop(video_id)
+            return VideoDeleteResponse(id=job.id, deleted=True)
+    elif job.status is VideoGenerationStatus.FAILED:
+        if job.file_name is not None:
+            try:
+                await STORAGE_MANAGER.delete(job.file_name)
+            except Exception:
+                logger.warning("Failed to delete stored artifact for failed video job %s", video_id, exc_info=True)
+
+        await VIDEO_STORE.pop(video_id)
+        return VideoDeleteResponse(id=job.id, deleted=True)
+
+    if job.file_name is None:
+        raise HTTPException(status_code=409, detail="Video output not yet available. Please try again later.")
+
+    await STORAGE_MANAGER.delete(job.file_name)
+    await VIDEO_STORE.pop(video_id)
+    return VideoDeleteResponse(id=job.id, deleted=True)
+
+
+@router.get("/v1/videos/{video_id}/content")
+async def download_video(video_id: str) -> FileResponse:
+    """Download the generated file for a completed video job.
+
+    Args:
+        video_id: Identifier of the video job whose output should be returned.
+
+    Returns:
+        A ``FileResponse`` streaming the generated video file from local
+        storage.
+
+    Raises:
+        HTTPException: If the job does not exist, is still in progress, or the
+        generated file is missing from disk.
+    """
+    job = await VIDEO_STORE.get(video_id)
+    if job is None:
+        raise HTTPException(status_code=404, detail="Video not found")
+
+    if job.status == VideoGenerationStatus.FAILED:
+        raise HTTPException(status_code=422, detail="Video generation failed. Check job status for error details.")
+    if not job.file_name:
+        raise HTTPException(status_code=404, detail="Generation is still in-progress")
+
+    full_path = STORAGE_MANAGER.get_full_file_path(job.file_name)
+    if not os.path.exists(full_path):
+        raise HTTPException(status_code=404, detail="Generated video file not found on disk")
+
+    return FileResponse(path=full_path, media_type=job.media_type, filename=job.file_name)
+
+
+@profiler_router.post("/start_profile")
+async def start_profile(raw_request: Request, request: ProfileRequest | None = None):
+    """Start profiling for the engine.
+
+    Args:
+        request: Optional request body with stages to profile.
+            - stages: List of stage IDs to profile. If None, profiles all stages.
+
+    Example:
+        POST /start_profile
+        {"stages": [0, 1]}  # Profile only stages 0 and 1
+    """
+    try:
+        stages = request.stages if request else None
+        logger.info("Starting profiler for stages: %s", stages if stages else "all")
+        engine_client = raw_request.app.state.engine_client
+        result = await engine_client.start_profile(stages=stages)
+        logger.info("Profiler started.")
+        return JSONResponse(content=result)
+    except Exception as e:
+        logger.exception("Failed to start profiler: %s", e)
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to start profiler: {str(e)}"
+        )
+
+
+@profiler_router.post("/stop_profile")
+async def stop_profile(raw_request: Request, request: ProfileRequest | None = None):
+    """Stop profiling for the engine.
+
+    Args:
+        request: Optional request body with stages to stop profiling.
+            - stages: List of stage IDs to stop profiling. If None, stops all stages.
+
+    Example:
+        POST /stop_profile
+        {"stages": [0, 1]}  # Stop profiling only stages 0 and 1
+    """
+    try:
+        stages = request.stages if request else None
+        logger.info("Stopping profiler for stages: %s", stages if stages else "all")
+        engine_client = raw_request.app.state.engine_client
+        result = await engine_client.stop_profile(stages=stages)
+        logger.info("Profiler stopped.")
+        return JSONResponse(content=result)
+    except Exception as e:
+        logger.exception("Failed to stop profiler: %s", e)
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to stop profiler: {str(e)}"
+        )

vllm_omni/entrypoints/openai/errors.py

@@ -0,0 +1,3 @@
+class InvalidInputReferenceError(ValueError):
+    def __init__(self, message: str = "Invalid input reference.") -> None:
+        super().__init__(message)

vllm_omni/entrypoints/openai/protocol/__init__.py

@@ -12,7 +12,6 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoData,
     VideoGenerationRequest,
     VideoGenerationResponse,
-    VideoResponseFormat,
 )
 
 __all__ = [
@@ -23,6 +22,5 @@ __all__ = [
     "VideoData",
     "VideoGenerationRequest",
     "VideoGenerationResponse",
-    "VideoResponseFormat",
     "OmniChatCompletionStreamResponse",
 ]

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -8,17 +8,40 @@ mirrors the OpenAI Images API shape, with vllm-omni extensions for diffusion
 video models (e.g., Wan2.2).
 """
 
+import mimetypes
+import time
+import uuid
 from enum import Enum
-from typing import Any
+from functools import lru_cache
+from typing import Annotated, Any, Literal
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, StringConstraints
+
+from vllm_omni.entrypoints.openai.image_api_utils import parse_size
 
 
-class VideoResponseFormat(str, Enum):
-    """Video response format."""
+class VideoGenerationStatus(str, Enum):
+    QUEUED = "queued"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    FAILED = "failed"
 
-    B64_JSON = "b64_json"
-    URL = "url"  # Not implemented in PoC
+
+SizeStr = Annotated[str, StringConstraints(pattern=r"^\d+x\d+$")]
+SecondStr = Annotated[str, StringConstraints(pattern=r"^[1-9]\d*$")]
+DEFAULT_FPS = 24
+
+
+@lru_cache
+def file_extension(media_type: str):
+    media_type = str(media_type).split(";", 1)[0].strip().lower()
+    ext = mimetypes.guess_extension(media_type, strict=False)
+
+    if ext is None:
+        raise ValueError(f"No recognized file extension for media_type {media_type}")
+
+    # Keep naming stable for unknown/unsupported MIME types.
+    return ext.lstrip(".")
 
 
 class VideoParams(BaseModel):
@@ -29,6 +52,26 @@ class VideoParams(BaseModel):
     num_frames: int | None = Field(default=None, ge=1, description="Number of frames")
     fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
 
+    @property
+    def size(self) -> str | None:
+        if self.width and self.height:
+            return f"{self.width}x{self.height}"
+        else:
+            return None
+
+
+class FileImageReference(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    file_id: str
+
+
+class UrlImageReference(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    image_url: str
+
+
+ImageReference = UrlImageReference | FileImageReference
+
 
 class VideoGenerationRequest(BaseModel):
     """
@@ -37,42 +80,37 @@ class VideoGenerationRequest(BaseModel):
     Follows the OpenAI Images API conventions with extensions for video.
     """
 
-    prompt: str = Field(..., description="Text description of the desired video(s)")
-
     # OpenAI standard fields
     model: str | None = Field(
         default=None,
         description="Model to use (optional, uses server's configured model if omitted)",
     )
-    n: int = Field(default=1, ge=1, le=4, description="Number of videos to generate")
-    seconds: int | str | None = Field(
+    prompt: str = Field(..., description="Text description of the desired video(s)")
+    seconds: SecondStr | None = Field(
         default=None,
         description="Clip duration in seconds (OpenAI-compatible, e.g., 4, 8, 12)",
     )
-    size: str | None = Field(
+    size: SizeStr | None = Field(
         default=None,
         description="Video dimensions in WIDTHxHEIGHT format (e.g., '1280x720')",
     )
-    response_format: VideoResponseFormat = Field(
-        default=VideoResponseFormat.B64_JSON, description="Format of the returned video"
+
+    image_reference: ImageReference | None = Field(
+        default=None,
+        description="Optional JSON-safe image reference that guides generation. Provide either image_url or file_id.",
     )
-    user: str | None = Field(default=None, description="User identifier for tracking")
-    stream: bool | None = Field(default=False, description="Whether to stream outputs (not yet supported)")
 
     # Video params block for extensibility
     video_params: VideoParams | None = Field(default=None, description="Optional video-specific parameters")
 
-    # Image-to-video input (OpenAI-style name)
-    input_reference: str | None = Field(
-        default=None,
-        description="Image input for I2V; provided via multipart form-data as a file upload.",
-    )
+    # User specific tracking field
+    user: str | None = Field(default=None, description="User identifier for tracking")
 
     # Video-specific fields (top-level for OpenAI-style compatibility)
     width: int | None = Field(default=None, ge=1, description="Video width in pixels")
     height: int | None = Field(default=None, ge=1, description="Video height in pixels")
-    num_frames: int | None = Field(default=None, ge=1, description="Number of frames to generate")
     fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
+    num_frames: int | None = Field(default=None, ge=1, description="Number of frames to generate")
 
     # vllm-omni extensions for diffusion control
     negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the video")
@@ -123,34 +161,25 @@ class VideoGenerationRequest(BaseModel):
         ),
     )
 
-    @field_validator("size")
-    @classmethod
-    def validate_size(cls, v):
-        if v is None:
-            return None
-        if not isinstance(v, str) or "x" not in v:
-            raise ValueError("size must be in format 'WIDTHxHEIGHT' (e.g., '1280x720')")
-        return v
+    def resolve_video_params(self) -> VideoParams:
+        vp = VideoParams(width=self.width, height=self.height, fps=self.fps, num_frames=self.num_frames)
 
-    @field_validator("response_format")
-    @classmethod
-    def validate_response_format(cls, v):
-        if v is not None and v != VideoResponseFormat.B64_JSON:
-            raise ValueError(f"Only 'b64_json' response format is supported, got: {v}")
-        return v
+        if self.video_params is not None:
+            vp.width = vp.width or self.video_params.width
+            vp.height = vp.height or self.video_params.height
+            vp.fps = vp.fps or self.video_params.fps
+            vp.num_frames = vp.num_frames or self.video_params.num_frames
 
-    @field_validator("seconds")
-    @classmethod
-    def validate_seconds(cls, v):
-        if v is None:
-            return None
-        if isinstance(v, str):
-            if not v.isdigit():
-                raise ValueError("seconds must be an integer or numeric string")
-            return int(v)
-        if isinstance(v, int):
-            return v
-        raise ValueError("seconds must be an integer or numeric string")
+        if self.size:
+            vp.width, vp.height = parse_size(self.size)
+
+        if vp.fps is None:
+            vp.fps = DEFAULT_FPS
+
+        if vp.num_frames is None and self.seconds is not None:
+            vp.num_frames = int(self.seconds) * int(vp.fps)
+
+        return vp
 
 
 class VideoData(BaseModel):
@@ -166,3 +195,74 @@ class VideoGenerationResponse(BaseModel):
 
     created: int = Field(..., description="Unix timestamp of when the generation completed")
     data: list[VideoData] = Field(..., description="Array of generated videos")
+
+
+class VideoError(BaseModel):
+    code: str = Field(..., description="A machine-readable error code that was returned.")
+    message: str = Field(..., description="A human-readable description of the error that was returned.")
+
+
+class VideoResponse(BaseModel):
+    """Stored metadata for an async video generation job."""
+
+    # OpenAI standard fields
+    model: str = Field(..., description="Model name used for video generation.")
+    prompt: str = Field(..., description="The prompt that was used to generate the video.")
+    id: str = Field(
+        default_factory=lambda: f"video_gen_{uuid.uuid4().hex}", description="Unique id for a video request"
+    )
+    object: Literal["video"] = Field(default="video", description="Object type identifier.")
+    status: VideoGenerationStatus = Field(
+        default=VideoGenerationStatus.QUEUED, description="Current lifecycle status of the video job."
+    )
+    size: SizeStr | None = Field(default=None, description="Requested output size in WIDTHxHEIGHT format.")
+    progress: int = Field(default=0, description="Best-effort progress indicator from 0 to 100.")
+    seconds: SecondStr = Field(default="4", description="Requested clip length in seconds.")
+    quality: str = Field(default="default", description="Requested quality level for generation.")
+    completed_at: int | None = Field(
+        default=None, description="Unix timestamp (seconds) for when the job completed, if finished."
+    )
+    created_at: int = Field(
+        default_factory=lambda: int(time.time()),
+        description="Unix timestamp (seconds) for when the job was created.",
+    )
+    remixed_from_video_id: str | None = Field(
+        default=None,
+        description="Optional source video id for remix/edit flows.",
+    )
+    error: VideoError | None = Field(
+        default=None,
+        description="Error payload that explains why generation failed, if applicable.",
+    )
+
+    # vLLM specific fields
+    media_type: Literal["video/mp4"] = Field(default="video/mp4", description="MIME type of the generated artifact.")
+
+    expires_at: int | None = Field(default=None, description="Unix timestamp when this record is considered expired.")
+    file_name: str | None = Field(
+        default=None,
+        description="Filename of the saved output video files for this job.",
+    )
+    inference_time_s: float | None = Field(default=None, description="End-to-end inference time in seconds.")
+
+    @property
+    def file_extension(self) -> str:
+        return file_extension(self.media_type)
+
+
+class VideoDeleteResponse(BaseModel):
+    id: str = Field(description="Identifier of the deleted video.")
+    deleted: bool = Field(description="Indicates that the video resource was deleted.")
+    object: Literal["video.deleted"] = Field(
+        default="video.deleted", description="The object type that signals the deletion response."
+    )
+
+
+class VideoListResponse(BaseModel):
+    """Paginated-style wrapper for listing stored video jobs."""
+
+    first_id: str | None = Field(..., description="The ID of the first item in the list.")
+    last_id: str | None = Field(..., description="The ID of the last item in the list.")
+    has_more: bool = Field(..., description="Whether there are more items available.")
+    data: list[VideoResponse] = Field(..., description="Array of video job records.")
+    object: Literal["list"] = Field(default="list", description="Object type identifier for list responses.")

vllm_omni/entrypoints/openai/serving_chat.py

@@ -311,17 +311,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 negative_prompt = extra_body.get("negative_prompt")
 
                 engine_prompt_image: dict[str, Any] | None = None
+                is_img2img = False
                 if reference_images:
                     # Best-effort decode first reference image for i2i.
                     try:
                         img_bytes = base64.b64decode(reference_images[0])
                         img = Image.open(BytesIO(img_bytes))
-                        engine_prompt_image = {"image": img}
+                        engine_prompt_image = {"img2img": img}
+                        is_img2img = True
                     except Exception:
                         engine_prompt_image = None
 
                 # Override the prompts produced by chat-template preprocessing.
                 tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
+                if is_img2img:
+                    tprompt["modalities"] = ["img2img"]
                 if negative_prompt is not None:
                     tprompt["negative_prompt"] = negative_prompt
                 # GLM-Image's _call_hf_processor expects target_h/target_w in mm_processor_kwargs
@@ -490,10 +494,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             )
 
         # Preserve a clean text prompt for downstream stages (e.g., GLM-Image diffusion).
-        # For /v1/chat/completions, `request_prompt` is often the rendered chat template.
-        # Diffusion models generally want the raw user caption instead.
-        output_modalities = getattr(self.engine_client, "output_modalities", None)
-        if output_modalities and ("image" in output_modalities):
+        # For image generation, we want the raw user caption instead of a rendered template.
+        # But for multimodal comprehension (img2text), we MUST keep the rendered prompt
+        # containing image tokens.
+        req_modalities = getattr(request, "modalities", [])
+        if req_modalities and ("image" in req_modalities):
             messages_as_dicts: list[dict[str, Any]] = []
             for msg in messages:
                 if hasattr(msg, "model_dump"):

vllm_omni/entrypoints/openai/serving_video.py

@@ -4,23 +4,22 @@
 from __future__ import annotations
 
 import time
-import uuid
+from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, cast
 
-from fastapi import HTTPException, Request
+from fastapi import HTTPException
 from PIL import Image
 from vllm.engine.protocol import EngineClient
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.entrypoints.openai.image_api_utils import parse_size
 from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoData,
     VideoGenerationRequest,
     VideoGenerationResponse,
 )
-from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference, encode_video_base64
+from vllm_omni.entrypoints.openai.video_api_utils import encode_video_base64
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
@@ -28,6 +27,13 @@ from vllm_omni.lora.utils import stable_lora_int_id
 logger = init_logger(__name__)
 
 
+@dataclass
+class ReferenceImage:
+    """Reference class for tracking additional metadata if needed"""
+
+    data: Image.Image
+
+
 class OmniOpenAIServingVideo:
     """OpenAI-style video generation handler for omni diffusion models."""
 
@@ -41,6 +47,18 @@ class OmniOpenAIServingVideo:
         self._model_name = model_name
         self._stage_configs = stage_configs
 
+    @property
+    def model_name(self) -> str | None:
+        return self._model_name
+
+    @property
+    def stage_configs(self) -> list[Any] | None:
+        return self._stage_configs
+
+    def set_stage_configs_if_missing(self, stage_configs: list[Any] | None) -> None:
+        if self._stage_configs is None and stage_configs is not None:
+            self._stage_configs = stage_configs
+
     @classmethod
     def for_diffusion(
         cls,
@@ -57,57 +75,32 @@ class OmniOpenAIServingVideo:
     async def generate_videos(
         self,
         request: VideoGenerationRequest,
-        raw_request: Request | None = None,
+        reference_id: str,
         *,
-        input_reference_bytes: bytes | None = None,
+        reference_image: ReferenceImage | None = None,
     ) -> VideoGenerationResponse:
-        if request.stream:
-            raise HTTPException(
-                status_code=HTTPStatus.BAD_REQUEST.value,
-                detail="Streaming video generation is not supported yet.",
-            )
-
-        request_id = f"video_gen_{uuid.uuid4().hex}"
-        model_name = self._resolve_model_name(raw_request)
-
-        if request.model is not None and model_name is not None and request.model != model_name:
-            logger.warning(
-                "Model mismatch: request specifies '%s' but server is running '%s'. Using server model.",
-                request.model,
-                model_name,
-            )
-
-        prompt: OmniTextPrompt = {"prompt": request.prompt}
+        prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt)
         if request.negative_prompt is not None:
             prompt["negative_prompt"] = request.negative_prompt
 
-        input_image = None
-        try:
-            input_image = decode_input_reference(request.input_reference, input_reference_bytes)
-        except ValueError as exc:
-            raise HTTPException(
-                status_code=HTTPStatus.BAD_REQUEST.value,
-                detail=str(exc),
-            ) from exc
+        gen_params = OmniDiffusionSamplingParams()
 
-        gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
-
-        width, height, num_frames, fps = self._resolve_video_params(request)
-        if input_image is not None and width is not None and height is not None:
-            target_size = (width, height)
+        input_image = None if reference_image is None else reference_image.data
+        vp = request.resolve_video_params()
+        if input_image is not None and vp.width is not None and vp.height is not None:
+            target_size = (vp.width, vp.height)
             if input_image.size != target_size:
                 input_image = input_image.resize(target_size, Image.Resampling.LANCZOS)
         if input_image is not None:
             prompt["multi_modal_data"] = {"image": input_image}
-
-        if width is not None and height is not None:
-            gen_params.width = width
-            gen_params.height = height
-        if num_frames is not None:
-            gen_params.num_frames = num_frames
-        if fps is not None:
-            gen_params.fps = fps
-            gen_params.frame_rate = float(fps)
+        if vp.width is not None and vp.height is not None:
+            gen_params.width = vp.width
+            gen_params.height = vp.height
+        if vp.num_frames is not None:
+            gen_params.num_frames = vp.num_frames
+        if vp.fps is not None:
+            gen_params.fps = vp.fps
+            gen_params.frame_rate = float(vp.fps)
 
         if request.num_inference_steps is not None:
             gen_params.num_inference_steps = request.num_inference_steps
@@ -140,12 +133,12 @@ class OmniOpenAIServingVideo:
             gen_params.seed,
         )
 
-        result = await self._run_generation(prompt, gen_params, request_id, raw_request)
+        result = await self._run_generation(prompt, gen_params, reference_id)
         _t_encode_start = time.perf_counter()
         videos = self._extract_video_outputs(result)
         audios = self._extract_audio_outputs(result, expected_count=len(videos))
-        output_fps = fps or 24
         audio_sample_rate = self._resolve_audio_sample_rate(result)
+        output_fps = vp.fps or 24
 
         video_data = [
             VideoData(
@@ -166,37 +159,6 @@ class OmniOpenAIServingVideo:
         logger.info("Video response encoding (MP4+base64): %.2f ms", _t_encode_ms)
         return VideoGenerationResponse(created=int(time.time()), data=video_data)
 
-    def _resolve_model_name(self, raw_request: Request | None) -> str | None:
-        if self._model_name:
-            return self._model_name
-        if raw_request is None:
-            return None
-        serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
-        if serving_models and getattr(serving_models, "base_model_paths", None):
-            base_paths = serving_models.base_model_paths
-            if base_paths:
-                return base_paths[0].name
-        return None
-
-    @staticmethod
-    def _resolve_video_params(request: VideoGenerationRequest) -> tuple[int | None, int | None, int | None, int | None]:
-        width = request.width or (request.video_params.width if request.video_params else None)
-        height = request.height or (request.video_params.height if request.video_params else None)
-        num_frames = request.num_frames or (request.video_params.num_frames if request.video_params else None)
-        fps = request.fps or (request.video_params.fps if request.video_params else None)
-        seconds = request.seconds
-
-        if request.size:
-            width, height = parse_size(request.size)
-
-        if fps is None:
-            fps = 24  # Default FPS if not specified
-
-        if num_frames is None and seconds is not None:
-            num_frames = int(seconds) * int(fps)
-
-        return width, height, num_frames, fps
-
     @staticmethod
     def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None:
         if lora_body is None:
@@ -237,22 +199,15 @@ class OmniOpenAIServingVideo:
         prompt: OmniTextPrompt,
         gen_params: OmniDiffusionSamplingParams,
         request_id: str,
-        raw_request: Request | None,
     ) -> Any:
         has_stage_list = hasattr(self._engine_client, "stage_list")
         logger.info(
             "Video generation routing: stage_configs=%s, has_stage_list=%s, engine_type=%s",
-            "present"
-            if (self._stage_configs or (getattr(raw_request.app.state, "stage_configs", None) if raw_request else None))
-            else "missing",
+            "present" if self._stage_configs else "missing",
             has_stage_list,
             type(self._engine_client).__name__,
         )
-        stage_configs = (
-            self._stage_configs
-            or (getattr(raw_request.app.state, "stage_configs", None) if raw_request else None)
-            or getattr(self._engine_client, "stage_configs", None)
-        )
+        stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None)
 
         if not stage_configs:
             if not hasattr(self._engine_client, "stage_list"):

vllm_omni/entrypoints/openai/storage.py

@@ -0,0 +1,70 @@
+import asyncio
+import os
+from abc import ABC, abstractmethod
+from tempfile import NamedTemporaryFile
+
+
+class StorageBaseManager(ABC):
+    @abstractmethod
+    async def save(self, *args, **kwargs) -> str:
+        pass
+
+    @abstractmethod
+    async def delete(self, *args, **kwargs) -> bool:
+        pass
+
+
+class LocalStorageManager(StorageBaseManager):
+    def __init__(self, storage_path: str | None = None):
+        if storage_path is None:
+            storage_path = os.getenv("VLLM_OMNI_STORAGE_PATH", "/tmp/storage")
+        self.storage_path = storage_path
+        os.makedirs(self.storage_path, exist_ok=True)
+
+        max_concurrency = int(os.getenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "4"))
+        self._io_semaphore = asyncio.Semaphore(max(1, max_concurrency))
+
+    def _save_sync(self, data: bytes, file_name: str) -> str:
+        filename = self.get_full_file_path(file_name)
+        tmp_name: str | None = None
+        try:
+            with NamedTemporaryFile("wb", dir=self.storage_path, delete=False) as f:
+                tmp_name = f.name
+                f.write(data)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp_name, filename)
+            return filename
+        except Exception:
+            if tmp_name is not None:
+                try:
+                    os.remove(tmp_name)
+                except OSError:
+                    pass
+            raise
+
+    async def save(self, data: bytes, file_name: str) -> str:
+        async with self._io_semaphore:
+            return await asyncio.to_thread(self._save_sync, data, file_name)
+
+    def _delete_sync(self, file_name: str) -> bool:
+        try:
+            os.remove(self.get_full_file_path(file_name))
+        except FileNotFoundError:
+            return False
+        return True
+
+    async def delete(self, file_name: str) -> bool:
+        async with self._io_semaphore:
+            return await asyncio.to_thread(self._delete_sync, file_name)
+
+    def exists(self, file_name: str) -> bool:
+        return os.path.exists(self.get_full_file_path(file_name))
+
+    def get_full_file_path(self, file_name: str) -> str:
+        return os.path.join(self.storage_path, file_name)
+
+
+# Should implement TTL management on local file storage
+# and recovery on restart in the *STORE's
+STORAGE_MANAGER = LocalStorageManager()

vllm_omni/entrypoints/openai/stores.py

@@ -0,0 +1,74 @@
+import asyncio
+from typing import Any, Generic, TypeVar
+
+from pydantic import BaseModel
+
+from vllm_omni.entrypoints.openai.protocol.videos import VideoResponse
+
+T = TypeVar("T", bound=BaseModel)
+U = TypeVar("U")
+
+
+class TaskRegistry:
+    def __init__(self) -> None:
+        self._tasks: dict[str, asyncio.Task[None]] = {}
+        self._lock = asyncio.Lock()
+
+    async def get(self, key: str) -> asyncio.Task[None] | None:
+        async with self._lock:
+            return self._tasks.get(key)
+
+    async def pop(self, key: str) -> asyncio.Task[None] | None:
+        async with self._lock:
+            return self._tasks.pop(key, None)
+
+    async def upsert(self, key: str, task: asyncio.Task[None]) -> None:
+        def _cleanup(_: asyncio.Task[None]) -> None:
+            asyncio.create_task(self.pop(key))
+
+        task.add_done_callback(_cleanup)
+
+        async with self._lock:
+            self._tasks[key] = task
+
+
+class AsyncDictStore(Generic[T]):
+    """A small async-safe in-memory key-value store for dict items.
+
+    This encapsulates the usual pattern of a module-level dict guarded by
+    an asyncio.Lock and provides simple CRUD methods that are safe to call
+    concurrently from FastAPI request handlers and background tasks.
+    """
+
+    def __init__(self) -> None:
+        self._items: dict[str, T] = {}
+        self._lock = asyncio.Lock()
+
+    async def upsert(self, key: str, value: T) -> None:
+        async with self._lock:
+            self._items[key] = value
+
+    async def update_fields(self, key: str, updates: dict[str, Any]) -> T | None:
+        async with self._lock:
+            item: T | None = self._items.get(key)
+            if item is None:
+                return None
+            new_item = item.model_copy(update=updates)
+            self._items[key] = new_item
+            return new_item
+
+    async def get(self, key: str) -> T | None:
+        async with self._lock:
+            return self._items.get(key)
+
+    async def pop(self, key: str) -> T | None:
+        async with self._lock:
+            return self._items.pop(key, None)
+
+    async def list_values(self) -> list[T]:
+        async with self._lock:
+            return list(self._items.values())
+
+
+VIDEO_STORE: AsyncDictStore[VideoResponse] = AsyncDictStore()
+VIDEO_TASKS: TaskRegistry = TaskRegistry()

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -7,32 +7,82 @@ Shared helper utilities for OpenAI-compatible video generation API.
 from __future__ import annotations
 
 import base64
+import binascii
 import os
 import tempfile
 from io import BytesIO
 from typing import Any
 
+import httpx
 import numpy as np
 import torch
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
+
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
+from vllm_omni.entrypoints.openai.protocol.videos import (
+    FileImageReference,
+    ImageReference,
+    UrlImageReference,
+)
 
 
-def decode_input_reference(input_reference: str | None, input_reference_bytes: bytes | None) -> Image.Image | None:
-    """Decode image input from multipart bytes or base64/data URL."""
-    if input_reference and input_reference_bytes:
-        raise ValueError("Provide input_reference either as file upload or base64, not both.")
-    if input_reference_bytes:
-        return Image.open(BytesIO(input_reference_bytes)).convert("RGB")
+def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
+    try:
+        return Image.open(BytesIO(image_bytes)).convert("RGB")
+    except (UnidentifiedImageError, OSError, ValueError) as exc:
+        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc
+
+
+def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
     if input_reference:
         if input_reference.startswith("data:image"):
             _, b64_data = input_reference.split(",", 1)
         else:
             b64_data = input_reference
+
         try:
             image_bytes = base64.b64decode(b64_data)
-        except Exception as exc:  # pragma: no cover - malformed base64
-            raise ValueError("Invalid base64 input_reference.") from exc
-        return Image.open(BytesIO(image_bytes)).convert("RGB")
+        except (binascii.Error, ValueError) as exc:  # pragma: no cover - malformed base64
+            raise InvalidInputReferenceError(f"Invalid {source}: image data is not valid base64.") from exc
+        return _decode_image_bytes(image_bytes, source=source)
+    raise InvalidInputReferenceError(f"Invalid {source}: image data is empty.")
+
+
+async def decode_image_url(image_url: str) -> Image.Image:
+    if image_url.startswith("data:image"):
+        return _decode_base64_image(image_url, source="image_reference.image_url")
+
+    if image_url.startswith(("http://", "https://")):
+        async with httpx.AsyncClient(timeout=60) as client:
+            try:
+                response = await client.get(image_url)
+                response.raise_for_status()
+            except httpx.HTTPError as exc:
+                raise InvalidInputReferenceError(
+                    "Invalid image_reference.image_url: failed to download image."
+                ) from exc
+        return _decode_image_bytes(response.content, source="image_reference.image_url")
+
+    raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")
+
+
+async def decode_input_reference(
+    image_reference: ImageReference | None,
+    input_reference_bytes: bytes | None,
+) -> Image.Image | None:
+    """Decode image input from multipart bytes, base64/data URL, or image_reference."""
+
+    if input_reference_bytes is not None and image_reference is not None:
+        raise InvalidInputReferenceError("Provide either input_reference or image_reference, not both.")
+
+    if isinstance(input_reference_bytes, bytes):
+        return _decode_image_bytes(input_reference_bytes, source="input_reference")
+
+    if isinstance(image_reference, UrlImageReference):
+        return await decode_image_url(image_reference.image_url)
+    elif isinstance(image_reference, FileImageReference):
+        raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")
+
     return None
 
 

vllm_omni/model_executor/models/bagel/bagel.py

@@ -202,6 +202,15 @@ class OmniBagelDataParser(MultiModalDataParser):
 
 
 class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]):
+    IMG2IMG_PLACEHOLDER = "<|fim_middle|>"
+
+    def _cached_apply_hf_processor(self, inputs, timing_ctx):
+        # img2img: prompt text must be modified based on mm data presence,
+        # so text and mm data cannot be tokenized separately — bypass cache.
+        if inputs.mm_data_items.get_all_counts().get("img2img", 0) > 0:
+            return self._apply_hf_processor(inputs, timing_ctx)
+        return super()._cached_apply_hf_processor(inputs, timing_ctx)
+
     def _get_mm_fields_config(self, hf_inputs, hf_processor_mm_kwargs):
         return {
             "pixel_values": MultiModalFieldConfig.batched("image"),
@@ -218,6 +227,9 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn
         has_image = "images" in mm_data
         has_img2img = "pixel_values_img2img" in mm_data
 
+        if has_img2img and self.IMG2IMG_PLACEHOLDER not in prompt:
+            prompt = f"{self.IMG2IMG_PLACEHOLDER}{prompt}"
+
         if has_image and has_img2img:
             outputs = BatchFeature()
 

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -78,7 +78,11 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalKwargsItems,
 )
-from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -115,6 +119,15 @@ class GlmImagePixelInputs(TensorSchema):
     image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
 
 
+class GlmImageDataParser(MultiModalDataParser):
+    """GLM-Image treats ``img2img`` input identically to ``image``."""
+
+    def _get_subparsers(self):
+        parsers = super()._get_subparsers()
+        parsers["img2img"] = self._parse_image_data
+        return parsers
+
+
 class GlmImageProcessingInfo(BaseProcessingInfo):
     """
     Processing information for GLM-Image model.
@@ -162,14 +175,19 @@ class GlmImageProcessingInfo(BaseProcessingInfo):
             **kwargs,
         )
 
+    def get_data_parser(self) -> GlmImageDataParser:
+        return GlmImageDataParser(
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         # GLM-Image is an image GENERATION model that supports:
         # - Text-to-image (t2i): no multimodal input needed
         # - Image-to-image (i2i): source images provided as input
         #
         # For i2i mode, we support up to 1 image as condition.
-        # The model architecture supports multiple images but typical usage is 1.
-        return {"image": 1}
+        # "img2img" is an alias used by the serving layer; parsed as "image".
+        return {"image": 1, "img2img": 1}
 
     def get_num_image_tokens(
         self,