File size: 5,995 Bytes
1eac567 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """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) # [(k, v)] post-RoPE, last window-1 tokens
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()
|