| """Streaming Mimi decoder for one stream: feed any number of frames at a time, get exactly their audio back. |
| |
| Arithmetic of upstream serve/fast_codec.StaticMimi (itself a port of the verified streaming reference) without CUDA |
| graphs or seat slots: causal-convolution history, transposed-convolution overlap tails (bias removed from the stored |
| tail), and sliding-window attention with absolute RoPE positions (a query at position p sees keys with |
| p - sliding_window < key_pos <= p). HF sub-modules are only called for their weights. |
| """ |
| import torch |
| import torch.nn.functional as F |
| from transformers.models.mimi.modeling_mimi import (MimiConv1d, MimiConvTranspose1d, MimiResnetBlock, |
| apply_rotary_pos_emb, repeat_kv) |
|
|
|
|
| class StreamingMimi: |
| def __init__(self, mimi): |
| cfg = mimi.config |
| if not cfg.use_causal_conv or cfg.pad_mode != "constant" or cfg.trim_right_ratio != 1: |
| raise ValueError("streaming decode requires causal constant padding and trim_right_ratio=1") |
| self.m = mimi |
| self.window = int(cfg.sliding_window) |
| attn = mimi.decoder_transformer.layers[0].self_attn |
| self.kv_groups, self.head_dim = attn.num_key_value_groups, attn.head_dim |
| rotary = attn.rotary_emb |
| self._inv_freq = rotary.inv_freq.detach().float().clone() |
| self._rope_scale = float(rotary.attention_scaling) |
| self.reset() |
|
|
| def reset(self): |
| self.state = {} |
| self.kv = [None] * len(self.m.decoder_transformer.layers) |
| self.tokens = 0 |
|
|
| def _layer(self, module, x, key): |
| if isinstance(module, MimiConv1d): |
| if module.conv.stride != (1,): |
| raise ValueError("only stride-one decoder convolutions are supported") |
| pad = (module.conv.kernel_size[0] - 1) * module.conv.dilation[0] |
| if pad: |
| prev = self.state.get(key) |
| if prev is None: |
| prev = x.new_zeros(x.shape[0], x.shape[1], pad) |
| x = torch.cat((prev, x), dim=-1) |
| self.state[key] = x[..., -pad:] |
| return module.conv(x) |
| if isinstance(module, MimiConvTranspose1d): |
| y = module.conv(x) |
| produced = x.shape[-1] * module.conv.stride[0] |
| if module.padding_right: |
| tail = y[..., produced:] |
| if module.conv.bias is not None: |
| tail = tail - module.conv.bias[None, :, None] |
| prev = self.state.get(key) |
| y = y[..., :produced].clone() |
| if prev is not None: |
| y[..., :prev.shape[-1]] += prev |
| self.state[key] = tail |
| return y |
| return y[..., :produced] |
| if isinstance(module, MimiResnetBlock): |
| shortcut = self._layer(module.shortcut, x, key + ".shortcut") |
| for i, child in enumerate(module.block): |
| x = self._layer(child, x, key + f".block.{i}") |
| return shortcut + x |
| return module(x) |
|
|
| @staticmethod |
| def _rvq(rvq, codes): |
| total = None |
| for i in range(codes.shape[1]): |
| e = F.embedding(codes[:, i], rvq.layers[i].codebook.embed).permute(0, 2, 1) |
| total = e if total is None else total + e |
| return rvq.output_proj(total) if rvq.output_proj is not None else total |
|
|
| def _quantizer(self, codes): |
| q = self.m.quantizer |
| n = q.num_semantic_quantizers |
| out = self._rvq(q.semantic_residual_vector_quantizer, codes[:, :n]) |
| if codes.shape[1] > n: |
| out = out + self._rvq(q.acoustic_residual_vector_quantizer, codes[:, n:]) |
| return out |
|
|
| def _transformer(self, h): |
| B, T, _ = h.shape |
| pos = torch.arange(self.tokens, self.tokens + T, device=h.device) |
| freqs = pos[:, None].float() * self._inv_freq.to(h.device)[None] |
| emb = torch.cat((freqs, freqs), dim=-1) |
| cos = (emb.cos() * self._rope_scale).to(h.dtype)[None] |
| sin = (emb.sin() * self._rope_scale).to(h.dtype)[None] |
| n_past = 0 if self.kv[0] is None else self.kv[0][0].shape[2] |
| key_pos = torch.arange(self.tokens - n_past, self.tokens + T, device=h.device) |
| mask = (key_pos[None] <= pos[:, None]) & (key_pos[None] > pos[:, None] - self.window) |
| for li, layer in enumerate(self.m.decoder_transformer.layers): |
| attn = layer.self_attn |
| x = layer.input_layernorm(h) |
| q = attn.q_proj(x).view(B, T, -1, self.head_dim).transpose(1, 2) |
| k = attn.k_proj(x).view(B, T, -1, self.head_dim).transpose(1, 2) |
| v = attn.v_proj(x).view(B, T, -1, self.head_dim).transpose(1, 2) |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
| if self.kv[li] is not None: |
| k = torch.cat((self.kv[li][0], k), dim=2) |
| v = torch.cat((self.kv[li][1], v), dim=2) |
| keep = self.window - 1 |
| self.kv[li] = (k[:, :, -keep:], v[:, :, -keep:]) |
| a = F.scaled_dot_product_attention(q, repeat_kv(k, self.kv_groups), repeat_kv(v, self.kv_groups), |
| attn_mask=mask[None, None]) |
| a = attn.o_proj(a.transpose(1, 2).reshape(B, T, -1)) |
| h = h + layer.self_attn_layer_scale(a) |
| h = h + layer.mlp_layer_scale(layer.mlp(layer.post_attention_layernorm(h))) |
| self.tokens += T |
| return h |
|
|
| @torch.inference_mode() |
| def decode(self, codes): |
| """codes: (32, T) long (one stream, T frames) -> float32 numpy (T * 1920,) PCM in [-1, 1].""" |
| x = self._quantizer(codes[None].long()) |
| x = self._layer(self.m.upsample, x, "upsample") |
| x = self._transformer(x.transpose(1, 2)).transpose(1, 2) |
| for i, layer in enumerate(self.m.decoder.layers): |
| x = self._layer(layer, x, f"decoder.{i}") |
| return x[0, 0].float().clamp(-1, 1).cpu().numpy() |
|
|
