333 lines
12 KiB
Python
333 lines
12 KiB
Python
"""Realtime client for vLLM-Omni /v1/realtime (audio + text events).
|
|
|
|
This client:
|
|
1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz),
|
|
2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events,
|
|
3) Receives response.audio.* and transcription.* events,
|
|
4) Saves synthesized audio to an output WAV file and optional text file.
|
|
|
|
By default each ``response.audio.delta`` is treated as an **incremental PCM**
|
|
chunk and all chunks are concatenated into the final ``--output-wav``.
|
|
|
|
Optional debugging: pass ``--delta-dump-dir DIR`` to write every
|
|
``response.audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, …
|
|
|
|
Usage:
|
|
python openai_realtime_client.py \
|
|
--url ws://localhost:8091/v1/realtime \
|
|
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
|
--input-wav input_16k_mono.wav \
|
|
--output-wav realtime_output.wav \
|
|
--delta-dump-dir ./rt_delta_wavs
|
|
|
|
Dependencies:
|
|
pip install websockets
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import wave
|
|
from pathlib import Path
|
|
|
|
try:
|
|
import websockets
|
|
except ImportError:
|
|
print("Please install websockets: pip install websockets")
|
|
raise SystemExit(1)
|
|
|
|
|
|
def _read_wav_pcm16(path: Path) -> bytes:
|
|
with wave.open(str(path), "rb") as wf:
|
|
nchannels = wf.getnchannels()
|
|
sampwidth = wf.getsampwidth()
|
|
framerate = wf.getframerate()
|
|
comptype = wf.getcomptype()
|
|
nframes = wf.getnframes()
|
|
|
|
if nchannels != 1:
|
|
raise ValueError(f"Input WAV must be mono (got {nchannels} channels).")
|
|
if sampwidth != 2:
|
|
raise ValueError(f"Input WAV must be 16-bit PCM (got sample width={sampwidth}).")
|
|
if framerate != 16000:
|
|
raise ValueError(f"Input WAV must be 16kHz (got {framerate} Hz).")
|
|
if comptype != "NONE":
|
|
raise ValueError(f"Input WAV must be uncompressed PCM (got comptype={comptype}).")
|
|
if nframes <= 0:
|
|
raise ValueError("Input WAV has no audio frames.")
|
|
|
|
return wf.readframes(nframes)
|
|
|
|
|
|
def _write_wav_pcm16(path: Path, pcm16_bytes: bytes, sample_rate_hz: int) -> None:
|
|
with wave.open(str(path), "wb") as wf:
|
|
wf.setnchannels(1)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(sample_rate_hz)
|
|
wf.writeframes(pcm16_bytes)
|
|
|
|
|
|
async def run_client(
|
|
url: str,
|
|
model: str,
|
|
input_wav: Path,
|
|
output_wav: Path,
|
|
output_text: Path | None,
|
|
chunk_ms: int,
|
|
send_delay_ms: int,
|
|
delta_dump_dir: Path | None,
|
|
request_idx: int = 1,
|
|
total_requests: int = 1,
|
|
) -> None:
|
|
log_prefix = f"[req {request_idx:02d}/{total_requests:02d}] " if total_requests > 1 else ""
|
|
pcm16 = _read_wav_pcm16(input_wav)
|
|
bytes_per_ms = 16000 * 2 // 1000 # mono PCM16 at 16kHz
|
|
chunk_bytes = max(bytes_per_ms * chunk_ms, 2)
|
|
|
|
incremental_pcm_parts: list[bytes] = []
|
|
output_sample_rate = 24000
|
|
delta_index = 0
|
|
text_chunks: list[str] = []
|
|
final_text: str = ""
|
|
|
|
if delta_dump_dir is not None:
|
|
delta_dump_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws:
|
|
# 1) Validate model.
|
|
await ws.send(
|
|
json.dumps(
|
|
{
|
|
"type": "session.update",
|
|
"model": model,
|
|
}
|
|
)
|
|
)
|
|
|
|
# 2) Start generation once (non-final commit).
|
|
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
|
|
|
|
# 3) Stream audio chunks.
|
|
for i in range(0, len(pcm16), chunk_bytes):
|
|
chunk = pcm16[i : i + chunk_bytes]
|
|
await ws.send(
|
|
json.dumps(
|
|
{
|
|
"type": "input_audio_buffer.append",
|
|
"audio": base64.b64encode(chunk).decode("utf-8"),
|
|
}
|
|
)
|
|
)
|
|
if send_delay_ms > 0:
|
|
await asyncio.sleep(send_delay_ms / 1000.0)
|
|
|
|
# 4) Final commit closes input stream.
|
|
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
|
|
|
|
# 5) Receive server events until audio done.
|
|
while True:
|
|
message = await ws.recv()
|
|
if isinstance(message, bytes):
|
|
# We only expect JSON text frames.
|
|
continue
|
|
|
|
event = json.loads(message)
|
|
event_type = event.get("type")
|
|
|
|
if event_type == "session.created":
|
|
continue
|
|
|
|
if event_type == "response.audio.delta":
|
|
sr = event.get("sample_rate_hz")
|
|
if isinstance(sr, int) and sr > 0:
|
|
output_sample_rate = sr
|
|
audio_b64 = event.get("audio", "")
|
|
if audio_b64:
|
|
pcm_delta = base64.b64decode(audio_b64)
|
|
incremental_pcm_parts.append(pcm_delta)
|
|
if delta_dump_dir is not None and pcm_delta:
|
|
delta_index += 1
|
|
dump_path = delta_dump_dir / f"delta_{delta_index:06d}.wav"
|
|
_write_wav_pcm16(dump_path, pcm_delta, output_sample_rate)
|
|
print(
|
|
f"{log_prefix}delta dump #{delta_index}: {dump_path} "
|
|
f"(pcm bytes={len(pcm_delta)}, sr={output_sample_rate})"
|
|
)
|
|
continue
|
|
|
|
if event_type == "transcription.delta":
|
|
delta = event.get("delta", "")
|
|
if delta:
|
|
text_chunks.append(delta)
|
|
print(delta, end="", flush=True)
|
|
continue
|
|
|
|
if event_type == "transcription.done":
|
|
final_text = event.get("text", "") or "".join(text_chunks)
|
|
usage = event.get("usage")
|
|
final_text_with_tag = f"Final transcription: {final_text}"
|
|
if text_chunks:
|
|
print()
|
|
print(f"{log_prefix}{final_text_with_tag}")
|
|
if usage:
|
|
print(f"{log_prefix}text usage: {usage}")
|
|
continue
|
|
|
|
if event_type == "response.audio.done":
|
|
break
|
|
|
|
if event_type == "error":
|
|
raise RuntimeError(f"Server error: {event}")
|
|
|
|
all_pcm16 = b"".join(incremental_pcm_parts)
|
|
if not all_pcm16:
|
|
raise RuntimeError("No audio received from server.")
|
|
|
|
output_wav.parent.mkdir(parents=True, exist_ok=True)
|
|
_write_wav_pcm16(output_wav, all_pcm16, output_sample_rate)
|
|
print(f"{log_prefix}Saved realtime audio to: {output_wav} (incremental chunks joined)")
|
|
|
|
if output_text is not None:
|
|
text_to_save = final_text if final_text else "".join(text_chunks)
|
|
output_text.parent.mkdir(parents=True, exist_ok=True)
|
|
output_text.write_text(text_to_save, encoding="utf-8")
|
|
print(f"{log_prefix}Saved realtime text to: {output_text}")
|
|
|
|
|
|
def _indexed_output_path(path: Path | None, index: int, total: int) -> Path | None:
|
|
if path is None or total <= 1:
|
|
return path
|
|
return path.with_name(f"{path.stem}_{index:02d}{path.suffix}")
|
|
|
|
|
|
async def run_clients_concurrent(
|
|
*,
|
|
url: str,
|
|
model: str,
|
|
input_wav: Path,
|
|
output_wav: Path,
|
|
output_text: Path | None,
|
|
chunk_ms: int,
|
|
send_delay_ms: int,
|
|
delta_dump_dir: Path | None,
|
|
num_requests: int,
|
|
concurrency: int,
|
|
) -> None:
|
|
sem = asyncio.Semaphore(concurrency)
|
|
|
|
async def _run_one(index: int) -> tuple[int, bool, str | None]:
|
|
per_output_wav = _indexed_output_path(output_wav, index, num_requests)
|
|
per_output_text = _indexed_output_path(output_text, index, num_requests)
|
|
per_delta_dir = None
|
|
if delta_dump_dir is not None:
|
|
per_delta_dir = delta_dump_dir / f"req_{index:02d}"
|
|
async with sem:
|
|
try:
|
|
await run_client(
|
|
url=url,
|
|
model=model,
|
|
input_wav=input_wav,
|
|
output_wav=per_output_wav,
|
|
output_text=per_output_text,
|
|
chunk_ms=chunk_ms,
|
|
send_delay_ms=send_delay_ms,
|
|
delta_dump_dir=per_delta_dir,
|
|
request_idx=index,
|
|
total_requests=num_requests,
|
|
)
|
|
return index, True, None
|
|
except Exception as exc:
|
|
return index, False, str(exc)
|
|
|
|
tasks = [asyncio.create_task(_run_one(i), name=f"rt-client-{i}") for i in range(1, num_requests + 1)]
|
|
results = await asyncio.gather(*tasks)
|
|
|
|
failed = [(idx, err) for idx, ok, err in results if not ok]
|
|
succeeded = num_requests - len(failed)
|
|
print(f"[summary] succeeded={succeeded}, failed={len(failed)}, total={num_requests}")
|
|
if failed:
|
|
for idx, err in failed:
|
|
print(f"[summary] req {idx:02d} failed: {err}")
|
|
raise RuntimeError(f"{len(failed)} concurrent request(s) failed")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Realtime audio/text client for vLLM-Omni")
|
|
parser.add_argument("--url", default="ws://localhost:8091/v1/realtime", help="WebSocket URL")
|
|
parser.add_argument(
|
|
"--model",
|
|
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
|
help="Model name for session.update",
|
|
)
|
|
parser.add_argument("--input-wav", required=True, type=Path, help="Input WAV (mono, PCM16, 16kHz)")
|
|
parser.add_argument("--output-wav", default=Path("realtime_output.wav"), type=Path, help="Output WAV path")
|
|
parser.add_argument(
|
|
"--output-text",
|
|
default=None,
|
|
type=Path,
|
|
help="Optional output text path for final transcription",
|
|
)
|
|
parser.add_argument("--chunk-ms", type=int, default=200, help="Input chunk size in milliseconds")
|
|
parser.add_argument(
|
|
"--send-delay-ms",
|
|
type=int,
|
|
default=0,
|
|
help="Delay between chunk sends; set >0 to simulate realtime upload",
|
|
)
|
|
parser.add_argument(
|
|
"--delta-dump-dir",
|
|
type=Path,
|
|
default=None,
|
|
help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory",
|
|
)
|
|
parser.add_argument("--num-requests", type=int, default=1, help="Total number of requests to send")
|
|
parser.add_argument(
|
|
"--concurrency",
|
|
type=int,
|
|
default=1,
|
|
help="Maximum number of concurrent websocket requests",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.num_requests <= 0:
|
|
raise ValueError("--num-requests must be >= 1")
|
|
if args.concurrency <= 0:
|
|
raise ValueError("--concurrency must be >= 1")
|
|
concurrency = min(args.concurrency, args.num_requests)
|
|
|
|
if args.num_requests == 1:
|
|
asyncio.run(
|
|
run_client(
|
|
url=args.url,
|
|
model=args.model,
|
|
input_wav=args.input_wav,
|
|
output_wav=args.output_wav,
|
|
output_text=args.output_text,
|
|
chunk_ms=args.chunk_ms,
|
|
send_delay_ms=args.send_delay_ms,
|
|
delta_dump_dir=args.delta_dump_dir,
|
|
)
|
|
)
|
|
else:
|
|
asyncio.run(
|
|
run_clients_concurrent(
|
|
url=args.url,
|
|
model=args.model,
|
|
input_wav=args.input_wav,
|
|
output_wav=args.output_wav,
|
|
output_text=args.output_text,
|
|
chunk_ms=args.chunk_ms,
|
|
send_delay_ms=args.send_delay_ms,
|
|
delta_dump_dir=args.delta_dump_dir,
|
|
num_requests=args.num_requests,
|
|
concurrency=concurrency,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|