From 71a951cbfff12db41b44191647c89f598ce4af00 Mon Sep 17 00:00:00 2001 From: zzehli Date: Sun, 21 Jun 2026 16:55:03 -0400 Subject: [PATCH 1/3] declare audiox extra params Signed-off-by: zzehli --- tests/model_extras/test_model_extras.py | 41 +++++++++++++++++++++++++ vllm_omni/model_extras/audiox.py | 27 ++++++++++++++++ vllm_omni/model_extras/registry.py | 8 +++++ 3 files changed, 76 insertions(+) create mode 100644 vllm_omni/model_extras/audiox.py diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 7eaf3edf612..b1c9556529c 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -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: diff --git a/vllm_omni/model_extras/audiox.py b/vllm_omni/model_extras/audiox.py new file mode 100644 index 00000000000..2ab7bc4d896 --- /dev/null +++ b/vllm_omni/model_extras/audiox.py @@ -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", + } +) diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 69608764a7a..d1f51b6f2a6 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -8,6 +8,10 @@ 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, @@ -110,6 +114,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, From a6ba746ab01384fb8e916f21e41f765fbfc9c4bb Mon Sep 17 00:00:00 2001 From: zzehli Date: Sun, 21 Jun 2026 22:12:41 -0400 Subject: [PATCH 2/3] refactor audiox examples, update e2e and recipes Signed-off-by: zzehli --- examples/offline_inference/audiox/README.md | 40 ---- examples/offline_inference/audiox/end2end.py | 110 ----------- .../offline_inference/text_to_audio/README.md | 41 +++- .../text_to_audio/text_to_audio.py | 183 ++++++++++++++---- examples/online_serving/audiox/README.md | 17 +- .../audiox/openai_chat_client.py | 12 +- recipes/audiox/AudioX.md | 75 +++++-- .../online_serving/test_audiox_expansion.py | 10 +- 8 files changed, 261 insertions(+), 227 deletions(-) delete mode 100644 examples/offline_inference/audiox/README.md delete mode 100644 examples/offline_inference/audiox/end2end.py diff --git a/examples/offline_inference/audiox/README.md b/examples/offline_inference/audiox/README.md deleted file mode 100644 index a0b46f9196f..00000000000 --- a/examples/offline_inference/audiox/README.md +++ /dev/null @@ -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 `/.wav` as 16-bit stereo WAV. diff --git a/examples/offline_inference/audiox/end2end.py b/examples/offline_inference/audiox/end2end.py deleted file mode 100644 index 34b5425724e..00000000000 --- a/examples/offline_inference/audiox/end2end.py +++ /dev/null @@ -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() diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md index 10a8ae37ed1..89ffc8bd8a9 100644 --- a/examples/offline_inference/text_to_audio/README.md +++ b/examples/offline_inference/text_to_audio/README.md @@ -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. diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 3072cc06f47..74d5d801182 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -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,70 @@ 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) + + 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, + ) - # Calculate audio end time - audio_end_in_s = args.audio_start + args.audio_length + # 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 +400,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 +426,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__": diff --git a/examples/online_serving/audiox/README.md b/examples/online_serving/audiox/README.md index aad65eb301e..bf068b3eb4f 100644 --- a/examples/online_serving/audiox/README.md +++ b/examples/online_serving/audiox/README.md @@ -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, diff --git a/examples/online_serving/audiox/openai_chat_client.py b/examples/online_serving/audiox/openai_chat_client.py index 8c8b572e50d..fd9aa8cc550 100755 --- a/examples/online_serving/audiox/openai_chat_client.py +++ b/examples/online_serving/audiox/openai_chat_client.py @@ -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, diff --git a/recipes/audiox/AudioX.md b/recipes/audiox/AudioX.md index 6365bd02136..a2e26073666 100644 --- a/recipes/audiox/AudioX.md +++ b/recipes/audiox/AudioX.md @@ -22,8 +22,9 @@ that produces stereo 44.1 kHz audio up to ~10 s per call. - vLLM-Omni weight bundle: - 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. diff --git a/tests/e2e/online_serving/test_audiox_expansion.py b/tests/e2e/online_serving/test_audiox_expansion.py index fc5bf06cf64..612f397e007 100644 --- a/tests/e2e/online_serving/test_audiox_expansion.py +++ b/tests/e2e/online_serving/test_audiox_expansion.py @@ -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) From ce4b4a9e349549aa6fc47a5fae6a6eb254615f68 Mon Sep 17 00:00:00 2001 From: zzehli Date: Wed, 24 Jun 2026 07:19:46 -0400 Subject: [PATCH 3/3] fix: linting issue Signed-off-by: zzehli --- examples/offline_inference/text_to_audio/text_to_audio.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 74d5d801182..3ad88a6f346 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -378,9 +378,7 @@ def main(): print(f"Total generation time: {generation_time:.2f} seconds") - output = args.output or ( - "audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav" - ) + 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"