| """Depth decoder (31 levels per frame) on MLX: serve/fast_depth.FastDepth._compute, fp16 by default. |
| |
| Same function as the reference: prefix position k is embedded from slot k alone, levels sharing a trunk group reuse the |
| group's K/V (exact because attention is causal), the last layer of a group's first pass evaluates only the final |
| position, and the speaker cross-attention is a per-voice constant (one reference token => softmax is exactly one). |
| """ |
| import math |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
| from .prompt import RMS_EPS |
|
|
| D, H, FFN, LAYERS, N_LEVELS = 1024, 16, 3072, 4, 32 |
| HEAD = D // H |
| GROUP_OF_LEVEL = [None] + list(range(8)) + [8] * 8 + [9] * 8 + [10] * 7 |
| EPS = 1e-5 |
|
|
|
|
| class Depth: |
| def __init__(self, w, dtype=mx.float16, temp=0.7, top_k=64, n_codes=2048): |
| c = lambda k: w[k].astype(dtype) |
| self.dtype, self.temp, self.top_k, self.n_codes = dtype, float(temp), int(top_k), n_codes |
| self.in_proj = c("depth.in_proj.weight") |
| level = c("depth.level_emb.weight") |
| self.level0 = level[0] |
| |
| lo = w["depth.code_lo.weight"].astype(mx.float32).reshape(N_LEVELS - 1, n_codes, -1) |
| table = (lo @ w["depth.code_up.weight"].astype(mx.float32).T).astype(dtype) |
| |
| self.embed = (table + level[1:N_LEVELS, None]).reshape(-1, D) |
| self.groups = [[{n: c(f"depth.groups.{g}.{l}.{n}") for n in |
| ("n1.weight", "n1.bias", "n2.weight", "n2.bias", "qkv.weight", "o.weight", "f1.weight", "f2.weight")} |
| for l in range(LAYERS)] for g in range(11)] |
| self.cross = [{n: c(f"depth.cross.{l}.{n}") for n in ("kv.weight", "o.weight")} for l in range(LAYERS)] |
| self.norm_w, self.norm_b = c("depth.norm.weight"), c("depth.norm.bias") |
| self.heads = [c(f"depth.heads.{i}.weight") for i in range(N_LEVELS - 1)] |
| self.spk_proj = w["depth.spk_proj_ecapa.weight"] |
| self.null_mem = c("depth.null_mem") |
| |
| self.fb_table, self.fb_norm = w["audio_in.table.weight"], w["audio_in.norm.weight"] |
| self._frame = mx.compile(self._compute) |
|
|
| def speaker(self, refs=None): |
| """Per-layer cross-attention constants for a voice: refs (192,) ECAPA or None (nulled memory).""" |
| |
| mem = self.null_mem[0] if refs is None else \ |
| (mx.array(refs, mx.float32)[None] @ self.spk_proj.astype(mx.float32).T).astype(self.dtype) |
| return mx.stack([(mem @ L["kv.weight"].T)[:, D:] @ L["o.weight"].T for L in self.cross]) |
|
|
| def feedback(self, codes): |
| lvl = mx.arange(N_LEVELS) * self.n_codes |
| e = self.fb_table[codes + lvl].sum(1) / math.sqrt(N_LEVELS) |
| return mx.fast.rms_norm(e, self.fb_norm, RMS_EPS) |
|
|
| def _trunk_layer(self, L, x, kv_prev, causal, last_only): |
| B, T, _ = x.shape |
| qkv = (mx.fast.layer_norm(x, L["n1.weight"], L["n1.bias"], EPS) @ L["qkv.weight"].T).reshape(B, T, 3, H, HEAD) |
| q, k, v = (qkv[:, :, i].transpose(0, 2, 1, 3) for i in range(3)) |
| if kv_prev is not None: |
| k, v = mx.concatenate([kv_prev[0], k], 2), mx.concatenate([kv_prev[1], v], 2) |
| if last_only: |
| q, x = q[:, :, -1:], x[:, -1:] |
| a = mx.fast.scaled_dot_product_attention(q, k, v, scale=HEAD ** -0.5, mask="causal" if causal else None) |
| x = x + a.transpose(0, 2, 1, 3).reshape(B, x.shape[1], D) @ L["o.weight"].T |
| x = x + nn.gelu(mx.fast.layer_norm(x, L["n2.weight"], L["n2.bias"], EPS) @ L["f1.weight"].T) @ L["f2.weight"].T |
| return x, (k, v) |
|
|
| def _compute(self, h, cb0, cross, key, force=None): |
| """h [B, 2048] backbone hidden, cb0 [B] int, cross [LAYERS, 1, D] -> (codes [B, 32], feedback [B, 2048], logits list). |
| force [B, 32]: teacher forcing (prefix built from these codes instead of the sampled ones; for parity checks).""" |
| B = h.shape[0] |
| prefix = [(h.astype(self.dtype) @ self.in_proj.T + self.level0)[:, None], self.embed[cb0][:, None]] |
| codes, logits_all, kv, previous = [cb0], [], None, None |
| for level in range(1, N_LEVELS): |
| group, n = GROUP_OF_LEVEL[level], level + 1 |
| incremental = group == previous |
| x = prefix[level] if incremental else mx.concatenate(prefix[:n], 1) |
| new_kv = [] |
| for layer, L in enumerate(self.groups[group]): |
| last = (not incremental) and layer == LAYERS - 1 |
| x, kvl = self._trunk_layer(L, x, kv[layer] if incremental else None, causal=not incremental and not last, |
| last_only=last) |
| new_kv.append(kvl) |
| x = x + cross[layer] |
| kv, previous = new_kv, group |
| y = mx.fast.layer_norm(x[:, -1], self.norm_w, self.norm_b, EPS) |
| logits = (y @ self.heads[level - 1].T).astype(mx.float32) / self.temp |
| logits_all.append(logits) |
| if force is not None: |
| code = force[:, level] |
| elif key is None: |
| code = mx.argmax(logits, -1) |
| else: |
| kth = mx.topk(logits, self.top_k, -1).min(-1, keepdims=True) |
| probs = mx.softmax(mx.where(logits < kth, -mx.inf, logits), -1) |
| key, sub = mx.random.split(key) |
| race = -mx.log(mx.random.uniform(shape=probs.shape, key=sub)) |
| code = mx.argmax(probs / race, -1) |
| codes.append(code) |
| if n < N_LEVELS: |
| prefix.append(self.embed[code + level * self.n_codes][:, None]) |
| codes = mx.stack(codes, -1) |
| return codes, self.feedback(codes), logits_all |
|
|
| def __call__(self, h, cb0, cross, key): |
| codes, fb, _ = self._frame(h, cb0, cross, key) |
| return codes, fb |
|
|
