Merge branch 'main' into dev/rebase-v0170

Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>

.pre-commit-config.yaml

@@ -8,6 +8,7 @@ repos:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: check-yaml
         args: ["--unsafe"]
+        exclude: "examples/online_serving/chart-helm/templates/"
       - id: debug-statements
       - id: end-of-file-fixer
       - id: mixed-line-ending

docker/Dockerfile.ci

@@ -7,7 +7,7 @@ COPY . .
 
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.rocm

@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
 
 # Step 1: Setup - Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.xpu

@@ -14,6 +14,7 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
 RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     curl \
+    espeak-ng \
     ffmpeg \
     git \
     libsndfile1 \

docs/features/custom_pipeline.md

@@ -115,7 +115,7 @@ python image_edit.py \
     --image cherry_blossom.jpg \
     --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
     --output output_image_edit.png \
-    --num_inference_steps 10
+    --num-inference-steps 10
 ```
 
 ## Advanced Usage

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

@@ -34,13 +34,13 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.16.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 ```
 
-The release 0.16.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.16.0
+git checkout v0.17.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
-The release 0.16.0 of vLLM requires ROCm 7.0 environment.
+The release 0.17.0 of vLLM requires ROCm 7.0 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.16.0
+git checkout v0.17.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```

docs/getting_started/quickstart.md

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

docs/serving/audio_generate_api.md

@@ -0,0 +1,338 @@
+# Audio Generate API
+
+vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
+
+Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
+
+Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### Generate Audio
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0,
+    },
+    timeout=300.0,
+)
+
+with open("cat.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+
+#### Diffusion Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with the appropriate `Content-Type` header:
+
+| `response_format` | Content-Type |
+|--------------------|--------------|
+| `wav` | `audio/wav` |
+| `mp3` | `audio/mpeg` |
+| `flac` | `audio/flac` |
+| `pcm` | `audio/pcm` |
+| `aac` | `audio/aac` |
+| `opus` | `audio/opus` |
+
+## Examples
+
+### Basic Generation
+
+Generate audio with only a text prompt (model defaults for all other parameters):
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves crashing on a beach"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+Specify an explicit audio length in seconds:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+Set a `seed` to get deterministic results across runs:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Full Control
+
+Combine all parameters for precise control over generation:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+### Python Client
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "Thunder and rain",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42,
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("thunder.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Parameter Tuning Guide
+
+### `guidance_scale`
+
+Controls how closely the generated audio follows the text prompt.
+
+| Range | Behaviour |
+|-------|-----------|
+| 3 - 5 | More creative / varied output |
+| 7 (default) | Balanced adherence |
+| 10+ | Strict adherence to the prompt |
+
+### `num_inference_steps`
+
+Controls the number of denoising steps in the diffusion process.
+
+| Steps | Quality | Speed | Use Case |
+|-------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 | Very Good | Medium | General purpose |
+| 150+ | Excellent | Slow | Final / critical audio |
+
+### `audio_length`
+
+Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
+
+### `negative_prompt`
+
+Describes characteristics to avoid. Common negative prompts include:
+
+- `"Low quality, distorted, noisy"`
+- `"Silence, static"`
+- `"Music"` (when generating sound effects only)
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid or missing parameters:
+
+```json
+{
+    "error": {
+        "message": "Audio generation model did not produce audio output.",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model mismatch:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+### 422 Unprocessable Entity
+
+Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
+
+```json
+{
+    "detail": [
+        {
+            "type": "literal_error",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            ...
+        }
+    ]
+}
+```
+
+## Troubleshooting
+
+### "Audio generation model did not produce audio output"
+
+The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
+
+### Server Not Responding
+
+```bash
+# Check if the server is healthy
+curl http://localhost:8000/health
+```
+
+### Audio Quality Issues
+
+- Increase `num_inference_steps` (e.g. 150).
+- Add a negative prompt: `"Low quality, distorted, noisy"`.
+- Increase `guidance_scale` for stronger prompt adherence.
+
+### Generation Timeout
+
+- Reduce `num_inference_steps`.
+- Reduce `audio_length`.
+- Check GPU memory with `nvidia-smi`.
+
+### Out of Memory
+
+- Lower `--gpu-memory-utilization` (e.g. 0.8).
+- Reduce `audio_length`.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni \
+    --uvicorn-log-level debug
+```

docs/user_guide/diffusion/quantization/gguf.md

@@ -136,9 +136,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --height 768 \
   --width 1360 \
   --seed 42 \
-  --cfg_scale 4.0 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 4 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 4 \
   --output outputs/flux2_klein_4b.png
 ```
 
@@ -152,9 +152,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --height 768 \
   --width 1360 \
   --seed 42 \
-  --cfg_scale 4.0 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 4 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 4 \
   --output outputs/flux2_klein_4b_gguf.png
 ```
 

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -0,0 +1,193 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
+
+This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Run all curl examples
+bash curl_examples.sh
+
+# Or execute directly
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+### Method 2: Using Python Client
+
+```bash
+cd examples/online_serving/stable_audio
+
+# Simple generation
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+
+# With custom duration
+python stable_audio_client.py \
+    --text "A dog barking" \
+    --audio_length 5.0
+
+# With all parameters
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+The Python client supports the following command-line arguments:
+
+- `--api_url`: API endpoint URL (default: `http://localhost:8000/v1/audio/generate`)
+- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
+- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
+- `--audio_start`: Audio start time in seconds (default: `0.0`)
+- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
+- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
+- `--num_inference_steps`: Number of inference steps (default: `100`)
+- `--seed`: Random seed for reproducibility (default: `None`)
+- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
+- `--output`: Output file path (default: `stable_audio_output.wav`)
+
+### Method 3: Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "The sound of ocean waves crashing on a beach",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+    },
+    timeout=300.0,
+)
+
+with open("ocean.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Request Format
+
+### Simple Generation
+
+```json
+{
+    "input": "The sound of ocean waves"
+}
+```
+
+### Generation with Parameters
+
+```json
+{
+    "input": "A piano playing a gentle melody",
+    "audio_length": 10.0,
+    "negative_prompt": "Low quality, distorted, noisy",
+    "guidance_scale": 8.0,
+    "num_inference_steps": 150,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Generation Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+## Tuning Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality (quick previews)
+   - 100 steps: Good balance (general purpose)
+   - 150+ steps: High quality, slower (final / critical audio)
+3. **Guidance Scale**:
+   - Lower (3 - 5): More creative / varied output
+   - Default (7): Good balance
+   - Higher (10+): Strict adherence to the prompt
+4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
+5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
+
+## File Description
+
+| File | Description |
+|------|-------------|
+| `curl_examples.sh` | Curl examples covering common use cases |
+| `stable_audio_client.py` | Python client with full CLI argument support |
+
+## Troubleshooting
+
+1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
+2. **Connection refused**: Make sure the server is running on the correct port.
+3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
+4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
+5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
+
+## Example materials
+
+??? abstract "stable_audio_client.py"
+    ``````py
+    --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
+    ``````
+??? abstract "curl_examples.sh"
+    ``````sh
+    --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
+    ``````

examples/offline_inference/bagel/end2end.py

@@ -1,8 +1,7 @@
 import argparse
 import os
-from typing import cast
 
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+from vllm_omni.inputs.data import OmniPromptType
 
 
 def parse_args():
@@ -58,6 +57,7 @@ def parse_args():
         choices=[1, 2, 3],
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
 
     args = parser.parse_args()
     return args
@@ -88,108 +88,80 @@ def main():
 
     from PIL import Image
 
-    if args.modality == "img2img":
-        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+    from vllm_omni.entrypoints.omni import Omni
 
-        print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})")
+    omni_kwargs = {}
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
 
-        client = OmniDiffusion(
-            model=model_name,
-            parallel_config={"cfg_parallel_size": args.cfg_parallel_size},
-        )
+    omni_kwargs.update(
+        {
+            "log_stats": args.log_stats,
+            "init_sleep_seconds": args.init_sleep_seconds,
+            "batch_timeout": args.batch_timeout,
+            "init_timeout": args.init_timeout,
+            "shm_threshold_bytes": args.shm_threshold_bytes,
+            "worker_backend": args.worker_backend,
+            "ray_address": args.ray_address,
+        }
+    )
 
-        if args.image_path:
-            if os.path.exists(args.image_path):
-                loaded_image = Image.open(args.image_path).convert("RGB")
-                prompts = [
-                    {
-                        "prompt": cast(str, p),
-                        "multi_modal_data": {"image": loaded_image},
-                    }
-                    for p in prompts
-                ]
-            else:
-                print(f"[Warning] Image path {args.image_path} does not exist.")
+    omni = Omni(model=model_name, **omni_kwargs)
 
-        result = client.generate(
-            prompts,
-            OmniDiffusionSamplingParams(
-                seed=52,
-                need_kv_receive=False,
-                num_inference_steps=args.steps,
-                extra_args={
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
-                },
-            ),
-        )
-
-        # Ensure result is a list for iteration
-        if not isinstance(result, list):
-            omni_outputs = [result]
-        else:
-            omni_outputs = result
-
-    else:
-        from vllm_omni.entrypoints.omni import Omni
-
-        omni_kwargs = {}
-        if args.stage_configs_path:
-            omni_kwargs["stage_configs_path"] = args.stage_configs_path
-
-        omni_kwargs.update(
-            {
-                "log_stats": args.log_stats,
-                "init_sleep_seconds": args.init_sleep_seconds,
-                "batch_timeout": args.batch_timeout,
-                "init_timeout": args.init_timeout,
-                "shm_threshold_bytes": args.shm_threshold_bytes,
-                "worker_backend": args.worker_backend,
-                "ray_address": args.ray_address,
+    formatted_prompts = []
+    for p in prompts:
+        if args.modality == "img2img":
+            if not args.image_path or not os.path.exists(args.image_path):
+                raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
+            loaded_image = Image.open(args.image_path).convert("RGB")
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
+            prompt_dict = {
+                "prompt": final_prompt_text,
+                "multi_modal_data": {"img2img": loaded_image},
+                "modalities": ["img2img"],
             }
-        )
-
-        omni = Omni(model=model_name, **omni_kwargs)
-
-        formatted_prompts = []
-        for p in args.prompts:
-            if args.modality == "img2text":
-                if args.image_path:
-                    loaded_image = Image.open(args.image_path).convert("RGB")
-                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                    prompt_dict = {
-                        "prompt": final_prompt_text,
-                        "multi_modal_data": {"image": loaded_image},
-                        "modalities": ["text"],
-                    }
-                    formatted_prompts.append(prompt_dict)
-            elif args.modality == "text2text":
-                final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
-                formatted_prompts.append(prompt_dict)
-            else:
-                # text2img
-                final_prompt_text = f"<|im_start|>{p}<|im_end|>"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
-                if args.negative_prompt is not None:
-                    prompt_dict["negative_prompt"] = args.negative_prompt
-                formatted_prompts.append(prompt_dict)
-
-        params_list = omni.default_sampling_params_list
-        if args.modality == "text2img":
-            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
-            if len(params_list) > 1:
-                diffusion_params = params_list[1]
-                diffusion_params.num_inference_steps = args.steps  # type: ignore
-                extra = {
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
+            if args.negative_prompt is not None:
+                prompt_dict["negative_prompt"] = args.negative_prompt
+            formatted_prompts.append(prompt_dict)
+        elif args.modality == "img2text":
+            if args.image_path:
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                prompt_dict = {
+                    "prompt": final_prompt_text,
+                    "multi_modal_data": {"image": loaded_image},
+                    "modalities": ["text"],
                 }
-                if args.negative_prompt is not None:
-                    extra["negative_prompt"] = args.negative_prompt
-                diffusion_params.extra_args = extra  # type: ignore
+                formatted_prompts.append(prompt_dict)
+        elif args.modality == "text2text":
+            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
+            formatted_prompts.append(prompt_dict)
+        else:
+            final_prompt_text = f"<|im_start|>{p}<|im_end|>"
+            prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
+            if args.negative_prompt is not None:
+                prompt_dict["negative_prompt"] = args.negative_prompt
+            formatted_prompts.append(prompt_dict)
 
-        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+    params_list = omni.default_sampling_params_list
+    if args.modality in ("text2img", "img2img"):
+        params_list[0].max_tokens = 1  # type: ignore
+        if len(params_list) > 1:
+            diffusion_params = params_list[1]
+            diffusion_params.num_inference_steps = args.steps  # type: ignore
+            diffusion_params.cfg_parallel_size = args.cfg_parallel_size  # type: ignore
+            if args.seed is not None:
+                diffusion_params.seed = args.seed  # type: ignore
+            extra = {
+                "cfg_text_scale": args.cfg_text_scale,
+                "cfg_img_scale": args.cfg_img_scale,
+            }
+            if args.negative_prompt is not None:
+                extra["negative_prompt"] = args.negative_prompt
+            diffusion_params.extra_args = extra  # type: ignore
+
+    omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
     for i, req_output in enumerate(omni_outputs):
         images = getattr(req_output, "images", None)

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -29,8 +29,8 @@ Accelerated with cache_dit:
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --cache_backend cache_dit \
-        --cache_dit_enable_taylorseer
+        --cache-backend cache_dit \
+        --cache-dit-enable-taylorseer
 
 Layered RGBA output:
     python image_edit.py \
@@ -66,38 +66,38 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
     parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
     parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
-    parser.add_argument("--negative_prompt", type=str, default=None)
+    parser.add_argument("--negative-prompt", type=str, default=None)
     parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--cfg_scale", type=float, default=4.0)
-    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--cfg-scale", type=float, default=4.0)
+    parser.add_argument("--guidance-scale", type=float, default=1.0)
     parser.add_argument("--output", type=str, default="output.png")
-    parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
-    parser.add_argument("--num_inference_steps", type=int, default=50)
-    parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
-    parser.add_argument("--ulysses_degree", type=int, default=1)
-    parser.add_argument("--ring_degree", type=int, default=1)
-    parser.add_argument("--tensor_parallel_size", type=int, default=1)
+    parser.add_argument("--num-outputs-per-prompt", type=int, default=1)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
     parser.add_argument("--layers", type=int, default=4)
     parser.add_argument("--resolution", type=int, default=640)
     parser.add_argument("--color-format", type=str, default="RGB")
 
     # Acceleration + Optimization Options
-    parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
-    parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
-    parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
-    parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
-    parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
-    parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
-    parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
+    parser.add_argument("--cache-dit-fn-compute-blocks", type=int, default=1)
+    parser.add_argument("--cache-dit-bn-compute-blocks", type=int, default=0)
+    parser.add_argument("--cache-dit-max-warmup-steps", type=int, default=4)
+    parser.add_argument("--cache-dit-residual-diff-threshold", type=float, default=0.24)
+    parser.add_argument("--cache-dit-max-continuous-cached-steps", type=int, default=3)
+    parser.add_argument("--cache-dit-enable-taylorseer", action="store_true", default=False)
+    parser.add_argument("--cache-dit-taylorseer-order", type=int, default=1)
     parser.add_argument(
-        "--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
+        "--cache-dit-scm-steps-mask-policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
     )
-    parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
-    parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
-    parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
-    parser.add_argument("--enforce_eager", action="store_true")
-    parser.add_argument("--vae_use_slicing", action="store_true")
-    parser.add_argument("--vae_use_tiling", action="store_true")
+    parser.add_argument("--cache-dit-scm-steps-policy", type=str, default="dynamic", choices=["dynamic", "static"])
+    parser.add_argument("--tea-cache-rel-l1-thresh", type=float, default=0.2)
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
+    parser.add_argument("--enforce-eager", action="store_true")
+    parser.add_argument("--vae-use-slicing", action="store_true")
+    parser.add_argument("--vae-use-tiling", action="store_true")
     parser.add_argument("--enable-cpu-offload", action="store_true")
 
     return parser.parse_args()

examples/offline_inference/custom_pipeline/image_to_image/run.sh

@@ -3,4 +3,4 @@ python image_edit.py \
         --image cherry_blossom.jpg \
         --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
         --output output_image_edit.png \
-        --num_inference_steps 10
+        --num-inference-steps 10

examples/offline_inference/image_to_video/image_to_video.py

@@ -22,7 +22,7 @@ Usage:
     python image_to_video.py --model /path/to/LTX-2 \
         --model-class-name LTX2ImageToVideoPipeline \
         --image input.jpg --prompt "A cinematic dolly shot of a boat" \
-        --num-frames 121 --num_inference_steps 40 --guidance_scale 4.0 \
+        --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 """
 

examples/online_serving/chart-helm/.helmignore

@@ -0,0 +1,6 @@
+*.png
+.git/
+ct.yaml
+lintconf.yaml
+values.schema.json
+/workflows

examples/online_serving/chart-helm/Chart.yaml

@@ -0,0 +1,15 @@
+apiVersion: v2
+name: chart-vllm-omni
+description: A Helm chart for deploying vLLM-Omni on Kubernetes for omni-modality model serving
+
+# Application chart that can be packaged and deployed
+type: application
+
+# Chart version - increment on each change to the chart
+version: 0.1.0
+
+# vllm-omni application version — keep in sync with image.tag in values.yaml
+appVersion: "0.16.0"
+
+maintainers:
+  - name: vllm-omni-team

examples/online_serving/chart-helm/README.md

@@ -0,0 +1,144 @@
+# vLLM-Omni Helm Chart
+
+Helm chart for deploying [vLLM-Omni](https://github.com/vllm-project/vllm-omni) on Kubernetes. vLLM-Omni extends vLLM with omni-modality model serving, supporting text-to-image, multimodal chat, text-to-speech, and more.
+
+## Prerequisites
+
+- Kubernetes 1.24+
+- Helm 3.x
+- NVIDIA GPU nodes with [NVIDIA Device Plugin](https://github.com/NVIDIA/k8s-device-plugin)
+
+## Quick Start
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Tongyi-MAI/Z-Image-Turbo
+```
+
+## Configuration
+
+### Model Selection
+
+Set the `model` value to any supported HuggingFace model ID:
+
+| Model | Type | GPUs | Notes |
+|-------|------|------|-------|
+| `Tongyi-MAI/Z-Image-Turbo` | text-to-image | 1 | Small, fast (default) |
+| `stabilityai/stable-diffusion-3.5-medium` | text-to-image | 1 | ~6GB VRAM |
+| `Qwen/Qwen-Image` | text-to-image | 1 | Large, ~40GB+ VRAM |
+| `Qwen/Qwen2.5-Omni-7B` | multimodal | 2 | Text + audio + image + video |
+| `Qwen/Qwen3-Omni-7B-Chat` | multimodal | 2 | Latest omni model |
+| `Qwen/Qwen3-TTS` | text-to-speech | 1 | TTS |
+
+### HuggingFace Token
+
+For gated models that require authentication:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen2.5-Omni-7B \
+  --set hfToken=hf_xxxxx \
+  --set resources.requests."nvidia\.com/gpu"=2 \
+  --set resources.limits."nvidia\.com/gpu"=2
+```
+
+### Omni-Specific Flags
+
+Enable VAE memory optimizations for diffusion models:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen-Image \
+  --set omniArgs.vaeUseSlicing=true \
+  --set omniArgs.vaeUseTiling=true
+```
+
+Enable CPU offloading:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen-Image \
+  --set omniArgs.enableCpuOffload=true
+```
+
+Pass additional raw CLI flags:
+
+```bash
+helm install my-release ./chart-helm \
+  --set model=Qwen/Qwen-Image \
+  --set omniArgs.extraArgs[0]="--enable-layerwise-offload"
+```
+
+### Model Cache
+
+By default, a PersistentVolumeClaim is created for the HuggingFace model cache to avoid re-downloading models on pod restarts:
+
+```yaml
+modelCache:
+  enabled: true
+  storageSize: "50Gi"
+  storageClassName: ""
+```
+
+To use an ephemeral volume instead:
+
+```bash
+helm install my-release ./chart-helm \
+  --set modelCache.enabled=false
+```
+
+### Custom Command Override
+
+To fully override the container command:
+
+```bash
+helm install my-release ./chart-helm \
+  --set image.command[0]=vllm \
+  --set image.command[1]=serve \
+  --set image.command[2]=my-model \
+  --set image.command[3]=--omni \
+  --set image.command[4]=--host \
+  --set image.command[5]=0.0.0.0
+```
+
+## API Endpoints
+
+Once deployed, vLLM-Omni exposes the following OpenAI-compatible endpoints:
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/health` | GET | Health check |
+| `/v1/models` | GET | List available models |
+| `/v1/chat/completions` | POST | Chat completions (text/multimodal) |
+| `/v1/images/generations` | POST | Image generation |
+| `/v1/images/edits` | POST | Image editing |
+| `/v1/audio/speech` | POST | Text-to-speech |
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `Chart.yaml` | Chart metadata (name, version, maintainers) |
+| `values.yaml` | Default configuration values |
+| `values.schema.json` | JSON schema for validating values |
+| `templates/_helpers.tpl` | Helper templates for common configurations |
+| `templates/deployment.yaml` | Kubernetes Deployment |
+| `templates/service.yaml` | Kubernetes Service (ClusterIP) |
+| `templates/secrets.yaml` | Secrets (generic + HuggingFace token) |
+| `templates/pvc.yaml` | PersistentVolumeClaim for model cache |
+| `templates/configmap.yaml` | Optional ConfigMap |
+| `templates/hpa.yaml` | HorizontalPodAutoscaler |
+| `templates/poddisruptionbudget.yaml` | PodDisruptionBudget |
+| `templates/custom-objects.yaml` | Custom Kubernetes objects |
+
+## Running Tests
+
+This chart includes unit tests using [helm-unittest](https://github.com/helm-unittest/helm-unittest). Install the plugin and run tests:
+
+```bash
+# Install plugin
+helm plugin install https://github.com/helm-unittest/helm-unittest
+
+# Run tests
+helm unittest .
+```

examples/online_serving/chart-helm/ct.yaml

@@ -0,0 +1,3 @@
+chart-dirs:
+  - charts
+validate-maintainers: false

examples/online_serving/chart-helm/lintconf.yaml

@@ -0,0 +1,42 @@
+---
+rules:
+  braces:
+    min-spaces-inside: 0
+    max-spaces-inside: 0
+    min-spaces-inside-empty: -1
+    max-spaces-inside-empty: -1
+  brackets:
+    min-spaces-inside: 0
+    max-spaces-inside: 0
+    min-spaces-inside-empty: -1
+    max-spaces-inside-empty: -1
+  colons:
+    max-spaces-before: 0
+    max-spaces-after: 1
+  commas:
+    max-spaces-before: 0
+    min-spaces-after: 1
+    max-spaces-after: 1
+  comments:
+    require-starting-space: true
+    min-spaces-from-content: 2
+  document-end: disable
+  document-start: disable
+  empty-lines:
+    max: 2
+    max-start: 0
+    max-end: 0
+  hyphens:
+    max-spaces-after: 1
+  indentation:
+    spaces: consistent
+    indent-sequences: whatever
+    check-multi-line-strings: false
+  key-duplicates: enable
+  line-length: disable
+  new-line-at-end-of-file: disable
+  new-lines:
+    type: unix
+  trailing-spaces: enable
+  truthy:
+    level: warning

examples/online_serving/chart-helm/templates/_helpers.tpl

@@ -0,0 +1,203 @@
+{{/*
+Define the vllm-omni serve command from model + omniArgs.
+If image.command is set, uses that as a full override.
+*/}}
+{{- define "chart.omni-command" -}}
+{{- if .Values.image.command }}
+{{-   toYaml .Values.image.command }}
+{{- else }}
+- "vllm"
+- "serve"
+- {{ .Values.model | quote }}
+- "--omni"
+- "--host"
+- "0.0.0.0"
+- "--port"
+- {{ include "chart.container-port" . | quote }}
+{{- if .Values.omniArgs.vaeUseSlicing }}
+- "--vae-use-slicing"
+{{- end }}
+{{- if .Values.omniArgs.vaeUseTiling }}
+- "--vae-use-tiling"
+{{- end }}
+{{- if .Values.omniArgs.enableCpuOffload }}
+- "--enable-cpu-offload"
+{{- end }}
+{{- if .Values.omniArgs.numGpus }}
+- "--num-gpus"
+- {{ .Values.omniArgs.numGpus | quote }}
+{{- end }}
+{{- if .Values.omniArgs.stageConfigsPath }}
+- "--stage-configs-path"
+- {{ .Values.omniArgs.stageConfigsPath | quote }}
+{{- end }}
+{{- if and .Values.omniArgs.cacheBackend (ne .Values.omniArgs.cacheBackend "none") }}
+- "--cache-backend"
+- {{ .Values.omniArgs.cacheBackend | quote }}
+{{- end }}
+{{- if .Values.omniArgs.defaultSamplingParams }}
+- "--default-sampling-params"
+- {{ .Values.omniArgs.defaultSamplingParams | quote }}
+{{- end }}
+{{- if and .Values.omniArgs.workerBackend (ne .Values.omniArgs.workerBackend "multi_process") }}
+- "--worker-backend"
+- {{ .Values.omniArgs.workerBackend | quote }}
+{{- end }}
+{{- range .Values.omniArgs.extraArgs }}
+- {{ . | quote }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+{{/*
+Define HuggingFace environment variables
+*/}}
+{{- define "chart.hf-env" -}}
+- name: HF_HOME
+  value: "/cache/huggingface"
+- name: HOME
+  value: "/cache"
+{{- if .Values.hfToken }}
+- name: HF_TOKEN
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Release.Name }}-hf-token
+      key: token
+{{- end }}
+{{- end }}
+
+{{/*
+Define ports for the pods
+*/}}
+{{- define "chart.container-port" -}}
+{{-  default "8000" .Values.containerPort }}
+{{- end }}
+
+{{/*
+Define service name
+*/}}
+{{- define "chart.service-name" -}}
+{{-  if .Values.serviceName -}}
+{{ .Values.serviceName | lower | trim }}
+{{-  else -}}
+{{ printf "%s-service" .Release.Name }}
+{{-  end -}}
+{{- end }}
+
+{{/*
+Define service port
+*/}}
+{{- define "chart.service-port" -}}
+{{-  if .Values.servicePort }}
+{{-    .Values.servicePort }}
+{{-  else }}
+{{-    include "chart.container-port" . }}
+{{-  end }}
+{{- end }}
+
+{{/*
+Define service port name
+*/}}
+{{- define "chart.service-port-name" -}}
+"service-port"
+{{- end }}
+
+{{/*
+Define container port name
+*/}}
+{{- define "chart.container-port-name" -}}
+"container-port"
+{{- end }}
+
+{{/*
+Define deployment strategy
+*/}}
+{{- define "chart.strategy" -}}
+strategy:
+{{-   if not .Values.deploymentStrategy }}
+  type: Recreate
+{{-   else }}
+{{      toYaml .Values.deploymentStrategy | indent 2 }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define additional ports
+*/}}
+{{- define "chart.extraPorts" }}
+{{-   with .Values.extraPorts }}
+{{      toYaml . }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define chart external ConfigMaps and Secrets
+*/}}
+{{- define "chart.externalConfigs" -}}
+{{-   with .Values.externalConfigs -}}
+{{      toYaml . }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define startup, liveness and readiness probes
+*/}}
+{{- define "chart.probes" -}}
+{{-   if .Values.startupProbe  }}
+startupProbe:
+{{-     with .Values.startupProbe }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{-   if .Values.readinessProbe  }}
+readinessProbe:
+{{-     with .Values.readinessProbe }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{-   if .Values.livenessProbe  }}
+livenessProbe:
+{{-     with .Values.livenessProbe }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define resources
+*/}}
+{{- define "chart.resources" -}}
+requests:
+  memory: {{ required "Value 'resources.requests.memory' must be defined !" .Values.resources.requests.memory | quote }}
+  cpu: {{ required "Value 'resources.requests.cpu' must be defined !" .Values.resources.requests.cpu | quote }}
+  {{- if and (gt (int (index .Values.resources.requests "nvidia.com/gpu")) 0) (gt (int (index .Values.resources.limits "nvidia.com/gpu")) 0) }}
+  nvidia.com/gpu: {{ required "Value 'resources.requests.nvidia.com/gpu' must be defined !" (index .Values.resources.requests "nvidia.com/gpu") }}
+  {{- end }}
+limits:
+  memory: {{ required "Value 'resources.limits.memory' must be defined !" .Values.resources.limits.memory | quote }}
+  cpu: {{ required "Value 'resources.limits.cpu' must be defined !" .Values.resources.limits.cpu | quote }}
+  {{- if and (gt (int (index .Values.resources.requests "nvidia.com/gpu")) 0) (gt (int (index .Values.resources.limits "nvidia.com/gpu")) 0) }}
+  nvidia.com/gpu: {{ required "Value 'resources.limits.nvidia.com/gpu' must be defined !" (index .Values.resources.limits "nvidia.com/gpu") }}
+  {{- end }}
+{{- end }}
+
+{{/*
+Define user for the main container
+*/}}
+{{- define "chart.user" }}
+{{-   if .Values.image.runAsUser  }}
+runAsUser:
+{{-     with .Values.image.runAsUser }}
+{{-       toYaml . | nindent 2 }}
+{{-     end }}
+{{-   end }}
+{{- end }}
+
+{{/*
+Define chart labels
+*/}}
+{{- define "chart.labels" -}}
+{{-   with .Values.labels -}}
+{{      toYaml . }}
+{{-   end }}
+{{- end }}

