Blame for ogg/model.py

"""Ported from FashionFlora/ogg-tts (Apache-2.0), ogg_tts/omnifish.py. Changes: no training-machine paths, the
prompt cross-attention switch is a constructor argument, from_omnivoice -> from_config_dir (config + tokenizer only).

OmniFish v2: OmniVoice (k2-fsa, Qwen3-0.6B, masked discrete diffusion, bidirectional) predicts the
Fish Audio Drama 3 / S2 codec's semantic codebook; a small codebook-MTP decoder predicts the 9 residual
codebooks, chained like Drama 3's audio decoder (layer k sees layer k-1's state and code k, predicts
code k+1), but ~170M params, trained from scratch, attending over all frames in both directions
(non-autoregressive: at inference every layer runs on all frames at once).

Block diffusion (as LLaDA2.0 / BD3-LM): the target is generated in blocks of frames, left to right; inside
a block, masked iterative unmasking. Training sequence:
  [<|lang_start|>pl<|lang_end|><|instruct_start|>None<|instruct_end|>] [<|text_start|>text<|text_end|>]
  [prompt frames] [clean target blocks, semantic ids] [clean target blocks, all 10 codes] [noisy target blocks]
attention (see block_attention): context (text + prompt) sees only context; a full block (10 codes) sees context
and full blocks <= b; a semantic-only block and a noisy block b see context, full blocks < b and themselves. The
three copies share positions. So the backbone predicts a block from the full codes (timbre included) of the
blocks before it, and the MTP reads the semantic-only copy. Length is the model's own call (Rainbow Padding): the speech is followed by one
END frame and the PAD_IDS cycle up to the block boundary; generation stops at the block holding END.
Audio frame inputs (all 10 codebook embeddings summed): prompt frames and full target frames carry their 10
codes; semantic-only and noisy target frames carry only the semantic id (or MASK_ID) and MASK_ID for codebooks
1..9 - the same at training and inference. Losses: semantic CE on masked noisy frames (END/PAD included) + residual CE of the MTP, trained as masked
prediction over frames (see mtp_unmask_loss): it decodes a block codebook by codebook, unmasking the frames of
each codebook in a few rounds while seeing every decided code of the earlier blocks.
"""
import json
import math
import os
import re

import torch
import torch.utils.checkpoint
import torch.nn.functional as F
from torch import nn

CODEBOOK_SIZES = [4096] + [1024] * 9
NUM_CODEBOOKS = len(CODEBOOK_SIZES)
MASK_ID = max(CODEBOOK_SIZES)          # 4096, shared input mask id for every codebook
END_ID = MASK_ID + 1                   # 4097, semantic id of the frame after the speech
PAD_IDS = list(range(END_ID + 1, END_ID + 8))  # 4098..4104, Rainbow Padding cycle after END
VOCAB = PAD_IDS[-1] + 1                # 4105: per-codebook input embedding width, semantic head width
OLD_VOCAB = MASK_ID + 1                # checkpoints before block diffusion (no END/PAD)
RESIDUAL_WEIGHTS = [8.0, 6.35, 5.11, 4.22, 3.63, 3.29, 3.08, 3.01, 3.0]
LANGUAGES = ("None", "en", "pl", "de", "fr", "es", "it", "pt", "ru")
LANG_DIM = 64
MTP_KV_HEADS = int(os.environ.get("OMNIFISH_MTP_KV_HEADS", "2"))
MTP_DIM = int(os.environ.get("OMNIFISH_MTP_DIM", "768"))      # MTP width (the backbone is 1024)
MTP_BLOCKS = int(os.environ.get("OMNIFISH_MTP_BLOCKS", "2"))  # transformer blocks per codebook
MTP_ARCH = os.environ.get("OMNIFISH_MTP_ARCH", "axial")
AXIAL_CKPT = os.environ.get("OMNIFISH_AXIAL_CKPT", "0") == "1"  # activation checkpointing of the AxialMTP layers


def rotary(x, positions, base=10000.0):
    """x [B, H, T, D], positions [B, T]: rotate-half RoPE"""
    d = x.shape[-1]
    inv = base ** (-torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)
    ang = positions.float()[:, None, :, None] * inv
    cos, sin = ang.cos(), ang.sin()
    x1, x2 = x[..., : d // 2].float(), x[..., d // 2:].float()
    return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), -1).to(x.dtype)


# speaker / recording IDs: one learned vector each (OmniFish.spk_embeddings / q_embeddings + projections), one token in the sequence instead of the
# sub-word pieces of "pl_20" / "q35"; id_vocab.json is append-only (a new ID gets the next row)
ID_VOCAB = json.load(open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "id_vocab.json")))
ID_INDEX = {n: i for i, n in enumerate(ID_VOCAB)}
N_SPK = sum(1 for n in ID_VOCAB if not re.fullmatch(r"q\d+", n))  # speakers (and language bags) first, then recordings
SPK_DIM, Q_DIM = 256, 64
_ID_RE = re.compile(r"(?<=speaker:)(?:[A-Za-z]{2}_\w+|spk_\w+)|(?<=quality:)(q\d+)|(?<=<\|quad_start\|>)(?:[A-Za-z]{2}_\w+|spk_\w+)(?=<\|quad_end\|>)")


