vh
0655a37bf6
- max_tokens default 2400->3500 (~42s) in wrapper + gateway-chat client, with a _cap() clamp so prompt+gen never exceeds MAX_CTX (4096) — a cloning ref block is ~1100 tokens, so an unclamped 3500 would overflow context on the clone path. - Staged clone voices: /voices dir of <name>.wav + <name>.txt, each encoded to its Orpheus reference block at startup; voice="<name>" zero-shot clones it. Beatrice (a chatterbox reference) staged as the first normal-voice clone. GET /voices lists baddy + clones. - compose: mount voices dir + pass MORPHEUS_MAX_LEN to the wrapper (clamp must match engine). vLLM concurrency (measured, --max-num-seqs 8, 250-tok reqs): near-linear batching — 8 concurrent finish in the same ~2.8s as 1 (707 tok/s, 8.1x single, flat per-req latency). Chunked-sentence production can fan out for ~8x throughput; CPU SNAC decode is the scale bottleneck, not generation.
246 lines
11 KiB
Python
246 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""mOrpheus TTS wrapper — turns the vLLM engine's Orpheus audio tokens into 24kHz WAV.
|
|
|
|
Architecture: this CPU service builds the Orpheus prompt (as raw token ids), calls the
|
|
vLLM engine (which serves the mOrpheus LLM), recovers the generated audio-token ids by
|
|
re-tokenizing the returned text (vLLM emits them as `<custom_token_N>` strings when
|
|
skip_special_tokens=false), then SNAC-decodes them to audio.
|
|
|
|
Token scheme (verified): audio-base 128266, 7-token SNAC frames (pos0->L1, pos1/4->L2,
|
|
pos2/3/5/6->L3, offsets k*4096). Named-speaker prompt: [SOH] "voice: text" [EOT][SOA].
|
|
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, 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, StreamingResponse
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
from transformers import AutoTokenizer
|
|
from snac import SNAC
|
|
|
|
MODEL_DIR = os.environ.get("MORPHEUS_MODEL_DIR", "/model")
|
|
SNAC_DIR = os.environ.get("SNAC_DIR", "/snac")
|
|
VLLM_URL = os.environ.get("VLLM_URL", "http://vllm-morpheus:8000/v1/completions")
|
|
VLLM_MODEL = os.environ.get("VLLM_MODEL", "morpheus")
|
|
SNAC_DEVICE = os.environ.get("SNAC_DEVICE", "cpu")
|
|
DEFAULT_VOICE = os.environ.get("MORPHEUS_DEFAULT_VOICE", "baddy")
|
|
VOICES = [v for v in os.environ.get("MORPHEUS_VOICES", "baddy").split(",") if v]
|
|
AUDIO_MAX = 156937
|
|
AUDIO_BASE, SOS, EOS_SP, SOH, SOA, EOT, BOS = 128266, 128257, 128258, 128259, 128260, 128009, 128000
|
|
MAX_CTX = int(os.environ.get("MORPHEUS_MAX_LEN", "4096")) # must match the engine --max-model-len
|
|
|
|
|
|
def _cap(prompt_ids: list[int], want: int) -> int:
|
|
# never let prompt + gen exceed the context (a cloning ref block is ~1100 tokens)
|
|
return max(64, min(want, MAX_CTX - len(prompt_ids) - 16))
|
|
|
|
tok = AutoTokenizer.from_pretrained(MODEL_DIR)
|
|
snac_model = SNAC.from_pretrained(SNAC_DIR).to(SNAC_DEVICE).eval()
|
|
app = FastAPI(title="mOrpheus TTS", version="0.1.0")
|
|
# INTERNAL RESEARCH: browser callers (e.g. gateway-chat on ana-docker:8091) fetch this
|
|
# cross-origin; allow all origins on this internal-only endpoint.
|
|
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"],
|
|
allow_headers=["*"], expose_headers=["X-Audio-Seconds"])
|
|
|
|
|
|
class TTSReq(BaseModel):
|
|
text: str
|
|
voice: str = DEFAULT_VOICE
|
|
temperature: float = 0.6
|
|
top_p: float = 0.95
|
|
max_tokens: int = 3500 # ~42s of audio (auto-clamped to fit MAX_CTX; see _cap)
|
|
repetition_penalty: float = 1.1 # LOAD-BEARING: rep 1.0 => model never stops (rambles to
|
|
# the cap); rep 1.1 => clean end-of-speech. Keep <=1.1 for
|
|
# cloning (higher penalizes the in-context ref audio tokens).
|
|
reference_audio_b64: str | None = None # optional zero-shot clone: base64 WAV
|
|
reference_text: str | None = None # transcript of the reference
|
|
|
|
|
|
def _encode_ref(wav_bytes: bytes, ref_text: str) -> list[int]:
|
|
wav, sr = sf.read(io.BytesIO(wav_bytes))
|
|
if wav.ndim > 1:
|
|
wav = wav.mean(1)
|
|
wav = wav.astype(np.float32)
|
|
if sr != 24000:
|
|
wav = resample_poly(wav, 24000, sr).astype(np.float32)
|
|
wt = torch.tensor(wav, device=SNAC_DEVICE).view(1, 1, -1)
|
|
with torch.inference_mode():
|
|
codes = snac_model.encode(wt)
|
|
L1, L2, L3 = [c.squeeze(0).tolist() for c in codes]
|
|
ids = []
|
|
for i in range(len(L1)):
|
|
ids += [L1[i] + AUDIO_BASE, L2[2 * i] + AUDIO_BASE + 4096, L3[4 * i] + AUDIO_BASE + 8192,
|
|
L3[4 * i + 1] + AUDIO_BASE + 12288, L2[2 * i + 1] + AUDIO_BASE + 16384,
|
|
L3[4 * i + 2] + AUDIO_BASE + 20480, L3[4 * i + 3] + AUDIO_BASE + 24576]
|
|
return [BOS, SOH] + tok(ref_text, add_special_tokens=False).input_ids + [EOT, SOA, SOS] + ids + [EOS_SP]
|
|
|
|
|
|
# Pre-staged clone voices: a /voices dir of <name>.wav + <name>.txt (its transcript). Each is
|
|
# encoded once to its Orpheus reference block at startup, so voice="<name>" zero-shot clones it.
|
|
VOICES_DIR = os.environ.get("MORPHEUS_VOICES_DIR", "/voices")
|
|
CLONE_REFS = {}
|
|
if os.path.isdir(VOICES_DIR):
|
|
for fn in sorted(os.listdir(VOICES_DIR)):
|
|
if fn.endswith(".wav") and os.path.exists(os.path.join(VOICES_DIR, fn[:-4] + ".txt")):
|
|
name = fn[:-4]
|
|
try:
|
|
CLONE_REFS[name] = _encode_ref(open(os.path.join(VOICES_DIR, fn), "rb").read(),
|
|
open(os.path.join(VOICES_DIR, name + ".txt")).read().strip())
|
|
print(f"[voices] staged clone voice '{name}'", flush=True)
|
|
except Exception as e:
|
|
print(f"[voices] FAILED to stage '{name}': {e}", flush=True)
|
|
|
|
|
|
def _build_prompt(req: TTSReq) -> list[int]:
|
|
if req.reference_audio_b64 and req.reference_text:
|
|
ref = _encode_ref(base64.b64decode(req.reference_audio_b64), req.reference_text)
|
|
return ref + [SOH] + tok(req.text, add_special_tokens=False).input_ids + [EOT, SOA]
|
|
if req.voice in CLONE_REFS: # pre-staged clone voice (e.g. beatrice)
|
|
return CLONE_REFS[req.voice] + [SOH] + tok(req.text, add_special_tokens=False).input_ids + [EOT, SOA]
|
|
return [SOH] + tok(f"{req.voice}: {req.text}").input_ids + [EOT, SOA]
|
|
|
|
|
|
def _decode(ids: list[int]):
|
|
if SOS in ids:
|
|
ids = ids[len(ids) - 1 - ids[::-1].index(SOS) + 1:]
|
|
codes = [t - AUDIO_BASE for t in ids if AUDIO_BASE <= t <= AUDIO_MAX]
|
|
l1, l2, l3 = [], [], []
|
|
for i in range(len(codes) // 7):
|
|
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):
|
|
continue
|
|
l1.append(c[0]); l2 += [c[1], c[4]]; l3 += [c[2], c[3], c[5], c[6]]
|
|
if not l1:
|
|
return None
|
|
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()
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "voices": VOICES + list(CLONE_REFS), "default": DEFAULT_VOICE, "engine": VLLM_URL}
|
|
|
|
|
|
@app.get("/voices")
|
|
def voices():
|
|
return {"voices": VOICES + list(CLONE_REFS), "cloned": list(CLONE_REFS), "default": DEFAULT_VOICE}
|
|
|
|
|
|
@app.post("/tts")
|
|
def tts(req: TTSReq):
|
|
prompt_ids = _build_prompt(req)
|
|
payload = {"model": VLLM_MODEL, "prompt": prompt_ids, "max_tokens": _cap(prompt_ids, 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}
|
|
try:
|
|
r = requests.post(VLLM_URL, json=payload, timeout=180)
|
|
except requests.RequestException as e:
|
|
raise HTTPException(502, f"engine unreachable: {e}")
|
|
if r.status_code != 200:
|
|
raise HTTPException(502, f"engine {r.status_code}: {r.text[:200]}")
|
|
text = r.json()["choices"][0]["text"]
|
|
gen_ids = tok(text, add_special_tokens=False).input_ids
|
|
audio = _decode(gen_ids)
|
|
if audio is None:
|
|
raise HTTPException(500, "no audio tokens generated")
|
|
buf = io.BytesIO()
|
|
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": _cap(prompt_ids, 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"})
|