"""indic-speak on MLX (4-bit LM) + torch SNAC/Vocos decode. from tts_mlx import MLXTTS tts = MLXTTS() # models/indic-speak-mlx-4bit wav, stats = tts("नमस्ते", speaker="Amit") # float32 @ 24 kHz """ from __future__ import annotations import re import sys import time from pathlib import Path import mlx.core as mx import numpy as np import torch from huggingface_hub import snapshot_download from mlx_lm import load, stream_generate from mlx_lm.sample_utils import make_logits_processors, make_sampler REPO = "bodhan-ai/indic-speak" MLX_PATH = Path(__file__).parent / "models" / "indic-speak-mlx-4bit" SR = 24_000 class MLXTTS: sample_rate = SR def __init__(self, mlx_path: str | Path = MLX_PATH, decode_device: str = "mps"): # Accept a local dir or a Hub repo id. Self-contained repos ship vocos/ and # inference.py; otherwise fall back to the original bodhan-ai/indic-speak snapshot. mlx_path = Path(mlx_path) if Path(mlx_path).is_dir() else Path(snapshot_download(str(mlx_path))) if (mlx_path / "vocos" / "best.pt").is_file(): snap = str(mlx_path) else: snap = snapshot_download(REPO) sys.path.insert(0, snap) from inference import build_prompt, ids_to_codes, _load_vocos # repo helpers from snac import SNAC self._build_prompt, self._ids_to_codes = build_prompt, ids_to_codes self.model, self.tok = load(str(mlx_path)) self._base = self.tok.convert_tokens_to_ids("<|snac_0|>") self._eos = self.tok.convert_tokens_to_ids("<|end_of_speech|>") self._eos_set = {self._eos, self.tok.eos_token_id} self.dev = decode_device self.snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(self.dev).eval() self.vocos = _load_vocos(str(Path(snap) / "vocos" / "best.pt"), self.dev) def __call__(self, text: str, speaker: str = "Amit", style: str = "", *, temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2, max_tokens: int = 2520, seed: int | None = None): if seed is not None: mx.random.seed(seed) prompt = self._build_prompt(self.tok, text, speaker, style) sampler = make_sampler(temp=temperature, top_p=top_p, top_k=top_k) procs = make_logits_processors(repetition_penalty=repetition_penalty) ids, t0, t_first = [], time.time(), None for r in stream_generate(self.model, self.tok, prompt, max_tokens=max_tokens, sampler=sampler, logits_processors=procs): if t_first is None: t_first = time.time() - t0 if r.token in self._eos_set: break ids.append(r.token) t_gen = time.time() - t0 with torch.no_grad(): codes = self._ids_to_codes(ids, self._base, self.dev) z_q = self.snac.quantizer.from_codes(codes) wav = self.vocos(z_q.float())[0, 0].clamp(-1, 1).float().cpu().numpy() t_total = time.time() - t0 stats = dict(tokens=len(ids), tok_per_s=len(ids) / t_gen, ttft=t_first, gen_s=t_gen, total_s=t_total, audio_s=len(wav) / SR, rtf=t_total / (len(wav) / SR), hit_cap=len(ids) >= max_tokens) return wav, stats def stream(self, text: str, speaker: str = "Amit", style: str = "", *, temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2, max_tokens: int = 2520, seed: int | None = None, chunk_frames: int = 8, ctx_frames: int = 12, lookahead: int = 4): """Yield float32 audio chunks while tokens are still being generated. Every `chunk_frames` SNAC frames, decode [emitted-ctx, generated] through Vocos and emit frames [emitted, generated-lookahead). Vocos is non-causal, so `lookahead` frames of right context keep the seams clean; `ctx_frames` of left context keeps the receptive field warm. 1 frame = 2048 samples. Yields (chunk, stats) tuples; the last stats has done=True. """ if seed is not None: mx.random.seed(seed) prompt = self._build_prompt(self.tok, text, speaker, style) sampler = make_sampler(temp=temperature, top_p=top_p, top_k=top_k) procs = make_logits_processors(repetition_penalty=repetition_penalty) hi = self._base + 7 * 4096 ids, t0, t_first, emitted = [], time.time(), None, 0 def decode(frames_from: int, frames_to: int) -> np.ndarray: with torch.no_grad(): codes = self._ids_to_codes(ids[frames_from * 7: frames_to * 7], self._base, self.dev) z = self.snac.quantizer.from_codes(codes) return self.vocos(z.float())[0, 0].clamp(-1, 1).float().cpu().numpy() for r in stream_generate(self.model, self.tok, prompt, max_tokens=max_tokens, sampler=sampler, logits_processors=procs): if r.token in self._eos_set or not (self._base <= r.token < hi): break ids.append(r.token) n_frames = len(ids) // 7 ready = n_frames - lookahead if ready - emitted >= chunk_frames: start = max(0, emitted - ctx_frames) wav = decode(start, n_frames) new = wav[(emitted - start) * 2048:(ready - start) * 2048] if t_first is None: t_first = time.time() - t0 emitted = ready yield new, dict(done=False, tokens=len(ids), ttfa=t_first, elapsed=time.time() - t0) n_frames = len(ids) // 7 if n_frames > emitted: start = max(0, emitted - ctx_frames) wav = decode(start, n_frames) new = wav[(emitted - start) * 2048:] if t_first is None: t_first = time.time() - t0 emitted = n_frames yield new, dict(done=False, tokens=len(ids), ttfa=t_first, elapsed=time.time() - t0) el = time.time() - t0 yield np.zeros(0, np.float32), dict(done=True, tokens=len(ids), ttfa=t_first, elapsed=el, tok_per_s=len(ids) / max(el, 1e-6), audio_s=emitted * 2048 / SR, hit_cap=len(ids) >= max_tokens) _SENT_SPLIT = re.compile(r"(?<=[।.!?])\s+") def split_sentences(text: str) -> list[str]: return [p.strip() for p in _SENT_SPLIT.split(text.strip()) if p.strip()] def trim_silence(wav: np.ndarray, thr: float = 0.01, keep_s: float = 0.15, frame: int = 480) -> np.ndarray: """Cut leading/trailing near-silence, keeping `keep_s` of padding.""" n = len(wav) // frame * frame loud = np.abs(wav[:n]).reshape(-1, frame).max(1) > thr if not loud.any(): return wav[: int(keep_s * SR)] i0 = max(0, (np.argmax(loud) * frame) - int(keep_s * SR)) i1 = min(len(wav), (len(loud) - np.argmax(loud[::-1])) * frame + int(keep_s * SR)) return wav[i0:i1] def speak_long(tts: "MLXTTS", text: str, gap_s: float = 0.25, on_chunk=None, **kw): """Sentence-chunked synthesis (4-bit models truncate on long inputs). Each sentence gets its own generation with a token cap of ~2x its expected duration, silence is trimmed, and chunks are joined with a short gap. `on_chunk(wav, stats)` is called per sentence -> use it for streaming playback. """ parts, stats = [], [] gap = np.zeros(int(gap_s * SR), np.float32) for sent in split_sentences(text): cap = int((2 * len(sent) / 12 + 3) * 98) # ~98 tok per audio-second wav, st = tts(sent, max_tokens=min(cap, 2520), **kw) wav = trim_silence(wav) st["audio_s"] = len(wav) / SR if on_chunk: on_chunk(wav, st) parts += [wav, gap] stats.append(st) return np.concatenate(parts[:-1]) if parts else np.zeros(0, np.float32), stats def stream_long(tts: "MLXTTS", text: str, gap_s: float = 0.25, **kw): """Sentence-chunked streaming: yields (chunk, stats) with stats['sentence'] index.""" gap = np.zeros(int(gap_s * SR), np.float32) sents = split_sentences(text) for i, sent in enumerate(sents): cap = int((2 * len(sent) / 12 + 3) * 98) for chunk, st in tts.stream(sent, max_tokens=min(cap, 2520), **kw): st["sentence"] = i; st["n_sentences"] = len(sents) yield chunk, st if i < len(sents) - 1: yield gap, dict(done=False, sentence=i, n_sentences=len(sents), gap=True) if __name__ == "__main__": import soundfile as sf text = sys.argv[1] if len(sys.argv) > 1 else "नमस्ते, आज हम विज्ञान पढ़ेंगे। प्रकाश की चाल लगभग तीन लाख किलोमीटर प्रति सेकंड होती है।" t = time.time(); tts = MLXTTS(); print(f"loaded in {time.time()-t:.1f}s") for i in range(2): wav, s = tts(text, seed=1234) print({k: round(v, 2) if isinstance(v, float) else v for k, v in s.items()}) Path("out").mkdir(exist_ok=True); sf.write("out/mlx.wav", wav, SR); print("wrote out/mlx.wav")