def encode_text(tok, s, text_vocab, ids=True):
    """tokenize s -> input_ids [1, L]; with ids, every known speaker / recording ID after "speaker:", "quality:" or
    inside a turn marker is one token text_vocab + its row in ID_VOCAB (unknown names stay text)"""
    if not ids:
        return tok(s, return_tensors="pt").input_ids
    out, last = [], 0
    for m in _ID_RE.finditer(s):
        name = m.group(0)
        if name not in ID_INDEX:
            continue
        if m.start() > last:
            out += tok(s[last:m.start()], add_special_tokens=False).input_ids
        out.append(text_vocab + ID_INDEX[name])
        last = m.end()
    if last < len(s):
        out += tok(s[last:], add_special_tokens=False).input_ids
    return torch.tensor([out], dtype=torch.long)



class IdTokenizer:
    """the text tokenizer with speaker / recording IDs as single tokens (encode_text); everything else (decode,
    pad_token_id, ...) is the wrapped tokenizer's"""

    def __init__(self, tok, text_vocab):
        self.tok, self.text_vocab = tok, text_vocab

    def __call__(self, s, return_tensors="pt", **kw):
        class _Out:
            pass
        out = _Out()
        out.input_ids = encode_text(self.tok, s, self.text_vocab)
        return out

    def __getattr__(self, name):
        if name == "tok" or name.startswith("__"):  # not set yet (unpickling) / dunder lookups
            raise AttributeError(name)
        return getattr(self.tok, name)


def context_text(text, before=(), after=()):
    """<|object_ref_start|>before<|object_ref_end|><|text_start|>text<|text_end|><|box_start|>after<|box_end|>"""
    b = f"<|object_ref_start|>{' '.join(before)}<|object_ref_end|>" if before else ""
    a = f"<|box_start|>{' '.join(after)}<|box_end|>" if after else ""
    return f"{b}<|text_start|>{text}<|text_end|>{a}"


def dialogue_text(turns):
    """[(speaker token, line), ...] -> "<|quad_start|>pl_3<|quad_end|>line line <|quad_start|>pl_7<|quad_end|>line";
    a marker only where the speaker changes"""
    out, last = [], None
    for spk, line in turns:
        out.append(line if spk == last else f"<|quad_start|>{spk}<|quad_end|>{line}")
        last = spk
    return " ".join(out)


def block_attention(kind, block):
    """kind [B, L] (0 context, 1 clean target with semantic ids only, 3 clean target with all 10 codes, 2 noisy
    target, -1 padding), block [B, L] -> [B, 1, L, L]. Without kind 3 (the default, --full-context off): context
    sees context, clean block b sees clean blocks <= b, noisy block b sees clean blocks < b and itself; earlier
    blocks are seen by their semantic ids only. With kind 3: context sees context; full block b (3) sees context and
    full blocks <= b; semantic-only block b (1, the MTP's input states) and noisy block b (2) see context, full
    blocks < b and themselves. So a block's semantic is predicted from everything decided before it, residual
    codes of earlier blocks included, and the MTP never sees the codes it predicts. Padding rows see every real
    position (no empty row)."""
    kq, kk, bq, bk = kind[:, :, None], kind[:, None, :], block[:, :, None], block[:, None, :]
    if (kind == 3).any():  # full-code context (3 copies)
        att = ((kk == 0) & (kq >= 0)) | ((kq == 3) & (kk == 3) & (bk <= bq)) \
            | ((kq == 1) & (kk == 3) & (bk < bq)) | ((kq == 1) & (kk == 1) & (bk == bq)) \
            | ((kq == 2) & (kk == 3) & (bk < bq)) | ((kq == 2) & (kk == 2) & (bk == bq)) | ((kq == -1) & (kk >= 0))
    else:  # semantic-only context (2 copies): clean block b sees clean blocks <= b, noisy block b clean < b and itself
        att = ((kk == 0) & (kq >= 0)) | ((kq == 1) & (kk == 1) & (bk <= bq)) \
            | ((kq == 2) & (kk == 1) & (bk < bq)) | ((kq == 2) & (kk == 2) & (bk == bq)) | ((kq == -1) & (kk >= 0))
    return att[:, None]


