120 lines
4.8 KiB
Python
Executable File
120 lines
4.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""AudioX OpenAI-compatible chat client.
|
|
|
|
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 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
|
|
python openai_chat_client.py --task tv2a --prompt "drum beating" --video clip.mp4 -o tv2a.wav
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import mimetypes
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
import soundfile
|
|
import torch
|
|
|
|
VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
|
|
TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"})
|
|
|
|
|
|
def _to_data_url(path: str) -> str:
|
|
mime, _ = mimetypes.guess_type(path)
|
|
mime = mime or "video/mp4"
|
|
with open(path, "rb") as f:
|
|
data = base64.b64encode(f.read()).decode("ascii")
|
|
return f"data:{mime};base64,{data}"
|
|
|
|
|
|
def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
|
|
audio = audio.to(torch.float32)
|
|
audio = audio / audio.abs().max().clamp(min=1e-8)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
# soundfile expects channels-last (T, C); project convention is (C, T).
|
|
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
|
|
|
|
|
|
def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
|
|
for choice in body.get("choices", []):
|
|
audio_obj = choice.get("message", {}).get("audio")
|
|
if not (isinstance(audio_obj, dict) and audio_obj.get("data")):
|
|
continue
|
|
data, sr = soundfile.read(io.BytesIO(base64.b64decode(audio_obj["data"])), dtype="float32", always_2d=True)
|
|
return torch.from_numpy(data).transpose(0, 1), sr
|
|
brief = {k: v for k, v in body.items() if k != "choices"}
|
|
raise RuntimeError(f"no audio in response message.audio: {brief}")
|
|
|
|
|
|
def main() -> int:
|
|
p = argparse.ArgumentParser(description="AudioX OpenAI chat client")
|
|
p.add_argument("--task", required=True, choices=["t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m"])
|
|
p.add_argument("--prompt", "-p", default="", help="Text prompt (required for t2*/tv2*).")
|
|
p.add_argument("--video", help="Video path or URL (required for v2*/tv2*).")
|
|
p.add_argument("--output", "-o", default="audiox_out.wav")
|
|
p.add_argument("--server", "-s", default="http://localhost:8099")
|
|
p.add_argument("--model", default="zhangj1an/AudioX")
|
|
p.add_argument("--steps", type=int, default=250)
|
|
p.add_argument("--guidance-scale", type=float, default=7.0)
|
|
p.add_argument("--seed", type=int, default=42)
|
|
p.add_argument("--seconds-total", type=float, default=10.0)
|
|
p.add_argument("--seconds-start", type=float, default=0.0)
|
|
p.add_argument("--sigma-min", type=float, default=0.03)
|
|
p.add_argument("--sigma-max", type=float, default=1000.0)
|
|
args = p.parse_args()
|
|
|
|
if args.task in VIDEO_TASKS and not args.video:
|
|
print(f"ERROR: task {args.task!r} requires --video", file=sys.stderr)
|
|
return 2
|
|
if args.task in TEXT_TASKS and not args.prompt.strip() and args.task not in {"v2a", "v2m"}:
|
|
print(f"ERROR: task {args.task!r} requires --prompt", file=sys.stderr)
|
|
return 2
|
|
|
|
content: list[dict] = [{"type": "text", "text": args.prompt}]
|
|
if args.task in VIDEO_TASKS:
|
|
vurl = args.video if args.video.startswith(("http://", "https://")) else _to_data_url(args.video)
|
|
content.append({"type": "video_url", "video_url": {"url": vurl}})
|
|
|
|
payload = {
|
|
"model": args.model,
|
|
"messages": [{"role": "user", "content": content}],
|
|
"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,
|
|
"sigma_min": args.sigma_min,
|
|
"sigma_max": args.sigma_max,
|
|
},
|
|
}
|
|
|
|
print(f"POST {args.server}/v1/chat/completions task={args.task} steps={args.steps}")
|
|
r = requests.post(
|
|
f"{args.server}/v1/chat/completions",
|
|
headers={"Content-Type": "application/json"},
|
|
json=payload,
|
|
timeout=600,
|
|
)
|
|
r.raise_for_status()
|
|
audio, sr = _decode_audio_from_response(r.json())
|
|
_save_wav(audio, Path(args.output), sr)
|
|
dur = audio.shape[-1] / sr
|
|
print(f"saved {args.output} sr={sr}Hz duration={dur:.2f}s channels={audio.shape[0]}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|