| """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) |
| END_ID = MASK_ID + 1 |
| PAD_IDS = list(range(END_ID + 1, END_ID + 8)) |
| VOCAB = PAD_IDS[-1] + 1 |
| OLD_VOCAB = MASK_ID + 1 |
| 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_BLOCKS = int(os.environ.get("OMNIFISH_MTP_BLOCKS", "2")) |
| MTP_ARCH = os.environ.get("OMNIFISH_MTP_ARCH", "axial") |
| AXIAL_CKPT = os.environ.get("OMNIFISH_AXIAL_CKPT", "0") == "1" |
|
|
|
|
| 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) |
|
|
|
|
| |
| |
| 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)) |
| 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("__"): |
| 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(): |
| 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: |
| 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: |
| 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) |
| 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, :] |
| 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)) |
| self.frm = nn.ModuleList(MTPBlock(dim, n_head, n_kv, hd, ffn, eps) for _ in range(layers)) |
| 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, :] |
| lmask = (causal | (block < 0).reshape(B * T, 1, 1))[:, None] |
| fmask = (block[:, None, :] <= block[:, :, None]) & valid[:, None, :] |
| fmask = (fmask | torch.eye(T, dtype=torch.bool, device=x.device)[None])[:, None] |
| fmask = fmask.repeat_interleave(L, 0) |
| 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 |
| 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: |
| 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(): |
| 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 |
|
|
| @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) |
| 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": [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) |
| 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) |
| 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] |
| |
| |
| 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) |
| 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: |
| 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] |
| 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]) |
| 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: |
| 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] |
| 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, :] |
| attn = attn | torch.eye(2 * T, dtype=torch.bool, device=ht.device)[None] |
| attn = attn[:, None] |
| pos2 = torch.cat([pos, pos], 1) |
| nb = int(block.max()) + 1 |
| tgt_frame = valid & (block >= 0) |
| first = (block < 0).sum(1).clamp(max=T - 1) |
| 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) |
| 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()} |
|
|