class MTPBlock(nn.Module):
    """one bidirectional transformer block over frames (pre-norm attention + SwiGLU)"""

    def __init__(self, dim, n_head, n_kv, head_dim, ffn, eps=1e-6, causal=False):
        super().__init__()
        self.n_head, self.n_kv, self.head_dim, self.causal = n_head, n_kv, head_dim, causal
        self.attn_norm, self.ffn_norm = nn.RMSNorm(dim, eps=eps), nn.RMSNorm(dim, eps=eps)
        self.wq = nn.Linear(dim, n_head * head_dim, bias=False)
        self.wkv = nn.Linear(dim, 2 * n_kv * head_dim, bias=False)
        self.wo = nn.Linear(n_head * head_dim, dim, bias=False)
        self.w1, self.w3 = nn.Linear(dim, ffn, bias=False), nn.Linear(dim, ffn, bias=False)
        self.w2 = nn.Linear(ffn, dim, bias=False)

    def forward(self, x, positions, valid, attn):
        B, T, _ = x.shape
        a = self.attn_norm(x)
        q = self.wq(a).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k, v = self.wkv(a).view(B, T, 2, self.n_kv, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k = rotary(q, positions), rotary(k, positions)
        rep = self.n_head // self.n_kv
        k, v = k.repeat_interleave(rep, 1), v.repeat_interleave(rep, 1)
        mask = valid[:, None, None, :] if attn is None else attn
        if attn is None and self.causal:  # autoregressive models: a frame sees only earlier frames
            mask = mask & torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        x = x + self.wo(o.transpose(1, 2).reshape(B, T, -1))
        return x + self.w2(F.silu(self.w1(self.ffn_norm(x))) * self.w3(self.ffn_norm(x)))


class MTPLayer(nn.Module):
    """Mix code k's embedding (and the already decided code k+1 of other frames, MASK where not decided) into the
    running state (in_dim: the backbone width for the first layer, the MTP width after it), MTP_BLOCKS
    bidirectional transformer blocks over frames, predict code k+1."""

    def __init__(self, dim, vocab_in, vocab_out, in_dim=None, n_head=8, n_kv=2, head_dim=None, ffn=None, blocks=1,
                 eps=1e-6, causal=False):
        super().__init__()
        in_dim = in_dim or dim
        head_dim = head_dim or dim // n_head
        ffn = ffn or int(dim * 3.5)
        self.causal = causal
        self.embedding = nn.Embedding(vocab_in, dim)
        self.revealed = nn.Embedding(vocab_out + 1, dim)  # code k+1 where already decided; index vocab_out = MASK
        self.mask_id = vocab_out
        self.enorm, self.hnorm = nn.RMSNorm(dim, eps=eps), nn.RMSNorm(in_dim, eps=eps)
        self.eh_proj = nn.Linear(dim + in_dim, dim, bias=False)
        self.blocks = nn.ModuleList(MTPBlock(dim, n_head, n_kv, head_dim, ffn, eps, causal) for _ in range(blocks))
        self.out_norm = nn.RMSNorm(dim, eps=eps)
        self.head = nn.Linear(dim, vocab_out, bias=False)

    def forward(self, h, codes, positions, valid, attn=None, revealed=None):
        """h [B, T, in_dim], codes [B, T], positions [B, T], valid [B, T], attn [B, 1, T, T] (optional, overrides
        valid/causal), revealed [B, T] code k+1 or mask_id (None: all masked) -> (state [B, T, dim], logits)"""
        if revealed is None:
            revealed = torch.full_like(codes, self.mask_id)
        e = self.embedding(codes) + self.revealed(revealed)
        x = self.eh_proj(torch.cat((self.enorm(e), self.hnorm(h)), -1))
        for b in self.blocks:
            x = b(x, positions, valid, attn)
        x = self.out_norm(x)
        return x, self.head(x)


class PromptCross(nn.Module):
    """attention from every grid token to all prompt tokens (every prompt frame, every level); output projection
    starts at zero, so a model trained without it is unchanged until this is learned"""

    def __init__(self, dim, n_head, n_kv, head_dim, eps=1e-6):
        super().__init__()
        self.n_head, self.n_kv, self.head_dim = n_head, n_kv, head_dim
        self.qnorm, self.knorm = nn.RMSNorm(dim, eps=eps), nn.RMSNorm(dim, eps=eps)
        self.wq = nn.Linear(dim, n_head * head_dim, bias=False)
        self.wkv = nn.Linear(dim, 2 * n_kv * head_dim, bias=False)
        self.wo = nn.Linear(n_head * head_dim, dim, bias=False)
        nn.init.zeros_(self.wo.weight)

    def forward(self, x, p, pmask):
        """x [B, N, D] queries, p [B, M, D] prompt tokens, pmask [B, M] valid prompt tokens"""
        B, N, _ = x.shape; M = p.shape[1]
        q = self.wq(self.qnorm(x)).view(B, N, self.n_head, self.head_dim).transpose(1, 2)
        k, v = self.wkv(self.knorm(p)).view(B, M, 2, self.n_kv, self.head_dim).permute(2, 0, 3, 1, 4)
        rep = self.n_head // self.n_kv
        k, v = k.repeat_interleave(rep, 1), v.repeat_interleave(rep, 1)
        has = pmask.any(1)
        m = (pmask | ~has[:, None])[:, None, None, :]  # rows without a prompt: any key, output dropped below
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=m)
        return x + self.wo(o.transpose(1, 2).reshape(B, N, -1)) * has[:, None, None].to(x.dtype)


