Merge branch 'main' into dev/rebase-v0170
This commit is contained in:
40 files changed
+3270
-69
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
```
|
||||
@@ -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"
|
||||
``````
|
||||
@@ -0,0 +1,6 @@
|
||||
*.png
|
||||
.git/
|
||||
ct.yaml
|
||||
lintconf.yaml
|
||||
values.schema.json
|
||||
/workflows
|
||||
@@ -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
|
||||
@@ -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 .
|
||||
```
|
||||
@@ -0,0 +1,3 @@
|
||||
chart-dirs:
|
||||
- charts
|
||||
validate-maintainers: false
|
||||
@@ -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
|
||||
@@ -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 }}
|
||||
@@ -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 -}}
|
||||
@@ -0,0 +1,6 @@
|
||||
{{- if .Values.customObjects }}
|
||||
{{- range .Values.customObjects }}
|
||||
{{- tpl (. | toYaml) $ }}
|
||||
---
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
+54
@@ -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!"
|
||||
+170
@@ -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()
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
]
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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={
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in new issue
Block a user