43 lines
1.8 KiB
Python
43 lines
1.8 KiB
Python
"""Drive the server's WebSocket like the browser does: stream a wav in real time in 2048-sample chunks, 3 turns."""
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import websockets
|
|
|
|
|
|
async def main(path, url="ws://127.0.0.1:8000/ws"):
|
|
audio, rate = sf.read(path, dtype="float32")
|
|
async with websockets.connect(url, max_size=None) as ws:
|
|
print("<", json.loads(await ws.recv())["type"])
|
|
for turn in range(3): # later turns exercise the cached history
|
|
await ws.send(json.dumps({"type": "start", "rate": rate}))
|
|
for i in range(0, len(audio), 2048):
|
|
await ws.send(audio[i:i + 2048].tobytes())
|
|
await asyncio.sleep(2048 / rate) # real-time pacing, like a mic
|
|
await ws.send(json.dumps({"type": "stop"}))
|
|
pcm_bytes = 0
|
|
while True:
|
|
msg = await ws.recv()
|
|
if isinstance(msg, bytes):
|
|
pcm_bytes += len(msg)
|
|
continue
|
|
m = json.loads(msg)
|
|
if m["type"] in ("token", "audio_chunk", "log"):
|
|
continue
|
|
if m["type"] == "done":
|
|
keys = ("first_text_ms", "first_frame_ms", "first_audio_ms", "flush_ms", "prefill_calls",
|
|
"prefill_busy_ms", "context_tokens", "context_tokens_after", "rtf", "gen_ms", "codec_decode_ms", "audio_out_s")
|
|
print("< done", m["text"][:80], {k: m["stats"].get(k) for k in keys})
|
|
elif m["type"] != "prompt":
|
|
print("<", m["type"], {k: v for k, v in m.items() if k != "type"})
|
|
if m["type"] in ("done", "error"):
|
|
break
|
|
print(f" received {pcm_bytes / 4 / 24000:.2f}s of PCM")
|
|
|
|
|
|
asyncio.run(main(sys.argv[1]))
|