Funktionierender STand!!!

2026-06-20 19:54:41 +02:00
parent 404c632cfc
commit cc3c0dad7b
6 changed files with 379 additions and 89 deletions
--- a/docker/ws_log_proxy.py
+++ b/docker/ws_log_proxy.py
@@ -1,83 +1,350 @@
 #!/usr/bin/env python3
-"""Logging WebSocket proxy for Qwen3-TTS streaming calls."""
+"""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      = 300
+LEADING_FRAMES  = math.ceil(LEADING_MS  / 1000 / FRAME_SECS)
+TRAILING_MS     = 2_000
+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():
+def now() -> str:
    return dt.datetime.now().isoformat(timespec="seconds")


-def shorten(value, limit=500):
-    text = value if isinstance(value, str) else repr(value)
-    return text if len(text) <= limit else text[:limit] + "..."
+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 proxy(client_ws, upstream_url):
-    conn_id = f"ws-{id(client_ws):x}"
-    print(f"[{now()}] {conn_id} open -> {upstream_url}", flush=True)
-    async with websockets.connect(upstream_url, max_size=None) as upstream_ws:
-        async def client_to_upstream():
-            async for msg in client_ws:
-                if isinstance(msg, str):
-                    try:
-                        data = json.loads(msg)
-                        typ = data.get("type")
-                        if typ == "session.config":
-                            log_data = dict(data)
-                            if "speaker_embedding" in log_data:
-                                emb = log_data["speaker_embedding"] or []
-                                log_data["speaker_embedding"] = f"<{len(emb)} floats>"
-                            print(f"[{now()}] {conn_id} C->S config {json.dumps(log_data, ensure_ascii=False, sort_keys=True)}", flush=True)
-                        elif typ == "input.text":
-                            print(f"[{now()}] {conn_id} C->S text {data.get('text')!r}", flush=True)
-                        else:
-                            print(f"[{now()}] {conn_id} C->S {shorten(msg)}", flush=True)
-                    except Exception:
-                        print(f"[{now()}] {conn_id} C->S {shorten(msg)}", flush=True)
-                else:
-                    print(f"[{now()}] {conn_id} C->S binary {len(msg)} bytes", flush=True)
-                await upstream_ws.send(msg)
-
-        async def upstream_to_client():
-            async for msg in upstream_ws:
-                if isinstance(msg, str):
-                    try:
-                        data = json.loads(msg)
-                        print(f"[{now()}] {conn_id} S->C {data}", flush=True)
-                    except Exception:
-                        print(f"[{now()}] {conn_id} S->C {shorten(msg)}", flush=True)
-                else:
-                    print(f"[{now()}] {conn_id} S->C binary {len(msg)} bytes", flush=True)
-                await client_ws.send(msg)
-
-        await asyncio.gather(client_to_upstream(), upstream_to_client())
+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()


-async def handler(client_ws):
+# ── 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:
-        await proxy(client_ws, handler.upstream_url)
-    except Exception as exc:
-        print(f"[{now()}] proxy error: {exc}", flush=True)
+        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:
-            await client_ws.close()
-        except Exception:
-            pass
+            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 main():
+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=8094)
-    parser.add_argument("--upstream", default="ws://localhost:8091/v1/audio/speech/stream")
+    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()
-    handler.upstream_url = args.upstream
-    print(f"WS log proxy: ws://{args.host}:{args.port}/v1/audio/speech/stream -> {args.upstream}", flush=True)
-    async with websockets.serve(handler, args.host, args.port, max_size=None):
-        await asyncio.Future()
+
+    _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__":