examples/online_serving/chart-helm/templates/configmap.yaml

@@ -0,0 +1,11 @@
+{{- if .Values.configs -}}
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: "{{ .Release.Name }}-configs"
+  namespace: {{ .Release.Namespace }}
+data:
+  {{- with .Values.configs }}
+  {{- toYaml . | nindent 2 }}
+  {{- end }}
+{{- end -}}

examples/online_serving/chart-helm/templates/custom-objects.yaml

@@ -0,0 +1,6 @@
+{{- if .Values.customObjects }}
+{{- range .Values.customObjects }}
+{{- tpl (. | toYaml) $ }}
+---
+{{- end }}
+{{- end }}

examples/online_serving/chart-helm/templates/deployment.yaml

@@ -0,0 +1,106 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: "{{ .Release.Name }}-deployment-vllm-omni"
+  namespace: {{ .Release.Namespace }}
+  labels:
+  {{- include "chart.labels" . | nindent 4 }}
+spec:
+  {{- if not .Values.autoscaling.enabled }}
+  replicas: {{ .Values.replicaCount }}
+  {{- end }}
+  {{- include "chart.strategy" . | nindent 2 }}
+  selector:
+    matchLabels:
+    {{- include "chart.labels" . | nindent 6 }}
+  progressDeadlineSeconds: 1800
+  template:
+    metadata:
+      labels:
+      {{- include "chart.labels" . | nindent 8 }}
+    spec:
+      containers:
+        - name: "vllm-omni"
+          image: "{{ required "Required value 'image.repository' must be defined !" .Values.image.repository }}:{{ required "Required value 'image.tag' must be defined !" .Values.image.tag }}"
+          command:
+            {{- include "chart.omni-command" . | nindent 12 }}
+          securityContext:
+            {{- if .Values.image.securityContext }}
+              {{- with .Values.image.securityContext }}
+              {{- toYaml . | nindent 12 }}
+              {{- end }}
+            {{- else }}
+            runAsNonRoot: false
+              {{- include "chart.user" . | indent 12 }}
+            {{- end }}
+          imagePullPolicy: IfNotPresent
+          env:
+            {{- include "chart.hf-env" . | nindent 12 }}
+            {{- if .Values.image.env }}
+            {{- toYaml .Values.image.env | nindent 12 }}
+            {{- end }}
+          {{- if or .Values.externalConfigs .Values.configs .Values.secrets }}
+          envFrom:
+            {{- if .Values.configs }}
+            - configMapRef:
+                name: "{{ .Release.Name }}-configs"
+            {{- end }}
+            {{- if .Values.secrets }}
+            - secretRef:
+                name: "{{ .Release.Name }}-secrets"
+            {{- end }}
+            {{- include "chart.externalConfigs" . | nindent 12 }}
+          {{- end }}
+          ports:
+            - name: {{ include "chart.container-port-name" . }}
+              containerPort: {{ include "chart.container-port" . }}
+            {{- include "chart.extraPorts" . | nindent 12 }}
+          {{- include "chart.probes" . | indent 10 }}
+          resources: {{- include "chart.resources" . | nindent 12 }}
+          volumeMounts:
+          - name: shm
+            mountPath: /dev/shm
+          - name: model-cache
+            mountPath: /cache
+
+        {{- with .Values.extraContainers }}
+        {{ toYaml . | nindent 8 }}
+        {{- end }}
+
+      {{- if .Values.extraInit.initContainers }}
+      initContainers:
+        {{- toYaml .Values.extraInit.initContainers | nindent 8 }}
+      {{- end }}
+
+      volumes:
+        - name: shm
+          emptyDir:
+            medium: Memory
+            sizeLimit: {{ .Values.shmSize }}
+        - name: model-cache
+          {{- if .Values.modelCache.enabled }}
+          persistentVolumeClaim:
+            claimName: {{ .Release.Name }}-model-cache
+          {{- else }}
+          emptyDir: {}
+          {{- end }}
+
+      {{- with .Values.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- if .Values.gpuModels }}
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                - key: nvidia.com/gpu.product
+                  operator: In
+                  values:
+                    {{- toYaml .Values.gpuModels | nindent 20 }}
+      {{- end }}

examples/online_serving/chart-helm/templates/hpa.yaml

