Spaces:
Running
Running
| """Real-time streaming TTS server — FastAPI + WebSocket + Web Audio. | |
| Client sends {text}; server runs the frontend (text->phone/tone/lang ids), then the | |
| streaming vocoder, pushing 16-bit PCM chunks over the WebSocket as each 24-frame | |
| chunk (~384 ms audio) is produced. The browser plays them via Web Audio as they | |
| arrive → first sound in ~tens of ms, not after the whole utterance. | |
| Backend: the sherpa-onnx fork's OfflineTtsMbistftStreamModel (C++/ORT) when available; | |
| falls back to the pure-onnxruntime split runner (bit-exact) otherwise. | |
| """ | |
| import os, json, struct, asyncio | |
| from pathlib import Path | |
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect | |
| from fastapi.responses import HTMLResponse | |
| from fastapi.staticfiles import StaticFiles | |
| import numpy as np | |
| HERE = Path(__file__).resolve().parent | |
| os.environ.setdefault("NLTK_DATA", str(HERE / "nltk_data")) | |
| # v2.1 streaming split (3-voice) is hosted in the model repo to avoid the Space's 1 GB LFS | |
| # limit; downloaded (and cached) at startup. Local override via ENC_ONNX/DEC_ONNX. | |
| def _model(fn): | |
| from huggingface_hub import hf_hub_download | |
| return hf_hub_download(os.environ.get("PRIMETTS_REPO", "Luigi/PrimeTTS"), f"v21_streaming/{fn}") | |
| ENC = os.environ.get("ENC_ONNX") or _model("v21_enc.onnx") | |
| DEC = os.environ.get("DEC_ONNX") or _model("v21_dec.onnx") | |
| SR = 16000 | |
| CHUNK, LEFT, RIGHT, HOP, CHAN = 24, 64, 16, 256, 192 # non-causal clean vocoder | |
| # ---- frontend (text -> phone/tone/lang ids, add_blank) ---- | |
| import sys | |
| sys.path.insert(0, str(HERE)) | |
| import frontend_bopomofo as F # bopomofo + arpabet, 88 syms / 6 tones / 2 langs | |
| F.text_to_ids("您好") # warm the g2pw model | |
| def _blank(seq): | |
| o = [0] * (2 * len(seq) + 1) | |
| o[1::2] = seq | |
| return o | |
| import re as _re | |
| # Split AFTER punctuation (keep the delimiter with its clause), then greedily merge | |
| # tiny fragments so each chunk is a reasonable phrase (fast enc, natural prosody). | |
| _CLAUSE_RE = _re.compile(r'(?<=[。！？；;!?\n,，、:：])') | |
| def _split_clauses(text, min_chars=12): | |
| parts = [p for p in _CLAUSE_RE.split(text) if p.strip()] | |
| out, cur = [], "" | |
| for p in parts: | |
| cur += p | |
| if len(cur.strip()) >= min_chars: | |
| out.append(cur); cur = "" | |
| if cur.strip(): | |
| out.append(cur) | |
| return out or [text] | |
| # ---- backend: sherpa-onnx C++ (preferred) or onnxruntime split runner ---- | |
| class Backend: | |
| def __init__(self): | |
| self.kind = None | |
| try: | |
| try: | |
| from sherpa_onnx import OfflineTtsMbistftStreamModel as M | |
| except Exception: | |
| from _sherpa_onnx import OfflineTtsMbistftStreamModel as M | |
| self.m = M(enc=ENC, dec=DEC, num_threads=2, right_lookahead=16) | |
| self.kind = "sherpa" | |
| except Exception as e: | |
| print(f"[backend] sherpa-onnx unavailable ({e}); using onnxruntime split runner") | |
| import onnxruntime as ort | |
| so = ort.SessionOptions(); so.intra_op_num_threads = 2; so.inter_op_num_threads = 1 | |
| self.enc = ort.InferenceSession(ENC, so, providers=["CPUExecutionProvider"]) | |
| self.dec = ort.InferenceSession(DEC, so, providers=["CPUExecutionProvider"]) | |
| self.kind = "onnx" | |
| print(f"[backend] {self.kind}") | |
| def stream(self, phone, tone, lang, noise_scale, length_scale, emit, sid=0): | |
| """Call emit(pcm_bytes) per audio chunk. emit returns False to stop. sid selects voice.""" | |
| x, tn, lg = _blank(phone), _blank(tone), _blank(lang) | |
| if self.kind == "sherpa": | |
| def cb(samples, progress): | |
| return emit(_pcm16(np.asarray(samples, np.float32))) | |
| self.m.generate(x=x, tone=tn, lang=lg, noise_scale=noise_scale, | |
| length_scale=length_scale, callback=cb) | |
| return | |
| # onnxruntime split: enc once (with speaker sid) + chunked overlap-save dec | |
| z = self.enc.run(None, { | |
| "x": np.array([x], np.int64), "tone": np.array([tn], np.int64), | |
| "lang": np.array([lg], np.int64), "x_lengths": np.array([len(x)], np.int64), | |
| "noise_scale": np.array([noise_scale], np.float32), | |
| "length_scale": np.array([length_scale], np.float32), | |
| "sid": np.array([sid], np.int64)})[0] | |
| Tf = z.shape[2] | |
| for a in range(0, Tf, CHUNK): | |
| b = min(a + CHUNK, Tf); s0 = max(0, a - LEFT); e = min(Tf, b + RIGHT) | |
| w = self.dec.run(None, {"z": z[:, :, s0:e]})[0].reshape(-1) | |
| off = (a - s0) * HOP; keep = (b - a) * HOP | |
| if not emit(_pcm16(w[off:off + keep])): | |
| break | |
| def _pcm16(x): | |
| x = np.clip(x, -1.0, 1.0) | |
| return (x * 32767.0).astype("<i2").tobytes() | |
| BACKEND = Backend() | |
| # ---- two preloaded gguf models: story (llama2.c-zh 15M) + chat (SmolLM-Chinese-180M) ---- | |
| import time, random, textgen | |
| import opencc | |
| _CC = opencc.OpenCC("s2twp") # simplified zh -> Taiwan traditional (+phrases) | |
| # LLM weights hosted in the model repo (not bundled in the Space), downloaded at startup. | |
| def _llm(fn): | |
| from huggingface_hub import hf_hub_download | |
| return hf_hub_download(os.environ.get("PRIMETTS_REPO", "Luigi/PrimeTTS"), f"streaming_llm/{fn}") | |
| STORY_GGUF = os.environ.get("STORY_GGUF") or _llm("story15m_q8.gguf") | |
| CHAT_GGUF = os.environ.get("CHAT_GGUF") or _llm("gemma270m_it_q8.gguf") | |
| _PHRASE_END = "。！？…；;!?\n，,、:：" # flush a clause to TTS at these | |
| _TYPE_DELAY = float(os.environ.get("TYPE_DELAY", "0.03")) # pacing for the template fallback | |
| # story mode has NO user input — seed with a random opener for variety each run. | |
| STORY_OPENERS = ["从前，", "很久以前，", "有一天，", "在一个小村庄里，", "从前有一个小男孩，", | |
| "有一只小猫，", "在一片大森林里，", "很久很久以前，有一位国王，"] | |
| def _load_warm(path, n_ctx): | |
| """Load a gguf and prime its compute graph (safe eval-warmup, not a streaming | |
| completion which corrupts state) so the first real request isn't cold.""" | |
| from llama_cpp import Llama | |
| m = Llama(model_path=path, n_ctx=n_ctx, | |
| n_threads=int(os.environ.get("LLM_THREADS", "2")), verbose=False) | |
| try: | |
| m.eval(m.tokenize(b"\xe4\xbd\xa0\xe5\xa5\xbd")) # "你好" | |
| m.reset() | |
| except Exception as e: | |
| print("[gen] warm-up skipped:", e) | |
| return m | |
| class Gen: | |
| def __init__(self): | |
| self.story = None | |
| self.chat = None | |
| for name, path, ctx, attr in [("story", STORY_GGUF, 512, "story"), | |
| ("chat", CHAT_GGUF, 768, "chat")]: | |
| if os.path.exists(path): | |
| try: | |
| setattr(self, attr, _load_warm(path, ctx)) | |
| print(f"[gen] {name} model loaded + warmed:", os.path.basename(path)) | |
| except Exception as e: | |
| print(f"[gen] {name} model unavailable ({e})") | |
| def has_story(self): | |
| return self.story is not None | |
| def has_chat(self): | |
| return self.chat is not None | |
| def stream_tokens(self, prompt, mode="story", min_chars=120, cap_chars=320): | |
| # mode="story" -> llama2.c-zh 15M, no input, random opener -> coherent tale. | |
| # mode="chat" -> SmolLM-180M, user input, ChatML -> fluent response. | |
| # else (no model) -> instant template composer fallback. | |
| if mode == "story" and self.story is not None: | |
| seed = random.choice(STORY_OPENERS) | |
| yield seed | |
| ctx, total = seed, 0 | |
| while total < min_chars and total < cap_chars: | |
| got = 0 | |
| for o in self.story(ctx, max_tokens=140, temperature=0.9, top_p=0.9, | |
| repeat_penalty=1.1, stream=True, stop=["<s>", "</s>"]): | |
| t = o["choices"][0]["text"] | |
| if t: | |
| yield t; ctx += t; total += len(t); got += len(t) | |
| if total >= cap_chars: | |
| break | |
| if got == 0: | |
| break | |
| elif mode == "chat" and self.chat is not None: | |
| # gemma-3-270m-it is an INSTRUCT model — use its chat format with a | |
| # zh-TW instruction so it answers fluently in Traditional Chinese. | |
| p = (prompt or "").strip() or "你好" | |
| instr = "請用繁體中文簡短、自然地回答。" + p | |
| ctx = f"<start_of_turn>user\n{instr}<end_of_turn>\n<start_of_turn>model\n" | |
| for o in self.chat(ctx, max_tokens=200, temperature=0.7, top_p=0.9, | |
| repeat_penalty=1.1, stream=True, | |
| stop=["<end_of_turn>", "<eos>", "<start_of_turn>"]): | |
| t = o["choices"][0]["text"] | |
| if t: | |
| yield t | |
| else: | |
| for tok in textgen.stream(): | |
| time.sleep(_TYPE_DELAY) | |
| yield tok | |
| CHAT = Gen() | |
| app = FastAPI() | |
| async def index(): | |
| return HTMLResponse((HERE / "static" / "index.html").read_text(encoding="utf-8")) | |
| VOICES = [{"sid": 0, "name": "Xinran (♀)"}, {"sid": 1, "name": "Anchen (♂)"}, | |
| {"sid": 2, "name": "Bowen (♂)"}] | |
| async def voices(): | |
| return {"voices": VOICES} | |
| async def healthz(): | |
| return {"backend": BACKEND.kind, "sr": SR, "chunk_ms": CHUNK * HOP // (SR // 1000)} | |
| async def ws(sock: WebSocket): | |
| await sock.accept() | |
| loop = asyncio.get_event_loop() | |
| try: | |
| while True: | |
| msg = json.loads(await sock.receive_text()) | |
| text = (msg.get("text") or "").strip() | |
| if not text: | |
| continue | |
| ns = float(msg.get("noise_scale", 0.667)); ls = float(msg.get("length_scale", 1.0)) | |
| sid = int(msg.get("sid", 0)) | |
| # Split into clauses (string-only, no g2p yet) so BOTH the g2p frontend and the | |
| # encoder run per-clause inside the worker: first audio depends on clause 1, not | |
| # the whole text. (g2pw is BERT-on-CPU — running it over all clauses upfront was | |
| # the real first-audio bottleneck on long text.) | |
| clauses = _split_clauses(text) | |
| await sock.send_text(json.dumps({"type": "start", "sr": SR})) | |
| q: asyncio.Queue = asyncio.Queue() | |
| def emit(pcm): # called from the worker thread | |
| loop.call_soon_threadsafe(q.put_nowait, pcm) | |
| return True | |
| def work(): | |
| for cl in clauses: | |
| o = F.text_to_ids(cl) # g2p this clause only | |
| if not o["phone_ids"]: | |
| continue | |
| BACKEND.stream(o["phone_ids"], o["tone_ids"], o["lang_ids"], ns, ls, emit, sid=sid) | |
| loop.call_soon_threadsafe(q.put_nowait, None) | |
| fut = loop.run_in_executor(None, work) | |
| while True: | |
| pcm = await q.get() | |
| if pcm is None: | |
| break | |
| await sock.send_bytes(pcm) | |
| await fut | |
| await sock.send_text(json.dumps({"type": "end"})) | |
| except WebSocketDisconnect: | |
| pass | |
| async def has_chat(): | |
| return {"chat": True, "story": CHAT.has_story(), "chat_llm": CHAT.has_chat(), | |
| "story_name": "llama2.c-zh-15M", "chat_name": "Gemma-3-270m-it"} | |
| async def chat(sock: WebSocket): | |
| """STREAMING INPUT demo: LLM streams tokens (text) → phrase-buffered into the | |
| TTS → audio streams out. Both over one socket: JSON {type:token,t} for the | |
| live transcript, binary PCM for audio.""" | |
| await sock.accept() | |
| loop = asyncio.get_event_loop() | |
| try: | |
| while True: | |
| msg = json.loads(await sock.receive_text()) | |
| prompt = (msg.get("prompt") or "").strip() | |
| mode = msg.get("mode", "story") | |
| sid = int(msg.get("sid", 0)) | |
| if mode == "chat" and not prompt: # story needs no input; chat does | |
| continue | |
| await sock.send_text(json.dumps({"type": "start", "sr": SR})) | |
| q: asyncio.Queue = asyncio.Queue() | |
| def push(item): # (kind, payload) from worker thread | |
| loop.call_soon_threadsafe(q.put_nowait, item) | |
| def synth_phrase(phrase): | |
| phrase = phrase.strip() | |
| if not phrase: | |
| return | |
| o = F.text_to_ids(phrase) | |
| if not o["phone_ids"]: | |
| return | |
| BACKEND.stream(o["phone_ids"], o["tone_ids"], o["lang_ids"], 0.667, 1.0, | |
| lambda pcm: (push(("pcm", pcm)), True)[1], sid=sid) | |
| def work(): | |
| buf = "" | |
| for tok in CHAT.stream_tokens(prompt, mode): | |
| buf += tok | |
| # at each clause boundary: convert to zh-TW, then display + speak it | |
| while any(c in buf for c in _PHRASE_END): | |
| i = min(buf.index(c) for c in _PHRASE_END if c in buf) | |
| clause = _CC.convert(buf[:i + 1]) # simplified -> Taiwan traditional | |
| push(("tok", clause)) # live zh-TW transcript | |
| synth_phrase(clause) # zh-TW audio | |
| buf = buf[i + 1:] | |
| if buf.strip(): | |
| clause = _CC.convert(buf); push(("tok", clause)); synth_phrase(clause) | |
| push(None) | |
| fut = loop.run_in_executor(None, work) | |
| while True: | |
| item = await q.get() | |
| if item is None: | |
| break | |
| kind, payload = item | |
| if kind == "tok": | |
| await sock.send_text(json.dumps({"type": "token", "t": payload})) | |
| else: | |
| await sock.send_bytes(payload) | |
| await fut | |
| await sock.send_text(json.dumps({"type": "end"})) | |
| except WebSocketDisconnect: | |
| pass | |
| app.mount("/static", StaticFiles(directory=str(HERE / "static")), name="static") | |