diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 5479f8ac1..7fa510716 100644 --- a/.buildkite/test-merge.yml +++ b/.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 diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py index 2c936d33f..8600fe393 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py +++ b/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, diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py index 642c83c8d..42d396f06 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py +++ b/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"] diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py index c5ac4bdf4..f607f2cb8 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py +++ b/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 diff --git a/docs/.nav.yml b/docs/.nav.yml index eb135c291..41b190f42 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 4c709425e..87cdccef1 100644 --- a/docs/models/supported_models.md +++ b/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 diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 505f9627f..38cd5e67c 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/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. diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 7f32e3eff..e9179b575 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/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 diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md new file mode 100644 index 000000000..59b993a72 --- /dev/null +++ b/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). diff --git a/docs/user_guide/examples/online_serving/image_to_video.md b/docs/user_guide/examples/online_serving/image_to_video.md index dcbb10067..a2b30320b 100644 --- a/docs/user_guide/examples/online_serving/image_to_video.md +++ b/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" diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index a632c49a7..309edb840 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/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": "" } - ] + "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 diff --git a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py b/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py new file mode 100644 index 000000000..0dbf402e9 --- /dev/null +++ b/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) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md new file mode 100644 index 000000000..59b993a72 --- /dev/null +++ b/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). diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py new file mode 100644 index 000000000..17d0f06c3 --- /dev/null +++ b/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() diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md index 7e67a15e4..fea99efa6 100644 --- a/examples/online_serving/image_to_video/README.md +++ b/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 +``` diff --git a/examples/online_serving/image_to_video/run_curl_image_to_video.sh b/examples/online_serving/image_to_video/run_curl_image_to_video.sh index 5e6faf477..f4c1496a6 100644 --- a/examples/online_serving/image_to_video/run_curl_image_to_video.sh +++ b/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}" diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index a8c69e87b..92afb1fc9 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/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": "" } - ] + "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 ``` diff --git a/examples/online_serving/text_to_video/run_curl_text_to_video.sh b/examples/online_serving/text_to_video/run_curl_text_to_video.sh index 43889c398..a95fa53f8 100644 --- a/examples/online_serving/text_to_video/run_curl_text_to_video.sh +++ b/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}" diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/comfyui/test_comfyui_integration.py index d28818ab0..24f16d5e9 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/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", diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index 63cb8743d..d5b5e39e5 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/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 diff --git a/tests/e2e/online_serving/test_video_generation_api.py b/tests/e2e/online_serving/test_video_generation_api.py new file mode 100644 index 000000000..0711a1048 --- /dev/null +++ b/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 diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 307780309..0f905d9e0 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/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 diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index b68418e17..a2acc7d75 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/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.", diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 634cb3294..3e44f861b 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index abd668466..87b621d7b 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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( diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index dd566b059..a71dc2e8e 100644 --- a/vllm_omni/diffusion/envs.py +++ b/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), diff --git a/vllm_omni/diffusion/hooks/base.py b/vllm_omni/diffusion/hooks/base.py index 3c975408c..cda4201cc 100644 --- a/vllm_omni/diffusion/hooks/base.py +++ b/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 diff --git a/vllm_omni/diffusion/models/dreamid_omni/__init__.py b/vllm_omni/diffusion/models/dreamid_omni/__init__.py new file mode 100644 index 000000000..f14db611e --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py new file mode 100644 index 000000000..2a4e485fa --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py new file mode 100644 index 000000000..f8074fee2 --- /dev/null +++ b/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)) diff --git a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py new file mode 100644 index 000000000..db289daf9 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index abdfe1e50..ef90e9804 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/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 diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py index c4a04febb..308bd35a1 100644 --- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py +++ b/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, diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index d5af2de93..77e7885c0 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 1ba95124a..b98679033 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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; \ diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index fbfcec1d4..d70aa5c52 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 82770e02d..4929f951c 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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)}" + ) diff --git a/vllm_omni/entrypoints/openai/errors.py b/vllm_omni/entrypoints/openai/errors.py new file mode 100644 index 000000000..fda9b6436 --- /dev/null +++ b/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) diff --git a/vllm_omni/entrypoints/openai/protocol/__init__.py b/vllm_omni/entrypoints/openai/protocol/__init__.py index 3788ee2af..c73203cc4 100644 --- a/vllm_omni/entrypoints/openai/protocol/__init__.py +++ b/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", ] diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index 49fa43189..1da6bdb08 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/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.") diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 88a5378e2..ada4a698b 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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"): diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index b663eb52f..6a16e7c70 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/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"): diff --git a/vllm_omni/entrypoints/openai/storage.py b/vllm_omni/entrypoints/openai/storage.py new file mode 100644 index 000000000..52aef8948 --- /dev/null +++ b/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() diff --git a/vllm_omni/entrypoints/openai/stores.py b/vllm_omni/entrypoints/openai/stores.py new file mode 100644 index 000000000..98f8696af --- /dev/null +++ b/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() diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 64672ef6f..2ed1fd3de 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/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 diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 24a1331f9..a053e98dd 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/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() diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index 90bf26ede..f90826fd3 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/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,