@@ -0,0 +1,31 @@
+{{- if .Values.autoscaling.enabled }}
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: "{{ .Release.Name }}-hpa"
+  namespace: {{ .Release.Namespace }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: "{{ .Release.Name }}-deployment-vllm-omni"
+  minReplicas: {{ .Values.autoscaling.minReplicas }}
+  maxReplicas: {{ .Values.autoscaling.maxReplicas }}
+  metrics:
+    {{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
+    {{- end }}
+    {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
+    - type: Resource
+      resource:
+        name: memory
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
+    {{- end }}
+{{- end }}

examples/online_serving/chart-helm/templates/ingress.yaml

@@ -0,0 +1,32 @@
+{{- if .Values.ingress.enabled }}
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  name: "{{ .Release.Name }}-ingress"
+  namespace: {{ .Release.Namespace }}
+  labels:
+  {{- include "chart.labels" . | nindent 4 }}
+  {{- with .Values.ingress.annotations }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+spec:
+  {{- if .Values.ingress.ingressClassName }}
+  ingressClassName: {{ .Values.ingress.ingressClassName }}
+  {{- end }}
+  {{- if .Values.ingress.tls }}
+  tls:
+    {{- toYaml .Values.ingress.tls | nindent 4 }}
+  {{- end }}
+  rules:
+    - host: {{ .Values.ingress.host }}
+      http:
+        paths:
+          - path: /
+            pathType: Prefix
+            backend:
+              service:
+                name: "{{ include "chart.service-name" . }}"
+                port:
+                  number: {{ include "chart.service-port" . }}
+{{- end }}

examples/online_serving/chart-helm/templates/poddisruptionbudget.yaml

@@ -0,0 +1,12 @@
+{{- if gt (int .Values.replicaCount) 1 }}
+apiVersion: policy/v1
+kind: PodDisruptionBudget
+metadata:
+  name: "{{ .Release.Name }}-pdb"
+  namespace: {{ .Release.Namespace }}
+spec:
+  maxUnavailable: {{ default 1 .Values.maxUnavailablePodDisruptionBudget }}
+  selector:
+    matchLabels:
+    {{- include "chart.labels" . | nindent 6 }}
+{{- end }}

examples/online_serving/chart-helm/templates/pvc.yaml

@@ -0,0 +1,16 @@
+{{- if .Values.modelCache.enabled }}
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: "{{ .Release.Name }}-model-cache"
+  namespace: {{ .Release.Namespace }}
+spec:
+  accessModes:
+    {{- toYaml .Values.modelCache.accessModes | nindent 4 }}
+  {{- if .Values.modelCache.storageClassName }}
+  storageClassName: {{ .Values.modelCache.storageClassName }}
+  {{- end }}
+  resources:
+    requests:
+      storage: {{ .Values.modelCache.storageSize }}
+{{- end }}

examples/online_serving/chart-helm/templates/secrets.yaml

@@ -0,0 +1,23 @@
+{{- if .Values.secrets }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: "{{ .Release.Name }}-secrets"
+  namespace: {{ .Release.Namespace }}
+type: Opaque
+data:
+  {{- range $key, $val := .Values.secrets }}
+  {{ $key }}: {{ $val | b64enc | quote }}
+  {{- end }}
+---
+{{- end }}
+{{- if .Values.hfToken }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: "{{ .Release.Name }}-hf-token"
+  namespace: {{ .Release.Namespace }}
+type: Opaque
+data:
+  token: {{ .Values.hfToken | b64enc | quote }}
+{{- end }}

examples/online_serving/chart-helm/templates/service.yaml

@@ -0,0 +1,14 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: "{{ include "chart.service-name" . }}"
+  namespace: {{ .Release.Namespace }}
+spec:
+  type: ClusterIP
+  ports:
+    - name: {{ include "chart.service-port-name" . }}
+      port: {{ include "chart.service-port" . }}
+      targetPort: {{ include "chart.container-port-name" . }}
+      protocol: TCP
+  selector:
+  {{- include "chart.labels" . | nindent 4 }}

examples/online_serving/chart-helm/tests/deployment_test.yaml

@@ -0,0 +1,181 @@
+suite: test deployment
+templates:
+  - deployment.yaml
+tests:
+  - it: should create deployment with default omni command
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Deployment
+      - equal:
+          path: spec.template.spec.containers[0].name
+          value: vllm-omni
+      - equal:
+          path: spec.template.spec.containers[0].command[0]
+          value: vllm
+      - equal:
+          path: spec.template.spec.containers[0].command[1]
+          value: serve
+      - equal:
+          path: spec.template.spec.containers[0].command[2]
+          value: Tongyi-MAI/Z-Image-Turbo
+      - equal:
+          path: spec.template.spec.containers[0].command[3]
+          value: "--omni"
+      - equal:
+          path: spec.template.spec.containers[0].command[4]
+          value: "--host"
+      - equal:
+          path: spec.template.spec.containers[0].command[5]
+          value: "0.0.0.0"
+      - equal:
+          path: spec.template.spec.containers[0].command[6]
+          value: "--port"
+      - equal:
+          path: spec.template.spec.containers[0].command[7]
+          value: "8000"
+
+  - it: should use custom model when set
+    set:
+      model: "Qwen/Qwen2.5-Omni-7B"
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command[2]
+          value: Qwen/Qwen2.5-Omni-7B
+
+  - it: should include omniArgs flags in command
+    set:
+      omniArgs:
+        vaeUseSlicing: true
+        vaeUseTiling: true
+        enableCpuOffload: true
+        numGpus: 2
+        cacheBackend: "tea_cache"
+        workerBackend: "ray"
+        extraArgs:
+          - "--enable-layerwise-offload"
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--vae-use-slicing"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--vae-use-tiling"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--enable-cpu-offload"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--num-gpus"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--cache-backend"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--worker-backend"
+      - contains:
+          path: spec.template.spec.containers[0].command
+          content: "--enable-layerwise-offload"
+
+  - it: should use full command override when image.command is set
+    set:
+      image:
+        command:
+          - "custom-binary"
+          - "--flag"
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command[0]
+          value: custom-binary
+      - equal:
+          path: spec.template.spec.containers[0].command[1]
+          value: "--flag"
+
+  - it: should mount shm and model-cache volumes
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].volumeMounts
+          content:
+            name: shm
+            mountPath: /dev/shm
+      - contains:
+          path: spec.template.spec.containers[0].volumeMounts
+          content:
+            name: model-cache
+            mountPath: /cache
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: shm
+            emptyDir:
+              medium: Memory
+              sizeLimit: 8Gi
+
+  - it: should use PVC for model-cache when modelCache is enabled
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: model-cache
+            persistentVolumeClaim:
+              claimName: RELEASE-NAME-model-cache
+
+  - it: should use emptyDir for model-cache when modelCache is disabled
+    set:
+      modelCache:
+        enabled: false
+        storageSize: "50Gi"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: model-cache
+            emptyDir: {}
+
+  - it: should set HF_HOME environment variable
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: HF_HOME
+            value: /cache/huggingface
+
+  - it: should set HF_TOKEN env when hfToken is provided
+    set:
+      hfToken: "hf_test_token_123"
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: HF_TOKEN
+            valueFrom:
+              secretKeyRef:
+                name: RELEASE-NAME-hf-token
+                key: token
+
+  - it: should create custom init containers when specified
+    set:
+      extraInit:
+        initContainers:
+          - name: my-init
+            image: busybox:latest
+            command: ["echo", "init"]
+    asserts:
+      - lengthEqual:
+          path: spec.template.spec.initContainers
+          count: 1
+      - equal:
+          path: spec.template.spec.initContainers[0].name
+          value: my-init
+      - equal:
+          path: spec.template.spec.initContainers[0].image
+          value: busybox:latest

examples/online_serving/chart-helm/tests/ingress_test.yaml

@@ -0,0 +1,86 @@
+suite: test ingress
+templates:
+  - ingress.yaml
+tests:
+  - it: should not create Ingress when disabled
+    set:
+      ingress:
+        enabled: false
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should create Ingress when enabled
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Ingress
+      - equal:
+          path: spec.rules[0].host
+          value: vllm-omni.example.com
+      - equal:
+          path: spec.rules[0].http.paths[0].path
+          value: /
+      - equal:
+          path: spec.rules[0].http.paths[0].pathType
+          value: Prefix
+      - equal:
+          path: spec.rules[0].http.paths[0].backend.service.name
+          value: RELEASE-NAME-service
+      - equal:
+          path: spec.rules[0].http.paths[0].backend.service.port.number
+          value: 80
+
+  - it: should set ingressClassName when specified
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        ingressClassName: "nginx"
+    asserts:
+      - equal:
+          path: spec.ingressClassName
+          value: nginx
+
+  - it: should not set ingressClassName when empty
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        ingressClassName: ""
+    asserts:
+      - notExists:
+          path: spec.ingressClassName
+
+  - it: should include annotations when specified
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        annotations:
+          nginx.ingress.kubernetes.io/proxy-body-size: "100m"
+    asserts:
+      - isNotEmpty:
+          path: metadata.annotations
+
+  - it: should include TLS configuration when specified
+    set:
+      ingress:
+        enabled: true
+        host: "vllm-omni.example.com"
+        tls:
+          - secretName: vllm-omni-tls
+            hosts:
+              - vllm-omni.example.com
+    asserts:
+      - equal:
+          path: spec.tls[0].secretName
+          value: vllm-omni-tls
+      - equal:
+          path: spec.tls[0].hosts[0]
+          value: vllm-omni.example.com

examples/online_serving/chart-helm/tests/pvc_test.yaml

@@ -0,0 +1,73 @@
+suite: test pvc
+templates:
+  - pvc.yaml
+tests:
+  - it: should create PVC when modelCache is enabled
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        storageClassName: ""
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: PersistentVolumeClaim
+      - equal:
+          path: spec.accessModes[0]
+          value: ReadWriteOnce
+      - equal:
+          path: spec.resources.requests.storage
+          value: 50Gi
+
+  - it: should not create PVC when modelCache is disabled
+    set:
+      modelCache:
+        enabled: false
+        storageSize: "50Gi"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should use custom storage size
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "100Gi"
+        storageClassName: ""
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - equal:
+          path: spec.resources.requests.storage
+          value: 100Gi
+
+  - it: should set storageClassName when specified
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        storageClassName: "fast-ssd"
+        accessModes:
+          - ReadWriteOnce
+    asserts:
+      - equal:
+          path: spec.storageClassName
+          value: fast-ssd
+
+  - it: should support ReadWriteMany access mode
+    set:
+      modelCache:
+        enabled: true
+        storageSize: "50Gi"
+        storageClassName: ""
+        accessModes:
+          - ReadWriteMany
+    asserts:
+      - equal:
+          path: spec.accessModes[0]
+          value: ReadWriteMany

examples/online_serving/chart-helm/tests/secrets_test.yaml

@@ -0,0 +1,48 @@
+suite: test secrets
+templates:
+  - secrets.yaml
+tests:
+  - it: should create HF token secret when hfToken is provided
+    set:
+      hfToken: "hf_test_token_123"
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Secret
+      - equal:
+          path: metadata.name
+          value: RELEASE-NAME-hf-token
+      - exists:
+          path: data.token
+
+  - it: should not create any secrets when hfToken is empty and no secrets defined
+    set:
+      hfToken: ""
+      secrets: {}
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should create both generic and HF token secrets
+    set:
+      hfToken: "hf_test_token_123"
+      secrets:
+        mykey: myvalue
+    asserts:
+      - hasDocuments:
+          count: 2
+
+  - it: should create only generic secrets when no hfToken
+    set:
+      hfToken: ""
+      secrets:
+        mykey: myvalue
+    asserts:
+      - hasDocuments:
+          count: 1
+      - isKind:
+          of: Secret
+      - equal:
+          path: metadata.name
+          value: RELEASE-NAME-secrets

examples/online_serving/chart-helm/values.yaml

@@ -0,0 +1,197 @@
+# -- Default values for chart vllm-omni
+# -- Declare variables to be passed into your templates.
+
+# -- HuggingFace model ID to serve
+model: "Tongyi-MAI/Z-Image-Turbo"
+
+# -- Image configuration
+image:
+  # -- Image repository
+  repository: "vllm/vllm-omni"
+  # -- Image tag
+  tag: "v0.16.0"
+  # -- Override the container command entirely. If empty, the chart constructs
+  # -- it automatically from model + omniArgs.
+  command: []
+  # -- Optional environment variables for the container
+  env: []
+  # -- Security context override
+  securityContext: {}
+
+# -- Container port
+containerPort: 8000
+# -- Service name (auto-generated from release name if empty)
+serviceName:
+# -- Service port
+servicePort: 80
+# -- Additional ports configuration
+extraPorts: []
+
+# -- Number of replicas
+replicaCount: 1
+
+# -- Deployment strategy configuration
+deploymentStrategy: {}
+
+# -- Resource configuration
+resources:
+  requests:
+    # -- Number of CPUs
+    cpu: 4
+    # -- CPU memory
+    memory: 24Gi
+    # -- Number of GPUs
+    nvidia.com/gpu: 1
+  limits:
+    # -- Number of CPUs
+    cpu: 4
+    # -- CPU memory
+    memory: 24Gi
+    # -- Number of GPUs
+    nvidia.com/gpu: 1
+
+# -- GPU model types for node affinity scheduling (optional)
+# -- Leave empty to schedule on any GPU node. Set to restrict to specific GPU types.
+# -- Example: ["NVIDIA-A100-SXM4-40GB"]
+gpuModels: []
+
+# -- vllm-omni specific CLI arguments
+omniArgs:
+  # -- Additional raw CLI flags appended to the serve command (list of strings)
+  extraArgs: []
+  # -- Enable VAE slicing for memory optimization
+  vaeUseSlicing: false
+  # -- Enable VAE tiling for memory optimization
+  vaeUseTiling: false
+  # -- Enable CPU offloading for diffusion models
+  enableCpuOffload: false
+  # -- Number of GPUs for diffusion inference (empty = auto)
+  numGpus:
+  # -- Path to stage configs file (empty = auto-detected from model)
+  stageConfigsPath:
+  # -- Cache backend: none, tea_cache, or cache_dit
+  cacheBackend: "none"
+  # -- Default sampling params as a JSON string
+  defaultSamplingParams:
+  # -- Worker backend: multi_process or ray
+  workerBackend: "multi_process"
+
+# -- HuggingFace token for gated models (optional)
+hfToken: ""
+
+# -- Shared memory size for PyTorch multiprocessing
+shmSize: "8Gi"
+
+# -- Model cache configuration (HuggingFace downloads)
+modelCache:
+  # -- Use a PersistentVolumeClaim for model cache (recommended for production)
+  enabled: true
+  # -- Storage size for the PVC
+  storageSize: "50Gi"
+  # -- Storage class name (empty = cluster default)
+  storageClassName: ""
+  # -- Access modes
+  accessModes:
+    - ReadWriteOnce
+
+# -- Autoscaling configuration
+autoscaling:
+  # -- Enable autoscaling
+  enabled: false
+  # -- Minimum replicas
+  minReplicas: 1
+  # -- Maximum replicas
+  maxReplicas: 10
+  # -- Target CPU utilization for autoscaling
+  targetCPUUtilizationPercentage: 80
+  # targetMemoryUtilizationPercentage: 80
+
+# -- ConfigMap data (key-value pairs injected as environment variables)
+configs: {}
+
+# -- Secrets data (key-value pairs, base64-encoded automatically)
+secrets: {}
+
+# -- External ConfigMaps/Secrets references
+externalConfigs: []
+
+# -- Custom Kubernetes objects
+customObjects: []
+
+# -- PodDisruptionBudget max unavailable
+maxUnavailablePodDisruptionBudget: ""
+
+# -- Additional init containers
+extraInit:
+  initContainers: []
+
+# -- Additional sidecar containers
+extraContainers: []
+
+# -- Startup probe configuration
+# -- Protects slow-starting containers. Liveness and readiness probes
+# -- do not start until the startup probe succeeds.
+# -- failureThreshold * periodSeconds = max startup time (40 * 30 = 1200s = 20 min)
+startupProbe:
+  # -- HTTP check configuration
+  httpGet:
+    # -- Path to check
+    path: /health
+    # -- Port to check (uses named port to stay in sync with containerPort)
+    port: "container-port"
+  # -- Failures before killing the container
+  failureThreshold: 40
+  # -- How often to perform the probe
+  periodSeconds: 30
+
+# -- Readiness probe configuration
+readinessProbe:
+  # -- How often to perform the probe
+  periodSeconds: 10
+  # -- Failures before marking unready
+  failureThreshold: 3
+  # -- HTTP check configuration
+  httpGet:
+    # -- Path to check
+    path: /health
+    # -- Port to check (uses named port to stay in sync with containerPort)
+    port: "container-port"
+
+# -- Liveness probe configuration
+livenessProbe:
+  # -- Failures before restarting
+  failureThreshold: 3
+  # -- How often to perform the probe
+  periodSeconds: 15
+  # -- HTTP check configuration
+  httpGet:
+    # -- Path to check
+    path: /health
+    # -- Port to check (uses named port to stay in sync with containerPort)
+    port: "container-port"
+
+# -- Ingress configuration
+ingress:
+  # -- Enable Ingress
+  enabled: false
+  # -- Ingress class name (e.g., nginx, traefik)
+  ingressClassName: ""
+  # -- Hostname for the Ingress rule
+  host: "vllm-omni.example.com"
+  # -- Additional annotations
+  annotations: {}
+  # -- TLS configuration
+  tls: []
+  #  - secretName: vllm-omni-tls
+  #    hosts:
+  #      - vllm-omni.example.com
+
+# -- Node selector for pod scheduling
+nodeSelector: {}
+
+# -- Tolerations for pod scheduling
+tolerations: []
+
+# -- Labels applied to all resources
+labels:
+  app: "vllm-omni"

examples/online_serving/stable_audio/README.md

@@ -0,0 +1,234 @@
+# Stable Audio Online Serving
+
+Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
+
+## Features
+
+- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
+- **Flexible control**: Adjust audio length, guidance scale, inference steps
+- **Quality control**: Use negative prompts to avoid unwanted characteristics
+- **Reproducible**: Set random seed for deterministic generation
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### 2. Generate Audio
+
+#### Using curl
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+#### Using Python Client
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring" \
+    --audio_length 10.0 \
+    --output cat.wav
+```
+
+#### Using Bash Script
+
+```bash
+bash curl_examples.sh
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+```
+
+### Request Body
+
+```json
+{
+    "input": "Text description of the audio",
+    "audio_length": 10.0,
+    "audio_start": 0.0,
+    "negative_prompt": "Low quality",
+    "guidance_scale": 7.0,
+    "num_inference_steps": 100,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+### Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducibility |
+| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Usage Examples
+
+### Basic Generation
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+## Python Client Examples
+
+### Simple Generation
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+```
+
+### Custom Parameters
+
+```bash
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+### Different Output Format
+
+```bash
+python stable_audio_client.py \
+    --text "Guitar playing" \
+    --response_format mp3 \
+    --output guitar.mp3
+```
+
+## Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality
+   - 100 steps: Good balance (default)
+   - 150+ steps: High quality, slower
+3. **Guidance Scale**:
+   - Lower (3-5): More creative/varied
+   - Default (7): Good balance
+   - Higher (10+): More literal to prompt
+4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
+5. **Seeds**: Use same seed for reproducible results
+
+## Performance
+
+| Inference Steps | Quality | Speed | Use Case |
+|----------------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 (default) | Very Good | Medium | Production |
+| 150+ | Excellent | Slow | Final/critical audio |
+
+## Troubleshooting
+
+### Server not responding
+- Check if server is running: `curl http://localhost:8000/health`
+- Check server logs for errors
+
+### Audio quality issues
+- Increase `num_inference_steps` (e.g., 150)
+- Add negative prompts: `"Low quality, distorted, noisy"`
+- Increase `guidance_scale` for more prompt adherence
+
+### Generation timeout
+- Reduce `num_inference_steps`
+- Reduce `audio_length`
+- Check GPU memory with `nvidia-smi`
+
+### Wrong audio length
+- Ensure `audio_length` is within model limits (~47s max)
+- Adjust `audio_start` if trimming is needed
+
+## See Also
+
+- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
+- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
+- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)

examples/online_serving/stable_audio/curl_examples.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
+
+# Example 1: Simple request with default parameters
+echo "Example 1: Simple request with default parameters"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound audience clapping and cheering in a stadium"
+    }' --output stadium.wav
+
+# Example 2: Request with custom audio_length
+echo "Example 2: Custom audio length (5 seconds)"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+
+# Example 3: Request with negative prompt for quality control
+echo "Example 3: With negative prompt"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy"
+    }' --output piano.wav
+
+# Example 4: Full control with all parameters
+echo "Example 4: Full control (custom length, guidance, steps, seed)"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+
+# Example 5: Quick generation with fewer steps (faster but lower quality)
+echo "Example 5: Quick generation (fewer steps)"
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Ocean waves crashing on a beach",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output ocean.wav
+
+echo "All examples completed!"

examples/online_serving/stable_audio/stable_audio_client.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
+
+This script demonstrates how to use the OpenAI-compatible speech API
+to generate audio from text using Stable Audio models.
+
+Examples:
+    # Simple generation
+    python stable_audio_client.py --text "The sound of a cat purring"
+
+    # With custom duration
+    python stable_audio_client.py --text "A dog barking" --audio_length 5.0
+
+    # With all parameters
+    python stable_audio_client.py --text "Thunder and rain" \
+        --audio_length 15.0 \
+        --negative_prompt "Low quality" \
+        --guidance_scale 7.0 \
+        --num_inference_steps 100 \
+        --seed 42 \
+        --output thunder.wav
+"""
+
+import argparse
+import sys
+
+import requests
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
+    parser.add_argument(
+        "--api_url",
+        default="http://localhost:8000/v1/audio/generate",
+        help="API endpoint URL",
+    )
+    parser.add_argument(
+        "--text",
+        default="The sound of a cat purring",
+        help="Text prompt for audio generation",
+    )
+    parser.add_argument(
+        "--audio_length",
+        type=float,
+        default=10.0,
+        help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
+    )
+    parser.add_argument(
+        "--audio_start",
+        type=float,
+        default=0.0,
+        help="Audio start time in seconds",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        default="Low quality",
+        help="Negative prompt for classifier-free guidance",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=7.0,
+        help="Guidance scale for diffusion (higher = more adherence to prompt)",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=100,
+        help="Number of inference steps (higher = better quality, slower)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--output",
+        default="stable_audio_output.wav",
+        help="Output file path",
+    )
+    parser.add_argument(
+        "--response_format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm"],
+        help="Audio output format",
+    )
+    return parser.parse_args()
+
+
+def generate_audio(args):
+    """Generate audio using the API."""
+
+    # Build request payload
+    payload = {
+        "input": args.text,
+        "audio_length": args.audio_length,
+        "audio_start": args.audio_start,
+        "response_format": args.response_format,
+    }
+
+    # Add optional parameters
+    if args.negative_prompt:
+        payload["negative_prompt"] = args.negative_prompt
+    if args.guidance_scale:
+        payload["guidance_scale"] = args.guidance_scale
+    if args.num_inference_steps:
+        payload["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        payload["seed"] = args.seed
+
+    print(f"\n{'=' * 60}")
+    print("Stable Audio - Text-to-Audio Generation")
+    print(f"{'=' * 60}")
+    print(f"API URL: {args.api_url}")
+    print(f"Prompt: {args.text}")
+    print(f"Audio length: {args.audio_length}s")
+    print(f"Negative prompt: {args.negative_prompt}")
+    print(f"Guidance scale: {args.guidance_scale}")
+    print(f"Inference steps: {args.num_inference_steps}")
+    if args.seed is not None:
+        print(f"Seed: {args.seed}")
+    print(f"Output: {args.output}")
+    print(f"{'=' * 60}\n")
+
+    try:
+        # Make the API request
+        print("Generating audio...")
+        response = requests.post(
+            args.api_url,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=300,  # 5 minute timeout for long generations
+        )
+
+        # Check for errors
+        if response.status_code != 200:
+            print(f"Error: API returned status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return False
+
+        # Save the audio
+        with open(args.output, "wb") as f:
+            f.write(response.content)
+
+        print(f"✓ Audio saved to {args.output}")
+        print(f"  File size: {len(response.content) / 1024:.1f} KB")
+        return True
+
+    except requests.exceptions.Timeout:
+        print("Error: Request timed out. Try reducing inference steps or audio length.")
+        return False
+    except requests.exceptions.ConnectionError:
+        print(f"Error: Could not connect to {args.api_url}")
+        print("Make sure the server is running.")
+        return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def main():
+    args = parse_args()
+    success = generate_audio(args)
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -51,6 +51,7 @@ dev = [
     "mooncake-transfer-engine==0.3.8.post1",
     "av", # for ComfyUI tests
     "openpyxl>=3.0.0", # for nightly CI
+    "pyttsx3>=2.99"
 ]
 
 docs = [

tests/conftest.py

@@ -12,6 +12,7 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
 
 import concurrent.futures
 import gc
+import multiprocessing
 import socket
 import subprocess
 import sys
@@ -217,102 +218,193 @@ def generate_synthetic_audio(
     sample_rate: int = 48000,  # Default use 48000Hz.
     save_to_file: bool = False,
 ) -> dict[str, Any]:
-    """ "Generate synthetic audio with rain."""
+    """
+    Generate TTS speech with pyttsx3 and return base64 string.
+    """
+    import tempfile
+
+    import pyttsx3
     import soundfile as sf
 
-    # Initialize audio data array
-    num_samples = int(sample_rate * duration)
+    def _pick_voice(engine: pyttsx3.Engine) -> str | None:
+        voices = engine.getProperty("voices")
+        if not voices:
+            return None
+
+        preferred_tokens = (
+            "natural",
+            "jenny",
+            "sonia",
+            "susan",
+            "zira",
+            "aria",
+            "hazel",
+            "samantha",
+            "ava",
+            "allison",
+            "female",
+            "woman",
+            "english-us",
+            "en-us",
+            "english",
+        )
+        discouraged_tokens = (
+            "espeak",
+            "robot",
+            "mbrola",
+            "microsoft david",
+            "male",
+            "man",
+        )
+
+        best_voice = voices[0]
+        best_score = float("-inf")
+        for voice in voices:
+            voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower()
+            voice_languages = " ".join(
+                lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang)
+                for lang in getattr(voice, "languages", [])
+            ).lower()
+            combined_text = f"{voice_text} {voice_languages}"
+            score = 0
+            for idx, token in enumerate(preferred_tokens):
+                if token in combined_text:
+                    score += 20 - idx
+            for token in discouraged_tokens:
+                if token in combined_text:
+                    score -= 10
+            if "english" in combined_text or "en_" in combined_text or "en-" in combined_text:
+                score += 4
+            if "en-us" in combined_text or "english-us" in combined_text:
+                score += 4
+            if score > best_score:
+                best_score = score
+                best_voice = voice
+
+        return best_voice.id
+
+    def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
+        if src_sr == dst_sr or len(audio) == 0:
+            return audio.astype(np.float32)
+
+        src_len = audio.shape[0]
+        dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr))))
+        src_idx = np.arange(src_len, dtype=np.float32)
+        dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32)
+
+        resampled_channels: list[np.ndarray] = []
+        for ch in range(audio.shape[1]):
+            resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32))
+        return np.stack(resampled_channels, axis=1)
+
+    def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray:
+        current_channels = audio.shape[1]
+        if current_channels == target_channels:
+            return audio.astype(np.float32)
+        if target_channels == 1:
+            return np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
+        if current_channels == 1:
+            return np.repeat(audio, target_channels, axis=1).astype(np.float32)
+
+        collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
+        return np.repeat(collapsed, target_channels, axis=1).astype(np.float32)
+
+    def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray:
+        if len(audio) == 0:
+            return audio
+        energy = np.max(np.abs(audio), axis=1)
+        voiced = np.where(energy > threshold)[0]
+        if len(voiced) == 0:
+            return audio
+        start = max(0, int(voiced[0]) - int(sample_rate * 0.02))
+        end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1)
+        return audio[start:end]
+
+    def _enhance_speech(audio: np.ndarray) -> np.ndarray:
+        if len(audio) == 0:
+            return audio.astype(np.float32)
+        enhanced = audio.astype(np.float32).copy()
+        enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32)
+        if len(enhanced) > 1:
+            preemphasis = enhanced.copy()
+            preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1]
+            enhanced = 0.7 * enhanced + 0.3 * preemphasis
+        # Mild dynamic-range compression for ASR/TTS robustness.
+        enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced))
+        # Light fade to avoid clicks after trimming/repeating.
+        fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01)))
+        if fade > 1:
+            ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32)
+            ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32)
+            enhanced[:fade] *= ramp_in[:, None]
+            enhanced[-fade:] *= ramp_out[:, None]
+        peak = float(np.max(np.abs(enhanced)))
+        if peak > 1e-8:
+            enhanced = enhanced / peak * 0.95
+        return enhanced.astype(np.float32)
+
+    phrase_text = "test"
+    num_samples = int(sample_rate * max(1, duration))
     audio_data = np.zeros((num_samples, num_channels), dtype=np.float32)
 
