| """Audio tokenizer: encoder hidden states (B, T, n_layers, D) -> (B, n_global + T / frame_stack, d_llm) LLM inputs.""" |
| import math |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class AudioTokenizer(nn.Module): |
| def __init__(self, d_llm, layer_mix, d_global=2048, d_local=2048, n_global=16, d_q=512, frame_stack=2, |
| n_frames=104): |
| super().__init__() |
| self.mix_logits = nn.Parameter(torch.zeros(layer_mix["n_mix"], layer_mix["n_layers"])) |
| self.n_global, self.d_q = n_global, d_q |
| self.frame_stack, self.n_frames = frame_stack, n_frames |
| self.n_local = math.ceil(n_frames / frame_stack) |
| self.g_mlp = nn.Sequential(nn.LayerNorm(d_global), nn.Linear(d_global, 1024), nn.GELU(), |
| nn.Linear(1024, n_global * d_q)) |
| self.g_proj = nn.Sequential(nn.LayerNorm(d_q), nn.Linear(d_q, d_llm)) |
| self.f_proj = nn.Sequential(nn.LayerNorm(d_local * frame_stack), nn.Linear(d_local * frame_stack, 1024), |
| nn.GELU(), nn.Linear(1024, d_llm)) |
|
|
| @property |
| def n_tokens(self): |
| return self.n_global + self.n_local |
|
|
| def forward(self, hs): |
| w = torch.softmax(self.mix_logits.float(), -1).to(hs.dtype) |
| loc = torch.einsum("btld,ml->btmd", hs, w).flatten(2) |
| g = loc.mean(1) |
| glob = self.g_proj(self.g_mlp(g).view(g.size(0), self.n_global, self.d_q)) |
| assert loc.size(1) == self.n_frames, (loc.shape, self.n_frames) |
| k = self.frame_stack |
| pad = self.n_local * k - loc.size(1) |
| x = nn.functional.pad(loc, (0, 0, 0, pad)) if pad else loc |
| return torch.cat([glob, self.f_proj(x.reshape(loc.size(0), self.n_local, k * loc.size(2)))], 1) |
|
|
