"""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)