codec.py

"""Mini neural audio codec — encoder + residual VQ + decoder.

Inspired by SoundStream / EnCodec but tractable: trains on Apple Silicon (MPS)
in a few minutes and produces real waveform reconstructions.

Frame rate: 24000 / 320 = 75 Hz. With 2 codebooks of size 512 (9 bits each),
nominal bitrate ~= 75 * 9 * 2 = 1350 bits / sec.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


SAMPLE_RATE = 24000
HOP = 320  # total downsampling factor of the encoder


class ResBlock(nn.Module):
    def __init__(self, ch, dilation):
        super().__init__()
        self.net = nn.Sequential(
            nn.GELU(),
            nn.Conv1d(ch, ch, kernel_size=3, padding=dilation, dilation=dilation),
            nn.GELU(),
            nn.Conv1d(ch, ch, kernel_size=1),
        )

    def forward(self, x):
        return x + self.net(x)


class Encoder(nn.Module):
    def __init__(self, ch=64, latent=128):
        super().__init__()
        # 4 downsampling stages: 2x, 4x, 5x, 8x => total 320x
        strides = [2, 4, 5, 8]
        layers = [nn.Conv1d(1, ch, kernel_size=7, padding=3)]
        c = ch
        for s in strides:
            layers += [
                ResBlock(c, dilation=1),
                ResBlock(c, dilation=3),
                nn.GELU(),
                nn.Conv1d(c, c * 2, kernel_size=2 * s, stride=s, padding=s // 2),
            ]
            c *= 2
        layers += [nn.GELU(), nn.Conv1d(c, latent, kernel_size=3, padding=1)]
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class Decoder(nn.Module):
    def __init__(self, ch=64, latent=128):
        super().__init__()
        strides = [8, 5, 4, 2]
        c = ch * (2 ** 4)
        layers = [nn.Conv1d(latent, c, kernel_size=3, padding=1)]
        for s in strides:
            layers += [
                nn.GELU(),
                nn.ConvTranspose1d(c, c // 2, kernel_size=2 * s, stride=s, padding=s // 2),
                ResBlock(c // 2, dilation=1),
                ResBlock(c // 2, dilation=3),
            ]
            c //= 2
        layers += [nn.GELU(), nn.Conv1d(c, 1, kernel_size=7, padding=3), nn.Tanh()]
        self.net = nn.Sequential(*layers)

    def forward(self, z):
        return self.net(z)


class VectorQuantizer(nn.Module):
    """VQ codebook with EMA updates + dead-entry restart.

    EMA codebook updates (van den Oord 2017, EnCodec/DAC) keep the codebook
    stable as the encoder distribution drifts during training. Periodic
    restart of dead entries (reinitialising them from random encoder outputs)
    prevents the "codebook collapse" failure mode where only a handful of
    entries ever get used.
    """

    def __init__(self, dim, codebook_size=512, decay=0.99, eps=1e-5,
                 restart_threshold=1.0):
        super().__init__()
        embed = torch.randn(codebook_size, dim) * 0.1
        self.register_buffer('codebook', embed)
        self.register_buffer('cluster_size', torch.ones(codebook_size) * eps)
        self.register_buffer('embed_avg', embed.clone() * eps)
        self.decay = decay
        self.eps = eps
        self.codebook_size = codebook_size
        self.restart_threshold = restart_threshold

    def forward(self, z):
        # z: [B, D, T] -> [B*T, D]
        B, D, T = z.shape
        flat = z.permute(0, 2, 1).reshape(-1, D)
        dist = (flat.pow(2).sum(1, keepdim=True)
                - 2 * flat @ self.codebook.t()
                + self.codebook.pow(2).sum(1))
        idx = dist.argmin(dim=1)
        q = self.codebook[idx].view(B, T, D).permute(0, 2, 1)

        if self.training:
            onehot = F.one_hot(idx, self.codebook_size).type(flat.dtype)
            cluster_new = onehot.sum(0)
            self.cluster_size.data.mul_(self.decay).add_(cluster_new, alpha=1 - self.decay)
            embed_sum = onehot.t() @ flat
            self.embed_avg.data.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay)
            n = self.cluster_size.sum()
            smoothed = (self.cluster_size + self.eps) / (n + self.codebook_size * self.eps) * n
            self.codebook.data.copy_(self.embed_avg / smoothed.unsqueeze(1))

            # Dead-entry restart: any code whose cluster size has decayed below
            # restart_threshold gets replaced with a randomly-chosen encoder vector
            # from this batch. This is what gets us off of codebook collapse.
            dead = self.cluster_size < self.restart_threshold
            n_dead = int(dead.sum().item())
            if n_dead > 0 and flat.shape[0] > 0:
                rand_idx = torch.randint(0, flat.shape[0], (n_dead,), device=flat.device)
                self.codebook.data[dead] = flat[rand_idx].detach()
                self.embed_avg.data[dead] = flat[rand_idx].detach()
                self.cluster_size.data[dead] = 1.0

        commit_loss = F.mse_loss(z, q.detach())
        codebook_loss = torch.zeros((), device=z.device)
        q_st = z + (q - z).detach()
        return q_st, idx.view(B, T), commit_loss, codebook_loss


class ResidualVQ(nn.Module):
    """Residual VQ with optional quantizer dropout (a la SoundStream/EnCodec).

    During training we randomly pick how many codebooks to use per forward pass,
    so at inference time you can choose any active count <= N and the decoder
    still produces a reasonable reconstruction at the corresponding bitrate.
    """

    def __init__(self, dim, n_codebooks=4, codebook_size=512):
        super().__init__()
        self.layers = nn.ModuleList([
            VectorQuantizer(dim, codebook_size) for _ in range(n_codebooks)
        ])
        self.n_codebooks = n_codebooks
        self.codebook_size = codebook_size

    def forward(self, z, n_active=None):
        if n_active is None:
            n_active = self.n_codebooks
        residual = z
        out = torch.zeros_like(z)
        commit = 0.0
        codebook = 0.0
        all_idx = []
        for i, vq in enumerate(self.layers):
            if i >= n_active:
                break
            q, idx, cl, cbl = vq(residual)
            out = out + q
            residual = residual - q.detach()
            commit = commit + cl
            codebook = codebook + cbl
            all_idx.append(idx)
        return out, torch.stack(all_idx, dim=1), commit, codebook

    def bitrate(self, frame_rate, n_active=None):
        n = n_active if n_active is not None else self.n_codebooks
        bits_per_frame = n * math.log2(self.codebook_size)
        return frame_rate * bits_per_frame


class Codec(nn.Module):
    def __init__(self, ch=32, latent=128, n_codebooks=4, codebook_size=512):
        super().__init__()
        self.encoder = Encoder(ch=ch, latent=latent)
        self.rvq = ResidualVQ(latent, n_codebooks=n_codebooks, codebook_size=codebook_size)
        self.decoder = Decoder(ch=ch, latent=latent)

    def forward(self, x, n_active=None):
        z = self.encoder(x)
        zq, idx, commit, codebook = self.rvq(z, n_active=n_active)
        y = self.decoder(zq)
        # match length (decoder may differ by a few samples due to padding)
        if y.shape[-1] != x.shape[-1]:
            T = min(y.shape[-1], x.shape[-1])
            y = y[..., :T]
            x = x[..., :T]
        return y, idx, commit, codebook