mesm-tts: ten tiny English voices (ONNX + pdn WebAssembly engine), inference, samples, benchmarks
931612f | #!/usr/bin/env python3 | |
| """mesm-tts text-to-speech with onnxruntime: ten English voices, 2.3-6.5 MB each, CPU only, no GPU or PyTorch. | |
| pip install -r requirements.txt | |
| python tts.py "Hello there. How are you?" --voice hadley --out hello.wav | |
| python tts.py --list | |
| from tts import MesmTTS | |
| tts = MesmTTS("miles") # downloads the voice's ONNX file on first use if it is not next to this script | |
| audio = tts.synth("Good evening.") # float32 numpy array, 24 kHz mono | |
| tts.save("Good evening.", "evening.wav") | |
| Pipeline (the same one the mesmer.tools browser engine runs, see web/): | |
| text -> sentences -> misaki G2P (US English, eSpeak NG fallback for words outside its lexicon; exactly the front end the | |
| models were trained on) -> phoneme ids (tokenizer.json, '$' = 0 at both ends, at most 512 per call) -> ONNX model | |
| (input_ids int64 [1, T], speed float32 [1] -> waveform float32 [1, samples] @ 24 kHz) -> natural sentence joins (Phonon | |
| voices) -> per-voice level (gain_db) with a soft knee above 0.9. | |
| The models draw their own noise, so two runs of the same text differ slightly in the noisy parts. | |
| """ | |
| import argparse, json, os, re, sys, time, wave | |
| import numpy as np | |
| import onnxruntime as ort | |
| REPO_ID = "mesmertech/mesm-tts" | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| SR = 24000 | |
| MAX_IDS = 512 | |
| def _file(path, repo_id=REPO_ID): | |
| """A repo file: next to this script if present (a clone or snapshot of the repo), else from the Hugging Face Hub.""" | |
| local = os.path.join(HERE, path) | |
| if os.path.exists(local): | |
| return local | |
| from huggingface_hub import hf_hub_download | |
| return hf_hub_download(repo_id, path) | |
| # ---- text: sentences, as web/v2/paradee.mjs splitSentences ------------------------------------------------------------ | |
| # Split after . ! ? … (and any closing quotes / brackets) followed by a capital, digit, quote or bracket (the unit the models | |
| # were trained on), except after titles, common abbreviations and initials ("Dr. Smith", "J. R. R. Tolkien", "U.S. Army"). | |
| _ABBREV = re.compile(r'(?:^|[\s"“(])(?:Mr|Mrs|Ms|Dr|St|Jr|Sr|Prof|Rev|Gen|Gov|Sen|Rep|Capt|Lt|Col|Sgt|Adm|Mt|Ft|No|Nos|vs|etc|' | |
| r'Inc|Ltd|Co|Corp|Bros|Jan|Feb|Mar|Apr|Jun|Jul|Aug|Sep|Sept|Oct|Nov|Dec|[A-Z]|(?:[A-Za-z]\.)+[A-Za-z])\.$') | |
| _END = re.compile(r'[.!?…]["”’\')\]]*$') | |
| _BREAK = re.compile(r'\s+(?=[A-Z"“(\d])') | |
| def split_sentences(text): | |
| t = re.sub(r"\s+", " ", text).strip() | |
| pieces, start = [], 0 | |
| for m in _BREAK.finditer(t): | |
| if _END.search(t, 0, m.start()): | |
| pieces.append(t[start:m.start()]); start = m.end() | |
| pieces.append(t[start:]) | |
| out = [] | |
| for p in pieces: | |
| if out and _ABBREV.search(out[-1]): | |
| out[-1] += " " + p | |
| else: | |
| out.append(p) | |
| return [s.strip() for s in out if s.strip()] | |
| def split_phonemes(ps, limit=MAX_IDS - 12): | |
| """A phoneme string longer than the model's 512-id window is cut at the last clause punctuation, else the last word break.""" | |
| out = [] | |
| while len(ps) > limit: | |
| head = ps[:limit] | |
| k = max(head.rfind(p) for p in (", ", "; ", ": ", "— ", "… ")) | |
| k = k + 1 if k > limit / 3 else head.rfind(" ") | |
| if k <= 0: | |
| k = len(head) | |
| out.append(ps[:k].strip()); ps = ps[k:].strip() | |
| if ps: | |
| out.append(ps) | |
| return out | |
| class Tokenizer: | |
| """tokenizer.json is a character map: drop characters outside its set, one id per character, '$' (0) at both ends.""" | |
| def __init__(self, path): | |
| j = json.load(open(path, encoding="utf-8")) | |
| self.vocab = j["model"]["vocab"] | |
| self.strip = re.compile(j["normalizer"]["pattern"]["Regex"]) | |
| def __call__(self, ps): | |
| s = self.strip.sub("", ps) | |
| return np.array(([0] + [self.vocab[c] for c in s if c in self.vocab] + [0])[:MAX_IDS], np.int64) | |
| def make_g2p(): | |
| """misaki's English G2P exactly as Kokoro's KPipeline(lang_code="a") builds it, which is what produced the training data.""" | |
| from misaki import en | |
| try: | |
| from misaki import espeak | |
| fallback = espeak.EspeakFallback(british=False) | |
| except Exception as e: # eSpeak NG missing: words outside misaki's lexicon (mostly rare names) are skipped | |
| print(f"note: no eSpeak NG fallback ({e}); words outside misaki's lexicon will be skipped", file=sys.stderr) | |
| fallback = None | |
| g2p = en.G2P(trf=False, british=False, fallback=fallback, unk="") | |
| return lambda sentence: g2p(sentence)[0] | |
| # ---- audio: natural sentence joins (web/pv/flow.mjs = the reference pv_flow.join), level --------------------------------- | |
| _FR = 240 # 10 ms frames | |
| _FADE = np.linspace(0, 1, _FR, dtype=np.float32) | |
| def _edges(y): | |
| """[first speech sample, end of the last speech frame]; speech = 10 ms frames within 35 dB of the loudest one.""" | |
| n = len(y) // _FR | |
| db = 10 * np.log10((y[:n * _FR].reshape(n, _FR) ** 2).mean(1) + np.float32(1e-12)) | |
| sp = np.where(db > db.max() - 35)[0] | |
| return sp[0] * _FR, (sp[-1] + 1) * _FR | |
| def join(parts, gap_scale=1.0, lead_keep_s=0.02): | |
| """Sentences 2..n keep lead_keep_s of leading silence; sentences 1..n-1 keep gap_scale of their trailing silence.""" | |
| if not parts: | |
| return np.zeros(0, np.float32) | |
| if gap_scale == 1.0: | |
| return np.concatenate(parts) | |
| lead, out = round(lead_keep_s * SR), [] | |
| for i, y in enumerate(parts): | |
| y = y.copy() | |
| a, _ = _edges(y) | |
| if i > 0 and a > lead: | |
| y = y[a - lead:]; y[:_FR] *= _FADE | |
| _, b = _edges(y) | |
| tail = len(y) - b; keep = int(tail * gap_scale) | |
| if i < len(parts) - 1 and keep < tail: | |
| y = y[:b + keep]; y[-_FR:] *= _FADE[::-1] | |
| out.append(y) | |
| return np.concatenate(out) | |
| def level(y, gain_db): | |
| """The voice's gain, then a soft knee above 0.9 (identity below; 0.9 + 0.1 tanh((|x| - 0.9) / 0.1) above): never clips.""" | |
| y = y * np.float32(10 ** (gain_db / 20)) | |
| a = np.abs(y) | |
| return np.where(a <= 0.9, y, np.sign(y) * (0.9 + 0.1 * np.tanh((a - 0.9) / 0.1))).astype(np.float32) | |
| def write_wav(path, y, sr=SR): | |
| pcm = (np.clip(y, -1, 1) * 32767).round().astype("<i2") | |
| with wave.open(path, "wb") as w: | |
| w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr); w.writeframes(pcm.tobytes()) | |
| # ---- the model ------------------------------------------------------------------------------------------------------------ | |
| def load_voices(repo_id=REPO_ID): | |
| return json.load(open(_file("voices.json", repo_id), encoding="utf-8"))["voices"] | |
| class MesmTTS: | |
| """One voice. voice: a key of voices.json ("heart", "wren", "everly", ...). model: Heart's file ("v2f_lite", "v2f", | |
| "v2g"; default v2f_lite). fp32_compute: the *_noint8 twin (Heart only). threads: onnxruntime intra-op threads.""" | |
| def __init__(self, voice="heart", model=None, fp32_compute=False, threads=1, repo_id=REPO_ID, g2p=None): | |
| voices = load_voices(repo_id) | |
| if voice not in voices: | |
| raise ValueError(f"unknown voice {voice!r}; one of: {', '.join(voices)}") | |
| v = self.info = voices[voice] | |
| if "models" in v: # Heart: several files of one voice | |
| m = v["models"][model or v["default_model"]] | |
| onnx = m["onnx_fp32_compute"] if fp32_compute else m["onnx"] | |
| else: | |
| if model: | |
| raise ValueError(f"{voice} has a single model file") | |
| onnx = v["onnx"] | |
| self.gain_db, self.join = v["gain_db"], v.get("join") | |
| so = ort.SessionOptions() | |
| so.intra_op_num_threads, so.inter_op_num_threads = threads, 1 | |
| self.session = ort.InferenceSession(_file(onnx, repo_id), so, providers=["CPUExecutionProvider"]) | |
| self.tokenize = Tokenizer(_file("tokenizer.json", repo_id)) | |
| self.g2p = g2p or make_g2p() | |
| def plan(self, text): | |
| """text -> [(sentence, phonemes, ids)], long sentences split to fit the model's window.""" | |
| out = [] | |
| for s in split_sentences(text): | |
| for ps in split_phonemes(self.g2p(s)): | |
| ids = self.tokenize(ps) | |
| if len(ids) > 2: | |
| out.append((s, ps, ids)) | |
| return out | |
| def synth_ids(self, ids, speed=1.0): | |
| """One window of phoneme ids -> raw float32 audio (no join, no gain).""" | |
| return self.session.run(None, {"input_ids": ids[None], "speed": np.array([speed], np.float32)})[0][0] | |
| def synth(self, text, speed=1.0): | |
| """Whole text -> float32 audio @ 24 kHz, joined and levelled as the browser plays it. speed > 1 talks faster.""" | |
| parts = [self.synth_ids(ids, speed) for _, _, ids in self.plan(text)] | |
| j = self.join or {} | |
| return level(join(parts, j.get("gap_scale", 1.0), j.get("lead_keep_s", 0.02)), self.gain_db) | |
| def save(self, text, path, speed=1.0): | |
| y = self.synth(text, speed) | |
| write_wav(path, y) | |
| return y | |
| def main(): | |
| ap = argparse.ArgumentParser(description="mesm-tts text-to-speech (onnxruntime, CPU).") | |
| ap.add_argument("text", nargs="?", help="text to speak (or --file)") | |
| ap.add_argument("-v", "--voice", default="heart", help="voice id (see --list); default heart") | |
| ap.add_argument("-o", "--out", default="out.wav") | |
| ap.add_argument("-f", "--file", help="read the text from this file") | |
| ap.add_argument("--model", help="Heart only: v2f_lite (default), v2f or v2g") | |
| ap.add_argument("--fp32-compute", action="store_true", help="Heart only: the *_noint8 twin (no int8 matrix products)") | |
| ap.add_argument("--speed", type=float, default=1.0, help="speaking rate (1 = as trained; 1.2 = 20%% faster)") | |
| ap.add_argument("--threads", type=int, default=1, help="onnxruntime threads (default 1; the models are small)") | |
| ap.add_argument("--repo", default=REPO_ID, help="Hugging Face repo to fetch files from when they are not local") | |
| ap.add_argument("--list", action="store_true", help="list the voices and exit") | |
| a = ap.parse_args() | |
| if a.list: | |
| for k, v in load_voices(a.repo).items(): | |
| print(f"{k:8s} {v['name']:7s} {v['gender']:6s} {v['accent'].upper():3s} {v['tag']:9s} from {v['teacher']}") | |
| return | |
| text = open(a.file, encoding="utf-8").read() if a.file else a.text | |
| if not text: | |
| ap.error("give the text to speak, or --file") | |
| t0 = time.perf_counter() | |
| tts = MesmTTS(a.voice, a.model, a.fp32_compute, a.threads, a.repo) | |
| t1 = time.perf_counter() | |
| plan = tts.plan(text) | |
| t2 = time.perf_counter() | |
| parts = [tts.synth_ids(ids, a.speed) for _, _, ids in plan] | |
| t3 = time.perf_counter() | |
| j = tts.join or {} | |
| y = level(join(parts, j.get("gap_scale", 1.0), j.get("lead_keep_s", 0.02)), tts.gain_db) | |
| write_wav(a.out, y) | |
| dur = len(y) / SR | |
| print(f"{a.out}: {dur:.1f} s of audio ({tts.info['name']}). load {1e3 * (t1 - t0):.0f} ms, phonemes {1e3 * (t2 - t1):.0f} ms, " | |
| f"synthesis {1e3 * (t3 - t2):.0f} ms = {dur / max(t3 - t2, 1e-9):.0f}x real time on {a.threads} thread(s)") | |
| if __name__ == "__main__": | |
| main() | |