352 lines
14 KiB
Python
352 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Combined HTTP + WebSocket proxy for Qwen3-TTS on a single port.
|
|
|
|
Routes on one port (default 8091):
|
|
POST /v1/audio/speech → REST proxy to internal vllm-omni
|
|
GET /v1/audio/speech/stream → WebSocket (PCM-paced, REST-sourced)
|
|
GET /health → health proxy to internal vllm-omni
|
|
|
|
Why REST-sourced WebSocket?
|
|
The vllm-omni streaming codec drops the onset of short utterances.
|
|
The batch REST endpoint (/v1/audio/speech) uses full decode and renders
|
|
every word cleanly. This proxy keeps the client-facing WS protocol intact
|
|
while sourcing audio sentence-by-sentence from REST, pacing the result
|
|
into one continuous real-time PCM stream via pcm_pacer.
|
|
|
|
Why instructions is NOT forwarded:
|
|
With a cloned ICL voice (ref_audio) the voice character comes from the
|
|
reference; an extra style prompt makes the Talker emit EOS far too early
|
|
and clips short words ("Hallo." → 0.16 s vs 0.40 s without).
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
import datetime as dt
|
|
import json
|
|
import math
|
|
import re
|
|
|
|
import aiohttp
|
|
import websockets
|
|
from aiohttp import web
|
|
|
|
# ---------------------------------------------------------------------------
|
|
SAMPLE_RATE = 24_000
|
|
BYTES_PER_SAMPLE = 2
|
|
BYTES_PER_SEC = SAMPLE_RATE * BYTES_PER_SAMPLE # 48 000
|
|
|
|
SEND_CHUNK = 960 # 20 ms at 48 kHz
|
|
FRAME_SECS = SEND_CHUNK / BYTES_PER_SEC # 0.020 s
|
|
|
|
LEADING_MS = 100
|
|
LEADING_FRAMES = math.ceil(LEADING_MS / 1000 / FRAME_SECS)
|
|
TRAILING_MS = 200
|
|
TRAILING_FRAMES = math.ceil(TRAILING_MS / 1000 / FRAME_SECS)
|
|
|
|
# Per-sentence synthesis timeout. A healthy sentence renders in <2 s.
|
|
SYNTH_TIMEOUT_S = 30
|
|
# Direct REST proxy timeout — longer for potentially bigger inputs.
|
|
REST_PROXY_TIMEOUT_S = 120
|
|
|
|
SILENCE_CHUNK = bytes(SEND_CHUNK)
|
|
|
|
SPLIT_RE = re.compile(r"(?<=[.!?])\s+|(?<=[。！？])")
|
|
MIN_SENTENCE_LEN = 2
|
|
|
|
# Fields forwarded to the internal REST synthesis request.
|
|
# `instructions` deliberately omitted — see module docstring.
|
|
REST_PASSTHROUGH = ("model", "voice", "task_type", "language",
|
|
"seed", "speaker_embedding")
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_internal_base_url: str = "" # e.g. http://localhost:8095
|
|
_internal_ws_url: str = "" # e.g. ws://localhost:8095/v1/audio/speech/stream
|
|
|
|
|
|
def now() -> str:
|
|
return dt.datetime.now().isoformat(timespec="seconds")
|
|
|
|
|
|
class SentenceSplitter:
|
|
"""Incremental splitter: buffers text, yields complete sentences."""
|
|
|
|
def __init__(self) -> None:
|
|
self.buf = ""
|
|
|
|
def add(self, text: str) -> list[str]:
|
|
self.buf += text
|
|
out: list[str] = []
|
|
while True:
|
|
m = SPLIT_RE.search(self.buf)
|
|
if not m:
|
|
break
|
|
idx = m.end()
|
|
head = self.buf[:idx].strip()
|
|
if len(head) < MIN_SENTENCE_LEN:
|
|
break
|
|
self.buf = self.buf[idx:]
|
|
if head:
|
|
out.append(head)
|
|
return out
|
|
|
|
def flush(self) -> list[str]:
|
|
s = self.buf.strip()
|
|
self.buf = ""
|
|
return [s] if s else []
|
|
|
|
|
|
async def synth_rest(http: aiohttp.ClientSession, cfg: dict, text: str) -> bytes:
|
|
"""Synthesise one sentence via the internal batch REST endpoint → raw PCM."""
|
|
payload = {"input": text, "response_format": "pcm"}
|
|
for k in REST_PASSTHROUGH:
|
|
if cfg.get(k) is not None:
|
|
payload[k] = cfg[k]
|
|
async with http.post(f"{_internal_base_url}/v1/audio/speech",
|
|
json=payload) as resp:
|
|
resp.raise_for_status()
|
|
return await resp.read()
|
|
|
|
|
|
# ── REST proxy ───────────────────────────────────────────────────────────────
|
|
|
|
async def handle_rest(request: web.Request) -> web.StreamResponse:
|
|
"""Forward POST /v1/audio/speech to the internal vllm-omni REST API."""
|
|
body = await request.read()
|
|
ct = request.headers.get("Content-Type", "application/json")
|
|
timeout = aiohttp.ClientTimeout(total=REST_PROXY_TIMEOUT_S)
|
|
async with aiohttp.ClientSession(timeout=timeout) as http:
|
|
async with http.post(
|
|
f"{_internal_base_url}/v1/audio/speech",
|
|
data=body,
|
|
headers={"Content-Type": ct},
|
|
) as resp:
|
|
out_ct = resp.headers.get("Content-Type", "application/octet-stream")
|
|
response = web.StreamResponse(
|
|
status=resp.status,
|
|
headers={"Content-Type": out_ct},
|
|
)
|
|
await response.prepare(request)
|
|
async for chunk in resp.content.iter_chunked(65536):
|
|
await response.write(chunk)
|
|
await response.write_eof()
|
|
return response
|
|
|
|
|
|
async def handle_health(request: web.Request) -> web.Response:
|
|
"""Proxy GET /health to internal vllm-omni."""
|
|
try:
|
|
async with aiohttp.ClientSession() as http:
|
|
async with http.get(f"{_internal_base_url}/health") as resp:
|
|
return web.Response(status=resp.status, text=await resp.text())
|
|
except Exception:
|
|
return web.Response(status=503, text="unavailable")
|
|
|
|
|
|
# ── WebSocket PCM session ────────────────────────────────────────────────────
|
|
|
|
async def proxy_session_pcm(client_ws: web.WebSocketResponse,
|
|
conn_id: str, cfg: dict) -> None:
|
|
"""REST-sourced, paced, single continuous PCM stream per session."""
|
|
pcm_queue: asyncio.Queue = asyncio.Queue()
|
|
sentence_queue: asyncio.Queue = asyncio.Queue()
|
|
splitter = SentenceSplitter()
|
|
started = {"audio": False}
|
|
|
|
def enqueue_pcm(data: bytes) -> None:
|
|
for i in range(0, len(data), SEND_CHUNK):
|
|
piece = data[i:i + SEND_CHUNK]
|
|
if len(piece) < SEND_CHUNK:
|
|
piece = piece + bytes(SEND_CHUNK - len(piece))
|
|
pcm_queue.put_nowait(piece)
|
|
|
|
async def pcm_pacer() -> bool:
|
|
"""Drain pcm_queue at 48 kHz. Returns True = fully drained, False = disconnected."""
|
|
loop = asyncio.get_event_loop()
|
|
t0 = loop.time()
|
|
sent = 0
|
|
try:
|
|
while True:
|
|
target = t0 + sent / BYTES_PER_SEC
|
|
delay = target - loop.time()
|
|
if delay > 0.001:
|
|
await asyncio.sleep(delay)
|
|
try:
|
|
chunk = pcm_queue.get_nowait()
|
|
if chunk is None:
|
|
break
|
|
except asyncio.QueueEmpty:
|
|
chunk = SILENCE_CHUNK
|
|
await client_ws.send_bytes(chunk)
|
|
sent += len(chunk)
|
|
except Exception as exc:
|
|
print(f"[{now()}] {conn_id} PACER send failed after {sent} B "
|
|
f"({sent/BYTES_PER_SEC:.2f}s): {exc!r}", flush=True)
|
|
return False
|
|
print(f"[{now()}] {conn_id} PACER done, {sent} B "
|
|
f"({sent/BYTES_PER_SEC:.2f}s)", flush=True)
|
|
return True
|
|
|
|
pacer_task = asyncio.create_task(pcm_pacer())
|
|
|
|
async def synth_worker() -> None:
|
|
timeout = aiohttp.ClientTimeout(total=SYNTH_TIMEOUT_S)
|
|
async with aiohttp.ClientSession(timeout=timeout) as http:
|
|
while True:
|
|
text = await sentence_queue.get()
|
|
if text is None:
|
|
break
|
|
try:
|
|
pcm = await synth_rest(http, cfg, text)
|
|
except Exception as exc:
|
|
print(f"[{now()}] {conn_id} REST synth failed for "
|
|
f"{text!r}: {exc!r}", flush=True)
|
|
continue
|
|
print(f"[{now()}] {conn_id} REST {text!r} -> {len(pcm)} B "
|
|
f"({len(pcm)/BYTES_PER_SEC:.2f}s)", flush=True)
|
|
if not started["audio"]:
|
|
started["audio"] = True
|
|
await client_ws.send_str(json.dumps(
|
|
{"type": "audio.start", "format": "pcm",
|
|
"sample_rate": SAMPLE_RATE}))
|
|
for _ in range(LEADING_FRAMES):
|
|
pcm_queue.put_nowait(SILENCE_CHUNK)
|
|
enqueue_pcm(pcm)
|
|
for _ in range(TRAILING_FRAMES):
|
|
pcm_queue.put_nowait(SILENCE_CHUNK)
|
|
pcm_queue.put_nowait(None)
|
|
|
|
worker_task = asyncio.create_task(synth_worker())
|
|
|
|
async def cleanup() -> None:
|
|
for t in (worker_task, pacer_task):
|
|
if not t.done():
|
|
t.cancel()
|
|
await asyncio.gather(worker_task, pacer_task, return_exceptions=True)
|
|
|
|
try:
|
|
async for msg in client_ws:
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
try:
|
|
data = json.loads(msg.data)
|
|
typ = data.get("type")
|
|
except Exception:
|
|
continue
|
|
if typ == "input.text":
|
|
text = data.get("text", "")
|
|
print(f"[{now()}] {conn_id} C text {text!r}", flush=True)
|
|
for sent in splitter.add(text):
|
|
sentence_queue.put_nowait(sent)
|
|
elif typ == "input.done":
|
|
print(f"[{now()}] {conn_id} C input.done", flush=True)
|
|
for sent in splitter.flush():
|
|
sentence_queue.put_nowait(sent)
|
|
sentence_queue.put_nowait(None)
|
|
drained_ok = await pacer_task
|
|
if not worker_task.done():
|
|
worker_task.cancel()
|
|
if drained_ok:
|
|
try:
|
|
await client_ws.send_str(
|
|
json.dumps({"type": "session.done"}))
|
|
except Exception:
|
|
pass
|
|
return
|
|
elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
|
|
break
|
|
finally:
|
|
await cleanup()
|
|
|
|
|
|
async def proxy_session_passthrough(client_ws: web.WebSocketResponse,
|
|
conn_id: str, first_msg_data: str) -> None:
|
|
"""Transparent WS passthrough to streaming upstream (non-PCM clients)."""
|
|
async with websockets.connect(_internal_ws_url, max_size=None) as upstream:
|
|
await upstream.send(first_msg_data)
|
|
|
|
async def c2u() -> None:
|
|
async for msg in client_ws:
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
await upstream.send(msg.data)
|
|
elif msg.type == aiohttp.WSMsgType.BINARY:
|
|
await upstream.send(msg.data)
|
|
else:
|
|
break
|
|
|
|
async def u2c() -> None:
|
|
async for m in upstream:
|
|
if isinstance(m, str):
|
|
await client_ws.send_str(m)
|
|
else:
|
|
await client_ws.send_bytes(m)
|
|
|
|
await asyncio.gather(c2u(), u2c())
|
|
|
|
|
|
async def handle_ws(request: web.Request) -> web.WebSocketResponse:
|
|
conn_id = f"ws-{id(request):x}"
|
|
client_ws = web.WebSocketResponse(max_msg_size=0)
|
|
await client_ws.prepare(request)
|
|
print(f"[{now()}] {conn_id} open", flush=True)
|
|
|
|
# First client message must be session.config.
|
|
first = await client_ws.receive()
|
|
if first.type != aiohttp.WSMsgType.TEXT:
|
|
return client_ws
|
|
try:
|
|
cfg = json.loads(first.data)
|
|
except Exception:
|
|
cfg = {}
|
|
|
|
if cfg.get("response_format", "pcm") == "pcm":
|
|
log = dict(cfg)
|
|
if log.get("speaker_embedding"):
|
|
log["speaker_embedding"] = f"<{len(log['speaker_embedding'])} floats>"
|
|
print(f"[{now()}] {conn_id} config "
|
|
f"{json.dumps(log, ensure_ascii=False, sort_keys=True)}", flush=True)
|
|
await proxy_session_pcm(client_ws, conn_id, cfg)
|
|
else:
|
|
print(f"[{now()}] {conn_id} non-PCM -> passthrough", flush=True)
|
|
await proxy_session_passthrough(client_ws, conn_id, first.data)
|
|
|
|
return client_ws
|
|
|
|
|
|
# ── main ─────────────────────────────────────────────────────────────────────
|
|
|
|
async def main() -> None:
|
|
global _internal_base_url, _internal_ws_url
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="0.0.0.0")
|
|
parser.add_argument("--port", type=int, default=8091)
|
|
parser.add_argument("--upstream", default="ws://localhost:8095/v1/audio/speech/stream")
|
|
parser.add_argument("--rest", default=None,
|
|
help="Internal REST base URL (default: derived from --upstream)")
|
|
args = parser.parse_args()
|
|
|
|
_internal_ws_url = args.upstream
|
|
if args.rest:
|
|
_internal_base_url = args.rest.rstrip("/")
|
|
else:
|
|
http = args.upstream.replace("ws://", "http://").replace("wss://", "https://")
|
|
_internal_base_url = http.rsplit("/v1/", 1)[0]
|
|
|
|
print(f"Proxy on {args.host}:{args.port} "
|
|
f"(REST: {_internal_base_url}/v1/audio/speech "
|
|
f"WS: {_internal_ws_url})", flush=True)
|
|
|
|
app = web.Application()
|
|
app.router.add_get("/v1/audio/speech/stream", handle_ws)
|
|
app.router.add_post("/v1/audio/speech", handle_rest)
|
|
app.router.add_get("/health", handle_health)
|
|
|
|
runner = web.AppRunner(app, access_log=None)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, args.host, args.port)
|
|
await site.start()
|
|
print(f"Listening on http+ws://{args.host}:{args.port}", flush=True)
|
|
await asyncio.Future()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|