-    # Configure parameters based on rain intensity
-    drop_density = 10  # Number of raindrops per second
-    drop_volume = 0.15  # Volume of individual raindrops
-    background_volume = 0.02  # Volume of background rain noise
+    engine = pyttsx3.init()
+    engine.setProperty("rate", 112)
+    engine.setProperty("volume", 1.0)
+    selected_voice = _pick_voice(engine)
+    if selected_voice is not None:
+        engine.setProperty("voice", selected_voice)
 
-    # Pink noise sounds more natural than white noise for rain
-    white_noise = np.random.randn(num_samples)
-    pink_noise = np.convolve(white_noise, np.ones(8) / 8, mode="same")
-    pink_noise = pink_noise / np.max(np.abs(pink_noise)) if np.max(np.abs(pink_noise)) > 0 else pink_noise
-    bg_noise = pink_noise * background_volume
+    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    temp_wav.close()
 
-    # Add background noise to all channels
-    for ch in range(num_channels):
-        audio_data[:, ch] += bg_noise
+    try:
+        engine.save_to_file(phrase_text, temp_wav.name)
+        engine.runAndWait()
+        engine.stop()
 
-    # Total number of raindrops = density × duration × channels for stereo effect
-    total_drops = int(drop_density * duration * num_channels)
+        ready = False
+        for _ in range(50):
+            if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44:
+                ready = True
+                break
+            time.sleep(0.1)
 
-    for _ in range(total_drops):
-        # Random timing for raindrop
-        drop_time = random.uniform(0, duration)
+        if not ready:
+            raise RuntimeError("pyttsx3 did not produce a WAV file in time.")
 
-        # Random duration of raindrop sound (0.01-0.05 seconds)
-        drop_duration = random.uniform(0.01, 0.05)
+        tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True)
+    finally:
+        if os.path.exists(temp_wav.name):
+            os.unlink(temp_wav.name)
 
-        # Random frequency gives variation in raindrop pitch
-        drop_freq = random.uniform(500, 5000)  # Hz
+    if len(tts_audio) == 0:
+        raise RuntimeError("pyttsx3 produced an empty WAV file.")
 
-        # Random channel selection for stereo positioning
-        channel = random.randint(0, num_channels - 1)
+    tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate)
+    tts_audio = _match_channels(tts_audio, num_channels)
+    tts_audio = _trim_silence(tts_audio, threshold=0.012)
+    tts_audio = _enhance_speech(tts_audio)
 
-        # Calculate sample positions for this raindrop
-        start_sample = int(drop_time * sample_rate)
-        drop_samples = int(drop_duration * sample_rate)
-        end_sample = min(start_sample + drop_samples, num_samples)
+    lead_silence = min(int(sample_rate * 0.02), num_samples // 8)
+    pause_samples = int(sample_rate * 0.18)
+    start = lead_silence
+    phrase_len = tts_audio.shape[0]
 
-        if start_sample < end_sample:
-            # Generate the raindrop sound
-            num_drop_samples = end_sample - start_sample
-            t = np.arange(num_drop_samples) / sample_rate
+    while start < num_samples:
+        take = min(phrase_len, num_samples - start)
+        audio_data[start : start + take] = tts_audio[:take]
+        start += phrase_len + pause_samples
 
-            # Basic sine wave for raindrop sound
-            drop_sound = drop_volume * np.sin(2 * math.pi * drop_freq * t)
-
-            # Apply envelope for natural attack and decay
-            envelope = np.ones(num_drop_samples)
-            attack_samples = int(num_drop_samples * 0.1)  # 10% of samples for attack
-            decay_samples = num_drop_samples - attack_samples
-
-            if attack_samples > 0:
-                # Linear attack: volume increases from 0 to 1
-                envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
-
-            if decay_samples > 0:
-                # Exponential decay for natural sound fade
-                decay = np.exp(-8 * t[attack_samples:] / drop_duration)
-                envelope[attack_samples:] = decay
-
-            # Apply envelope to raindrop sound
-            drop_sound *= envelope
-
-            # Add raindrop sound to selected channel
-            audio_data[start_sample:end_sample, channel] += drop_sound
-
-    # Step 3: Add simple reverb effect for realism
-    # Reverb simulates sound reflections in environment
-    if duration > 2:
-        # Single delay reverb (100ms delay)
-        delay_samples = int(0.1 * sample_rate)
-        if delay_samples < num_samples:
-            for ch in range(num_channels):
-                delayed = np.zeros(num_samples)
-                delayed[delay_samples:] = audio_data[:-delay_samples, ch] * 0.3
-                audio_data[:, ch] += delayed
-
-    # Step 4: Normalize audio to prevent clipping
-    # Find maximum amplitude and scale to 80% of maximum volume
-    max_amp = np.max(np.abs(audio_data))
+    max_amp = float(np.max(np.abs(audio_data)))
     if max_amp > 0:
-        audio_data = audio_data / max_amp * 0.8
+        audio_data = audio_data / max_amp * 0.95
 
-    audio_array = audio_data.copy()
-    result = {
-        "np_array": audio_array,
+    audio_bytes: bytes | None = None
+    output_path: str | None = None
+    result: dict[str, Any] = {
+        "np_array": audio_data.copy(),
     }
 
-    # Handle file saving
-    audio_bytes = None
-
     if save_to_file:
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         output_path = f"audio_{num_channels}ch_{timestamp}.wav"
@@ -617,24 +709,38 @@ def convert_audio_to_text(audio_data):
     return text
 
 
-def convert_audio_file_to_text(output_path):
+def _whisper_transcribe_in_current_process(output_path: str) -> str:
     import whisper
 
-    model = whisper.load_model("small")
-    text = model.transcribe(
-        output_path,
-        temperature=0.0,
-        word_timestamps=True,
-        condition_on_previous_text=False,
-    )["text"]
-    del model
-    if torch.cuda.is_available():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = whisper.load_model("small", device=device)
+    try:
+        text = model.transcribe(
+            output_path,
+            temperature=0.0,
+            word_timestamps=True,
+            condition_on_previous_text=False,
+        )["text"]
+    finally:
+        # Sync GPU so in-flight ops finish before we free the model; otherwise
+        # freed memory may not show up until those ops complete.
+        if torch is not None and torch.cuda.is_available():
+            torch.cuda.synchronize()
+        del model
         gc.collect()
-        torch.cuda.empty_cache()
-    if text:
-        return text
-    else:
-        return ""
+        if torch is not None and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    return text or ""
+
+
+def convert_audio_file_to_text(output_path: str) -> str:
+    """Convert an audio file to text in an isolated subprocess."""
+    # Import locally to avoid impacting test module import time.
+    ctx = multiprocessing.get_context("spawn")
+    with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
+        future = executor.submit(_whisper_transcribe_in_current_process, output_path)
+        return future.result()
 
 
 def merge_base64_and_convert_to_text(base64_list):
@@ -1105,7 +1211,7 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
     assert response.success, "The request failed."
     e2e_latency = response.e2e_latency
     if e2e_latency is not None:
-        print(f"the avg e2e is: {e2e_latency}")
+        print(f"the avg e2e latency is: {e2e_latency}")
 
     modalities = request_config.get("modalities", ["text", "audio"])
 

tests/distributed/omni_connectors/test_kv_flow.py

@@ -6,6 +6,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
     OmniKVCacheConfig,
     OmniKVTransferManager,
 )
+from vllm_omni.distributed.omni_connectors.utils.kv_utils import normalize_layer_kv
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.cache]
@@ -192,8 +193,7 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants,
     else:
         layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor")
 
