"""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): # The extra terminal frame supplies overlap at end of stream. During # streaming the next latent is needed to finish the preceding hop. 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])