| """Phonon (gradium/phonon-31m-en) building blocks for recipes/phonon-31m-en.py: weights, eager reference, export programs. |
| |
| The eager reference follows the official runtime (the `ptts` crate, src/{flow_lm,tts_model, |
| transformer,mlp,conditioners,mimi,seanet,conv}.rs) operation by operation: |
| |
| voice speaker_wavs [1,512,125] -> transpose -> speaker_proj (512x512, no bias) -> [1,125,512] |
| prime transformer(voice) then transformer(text_embed(tokens)), both written to the KV cache |
| step x = input_linear(prev latent or bos_emb) + cond (cond = num_speakers("1") LUT |
| + duration_delta(0.0) sinusoid, each through its output_proj) |
| h = out_norm(transformer(x))[:, -1]; eos_logit = out_eos(h) |
| latent = noise + flow_net(h, s=0, t=1, noise) (LSD, 1 step; noise ~ N(0, temp)) |
| decode (latent * emb_std + emb_mean) -> 1x1 conv 32->512 -> depthwise convtr x16 -> 2-layer |
| transformer (window 250, layer scale) -> SEANet decoder (ratios 6,5,4) -> 1920 samples |
| |
| Export programs (WebGPU): |
| FlowLMProgram forward(embeds[1,S,512] S<=128, input_pos[1] i64, noise[1,32]) |
| -> latent[1,32], eos_logit[1,1], next_embeds[1,1,512] (KV cache in the program) |
| MimiStepProgram forward(latent[1,1,32], *state) -> pcm[1,1,1920], *new_state (state explicit) |
| |
| Rewrites for the public WebGPU runtime (same function, checked against the eager reference): |
| * RoPE: q/k rows of in_proj permuted from interleaved pairs to halves, so the runtime's fused |
| HF RoPE (et_vk.apply_rotary_emb_hf) applies; q.k is unchanged. |
| * flow-LM LayerNorm (dynamic S): mean as a 1/512 linear + the rms_norm pattern, because |
| native_layer_norm has no resize hook and mean.dim's resize drops a keepdim'ed dim. |
| * ELU = relu(x) + exp(min(x, 0)) - 1 and SiLU = x * sigmoid(x) (no aten.elu / aten.silu). |
| * ConvTranspose1d as a 4-D (H = 1) conv_transpose2d (3-D convolution has no transposed path). |
| * Mimi streaming state as explicit tensors: mutable buffers other than the llama KV cache are |
| not kept live by the delegate. |
| """ |
|
|
| import math |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| D = 512 |
| H = 8 |
| HD = 64 |
| LDIM = 32 |
| FLOW_LAYERS = 6 |
| MIMI_LAYERS = 2 |
| MIMI_CONTEXT = 250 |
| FRAME = 1920 |
| UPS = 16 |
| NEG = -1.0e30 |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _p(*shape): |
| return nn.Parameter(torch.zeros(*shape), requires_grad=False) |
|
|
|
|
| class _Lin(nn.Module): |
| def __init__(self, i, o, bias): |
| super().__init__() |
| self.weight = _p(o, i) |
| if bias: |
| self.bias = _p(o) |
|
|
|
|
| class _LN(nn.Module): |
| def __init__(self, n): |
| super().__init__() |
| self.weight = _p(n) |
| self.bias = _p(n) |
|
|
|
|
| class _Attn(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.in_proj = _Lin(D, 3 * D, False) |
| self.out_proj = _Lin(D, D, False) |
|
|
|
|
| class _Scale(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.scale = _p(D) |
|
|
|
|
| class _Layer(nn.Module): |
| def __init__(self, layer_scale): |
| super().__init__() |
| self.self_attn = _Attn() |
| self.norm1 = _LN(D) |
| self.norm2 = _LN(D) |
| self.linear1 = _Lin(D, 2048, False) |
| self.linear2 = _Lin(2048, D, False) |
| if layer_scale: |
| self.layer_scale_1 = _Scale() |
| self.layer_scale_2 = _Scale() |
|
|
|
|
| class _Transformer(nn.Module): |
| def __init__(self, n, layer_scale): |
| super().__init__() |
| self.layers = nn.ModuleList(_Layer(layer_scale) for _ in range(n)) |
|
|
|
|
| class _Seq(nn.Module): |
| """nn.Sequential-style numbered children with only the indices that hold weights.""" |
|
|
| def __init__(self, items): |
| super().__init__() |
| for k, m in items.items(): |
| self.add_module(str(k), m) |
|
|
| def __getitem__(self, i): |
| return getattr(self, str(i)) |
|
|
|
|
| class _TimeEmbed(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.freqs = _p(128) |
| self.mlp = _Seq({0: _Lin(256, D, True), 2: _Lin(D, D, True), 3: _Alpha()}) |
|
|
|
|
| class _Alpha(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.alpha = _p(D) |
|
|
|
|
| class _ResBlock(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.in_ln = _LN(D) |
| self.mlp = _Seq({0: _Lin(D, D, True), 2: _Lin(D, D, True)}) |
| self.adaLN_modulation = _Seq({1: _Lin(D, 3 * D, True)}) |
|
|
|
|
| class _Final(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.linear = _Lin(D, LDIM, True) |
| self.adaLN_modulation = _Seq({1: _Lin(D, 2 * D, True)}) |
|
|
|
|
| class _FlowNet(nn.Module): |
| def __init__(self, depth=6): |
| super().__init__() |
| self.time_embed = nn.ModuleList([_TimeEmbed(), _TimeEmbed()]) |
| self.cond_embed = _Lin(D, D, True) |
| self.input_proj = _Lin(LDIM, D, True) |
| self.res_blocks = nn.ModuleList(_ResBlock() for _ in range(depth)) |
| self.final_layer = _Final() |
|
|
|
|
| class _Embed(nn.Module): |
| def __init__(self, n, d): |
| super().__init__() |
| self.weight = _p(n, d) |
|
|
|
|
| class _Cond(nn.Module): |
| def __init__(self, embed=None, proj_in=None): |
| super().__init__() |
| self.learnt_padding = _p(1, 1, D) |
| if embed is not None: |
| self.embed = _Embed(*embed) |
| if proj_in is not None: |
| self.output_proj = _Lin(proj_in, D, False) |
|
|
|
|
| class _Conditioners(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.duration_delta = _Cond(proj_in=128) |
| self.num_speakers = _Cond(embed=(32, 16), proj_in=16) |
| self.speaker_wavs = _Cond(proj_in=512) |
| self.transcript_in_segment = _Cond(embed=(4001, D)) |
|
|
|
|
| class _CondProvider(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.conditioners = _Conditioners() |
|
|
|
|
| class _WST(nn.Module): |
| """flow.w_s_t: a training-time loss-weighting net; the runtime drops it (loader.rs remap_key).""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.w_s_t = _Seq( |
| { |
| 0: _Lin(2, 32, True), |
| 3: _Lin(32, 32, True), |
| 6: _Lin(32, 32, True), |
| 9: _Lin(32, 1, True), |
| } |
| ) |
|
|
|
|
| class FlowLMWeights(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.bos_emb = _p(LDIM) |
| self.emb_mean = _p(LDIM) |
| self.emb_std = _p(LDIM) |
| self.condition_provider = _CondProvider() |
| self.flow = _WST() |
| self.flow_net = _FlowNet() |
| self.input_linear = _Lin(LDIM, D, False) |
| self.out_eos = _Lin(D, 1, True) |
| self.out_norm = _LN(D) |
| self.transformer = _Transformer(FLOW_LAYERS, False) |
|
|
|
|
| class _Conv(nn.Module): |
| def __init__(self, i, o, k): |
| super().__init__() |
| self.conv = _ConvW(i, o, k, bias=True) |
|
|
|
|
| class _ConvW(nn.Module): |
| def __init__(self, i, o, k, bias): |
| super().__init__() |
| self.weight = _p(o, i, k) |
| if bias: |
| self.bias = _p(o) |
|
|
|
|
| class _ConvTr(nn.Module): |
| def __init__(self, i, o, k): |
| super().__init__() |
| self.convtr = _ConvTrW(i, o, k, True) |
|
|
|
|
| class _ConvTrW(nn.Module): |
| def __init__(self, i, o_per_group, k, bias): |
| super().__init__() |
| self.weight = _p(i, o_per_group, k) |
| if bias: |
| self.bias = _p(o_per_group) |
|
|
|
|
| class _Res(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.block = _Seq({1: _Conv(dim, dim // 2, 3), 3: _Conv(dim // 2, dim, 1)}) |
|
|
|
|
| class _SEANetDecoder(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.model = _Seq( |
| { |
| 0: _Conv(512, 512, 7), |
| 2: _ConvTr(512, 256, 12), |
| 3: _Res(256), |
| 5: _ConvTr(256, 128, 10), |
| 6: _Res(128), |
| 8: _ConvTr(128, 64, 8), |
| 9: _Res(64), |
| 11: _Conv(64, 1, 3), |
| } |
| ) |
|
|
|
|
| class _ProjT(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.transformer = _Transformer(MIMI_LAYERS, True) |
|
|
|
|
| class _Quantizer(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.output_proj = _ConvW(LDIM, D, 1, bias=False) |
|
|
|
|
| class _Upsample(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.convtr = _Seq({}) |
| self.convtr.convtr = _ConvTrW(D, 1, 2 * UPS, False) |
|
|
|
|
| class MimiWeights(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.decoder = _SEANetDecoder() |
| self.decoder_transformer = _ProjT() |
| self.quantizer = _Quantizer() |
| self.upsample = _Upsample() |
|
|
|
|
| class PhononWeights(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.flow_lm = FlowLMWeights() |
| self.mimi = MimiWeights() |
|
|
|
|
| def load_weights(safetensors_path): |
| """Every checkpoint tensor, as float32 (bf16 -> f32 is exact), loaded with strict=True.""" |
| from safetensors.torch import load_file |
|
|
| sd = {k: v.float() for k, v in load_file(str(safetensors_path)).items()} |
| w = PhononWeights() |
| w.load_state_dict(sd, strict=True) |
| return w.eval() |
|
|
|
|
| def load_voice(path): |
| """speaker_wavs [1, 512, T] (stored speaker-Mimi latents) from a voice .safetensors file.""" |
| from safetensors.torch import load_file |
|
|
| t = load_file(str(path)) |
| assert list(t) == ["speaker_wavs"], list(t) |
| return t["speaker_wavs"].float() |
|
|
|
|
| |
| |
| |
|
|
|
|
| def lin(m, x): |
| """nn.Linear on a 2-D view (the public WebGPU runtime's fp32 aten.linear takes 2-D input).""" |
| shp = x.shape |
| y = F.linear(x.reshape(-1, shp[-1]), m.weight, getattr(m, "bias", None)) |
| return y.reshape(*shp[:-1], y.shape[-1]) |
|
|
|
|
| def lin3(m, x): |
| """nn.Linear on the tensor as it is (any rank); the Vulkan lowering adds its own 2-D views.""" |
| return F.linear(x, m.weight, getattr(m, "bias", None)) |
|
|
|
|
| def layer_norm_dyn(m, x, eps, mean_w): |
| """LayerNorm for the dynamic-length flow-LM program, from ops that follow a resized sequence |
| length on the public WebGPU runtime (its native_layer_norm has no resize hook, and its mean.dim |
| resize drops a keepdim'ed dim): the mean is a linear with weight 1/512 (exact: a power of two), |
| the rest matches the Vulkan rms_norm pattern (fused into et_vk.rms_norm).""" |
| xc = x - F.linear(x, mean_w) |
| return xc * torch.rsqrt((xc * xc).mean(-1, keepdim=True) + eps) * m.weight + m.bias |
|
|
|
|
| def layer_norm(m, x, eps): |
| return F.layer_norm(x, (x.shape[-1],), m.weight, m.bias, eps) |
|
|
|
|
| def silu(x): |
| return x * torch.sigmoid(x) |
|
|
|
|
| def elu(x): |
| """ELU(alpha=1) from ops the public WebGPU runtime registers (it has no aten.elu).""" |
| return torch.relu(x) + torch.exp(torch.clamp(x, max=0.0)) - 1.0 |
|
|
|
|
| def gelu(x): |
| return F.gelu(x) |
|
|
|
|
| def inv_freq(): |
| """RoPE frequencies as the runtime computes them (rope.rs): 1 / max_period^(i / half) in f32.""" |
| half = HD // 2 |
| i = torch.arange(half, dtype=torch.float32) |
| return 1.0 / torch.pow(torch.tensor(10000.0, dtype=torch.float32), i / half) |
|
|
|
|
| def rope_interleaved(x, pos): |
| """rope_i on [B, T, H, HD] with absolute positions pos [T] (the runtime's convention: |
| pairs (2i, 2i+1) rotate by pos * inv_freq[i]).""" |
| ang = pos.to(torch.float32)[:, None] * inv_freq()[None, :] |
| c, s = torch.cos(ang)[None, :, None, :], torch.sin(ang)[None, :, None, :] |
| xr, xi = x[..., 0::2], x[..., 1::2] |
| return torch.stack([xr * c - xi * s, xr * s + xi * c], dim=-1).flatten(-2) |
|
|
|
|
| def time_embed_rms(te, t): |
| """TimestepEmbedder with TimeRmsNorm::Rms (phonon config time_rms_norm: rms).""" |
| args = t * te.freqs |
| emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| x = lin(te.mlp[2], silu(lin(te.mlp[0], emb))) |
| ms = (x * x).mean(-1, keepdim=True) |
| return x * torch.rsqrt(ms + 1e-5) * te.mlp[3].alpha |
|
|
|
|
| def flow_net(fn, c, x, t_comb): |
| """SimpleMLPAdaLN: c [N, 512] backbone output, x [N, 32] noise, t_comb [N or 1, 512].""" |
| x = lin(fn.input_proj, x) |
| y = t_comb + lin(fn.cond_embed, c) |
| sy = silu(y) |
| for b in fn.res_blocks: |
| ada = lin(b.adaLN_modulation[1], sy) |
| shift, scale, gate = ada[..., :D], ada[..., D : 2 * D], ada[..., 2 * D :] |
| h = layer_norm(b.in_ln, x, 1e-6) * (scale + 1.0) + shift |
| h = lin(b.mlp[2], silu(lin(b.mlp[0], h))) |
| x = x + gate * h |
| fl = fn.final_layer |
| ada = lin(fl.adaLN_modulation[1], sy) |
| shift, scale = ada[..., :D], ada[..., D:] |
| x = F.layer_norm(x, (D,), None, None, 1e-6) * (scale + 1.0) + shift |
| return lin(fl.linear, x) |
|
|
|
|
| def lsd_time(fn): |
| """(time_embed[0](s=0) + time_embed[1](t=1)) / 2, the flow head's constant time condition |
| for lsd_decode_steps = 1.""" |
| s = torch.zeros(1, 1) |
| t = torch.ones(1, 1) |
| return ( |
| time_embed_rms(fn.time_embed[0], s) + time_embed_rms(fn.time_embed[1], t) |
| ) / 2 |
|
|
|
|
| def sin_embedding(position, dim, max_period): |
| half = dim // 2 |
| i = torch.arange(half, dtype=torch.float32) |
| ph = position / torch.pow( |
| torch.tensor(max_period, dtype=torch.float32), i / (half - 1) |
| ) |
| return torch.cat([torch.cos(ph), torch.sin(ph)]) |
|
|
|
|
| def summed_conditions(flw, num_speakers="1", duration_delta=0.0): |
| """conditioners.rs Conditioners::sum with the runtime's defaults, [1, 1, 512].""" |
| c = flw.condition_provider.conditioners |
| idx = ["0", "1", "2"].index(num_speakers) |
| ns = F.linear( |
| c.num_speakers.embed.weight[idx : idx + 1], c.num_speakers.output_proj.weight |
| ) |
| dd = F.linear( |
| sin_embedding(1000.0 * duration_delta, 128, 10000.0)[None], |
| c.duration_delta.output_proj.weight, |
| ) |
| return (ns + dd).view(1, 1, D) |
|
|
|
|
| def voice_embedding(flw, speaker_wavs): |
| """[1, 512, T] speaker latents -> [1, T, 512] voice prompt (loader.rs project_latents).""" |
| w = flw.condition_provider.conditioners.speaker_wavs.output_proj.weight |
| return F.linear(speaker_wavs.transpose(1, 2), w) |
|
|
|
|
| def text_embedding(flw, tokens): |
| emb = flw.condition_provider.conditioners.transcript_in_segment.embed.weight |
| return emb[torch.as_tensor(tokens, dtype=torch.long)][None] |
|
|
|
|
| |
| |
| |
|
|
|
|
| class FlowLMRef: |
| def __init__(self, w: PhononWeights, num_speakers="1", duration_delta=0.0): |
| self.w = w.flow_lm |
| self.cond = summed_conditions(self.w, num_speakers, duration_delta) |
| self.t_comb = lsd_time(self.w.flow_net) |
| self.reset() |
|
|
| def reset(self): |
| self.k = [None] * FLOW_LAYERS |
| self.v = [None] * FLOW_LAYERS |
| self.pos = 0 |
|
|
| @torch.no_grad() |
| def backbone(self, x): |
| B, T, _ = x.shape |
| pos = torch.arange(self.pos, self.pos + T) |
| for i, L in enumerate(self.w.transformer.layers): |
| h = layer_norm(L.norm1, x, 1e-5) |
| qkv = lin(L.self_attn.in_proj, h).view(B, T, 3, H, HD) |
| q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] |
| q, k = rope_interleaved(q, pos), rope_interleaved(k, pos) |
| self.k[i] = k if self.k[i] is None else torch.cat([self.k[i], k], 1) |
| self.v[i] = v if self.v[i] is None else torch.cat([self.v[i], v], 1) |
| kk, vv = self.k[i].transpose(1, 2), self.v[i].transpose(1, 2) |
| qq = q.transpose(1, 2) |
| Tk = kk.shape[2] |
| qpos = torch.arange(Tk - T, Tk)[:, None] |
| mask = torch.where(torch.arange(Tk)[None, :] <= qpos, 0.0, float("-inf")) |
| att = torch.softmax(qq @ kk.transpose(-1, -2) / math.sqrt(HD) + mask, -1) |
| o = (att @ vv).transpose(1, 2).reshape(B, T, D) |
| x = x + lin(L.self_attn.out_proj, o) |
| h = layer_norm(L.norm2, x, 1e-5) |
| x = x + lin(L.linear2, gelu(lin(L.linear1, h))) |
| self.pos += T |
| return x |
|
|
| def prompt(self, embeds): |
| self.backbone(embeds) |
|
|
| @torch.no_grad() |
| def step(self, prev_latent, noise): |
| """prev_latent None (BOS) or [1, 1, 32]; noise [1, 32]. -> latent [1,1,32], eos_logit.""" |
| seq = self.w.bos_emb.view(1, 1, LDIM) if prev_latent is None else prev_latent |
| x = lin(self.w.input_linear, seq) + self.cond |
| h = layer_norm(self.w.out_norm, self.backbone(x), 1e-5)[:, -1] |
| eos = lin(self.w.out_eos, h) |
| latent = noise + flow_net(self.w.flow_net, h, noise, self.t_comb) |
| return latent.view(1, 1, LDIM), eos |
|
|
|
|
| def _conv1d_stream(m, x, prev, dilation=1): |
| xx = torch.cat([prev, x], 2) if prev is not None else x |
| y = F.conv1d(xx, m.conv.weight, m.conv.bias, dilation=dilation) |
| tp = prev.shape[2] if prev is not None else 0 |
| return y, (xx[:, :, xx.shape[2] - tp :] if tp else None) |
|
|
|
|
| def _convtr_stream(w, bias, x, partial, stride, groups=1): |
| y = F.conv_transpose1d(x, w, bias, stride=stride, groups=groups) |
| pt = partial.shape[2] |
| y = torch.cat([y[:, :, :pt] + partial, y[:, :, pt:]], 2) |
| new_partial = y[:, :, y.shape[2] - pt :] |
| if bias is not None: |
| new_partial = new_partial - bias.view(1, -1, 1) |
| return y[:, :, : y.shape[2] - pt], new_partial |
|
|
|
|
| class MimiRef: |
| """Streaming Mimi decoder (MimiDecoder::decode_from_latent_step), one or more frames per call.""" |
|
|
| def __init__(self, w: PhononWeights): |
| self.w = w.mimi |
| self.emb_std = w.flow_lm.emb_std |
| self.emb_mean = w.flow_lm.emb_mean |
| self.reset() |
|
|
| def reset(self): |
| z = lambda *s: torch.zeros(1, *s) |
| self.up = z(D, UPS) |
| self.kv = [[None, None] for _ in range(MIMI_LAYERS)] |
| self.offset = 0 |
| self.prev = { |
| "init": z(512, 6), |
| "r1": z(256, 2), |
| "r2": z(128, 2), |
| "r3": z(64, 2), |
| "final": z(64, 2), |
| } |
| self.partial = {"u1": z(256, 6), "u2": z(128, 5), "u3": z(64, 4)} |
|
|
| @torch.no_grad() |
| def transformer(self, x): |
| T = x.shape[1] |
| pos = torch.arange(self.offset, self.offset + T) |
| for i, L in enumerate(self.w.decoder_transformer.transformer.layers): |
| h = layer_norm(L.norm1, x, 1e-5) |
| qkv = lin(L.self_attn.in_proj, h).view(1, T, 3, H, HD) |
| q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] |
| q, k = rope_interleaved(q, pos), rope_interleaved(k, pos) |
| q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) |
| pk, pv = self.kv[i] |
| kk = k if pk is None else torch.cat([pk, k], 2) |
| vv = v if pv is None else torch.cat([pv, v], 2) |
| kvlen = kk.shape[2] - T |
| qi = torch.arange(T)[:, None] + kvlen |
| kj = torch.arange(kk.shape[2])[None, :] |
| ok = (kj <= qi) & (kj >= qi - MIMI_CONTEXT) |
| mask = torch.where(ok, 0.0, float("-inf")) |
| att = torch.softmax(q @ kk.transpose(-1, -2) / math.sqrt(HD) + mask, -1) |
| o = (att @ vv).transpose(1, 2).reshape(1, T, D) |
| self.kv[i] = [kk[:, :, -MIMI_CONTEXT:], vv[:, :, -MIMI_CONTEXT:]] |
| x = x + L.layer_scale_1.scale * lin(L.self_attn.out_proj, o) |
| h = layer_norm(L.norm2, x, 1e-5) |
| x = x + L.layer_scale_2.scale * lin(L.linear2, gelu(lin(L.linear1, h))) |
| self.offset += T |
| return x |
|
|
| def _res(self, blk, x, key): |
| v, self.prev[key] = _conv1d_stream(blk.block[1], elu(x), self.prev[key]) |
| v, _ = _conv1d_stream(blk.block[3], elu(v), None) |
| return x + v |
|
|
| @torch.no_grad() |
| def decode(self, latent): |
| """latent [1, T, 32] (flow-LM space) -> pcm [1, 1, T * 1920].""" |
| x = (latent * self.emb_std + self.emb_mean).transpose(1, 2) |
| x = F.conv1d(x, self.w.quantizer.output_proj.weight) |
| x, self.up = _convtr_stream( |
| self.w.upsample.convtr.convtr.weight, None, x, self.up, UPS, groups=D |
| ) |
| x = self.transformer(x.transpose(1, 2)).transpose(1, 2) |
| m = self.w.decoder.model |
| x, self.prev["init"] = _conv1d_stream(m[0], x, self.prev["init"]) |
| for up, res, s, ku, kr in [ |
| (m[2], m[3], 6, "u1", "r1"), |
| (m[5], m[6], 5, "u2", "r2"), |
| (m[8], m[9], 4, "u3", "r3"), |
| ]: |
| x, self.partial[ku] = _convtr_stream( |
| up.convtr.weight, up.convtr.bias, elu(x), self.partial[ku], s |
| ) |
| x = self._res(res, x, kr) |
| x, self.prev["final"] = _conv1d_stream(m[11], elu(x), self.prev["final"]) |
| return x |
|
|
|
|
| |
| |
| |
|
|
|
|
| def prepare_text_prompt(text): |
| text = " ".join(text.strip().replace("\n", " ").replace("\r", " ").split()) |
| if not text: |
| return text, 3 |
| frames_after_eos = 3 if len(text.split()) <= 4 else 1 |
| text = text[0].upper() + text[1:] |
| if text[-1].isalnum(): |
| text += "." |
| return text, frames_after_eos |
|
|
|
|
| def frame_budget(num_tokens, frame_rate=12.5): |
| return math.ceil((num_tokens / 3.0 + 2.0) * frame_rate) |
|
|
|
|
| class EosPolicy: |
| def __init__(self, frames_after_eos): |
| self.n, self.countdown = frames_after_eos, None |
|
|
| def should_stop(self, is_eos): |
| if is_eos and self.countdown is None: |
| self.countdown = self.n |
| if self.countdown is None: |
| return False |
| if self.countdown == 0: |
| return True |
| self.countdown -= 1 |
| return False |
|
|
|
|
| def generate(step_fn, tokens, frames_after_eos, noise_fn, eos_threshold=4.0): |
| """The runtime's loop (synth.rs run_inline): BOS first, the EOS frame and its tail are kept. |
| step_fn(prev_latent_or_None, noise) -> (latent [1,1,32], eos_logit). Returns latents, eos logits.""" |
| eos = EosPolicy(frames_after_eos) |
| prev, lats, logits = None, [], [] |
| for i in range(frame_budget(len(tokens))): |
| lat, e = step_fn(prev, noise_fn(i)) |
| lats.append(lat) |
| logits.append(float(e.reshape(-1)[0])) |
| if eos.should_stop(logits[-1] > eos_threshold): |
| break |
| prev = lat |
| return lats, logits |
|
|
|
|
| def tokenizer(tokenizer_json): |
| from tokenizers import Tokenizer |
|
|
| return Tokenizer.from_file(str(tokenizer_json)) |
|
|
|
|
| def tokenize(tok, text): |
| prepared, frames_after_eos = prepare_text_prompt(text) |
| return ( |
| tok.encode(prepared, add_special_tokens=False).ids, |
| frames_after_eos, |
| prepared, |
| ) |
|
|
|
|
| |
| |
| |
|
|
| MAX_CTX = 512 |
| PREFILL_MAX = 128 |
|
|
|
|
| def _half_split_perm(): |
| """Within each head, reorder q/k channels from interleaved pairs (2i, 2i+1) to halves |
| (i, i + 32): rope_i on the old layout equals rotate_half RoPE on the new one, and q.k is |
| unchanged because q and k are permuted alike.""" |
| per = torch.cat([torch.arange(0, HD, 2), torch.arange(1, HD, 2)]) |
| return torch.cat([h * HD + per for h in range(H)]) |
|
|
|
|
| def _permuted_in_proj(w): |
| """in_proj weight [3*512, 512] with the q and k row blocks permuted, v untouched.""" |
| perm = _half_split_perm() |
| return torch.cat([w[:D][perm], w[D : 2 * D][perm], w[2 * D :]], 0).contiguous() |
|
|
|
|
| def rope_tables(n): |
| """cos/sin [n, 64] in the half-split layout, angle = f32(pos) * f32(inv_freq) as rope.rs.""" |
| ang = torch.arange(n, dtype=torch.float32)[:, None] * inv_freq()[None, :] |
| c, s = torch.cos(ang), torch.sin(ang) |
| return torch.cat([c, c], -1).contiguous(), torch.cat([s, s], -1).contiguous() |
|
|
|
|
| def _rope_hf(q, k, cos, sin): |
| from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingPattern |
|
|
| return HfRotaryEmbeddingPattern()(q, k, cos, sin) |
|
|
|
|
| class _LinP(nn.Module): |
| def __init__(self, weight, bias=None): |
| super().__init__() |
| self.weight = nn.Parameter(weight.detach().clone().contiguous(), requires_grad=False) |
| if bias is not None: |
| self.bias = nn.Parameter(bias.detach().clone().contiguous(), requires_grad=False) |
|
|
|
|
| class FlowLMProgram(nn.Module): |
| """forward(embeds[1,S,512], input_pos[1] i64, noise[1,32]) |
| -> latent[1,32], eos_logit[1,1], next_embeds[1,1,512]. |
| |
| Writes the S positions [input_pos, input_pos + S) of the in-program KV cache (llama |
| sdpa_with_kv_cache, causal), then runs the head on the last position: eos_logit = out_eos(h), |
| latent = noise + flow_net(h, s=0, t=1, noise) and next_embeds = input_linear(latent) + cond, |
| the embedding of the next step's input. Prompt calls (voice, text) ignore the outputs.""" |
|
|
| def __init__(self, w: PhononWeights, num_speakers="1", duration_delta=0.0, attn="sdpa_kv"): |
| super().__init__() |
| from executorch.extension.llm.custom_ops import custom_ops |
|
|
| f = w.flow_lm |
| self.attn = attn |
| self.layers = nn.ModuleList() |
| for L in f.transformer.layers: |
| m = nn.Module() |
| m.norm1, m.norm2 = L.norm1, L.norm2 |
| m.in_proj = _LinP(_permuted_in_proj(L.self_attn.in_proj.weight)) |
| m.out_proj, m.linear1, m.linear2 = L.self_attn.out_proj, L.linear1, L.linear2 |
| self.layers.append(m) |
| self.out_norm, self.out_eos, self.input_linear = f.out_norm, f.out_eos, f.input_linear |
| self.fn = f.flow_net |
| with torch.no_grad(): |
| self.register_buffer("cond", summed_conditions(f, num_speakers, duration_delta).clone(), persistent=False) |
| self.register_buffer("t_comb", lsd_time(f.flow_net).clone(), persistent=False) |
| self.register_buffer("mean_w", torch.full((1, D), 1.0 / D), persistent=False) |
| cos, sin = rope_tables(MAX_CTX) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
| for i in range(FLOW_LAYERS): |
| self.register_buffer(f"k_cache_{i}", torch.zeros(1, MAX_CTX, H, HD), persistent=False) |
| self.register_buffer(f"v_cache_{i}", torch.zeros(1, MAX_CTX, H, HD), persistent=False) |
|
|
| def reset(self): |
| for i in range(FLOW_LAYERS): |
| getattr(self, f"k_cache_{i}").zero_() |
| getattr(self, f"v_cache_{i}").zero_() |
|
|
| def forward(self, embeds, input_pos, noise): |
| S = embeds.shape[1] |
| p = input_pos[0].item() |
| torch._check_is_size(p) |
| torch._check(p + S <= MAX_CTX) |
| cos = self.rope_cos.narrow(0, p, S) |
| sin = self.rope_sin.narrow(0, p, S) |
| x = embeds |
| for i, L in enumerate(self.layers): |
| h = layer_norm_dyn(L.norm1, x, 1e-5, self.mean_w) |
| qkv = lin3(L.in_proj, h) |
| q = qkv[..., :D].reshape(1, S, H, HD) |
| k = qkv[..., D : 2 * D].reshape(1, S, H, HD) |
| v = qkv[..., 2 * D :].reshape(1, S, H, HD) |
| q, k = _rope_hf(q, k, cos, sin) |
| kc, vc = getattr(self, f"k_cache_{i}"), getattr(self, f"v_cache_{i}") |
| o = torch.ops.llama.sdpa_with_kv_cache(q, k, v, kc, vc, p, S, None, 0.0, True, None) |
| x = x + lin3(L.out_proj, o.reshape(1, S, D)) |
| h = layer_norm_dyn(L.norm2, x, 1e-5, self.mean_w) |
| x = x + lin3(L.linear2, gelu(lin3(L.linear1, h))) |
| h = layer_norm(self.out_norm, x[:, -1, :], 1e-5) |
| eos = lin3(self.out_eos, h) |
| latent = noise + flow_net(self.fn, h, noise, self.t_comb) |
| nxt = lin3(self.input_linear, latent).reshape(1, 1, D) + self.cond |
| return latent, eos, nxt |
|
|
|
|
| MIMI_STATE = [ |
| ("upsample_partial", (1, D, UPS)), |
| ("k0", (1, H, MIMI_CONTEXT, HD)), |
| ("v0", (1, H, MIMI_CONTEXT, HD)), |
| ("k1", (1, H, MIMI_CONTEXT, HD)), |
| ("v1", (1, H, MIMI_CONTEXT, HD)), |
| ("kv_seen", (1, 1, 1, MIMI_CONTEXT)), |
| ("rope_offset", (1,)), |
| ("init_conv_prev", (1, 512, 6)), |
| ("up1_partial", (1, 256, 6)), |
| ("res1_prev", (1, 256, 2)), |
| ("up2_partial", (1, 128, 5)), |
| ("res2_prev", (1, 128, 2)), |
| ("up3_partial", (1, 64, 4)), |
| ("res3_prev", (1, 64, 2)), |
| ("final_conv_prev", (1, 64, 2)), |
| ] |
|
|
|
|
| def mimi_zero_state(): |
| return [torch.zeros(s) for _, s in MIMI_STATE] |
|
|
|
|
| def _convtr(x, w, b, s, groups=1): |
| """ConvTranspose1d as a 4-D (H = 1) conv_transpose2d: the public runtime's 3-D convolution |
| has no transposed path, its 4-D one does (with groups).""" |
| y = F.conv_transpose2d(x.unsqueeze(2), w.unsqueeze(2), b, stride=(1, s), groups=groups) |
| return y.squeeze(2) |
|
|
|
|
| def _convtr_step(x, w, b, partial, s, groups=1): |
| y = _convtr(x, w, b, s, groups) |
| pt = partial.shape[2] |
| n = y.shape[2] |
| head = y[:, :, :pt] + partial |
| out = torch.cat([head, y[:, :, pt : n - pt]], 2) |
| tail = y[:, :, n - pt :] |
| if b is not None: |
| tail = tail - b.view(1, -1, 1) |
| return out, tail |
|
|
|
|
| def _conv_step(m, x, prev): |
| xx = torch.cat([prev, x], 2) |
| y = F.conv1d(xx, m.conv.weight, m.conv.bias) |
| return y, xx[:, :, xx.shape[2] - prev.shape[2] :] |
|
|
|
|
| class MimiStepProgram(nn.Module): |
| """forward(latent[1,1,32], *state) -> pcm[1,1,1920], *new_state (MIMI_STATE order). |
| |
| One latent frame per call, exactly the runtime's streaming decoder: every piece of streaming |
| state (conv tails, transposed-conv partial sums, the 2-layer transformer's last 250 keys and |
| values, which of them are filled, and the RoPE offset) is an explicit input and output. The |
| first call takes all-zero state.""" |
|
|
| def __init__(self, w: PhononWeights): |
| super().__init__() |
| m = w.mimi |
| self.emb_std = nn.Parameter(w.flow_lm.emb_std.detach().clone().view(1, 1, LDIM), requires_grad=False) |
| self.emb_mean = nn.Parameter(w.flow_lm.emb_mean.detach().clone().view(1, 1, LDIM), requires_grad=False) |
| self.qproj = _LinP(m.quantizer.output_proj.weight.view(D, LDIM)) |
| self.up_w = nn.Parameter(m.upsample.convtr.convtr.weight.detach().clone(), requires_grad=False) |
| self.layers = nn.ModuleList() |
| for L in m.decoder_transformer.transformer.layers: |
| q = nn.Module() |
| q.norm1, q.norm2 = L.norm1, L.norm2 |
| q.in_proj = _LinP(_permuted_in_proj(L.self_attn.in_proj.weight)) |
| q.out_proj, q.linear1, q.linear2 = L.self_attn.out_proj, L.linear1, L.linear2 |
| q.ls1 = nn.Parameter(L.layer_scale_1.scale.detach().clone(), requires_grad=False) |
| q.ls2 = nn.Parameter(L.layer_scale_2.scale.detach().clone(), requires_grad=False) |
| self.layers.append(q) |
| self.dec = m.decoder.model |
| self.register_buffer("inv_freq", inv_freq().view(1, HD // 2), persistent=False) |
| self.register_buffer("steps", torch.arange(UPS, dtype=torch.float32).view(UPS, 1), persistent=False) |
| i = torch.arange(UPS)[:, None] |
| j = torch.arange(MIMI_CONTEXT + UPS)[None, :] |
| band = torch.where((j >= i) & (j <= i + MIMI_CONTEXT), 0.0, NEG) |
| self.register_buffer("band", band.view(1, 1, UPS, MIMI_CONTEXT + UPS), persistent=False) |
| self.register_buffer("ones_seen", torch.ones(1, 1, 1, UPS), persistent=False) |
|
|
| def _attn(self, L, x, kc, vc, seen, cos, sin): |
| T = UPS |
| h = layer_norm(L.norm1, x, 1e-5) |
| qkv = lin(L.in_proj, h) |
| q = qkv[..., :D].reshape(1, T, H, HD) |
| k = qkv[..., D : 2 * D].reshape(1, T, H, HD) |
| v = qkv[..., 2 * D :].reshape(1, T, H, HD) |
| rh = lambda t: torch.cat([-t[..., HD // 2 :], t[..., : HD // 2]], -1) |
| q = q * cos + rh(q) * sin |
| k = k * cos + rh(k) * sin |
| q = q.permute(0, 2, 1, 3) |
| k = k.permute(0, 2, 1, 3) |
| v = v.permute(0, 2, 1, 3) |
| kk = torch.cat([kc, k], 2) |
| vv = torch.cat([vc, v], 2) |
| seen_all = torch.cat([seen, self.ones_seen], 3) |
| mask = self.band + (seen_all - 1.0) * 1.0e30 |
| sc = torch.bmm(q.reshape(H, T, HD), kk.reshape(H, MIMI_CONTEXT + T, HD).transpose(1, 2)) * (HD**-0.5) |
| pr = torch.softmax(sc.reshape(1, H, T, MIMI_CONTEXT + T) + mask, -1) |
| o = torch.bmm(pr.reshape(H, T, MIMI_CONTEXT + T), vv.reshape(H, MIMI_CONTEXT + T, HD)) |
| o = o.reshape(1, H, T, HD).permute(0, 2, 1, 3).reshape(1, T, D) |
| x = x + L.ls1 * lin(L.out_proj, o) |
| h = layer_norm(L.norm2, x, 1e-5) |
| x = x + L.ls2 * lin(L.linear2, gelu(lin(L.linear1, h))) |
| return x, kk[:, :, UPS:], vv[:, :, UPS:], seen_all[..., UPS:] |
|
|
| def _res(self, blk, x, prev): |
| v, prev = _conv_step(blk.block[1], elu(x), prev) |
| v = F.conv1d(elu(v), blk.block[3].conv.weight, blk.block[3].conv.bias) |
| return x + v, prev |
|
|
| def forward(self, latent, up_p, k0, v0, k1, v1, seen, off, c0, u1, r1, u2, r2, u3, r3, cf): |
| x = latent * self.emb_std + self.emb_mean |
| x = lin(self.qproj, x).reshape(1, D, 1) |
| x, up_p = _convtr_step(x, self.up_w, None, up_p, UPS, groups=D) |
| pos = self.steps + off |
| ang = pos * self.inv_freq |
| cos = torch.cat([torch.cos(ang), torch.cos(ang)], -1).reshape(1, UPS, 1, HD) |
| sin = torch.cat([torch.sin(ang), torch.sin(ang)], -1).reshape(1, UPS, 1, HD) |
| x = x.permute(0, 2, 1) |
| x, k0, v0, seen0 = self._attn(self.layers[0], x, k0, v0, seen, cos, sin) |
| x, k1, v1, _ = self._attn(self.layers[1], x, k1, v1, seen, cos, sin) |
| x = x.permute(0, 2, 1) |
| m = self.dec |
| x, c0 = _conv_step(m[0], x, c0) |
| x, u1 = _convtr_step(elu(x), m[2].convtr.weight, m[2].convtr.bias, u1, 6) |
| x, r1 = self._res(m[3], x, r1) |
| x, u2 = _convtr_step(elu(x), m[5].convtr.weight, m[5].convtr.bias, u2, 5) |
| x, r2 = self._res(m[6], x, r2) |
| x, u3 = _convtr_step(elu(x), m[8].convtr.weight, m[8].convtr.bias, u3, 4) |
| x, r3 = self._res(m[9], x, r3) |
| pcm, cf = _conv_step(m[11], elu(x), cf) |
| return pcm, up_p, k0, v0, k1, v1, seen0, off + float(UPS), c0, u1, r1, u2, r2, u3, r3, cf |
|
|