-    manager = OmniKVTransferManager(kv_config)
-    normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
+    normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
     assert normalized is None
 
 

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -13,7 +13,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -46,7 +46,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -13,7 +13,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -45,7 +45,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -27,7 +27,12 @@ def inference(model_name: str, offload: bool = True):
     torch.cuda.reset_peak_memory_stats(device=device_index)
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
-    m = Omni(model=model_name, enable_cpu_offload=offload)
+    m = Omni(
+        model=model_name,
+        # TODO: we might want to add overlapped feature e2e tests
+        # cache_backend="cache_dit",
+        enable_cpu_offload=offload,
+    )
     torch.cuda.reset_peak_memory_stats(device=device_index)
     height = 256
     width = 256

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -38,6 +38,8 @@ def run_inference(
     m = Omni(
         model=model_name,
         enable_layerwise_offload=layerwise_offload,
+        # TODO: we might want to add overlapped feature e2e tests
+        # cache_backend="cache_dit",
         boundary_ratio=0.875,
         flow_shift=5.0,
     )

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -10,12 +10,29 @@ from tests.conftest import (
     generate_synthetic_audio,
     generate_synthetic_image,
     generate_synthetic_video,
+    modify_stage_config,
 )
 from tests.utils import hardware_test
 from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 
+
+def get_cuda_graph_config():
+    path = modify_stage_config(
+        str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml"),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+                1: {"engine_args.enforce_eager": "true"},
+            },
+        },
+    )
+    return path
+
+
 # CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
 if current_omni_platform.is_npu():
     stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
@@ -23,7 +40,7 @@ elif current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
 else:
-    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
+    stage_config = get_cuda_graph_config()
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models]

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -13,18 +13,35 @@ import pytest
 
 from tests.conftest import (
     generate_synthetic_video,
+    modify_stage_config,
 )
 from tests.utils import hardware_test
 from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
+
+def get_cuda_graph_config():
+    path = modify_stage_config(
+        str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+                1: {"engine_args.enforce_eager": "true"},
+            },
+        },
+    )
+    return path
+
+
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
 if current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")]
+    stage_configs = [get_cuda_graph_config()]
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]

tests/e2e/online_serving/test_qwen3_omni.py

@@ -114,7 +114,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
         "messages": messages,
         "stream": True,
         "key_words": {
-            "audio": ["water", "chirping", "crackling", "rain"],
+            "audio": ["test"],
             "image": ["square", "quadrate"],
         },
     }
@@ -126,13 +126,14 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_001(omni_server, openai_client) -> None:
     """
     Test text input processing and text/audio output generation via OpenAI API.
     Deploy Setting: default yaml
     Input Modal: text
-    Output Modal: text + audio
+    Output Modal: text
     Datasets: few requests
     """
     messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -22,7 +22,7 @@ from tests.utils import hardware_test
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
-AUDIO_KEY = ["water", "chirping", "crackling", "rain"]
+AUDIO_KEY = ["test"]
 IMAGE_KEY = ["square", "quadrate"]
 VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"]
 
@@ -90,12 +90,19 @@ def get_max_batch_size(size_type="few"):
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
     messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
 
     request_config = {
         "model": omni_server.model,
         "messages": messages,
         "modalities": ["audio"],
+        "stream": True,
         "key_words": {"text": ["beijing"]},
     }
 
@@ -107,6 +114,12 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: few requests
+    """
     messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
 
     request_config = {
@@ -115,7 +128,7 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"text": ["beijing"]},
     }
 
-    openai_client.send_request(request_config)
+    openai_client.send_request(request_config, request_num=get_max_batch_size())
 
 
 @pytest.mark.advanced_model
@@ -123,6 +136,12 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_image_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: image
+    Output Modal: text
+    Input Setting: stream=True
+    Datasets: single request
+    """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
     messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
 
@@ -130,6 +149,7 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
         "model": omni_server.model,
         "messages": messages,
         "modalities": ["text"],
+        "stream": True,
         "key_words": {"image": IMAGE_KEY},
     }
 
@@ -141,6 +161,12 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_image_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: image
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
     messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
 
@@ -159,6 +185,12 @@ def test_image_to_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_image_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: image
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: few requests
+    """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
 
     messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
@@ -177,6 +209,12 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_video_to_text_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: video
+    Output Modal: text
+    Input Setting: stream=False
+    Datasets: single request
+    """
     video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
     messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
 
@@ -195,6 +233,12 @@ def test_video_to_text_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_video_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: video
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
     video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
     messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
 
@@ -213,6 +257,12 @@ def test_video_to_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_video_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: video
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: few requests
+    """
     video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
 
     messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
@@ -231,6 +281,12 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text, audio
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
     audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
     messages = dummy_messages_from_mix_data(
         audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
@@ -239,7 +295,6 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
     request_config = {
         "model": omni_server.model,
         "messages": messages,
-        "modalities": ["text"],
         "key_words": {"audio": AUDIO_KEY},
     }
 
@@ -251,6 +306,12 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text, image
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
 
     messages = dummy_messages_from_mix_data(
@@ -260,19 +321,23 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
     request_config = {
         "model": omni_server.model,
         "messages": messages,
-        "modalities": ["audio"],
         "key_words": {"image": IMAGE_KEY},
     }
 
     openai_client.send_request(request_config)
 
 
-@pytest.mark.skip(reason="There is a known issue with oom error.")
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text, video
+    Output Modal: text, audio
+    Input Setting: stream=True
+    Datasets: single requests
+    """
     video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 30)['base64']}"
 
     messages = dummy_messages_from_mix_data(
@@ -282,10 +347,11 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
     request_config = {
         "model": omni_server.model,
         "messages": messages,
+        "stream": True,
         "key_words": {"video": VIDEO_KEY},
     }
 
-    openai_client.send_request(request_config, request_num=get_max_batch_size())
+    openai_client.send_request(request_config)
 
 
 @pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
@@ -294,6 +360,12 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text, audio, image, video
+    Output Modal: text, audio
+    Input Setting: stream=True
+    Datasets: few requests
+    """
     video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
     audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"

tests/entrypoints/openai_api/test_serving_audio_generate.py

@@ -0,0 +1,509 @@
+# tests/entrypoints/openai_api/test_serving_audio_generate.py
+import logging
+from inspect import Signature, signature
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    CreateAudio,
+    OpenAICreateAudioGenerateRequest,
+)
+from vllm_omni.entrypoints.openai.serving_audio_generate import (
+    OmniOpenAIServingAudioGenerate,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+logger = logging.getLogger(__name__)
+
+
+# Helper: create a mock audio output for endpoint tests
+def create_mock_audio_output(
+    request_id: str = "audiogen-mock-123",
+    sample_rate: int = 44100,
+    num_samples: int = 44100,
+    audio_key: str = "audio",
+) -> OmniRequestOutput:
+    """Return an OmniRequestOutput mimicking diffusion audio model output."""
+
+    audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
+
+    return OmniRequestOutput.from_diffusion(
+        request_id=request_id,
+        images=[],
+        prompt=None,
+        metrics={},
+        multimodal_output={audio_key: audio_tensor, "sr": sample_rate},
+    )
+
+
+def _make_engine_client(*, audio_key: str = "audio", sample_rate: int = 44100):
+    """Build a mock engine client producing audio output."""
+    mock_engine_client = MagicMock()
+    mock_engine_client.errored = False
+    mock_engine_client.model_type = "StableAudioPipeline"
+    mock_engine_client.default_sampling_params_list = [{}]
+
+    async def mock_generate_fn(*args, **kwargs):
+        yield create_mock_audio_output(
+            request_id=kwargs.get("request_id", "audiogen-mock"),
+            sample_rate=sample_rate,
+            audio_key=audio_key,
+        )
+
+    mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn)
+    return mock_engine_client
+
+
+def _make_server(engine_client=None):
+    """Build an OmniOpenAIServingAudioGenerate with mocks."""
+    if engine_client is None:
+        engine_client = _make_engine_client()
+
+    mock_models = MagicMock()
+    mock_models.is_base_model.return_value = True
+
+    return OmniOpenAIServingAudioGenerate(
+        engine_client=engine_client,
+        models=mock_models,
+        request_logger=MagicMock(),
+    )
+
+
+@pytest.fixture
+def test_app():
+    server = _make_server()
+
+    original_fn = server.create_audio_generate
+    sig = signature(original_fn)
+    new_params = [p for name, p in sig.parameters.items() if name != "raw_request"]
+    new_sig = Signature(parameters=new_params, return_annotation=sig.return_annotation)
+
+    async def patched_create_audio_generate(*args, **kwargs):
+        return await original_fn(*args, **kwargs)
+
+    patched_create_audio_generate.__signature__ = new_sig
+    server.create_audio_generate = patched_create_audio_generate
+
+    app = FastAPI()
+    app.add_api_route(
+        "/v1/audio/generate",
+        server.create_audio_generate,
+        methods=["POST"],
+        response_model=None,
+    )
+    return app
+
+
+@pytest.fixture
+def client(test_app):
+    return TestClient(test_app)
+
+
+# Request Validation (Pydantic model)
+class TestRequestValidation:
+    """Validate OpenAICreateAudioGenerateRequest pydantic constraints."""
+
+    def test_valid_minimal_request(self):
+        req = OpenAICreateAudioGenerateRequest(input="A calm piano melody")
+        assert req.input == "A calm piano melody"
+        assert req.response_format == "wav"
+        assert req.speed == 1.0
+
+    def test_fields_are_wired_correctly(self):
+        req = OpenAICreateAudioGenerateRequest(
+            input="rain sounds",
+            model="stable-audio",
+            response_format="flac",
+            speed=1.5,
+            audio_length=10.0,
+            audio_start=2.0,
+            negative_prompt="noise",
+            guidance_scale=7.5,
+            num_inference_steps=100,
+            seed=42,
+        )
+        assert req.input == "rain sounds"
+        assert req.model == "stable-audio"
+        assert req.response_format == "flac"
+        assert req.speed == 1.5
+        assert req.audio_length == 10.0
+        assert req.audio_start == 2.0
+        assert req.negative_prompt == "noise"
+        assert req.guidance_scale == 7.5
+        assert req.num_inference_steps == 100
+        assert req.seed == 42
+
+    def test_invalid_response_format(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", response_format="invalid_format")
+
+    def test_speed_lower_bound(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", speed=0.1)
+
+    def test_speed_upper_bound(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", speed=5.0)
+
+    def test_speed_at_boundaries(self):
+        req_low = OpenAICreateAudioGenerateRequest(input="test", speed=0.25)
+        assert req_low.speed == 0.25
+        req_high = OpenAICreateAudioGenerateRequest(input="test", speed=4.0)
+        assert req_high.speed == 4.0
+
+    def test_stream_format_sse_rejected(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", stream_format="sse")
+
+    def test_stream_format_audio_accepted(self):
+        req = OpenAICreateAudioGenerateRequest(input="test", stream_format="audio")
+        assert req.stream_format == "audio"
+
+
+# Constructor & Class Methods
+class TestConstructor:
+    def test_default_init(self):
+        server = _make_server()
+        assert server.diffusion_mode is False
+
+    def test_for_diffusion_factory(self):
+        engine_client = _make_engine_client()
+        mock_models = MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        server = OmniOpenAIServingAudioGenerate.for_diffusion(
+            engine_client=engine_client,
+            models=mock_models,
+            request_logger=MagicMock(),
+        )
+        assert server.diffusion_mode is True
+
+    def test_is_stable_audio_model_true(self):
+        server = _make_server()
+        assert server._is_stable_audio_model() is True
+
+    def test_is_stable_audio_model_false(self):
+        engine = _make_engine_client()
+        engine.model_type = "SomeOtherModel"
+        server = _make_server(engine_client=engine)
+        assert server._is_stable_audio_model() is False
+
+
+# Parameter Wiring — verify request params reach the engine
+class TestParameterWiring:
+    """Ensure request parameters are correctly forwarded to the engine."""
+
+    @pytest.fixture
+    def server_and_engine(self):
+        engine = _make_engine_client()
+        server = _make_server(engine_client=engine)
+        return server, engine
+
+    @pytest.mark.asyncio
+    async def test_prompt_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="birds chirping")
+        await server.create_audio_generate(req)
+
+        engine.generate.assert_called_once()
+        call_kwargs = engine.generate.call_args[1]
+        assert call_kwargs["prompt"]["prompt"] == "birds chirping"
+        assert call_kwargs["output_modalities"] == ["audio"]
+
+    @pytest.mark.asyncio
+    async def test_negative_prompt_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="a calm ocean", negative_prompt="noise distortion")
+        await server.create_audio_generate(req)
+
+        call_kwargs = engine.generate.call_args[1]
+        assert call_kwargs["prompt"]["negative_prompt"] == "noise distortion"
+
+    @pytest.mark.asyncio
+    async def test_negative_prompt_absent(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="a calm ocean")
+        await server.create_audio_generate(req)
+
+        call_kwargs = engine.generate.call_args[1]
+        assert "negative_prompt" not in call_kwargs["prompt"]
+
+    @pytest.mark.asyncio
+    async def test_guidance_scale_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", guidance_scale=12.0)
+        await server.create_audio_generate(req)
+
+        call_kwargs = engine.generate.call_args[1]
+        sp = call_kwargs["sampling_params_list"][0]
+        assert isinstance(sp, OmniDiffusionSamplingParams)
+        assert sp.guidance_scale == 12.0
+
+    @pytest.mark.asyncio
+    async def test_num_inference_steps_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", num_inference_steps=200)
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.num_inference_steps == 200
+
+    @pytest.mark.asyncio
+    async def test_seed_creates_generator(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", seed=42)
+
+        with patch("vllm_omni.entrypoints.openai.serving_audio_generate.torch") as mock_torch:
+            mock_gen = MagicMock()
+            mock_gen.manual_seed.return_value = mock_gen
+            mock_torch.Generator.return_value = mock_gen
+
+            await server.create_audio_generate(req)
+
+            mock_torch.Generator.assert_called_once()
+            mock_gen.manual_seed.assert_called_once_with(42)
+
+    @pytest.mark.asyncio
+    async def test_seed_none_skips_generator(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test")
+
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.generator is None
+
+    @pytest.mark.asyncio
+    async def test_audio_length_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", audio_length=10.0, audio_start=2.0)
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.extra_args["audio_start_in_s"] == 2.0
+        assert sp.extra_args["audio_end_in_s"] == 12.0  # start + length
+
+    @pytest.mark.asyncio
+    async def test_audio_length_default_start(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", audio_length=5.0)
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.extra_args["audio_start_in_s"] == 0.0
+        assert sp.extra_args["audio_end_in_s"] == 5.0
+
+    @pytest.mark.asyncio
+    async def test_no_audio_length_skips_extra_args(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.extra_args == {}
+
+    @pytest.mark.asyncio
+    async def test_defaults_not_set_when_omitted(self, server_and_engine):
+        """Guidance scale and num_inference_steps keep dataclass defaults when not in request."""
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        defaults = OmniDiffusionSamplingParams()
+        assert sp.guidance_scale == defaults.guidance_scale
+        assert sp.num_inference_steps == defaults.num_inference_steps
+
+
+# Audio Response Format
+class TestAudioResponseFormat:
+    def test_wav_response(self, client):
+        payload = {"input": "a gentle rain", "response_format": "wav"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/wav"
+        assert len(response.content) > 0
+
+    def test_mp3_response(self, client):
+        payload = {"input": "a gentle rain", "response_format": "mp3"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/mpeg"
+        assert len(response.content) > 0
+
+    def test_flac_response(self, client):
+        payload = {"input": "a gentle rain", "response_format": "flac"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/flac"
+        assert len(response.content) > 0
+
+    def test_invalid_format_rejected(self, client):
+        payload = {"input": "test", "response_format": "banana"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 422
+
+    @patch("vllm_omni.entrypoints.openai.serving_audio_generate.OmniOpenAIServingAudioGenerate.create_audio")
+    def test_speed_parameter_forwarded(self, mock_create_audio, test_app):
+        mock_audio_response = MagicMock()
+        mock_audio_response.audio_data = b"dummy_audio"
+        mock_audio_response.media_type = "audio/wav"
+        mock_create_audio.return_value = mock_audio_response
+
+        c = TestClient(test_app)
+        payload = {"input": "test", "response_format": "wav", "speed": 2.5}
+        c.post("/v1/audio/generate", json=payload)
+
+        mock_create_audio.assert_called_once()
+        audio_obj = mock_create_audio.call_args[0][0]
+        assert isinstance(audio_obj, CreateAudio)
+        assert audio_obj.speed == 2.5
+
+    @patch("vllm_omni.entrypoints.openai.serving_audio_generate.OmniOpenAIServingAudioGenerate.create_audio")
+    def test_sample_rate_from_output(self, mock_create_audio, test_app):
+        mock_audio_response = MagicMock()
+        mock_audio_response.audio_data = b"dummy"
+        mock_audio_response.media_type = "audio/wav"
+        mock_create_audio.return_value = mock_audio_response
+
+        c = TestClient(test_app)
+        payload = {"input": "test"}
+        c.post("/v1/audio/generate", json=payload)
+
+        audio_obj = mock_create_audio.call_args[0][0]
+        assert audio_obj.sample_rate == 44100  # Stable Audio default
+
+
+# Error Handling
+class TestErrorHandling:
+    @pytest.mark.asyncio
+    async def test_no_output_returns_error(self):
+        engine = _make_engine_client()
+
+        async def empty_gen(*args, **kwargs):
+            return
+            yield  # noqa: unreachable – makes this an async generator
+
+        engine.generate = MagicMock(side_effect=empty_gen)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        # create_error_response returns an ErrorResponse with .error.message
+        assert "No output generated" in resp.error.message
+
+    @pytest.mark.asyncio
+    async def test_no_audio_in_output_returns_error(self):
+        engine = _make_engine_client()
+
+        async def gen_without_audio(*args, **kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="test",
+                images=[],
+                prompt=None,
+                metrics={},
+                multimodal_output={},  # no audio key
+            )
+
+        engine.generate = MagicMock(side_effect=gen_without_audio)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        assert "did not produce audio" in resp.error.message
+
+    @pytest.mark.asyncio
+    async def test_engine_errored_raises(self):
+        engine = _make_engine_client()
+        engine.errored = True
+        engine.dead_error = RuntimeError("engine is dead")
+        server = _make_server(engine_client=engine)
+
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        with pytest.raises(RuntimeError, match="engine is dead"):
+            await server.create_audio_generate(req)
+
+    @pytest.mark.asyncio
+    async def test_model_outputs_key_fallback(self):
+        """Audio data under 'model_outputs' key should be accepted."""
+        engine = _make_engine_client(audio_key="model_outputs")
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        # Should succeed and return a Response with audio bytes
+        assert hasattr(resp, "body")
+        assert len(resp.body) > 0
+
+    @pytest.mark.asyncio
+    async def test_value_error_returns_error_response(self):
+        engine = _make_engine_client()
+
+        async def gen_value_error(*args, **kwargs):
+            raise ValueError("bad value")
+            yield  # noqa: unreachable
+
+        engine.generate = MagicMock(side_effect=gen_value_error)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        assert "bad value" in resp.error.message
+
+    @pytest.mark.asyncio
+    async def test_generic_exception_returns_error_response(self):
+        engine = _make_engine_client()
+
+        async def gen_runtime_error(*args, **kwargs):
+            raise RuntimeError("something went wrong")
+            yield  # noqa: unreachable
+
+        engine.generate = MagicMock(side_effect=gen_runtime_error)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        assert "Audio generation failed" in resp.error.message
+
+
+# End-to-End via TestClient
+class TestAudioGenerateAPI:
+    def test_basic_success(self, client):
+        payload = {"input": "ambient forest sounds"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert len(response.content) > 0
+
+    def test_with_all_params(self, client):
+        payload = {
+            "input": "gentle piano",
+            "response_format": "wav",
+            "speed": 1.0,
+            "audio_length": 5.0,
+            "audio_start": 0.0,
+            "negative_prompt": "noise",
+            "guidance_scale": 7.0,
+            "num_inference_steps": 50,
+            "seed": 123,
+        }
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/wav"
+
+    def test_missing_input_rejected(self, client):
+        payload = {}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 422
+
+    def test_extra_unknown_fields_ignored(self, client):
+        payload = {"input": "test", "unknown_field": "value"}
+        response = client.post("/v1/audio/generate", json=payload)
+        # Pydantic v2 ignores extra fields by default
+        assert response.status_code == 200

tests/entrypoints/test_omni_diffusion.py

@@ -612,7 +612,7 @@ def test_initialize_stage_configs_called_when_none(
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -687,7 +687,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, moc
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -742,7 +742,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mo
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -848,7 +848,7 @@ def test_generate_pipeline_with_batch_input(monkeypatch: pytest.MonkeyPatch, moc
     stage_cfg1["stage_id"] = 1
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -969,7 +969,7 @@ def test_generate_no_final_output_returns_empty(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1062,7 +1062,7 @@ def test_generate_sampling_params_none_use_default(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1140,7 +1140,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker:
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1202,7 +1202,7 @@ def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker
     """Test that generate handles error messages from stages correctly."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1283,7 +1283,7 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
     """Test that close() sends shutdown signal to all input queues."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,

tests/entrypoints/test_omni_llm.py

@@ -577,7 +577,7 @@ def test_initialize_stage_configs_called_when_none(
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -654,7 +654,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, moc
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -711,7 +711,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mo
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -819,7 +819,7 @@ def test_generate_no_final_output_returns_empty(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -911,7 +911,7 @@ def test_generate_sampling_params_none_use_default(
     stage_cfg1["final_output"] = False
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -988,7 +988,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker:
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1052,7 +1052,7 @@ def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker
     """Test that generate handles error messages from stages correctly."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,
@@ -1135,7 +1135,7 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
     """Test that close() sends shutdown signal to all input queues."""
 
     def _fake_loader(
-        model: str,
+        config_path: str,
         stage_configs_path: str | None = None,
         base_engine_args: dict | None = None,
         default_stage_cfg_factory=None,

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -7,7 +7,10 @@ from types import SimpleNamespace
 import pytest
 import torch
 
-from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_async_chunk
+from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+    talker2code2wav,
+    talker2code2wav_async_chunk,
+)
 
 _FRAME = [1, 2, 3, 4]  # 4-codebook frame
 _Q = len(_FRAME)  # num quantizers
@@ -29,6 +32,7 @@ def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
     return SimpleNamespace(
         code_prompt_token_ids=defaultdict(list),
         put_req_chunk=defaultdict(int),
+        request_payload={},
         connector=SimpleNamespace(
             config={
                 "extra": {
@@ -152,3 +156,113 @@ def test_per_request_override_wins_over_stage_config():
     tm = _tm(initial_chunk=5)
     payload = _call(tm, "r-override2", n_frames=10, put_req=0, req_ic=15)
     assert payload is None
+
+
+def test_first_streaming_chunk_prepends_ref_code_context():
+    tm = _tm(initial_chunk=10)
+    rid = "r-ref"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(10)]
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 2
+    assert len(payload["code_predictor_codes"]) == _Q * 12
+
+
+def test_ref_code_context_only_applies_to_first_streaming_chunk():
+    tm = _tm(initial_chunk=10)
+    rid = "r-ref2"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)]
+    tm.put_req_chunk[rid] = 1
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 10
+    assert len(payload["code_predictor_codes"]) == _Q * 20
+
+
+def test_ref_code_context_can_be_buffered_before_first_emit():
+    tm = _tm(initial_chunk=10)
+    rid = "r-ref-buffered"
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+    first_payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]]), "ref_code": ref_code},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+    assert first_payload is None
+    assert rid in tm.request_payload
+
+    for _ in range(8):
+        talker2code2wav_async_chunk(
+            transfer_manager=tm,
+            pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+            request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+            is_finished=False,
+        )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 2
+    assert len(payload["code_predictor_codes"]) == _Q * 12
+    assert rid not in tm.request_payload
+
+
+def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    audio_codes = torch.tensor(
+        [
+            [0, 0, 0, 0],
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+        ],
+        dtype=torch.long,
+    )
+    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
+    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+
+    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
+
+    assert len(prompts) == 1
+    prompt = prompts[0]
+    assert prompt["additional_information"] == {"left_context_size": [2]}
+    assert prompt["prompt_token_ids"] == [
+        9,
+        8,
+        1,
+        5,
+        9,
+        8,
+        2,
+        6,
+        9,
+        8,
+        3,
+        7,
+        9,
+        8,
+        4,
+        8,
+    ]

