File size: 16,940 Bytes
23c283a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 | """Maya 2 generation for one request at a time: the official engine's per-piece loop (maya3/engine.py) on MLX.
from maya2_mlx import Maya2
tts = Maya2() # models/maya2-8bit: dynamic 8-bit backbone, fp16 depth decoder, fp32 Mimi
wav = tts.generate("Hello there.", lang="en", voice="Anirban", seed=7) # float32, 24 kHz, in [-1, 1]
Same request semantics as Maya3TTS.generate (api.py:224-230 -> engine.py): split into ~20 s pieces, one fresh
backbone context per piece, controlled stop, duration cap and loop guard, 0.15 s gap between pieces, int16 PCM / 32767.
"""
import math
import os
import random
import time
import mlx.core as mx
import numpy as np
from . import philox
from . import quant as Q
from . import weights as W
from .backbone import Backbone
from .depth import Depth
from .mimi import Mimi
from .prompt import Packer, VoiceRegistry, expected_frames, load_named_voices, load_prior, split_text
PORT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SR = 24000 # engine.py:42
FRAME = 1920 # engine.py:43
AEOS = 2048 # tts/tokens.py:23
REASONS = {1: "aeos", 2: "aeos_argmax", 3: "aeos_tau", 4: "aeos_hazard"} # engine.py:46
# EngineConfig defaults (engine.py:49-79) that the single-request loop reads
CB0_TEMP, CB0_TOP_K, CB0_TOP_P, CB0_REP, CB0_REP_WIN = 0.9, 64, 1.0, 0.8, 12
DEPTH_TEMP, DEPTH_TOP_K = 0.7, 64
HARD_CAP, CAPACITY, SEED = 1500, 2048, 20261024
GAP_S, CHUNK_S = 0.15, 20.0
WARMUP_PIECES = 2 # Engine.warmup runs two one-piece requests before any user request (engine.py:243-249)
class GenerationError(RuntimeError):
"""api.py:31-34"""
def __init__(self, code, msg):
super().__init__(f"{code}: {msg}")
self.code = code
def check_sampling(sp):
"""engine.py:110-128"""
if not sp:
return {}
allowed = {"seed", "temperature", "top_k", "top_p", "repetition_penalty"}
bad = set(sp) - allowed
if bad:
raise ValueError(f"unknown sampling keys {sorted(bad)}; allowed {sorted(allowed)}")
out = {k: v for k, v in sp.items() if v is not None}
if "temperature" in out and not 0.05 <= float(out["temperature"]) <= 2.0:
raise ValueError("temperature must be in [0.05, 2.0]")
if "top_k" in out and not 1 <= int(out["top_k"]) <= 2049:
raise ValueError("top_k must be in [1, 2049]")
if "top_p" in out and not 0.0 < float(out["top_p"]) <= 1.0:
raise ValueError("top_p must be in (0, 1]")
if "repetition_penalty" in out and not 0.0 <= float(out["repetition_penalty"]) <= 5.0:
raise ValueError("repetition_penalty must be in [0, 5]")
if "seed" in out:
out["seed"] = int(out["seed"])
return out
class Maya2:
def __init__(self, models_dir=W.DEFAULT_DIR, strict=False):
"""strict=False: production precision (backbone bf16, depth fp16, Mimi fp32 without TF32).
strict=True: fp32 everywhere (the parity gates)."""
from maya3 import guards as G # the official file, unchanged
self.G = G
t0 = time.perf_counter()
meta = W.load_meta(models_dir)
self.meta = meta
self.models_dir = models_dir
self.bdt = mx.float32 if strict else mx.bfloat16 # engine.py:158 (model dtype)
self.ddt = mx.float32 if strict else mx.float16 # engine.py:192, depth_dtype "fp16"
self.backbone = Backbone(meta, self.bdt, models_dir, capacity=CAPACITY)
self.recipe = Q.load_recipe(models_dir)
if strict and self.recipe:
raise ValueError("strict (fp32) runs use the bf16 release export, not quantized weights")
aud = W.load(["audio_in.safetensors"], models_dir)
table = Q.cast(Q.take(aud, "table.weight", self.recipe), self.bdt)
fbw = aud["norm.weight"].astype(self.bdt)
self.depth = Depth(meta, (table, fbw), self.ddt, DEPTH_TEMP, DEPTH_TOP_K, models_dir)
self.depth.set_speaker(None) # named / plain: the null speaker memory
self.mimi = Mimi(meta, models_dir)
self.anchor = self._anchor(W.load(["anchor.safetensors"], models_dir))
self.packer = Packer(os.path.join(models_dir, "tokenizer.json"), meta["token_ids"])
self.prior = load_prior(os.path.join(PORT, "configs", "text_frames_prior.json"))
self.dur = G.DurationPrior()
self.gcfg = G.GuardConfig()
self.registry = VoiceRegistry()
load_named_voices(self.registry, os.path.join(models_dir, "voices.json"))
self._seed_rng = random.Random(SEED) # engine.py:224-225
for _ in range(WARMUP_PIECES):
self._seed_rng.randrange(1, 2 ** 31 - 1)
self.load_s = time.perf_counter() - t0
self.last_info = None
# ------------------------------------------------------------------ constants
def _anchor(self, w):
"""The speaker anchor of a row with no speaker embedding (anchors.py:86-109): fc2(gelu(fc1(null_vec))),
Linear = fp32 accumulation + bias, one rounding; GELU in fp32, one rounding. Model dtype."""
dt = self.bdt
x = w["null_vec"].astype(dt)[None]
h = (x.astype(mx.float32) @ w["fc1.weight"].astype(dt).astype(mx.float32).T
+ w["fc1.bias"].astype(dt).astype(mx.float32)).astype(dt)
hf = h.astype(mx.float32)
h = (hf * 0.5 * (1.0 + mx.erf(hf * math.sqrt(0.5)))).astype(dt)
a = (h.astype(mx.float32) @ w["fc2.weight"].astype(dt).astype(mx.float32).T
+ w["fc2.bias"].astype(dt).astype(mx.float32)).astype(dt)
mx.eval(a)
return a[0]
@property
def voices(self):
return self.registry.ids()
# ------------------------------------------------------------------ public
def generate(self, text, lang="en", voice=None, max_frames=None, return_info=False, **kw):
"""Synthesise text -> float32 waveform (24 kHz), as Maya3TTS.generate (api.py:224-230).
kw: normalize, seed, temperature, top_k, top_p, repetition_penalty. (The official positional arguments
ref_audio, ref_text, ref_lang and mode belong to cloning, which is not ported; pass the rest by keyword.)"""
info = {}
chunks = list(self.stream(text, lang, voice, max_frames=max_frames, info=info, **kw))
wav = (np.concatenate(chunks).astype(np.float32) / 32767.0) if chunks else np.zeros(0, np.float32)
return (wav, info) if return_info else wav
def stream(self, text, lang="en", voice=None, max_frames=None, normalize=False, info=None, seed=None,
temperature=None, top_k=None, top_p=None, repetition_penalty=None):
"""Yield int16 PCM chunks (24 kHz) as they are generated (api.py:196-222); `info` (a dict) receives the end
payload. Errors while preparing a piece raise GenerationError("invalid_request", ...) as the engine reports
them (engine.py:462-468, api.py:219-220)."""
v = self.registry.get(voice or "plain") # api.py:203 (before the sampling check)
if normalize: # api.py:188-194
from maya3.text import normalise_text
text = [normalise_text(t, lang) for t in text] if isinstance(text, (list, tuple)) \
else normalise_text(text, lang)
sampling = check_sampling({"seed": seed, "temperature": temperature, "top_k": top_k, "top_p": top_p,
"repetition_penalty": repetition_penalty}) # engine.py:293
if not isinstance(text, (list, tuple)): # engine.py:296-299
if not isinstance(text, str) or not text.strip():
raise ValueError("text must be a non-empty string")
text = text.strip()
pieces = split_text(text, lang, self.prior, CHUNK_S) if CHUNK_S > 0 else [text]
else:
pieces = list(text)
req = {"frames": 0, "stops": [], "t0": time.perf_counter(), "t_first": None}
k = 0
while True:
try:
plan = self._prepare(pieces[k], lang, v, max_frames, req)
except Exception as exc: # engine.py:464-468
raise GenerationError("invalid_request", f"{type(exc).__name__}: {exc}") from exc
for pcm in self._piece(plan, k, sampling, max_frames, req):
yield pcm
k += 1
if k < len(pieces) and not (max_frames and req["frames"] >= max_frames): # engine.py:641-645
yield np.zeros(int(GAP_S * SR), np.int16)
else:
break
end = {"reason": req["stops"][-1]["reason"], "frames": req["frames"], "audio_s": req["frames"] * FRAME / SR,
"ttfa_s": (req["t_first"] - req["t0"]) if req["t_first"] else None,
"wall_s": time.perf_counter() - req["t0"], "pieces": len(pieces), "stops": req["stops"]}
if info is not None:
info.update(end)
self.last_info = end
# ------------------------------------------------------------------ one piece
def _plan(self, text, lang, max_frames, req):
"""engine.py:394-418 + the capacity rule of engine.py:469-476."""
g = self.gcfg
ef, er = expected_frames(self.prior, lang, len(text))
loose = min(HARD_CAP, int(2.0 * (ef + er) + 40))
cap = loose
if g.enabled:
cap = min(cap, max(int(ef * 0.5) + 8, self.dur.cap_frames(text, lang, g.dur_k, g.dur_margin_s)))
if max_frames:
cap = min(cap, max(1, int(max_frames) - req["frames"]))
minimum = max(4, int(0.5 * ef))
if g.enabled and g.loop and g.aware_v2:
suspend = self.G.suspend_windows_v2(text, lang, self.dur)
else:
suspend = self.G.tag_windows(text, lang, self.dur, g.tag_hold_s) \
if (g.enabled and g.loop and g.tag_aware and "[" in text) else ()
return cap, loose, minimum, suspend
def _prompt(self, voice, text, lang):
"""Token ids and the composed prompt embeddings (s2_serve_controls.py:197-223, tts_model_v2.py:128-170):
text embeddings, plus the null speaker anchor at the <|audio|> sentinel (the last position)."""
ids = self.packer.prompt_ids(voice.render(text), lang)
emb = self.backbone.embed_tokens(ids)
last = emb[-1] + self.anchor # index_put accumulate: one add
emb = mx.concatenate([emb[:-1], last[None]], 0)
return ids, emb
def _cb0(self, logits, hist, hz, can, sp, aeos_bias, noise):
"""Engine._cb0 (engine.py:652-694) for one seat. logits fp32 [2049]; hist int [12] (-1 empty);
hz fp32 scalar; can bool. Returns (stop code, cb0, new hz) as mx arrays."""
V = logits.shape[0]
idx = mx.arange(V)
temp = float(sp.get("temperature") or CB0_TEMP)
top_k = int(sp.get("top_k") or CB0_TOP_K)
top_p = float(sp.get("top_p") or CB0_TOP_P)
rep = float(sp["repetition_penalty"] if sp.get("repetition_penalty") is not None else CB0_REP)
logits = logits / mx.array(temp, mx.float32)
logits = logits + mx.where(idx == AEOS, mx.array(aeos_bias, mx.float32), mx.array(0.0, mx.float32))
if CB0_REP_WIN > 0:
h = mx.array(hist, mx.int32)
pen = mx.zeros((V,), mx.float32).at[mx.maximum(h, 0)].add((h >= 0).astype(mx.float32))
pen = mx.where(idx == AEOS, 0.0, pen)
logits = logits - mx.array(rep, mx.float32) * (pen > 0).astype(mx.float32)
srt = mx.sort(logits) # ascending: desc index i = V-1-i
kth = srt[V - 1 - min(max(top_k - 1, 0), V - 1)]
logits = mx.where(logits < kth, -mx.inf, logits)
if top_p < 1.0:
p = mx.softmax(logits, precise=True)
si = mx.argsort(-p)
sp_ = p[si]
drop_sorted = (mx.cumsum(sp_) - sp_) > top_p
drop = mx.zeros((V,), mx.int32).at[si].add(drop_sorted.astype(mx.int32)) > 0 # scatter back by si
logits = mx.where(drop, -mx.inf, logits)
probs = mx.softmax(logits, precise=True)
pa = probs[AEOS]
arg = mx.argmax(probs)
hzv = mx.array(hz, mx.float32)
if can:
hzv = 1 - (1 - hzv) * (1 - pa)
code = mx.array(0)
if can:
code = mx.where(hzv >= 0.6, 4, code)
code = mx.where(pa >= 0.35, 3, code)
code = mx.where(arg == AEOS, 2, code)
q = mx.where(idx == AEOS, 0.0, probs)
q = q / mx.maximum(q.sum(), 1e-12)
cb0 = mx.minimum(mx.argmax(q / noise), AEOS - 1).astype(mx.int32)
return code, cb0, hzv
def _prepare(self, text, lang, voice, max_frames, req):
"""_plan (engine.py:400-418: the prompt is built first) and the capacity rule of _admit (:469-476)."""
ids, emb = self._prompt(voice, text, lang)
cap, loose, minimum, suspend = self._plan(text, lang, max_frames, req)
n = len(ids)
if n + cap + 2 > CAPACITY:
cap = CAPACITY - n - 2
if cap < 8:
raise ValueError(f"prompt of {n} tokens leaves no room in KV capacity {CAPACITY}")
return emb, cap, loose, minimum, suspend
def _piece(self, plan, k, sp, max_frames, req):
g = self.gcfg
emb, cap, loose, minimum, suspend = plan
if "seed" in sp and sp["seed"] is not None: # engine.py:430-437
seed = (int(sp["seed"]) * 1000003 + k * 7919) % (2 ** 31 - 1)
else:
seed = self._seed_rng.randrange(1, 2 ** 31 - 1)
row_seed = (seed * 69069 + 1) % (2 ** 31 - 1)
self.backbone.reset()
self.mimi.reset()
hist = np.full(CB0_REP_WIN, -1, np.int64) # cb0 presence-penalty window
ghist = np.full(g.hist, -1, np.int64) # loop-guard cb0 tail (_Seat.hist)
hz = np.float32(0.0)
aeos_bias = 0.0
frames, boost_at, n_boost, ctr = 0, None, 0, 0
h = self.backbone.forward(emb)[-1]
reason = None
while reason is None:
noise = mx.array(philox.cb0_noise(seed, ctr))
uni = mx.array(philox.depth_uniform(row_seed, ctr))
code, cb0, hz_new = self._cb0(self.backbone.cb0_logits(h), hist, hz, frames >= minimum, sp, aeos_bias,
noise)
codes, fb = self.depth.frame(h, cb0, uni)
pcm = self.mimi.decode_frame(codes)
mx.eval(code, cb0, hz_new, codes, fb, pcm)
ctr += 1
c0 = int(cb0.item())
hist = np.concatenate([hist[1:], [c0]])
hz = np.float32(hz_new.item())
stop = int(code.item())
if stop: # engine.py:576-578: not emitted
reason = REASONS[stop]
break
frames += 1
req["frames"] += 1
if req["t_first"] is None:
req["t_first"] = time.perf_counter()
pcm16 = np.rint(np.clip(np.array(pcm, np.float32), -1, 1) * np.float32(32767)).astype(np.int16)
yield pcm16
ghist[:-1] = ghist[1:]
ghist[-1] = c0
if frames >= cap: # engine.py:586-593
if max_frames and req["frames"] >= max_frames:
reason = "max_frames"
elif g.enabled and cap < loose:
reason = "cap_guard"
else:
reason = "cap"
break
if g.enabled and g.loop and frames >= g.start and frames % g.every == 0 and \
not any(a <= frames <= b for a, b in suspend): # engine.py:594-617
hit = self.G.loop_detect_batch(ghist[None], np.array([frames]), g)[0]
if hit is not None and g.mode == "boost":
if boost_at is None:
boost_at = frames
n_boost += 1
aeos_bias = g.boost
if frames - boost_at >= g.backstop:
reason = "loop_backstop"
break
elif hit is not None:
reason = "loop_" + hit.split(":")[0]
break
elif boost_at is not None:
aeos_bias = 0.0
boost_at = None
h = self.backbone.forward(fb[None])[0]
# engine.py:639, plus the piece's seed (port addition: the official stream is not reported)
req["stops"].append({"reason": reason, "frames": frames, "cap": cap, "n_boost": n_boost, "seed": seed})
|