| """Mach Boom runtime: speech straight from the coded weights on a CPU (numpy and a small C library). |
| |
| from mach import MachBoom |
| tts = MachBoom("mach-boom.bin", voices_dir="voices", tokenizer="tokenizer.json") |
| wav = tts.synth("Your booking reference is AB-4471-Z.", "Zoey") |
| """ |
| import json, math, mmap, os, sys, threading |
| import numpy as np |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, HERE) |
| import ctypes |
| from packfmt import K, ptr, tlut10, Chunk, Lin, LMS, FlowS |
|
|
| MAGIC = b"PHMACH01" |
| AMX_DEFAULT = "1" |
|
|
|
|
| def read_safetensor(path, name): |
| raw = open(path, "rb").read() |
| n = int.from_bytes(raw[:8], "little") |
| e = json.loads(raw[8:8 + n])[name] |
| dt = {"F32": "<f4", "F16": "<f2", "BF16": None}[e["dtype"]] |
| o0, o1 = e["data_offsets"] |
| return np.frombuffer(raw[8 + n + o0:8 + n + o1], dt).reshape(e["shape"]) |
|
|
|
|
| def _ln(x, w, b, eps): |
| mu = x.mean(-1, keepdims=True) |
| var = ((x - mu) ** 2).mean(-1, keepdims=True) |
| y = (x - mu) / np.sqrt(var + eps) |
| return y * w + b if w is not None else y |
|
|
|
|
| def _until_none(q): |
| while (c := q.get()) is not None: |
| yield c |
|
|
|
|
| class Packed: |
| def __init__(self, rt, n, m, K=1, buf=None): |
| self.lib, self.n, self.m, self.K = rt.k.lib, n, m, K |
| self.n32, self.m32 = -(-n // 32) * 32, -(-m // 32) * 32 |
| self.nk = K * self.n32 |
| self.nkb = self.nk // 32 |
| size = self.m32 // 16 * self.nkb * 512 |
| self.B = (rt.scr(buf, size // 2) if buf else np.empty(size // 2, np.float32)).view(np.uint16) |
|
|
| def put_t(self, k, Wt, j0=0): |
| self.lib.tmk_bf16_pack_k(ptr(Wt), Wt.strides[0] // 4, self.n, Wt.shape[1], self.B.ctypes.data + j0 * self.nkb * 64, |
| k * self.n32 // 32, self.nkb) |
|
|
| def put_coded(self, k, rt, step, tl, c, j0=0): |
| codes, su, sv, ws, m, n = c |
| self.lib.tmk_wpack(ptr(codes), step, m, n, ptr(tl), ws, ptr(su), ptr(sv), ptr(rt.scr("s1", m * n)), |
| ptr(rt.scr("s2", m * n)), self.B.ctypes.data + j0 * self.nkb * 64, k * self.n32 // 32, self.nkb) |
|
|
| def put(self, k, W, j0=0): |
| self.lib.tmk_bf16_pack_j(ptr(W), W.strides[0] // 4, self.n, W.shape[0], self.B.ctypes.data + j0 * self.nkb * 64, |
| k * self.n32 // 32, self.nkb) |
|
|
|
|
| class Coded: |
| def __init__(self, rt, ent): |
| self.rt, self.K, self.kind, self.convT = rt, ent["K"], ent["kind"], ent["convT"] |
| self.step = 2 * self.K |
| self.tl = rt.tlut[self.K] |
| self.chunks = [(rt.blob(c["codes"]), rt.blob(c["su"]), rt.blob(c["sv"]), float(c["ws"]), c["m"], c["n"]) |
| for c in ent["chunks"]] |
| self.n = self.chunks[0][5] |
| self.m = sum(c[4] for c in self.chunks) if self.kind == "rows" else self.chunks[0][4] |
|
|
| def vec(self, x, out=None): |
| y = np.empty(self.m, np.float32) if out is None else out |
| x = np.ascontiguousarray(x, np.float32) |
| lib, tmp, r0 = self.rt.k.lib, self.rt.scr("tmp", 4096), 0 |
| for codes, su, sv, ws, m, n in self.chunks: |
| lib.tmk_linear_vec(ptr(codes), self.step, m, n, ptr(self.tl), ptr(su), ptr(sv), ws, ptr(x), |
| y.ctypes.data + 4 * r0, ptr(tmp)) |
| r0 += m |
| return y |
|
|
| def _chunk_dense(self, c, out=None): |
| codes, su, sv, ws, m, n = c |
| lib = self.rt.k.lib |
| A = self.rt.scr("s1", m * n).reshape(m, n) |
| lib.tmk_qdense(ptr(codes), self.step, m, n, ptr(self.tl), ws, ptr(A)) |
| lib.tmk_rot(ptr(A), m, n, None, ptr(su), 1.0 / math.sqrt(n)) |
| B = np.empty((n, m), np.float32) if out is None else out |
| B[...] = A.T |
| lib.tmk_rot(ptr(B), n, m, None, ptr(sv), 1.0 / math.sqrt(m)) |
| return B |
|
|
| def _chunk_mat(self, c, X): |
| codes, su, sv, ws, m, n = c |
| lib = self.rt.k.lib |
| T = X.shape[0] |
| lm, ln_ = m.bit_length() - 1, n.bit_length() - 1 |
| if T * (n * ln_ + m * lm) > m * n * (lm + ln_): |
| return X @ self._chunk_dense(c, self.rt.scr("s2", m * n).reshape(n, m)) |
| Xr = np.array(X, np.float32, order="C") |
| lib.tmk_rot(ptr(Xr), T, n, ptr(su), None, 1.0 / math.sqrt(n)) |
| Q = self.rt.scr("s1", m * n).reshape(m, n) |
| lib.tmk_qdense(ptr(codes), self.step, m, n, ptr(self.tl), ws, ptr(Q)) |
| Z = Xr @ Q.T |
| lib.tmk_rot(ptr(Z), T, m, None, ptr(sv), 1.0 / math.sqrt(m)) |
| return Z |
|
|
| def dense_t(self): |
| if self.kind == "taps": |
| return [self._chunk_dense(c) for c in self.chunks] |
| out, r0 = np.empty((self.n, self.m), np.float32), 0 |
| for c in self.chunks: |
| out[:, r0:r0 + c[4]] = self._chunk_dense(c, self.rt.scr("s2", c[4] * c[5]).reshape(c[5], c[4])); r0 += c[4] |
| return out |
|
|
| def dense(self): |
| return np.ascontiguousarray(self.dense_t().T) |
|
|
| def mat_amx(self, X, act=0): |
| rt = self.rt |
| lib, T = rt.k.lib, X.shape[0] |
| Y, r0 = np.empty((T, self.m), np.float32), 0 |
| for c in self.chunks: |
| codes, su, sv, ws, m, n = c |
| lm, ln_ = m.bit_length() - 1, n.bit_length() - 1 |
| P = Packed(rt, n, m, 1, "bp") |
| if T * (n * ln_ + m * lm) > m * n * (lm + ln_): |
| P.put_coded(0, rt, self.step, self.tl, c) |
| rt.amx_mm(X, P, out=Y[:, r0:r0 + m], act=act) |
| else: |
| Xr = np.array(X, np.float32, order="C") |
| if act: |
| rt.k.lib.tmk_elu(ptr(Xr), Xr.size) |
| lib.tmk_rot(ptr(Xr), T, n, ptr(su), None, 1.0 / math.sqrt(n)) |
| rb = max(32, (1 << 18) // n // 32 * 32) |
| Q = rt.scr("s1", min(rb, m) * n) |
| for j in range(0, m, rb): |
| mb = min(rb, m - j) |
| lib.tmk_qdense(ptr(codes[j * n // 16 * self.step:]), self.step, mb, n, ptr(self.tl), ws, ptr(Q)) |
| P.put(0, Q[:mb * n].reshape(mb, n), j0=j) |
| Z = rt.amx_mm(Xr, P, out=Y if len(self.chunks) == 1 else None) |
| lib.tmk_rot(ptr(Z), T, m, None, ptr(sv), 1.0 / math.sqrt(m)) |
| if Z is not Y: |
| Y[:, r0:r0 + m] = Z |
| r0 += m |
| return Y |
|
|
| def packed_tap(self, c): |
| P = Packed(self.rt, self.n, self.m, 1, "bp") |
| P.put_coded(0, self.rt, self.step, self.tl, c) |
| return P |
|
|
| def taps_amx(self, x, T, dil=1, act=0): |
| y = np.empty((T, self.m), np.float32) |
| for k, c in enumerate(self.chunks): |
| self.rt.amx_mm(x[k * dil:], self.packed_tap(c), T=T, out=y, acc=k > 0, act=act) |
| return y |
|
|
| def packed_taps(self, buf=None): |
| n, m = self.chunks[0][5], self.chunks[0][4] |
| assert self.kind == "taps" or len(self.chunks) == 1 |
| P = Packed(self.rt, n, m, len(self.chunks), buf) |
| for k, c in enumerate(self.chunks): |
| P.put_coded(k, self.rt, self.step, self.tl, c) |
| return P |
|
|
| def packed_convT(self, buf=None): |
| n, m, K = self.chunks[0][5], self.chunks[0][4], len(self.chunks) |
| assert m % 32 == 0 |
| P = Packed(self.rt, n, K * m, 1, buf) |
| for k, c in enumerate(self.chunks): |
| P.put_coded(0, self.rt, self.step, self.tl, c, j0=k * m) |
| return P |
|
|
| def mat(self, X, amx=False, act=0): |
| if amx: |
| return self.mat_amx(X, act) |
| return np.concatenate([self._chunk_mat(c, X) for c in self.chunks], 1) if len(self.chunks) > 1 \ |
| else self._chunk_mat(self.chunks[0], X) |
|
|
|
|
| def _one_arena(): |
| n = int(os.environ.get("PHM_ARENA_MAX", "1")) |
| if n > 0 and sys.platform.startswith("linux"): |
| try: |
| ctypes.CDLL(None).mallopt(-8, n) |
| except (OSError, AttributeError): |
| pass |
|
|
|
|
| class MachBoom: |
| def __init__(self, path, voices_dir=None, tokenizer=None, threads=None, so=None, lm_threads=1, amx=None): |
| _one_arena() |
| with open(path, "rb") as f: |
| raw = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) |
| assert raw[:8] == MAGIC, "not a Mach Boom file" |
| hl = int.from_bytes(raw[8:16], "little") |
| H = json.loads(raw[16:16 + hl]) |
| base = 16 + hl + ((-(16 + hl)) % 64) |
| self.buf = np.frombuffer(raw, np.uint8) |
| self.base, self.cfg, self.T = base, H["config"], H["tensors"] |
| self.k = K(so) |
| if amx is None: |
| amx = os.environ.get("PHM_AMX", AMX_DEFAULT) != "0" |
| self.amx = bool(amx) and bool(self.k.lib.tmk_amx_ok()) |
| self.lm_threads = int(lm_threads) |
| self.k.lib.tmk_threads(self.lm_threads) |
| self.tlut = {int(k): tlut10(self.blob(v).view(np.float32)) for k, v in H["tlut"].items()} |
| self._tls = threading.local() |
| self._coded, self._small = {}, {} |
| c = self.cfg["lm"] |
| self.D, self.NL, self.NH, self.HD = c["hidden"], c["layers"], c["heads"], c["head_dim"] |
| self.inv_freq = (1.0 / (c["rope_theta"] ** (np.arange(0, self.HD, 2) / self.HD))).astype(np.float32) |
| d = self.cfg["dec"] |
| self.dinv_freq = (1.0 / (d["rope_theta"] ** (np.arange(0, d["head_dim"], 2) / d["head_dim"]))).astype(np.float32) |
| self.voices_dir = voices_dir |
| self._voice_cache = {} |
| self.tok = None |
| if tokenizer: |
| from tokenizers import Tokenizer |
| self.tok = Tokenizer.from_file(tokenizer) |
|
|
| def scr(self, which, n): |
| d = self._tls.__dict__ |
| a = d.get(which) |
| if a is None or a.size < n: |
| a = d[which] = np.empty(n, np.float32) |
| return a[:n] |
|
|
| def amx_mm(self, x, P, T=None, dil=1, out=None, acc=False, act=0): |
| x = np.ascontiguousarray(x, np.float32) |
| T = x.shape[0] - (P.K - 1) * dil if T is None else T |
| T32 = -(-T // 32) * 32 |
| Xb = self.scr("xb", T32 * P.nk // 2).view(np.uint16) |
| lib = self.k.lib |
| lib.tmk_bf16_x(ptr(x), x.shape[1], T, P.n, P.K, dil, ptr(Xb), int(act)) |
| y = np.empty((T, P.m), np.float32) if out is None else out |
| assert y.strides[1] == 4 and y.shape == (T, P.m) |
| lib.tmk_amx_gemm(ptr(Xb), T, P.nk, ptr(P.B), P.m, y.ctypes.data, y.strides[0] // 4, int(acc)) |
| return y |
|
|
| def blob(self, ol): |
| o, n = ol |
| return self.buf[self.base + o:self.base + o + n] |
|
|
| def f32(self, name): |
| e = self.T[name] |
| return self.blob(e["data"]).view(np.float32).reshape(e["shape"]) |
|
|
| def coded(self, name): |
| if name not in self._coded: |
| self._coded[name] = Coded(self, self.T[name]) |
| return self._coded[name] |
|
|
| def small(self, name): |
| if name not in self._small: |
| e = self.T[name] |
| q = self.blob(e["q"]).view(np.int8).reshape(e["rows"], e["cols"]).astype(np.float32) |
| s = self.blob(e["s"]).view(np.float16).astype(np.float32).reshape(e["rows"], -1) |
| self._small[name] = (q.reshape(e["rows"], -1, e["g"]) * s[..., None]).reshape(e["rows"], e["cols"]) |
| return self._small[name] |
|
|
| def embed_row(self, i): |
| e = self.T["lm:model.embed_tokens"] |
| n, g = e["cols"], e["g"] |
| b = self.blob(e["q"])[i * n // 2:(i + 1) * n // 2] |
| q = np.empty(n, np.int8); q[0::2] = (b & 0xF).astype(np.int8); q[1::2] = (b >> 4).astype(np.int8) |
| q = np.where(q > 7, q - 16, q).astype(np.float32) |
| s = self.blob(e["s"]).view(np.float16)[i * n // g:(i + 1) * n // g].astype(np.float32) |
| return (q.reshape(-1, g) * s[:, None]).reshape(n) |
|
|
| def lin(self, name, x, amx=False): |
| root = name |
| if root in self.T and self.T[root]["type"] == "trellis": |
| c = self.coded(root) |
| if x.ndim == 1: |
| y = c.vec(x) |
| elif x.shape[0] <= 2: |
| y = np.stack([c.vec(r) for r in x]) |
| else: |
| y = c.mat(x, amx) |
| elif root in self.T: |
| y = x @ self.small(root).T |
| else: |
| y = x @ self.f32(root + ".weight").reshape(self.T[root + ".weight"]["shape"][0], -1).T |
| bn = root + ".bias" |
| return y + self.f32(bn) if bn in self.T else y |
|
|
| def _clin(self, name): |
| L = Lin() |
| b = name + ".bias" |
| L.bias = ptr(self.f32(b)) if b in self.T else None |
| e = self.T[name] |
| if e["type"] == "trellis": |
| c = self.coded(name) |
| arr = (Chunk * len(c.chunks))() |
| for i, (codes, su, sv, ws, m, n) in enumerate(c.chunks): |
| arr[i] = Chunk(ptr(codes), ptr(su), ptr(sv), ws, m, n) |
| self._keep.append(arr) |
| L.ch, L.nch, L.step, L.tlut, L.m, L.n = arr, len(c.chunks), c.step, ptr(c.tl), c.m, c.n |
| else: |
| W = np.ascontiguousarray(self.small(name), np.float32) |
| self._keep.append(W) |
| L.nch, L.dense, L.m, L.n = 0, ptr(W), W.shape[0], W.shape[1] |
| return L |
|
|
| def _cmodel(self): |
| if getattr(self, "_cm", None) is not None: |
| return self._cm |
| self._keep = [] |
| c = self.cfg["lm"] |
| lm = LMS() |
| lm.D, lm.NL, lm.NH, lm.HD, lm.FF, lm.eps = self.D, self.NL, self.NH, self.HD, c["ffn"], c["ln_eps"] |
| lm.inv_freq = ptr(self.inv_freq) |
| for l in range(self.NL): |
| pre = f"lm:model.layers.{l}." |
| for i, nm in enumerate(("input_layernorm.weight", "input_layernorm.bias", "post_attention_layernorm.weight", |
| "post_attention_layernorm.bias")): |
| lm.ln[l][i] = ptr(self.f32(pre + nm)) |
| for i, nm in enumerate(("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", |
| "mlp.fc1", "mlp.fc2")): |
| lm.lin[l][i] = self._clin(pre + nm) |
| lm.lnf_w, lm.lnf_b = ptr(self.f32("lm:model.norm.weight")), ptr(self.f32("lm:model.norm.bias")) |
| lm.in_proj, lm.in_add = self._clin("lm:model.audio_latent_proj"), ptr(self.f32("const:cond")) |
| fl = FlowS() |
| fl.D, fl.NL, fl.LD, fl.eps = self.D, c["flow_layers"], c["latent_dim"], c["flow_eps"] |
| fl.inp, fl.cond = self._clin("lm:flow_head.input_proj"), self._clin("lm:flow_head.cond_proj") |
| for l in range(c["flow_layers"]): |
| pre = f"lm:flow_head.layers.{l}." |
| fl.ada[l], fl.fc1[l], fl.fc2[l] = self._clin(pre + "adaln_proj"), self._clin(pre + "fc1"), self._clin(pre + "fc2") |
| fl.norm[l][0], fl.norm[l][1] = ptr(self.f32(pre + "norm.weight")), ptr(self.f32(pre + "norm.bias")) |
| fl.fada, fl.fproj = self._clin("lm:flow_head.final_layer.adaln_proj"), self._clin("lm:flow_head.final_layer.proj") |
| fl.time = ptr(self.f32("const:flow_time")) |
| self._fwork = np.empty(8 * self.D + 4096, np.float32) |
| self._cm = (lm, fl) |
| return self._cm |
|
|
| def flow_c(self, h, noise): |
| lm, fl = self._cmodel() |
| h = np.ascontiguousarray(h, np.float32); noise = np.ascontiguousarray(noise, np.float32) |
| out = np.empty(self.cfg["lm"]["latent_dim"], np.float32) |
| self.k.lib.tmk_flow_step(ctypes.byref(fl), ptr(h), ptr(noise), ptr(out), ptr(self._fwork)) |
| return out |
|
|
| def _rope(self, x, pos, inv): |
| f = pos[:, None].astype(np.float32) * inv[None] |
| cos, sin = np.cos(np.concatenate([f, f], -1))[:, None], np.sin(np.concatenate([f, f], -1))[:, None] |
| h = x.shape[-1] // 2 |
| rot = np.concatenate([-x[..., h:], x[..., :h]], -1) |
| return x * cos + rot * sin |
|
|
| def _layer(self, l, x, pos, cache, P0): |
| c = self.cfg["lm"] |
| pre = f"lm:model.layers.{l}." |
| h = _ln(x, self.f32(pre + "input_layernorm.weight"), self.f32(pre + "input_layernorm.bias"), c["ln_eps"]) |
| T = x.shape[0] |
| q = self.lin(pre + "self_attn.q_proj", h).reshape(T, self.NH, self.HD) |
| k = self.lin(pre + "self_attn.k_proj", h).reshape(T, self.NH, self.HD) |
| v = self.lin(pre + "self_attn.v_proj", h).reshape(T, self.NH, self.HD) |
| q, k = self._rope(q, pos, self.inv_freq), self._rope(k, pos, self.inv_freq) |
| Kc, Vc = cache |
| Kc[:, P0:P0 + T] = k.transpose(1, 0, 2); Vc[:, P0:P0 + T] = v.transpose(1, 0, 2) |
| Kh, Vh = Kc[:, :P0 + T], Vc[:, :P0 + T] |
| s = np.matmul(q.transpose(1, 0, 2), Kh.transpose(0, 2, 1)) * np.float32(1 / math.sqrt(self.HD)) |
| if T > 1: |
| s = s + np.where(np.arange(P0 + T)[None, :] > pos[:, None], -np.inf, 0.0).astype(np.float32)[None] |
| s = np.exp(s - s.max(-1, keepdims=True)); s /= s.sum(-1, keepdims=True) |
| a = np.matmul(s, Vh).transpose(1, 0, 2).reshape(T, -1) |
| x = x + self.lin(pre + "self_attn.o_proj", np.ascontiguousarray(a, np.float32)) |
| h = _ln(x, self.f32(pre + "post_attention_layernorm.weight"), self.f32(pre + "post_attention_layernorm.bias"), c["ln_eps"]) |
| u = np.ascontiguousarray(self.lin(pre + "mlp.fc1", h), np.float32) |
| self.k.lib.tmk_gelu(ptr(u), u.size) |
| return x + self.lin(pre + "mlp.fc2", u) |
|
|
| def backbone(self, x, P0, caches): |
| T = x.shape[0] |
| pos = np.arange(P0, P0 + T) |
| for l in range(self.NL): |
| x = self._layer(l, x, pos, caches[l], P0) |
| c = self.cfg["lm"] |
| return _ln(x, self.f32("lm:model.norm.weight"), self.f32("lm:model.norm.bias"), c["ln_eps"]) |
|
|
| def flow(self, h, noise): |
| c = self.cfg["lm"] |
| x = self.lin("lm:flow_head.input_proj", noise) |
| cond = self.lin("lm:flow_head.cond_proj", h) + self.f32("const:flow_time") |
| sc = np.ascontiguousarray(cond, np.float32).copy() |
| self.k.lib.tmk_silu(ptr(sc), sc.size) |
| D = self.D |
| for l in range(c["flow_layers"]): |
| pre = f"lm:flow_head.layers.{l}." |
| ada = self.lin(pre + "adaln_proj", sc) |
| shift, scale, gate = ada[:D], ada[D:2 * D], ada[2 * D:] |
| m = _ln(x, self.f32(pre + "norm.weight"), self.f32(pre + "norm.bias"), c["flow_eps"]) * (1 + scale) + shift |
| u = np.ascontiguousarray(self.lin(pre + "fc1", m), np.float32) |
| self.k.lib.tmk_silu(ptr(u), u.size) |
| x = x + gate * self.lin(pre + "fc2", u) |
| ada = self.lin("lm:flow_head.final_layer.adaln_proj", sc) |
| m = _ln(x, None, None, c["flow_eps"]) * (1 + ada[D:]) + ada[:D] |
| return noise + self.lin("lm:flow_head.final_layer.proj", m) |
|
|
| def voices(self): |
| for p in (os.path.join(self.voices_dir, "voices.jsonl"), os.path.join(self.voices_dir, "..", "voices.jsonl")): |
| if os.path.exists(p): |
| return {r["name"]: r["voice_id"] for r in map(json.loads, open(p, encoding="utf-8"))} |
| return {} |
|
|
| def voice_latents(self, voice): |
| if isinstance(voice, np.ndarray): |
| return voice |
| path = os.path.join(self.voices_dir, f"{voice}.safetensors") |
| if not os.path.exists(path): |
| ids = {k.lower(): v for k, v in self.voices().items()} |
| if voice.lower() not in ids: |
| raise KeyError(f"no voice {voice!r} in {self.voices_dir}") |
| path = os.path.join(self.voices_dir, f"{ids[voice.lower()]}.safetensors") |
| return read_safetensor(path, "speaker_wavs")[0].T.astype(np.float32) |
|
|
| def _alloc(self, P): |
| return [(np.zeros((self.NH, P, self.HD), np.float32), np.zeros((self.NH, P, self.HD), np.float32)) |
| for _ in range(self.NL)] |
|
|
| def _voice_prefix(self, voice): |
| key = voice if isinstance(voice, str) else id(voice) |
| if key not in self._voice_cache: |
| lat = self.voice_latents(voice) |
| caches = self._alloc(lat.shape[0]) |
| self.backbone(self.lin("lm:model.audio_prompt_proj", lat), 0, caches) |
| self._voice_cache[key] = caches |
| return self._voice_cache[key] |
|
|
| def generate_latents(self, ids, voice, seed=0, temperature=None, num_words=None, max_frames=None, forced=None, |
| trace=None, on_chunk=None, chunk=16): |
| out = [] |
| for c in self.latent_chunks(ids, voice, seed=seed, temperature=temperature, num_words=num_words, |
| max_frames=max_frames, forced=forced, trace=trace, chunk=chunk): |
| out.append(c) |
| if on_chunk is not None: |
| on_chunk(c) |
| return np.concatenate(out) |
|
|
| def latent_chunks(self, ids, voice, seed=0, temperature=None, num_words=None, max_frames=None, forced=None, |
| trace=None, chunk=16, first=None): |
| g = self.cfg["gen"] |
| temperature = g["temperature"] if temperature is None else temperature |
| rng = np.random.default_rng(seed) |
| pre = self._voice_prefix(voice) |
| Pv = pre[0][0].shape[1] |
| nt = len(ids) |
| max_frames = max_frames or math.ceil((nt / 3 + 2) * self.cfg["lm"]["frame_rate"]) |
| P = Pv + nt + 1 + max_frames |
| caches = self._alloc(P) |
| for (K, V), (k0, v0) in zip(caches, pre): |
| K[:, :Pv] = k0; V[:, :Pv] = v0 |
| x = np.stack([self.embed_row(i) for i in ids] + [self.f32("const:bos")]).astype(np.float32) |
| h = self.backbone(x, Pv, caches)[-1] |
| self._tls.__dict__.pop("s1", None) |
| fae = g["frames_after_eos_short_text"] if (num_words is not None and num_words <= g["short_text_max_words"]) \ |
| else g["frames_after_eos"] |
| lm, fl = self._cmodel() |
| lm.Pmax = P |
| lib = self.k.lib |
| Kp = (ctypes.c_void_p * self.NL)(*[k.ctypes.data for k, _ in caches]) |
| Vp = (ctypes.c_void_p * self.NL)(*[v.ctypes.data for _, v in caches]) |
| work = np.empty(lib.tmk_lm_work(ctypes.byref(lm)), np.float32) |
| h = np.ascontiguousarray(h, np.float32) |
| ew, eb = self.f32("lm:eos_head.weight").reshape(-1), float(self.f32("lm:eos_head.bias")[0]) |
| LD = self.cfg["lm"]["latent_dim"] |
| sd = math.sqrt(temperature) |
| lats, remaining, p, sent = [], -1, Pv + nt + 1, 0 |
| std, mean = self.f32("lm:latent_std"), self.f32("lm:latent_mean") |
| for f in range(max_frames): |
| eos = float(h @ ew) + eb |
| noise = (rng.standard_normal(LD) * sd).astype(np.float32) |
| if forced is None: |
| lat = np.empty(LD, np.float32) |
| lib.tmk_flow_step(ctypes.byref(fl), ptr(h), ptr(noise), ptr(lat), ptr(self._fwork)) |
| else: |
| lat = np.ascontiguousarray(forced[f], np.float32) |
| if trace is not None: |
| trace.append((h.copy(), eos)) |
| lats.append(lat) |
| n = first if (first and sent == 0) else chunk |
| if len(lats) - sent >= n: |
| yield np.stack(lats[sent:sent + n]) * std + mean |
| sent += n |
| if forced is not None and f == len(forced) - 1: |
| break |
| if remaining < 0 and eos > g["eos_threshold"]: |
| remaining = fae |
| if remaining == 0: |
| break |
| if remaining > 0: |
| remaining -= 1 |
| if f == max_frames - 1: |
| break |
| hn = np.empty(self.D, np.float32) |
| lib.tmk_lm_step(ctypes.byref(lm), ptr(lat), None, p, Kp, Vp, ptr(hn), ptr(work)) |
| h = hn |
| p += 1 |
| if sent < len(lats): |
| yield np.stack(lats[sent:]) * std + mean |
|
|
| def decoder(self): |
| return MimiStream(self) |
|
|
| def decode(self, latents, chunk=16): |
| st, out = MimiStream(self), [] |
| for i in range(0, len(latents), chunk): |
| out.append(st.push(latents[i:i + chunk])) |
| return np.concatenate(out) |
|
|
| def text_ids(self, text): |
| enc = self.tok.encode(text.strip(), add_special_tokens=False) |
| return enc.ids, len(self.tok.decode(enc.ids, skip_special_tokens=True).split()) |
|
|
| def synth(self, text, voice, seed=0, frontend=True, temperature=None, threads=1, chunk=None): |
| return np.concatenate(list(self.synth_stream(text, voice, seed, frontend, temperature, threads, chunk))) |
|
|
| def synth_stream(self, text, voice, seed=0, frontend=True, temperature=None, threads=1, chunk=None, first=None): |
| if frontend: |
| try: |
| from frontend_en import normalize |
| except ImportError: |
| sys.path.append(os.path.dirname(HERE)) |
| from frontend_en import normalize |
| text = normalize(text) |
| ids, nw = self.text_ids(text) |
| chunk = chunk or (16 if threads <= 2 else 8) |
| lats = self.latent_chunks(ids, voice, seed=seed, num_words=nw, temperature=temperature, chunk=chunk, first=first) |
| st = MimiStream(self) |
| if threads <= 1: |
| self.k.lib.tmk_threads(self.lm_threads) |
| for c in lats: |
| yield st.push(c) |
| return |
| import queue |
| self.k.lib.tmk_threads(max(self.lm_threads, 2 if threads >= 4 else 1)) |
| stages = [lambda c: c] + {2: [st.push], 3: [st.push_a, st.push_b]}.get(threads, [st.push_t, st.push_s0, st.push_b]) |
| qs = [queue.Queue() for _ in stages] |
| err, stop = [], threading.Event() |
|
|
| def work(i): |
| try: |
| for c in (lats if i == 0 else _until_none(qs[i - 1])): |
| if stop.is_set() or err: |
| break |
| qs[i].put(stages[i](c)) |
| except BaseException as e: |
| err.append(e) |
| finally: |
| qs[i].put(None) |
| sw = sys.getswitchinterval() |
| sys.setswitchinterval(0.001) |
| ths = [threading.Thread(target=work, args=(i,), daemon=True) for i in range(len(stages))] |
| for t in ths: |
| t.start() |
| try: |
| yield from _until_none(qs[-1]) |
| if err: |
| raise err[0] |
| finally: |
| stop.set() |
| for t in ths: |
| t.join() |
| sys.setswitchinterval(sw) |
|
|
|
|
| class MimiStream: |
| def __init__(self, rt): |
| self.rt, d = rt, rt.cfg["dec"] |
| self.d = d |
| D = d["hidden"] |
| self.up_tail = np.zeros((d["latent_stride"], D), np.float32) |
| W = d["window"] |
| self.kv = [(np.zeros((d["heads"], 0, d["head_dim"]), np.float32),) * 2 for _ in range(d["layers"])] |
| self.pos = 0 |
| self.state, self.pd = {}, {} |
| self.piece = 32 |
|
|
| HI = ("5.conv", "6.block.1.conv", "6.block.3.conv", "8.conv", "9.block.1.conv", "9.block.3.conv", "11.conv") |
| CONVT = ("2.conv", "5.conv", "8.conv") |
| SEA = "dec:decoder.decoder_seanet.layers." |
|
|
| def _dense_push(self): |
| if self.rt.amx: |
| return self._dense_push_amx() |
| rt, pd = self.rt, {} |
| for nm in self.HI: |
| name = self.SEA + nm |
| e = rt.T.get(name) |
| if e is None: |
| w = rt.f32(name + ".weight") |
| pd[name] = [np.ascontiguousarray(w[:, :, k].T) for k in range(w.shape[2])] |
| continue |
| c = rt.coded(name) |
| if nm in self.CONVT: |
| K = e["shape"][2] |
| if c.kind == "rows": |
| W = c.dense_t().T.reshape(c.m, -1, K) |
| pd[name] = np.ascontiguousarray(W.transpose(0, 2, 1).reshape(c.m, -1)) |
| else: |
| pd[name] = np.concatenate(c.dense_t(), 1) |
| else: |
| pd[name] = c.dense_t() if c.kind == "taps" else [c.dense_t()] |
| self.pd = pd |
|
|
| def _dense_push_amx(self): |
| rt, pd = self.rt, {} |
| for nm in self.HI: |
| name = self.SEA + nm |
| e = rt.T.get(name) |
| if e is not None and e["kind"] == "taps" and nm in self.CONVT: |
| pd[name] = rt.coded(name).packed_convT("pk:" + name) |
| elif e is not None and nm not in self.CONVT: |
| pd[name] = rt.coded(name).packed_taps("pk:" + name) |
| else: |
| pd[name] = None |
| self.pd = pd |
| for nm in self.HI: |
| name = self.SEA + nm |
| if pd[name] is not None: |
| continue |
| e = rt.T.get(name) |
| if e is None: |
| w = rt.f32(name + ".weight") |
| taps = [np.ascontiguousarray(w[:, :, k].T) for k in range(w.shape[2])] |
| else: |
| c, K = rt.coded(name), e["shape"][2] |
| W = c.dense_t().T.reshape(c.m, -1, K) |
| taps = [np.ascontiguousarray(W.transpose(0, 2, 1).reshape(c.m, -1))] |
| n, m = taps[0].shape |
| if m < 16: |
| pd[name] = taps |
| continue |
| P = Packed(rt, n, m, len(taps), "pk:" + name) |
| for k, t in enumerate(taps): |
| P.put_t(k, t) |
| pd[name] = P |
|
|
| def _bias(self, name): |
| b = name + ".bias" |
| return self.rt.f32(b) if b in self.rt.T else None |
|
|
| def _taps_conv(self, name, x, K, dil=1, act=0): |
| rt = self.rt |
| if act and not rt.amx: |
| x = self._elu(x) |
| pad = (K - 1) * dil |
| T = x.shape[0] |
| if pad: |
| hist = self.state.get(name) |
| xb = np.concatenate([hist if hist is not None else np.zeros((pad, x.shape[1]), np.float32), x]) |
| self.state[name] = xb[T:].copy() |
| else: |
| xb = x |
| taps = self.pd.get(name) |
| if rt.amx and act and not isinstance(taps, Packed) and not (taps is None and name in rt.T): |
| xb = self._elu(xb) |
| act = 0 |
| if isinstance(taps, Packed): |
| y = rt.amx_mm(xb, taps, T=T, dil=dil, act=act) |
| elif taps is None and rt.amx and rt.T[name]["kind"] == "taps": |
| y = rt.coded(name).taps_amx(xb, T, dil, act=act) |
| elif taps is not None: |
| y = xb[:T] @ taps[0] |
| for k in range(1, K): |
| y += xb[k * dil:k * dil + T] @ taps[k] |
| else: |
| e = rt.T[name] |
| c = rt.coded(name) |
| if e["kind"] == "taps": |
| y = c._chunk_mat(c.chunks[0], xb[:T]) |
| for k in range(1, K): |
| y += c._chunk_mat(c.chunks[k], xb[k * dil:k * dil + T]) |
| else: |
| y = c.mat(x, rt.amx, act) |
| b = self._bias(name) |
| if b is not None: |
| y += b |
| return y |
|
|
| def _convT(self, name, x, K, s, act=0): |
| rt = self.rt |
| T = x.shape[0] |
| W = self.pd.get(name) |
| amx_taps = rt.amx and W is None and rt.coded(name).kind == "taps" and rt.coded(name).m % 32 == 0 |
| if act and not (isinstance(W, Packed) or amx_taps): |
| x = self._elu(x, inplace=True) |
| act = 0 |
| if amx_taps: |
| c = rt.coded(name) |
| out, Z2 = np.empty((T, s * c.m), np.float32), np.empty((T, s * c.m), np.float32) |
| for k, ch in enumerate(c.chunks): |
| j = k % s * c.m |
| rt.amx_mm(x, c.packed_tap(ch), out=(out, Z2)[k // s][:, j:j + c.m], act=act) |
| out += self._tbias(name, s) |
| else: |
| if isinstance(W, Packed): |
| Z = rt.amx_mm(x, W, act=act) |
| elif W is not None: |
| Z = x @ W |
| else: |
| c = rt.coded(name) |
| co = c.m |
| Z = np.empty((T, K * co), np.float32) |
| for k, ch in enumerate(c.chunks): |
| Z[:, k * co:(k + 1) * co] = c._chunk_mat(ch, x) |
| sc = Z.shape[1] // 2 |
| out, Z2 = Z[:, :sc] + self._tbias(name, s), Z[:, sc:] |
| out[1:] += Z2[:-1] |
| carry = self.state.get(name) |
| if carry is not None: |
| out[0] += carry |
| self.state[name] = Z2[-1].copy() |
| return out.reshape(T * s, -1) |
|
|
| def _tbias(self, name, s): |
| k = ("tb", name) |
| if k not in self.state: |
| self.state[k] = np.tile(self.rt.f32(name + ".bias"), s) |
| return self.state[k] |
|
|
| def _elu(self, x, inplace=False): |
| x = np.ascontiguousarray(x, np.float32) |
| if not inplace: |
| x = x.copy() |
| self.rt.k.lib.tmk_elu(ptr(x), x.size) |
| return x |
|
|
| def _dtf_args(self, l): |
| cache = self.rt.__dict__.setdefault("_dtf", {}) |
| if l not in cache: |
| rt = self.rt |
| if not hasattr(rt, "_keep"): |
| rt._keep = [] |
| k0 = len(rt._keep) |
| pre = f"dec:decoder.decoder_transformer.layers.{l}." |
| L = (Lin * 6)() |
| for i, nm in enumerate(("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", |
| "mlp.fc1", "mlp.fc2")): |
| L[i] = rt._clin(pre + nm) |
| f = [rt.f32(pre + nm) for nm in ("input_layernorm.weight", "input_layernorm.bias", |
| "post_attention_layernorm.weight", "post_attention_layernorm.bias", |
| "self_attn_layer_scale.scale", "mlp_layer_scale.scale")] |
| cache[l] = (L, f, rt._keep[k0:]) |
| return cache[l] |
|
|
| def _tf_layer_c(self, l, x, pos): |
| rt, d = self.rt, self.d |
| T, D, H, hd, FF = x.shape[0], d["hidden"], d["heads"], d["head_dim"], d["ffn"] |
| x = np.array(x, np.float32, order="C") |
| K0, V0 = self.kv[l] |
| P0 = K0.shape[1] |
| Kc = np.empty((H, P0 + T, hd), np.float32); Kc[:, :P0] = K0 |
| Vc = np.empty((H, P0 + T, hd), np.float32); Vc[:, :P0] = V0 |
| L, f, _ = self._dtf_args(l) |
| lib = rt.k.lib |
| work = rt.scr("dtf", lib.tmk_dtf_work(T, D, FF, d["window"])) |
| lib.tmk_dtf_layer(ptr(x), T, D, H, FF, d["eps"], ptr(f[0]), ptr(f[1]), ptr(f[2]), ptr(f[3]), ctypes.addressof(L), |
| ptr(f[4]), ptr(f[5]), ptr(Kc), ptr(Vc), P0, d["window"], ptr(rt.dinv_freq), int(pos[0]), ptr(work)) |
| keep = d["window"] - 1 |
| self.kv[l] = (Kc[:, -keep:].copy(), Vc[:, -keep:].copy()) |
| return x |
|
|
| def _attn(self, l, pre, h, pos): |
| rt, d = self.rt, self.d |
| T, H, hd = h.shape[0], d["heads"], d["head_dim"] |
| q = rt.lin(pre + "self_attn.q_proj", h, rt.amx).reshape(T, H, hd) |
| k = rt.lin(pre + "self_attn.k_proj", h, rt.amx).reshape(T, H, hd) |
| v = rt.lin(pre + "self_attn.v_proj", h, rt.amx).reshape(T, H, hd) |
| q, k = rt._rope(q, pos, rt.dinv_freq), rt._rope(k, pos, rt.dinv_freq) |
| K0, V0 = self.kv[l] |
| Kc = np.concatenate([K0, k.transpose(1, 0, 2)], 1); Vc = np.concatenate([V0, v.transpose(1, 0, 2)], 1) |
| P = Kc.shape[1] |
| qt, Kt, a = q.transpose(1, 0, 2), Kc.transpose(0, 2, 1), np.empty((T, H, hd), np.float32) |
| g = max(1, (1 << 18) // (T * P)) |
| for h0 in range(0, H, g): |
| s = np.matmul(qt[h0:h0 + g], Kt[h0:h0 + g]) |
| s *= np.float32(1 / math.sqrt(hd)) |
| rt.k.lib.tmk_softmax_band(ptr(s), s.shape[0] * T, P, T, P - T, d["window"]) |
| a[:, h0:h0 + g] = np.matmul(s, Vc[h0:h0 + g]).transpose(1, 0, 2) |
| keep = d["window"] - 1 |
| self.kv[l] = (Kc[:, -keep:].copy(), Vc[:, -keep:].copy()) |
| return a.reshape(T, -1) |
|
|
| def _tf_layer(self, l, x, pos): |
| if self.rt.amx: |
| return self._tf_layer_c(l, x, pos) |
| rt, d = self.rt, self.d |
| pre = f"dec:decoder.decoder_transformer.layers.{l}." |
| h = _ln(x, rt.f32(pre + "input_layernorm.weight"), rt.f32(pre + "input_layernorm.bias"), d["eps"]) |
| a = self._attn(l, pre, h, pos) |
| x = x + rt.f32(pre + "self_attn_layer_scale.scale") * rt.lin(pre + "self_attn.o_proj", a, rt.amx) |
| h = _ln(x, rt.f32(pre + "post_attention_layernorm.weight"), rt.f32(pre + "post_attention_layernorm.bias"), d["eps"]) |
| u = np.ascontiguousarray(rt.lin(pre + "mlp.fc1", h, rt.amx), np.float32) |
| rt.k.lib.tmk_gelu(ptr(u), u.size) |
| return x + rt.f32(pre + "mlp_layer_scale.scale") * rt.lin(pre + "mlp.fc2", u, rt.amx) |
|
|
| def push(self, lat): |
| return self.push_b(self.push_a(lat)) |
|
|
| def push_a(self, lat): |
| return self.push_s0(self.push_t(lat)) |
|
|
| def push_t(self, lat): |
| rt, d = self.rt, self.d |
| F = lat.shape[0] |
| x = lat @ rt.small("dec:latent_proj").T |
| s = d["latent_stride"] |
| w = rt.small("dec:decoder.upsample.conv").reshape(-1, 2 * s) |
| up = np.zeros((F * s + s, x.shape[1]), np.float32) |
| for k in range(2 * s): |
| up[k:k + F * s:s] += x * w[:, k][None] |
| up[:s] += self.up_tail |
| self.up_tail = up[F * s:].copy() |
| h = up[:F * s] |
| pos = np.arange(self.pos, self.pos + h.shape[0]) |
| self.pos += h.shape[0] |
| for l in range(d["layers"]): |
| h = self._tf_layer(l, h, pos) |
| return h |
|
|
| def push_s0(self, h): |
| return self._stage(0, self._taps_conv(self.SEA + "0.conv", h, self.d["kernel"])) |
|
|
| def push_b(self, y): |
| d = self.d |
| self._dense_push() |
| piece = self.piece * d["ratios"][0] |
| out = np.concatenate([self._tail(y[i:i + piece]) for i in range(0, y.shape[0], piece)]) |
| self.pd = {} |
| return out |
|
|
| def _stage(self, j, y): |
| d, p = self.d, self.SEA |
| ratio, li = d["ratios"][j], 1 + 3 * j |
| if self.rt.amx: |
| y = self._convT(f"{p}{li + 1}.conv", y, 2 * ratio, ratio, act=1) |
| r = self._taps_conv(f"{p}{li + 2}.block.1.conv", y, d["res_kernel"], act=1) |
| y += self._taps_conv(f"{p}{li + 2}.block.3.conv", r, 1, act=1) |
| return y |
| y = self._elu(y, inplace=True) |
| y = self._convT(f"{p}{li + 1}.conv", y, 2 * ratio, ratio) |
| r = self._taps_conv(f"{p}{li + 2}.block.1.conv", self._elu(y), d["res_kernel"]) |
| y += self._taps_conv(f"{p}{li + 2}.block.3.conv", self._elu(r, inplace=True), 1) |
| return y |
|
|
| def _tail(self, y): |
| d = self.d |
| for j in range(1, len(d["ratios"])): |
| y = self._stage(j, y) |
| li = 1 + 3 * len(d["ratios"]) |
| y = self._taps_conv(f"{self.SEA}{li + 1}.conv", self._elu(y, inplace=True), d["last_kernel"]) |
| return y[:, 0].astype(np.float32) |
|
|