vllm_omni/diffusion/hooks/base.py

@@ -107,7 +107,7 @@ class ModelHook:
             The output of the forward pass.
         """
         args, kwargs = self.pre_forward(module, *args, **kwargs)
-        output = module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        output = module._omni_original_forward(*args, **kwargs)  # type: ignore[attr-defined]
         return self.post_forward(module, output)
 
     def reset_state(self, module: nn.Module) -> nn.Module:
@@ -131,7 +131,7 @@ class _WrappedForward:
     def __call__(self, *args: Any, **kwargs: Any):
         registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
         if registry is None or not registry._hooks:
-            return self.module._original_forward(*args, **kwargs)
+            return self.module._omni_original_forward(*args, **kwargs)
         return registry.dispatch(*args, **kwargs)
 
 
@@ -162,8 +162,10 @@ class HookRegistry:
             setattr(module, "_hook_registry", registry)
 
             # Wrap module.forward once so hooks can intercept calls.
-            if not hasattr(module, "_original_forward"):
-                module._original_forward = module.forward  # type: ignore[attr-defined]
+            # NOTE: Use `_omni_original_forward` to avoid collision with cache-dit's
+            # `_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]
 
         return registry
@@ -212,7 +214,7 @@ class HookRegistry:
             The output of the forward pass.
         """
         if not self._hooks:
-            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+            return self.module._omni_original_forward(*args, **kwargs)  # type: ignore[attr-defined]
 
         # For single hook case, call directly
         if len(self._hooks) == 1:
@@ -228,7 +230,7 @@ class HookRegistry:
             args, kwargs = hook.pre_forward(self.module, *args, **kwargs)
 
         # Call original forward
-        output = self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        output = self.module._omni_original_forward(*args, **kwargs)  # type: ignore[attr-defined]
 
         # Apply all post_forward hooks in reverse order
         for _, hook in reversed(sorted_hooks):

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -332,6 +332,9 @@ class BagelPipeline(nn.Module):
             else:
                 gen_context["ropes"] = [seq_len]
 
+            if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata:
+                image_shape = tuple(req.sampling_params.kv_metadata["image_shape"])
+
             cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None)
             if cfg_text_kv is not None:
                 logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache")

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -356,7 +356,6 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
         negative_prompt: str | list[str] | None = None,
         audio_end_in_s: float | None = None,
         audio_start_in_s: float = 0.0,
-        num_inference_steps: int = 100,
         guidance_scale: float = 7.0,
         num_waveforms_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -374,7 +373,6 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
             negative_prompt: Negative prompt for CFG
             audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0)
             audio_start_in_s: Audio start time in seconds
-            num_inference_steps: Number of denoising steps
             guidance_scale: CFG scale
             num_waveforms_per_prompt: Number of audio outputs per prompt
             generator: Random generator for reproducibility
@@ -395,7 +393,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_inference_steps = req.sampling_params.num_inference_steps
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
 

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -50,6 +50,9 @@ class LayerwiseOffloadHook(ModelHook):
         # Per-block synchronization primitive: set after H2D copy completes.
         self._prefetch_done: current_omni_platform.Event | None = None
 
+        # Backward link to the hook that is responsible for prefetching *this* block's weights
+        self._prev_hook: LayerwiseOffloadHook | None = None
+
         self.next_block_parameters: dict[str, nn.Parameter] = {}
         self.next_block_buffers: dict[str, torch.Tensor] = {}
         self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
@@ -128,6 +131,14 @@ class LayerwiseOffloadHook(ModelHook):
 
         return dtype_cpu_flattened_weights, dtype_metadata
 
+    @property
+    def is_materialized(self) -> bool:
+        """Check whether this block's parameters hold real data on device."""
+        for param in self.block_parameters.values():
+            return param.data.dim() > 0
+
+        return True
+
     @torch.compiler.disable
     def prefetch_layer(self, non_blocking: bool = True) -> None:
         """Copy layer weights from CPU -> GPU.
@@ -185,6 +196,12 @@ class LayerwiseOffloadHook(ModelHook):
             buf.data = torch.empty((), device=self.device, dtype=buf.dtype)
 
     def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
+        # if the previous hook was skipped and the weights are not on device,
+        # (e.g. by cache-dit block caching), ask the previous hook to
+        # synchronously prefetch *this* block's weights before computation
+        if not self.is_materialized and self._prev_hook is not None:
+            self._prev_hook.prefetch_layer(non_blocking=False)
+
         self.prefetch_layer(non_blocking=True)
 
         return args, kwargs
@@ -293,13 +310,23 @@ class LayerWiseOffloadBackend(OffloadBackend):
             # For subsequent requests, the first layer/block will be pre-fetched
             # during the last layer compute of the previous request.
             last_block, first_block = blocks[-1], blocks[0]
-            hook = apply_block_hook(last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
-            hook.prefetch_layer(non_blocking=False)
+            last_hook = apply_block_hook(
+                last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory
+            )
+            last_hook.prefetch_layer(non_blocking=False)
 
+            block_hooks: list[LayerwiseOffloadHook] = [last_hook]
             # Register hook for each of blocks
             for i, block in enumerate(blocks[:-1]):
                 next_block = blocks[(i + 1) % num_blocks]
-                apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
+                hook = apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
+                block_hooks.append(hook)
+
+            # NOTE(yuanheng-zhao): We make each hook gets a backward reference to the hook
+            # that is responsible for prefetching its block's weights. This is specifically a
+            # workaround for that arbitrary blocks are skipped by caching systems (e.g., cache-dit)
+            for i in range(len(block_hooks)):
+                block_hooks[i]._prev_hook = block_hooks[i - 1]
 
             logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)")
 

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -35,6 +35,23 @@ class SequentialOffloadHook(ModelHook):
         self.device = device
         self.pin_memory = pin_memory
 
+    @staticmethod
+    def _move_params(module: nn.Module, device: torch.device) -> None:
+        """Move module parameters and buffers to device.
+
+        This cls method specifically prevents recursion device movement,
+        E.g., Cache-DiT CachedBlocks has attr `transformer` as a ref to original
+        transformer blocks, thus `module.to(device)` will fail for recursion calling,
+        refer to
+        https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83
+        """
+        for p in module.parameters():
+            if p.data.device != device:
+                p.data = p.data.to(device, non_blocking=True)
+        for b in module.buffers():
+            if b.device != device:
+                b.data = b.data.to(device, non_blocking=True)
+
     def _to_cpu(self, module: nn.Module) -> None:
         """Move module to CPU."""
         try:
@@ -42,13 +59,12 @@ class SequentialOffloadHook(ModelHook):
         except StopIteration:
             return
 
-        previous_device = param.device
         # Skip if already on CPU
-        if previous_device.type == "cpu":
+        if param.device.type == "cpu":
             return
 
-        module.to("cpu")
-        torch.cuda.empty_cache()
+        self._move_params(module, torch.device("cpu"))
+        current_omni_platform.empty_cache()
 
         if self.pin_memory:
             for p in module.parameters():
@@ -64,7 +80,7 @@ class SequentialOffloadHook(ModelHook):
         except StopIteration:
             return
 
-        module.to(self.device)
+        self._move_params(module, self.device)
 
     def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
         # Offload target modules to CPU
@@ -186,19 +202,6 @@ class ModelLevelOffloadBackend(OffloadBackend):
             except Exception as exc:
                 logger.debug("Failed to move VAE to GPU: %s", exc)
 
-        # Initial state: keep DiT modules on CPU (encoders typically run first)
-        # TODO: This part seems to be unnecessary, remove it after testing
-        # for dit_mod in modules.dits:
-        #     dit_mod.to("cpu")
-
-        # torch.cuda.empty_cache()
-
-        # if self.config.pin_cpu_memory:
-        #     for dit_mod in modules.dits:
-        #         for p in dit_mod.parameters():
-        #             if p.data.device.type == "cpu" and not p.data.is_pinned():
-        #                 p.data = p.data.pin_memory()
-
         # Apply sequential offloading hooks
         apply_sequential_offload(
             dit_modules=modules.dits,

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec
+from .utils.kv_utils import normalize_layer_kv
 
 logger = init_logger(__name__)
 
@@ -200,8 +201,12 @@ class OmniKVTransferManager:
                     logger.warning(f"Request {req_id} has no block IDs, skipping")
                     continue
 
+                custom_metadata = data.get("custom_metadata")
+
                 # Extract KV cache from GPU blocks -> CPU tensors
-                kv_data = self._extract_kv_cache(req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype)
+                kv_data = self._extract_kv_cache(
+                    req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata
+                )
                 if kv_data:
                     # Resolve global request ID if available
                     transfer_req_id = request_id_resolver(req_id) if request_id_resolver else req_id
@@ -224,6 +229,7 @@ class OmniKVTransferManager:
         kv_caches: list[LayerKV],
         block_size: int,
         cache_dtype: str,
+        custom_metadata: dict[str, Any] | None = None,
     ) -> KVCacheTransferData | None:
         """Extract KV cache from GPU blocks for a single request.
 
