from __future__ import annotations import re import time from pathlib import Path from typing import Any import numpy as np import psutil from .assets import PromptAssets, default_assets_dir from .backbone import Backbone from .depth import Depth from .codec import MimiDecoder from .prompt import AudioIn, build_prompt from .sampling import gumbel from .control import ( expected_frames, frame_cap, min_frames, split_text, StopPolicy, sample_cb0, ) # Speakable text check: must contain at least one alphanumeric character ALPHANUM_PATTERN = re.compile(r"[\w\d]", re.UNICODE) def _softmax_topk(logits: np.ndarray, temp: float = 0.7, top_k: int = 64) -> np.ndarray: logits = np.asarray(logits, dtype=np.float32) k = min(top_k, logits.shape[-1]) topk_idx = np.argpartition(logits, -k)[-k:] masked = np.full_like(logits, -np.inf) masked[topk_idx] = logits[topk_idx] / temp max_val = np.max(masked[topk_idx]) exp_vals = np.exp(masked - max_val) exp_vals[masked == -np.inf] = 0.0 sum_exp = np.sum(exp_vals) if sum_exp > 0: return exp_vals / sum_exp probs = np.zeros_like(logits) probs[topk_idx] = 1.0 / len(topk_idx) return probs class MayaFast: def __init__( self, assets_dir: str | Path | None = None, threads: int | None = None, ) -> None: if assets_dir is None: self.assets_dir = default_assets_dir() else: self.assets_dir = Path(assets_dir) self.threads = threads self.assets = PromptAssets(self.assets_dir, verify=True) self.shapes = self.assets.shapes self.controls = self.assets.controls self.sampling_cfg = self.controls.get("sampling", {}) # Load models. Depth first: its session peaks ~1 GB above steady state # during graph optimization, so loading it before the backbone keeps # that transient off the combined peak. # Graph file names come from the manifest ("graphs"), defaulting to the # v1 names (see assets.DEFAULT_GRAPHS). self.graphs = self.assets.graphs depth_path = self.assets_dir / self.graphs["depth"] self.depth = Depth(depth_path, threads=self.threads) backbone_path = self.assets_dir / self.graphs["backbone"] self.backbone = Backbone( backbone_path, threads=self.threads, config=self.shapes, ) mimi_path = self.assets_dir / self.graphs["mimi"] self.mimi = MimiDecoder(mimi_path, threads=self.threads) self.audio_in = self.assets.audio_in @property def voices(self) -> list[str]: return sorted(list(self.assets.voices_info.keys())) def synthesize( self, text: str, voice: str = "Charvi", seed: int = 0, chunk_s: float | None = None, lang: str = "mr", ) -> tuple[np.ndarray, dict[str, Any]]: """ Synthesize text into 24 kHz float32 audio. Returns: (wav, metrics) """ # Review focus #2: Empty or punctuation-only text: expected ValueError("no speakable text") if not text or not ALPHANUM_PATTERN.search(text): raise ValueError("no speakable text") if voice not in self.assets.voices_info: raise KeyError(f"Unknown voice {voice!r}. Available: {self.voices}") t_wall_start = time.perf_counter() prior = self.controls.get("text_frames_prior", {}) sampling_cfg = self.sampling_cfg q_serve = int(sampling_cfg.get("q_serve", 16)) temp = float(sampling_cfg.get("temp", 0.7)) top_k = int(sampling_cfg.get("top_k", 64)) tau = float(sampling_cfg.get("tau", 0.35)) hazard = float(sampling_cfg.get("hazard", 0.6)) aeos_class = int(self.shapes.get("n_codes", 2048)) # Split text into chunks c_s = 12.0 if chunk_s is None else chunk_s chunks = split_text(text, lang, prior, chunk_s=c_s) rng = np.random.default_rng(seed) # Speaker memory for Depth: named voices use a null memory (use_mem False) d_mem = int(self.shapes.get("d_mem", 1024)) mem = np.zeros((1, 1, d_mem), dtype=np.float32) gen_codes: list[np.ndarray] = [] reasons: list[str] = [] caps: list[int] = [] max_past_len = 0 max_chunk_bound = 0 backbone_time = 0.0 depth_time = 0.0 mimi_time = 0.0 for chunk_text in chunks: # Review focus #1: Each piece a fresh KV session self.backbone.reset() prompt = build_prompt(chunk_text, voice, self.assets, lang=lang) ef, er = expected_frames(prior, lang, len(chunk_text)) cap = frame_cap(chunk_text, lang, ef, er, self.controls) min_f = min_frames(ef) caps.append(cap) chunk_bound = prompt.n_prompt + cap if chunk_bound > max_chunk_bound: max_chunk_bound = chunk_bound # Initial backbone step on prompt t0 = time.perf_counter() h, cb0_logits = self.backbone.step(prompt.embeds) backbone_time += time.perf_counter() - t0 if self.backbone.past_len > max_past_len: max_past_len = self.backbone.past_len stop_policy = StopPolicy(min_frames=min_f, tau=tau, hazard=hazard) chunk_frames: list[np.ndarray] = [] reason = "cap" for t in range(cap): # 1. Softmax with top_k mask over cb0_logits p = _softmax_topk(cb0_logits, temp=temp, top_k=top_k) # 2. StopPolicy check stop_reason = stop_policy.check(t, p, aeos_class=aeos_class) if stop_reason is not None: reason = stop_reason break # 3. Sample cb0 (aeos zeroed, Gumbel-max over the rest) cb0_noise = gumbel(rng, (cb0_logits.shape[-1],)) cb0 = sample_cb0(p, cb0_noise, aeos_class=aeos_class) # 4. Depth frame depth_noise = gumbel(rng, (15, int(self.shapes.get("n_codes", 2048)))) t_d = time.perf_counter() frame_codes = self.depth.frame(h, cb0, mem, depth_noise) depth_time += time.perf_counter() - t_d chunk_frames.append(frame_codes) # 5. AudioIn feedback fb = self.audio_in.frames(frame_codes[None, :], q_use=q_serve) # 6. Next backbone step t_b = time.perf_counter() h, cb0_logits = self.backbone.step(fb) backbone_time += time.perf_counter() - t_b if self.backbone.past_len > max_past_len: max_past_len = self.backbone.past_len reasons.append(reason) # Note on upstream aeos step: # Upstream does one self._step(audio_in.aeos(1)) after the loop if reason.startswith("aeos"). # However, because each chunk resets the KV session and no further steps read from this KV session, # running the final aeos step changes nothing returned and wastes compute. We skip it here. if chunk_frames: # chunk_frames: list of (16,) chunk_codes_arr = np.stack(chunk_frames, axis=1) # (16, F_chunk) gen_codes.append(chunk_codes_arr) if gen_codes: all_codes = np.concatenate(gen_codes, axis=1) # (16, Total_F) total_frames = all_codes.shape[1] t_m = time.perf_counter() wav = self.mimi.decode(all_codes) mimi_time += time.perf_counter() - t_m else: total_frames = 0 wav = np.zeros(0, dtype=np.float32) wall_s = time.perf_counter() - t_wall_start audio_s = len(wav) / 24000.0 rtf = wall_s / max(1e-6, audio_s) # Peak RSS in MB process = psutil.Process() peak_rss_mb = process.memory_info().rss / (1024 * 1024) metrics = { "audio_s": audio_s, "wall_s": wall_s, "rtf": rtf, "frames": total_frames, "reasons": reasons, "caps": caps, "backbone_s": backbone_time, "depth_s": depth_time, "mimi_s": mimi_time, "peak_rss_mb": peak_rss_mb, "max_past_len": max_past_len, "max_chunk_bound": max_chunk_bound, } return wav, metrics