#!/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 `` 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] [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 .wav + .txt (its transcript). Each is # encoded once to its Orpheus reference block at startup, so voice="" 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 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"") 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"})