@@ -234,6 +240,7 @@ class OmniKVTransferManager:
             kv_caches: List of KV cache (tensor or tuple) per layer
             block_size: Size of each cache block
             cache_dtype: Data type of the cache
+            custom_metadata: Optional custom metadata to include
 
         Note: If key/value block counts differ, extraction uses only the overlapping
         block range. Extra key/value blocks are ignored, so returned KV may be partial.
@@ -246,7 +253,7 @@ class OmniKVTransferManager:
         value_cache: list[torch.Tensor | None] = [None] * num_layers
 
         for layer_idx, layer_kv in enumerate(kv_caches):
-            kv_pair = self._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
+            kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
             if kv_pair is None:
                 continue
             key_blocks, value_blocks = kv_pair
@@ -289,57 +296,10 @@ class OmniKVTransferManager:
                 "num_layers": num_layers,
                 "dtype": str(cache_dtype),
                 "seq_len": seq_len,
+                **(custom_metadata or {}),
             },
         )
 
-    def _normalize_layer_kv(
-        self,
-        layer_kv: LayerKV,
-        req_id: str,
-        layer_idx: int,
-    ) -> tuple[torch.Tensor, torch.Tensor] | None:
-        """Normalize one layer KV cache to a `(key_blocks, value_blocks)` tuple.
-
-        Args:
-            layer_kv: The raw KV cache (tensor or tuple) for the layer
-            req_id: Request ID for logging
-            layer_idx: Layer index for logging
-
-        Returns:
-            Tuple of (key_blocks, value_blocks) if valid, None otherwise
-        """
-        if isinstance(layer_kv, torch.Tensor):
-            if layer_kv.ndim < 3 or layer_kv.shape[0] != 2:
-                logger.warning(
-                    f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
-                    f"expected [2, blocks, block_size, ...], got {tuple(layer_kv.shape)}"
-                )
-                return None
-            key_blocks = layer_kv[0]
-            value_blocks = layer_kv[1]
-        elif isinstance(layer_kv, tuple):
-            if len(layer_kv) != 2:
-                logger.warning(
-                    f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)"
-                )
-                return None
-            key_blocks, value_blocks = layer_kv
-            if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
-                logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
-                return None
-        else:
-            logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
-            return None
-        # ensure key/value blocks are at least 2D for block indexing
-        if key_blocks.ndim < 2 or value_blocks.ndim < 2:
-            logger.warning(
-                f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
-                f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
-            )
-            return None
-
-        return key_blocks, value_blocks
-
     def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None:
         """Transfer KV cache data to downstream stage via OmniConnector.
 
@@ -496,6 +456,8 @@ class OmniKVTransferManager:
 
         if "metadata" in data:
             req.kv_metadata = data["metadata"]
+            if hasattr(req, "sampling_params") and req.sampling_params is not None:
+                req.sampling_params.kv_metadata = data["metadata"]
 
     # Legacy compatibility method
     def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:

vllm_omni/distributed/omni_connectors/utils/kv_utils.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Utility helpers for KV cache manipulation."""
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+
+
+def normalize_layer_kv(
+    layer_kv: LayerKV,
+    *,
+    req_id: str = "",
+    layer_idx: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Normalize one layer KV cache to a ``(key_blocks, value_blocks)`` tuple.
+
+    In vLLM, different attention backends return paged-attention KV blocks
+    with different layouts.  For example:
+
+    * **FlashAttention** (``vllm/v1/attention/backends/flash_attn.py``)
+      returns shape ``(2, num_blocks, block_size, num_kv_heads, head_size)``
+      – the key/value dimension is at **dim 0**.
+    * **FlashInfer** (``vllm/v1/attention/backends/flashinfer.py``)
+      returns shape ``(num_blocks, 2, block_size, num_kv_heads, head_size)``
+      – the key/value dimension is at **dim 1**.
+
+    This utility handles both layouts (and the tuple case) so that
+    downstream code can work with any backend.
+
+    Supported layouts:
+
+    * **Stacked tensor** ``[2, num_blocks, block_size, n_heads, head_dim]`` –
+      dim-0 selects key / value.
+    * **Stacked tensor** ``[num_blocks, 2, block_size, n_heads, head_dim]`` –
+      dim-1 selects key / value.
+    * **Tuple** ``(key_tensor, value_tensor)`` – returned as-is after
+      validation.
+
+    Args:
+        layer_kv: The raw KV cache (tensor or tuple) for the layer.
+        req_id: Request ID used only for diagnostic log messages.
+        layer_idx: Layer index used only for diagnostic log messages.
+
+    Returns:
+        ``(key_blocks, value_blocks)`` if *layer_kv* is valid, ``None``
+        otherwise.
+    """
+    if isinstance(layer_kv, torch.Tensor):
+        if layer_kv.ndim >= 3 and layer_kv.shape[0] == 2:
+            key_blocks = layer_kv[0]
+            value_blocks = layer_kv[1]
+        elif layer_kv.ndim >= 3 and layer_kv.shape[1] == 2:
+            key_blocks = layer_kv[:, 0]
+            value_blocks = layer_kv[:, 1]
+        else:
+            logger.warning(
+                f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
+                f"expected [2, ...] or [..., 2, ...] at dim 0/1, got {tuple(layer_kv.shape)}"
+            )
+            return None
+    elif isinstance(layer_kv, tuple):
+        if len(layer_kv) != 2:
+            logger.warning(f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)")
+            return None
+        key_blocks, value_blocks = layer_kv
+        if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
+            logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
+            return None
+    else:
+        logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
+        return None
+
+    # ensure key/value blocks are at least 2D for block indexing
+    if key_blocks.ndim < 2 or value_blocks.ndim < 2:
+        logger.warning(
+            f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
+            f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
+        )
+        return None
+
+    return key_blocks, value_blocks

vllm_omni/entrypoints/omni_llm.py

@@ -28,6 +28,7 @@ from vllm_omni.entrypoints.utils import (
     load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
+    resolve_model_type,
 )
 
 logger = init_logger(__name__)
@@ -84,12 +85,13 @@ class OmniLLM(LLM):
         self.worker_backend = kwargs.get("worker_backend", "multi_process")
         self.ray_address = kwargs.get("ray_address", None)
         self.batch_timeout = batch_timeout
+        self.model_type = resolve_model_type(model)
         self.log_stats: bool = bool(log_stats)
 
         # Load stage configurations
         if stage_configs_path is None:
-            self.config_path = resolve_model_config_path(model)
-            self.stage_configs = load_stage_configs_from_model(model)
+            self.config_path = resolve_model_config_path(self.model_type)
+            self.stage_configs = load_stage_configs_from_model(config_path=self.config_path)
         else:
             self.config_path = stage_configs_path
             self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)

vllm_omni/entrypoints/openai/api_server.py

@@ -87,7 +87,10 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
     encode_image_base64,
     parse_size,
 )
-from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    OpenAICreateAudioGenerateRequest,
+    OpenAICreateSpeechRequest,
+)
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
     ImageGenerationRequest,
@@ -97,6 +100,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoGenerationRequest,
     VideoGenerationResponse,
 )
+from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServingAudioGenerate
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
@@ -175,8 +179,29 @@ class _DiffusionServingModels:
     provide a lightweight fallback.
     """
 
+    class _NullModelConfig:
+        def __getattr__(self, name):
+            return None
+
+    class _Unsupported:
+        def __init__(self, name: str):
+            self.name = name
+
+        def __call__(self, *args, **kwargs):
+            raise NotImplementedError(f"{self.name} is not supported in diffusion mode")
+
+        def __getattr__(self, attr):
+            raise NotImplementedError(f"{self.name}.{attr} is not supported in diffusion mode")
+
     def __init__(self, base_model_paths: list[BaseModelPath]) -> None:
         self._base_model_paths = base_model_paths
+        self.model_config = self._NullModelConfig()
+
+    def __getattr__(self, name):
+        """Return a sentinel that raises NotImplementedError if called or
+        accessed, so any use of unsupported OpenAIServingModels features in
+        diffusion mode fails loudly with a descriptive message."""
+        return self._Unsupported(name)
 
     async def show_available_models(self) -> ModelList:
         return ModelList(
@@ -434,10 +459,10 @@ async def omni_init_app_state(
 
     # For omni models
     state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
+    model_name = served_model_names[0] if served_model_names else args.model
 
     # Pure Diffusion mode: use simplified initialization logic
     if is_pure_diffusion:
-        model_name = served_model_names[0] if served_model_names else args.model
         state.vllm_config = None
         state.diffusion_engine = engine_client
         state.openai_serving_models = _DiffusionServingModels(base_model_paths)
@@ -450,6 +475,16 @@ async def omni_init_app_state(
             model_name=model_name,
         )
 
+        # audio related
+        state.openai_serving_speech = None
+        state.openai_serving_audio_generate = OmniOpenAIServingAudioGenerate.for_diffusion(
+            engine_client,
+            state.openai_serving_models,
+            request_logger=request_logger,
+            model_name=model_name,
+        )
+
+        # video related
         diffusion_stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
         state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
             diffusion_engine=engine_client,  # type: ignore
@@ -457,7 +492,6 @@ async def omni_init_app_state(
             stage_configs=diffusion_stage_configs,
         )
 
-        state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False)
         state.server_load_metrics = 0
         logger.info("Pure diffusion API server initialized for model: %s", model_name)
         return
@@ -732,7 +766,11 @@ async def omni_init_app_state(
     )
 
     state.openai_serving_speech = OmniOpenAIServingSpeech(
-        engine_client, state.openai_serving_models, request_logger=request_logger
+        engine_client, state.openai_serving_models, request_logger=request_logger, model_name=model_name
+    )
+
+    state.openai_serving_audio_generate = OmniOpenAIServingAudioGenerate(
+        engine_client, state.openai_serving_models, request_logger=request_logger, model_name=model_name
     )
 
     state.openai_serving_video = OmniOpenAIServingVideo(
@@ -757,6 +795,10 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
 
+def OmniAudioGenerate(request: Request) -> OmniOpenAIServingAudioGenerate | None:
+    return getattr(request.app.state, "openai_serving_audio_generate", None)
+
+
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(validate_json_request)],
@@ -867,6 +909,34 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
 
 
+@router.post(
+    "/v1/audio/generate",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.OK.value: {"content": {"audio/*": {}}},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+@with_cancellation
+@load_aware_call
+async def create_audio_generate(request: OpenAICreateAudioGenerateRequest, raw_request: Request):
+    handler = OmniAudioGenerate(raw_request)
+    if handler is None:
+        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        if base_server is None:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND.value,
+                detail="The model does not support Audio Generate API",
+            )
+        return base_server.create_error_response(message="The model does not support Audio Generate API")
+    try:
+        return await handler.create_audio_generate(request, raw_request)
+    except Exception as e:
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
+
+
 @router.get(
     "/v1/audio/voices",
     responses={

vllm_omni/entrypoints/openai/audio_utils_mixin.py

@@ -45,6 +45,10 @@ class AudioMixin:
                 "Only mono (1D) and stereo (2D) are supported."
             )
 
+        if audio_tensor.ndim == 2 and audio_tensor.shape[0] == 2:
+            # Convert from [channels, samples] to [samples, channels]
+            audio_tensor = audio_tensor.T
+
         audio_tensor, sample_rate = self._apply_speed_adjustment(audio_tensor, speed, sample_rate)
 
         supported_formats = {

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -85,6 +85,53 @@ class OpenAICreateSpeechRequest(BaseModel):
         return self
 
 
+class OpenAICreateAudioGenerateRequest(BaseModel):
+    """Request model for audio generation via diffusion models (e.g. Stable Audio)."""
+
+    input: str = Field(
+        description="Text prompt describing the audio to generate",
+    )
+    model: str | None = None
+    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    speed: float | None = Field(
+        default=1.0,
+        ge=0.25,
+        le=4.0,
+    )
+    stream_format: Literal["sse", "audio"] | None = "audio"
+    audio_length: float | None = Field(
+        default=None,
+        description="Audio length in seconds",
+    )
+    audio_start: float | None = Field(
+        default=0.0,
+        description="Audio start time in seconds",
+    )
+    negative_prompt: str | None = Field(
+        default=None,
+        description="Negative prompt for classifier-free guidance",
+    )
+    guidance_scale: float | None = Field(
+        default=None,
+        description="Guidance scale for diffusion models",
+    )
+    num_inference_steps: int | None = Field(
+        default=None,
+        description="Number of inference steps",
+    )
+    seed: int | None = Field(
+        default=None,
+        description="Random seed for reproducibility",
+    )
+
+    @field_validator("stream_format")
+    @classmethod
+    def validate_stream_format(cls, v: str) -> str:
+        if v == "sse":
+            raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.")
+        return v
+
+
 class CreateAudio(BaseModel):
     audio_tensor: np.ndarray
     sample_rate: int = 24000

vllm_omni/entrypoints/openai/serving_audio_generate.py

@@ -0,0 +1,168 @@
+import asyncio
+
+import torch
+from fastapi import Request
+from fastapi.responses import Response
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
+from vllm.logger import init_logger
+from vllm.utils import random_uuid
+
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    AudioResponse,
+    CreateAudio,
+    OpenAICreateAudioGenerateRequest,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+logger = init_logger(__name__)
+
+
+class OmniOpenAIServingAudioGenerate(OpenAIServing, AudioMixin):
+    """Serving class for audio generation via diffusion models (e.g. Stable Audio)."""
+
+    def __init__(self, *args, **kwargs):
+        self.model_name = kwargs.pop("model_name", None)
+        super().__init__(*args, **kwargs)
+        self.diffusion_mode = False
+
+    @classmethod
+    def for_diffusion(cls, *args, **kwargs) -> "OmniOpenAIServingAudioGenerate":
+        """Create an instance configured to run in diffusion mode."""
+        instance = cls(*args, **kwargs)
+        instance.diffusion_mode = True
+        return instance
+
+    def _is_stable_audio_model(self) -> bool:
+        return self.engine_client.model_type == "StableAudioPipeline"
+
+    async def create_audio_generate(
+        self,
+        request: OpenAICreateAudioGenerateRequest,
+        raw_request: Request | None = None,
+    ):
+        """
+        Generate audio using diffusion-based models (e.g. Stable Audio).
+
+        This endpoint is designed for audio generation models as
+        opposed to TTS models that specifically generate speech.
+        """
+
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            logger.error("Error with model %s", error_check_ret)
+            return error_check_ret
+
+        if self.engine_client.errored:
+            raise self.engine_client.dead_error
+
+        request_id = f"audiogen-{random_uuid()}"
+
+        try:
+            default_sr = 44100  # Default sample rate for Stable Audio
+
+            # Build prompt for diffusion audio generation
+            prompt = {
+                "prompt": request.input,
+            }
+            if request.negative_prompt:
+                prompt["negative_prompt"] = request.negative_prompt
+
+            # Build sampling params for diffusion
+            sampling_params_list = [OmniDiffusionSamplingParams(num_outputs_per_prompt=1)]
+
+            # Create generator if seed provided
+            if request.seed is not None:
+                from vllm_omni.platforms import current_omni_platform
+
+                generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request.seed)
+                sampling_params_list[0].generator = generator
+
+            if request.guidance_scale is not None:
+                sampling_params_list[0].guidance_scale = request.guidance_scale
+
+            if request.num_inference_steps is not None:
+                sampling_params_list[0].num_inference_steps = request.num_inference_steps
+
+            # Set up audio duration parameters
+            if request.audio_length is not None:
+                audio_length = request.audio_length
+                audio_start = request.audio_start if request.audio_start is not None else 0.0
+                audio_end_in_s = audio_start + audio_length
+                sampling_params_list[0].extra_args = {
+                    "audio_start_in_s": audio_start,
+                    "audio_end_in_s": audio_end_in_s,
+                }
+
+            logger.info(
+                "Audio generation request %s: prompt=%r",
+                request_id,
+                request.input[:50] + "..." if len(request.input) > 50 else request.input,
+            )
+
+            generator = self.engine_client.generate(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params_list=sampling_params_list,
+                output_modalities=["audio"],
+            )
+
+            final_output: OmniRequestOutput | None = None
+            async for res in generator:
+                final_output = res
+
+            if final_output is None:
+                return self.create_error_response("No output generated from the model.")
+
+            # Extract audio from output
+            audio_output = None
+            if hasattr(final_output, "multimodal_output") and final_output.multimodal_output:
+                audio_output = final_output.multimodal_output
+            if not audio_output and hasattr(final_output, "request_output"):
+                if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"):
+                    audio_output = final_output.request_output.multimodal_output
+
+            # Check for audio data using either "audio" or "model_outputs" key
+            audio_key = None
+            if audio_output:
+                if "audio" in audio_output:
+                    audio_key = "audio"
+                elif "model_outputs" in audio_output:
+                    audio_key = "model_outputs"
+
+            if not audio_output or audio_key is None:
+                return self.create_error_response("Audio generation model did not produce audio output.")
+
+            audio_tensor = audio_output[audio_key]
+            sample_rate = audio_output.get("sr", default_sr)
+            if hasattr(sample_rate, "item"):
+                sample_rate = sample_rate.item()
+
+            # Convert tensor to numpy
+            if hasattr(audio_tensor, "float"):
+                audio_tensor = audio_tensor.float().detach().cpu().numpy()
+
+            # Squeeze batch dimension if present, but preserve channel dimension for stereo
+            if audio_tensor.ndim > 1:
+                audio_tensor = audio_tensor.squeeze()
+
+            audio_obj = CreateAudio(
+                audio_tensor=audio_tensor,
+                sample_rate=int(sample_rate),
+                response_format=request.response_format or "wav",
+                speed=request.speed or 1.0,
+                stream_format=request.stream_format,
+                base64_encode=False,
+            )
+
+            audio_response: AudioResponse = self.create_audio(audio_obj)
+            return Response(content=audio_response.audio_data, media_type=audio_response.media_type)
+
+        except asyncio.CancelledError:
+            return self.create_error_response("Client disconnected")
+        except ValueError as e:
+            return self.create_error_response(e)
+        except Exception as e:
+            logger.exception("Audio generation failed: %s", e)
+            return self.create_error_response(f"Audio generation failed: {e}")

vllm_omni/entrypoints/openai/serving_speech.py

@@ -11,6 +11,7 @@ from typing import Any
 import numpy as np
 from fastapi import Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
+from transformers.utils.hub import cached_file
 from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
@@ -83,6 +84,7 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:
 
 class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     def __init__(self, *args, **kwargs):
+        self.model_name = kwargs.pop("model_name", None)
         super().__init__(*args, **kwargs)
         # Initialize uploaded speakers storage
         speech_voice_samples_dir = os.environ.get("SPEECH_VOICE_SAMPLES", "/tmp/voice_samples")
@@ -120,7 +122,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         try:
             model_path = self.engine_client.model_config.model
             st_config_path = os.path.join(model_path, "speech_tokenizer", "config.json")
-            if os.path.exists(st_config_path):
+            if not os.path.exists(st_config_path):
+                st_config_path = cached_file(model_path, "speech_tokenizer/config.json")
+            if st_config_path is not None and os.path.exists(st_config_path):
                 with open(st_config_path) as f:
                     st_config = json.load(f)
                 output_sr = st_config.get("output_sample_rate")
@@ -753,6 +757,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id = f"speech-{random_uuid()}"
 
         try:
+            sampling_params_list = self.engine_client.default_sampling_params_list
+            default_sr = 24000  # Default sample rate for TTS models
             if self._is_tts:
                 # Validate TTS parameters
                 validation_error = self._validate_tts_request(request)
@@ -778,8 +784,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 tts_params.get("task_type", ["unknown"])[0],
             )
 
-            sampling_params_list = self.engine_client.default_sampling_params_list
-
             generator = self.engine_client.generate(
                 prompt=prompt,
                 request_id=request_id,
@@ -806,7 +810,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 return self.create_error_response("TTS model did not produce audio output.")
 
             audio_tensor = audio_output[audio_key]
-            sr_raw = audio_output.get("sr", 24000)
+            sr_raw = audio_output.get("sr", default_sr)
             sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
             sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
 

vllm_omni/entrypoints/utils.py

@@ -169,22 +169,46 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     return obj
 
 
-def resolve_model_config_path(model: str) -> str:
-    """Resolve the stage config file path from the model name.
+def resolve_model_config_path(model_type: str) -> str | None:
+    """Resolve the stage config file path from the model type.
 
     Resolves stage configuration path based on the model type and device type.
     First tries to find a device-specific YAML file from stage_configs/{device_type}/
     directory. If not found, falls back to the default config file.
 
     Args:
-        model: Model name or path (used to determine model_type)
+        model_type: Model type string
 
     Returns:
-        String path to the stage configuration file
+        String path to the stage configuration file if found, None otherwise
+    """
+    default_config_path = current_omni_platform.get_default_stage_config_path()
+    config_file_name = f"{model_type}.yaml"
+    complete_config_path = PROJECT_ROOT / default_config_path / config_file_name
+    if os.path.exists(complete_config_path):
+        return str(complete_config_path)
+
+    # Fall back to default config
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        return None
+    return str(stage_config_path)
+
+
+def resolve_model_type(model: str) -> str:
+    """Resolve the model type from the model name.
+
+    Args:
+        model: Model name or path
+
+    Returns:
+        Model type string (e.g. ``"Qwen3TTSForConditionalGeneration"``,
+        ``"StableAudioPipeline"``).
 
     Raises:
-        ValueError: If model_type cannot be determined
-        FileNotFoundError: If no stage config file exists for the model type
+        ValueError: If the model type cannot be determined from any
+            available configuration file.
     """
     # Try to get config from standard transformers format first
     try:
@@ -217,42 +241,26 @@ def resolve_model_config_path(model: str) -> str:
                 f"Please ensure the model has proper configuration files with 'model_type' field"
             )
 
-    default_config_path = current_omni_platform.get_default_stage_config_path()
-    model_type_str = f"{model_type}.yaml"
-    complete_config_path = PROJECT_ROOT / default_config_path / model_type_str
-    if os.path.exists(complete_config_path):
-        return str(complete_config_path)
-
-    # Fall back to default config
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if not os.path.exists(stage_config_path):
-        return None
-    return str(stage_config_path)
+    return model_type
 
 
-def load_stage_configs_from_model(model: str, base_engine_args: dict | None = None) -> list:
-    """Load stage configurations from model's default config file.
-
-    Loads stage configurations based on the model type and device type.
-    First tries to load a device-specific YAML file from stage_configs/{device_type}/
-    directory. If not found, falls back to the default config file.
+def load_stage_configs_from_model(config_path: str | None, base_engine_args: dict | None = None) -> list:
+    """Load stage configurations from a resolved config file path.
 
     Args:
-        model: Model name or path (used to determine model_type)
+        config_path: Path to the YAML configuration file, or None.
+            When None, returns an empty list.
+        base_engine_args: Optional engine arguments to merge with stage configs.
 
     Returns:
-        List of stage configuration dictionaries
-
-    Raises:
-        FileNotFoundError: If no stage config file exists for the model type
+        List of stage configuration dictionaries, or empty list if
+        config_path is None.
     """
     if base_engine_args is None:
         base_engine_args = {}
-    stage_config_path = resolve_model_config_path(model)
-    if stage_config_path is None:
+    if config_path is None:
         return []
-    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path, base_engine_args=base_engine_args)
+    stage_configs = load_stage_configs_from_yaml(config_path=config_path, base_engine_args=base_engine_args)
     return stage_configs
 
 
@@ -306,9 +314,10 @@ def load_and_resolve_stage_configs(
     Returns:
         Tuple of (config_path, stage_configs)
     """
+    model_type = resolve_model_type(model)
     if stage_configs_path is None:
-        config_path = resolve_model_config_path(model)
-        stage_configs = load_stage_configs_from_model(model, base_engine_args=kwargs)
+        config_path = resolve_model_config_path(model_type)
+        stage_configs = load_stage_configs_from_model(config_path, base_engine_args=kwargs)
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()

vllm_omni/model_executor/models/__init__.py

@@ -1,4 +1,8 @@
+from .bagel.bagel import OmniBagelForConditionalGeneration
 from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
 from .registry import OmniModelRegistry  # noqa: F401
 
-__all__ = ["Qwen3OmniMoeForConditionalGeneration"]
+__all__ = [
+    "Qwen3OmniMoeForConditionalGeneration",
+    "OmniBagelForConditionalGeneration",
+]

vllm_omni/model_executor/models/bagel/__init__.py

@@ -0,0 +1,3 @@
+from .bagel import OmniBagelForConditionalGeneration
+
+__all__ = ["OmniBagelForConditionalGeneration"]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -447,6 +447,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
         audio_codes_list: list[torch.Tensor] = []
         ref_code_len_list: list[torch.Tensor] = []
+        ref_code_tensor: torch.Tensor | None = None
         codec_streaming_list: list[torch.Tensor] = []
         for info in info_dicts:
             if not isinstance(info, dict):
@@ -459,6 +460,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     codec_streaming_list.append(
                         torch.full((int(ac.shape[0]),), int(cs), dtype=torch.int8, device=ac.device)
                     )
+            ref_code = info.get("ref_code")
+            if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+                ref_code_tensor = ref_code
             ref_len = info.get("ref_code_len")
             if ref_len is None:
                 continue
@@ -487,6 +491,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         mm: dict[str, torch.Tensor] = {"audio_codes": audio_codes}
         if ref_code_len_list:
             mm["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len]
+        if ref_code_tensor is not None:
+            mm["ref_code"] = [ref_code_tensor]
         if codec_streaming_list:
             mm["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len]
         return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
@@ -538,8 +544,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round.
             is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2
             if is_first_prefill:
-                full_prompt_embeds, tailing_text_hidden, tts_pad_embed, ref_code_len = self._build_prompt_embeds(
-                    task_type=task_type, info_dict=info_dict
+                full_prompt_embeds, tailing_text_hidden, tts_pad_embed, ref_code_len, ref_code = (
+                    self._build_prompt_embeds(task_type=task_type, info_dict=info_dict)
                 )
                 # Store full prompt embeddings + trailing queue on CPU for later chunks/steps.
                 prompt_embeds_cpu = full_prompt_embeds.detach().to("cpu").contiguous()
@@ -550,6 +556,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     "talker_prefill_offset": 0,
                     "codec_streaming": codec_streaming,
                 }
+                if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+                    info_update["ref_code"] = ref_code.detach().to("cpu").contiguous()
                 if ref_code_len is not None:
                     info_update["ref_code_len"] = int(ref_code_len)
                 # Always return a span_len slice; if the scheduled placeholder is longer, pad with tts_pad_embed.
@@ -773,7 +781,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
                 if ref_code_len is None and estimate_ref_code_len is not None:
                     ref_code_len = estimate_ref_code_len(info.get("ref_audio"))
-
                 if ref_code_len is None:
                     raise ValueError(
                         "Base in-context voice cloning requires either `voice_clone_prompt.ref_code` "
@@ -1183,7 +1190,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         *,
         task_type: str,
         info_dict: dict[str, Any],
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int | None]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int | None, torch.Tensor | None]:
         text = (info_dict.get("text") or [""])[0]
         language = (info_dict.get("language") or ["Auto"])[0]
         non_streaming_mode_val = info_dict.get("non_streaming_mode")
@@ -1268,6 +1275,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # Speaker embedding/token (task-dependent)
         speaker_embed = None
         ref_code_len: int | None = None
+        ref_code_prompt: torch.Tensor | None = None
 
         def _as_singleton(x: object) -> object:
             if isinstance(x, list):
@@ -1333,6 +1341,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 wav_np, sr = self._normalize_ref_audio(ref_audio_list[0])
                 ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device)
                 ref_code_len = int(ref_code_t.shape[0])
+            if isinstance(ref_code_t, torch.Tensor):
+                ref_code_prompt = ref_code_t
 
             # Speaker embedding: use prompt embed if provided; otherwise extract from audio.
             spk = None
@@ -1521,6 +1531,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             trailing_text_hidden.squeeze(0),  # [T, H]
             tts_pad_embed.squeeze(0),  # [1, H]
             ref_code_len,
+            ref_code_prompt.contiguous() if isinstance(ref_code_prompt, torch.Tensor) else None,
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm_omni/model_executor/models/registry.py

@@ -119,6 +119,11 @@ _OMNI_MODELS = {
         "glm_image_ar",
         "GlmImageForConditionalGeneration",
     ),
+    "OmniBagelForConditionalGeneration": (
+        "bagel",
+        "bagel",
+        "OmniBagelForConditionalGeneration",
+    ),
     "HunyuanImage3ForCausalMM": (
         "hunyuan_image3",
         "hunyuan_image3",

vllm_omni/model_executor/stage_configs/bagel.yaml

@@ -6,15 +6,14 @@ stage_args:
     prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
-      # 2 = 1 user prompt + 1 CFG companion (text-unconditional).
-      # For multi-prompt batches this should scale as batch_size × 2.
-      max_batch_size: 2
+      # 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional).
+      max_batch_size: 3
     engine_args:
       model_stage: thinker
-      model_arch: BagelForConditionalGeneration
+      model_arch: OmniBagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -46,7 +45,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml

@@ -11,7 +11,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -45,7 +45,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

vllm_omni/model_executor/stage_input_processors/bagel.py

@@ -20,6 +20,7 @@ from typing import Any
 logger = logging.getLogger(__name__)
 
 CFG_TEXT_SUFFIX = "__cfg_text"
+CFG_IMG_SUFFIX = "__cfg_img"
 
 
 @dataclass
@@ -73,7 +74,37 @@ def expand_cfg_prompts(
             ),
         ]
 
-    # img2img: more complex (3 distinct caches). Reserve for future.
+    if "img2img" in modalities:
+        IMG2IMG_PLACEHOLDER = "<|fim_middle|>"
+
+        cfg_text_dict: dict[str, Any] = {
+            "prompt": f"{IMG2IMG_PLACEHOLDER}{neg_prompt}",
+            "modalities": ["img2img"],
+        }
+        mm_data = prompt.get("multi_modal_data")
+        if mm_data:
+            cfg_text_dict["multi_modal_data"] = mm_data
+
+        original_text = prompt.get("prompt", "")
+        cfg_img_text = original_text.replace(IMG2IMG_PLACEHOLDER, "")
+        cfg_img_dict: dict[str, Any] = {
+            "prompt": cfg_img_text,
+            "modalities": ["img2img"],
+        }
+
+        return [
+            ExpandedPrompt(
+                prompt=cfg_text_dict,
+                role="cfg_text",
+                request_id_suffix=CFG_TEXT_SUFFIX,
+            ),
+            ExpandedPrompt(
+                prompt=cfg_img_dict,
+                role="cfg_img",
+                request_id_suffix=CFG_IMG_SUFFIX,
+            ),
+        ]
+
     return []
 
 

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -27,13 +27,24 @@ def talker2code2wav(
         # Filter zero-padded frames (EOS/invalid steps), matching _extract_last_frame behavior
         valid_mask = audio_codes.any(dim=1)
         audio_codes = audio_codes[valid_mask]
+        ref_code = output.multimodal_output.get("ref_code")
+        if isinstance(ref_code, list):
+            ref_code = ref_code[0] if ref_code else None
+        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+            ref_code = ref_code.to(torch.long).cpu().contiguous()
+            ref_code_len = int(ref_code.shape[0])
+            audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
+        else:
+            ref_code_len = 0
         # Code2Wav expects codebook-major flat: [Q*num_frames]
         codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
+        additional_information = {"left_context_size": [ref_code_len]} if ref_code_len > 0 else None
         code2wav_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=codec_codes,
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
+                additional_information=additional_information,
             )
         )
     return code2wav_inputs
@@ -61,12 +72,19 @@ def talker2code2wav_async_chunk(
 ) -> dict[str, Any] | None:
     request_id = request.external_req_id
     finished = bool(is_finished or request.is_finished())
+    request_payload = getattr(transfer_manager, "request_payload", None)
+    if request_payload is None:
+        request_payload = {}
+        transfer_manager.request_payload = request_payload
 
     if isinstance(pooling_output, dict):
         frame = _extract_last_frame(pooling_output)
         if frame is not None:
             codec_codes = frame.cpu().tolist()
             transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
+        ref_code = pooling_output.get("ref_code")
+        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None:
+            request_payload[request_id] = ref_code.to(torch.long).cpu().contiguous()
     elif not finished:
         # Some steps may not produce pooling_output. Only flush on finish.
         return None
@@ -139,6 +157,14 @@ def talker2code2wav_async_chunk(
         left_context_size = max(0, int(end_index - context_length))
         window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
 
+    # Match offline Base ICL decoding by prepending ref_code only once to the first decoder window.
+    if transfer_manager.put_req_chunk[request_id] == 0:
+        ref_code = request_payload.pop(request_id, None)
+        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+            ref_frames = ref_code.tolist()
+            window_frames = ref_frames + window_frames
+            left_context_size += len(ref_frames)
+
     # Pack context + chunk into codebook-major flat codes for adapter.
     code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist()
 

vllm_omni/patch.py

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

vllm_omni/worker/gpu_ar_model_runner.py

@@ -98,9 +98,25 @@ class GPUARModelRunner(OmniGPUModelRunner):
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
+        if not getattr(self, "_warmup_state_cleared", False):
+            self._warmup_state_cleared = True
+            if hasattr(self.model, "_clear_warmup_state"):
+                self.model._clear_warmup_state()
+
         # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
+        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
+        if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
+            for req_id, data in finished_reqs.items():
+                try:
+                    model_meta = self.model.get_kv_transfer_metadata(req_id)
+                    if model_meta:
+                        existing = data.get("custom_metadata") or {}
+                        existing.update(model_meta)
+                        data["custom_metadata"] = existing
+                except Exception as e:
+                    logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}")
         self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
-            finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
+            finished_reqs=finished_reqs,
             kv_caches=self.kv_caches,
             block_size=self.cache_config.block_size,
             cache_dtype=str(self.cache_config.cache_dtype),
@@ -290,6 +306,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 sampler=self.sampler,
             )
 
+            # [Omni] Map pending ropes metadata to req_ids.
+            if hasattr(self.model, "flush_pending_metadata"):
+                self.model.flush_pending_metadata(list(req_ids))
+
         with record_function_or_nullcontext("gpu_model_runner: postprocess"):
             if self.use_aux_hidden_state_outputs:
                 # True when EAGLE 3 is used.