feat(morpheus,gateway-chat): streaming decode — TTFA ~4.5s -> ~0.8s
Wrapper gains POST /tts/stream: reads the vLLM token stream, decodes SNAC in WINDOWED
CHUNKS (every 6 frames, decode [2 ctx | 6 | 2 ctx] and emit only the middle 6 — context
both sides => seamless), and streams raw PCM16 (24kHz mono) as it generates. Windowed
(not per-frame) because per-frame CPU decode's per-call overhead x ~60 frames serialized
to ~7s (RTF 2.2); windowed keeps up (RTF ~0.97). Whole-clip /tts kept for non-browser use.
gateway-chat plays the stream via the Web Audio API (fetch reader -> int16->float32 ->
scheduled AudioBufferSourceNodes on a running clock; a new reply supersedes the prior
stream via a generation counter; 🔊 replays). Measured: TTFA 0.80s (was ~4.5s whole-clip),
RTF 0.97, full-duration match. CORS already covers the new route.
Deployed: tts rebuilt on irv-ml1, page pushed to ana-docker.
This commit is contained in:
@@ -11,11 +11,11 @@ pos2/3/5/6->L3, offsets k*4096). Named-speaker prompt: [SOH] "voice: text" [EOT]
|
||||
Zero-shot cloning: [BOS][SOH] ref_text [EOT][SOA][SOS] <ref audio tokens> [EOS_sp] then
|
||||
[SOH] text [EOT][SOA]. Generation stops at end-of-speech (128258).
|
||||
"""
|
||||
import os, io, base64
|
||||
import os, io, re, json, base64
|
||||
import numpy as np, torch, soundfile as sf, requests
|
||||
from scipy.signal import resample_poly
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import Response
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel
|
||||
from transformers import AutoTokenizer
|
||||
@@ -128,3 +128,92 @@ def tts(req: TTSReq):
|
||||
sf.write(buf, audio, 24000, format="WAV", subtype="PCM_16")
|
||||
return Response(content=buf.getvalue(), media_type="audio/wav",
|
||||
headers={"X-Audio-Seconds": f"{len(audio)/24000:.2f}"})
|
||||
|
||||
|
||||
# --- streaming decode: emit raw PCM16 (24kHz mono) as the engine generates -----------
|
||||
# vLLM emits audio tokens as <custom_token_N> strings (token_id = 128256 + N). We decode in
|
||||
# WINDOWED CHUNKS: every CHUNK new frames, decode a small window [CTX left | CHUNK | CTX
|
||||
# right] and emit only the middle CHUNK frames (context on both sides => seamless, no
|
||||
# clicks). O(window) per call and few calls => decode keeps up with gen on CPU (per-frame
|
||||
# decode did NOT — per-call overhead x ~60 frames serialized to seconds). TTFA ~0.7s.
|
||||
CUSTOM_RE = re.compile(r"<custom_token_(\d+)>")
|
||||
FRAME = 2048 # audio samples per SNAC frame
|
||||
|
||||
|
||||
def _decode_frames(ids: list[int]):
|
||||
codes = [t - AUDIO_BASE for t in ids if AUDIO_BASE <= t <= AUDIO_MAX]
|
||||
n = len(codes) // 7
|
||||
if n == 0:
|
||||
return None
|
||||
l1, l2, l3 = [], [], []
|
||||
for i in range(n):
|
||||
f = codes[7 * i:7 * i + 7]
|
||||
c = [f[0], f[1] - 4096, f[2] - 8192, f[3] - 12288, f[4] - 16384, f[5] - 20480, f[6] - 24576]
|
||||
if any(x < 0 or x > 4095 for x in c):
|
||||
return None
|
||||
l1.append(c[0]); l2 += [c[1], c[4]]; l3 += [c[2], c[3], c[5], c[6]]
|
||||
ct = [torch.tensor(l1).unsqueeze(0).to(SNAC_DEVICE),
|
||||
torch.tensor(l2).unsqueeze(0).to(SNAC_DEVICE),
|
||||
torch.tensor(l3).unsqueeze(0).to(SNAC_DEVICE)]
|
||||
with torch.inference_mode():
|
||||
return snac_model.decode(ct).squeeze().cpu().numpy()
|
||||
|
||||
|
||||
def _pcm16(arr) -> bytes:
|
||||
return (np.clip(arr, -1, 1) * 32767).astype(np.int16).tobytes()
|
||||
|
||||
|
||||
@app.post("/tts/stream")
|
||||
def tts_stream(req: TTSReq):
|
||||
prompt_ids = _build_prompt(req)
|
||||
payload = {"model": VLLM_MODEL, "prompt": prompt_ids, "max_tokens": req.max_tokens,
|
||||
"temperature": req.temperature, "top_p": req.top_p, "skip_special_tokens": False,
|
||||
"stop_token_ids": [EOS_SP], "repetition_penalty": req.repetition_penalty, "stream": True}
|
||||
CTX, CHUNK = 2, 6 # context frames each side, frames emitted per decode
|
||||
|
||||
def emit(emitted, total, final):
|
||||
# decode [emitted-CTX : end] with left/right context; return (pcm_bytes, new_emitted)
|
||||
end = total if final else emitted + CHUNK
|
||||
ws = max(0, emitted - CTX)
|
||||
audio = _decode_frames(buf[7 * ws:7 * (end + (0 if final else CTX))])
|
||||
if audio is None:
|
||||
return b"", end
|
||||
off = (emitted - ws) * FRAME
|
||||
chunk = audio[off:] if final else audio[off:off + CHUNK * FRAME]
|
||||
return (_pcm16(chunk) if len(chunk) else b""), end
|
||||
|
||||
buf = []
|
||||
|
||||
def gen():
|
||||
acc, cursor, emitted = "", 0, 0
|
||||
with requests.post(VLLM_URL, json=payload, stream=True, timeout=180) as r:
|
||||
for raw in r.iter_lines():
|
||||
if not raw:
|
||||
continue
|
||||
s = raw.decode("utf-8", "ignore")
|
||||
if not s.startswith("data:"):
|
||||
continue
|
||||
d = s[5:].strip()
|
||||
if d == "[DONE]":
|
||||
break
|
||||
try:
|
||||
j = json.loads(d)
|
||||
except Exception:
|
||||
continue
|
||||
acc += j["choices"][0].get("text", "")
|
||||
for m in CUSTOM_RE.finditer(acc, cursor):
|
||||
cursor = m.end()
|
||||
tid = 128256 + int(m.group(1))
|
||||
if AUDIO_BASE <= tid <= AUDIO_MAX:
|
||||
buf.append(tid)
|
||||
while (len(buf) // 7) - emitted >= CHUNK + CTX: # enough new frames + right ctx
|
||||
b, emitted = emit(emitted, len(buf) // 7, final=False)
|
||||
if b:
|
||||
yield b
|
||||
if (len(buf) // 7) > emitted: # flush the tail
|
||||
b, emitted = emit(emitted, len(buf) // 7, final=True)
|
||||
if b:
|
||||
yield b
|
||||
|
||||
return StreamingResponse(gen(), media_type="application/octet-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Sample-Rate": "24000"})
|
||||
|
||||
Reference in New Issue
Block a user