| """Paper-inspired small speech codec; not the authors' TinyCall implementation. |
| |
| Architecture choices missing from arXiv:2609.25379 are explicit here. The |
| encoder uses causal residual convolutions with a total stride of 320. Spectral |
| synthesis uses a 20 ms overlap and therefore has one frame of output delay. |
| """ |
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| class CausalConv(nn.Conv1d): |
| def forward(self, x): |
| stride = self.stride[0] |
| history = self.dilation[0] * (self.kernel_size[0] - 1) + 1 - stride |
| return super().forward(F.pad(x, (history, (-x.shape[-1]) % stride))) |
|
|
|
|
| class Residual(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.net = nn.Sequential(nn.ELU(), CausalConv(channels, channels // 2, 3), |
| nn.ELU(), CausalConv(channels // 2, channels, 1)) |
|
|
| def forward(self, x): |
| return x + self.net(x) |
|
|
|
|
| class TinyCallEncoder(nn.Module): |
| sample_rate = 16000 |
| hop = 320 |
| dimension = 64 |
|
|
| def __init__(self): |
| super().__init__() |
| layers = [CausalConv(1, 8, 7)] |
| channels = 8 |
| for stride in (2, 4, 5, 8): |
| layers.extend([Residual(channels), nn.ELU(), |
| CausalConv(channels, channels * 2, 2 * stride, |
| stride=stride)]) |
| channels *= 2 |
| layers.extend([nn.ELU(), CausalConv(channels, self.dimension, 3)]) |
| self.net = nn.Sequential(*layers) |
|
|
| def forward(self, audio): |
| if audio.ndim != 2 or audio.shape[-1] == 0: |
| raise ValueError("Expected nonempty [batch,samples] mono audio") |
| return self.net(audio[:, None]) |
|
|
|
|
| class ResidualFSQ(nn.Module): |
| """Four six-coordinate scalar stages, 46 bits/frame, with STE training. |
| |
| Projection widths, scale initialization, and bounded-grid formula are our |
| implementation choices. This does not assert the paper's unpublished RFSQ |
| implementation or its normalization has been reproduced exactly. |
| """ |
| bits_per_frame = 46 |
|
|
| def __init__(self, dimension=64): |
| super().__init__() |
| self.down = nn.ModuleList(nn.Conv1d(dimension, 6, 1) for _ in range(4)) |
| self.up = nn.ModuleList(nn.Conv1d(6, dimension, 1) for _ in range(4)) |
| self.log_scales = nn.Parameter(torch.tensor([0., -0.6931, -1.3863, -2.0794])) |
| self.register_buffer("levels", torch.tensor([[8, 8, 8, 4, 4, 4]] * 2 + |
| [[4, 4, 2, 2, 2, 2]] * 2)) |
|
|
| def forward(self, x): |
| residual, result, codes = x, torch.zeros_like(x), [] |
| for stage, (down, up) in enumerate(zip(self.down, self.up)): |
| scale = self.log_scales[stage].exp() |
| levels = self.levels[stage][None, :, None].to(x.dtype) |
| bounded = down(residual / scale).tanh() |
| indices = ((bounded + 1) * (levels - 1) / 2).round().long() |
| discrete = indices.to(x.dtype) * 2 / (levels - 1) - 1 |
| values = bounded + (discrete - bounded).detach() |
| reconstructed = scale * up(values) |
| residual = residual - reconstructed |
| result = result + reconstructed |
| codes.append(indices) |
| return result, torch.stack(codes, dim=1) |
|
|
| def decode(self, codes): |
| if codes.ndim != 4 or codes.shape[1:3] != (4, 6): |
| raise ValueError("Expected [batch,4,6,frames] scalar indices") |
| outputs = [] |
| for stage, up in enumerate(self.up): |
| indices = codes[:, stage] |
| levels = self.levels[stage][None, :, None] |
| if (indices < 0).any() or (indices >= levels).any(): |
| raise ValueError("Scalar index outside its stage's grid") |
| values = indices.to(up.weight.dtype) * 2 / (levels - 1) - 1 |
| outputs.append(self.log_scales[stage].exp() * up(values)) |
| return torch.stack(outputs).sum(0) |
|
|
|
|
| class SpectralBlock(nn.Module): |
| def __init__(self, channels=128): |
| super().__init__() |
| self.depthwise = CausalConv(channels, channels, 7, groups=channels) |
| self.norm = nn.LayerNorm(channels) |
| self.up = nn.Linear(channels, 3 * channels) |
| self.down = nn.Linear(3 * channels, channels) |
|
|
| def forward(self, x): |
| h = self.norm(self.depthwise(x).transpose(1, 2)) |
| return x + self.down(F.gelu(self.up(h))).transpose(1, 2) |
|
|
|
|
| class SpectralDecoder(nn.Module): |
| hop = 320 |
| n_fft = 640 |
|
|
| def __init__(self, dimension=64): |
| super().__init__() |
| self.input = nn.Conv1d(dimension, 128, 1) |
| self.blocks = nn.Sequential(*(SpectralBlock() for _ in range(5))) |
| self.spectrum = nn.Conv1d(128, 2 * (self.n_fft // 2 + 1), 1) |
| self.register_buffer("window", torch.hann_window(self.n_fft)) |
|
|
| def forward(self, features, samples=None): |
| |
| |
| features = F.pad(features, (0, 1)) |
| values = self.spectrum(self.blocks(self.input(features))) |
| log_magnitude, phase = values.chunk(2, dim=1) |
| magnitude = log_magnitude.clamp(-10, 5).exp() |
| phase = math.pi * phase.tanh() |
| spectrum = torch.complex(magnitude * phase.cos(), magnitude * phase.sin()) |
| available = (features.shape[-1] - 1) * self.hop |
| audio = torch.istft(spectrum, self.n_fft, self.hop, window=self.window, |
| center=True, length=available) |
| if samples is not None: |
| if samples > available: |
| raise ValueError("Requested output exceeds latent duration") |
| audio = audio[:, :samples] |
| return audio |
|
|
|
|
| class TinyCallCodec(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.encoder = TinyCallEncoder() |
| self.quantizer = ResidualFSQ(self.encoder.dimension) |
| self.decoder = SpectralDecoder(self.encoder.dimension) |
|
|
| def forward(self, audio, quantized_fraction=0.): |
| if not 0 <= quantized_fraction <= 1: |
| raise ValueError("quantized_fraction must be in [0,1]") |
| features = self.encoder(audio) |
| if quantized_fraction: |
| quantized, _ = self.quantizer(features) |
| features = features.lerp(quantized, quantized_fraction) |
| return self.decoder(features, audio.shape[-1]) |
|
|
| def accounting(self): |
| parts = {name: sum(p.numel() for p in module.parameters()) |
| for name, module in self.named_children()} |
| return {**parts, "total": sum(parts.values()), "sample_rate": 16000, |
| "feature_rate": 50, "feature_dimension": 64, |
| "quantized_bits_per_frame": 46, "quantized_bits_per_second": 2300, |
| "implementation": "independent paper-inspired reconstruction"} |
|
|
|
|
| class RecognitionHead(nn.Module): |
| def __init__(self, dimension=64, hidden=96, classes=29): |
| super().__init__() |
| self.gru = nn.GRU(dimension, hidden, batch_first=True) |
| self.output = nn.Linear(hidden, classes) |
| with torch.no_grad(): |
| self.output.bias[0] = -2. |
|
|
| def forward(self, features): |
| return self.output(self.gru(features)[0]) |
|
|
