File size: 7,292 Bytes
b1882a9 | 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """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])
|