Upload folder using huggingface_hub

README.md

@@ -1,3 +1 @@
----
-license: mit
----
+Soundstream implementation. Sample rate $16000$.

config.json

@@ -0,0 +1,6 @@
+{
+  "channels": 32,
+  "codebook_size": 1024,
+  "latent_dim": 512,
+  "num_quantizers": 8
+}

model.py

@@ -0,0 +1,386 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class CausalConv1d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        self.left_padding = dilation * (kernel_size - 1)
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+        )
+
+    def forward(self, x):
+        x = F.pad(x, (self.left_padding, 0))
+        return self.conv(x)
+
+
+class CausalConvTranspose1d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+    ):
+        super().__init__()
+        self.stride = stride
+        self.conv_transpose = nn.ConvTranspose1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+        )
+
+    def forward(self, x):
+        target_length = x.shape[-1] * self.stride
+        x = self.conv_transpose(x)
+        return x[..., :target_length]
+
+
+class ResidualUnit(nn.Module):
+
+    def __init__(self, channels: int, dilation: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.ELU(),
+            CausalConv1d(kernel_size=7, in_channels=channels, out_channels=channels, dilation=dilation),
+            nn.ELU(),
+            nn.Conv1d(kernel_size=1, in_channels=channels, out_channels=channels)
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class EncoderBlock(nn.Module):
+
+    def __init__(self, channels: int, s: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(channels=channels // 2, dilation=1),
+            ResidualUnit(channels=channels // 2, dilation=3),
+            ResidualUnit(channels=channels // 2, dilation=9),
+            CausalConv1d(kernel_size=2 * s, in_channels=channels // 2, out_channels=channels, stride=s)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+
+    def __init__(self, channels: int, s: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            CausalConvTranspose1d(kernel_size=2 * s, in_channels=channels, out_channels=channels // 2, stride=s),
+            ResidualUnit(channels=channels // 2, dilation=1),
+            ResidualUnit(channels=channels // 2, dilation=3),
+            ResidualUnit(channels=channels // 2, dilation=9)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+
+    def __init__(self, channels: int = 16, dim: int = 512):
+        super().__init__()
+        # NB: attribute name "encoder" matches training checkpoint keys
+        self.encoder = nn.Sequential(
+            CausalConv1d(kernel_size=7, in_channels=1, out_channels=channels),
+            EncoderBlock(channels=2 * channels, s=2),
+            EncoderBlock(channels=4 * channels, s=4),
+            EncoderBlock(channels=8 * channels, s=5),
+            EncoderBlock(channels=16 * channels, s=5),
+            CausalConv1d(kernel_size=3, in_channels=16 * channels, out_channels=dim)
+        )
+
+    def forward(self, audio):
+        return self.encoder(audio)
+
+
+class Decoder(nn.Module):
+
+    def __init__(self, channels: int = 16, dim: int = 512):
+        super().__init__()
+        # NB: attribute name "decoder" matches training checkpoint keys
+        self.decoder = nn.Sequential(
+            CausalConv1d(kernel_size=7, in_channels=dim, out_channels=16 * channels),
+            DecoderBlock(channels=16 * channels, s=5),
+            DecoderBlock(channels=8 * channels, s=5),
+            DecoderBlock(channels=4 * channels, s=4),
+            DecoderBlock(channels=2 * channels, s=2),
+            CausalConv1d(kernel_size=7, in_channels=channels, out_channels=1)
+        )
+
+    def forward(self, quantized):
+        return self.decoder(quantized)
+
+
+@torch.no_grad()
+def _k_means(vectors, num_clusters, num_iters):
+    n = vectors.size(0)
+    device = vectors.device
+
+    if n >= num_clusters:
+        init_indices = torch.randperm(n, device=device)[:num_clusters]
+    else:
+        init_indices = torch.randint(0, n, (num_clusters,), device=device)
+
+    centroids = vectors[init_indices].clone()
+
+    for _ in range(num_iters):
+        dists = (
+            vectors.pow(2).sum(1, keepdim=True)
+            - 2 * vectors @ centroids.t()
+            + centroids.pow(2).sum(1)
+        )
+        assignments = dists.argmin(1)
+
+        counts = torch.bincount(assignments, minlength=num_clusters).to(vectors.dtype)
+        sums = torch.zeros_like(centroids)
+        sums.index_add_(0, assignments, vectors)
+
+        non_empty = counts > 0
+        if non_empty.any():
+            centroids[non_empty] = sums[non_empty] / counts[non_empty].unsqueeze(1)
+
+        empty = ~non_empty
+        if empty.any():
+            centroids[empty] = vectors[torch.randint(0, n, (int(empty.sum()),), device=device)]
+
+    dists = (
+        vectors.pow(2).sum(1, keepdim=True)
+        - 2 * vectors @ centroids.t()
+        + centroids.pow(2).sum(1)
+    )
+    counts = torch.bincount(dists.argmin(1), minlength=num_clusters).to(vectors.dtype)
+    return centroids, counts
+
+
+class VectorQuantizer(nn.Module):
+
+    def __init__(
+        self,
+        codebook_size: int,
+        latent_dim: int,
+        decay: float = 0.99,
+        dead_code_threshold: float = 2.0,
+        kmeans_iters: int = 50,
+    ):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.latent_dim = latent_dim
+        self.decay = decay
+        self.dead_code_threshold = dead_code_threshold
+        self.kmeans_iters = kmeans_iters
+        self.eps = 1e-8
+
+        self.register_buffer("initialized", torch.tensor(False, dtype=torch.bool))
+        self.register_buffer("embedding", torch.randn(codebook_size, latent_dim))
+        self.register_buffer("ema_n", torch.zeros(codebook_size))
+        self.register_buffer("ema_s", torch.zeros(codebook_size, latent_dim))
+
+    def forward(self, latent):
+        B, D, T = latent.shape
+        flat = latent.transpose(1, 2).reshape(-1, D)
+
+        if self.training and not self.initialized:
+            self._init_codebook(flat)
+            self.initialized.fill_(True)
+
+        idx, quantized = self._nearest(flat)
+
+        if self.training:
+            self._update_ema(flat, idx)
+            self._replace_dead_codes(flat)
+
+        commit_loss = F.mse_loss(flat, quantized.detach())
+        quantized_ste = flat + (quantized - flat).detach()
+
+        return {
+            "quantized": quantized_ste.reshape(B, T, D).transpose(1, 2).contiguous(),
+            "indices": idx.reshape(B, T),
+            "commitment_loss": commit_loss,
+        }
+
+    def _nearest(self, flat):
+        dists = (
+            flat.pow(2).sum(1, keepdim=True)
+            - 2 * flat @ self.embedding.t()
+            + self.embedding.pow(2).sum(1)
+        )
+        idx = dists.argmin(1)
+        return idx, self.embedding[idx]
+
+    @torch.no_grad()
+    def _init_codebook(self, flat):
+        centroids, counts = _k_means(flat, self.codebook_size, self.kmeans_iters)
+        counts = counts.clamp_min(1.0)
+        w = counts / counts.mean()
+        self.embedding.copy_(centroids)
+        self.ema_n.copy_(w)
+        self.ema_s.copy_(centroids * w.unsqueeze(1))
+
+    @torch.no_grad()
+    def _update_ema(self, flat, indices):
+        bins = torch.bincount(indices, minlength=self.codebook_size).to(flat.dtype)
+        sums = torch.zeros_like(self.ema_s)
+        sums.index_add_(0, indices, flat)
+        self.ema_n.mul_(self.decay).add_(bins, alpha=1 - self.decay)
+        self.ema_s.mul_(self.decay).add_(sums, alpha=1 - self.decay)
+        self.embedding.copy_(self.ema_s / (self.ema_n.unsqueeze(1) + self.eps))
+
+    @torch.no_grad()
+    def _replace_dead_codes(self, flat):
+        dead = self.ema_n < self.dead_code_threshold
+        if not dead.any():
+            return
+        n = int(dead.sum())
+        picks = flat[torch.randint(0, flat.size(0), (n,), device=flat.device)]
+        self.embedding[dead] = picks
+        self.ema_s[dead] = picks
+        self.ema_n[dead] = self.dead_code_threshold
+
+    @torch.no_grad()
+    def quantize(self, latent):
+        B, D, T = latent.shape
+        flat = latent.transpose(1, 2).reshape(-1, D)
+        idx, quantized = self._nearest(flat)
+        return {
+            "quantized": quantized.reshape(B, T, D).transpose(1, 2).contiguous(),
+            "indices": idx.reshape(B, T),
+        }
+
+    def decode_indices(self, indices):
+        return self.embedding[indices].transpose(1, 2).contiguous()
+
+
+class ResidualVectorQuantizer(nn.Module):
+
+    def __init__(
+        self,
+        latent_dim: int,
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+    ):
+        super().__init__()
+        self.num_quantizers = num_quantizers
+        self.quantizers = nn.ModuleList(
+            VectorQuantizer(codebook_size, latent_dim) for _ in range(num_quantizers)
+        )
+
+    def forward(self, latent):
+        residual = latent
+        quantized = torch.zeros_like(latent)
+        total_commit = latent.new_zeros(())
+        all_indices = []
+
+        for vq in self.quantizers:
+            out = vq(residual)
+            quantized = quantized + out["quantized"]
+            residual = residual - out["quantized"].detach()
+            total_commit = total_commit + out["commitment_loss"]
+            all_indices.append(out["indices"])
+
+        return {
+            "quantized": quantized,
+            "indices": torch.stack(all_indices, dim=1),
+            "commitment_loss": total_commit,
+        }
+
+    @torch.no_grad()
+    def encode(self, latent):
+        residual = latent
+        all_indices = []
+        for vq in self.quantizers:
+            out = vq.quantize(residual)
+            all_indices.append(out["indices"])
+            residual = residual - out["quantized"]
+        return torch.stack(all_indices, dim=1)
+
+    @torch.no_grad()
+    def decode(self, indices):
+        quantized = None
+        for i, vq in enumerate(self.quantizers):
+            stage = vq.decode_indices(indices[:, i])
+            quantized = stage if quantized is None else quantized + stage
+        return quantized
+
+
+class SoundStreamCodec(
+    nn.Module,
+    PyTorchModelHubMixin,
+    library_name="soundstream-impl",
+    license="mit",
+):
+    def __init__(
+        self,
+        channels: int = 32,
+        latent_dim: int = 512,
+        codebook_size: int = 1024,
+        num_quantizers: int = 8,
+    ):
+        super().__init__()
+        self.strides = (2, 4, 5, 5)
+        self.downsampling_factor = 1
+        for s in self.strides:
+            self.downsampling_factor *= s
+
+        self.encoder = Encoder(channels=channels, dim=latent_dim)
+        self.quantizer = ResidualVectorQuantizer(
+            latent_dim=latent_dim,
+            codebook_size=codebook_size,
+            num_quantizers=num_quantizers,
+        )
+        self.decoder = Decoder(channels=channels, dim=latent_dim)
+
+    def forward(self, audio, **kwargs):
+        original_length = audio.size(-1)
+        audio = self._pad_to_stride(audio)
+
+        latent = self.encoder(audio)
+        q_out = self.quantizer(latent)
+        reconstructed = self.decoder(q_out["quantized"])
+        reconstructed = reconstructed[..., :original_length]
+
+        return {
+            "reconstructed_audio": reconstructed,
+            "latent": latent,
+            **q_out,
+        }
+
+    @torch.no_grad()
+    def encode(self, audio):
+        audio = self._pad_to_stride(audio)
+        return self.quantizer.encode(self.encoder(audio))
+
+    @torch.no_grad()
+    def decode(self, indices, original_length=None):
+        out = self.decoder(self.quantizer.decode(indices))
+        if original_length is not None:
+            out = out[..., :original_length]
+        return out
+
+    def _pad_to_stride(self, audio):
+        remainder = audio.size(-1) % self.downsampling_factor
+        if remainder == 0:
+            return audio
+        return F.pad(audio, (0, self.downsampling_factor - remainder), mode="replicate")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eaf8b7aa910c717a05558fd1bf214823564331504c29a9cf2cf50fc3bf29e452
+size 74533004