anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Migrate audiox examples (#4607)

Browse Source 
Signed-off-by: zzehli <jaeli_ottawa@outlook.com>
This commit is contained in:
Jae authored and GitHub committed 2026-06-30 04:41:30 +00:00
1 parent 2fc00efa49
commit 0899a1a666
11 files changed
+335 -227
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

examples/offline_inference/audiox/README.md
-40
View File
@@ -1,40 +0,0 @@
# AudioX offline inference
Generate audio with the [AudioX](https://zeyuet.github.io/AudioX/) MMDiT diffusion
pipeline (`AudioXPipeline`). Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`.
## Prerequisites
Download a vLLM-Omni weight bundle (component-sharded safetensors):
```bash
huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
```
The Hugging Face id `zhangj1an/AudioX` also works directly without prefetching.
## Usage
```bash
# Text-to-audio only (default uses zhangj1an/AudioX from the Hub):
python end2end.py --tasks t2a
# All six tasks against a local bundle and a sample video for v2*/tv2*:
python end2end.py \
--model ./audiox_weights \
--video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4
# Subset of tasks, custom seed and steps:
python end2end.py --tasks t2a tv2a --num-inference-steps 100 --seed 0
```
## Arguments
- `--model`: HF id or local bundle path (default: `zhangj1an/AudioX`).
- `--tasks`: any subset of `t2a t2m v2a v2m tv2a tv2m` (default: all).
- `--video`: video file/URL — required for `v2*` and `tv2*`.
- `--reference-audio`: optional audio prompt (audio-conditioned generation).
- `--num-inference-steps`, `--guidance-scale`, `--seed`, `--seconds-total`,
`--sample-rate`, `--output-dir`: generation knobs.
Outputs land in `<output-dir>/<task>.wav` as 16-bit stereo WAV.
examples/offline_inference/audiox/end2end.py
-110
View File
@@ -1,110 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end AudioX offline example covering the 6 t2*/v2*/tv2* tasks.
Provide a directory with the **vLLM-Omni AudioX safetensors bundle** (e.g. from
``zhangj1an/AudioX`` on Hugging Face)::
huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
python end2end.py --model ./audiox_weights
python end2end.py --model ./audiox_weights --tasks t2a tv2a
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
import soundfile
import torch
import torchaudio.functional as TF
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
ROOT = Path(__file__).resolve().parent
SAMPLE_PROMPTS: dict[str, str] = {
"t2a": "Fireworks burst twice, followed by a period of silence before a clock begins ticking.",
"t2m": "Uplifting ukulele tune for a travel vlog",
"v2a": "",
"v2m": "",
"tv2a": "drum beating sound and human talking",
"tv2m": "uplifting music matching the scene",
}
ALL_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"})
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="AudioX offline end-to-end (6 t2*/v2*/tv2* tasks).")
p.add_argument("--model", default="zhangj1an/AudioX", help="HF id or local AudioX bundle path.")
p.add_argument("--tasks", nargs="+", default=list(ALL_TASKS), choices=ALL_TASKS)
p.add_argument("--video", default="", help="Video path / URL (required for v2*/tv2*).")
p.add_argument("--reference-audio", default="", help="Optional audio prompt for audio-conditioned generation.")
p.add_argument("--output-dir", default=str(ROOT / "audiox_task_outputs"))
p.add_argument("--num-inference-steps", type=int, default=250)
p.add_argument("--seconds-total", type=float, default=10.0)
p.add_argument("--guidance-scale", type=float, default=6.0)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--sample-rate", type=int, default=48000, help="Output WAV rate (resampled if != model rate).")
return p.parse_args()
def save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
"""Write 16-bit PCM WAV. ``audio`` is ``[channels, samples]`` float in [-1, 1]."""
path.parent.mkdir(parents=True, exist_ok=True)
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
def main() -> None:
args = parse_args()
omni = Omni(model=args.model, model_class_name="AudioXPipeline")
for task in args.tasks:
if task in VIDEO_TASKS and not args.video:
raise SystemExit(f"task={task!r} requires --video")
prompt = SAMPLE_PROMPTS[task] if task in TEXT_TASKS else ""
extra: dict = {"audiox_task": task, "seconds_start": 0.0, "seconds_total": float(args.seconds_total)}
if task in VIDEO_TASKS:
extra["video_path"] = args.video
if args.reference_audio:
extra["audio_path"] = args.reference_audio
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
t0 = time.perf_counter()
outputs = omni.generate(
prompt,
OmniDiffusionSamplingParams(
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
seed=args.seed,
extra_args=extra,
),
)
audio = outputs[0].request_output.multimodal_output.get("audio")
if audio is None:
raise RuntimeError(f"No audio produced for task {task!r}")
audio = torch.as_tensor(audio).detach().cpu().float()
if audio.ndim == 3:
audio = audio[0]
model_sr = int(outputs[0].request_output.multimodal_output.get("audio_sample_rate") or 44100)
if model_sr != args.sample_rate:
audio = TF.resample(audio, model_sr, args.sample_rate)
out_path = Path(args.output_dir) / f"{task}.wav"
save_wav(audio, out_path, args.sample_rate)
print(f"[{task}] saved {out_path} ({time.perf_counter() - t0:.2f}s)")
omni.close()
if __name__ == "__main__":
main()
examples/offline_inference/text_to_audio/README.md
+40 -1
View File
@@ -1,5 +1,12 @@
# Text-To-Audio
A unified script for text/video-to-audio generation. Supported models:
| Model | Tasks | Notes |
|-------|-------|-------|
| `stabilityai/stable-audio-open-1.0` | text-to-audio | gated; uses `--audio-length` |
| `zhangj1an/AudioX` | `t2a` / `t2m` / `v2a` / `v2m` / `tv2a` / `tv2m` | pass `--task`; video tasks need `--video` |
The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
## Prerequisites
@@ -43,13 +50,45 @@ python text_to_audio.py \
--output stable_audio_output.wav
```
### AudioX
AudioX supports six tasks. Sampler and reference knobs (declared in
`vllm_omni/model_extras/audiox.py`) are passed via the generic `--extra-body`
JSON flag, routed into sampling `extra_args`.
Text tasks (`t2a` / `t2m`):
```bash
python text_to_audio.py \
--model zhangj1an/AudioX --task t2a \
--prompt "Fireworks burst twice, followed by a clock ticking." \
--num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 --seed 42 \
--extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \
--output t2a.wav
```
Video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) require `--video`:
```bash
python text_to_audio.py \
--model zhangj1an/AudioX --task tv2a \
--prompt "drum beating sound and human talking" \
--video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \
--num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
--output tv2a.wav
```
Key arguments:
- `--prompt`: text description (string).
- `--task`: [AudioX] one of `t2a`/`t2m`/`v2a`/`v2m`/`tv2a`/`tv2m`.
- `--video`: [AudioX `v2*`/`tv2*`] video file/URL for conditioning (→ `video_path`).
- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s` for Stable Audio, `seconds_start` for AudioX).
- `--audio-length`: audio duration in seconds (audio length for Stable Audio, `seconds_total` for AudioX).
- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. For AudioX, sampler/reference knobs go here, e.g. `'{"sigma_min": 0.03, "sigma_max": 1000.0, "cfg_rescale": 0.0, "audio_path": "ref.wav"}'`.
- `--negative-prompt`: negative prompt for classifier-free guidance.
- `--seed`: integer seed for deterministic generation.
- `--guidance-scale`: classifier-free guidance scale.
- `--audio-length`: audio duration in seconds.
- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
- `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT.
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
examples/offline_inference/text_to_audio/text_to_audio.py
+141 -40
View File
@@ -2,37 +2,70 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example script for text-to-audio generation using Stable Audio Open.
Example script for text/video-to-audio generation with diffusion audio models.
This script demonstrates how to generate audio from text prompts using
the Stable Audio Open model with vLLM-Omni.
This script supports:
* Stable Audio Open (`stabilityai/stable-audio-open-1.0`) — text-to-audio.
* AudioX (`zhangj1an/AudioX`) — six tasks: t2a / t2m / v2a / v2m / tv2a / tv2m.
Model-specific generation knobs are declared in `vllm_omni/model_extras` and
routed into `sampling_params.extra_args` via `apply_declared_extra_args`, so
the same script serves both models without per-model branching in the engine.
Usage:
# Stable Audio Open
python text_to_audio.py --prompt "The sound of a dog barking"
python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache
# AudioX (text-to-audio / text-to-music)
python text_to_audio.py --model zhangj1an/AudioX --task t2a \
--prompt "Fireworks burst twice, then a clock ticking." \
--num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0
# AudioX video-conditioned (v2*/tv2*) — pass a video file/URL
python text_to_audio.py --model zhangj1an/AudioX --task tv2a \
--prompt "drum beating sound and human talking" --video clip.mp4 \
--extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}'
"""
import argparse
import json
import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.platforms import current_omni_platform
AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
AUDIOX_VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
def parse_extra_body(value: str) -> dict[str, Any]:
"""Parse a JSON object string for --extra-body."""
try:
obj = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e
if not isinstance(obj, dict):
raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
return obj
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate audio with Stable Audio Open.")
parser = argparse.ArgumentParser(description="Generate audio with supported diffusion audio models.")
parser.add_argument(
"--model",
default="stabilityai/stable-audio-open-1.0",
help="Stable Audio model name or local path.",
help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0, zhangj1an/AudioX.",
)
parser.add_argument(
"--prompt",
@@ -41,8 +74,11 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument(
"--negative-prompt",
default="Low quality.",
help="Negative prompt for classifier-free guidance.",
default=None,
help=(
"Negative prompt for classifier-free guidance. Off by default; "
'recommended for Stable Audio, e.g. --negative-prompt "Low quality".'
),
)
parser.add_argument(
"--seed",
@@ -60,13 +96,19 @@ def parse_args() -> argparse.Namespace:
"--audio-start",
type=float,
default=0.0,
help="Audio start time in seconds.",
help=(
"Audio start offset in seconds. Applies to both models "
"(audio_start_in_s for Stable Audio, seconds_start for AudioX)."
),
)
parser.add_argument(
"--audio-length",
type=float,
default=10.0,
help="Audio length in seconds (max ~47s for stable-audio-open-1.0).",
help=(
"Audio duration in seconds. Maps to audio length for Stable Audio "
"(max ~47s for stable-audio-open-1.0) and seconds_total for AudioX."
),
)
parser.add_argument(
"--num-inference-steps",
@@ -80,11 +122,41 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of audio waveforms to generate for the given prompt.",
)
parser.add_argument(
"--task",
type=str,
default=None,
choices=list(AUDIOX_TASKS),
help=(
"[AudioX only] Generation task: t2a/t2m/v2a/v2m/tv2a/tv2m. "
"v2*/tv2* require --video. Ignored for Stable Audio Open."
),
)
parser.add_argument(
"--video",
type=str,
default=None,
help="[AudioX v2*/tv2*] Video file path or URL used for conditioning.",
)
parser.add_argument(
"--extra-body",
type=parse_extra_body,
default=None,
help=(
"Model-specific generation params as a JSON object, e.g. "
'\'{"sigma_min": 0.03, "sigma_max": 1000.0}\'. Each key is filtered '
"against the model's declared extra_body_params (see vllm_omni/model_extras), "
"so unknown keys are silently dropped. Values here override the equivalent flags."
),
)
parser.add_argument(
"--output",
type=str,
default="stable_audio_output.wav",
help="Path to save the generated audio (WAV format).",
default=None,
help=(
"Path to save the generated audio (WAV format). Defaults to "
"audiox_output.wav for AudioX and stable_audio_output.wav otherwise."
),
)
parser.add_argument(
"--sample-rate",
@@ -203,6 +275,8 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
def main():
args = parse_args()
model_name = str(args.model).lower() if args.model else ""
is_audiox = "audiox" in model_name
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
cache_config = None
if args.cache_backend == "tea_cache":
@@ -211,12 +285,14 @@ def main():
}
print(f"\n{'=' * 60}")
print("Stable Audio Open - Text-to-Audio Generation")
print("Text/Video-to-Audio Generation")
print(f"{'=' * 60}")
print(f" Model: {args.model}")
if args.task is not None:
print(f" Task: {args.task}")
print(f" Prompt: {args.prompt}")
print(f" Negative prompt: {args.negative_prompt}")
print(f" Audio length: {args.audio_length}s")
print(f" Duration: {args.audio_length}s")
print(f" Inference steps: {args.num_inference_steps}")
print(f" Guidance scale: {args.guidance_scale}")
print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
@@ -235,8 +311,9 @@ def main():
hsdp_replicate_size=args.hsdp_replicate_size,
)
# Initialize Omni with Stable Audio model
omni = Omni(
# Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio
# is resolved from its config.
omni_kwargs: dict[str, Any] = dict(
model=args.model,
parallel_config=parallel_config,
cache_backend=args.cache_backend,
@@ -245,43 +322,68 @@ def main():
enable_cpu_offload=args.enable_cpu_offload,
enable_layerwise_offload=args.enable_layerwise_offload,
)
if is_audiox:
omni_kwargs["model_class_name"] = "AudioXPipeline"
omni = Omni(**omni_kwargs)
model_class_name = get_model_class_name(omni)
declared_extra_body_params = get_extra_body_params(model_class_name)
# Calculate audio end time
audio_end_in_s = args.audio_start + args.audio_length
diffusion_params = OmniDiffusionSamplingParams(
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_waveforms,
seed=args.seed,
)
# Negative prompt is opt-in for both models (off by default). Including a
# non-empty negative prompt enables the separate CFG conditioning branch.
prompt: dict[str, Any] = {"prompt": args.prompt}
if args.negative_prompt:
prompt["negative_prompt"] = args.negative_prompt
if model_class_name == "AudioXPipeline":
if args.task is None:
raise ValueError("AudioX requires --task (one of t2a/t2m/v2a/v2m/tv2a/tv2m).")
if args.task in AUDIOX_VIDEO_TASKS and not args.video:
raise ValueError(f"AudioX task {args.task!r} requires --video.")
# AudioX core knobs are surfaced as shared flags; sampler/reference knobs
# (sigma_min/sigma_max/cfg_rescale/audio_path) arrive via --extra-body and
# are filtered against the declared extra_body contract.
user_extra: dict[str, Any] = {
"audiox_task": args.task,
"seconds_start": args.audio_start,
"seconds_total": args.audio_length,
"video_path": args.video,
}
if args.extra_body:
user_extra.update(args.extra_body)
apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra)
else:
# Stable Audio Open path (backward-compatible).
audio_end_in_s = args.audio_start + args.audio_length
diffusion_params.extra_args = {
"audio_start_in_s": args.audio_start,
"audio_end_in_s": audio_end_in_s,
}
if args.extra_body:
diffusion_params.extra_args.update(args.extra_body)
# Time profiling for generation
generation_start = time.perf_counter()
# Generate audio
outputs = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
},
OmniDiffusionSamplingParams(
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_waveforms,
extra_args={
"audio_start_in_s": args.audio_start,
"audio_end_in_s": audio_end_in_s,
},
),
)
outputs = omni.generate(prompt, diffusion_params)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
print(f"Total generation time: {generation_time:.2f} seconds")
# Process and save audio
output_path = Path(args.output)
output = args.output or ("audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav")
output_path = Path(output)
output_path.parent.mkdir(parents=True, exist_ok=True)
suffix = output_path.suffix or ".wav"
stem = output_path.stem or "stable_audio_output"
# Extract audio from omni.generate() outputs
if not outputs:
raise ValueError("No output generated from omni.generate()")
@@ -296,7 +398,6 @@ def main():
if audio is None:
raise ValueError("No audio output found in request_output")
# Handle different output formats
if isinstance(audio, torch.Tensor):
audio = audio.cpu().float().numpy()
@@ -323,7 +424,7 @@ def main():
save_audio(audio, str(output_path), args.sample_rate)
print(f"Saved generated audio to {output_path}")
print(f"\nGenerated {args.audio_length}s of audio at {args.sample_rate} Hz")
print(f"\nGenerated ~{args.audio_length}s of audio at {args.sample_rate} Hz")
if __name__ == "__main__":
examples/online_serving/audiox/README.md
+8 -9
View File
@@ -37,11 +37,10 @@ python openai_chat_client.py --task tv2a \
The client sends:
- `num_inference_steps`, `guidance_scale`, `seed` as first-class OpenAI chat-completion fields
- `audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max` nested under
`extra_args` (a reserved dict on the request body that the server forwards verbatim into
the pipeline's `sampling_params.extra_args` — the same escape hatch `serving_video.py` exposes
as `extra_params` on /v1/videos)
- `num_inference_steps`, `guidance_scale`, `seed`, `audiox_task`, `seconds_start`,
`seconds_total`, `sigma_min`, `sigma_max` as flattened keys under `extra_body`. The server
routes the declared knobs (see `vllm_omni/model_extras/audiox.py`) into the pipeline's
`sampling_params.extra_args`.
- For `v2*` / `tv2*` tasks, the video as a `video_url` content item (data URI for local files)
## curl
@@ -52,10 +51,10 @@ curl -sS -X POST http://localhost:8099/v1/chat/completions \
-d '{
"model": "zhangj1an/AudioX",
"messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}],
"num_inference_steps": 250,
"guidance_scale": 7.0,
"seed": 42,
"extra_args": {
"extra_body": {
"num_inference_steps": 250,
"guidance_scale": 7.0,
"seed": 42,
"audiox_task": "t2m",
"seconds_total": 10.0,
"sigma_min": 0.3,
examples/online_serving/audiox/openai_chat_client.py
+6 -6
View File
@@ -4,8 +4,8 @@
AudioX supports 6 tasks (t2a, t2m, v2a, v2m, tv2a, tv2m). Text-only tasks send the prompt as the
chat message; video-conditioned tasks additionally attach the video via a ``video_url`` content
item (data URI for local files). Task + generation knobs (steps, cfg, sigma range, seconds, seed)
are sent via the OpenAI SDK's ``extra_body`` as ``extra_args`` — the same pipeline-agnostic escape
hatch used by the /v1/videos endpoint's ``extra_params`` field.
are sent as flattened keys under ``extra_body``; the server routes the declared knobs into
``sampling_params.extra_args`` (see ``vllm_omni/model_extras/audiox.py``).
Usage:
python openai_chat_client.py --task t2a --prompt "Fireworks burst twice..." --output t2a.wav
@@ -88,10 +88,10 @@ def main() -> int:
payload = {
"model": args.model,
"messages": [{"role": "user", "content": content}],
"num_inference_steps": args.steps,
"guidance_scale": args.guidance_scale,
"seed": args.seed,
"extra_args": {
"extra_body": {
"num_inference_steps": args.steps,
"guidance_scale": args.guidance_scale,
"seed": args.seed,
"audiox_task": args.task,
"seconds_start": args.seconds_start,
"seconds_total": args.seconds_total,
recipes/audiox/AudioX.md
+55 -20
View File
@@ -22,8 +22,9 @@ that produces stereo 44.1 kHz audio up to ~10 s per call.
- vLLM-Omni weight bundle: <https://huggingface.co/zhangj1an/AudioX>
- Pipeline: `vllm_omni.diffusion.models.audiox.pipeline_audiox.AudioXPipeline`
- Input transforms: `vllm_omni.transformers_utils.processors.audiox`
- Offline example: [`examples/offline_inference/audiox/`](../../examples/offline_inference/audiox/)
- Online example: [`examples/online_serving/audiox/`](../../examples/online_serving/audiox/)
- Param contract: `vllm_omni/model_extras/audiox.py` (declared `extra_body` knobs)
- Offline example: [`examples/offline_inference/text_to_audio/text_to_audio.py`](../../examples/offline_inference/text_to_audio/text_to_audio.py)
- Online: standard OpenAI chat-completions endpoint (see commands below)
## Hardware Support
@@ -41,27 +42,63 @@ that produces stereo 44.1 kHz audio up to ~10 s per call.
#### Command
Offline (text-to-audio):
AudioX uses the **standard** `text_to_audio` example offline and the standard
OpenAI chat-completions endpoint online. All model-specific knobs
(`audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max`,
`cfg_rescale`, `video_path`, `audio_path`) are declared in
`vllm_omni/model_extras/audiox.py` and routed through `extra_body` (online) /
`--extra-body` (offline). The offline example also exposes `--task`, `--video`,
`--audio-start`, and `--audio-length` shortcuts for the most common knobs.
Offline — text tasks (`t2a` / `t2m`):
```bash
huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
python examples/offline_inference/audiox/end2end.py \
--model ./audiox_weights \
--tasks t2a \
--num-inference-steps 250 \
--seconds-total 10
python examples/offline_inference/text_to_audio/text_to_audio.py \
--model ./audiox_weights --task t2a \
--prompt "Fireworks burst twice, followed by a clock ticking." \
--num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
--extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \
--output t2a.wav
```
Offline — video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) need `--video`:
```bash
python examples/offline_inference/text_to_audio/text_to_audio.py \
--model ./audiox_weights --task tv2a \
--prompt "drum beating sound and human talking" \
--video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \
--num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
--output tv2a.wav
```
Online:
```bash
bash examples/online_serving/audiox/run_server.sh
python examples/online_serving/audiox/openai_chat_client.py \
--task t2a \
--prompt "Fireworks burst twice, followed by a clock ticking." \
--output t2a.wav
DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \
vllm serve zhangj1an/AudioX --omni --model-class-name AudioXPipeline --port 8099
curl -sS -X POST http://localhost:8099/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "zhangj1an/AudioX",
"messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}],
"extra_body": {
"num_inference_steps": 250,
"guidance_scale": 7.0,
"seed": 42,
"audiox_task": "t2m",
"seconds_total": 10.0,
"sigma_min": 0.3,
"sigma_max": 500.0
}
}' > t2m.json
```
For `v2*` / `tv2*` online, attach the video as a `video_url` content item
(data URI for local files) in the message `content`.
#### Verification
```bash
@@ -75,13 +112,11 @@ ffprobe t2a.wav
#### Notes
- Memory usage: ~10 GB peak with `num_inference_steps=250`, 10 s of audio.
- Output rate: 44.1 kHz stereo, regardless of `--sample-rate` (resampled in the example
script if requested).
- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via
`extra_args["audiox_task"]` (offline) or the `extra_args` field in the OpenAI
chat-completions body (online).
- Video conditioning: `v2*` and `tv2*` require a video file; the online client
attaches it as an OpenAI `video_url` content item (data URI for local files).
- Output rate: 44.1 kHz stereo regardless of requested rate.
- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via the
`audiox_task` knob (in `extra_body` online, or `--task` offline).
- Video conditioning: `v2*` and `tv2*` require a video; offline via `--video`
(→ `video_path` knob), online via a `video_url` content item.
- Cache acceleration is **not** supported (AudioXPipeline is in `_NO_CACHE_ACCELERATION`).
- Tensor parallelism is supported via `--tensor-parallel-size` (DiT QKV is sharded with
`QKVParallelLinear`); cross-attention K/V is also TP-sharded.
tests/e2e/online_serving/test_audiox_expansion.py
+9 -1
View File
@@ -39,7 +39,13 @@ def _audiox_server_cases(model: str):
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True)
def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
"""AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`."""
"""AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`.
AudioX is registered in ``vllm_omni/model_extras`` (``AudioXPipeline``), so the
``audiox_task`` / ``seconds_*`` / ``sigma_*`` keys below flow through the standard
``apply_declared_extra_args`` server path into ``sampling_params.extra_args`` rather
than the old top-level ``extra_args`` escape hatch.
"""
request_config = {
"model": omni_server.model,
"messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT),
@@ -50,6 +56,8 @@ def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientH
"audiox_task": "t2a",
"seconds_start": 0.0,
"seconds_total": 2.0,
"sigma_min": 0.03,
"sigma_max": 1000.0,
},
}
openai_client.send_diffusion_request(request_config)
tests/model_extras/test_model_extras.py
+41
View File
@@ -147,6 +147,47 @@ def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
) == {"prompt": "a red sports car", "modalities": ["image"]}
@pytest.mark.core_model
@pytest.mark.cpu
def test_audiox_extra_registry_declares_request_and_response_params() -> None:
assert get_extra_body_params("AudioXPipeline") == frozenset(
{
"audiox_task",
"seconds_start",
"seconds_total",
"sigma_min",
"sigma_max",
"cfg_rescale",
"video_path",
"audio_path",
}
)
assert get_extra_output_params("AudioXPipeline") == frozenset({"audiox_task"})
assert should_init_extra_args_for_non_diffusion_stages("AudioXPipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_audiox_declared_extra_args_route_into_sampling_params() -> None:
params = OmniDiffusionSamplingParams()
declared = get_extra_body_params("AudioXPipeline")
apply_declared_extra_args(
params,
declared,
{
"audiox_task": "t2a",
"seconds_total": 10.0,
"sigma_min": 0.03,
"unknown": "ignored",
},
)
assert params.extra_args == {
"audiox_task": "t2a",
"seconds_total": 10.0,
"sigma_min": 0.03,
}
@pytest.mark.core_model
@pytest.mark.cpu
def test_helios_extra_registry_declares_request_and_response_params() -> None:
vllm_omni/model_extras/audiox.py
+27
View File
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
# Request params the AudioX pipeline reads from ``sampling_params.extra_args``
# (see ``pipeline_audiox.py`` ``forward``). Declaring them here routes
# ``extra_body`` (online) / ``--extra-body`` (offline) keys into ``extra_args``
# instead of silently dropping them.
AUDIOX_EXTRA_BODY_PARAMS = frozenset(
{
"audiox_task", # t2a/t2m/v2a/v2m/tv2a/tv2m
"seconds_start",
"seconds_total",
"sigma_min",
"sigma_max",
"cfg_rescale",
"video_path", # v2*/tv2* conditioning
"audio_path", # optional reference-audio conditioning
}
)
# Echoed back in ``DiffusionOutput.custom_output`` (pipeline_audiox.py).
AUDIOX_EXTRA_OUTPUT_PARAMS = frozenset(
{
"audiox_task",
}
)
vllm_omni/model_extras/registry.py
+8
View File
@@ -8,6 +8,10 @@ from typing import Any
from PIL import Image
from vllm_omni.model_extras.audiox import (
AUDIOX_EXTRA_BODY_PARAMS,
AUDIOX_EXTRA_OUTPUT_PARAMS,
)
from vllm_omni.model_extras.bagel import (
BAGEL_EXTRA_BODY_PARAMS,
BAGEL_EXTRA_OUTPUT_PARAMS,
@@ -118,6 +122,10 @@ def default_image_to_video_prompt(
_EXTRA_SPECS: dict[str, dict[str, Any]] = {
"AudioXPipeline": {
"extra_body_params": AUDIOX_EXTRA_BODY_PARAMS,
"extra_output_params": AUDIOX_EXTRA_OUTPUT_PARAMS,
},
"BagelPipeline": {
"extra_body_params": BAGEL_EXTRA_BODY_PARAMS,
"extra_output_params": BAGEL_EXTRA_OUTPUT_PARAMS,