class AxialMTP(nn.Module):
    """codebooks 1..9 by attention over a (frame x level) grid: token (t, j) carries code j of frame t (+ the frame's
    backbone state, + a level embedding) and predicts code j+1. Every layer: attention over the levels of a frame
    (causal: j sees <= j; a prompt frame sees all its levels, it is fully known), attention over frames at the same
    level (a frame sees the prompt, earlier blocks and its own block), SwiGLU. Decoding: 9 steps per block, all frames
    at once (level j+1 of the block from the tokens up to level j). Same interface as MTPDecoder (new_state, prime,
    decode_block)."""

    def __init__(self, in_dim=1024, dim=768, layers=6, n_head=8, n_kv=2, ffn=2688, eps=1e-6, ctx_frames=None, prompt_x=False):
        super().__init__()
        self.prompt_x = prompt_x
        self.dim, self.in_dim, self.ctx_frames = dim, in_dim, ctx_frames
        self.sizes = [VOCAB] + list(CODEBOOK_SIZES[1:])
        self.register_buffer("offsets", torch.tensor([0] + list(torch.tensor(self.sizes).cumsum(0)[:-1].tolist())), persistent=False)
        self.code_emb = nn.Embedding(sum(self.sizes), dim)
        self.level_emb = nn.Embedding(NUM_CODEBOOKS, dim)
        self.in_norm = nn.RMSNorm(in_dim, eps=eps)
        self.in_proj = nn.Linear(in_dim, dim, bias=False)
        hd = dim // n_head
        self.lvl = nn.ModuleList(MTPBlock(dim, n_head, n_kv, hd, ffn, eps) for _ in range(layers))  # level attention + FFN
        self.frm = nn.ModuleList(MTPBlock(dim, n_head, n_kv, hd, ffn, eps) for _ in range(layers))  # frame attention + FFN
        self.px = nn.ModuleList(PromptCross(dim, n_head, n_kv, hd, eps) for _ in range(layers)) if prompt_x else None
        self.out_norm = nn.RMSNorm(dim, eps=eps)
        self.heads = nn.ModuleList(nn.Linear(dim, CODEBOOK_SIZES[j + 1], bias=False) for j in range(NUM_CODEBOOKS - 1))

    def _tokens(self, h, codes):
        """h [B, T, in_dim], codes [B, 10, T] -> [B, T, 10, dim]"""
        c = codes.clamp(min=0) + self.offsets.view(1, -1, 1)
        x = self.code_emb(c).permute(0, 2, 1, 3) + self.level_emb.weight[None, None]
        return x + self.in_proj(self.in_norm(h))[:, :, None]

    def _run(self, x, pos, valid, block):
        """x [B, T, L, D]; masks from block (-1 prompt) / valid -> [B, T, L, D]"""
        B, T, L, D = x.shape
        lv = torch.arange(L, device=x.device)
        causal = lv[None, :, None] >= lv[None, None, :]                                       # [1, L, L]
        lmask = (causal | (block < 0).reshape(B * T, 1, 1))[:, None]                            # prompt frames: all levels
        fmask = (block[:, None, :] <= block[:, :, None]) & valid[:, None, :]                   # [B, T(q), T(k)]
        fmask = (fmask | torch.eye(T, dtype=torch.bool, device=x.device)[None])[:, None]       # a row always sees itself
        fmask = fmask.repeat_interleave(L, 0)                                                   # [B*L, 1, T, T]
        lpos = lv[None].expand(B * T, L)
        fpos = pos.repeat_interleave(L, 0)
        Pm = int((block < 0).sum(1).max()) if self.px is not None else 0  # prompt frames lead every row
        pmask = ((block[:, :Pm] < 0) & valid[:, :Pm])[:, :, None].expand(B, Pm, L).reshape(B, Pm * L) if Pm else None

        def layer(x, la, fa, px):
            x = la(x.reshape(B * T, L, D), lpos, None, lmask).reshape(B, T, L, D)
            y = fa(x.permute(0, 2, 1, 3).reshape(B * L, T, D), fpos, None, fmask)
            x = y.reshape(B, L, T, D).permute(0, 2, 1, 3)
            if px is not None and Pm:  # every token also reads the whole prompt (all frames, all levels)
                x = px(x.reshape(B, T * L, D), x[:, :Pm].reshape(B, Pm * L, D), pmask).reshape(B, T, L, D)
            return x
        for n, (la, fa) in enumerate(zip(self.lvl, self.frm)):
            px = self.px[n] if self.px is not None else None
            if AXIAL_CKPT and self.training and torch.is_grad_enabled():  # optional: recompute in backward (less memory)
                x = torch.utils.checkpoint.checkpoint(layer, x, la, fa, px, use_reentrant=False)
            else:
                x = layer(x, la, fa, px)
        return self.out_norm(x)

    def loss(self, h, codes, pos, valid, block):
        """teacher forcing: h [B, T, in_dim], codes [B, 10, T] true codes, block [B, T] (-1: prompt frames, context
        only) -> 9 losses (codebooks 1..9) on the valid target frames"""
        x = self._run(self._tokens(h, codes), pos, valid, block)
        tgt = valid & (block >= 0)
        res = []
        for j, head in enumerate(self.heads):
            logits = head(x[:, :, j][tgt])
            res.append(F.cross_entropy(logits.float(), codes[:, j + 1][tgt]))
        return torch.stack(res)

    def new_state(self, T, dim, device, dtype):
        return {"codes": torch.zeros(NUM_CODEBOOKS, T, dtype=torch.long, device=device),
                "states": [torch.zeros(1, T, dim, dtype=dtype, device=device)]}

    def prime(self, S, P):
        pass  # the prompt frames are read straight from S at every step

    @torch.no_grad()
    def decode_block(self, S, lo, hi, block, rounds=4, temperature=0.0, top_k=0, prefix=0):
        """fill codebooks 1..9 of frames lo:hi (S["states"][0] backbone states, S["codes"][0] semantic ids set;
        frames < lo decided; the first `prefix` frames are the prompt). Context: the prompt + every decided frame
        (as in training; ctx_frames caps it if set)."""
        codes, dev = S["codes"], S["codes"].device
        start = prefix if self.ctx_frames is None else max(prefix, lo - self.ctx_frames)  # default: all, as in training
        keep = torch.cat([torch.arange(prefix, device=dev), torch.arange(start, hi, device=dev)])
        i = keep
        blk = torch.where(i < prefix, -1, (i - prefix) // block)[None]
        h = S["states"][0][:, keep]
        pos = i[None]
        valid = torch.ones_like(blk, dtype=torch.bool)
        cur = (i >= lo)
        for j in range(NUM_CODEBOOKS - 1):
            c = codes[:, keep][None].clone()
            x = self._run(self._tokens(h, c), pos, valid, blk)
            lg = self.heads[j](x[0, cur, j]).float()
            if temperature > 0:
                if top_k > 0:
                    lg = lg.masked_fill(lg < lg.topk(top_k, -1).values[..., -1:], -float("inf"))
                nxt = torch.distributions.Categorical(logits=lg / temperature).sample()
            else:
                nxt = lg.argmax(-1)
            codes[j + 1, lo:hi] = nxt


class MTPDecoder(nn.Module):
    """codebooks 1..9 from the backbone state and the semantic codes: layer k embeds code k, predicts k+1"""

    def __init__(self, dim=1024, causal=False):
        """dim: the backbone width (the first layer's input); the MTP runs at MTP_DIM with MTP_BLOCKS blocks per
        codebook"""
        super().__init__()
        self.dim, self.in_dim = MTP_DIM, dim
        self.layers = nn.ModuleList(MTPLayer(MTP_DIM, CODEBOOK_SIZES[k], CODEBOOK_SIZES[k + 1], in_dim=dim if k == 0 else MTP_DIM,
                                             n_kv=MTP_KV_HEADS, blocks=MTP_BLOCKS, causal=causal)
                                    for k in range(NUM_CODEBOOKS - 1))

    def forward(self, h, codes, positions, valid, attn=None):
        """teacher forcing: codes [B, 10, T] true codes -> list of 9 logits [B, T, 1024] (codebooks 1..9)"""
        out = []
        for k, layer in enumerate(self.layers):
            h, logits = layer(h, codes[:, k], positions, valid, attn)
            out.append(logits)
        return out

    def new_state(self, T, dim, device, dtype):
        """buffers for block-by-block decoding of up to T frames: codes [10, T], clean input states of every layer"""
        return {"codes": torch.zeros(NUM_CODEBOOKS, T, dtype=torch.long, device=device),  # states[0]: backbone width
                "states": [torch.zeros(1, T, dim if k == 0 else self.dim, dtype=dtype, device=device) for k in range(len(self.layers) + 1)]}

    @torch.no_grad()
    def prime(self, S, P):
        """the P prompt frames at the start of S (backbone states + all 10 codes set): run every layer's clean pass
        over them once, so the target blocks see them as decided frames (block -1)"""
        codes, states, dev = S["codes"], S["states"], S["codes"].device
        pos = torch.arange(P, device=dev)[None]
        full = torch.ones(1, 1, P, P, dtype=torch.bool, device=dev)
        for k, layer in enumerate(self.layers):
            st, _ = layer(states[k][:, :P], codes[k, :P][None], pos, None, full, codes[k + 1, :P][None])
            states[k + 1][:, :P] = st.to(states[k + 1].dtype)

    def decode_block(self, S, lo, hi, block, rounds=4, temperature=0.0, top_k=0, prefix=0):
        """fill codebooks 1..9 of frames lo:hi (S["states"][0] = backbone states, S["codes"][0] = semantic ids
        already set there; frames < lo decided). Codebook by codebook, the frames of a codebook are unmasked over
        `rounds` iterations (most confident first, cosine schedule), each seeing the decided codes of the earlier
        blocks and of the frames revealed so far; layer k's input state is layer k-1's clean pass (all decided)."""
        codes, states, dev = S["codes"], S["states"], S["codes"].device
        pos = torch.arange(hi, device=dev)[None]
        i = torch.arange(hi, device=dev)
        bb = torch.where(i < prefix, -1, (i - prefix) // block)  # the first `prefix` frames: the prompt (block -1)
        kind = (torch.arange(hi, device=dev) >= lo).long()
        attn = (((bb[None, :, None] >= bb[None, None, :]) & (kind[None, None, :] == 0)) |
                ((kind[None, :, None] == 1) & (kind[None, None, :] == 1)))[:, None]
        cattn = (bb[None, :, None] >= bb[None, None, :])[:, None]
        for k, layer in enumerate(self.layers):
            hk = states[k][:, :hi]
            cur = torch.full((hi - lo,), layer.mask_id, dtype=torch.long, device=dev)
            for r in range(rounds):
                rev = torch.cat([codes[k + 1, :lo], cur])[None]
                _, logits = layer(hk, codes[k, :hi][None], pos, None, attn, rev)
                lg = logits[0, lo:hi].float()
                masked = cur == layer.mask_id
                probs = F.softmax(lg / max(temperature, 1e-5), -1)
                if temperature > 0:
                    if top_k > 0:
                        lg = lg.masked_fill(lg < lg.topk(top_k, -1).values[..., -1:], -float("inf"))
                    sample = torch.distributions.Categorical(logits=lg / temperature).sample()
                else:
                    sample = lg.argmax(-1)
                conf = torch.where(masked, probs.gather(1, sample[:, None])[:, 0], torch.full_like(probs[:, 0], -1.0))
                keep = int(math.floor((hi - lo) * math.cos(math.pi / 2 * (r + 1) / rounds)))
                n_fix = int(masked.sum()) - keep
                if n_fix > 0:
                    idx = conf.topk(n_fix).indices
                    cur[idx] = sample[idx]
            codes[k + 1, lo:hi] = cur
            st, _ = layer(hk, codes[k, :hi][None], pos, None, cattn, codes[k + 1, :hi][None])
            states[k + 1][:, lo:hi] = st[:, lo:hi].to(states[k + 1].dtype)

    @torch.no_grad()
    def generate(self, h, semantic, block, rounds=4, temperature=0.0, top_k=0):
        """semantic [T] codes of one utterance, h [T, dim] its backbone states -> codes [10, T], block by block"""
        T = semantic.shape[0]
        S = self.new_state(T, h.shape[-1], h.device, h.dtype)
        S["codes"][0], S["states"][0][0] = semantic, h
        for lo in range(0, T, block):
            self.decode_block(S, lo, min(T, lo + block), block, rounds, temperature, top_k)
        return S["codes"]


class OmniFish(nn.Module):
    def __init__(self, llm, prompt_x=False):
        super().__init__()
        self.llm = llm
        dim = llm.config.hidden_size
        self.audio_embeddings = nn.Embedding(NUM_CODEBOOKS * VOCAB, dim)
        self.semantic_head = nn.Linear(dim, VOCAB, bias=False)  # semantic codes, MASK (never a target), END, PADs
        self.mtp = AxialMTP(dim, prompt_x=prompt_x) if MTP_ARCH == "axial" else MTPDecoder(dim)
        self.text_vocab = llm.get_input_embeddings().weight.shape[0]
        # speaker / recording IDs (see encode_text): small vectors in a shared space, one projection each to the
        # model width (rare speakers learn 256 numbers, the projection learns from all of them)
        self.spk_embeddings = nn.Embedding(N_SPK, SPK_DIM)
        self.spk_proj = nn.Linear(SPK_DIM, dim)
        self.q_embeddings = nn.Embedding(len(ID_VOCAB) - N_SPK, Q_DIM)
        self.q_proj = nn.Linear(Q_DIM, dim)
        self.lang_embeddings = nn.Embedding(len(LANGUAGES), LANG_DIM)
        self.lang_proj = nn.Linear(LANG_DIM, dim)
        self.lang_start_id = None
        self.lang_token_to_index = {}
        std = float(llm.get_input_embeddings().weight.detach().std())
        nn.init.normal_(self.audio_embeddings.weight, std=std)
        nn.init.normal_(self.semantic_head.weight, std=1e-3)  # start from ~uniform semantic logits
        for head in (self.mtp.heads if MTP_ARCH == "axial" else [layer.head for layer in self.mtp.layers]):
            nn.init.normal_(head.weight, std=1e-3)
        self.register_buffer("codebook_offsets", torch.arange(NUM_CODEBOOKS) * VOCAB, persistent=False)
        w = torch.tensor(RESIDUAL_WEIGHTS, dtype=torch.float)
        self.register_buffer("residual_weights", w / w.sum(), persistent=False)

    @classmethod
    def from_config_dir(cls, path, dtype=torch.float32, attn="sdpa", prompt_x=False):
        """Build the model (no weights) from an OmniVoice folder holding config.json and the tokenizer files.
        The released OGG file carries the backbone weights too, so OmniVoice's own weights are not read."""
        from transformers import AutoModel, AutoTokenizer, Qwen3Config

        cfg = json.load(open(os.path.join(path, "config.json")))
        llm_cfg = Qwen3Config(**{k: v for k, v in cfg["llm_config"].items() if k != "model_type"})
        llm_cfg._attn_implementation = attn
        try:
            llm = AutoModel.from_config(llm_cfg, dtype=dtype)
        except TypeError:  # older transformers
            llm = AutoModel.from_config(llm_cfg, torch_dtype=dtype)
        tok = AutoTokenizer.from_pretrained(path)
        model = cls(llm, prompt_x=prompt_x).to(dtype)
        model.configure_languages(tok)
        return model, tok

    def configure_languages(self, tok):
        """Map the single token immediately after <|lang_start|> to the learned language table."""
        self.lang_start_id = tok.convert_tokens_to_ids("<|lang_start|>")
        self.lang_token_to_index = {}
        for i, lang in enumerate(LANGUAGES):
            ids = tok(lang, add_special_tokens=False).input_ids
            if len(ids) == 1:
                self.lang_token_to_index[int(ids[0])] = i

    def embed(self, input_ids, audio_mask):
        """input_ids [B, 10, L] (text ids repeated on every channel), audio_mask [B, L] -> [B, L, H]"""
        t = input_ids[:, 0]
        is_id = (t >= self.text_vocab) & ~audio_mask
        text = self.llm.get_input_embeddings()(torch.where(is_id, 0, t) * ~audio_mask)
        if is_id.any():
            k = (t - self.text_vocab).clamp(0, len(ID_VOCAB) - 1)
            spk = self.spk_proj(self.spk_embeddings(k.clamp(max=N_SPK - 1)))
            q = self.q_proj(self.q_embeddings((k - N_SPK).clamp(min=0)))
            text = torch.where((is_id & (k < N_SPK))[..., None], spk.to(text.dtype), text)
            text = torch.where((is_id & (k >= N_SPK))[..., None], q.to(text.dtype), text)
        if self.lang_start_id is not None and self.lang_token_to_index:
            after_start = torch.zeros_like(t, dtype=torch.bool)
            after_start[:, 1:] = t[:, :-1] == self.lang_start_id
            lang_mask = after_start & ~audio_mask
            lang_index = torch.zeros_like(t)
            for token_id, index in self.lang_token_to_index.items():
                lang_index = torch.where(t == token_id, index, lang_index)
            lang = self.lang_proj(self.lang_embeddings(lang_index))
            text = torch.where(lang_mask[..., None], lang.to(text.dtype), text)
        shifted = input_ids * audio_mask[:, None] + self.codebook_offsets.view(1, -1, 1)
        return torch.where(audio_mask[..., None], self.audio_embeddings(shifted).sum(1), text)

    def load_weights(self, state, tok):
        """load a (possibly older) state dict; a checkpoint from before the ID tables starts them from the text
        pieces of each ID (init_id_embeddings). -> True when the tables were new"""
        state = OmniFish.upgrade_state(state)
        fresh = "spk_embeddings.weight" not in state
        if fresh:
            state = dict(state, **{k: v.detach().cpu() for k, v in self.state_dict().items()
                                   if k.startswith(("spk_", "q_"))})
        own = self.state_dict()
        missing_px = [k for k in own if ".px." in k and k not in state]  # new prompt cross-attention: fresh init
        if missing_px:
            state = dict(state, **{k: own[k] for k in missing_px})
        self.load_state_dict(state)
        if fresh:
            self.init_id_embeddings(tok)
        return fresh

    @torch.no_grad()
    def init_id_embeddings(self, tok):
        """start the ID tables so that projection(vector) is the best rank-SPK_DIM / rank-Q_DIM fit (SVD) of the
        mean text embedding of each ID's sub-word pieces: the model keeps what it learned about "pl_20" as text"""
        emb = self.llm.get_input_embeddings().weight.float()
        for names, table, proj, r in ((ID_VOCAB[:N_SPK], self.spk_embeddings, self.spk_proj, SPK_DIM),
                                      (ID_VOCAB[N_SPK:], self.q_embeddings, self.q_proj, Q_DIM)):
            M = torch.stack([emb[tok(n, add_special_tokens=False).input_ids].mean(0) for n in names])  # [n, dim]
            mu = M.mean(0)
            U, S, Vh = torch.linalg.svd(M - mu, full_matrices=False)
            k = min(r, S.shape[0])
            table.weight.zero_(); proj.weight.zero_()
            table.weight[:, :k] = (U[:, :k] * S[:k].sqrt()).to(table.weight.dtype)
            proj.weight[:, :k] = (Vh[:k].T * S[:k].sqrt()).to(proj.weight.dtype)
            proj.bias.copy_(mu.to(proj.bias.dtype))

    def backbone(self, input_ids, audio_mask, attention_mask, position_ids):
        return self.llm(inputs_embeds=self.embed(input_ids, audio_mask), attention_mask=attention_mask,
                        position_ids=position_ids, return_dict=True)[0]

    @staticmethod
    def upgrade_state(state):
        """weights saved before END/PAD (input/semantic vocab 4097/4096) -> current shapes; new rows freshly
        initialised. Optimizer moments for those two tensors are remapped the same way (zeros for new rows)."""
        out = dict(state)
        emb = state["audio_embeddings.weight"]
        if emb.shape[0] == NUM_CODEBOOKS * VOCAB:
            return out
        dim = emb.shape[1]
        new = torch.randn(NUM_CODEBOOKS, VOCAB, dim, dtype=emb.dtype) * emb.std()
        new[:, :OLD_VOCAB] = emb.view(NUM_CODEBOOKS, OLD_VOCAB, dim)
        out["audio_embeddings.weight"] = new.view(-1, dim)
        head = state["semantic_head.weight"]
        newh = torch.randn(VOCAB, dim, dtype=head.dtype) * 1e-3
        newh[:head.shape[0]] = head
        out["semantic_head.weight"] = newh
        return out

    @staticmethod
    def upgrade_moment(name, t):
        """an AdamW moment of an old-shape parameter -> current shape, zeros for the new rows"""
        if t.dim() < 2:  # norms, biases: never reshaped
            return t
        dim = t.shape[1]
        if name == "audio_embeddings.weight" and t.shape[0] != NUM_CODEBOOKS * VOCAB:
            new = torch.zeros(NUM_CODEBOOKS, VOCAB, dim, dtype=t.dtype)
            new[:, :OLD_VOCAB] = t.view(NUM_CODEBOOKS, OLD_VOCAB, dim)
            return new.view(-1, dim)
        if name == "semantic_head.weight" and t.shape[0] != VOCAB:
            new = torch.zeros(VOCAB, dim, dtype=t.dtype)
            new[:t.shape[0]] = t
            return new
        return t

    @staticmethod
    def gather_targets(h, codes, target_mask):
        """target frames of every row, left-aligned: h [B, T, H], codes [B, 10, T], positions, valid [B, T]"""
        B = h.shape[0]
        n = target_mask.sum(1)
        T = int(n.max())
        order = torch.argsort((~target_mask).to(torch.int8), dim=1, stable=True)[:, :T]  # target positions first
        valid = torch.arange(T, device=h.device)[None] < n[:, None]
        ht = h.gather(1, order[..., None].expand(B, T, h.shape[-1]))
        ct = codes.gather(2, order[:, None].expand(B, codes.shape[1], T))
        positions = torch.arange(T, device=h.device)[None].expand(B, T)
        return ht, ct, positions, valid, order

    def mtp_unmask_loss(self, ht, ct, pos, valid, block):
        """MTP as masked prediction over frames: per layer, every block masks code k+1 of its frames at its own
        uniform random ratio; [clean frames | noisy frames]: clean frame i sees clean frames of blocks <= its block;
        noisy frame i sees clean frames of earlier blocks and noisy frames of its own block. Loss on masked noisy
        frames. -> 9 losses"""
        B, T = valid.shape
        kind = torch.cat([torch.zeros(B, T, dtype=torch.long, device=ht.device), torch.ones(B, T, dtype=torch.long, device=ht.device)], 1)
        bb = torch.cat([block, block], 1)
        v2 = torch.cat([valid, valid], 1)
        kq, kk, bq, bk = kind[:, :, None], kind[:, None, :], bb[:, :, None], bb[:, None, :]
        attn = ((kq == 0) & (kk == 0) & (bk <= bq)) | ((kq == 1) & (kk == 0) & (bk < bq)) | ((kq == 1) & (kk == 1) & (bk == bq))
        attn = (attn | ~v2[:, :, None]) & v2[:, None, :]  # padding rows see the valid frames (no empty row)
        attn = attn | torch.eye(2 * T, dtype=torch.bool, device=ht.device)[None]  # a row always sees itself
        attn = attn[:, None]
        pos2 = torch.cat([pos, pos], 1)
        nb = int(block.max()) + 1
        tgt_frame = valid & (block >= 0)  # block -1: the prompt's decided frames (context only, no loss)
        first = (block < 0).sum(1).clamp(max=T - 1)  # first target frame of each row
        state, res = ht, []
        for k, layer in enumerate(self.mtp.layers):
            ratio = torch.rand(B, nb, device=ht.device)
            m = (torch.rand(B, T, device=ht.device) < ratio.gather(1, block.clamp(min=0))) & tgt_frame
            none = ~m.any(1)  # at least one masked target frame per row
            m[torch.arange(B, device=ht.device)[none], first[none]] = tgt_frame[torch.arange(B, device=ht.device)[none], first[none]]
            tgt = ct[:, k + 1]
            rev = torch.cat([tgt, torch.where(m, torch.full_like(tgt, layer.mask_id), tgt)], 1)
            st, logits = layer(torch.cat([state, state], 1), torch.cat([ct[:, k], ct[:, k]], 1), pos2, v2, attn, rev)
            res.append(F.cross_entropy(logits[:, T:].float()[m], tgt[m]))
            state = st[:, :T]
        return torch.stack(res)

    @staticmethod
    def frame_attention(block, valid):
        """MTP mask over frames: frame i sees valid frame j iff block_j <= block_i -> [B, 1, T, T] (padding rows
        see the valid frames too, so no row is empty)"""
        return ((block[:, None, :, None] >= block[:, None, None, :]) | ~valid[:, None, :, None]) & valid[:, None, None, :]

    def forward(self, input_ids, audio_mask, attention_mask, position_ids, semantic_labels, codes, target_mask,
                target_block):
        """semantic_labels [B, L] (-100 except masked noisy frames), codes [B, 10, L] true codes, target_mask
        [B, L] speech frames of the clean copy (MTP targets), target_block [B, L] their block index"""
        h = self.backbone(input_ids, audio_mask, attention_mask, position_ids)
        sem_logits = self.semantic_head(h)
        sem_loss = F.cross_entropy(sem_logits.float().flatten(0, 1), semantic_labels.flatten(), ignore_index=-100)
        ht, ct, pos, valid, order = self.gather_targets(h, codes, target_mask)
        blk = target_block.gather(1, order)
        res = self.mtp.loss(ht, ct, pos, valid, blk) if MTP_ARCH == "axial" else self.mtp_unmask_loss(ht, ct, pos, valid, blk)
        loss = sem_loss + (res * self.residual_weights).sum()
        return {"loss": loss, "semantic": sem_loss.detach(), "residual": res.detach()}