Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import weight_norm
+
+
+# Snake activation
+
+@torch.jit.script
+def snake(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    shape = x.shape # [B, C, T]
+    x = x.reshape(shape[0], shape[1], -1) # [B, C, T]
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape) # [B, C, T]
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1)) # [1, C, 1] one for each channel
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+# Weight-normalized convolutions
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class VQ(nn.Module):
+    def __init__(self, latent_ch, K=1024, codebook_dim=8):
+        super().__init__()
+        self.in_proj = nn.Linear(latent_ch, codebook_dim, bias=False)
+        self.out_proj = nn.Linear(codebook_dim, latent_ch, bias=False)
+        self.codebook = nn.Embedding(K, codebook_dim)
+
+    def forward(self, z: torch.tensor):
+        # z -> [N, C] 2d tensor flattened
+
+        # project to low-dim codebook space
+        z_e = self.in_proj(z)  # [N, codebook_dim]
+
+        # L2 normalise for cosine similarity matching
+        z_e_norm = F.normalize(z_e, dim=-1)              # [N, codebook_dim]
+        cb_norm = F.normalize(self.codebook.weight, dim=-1)  # [K, codebook_dim]
+
+        # euclidean distance between two unit vectors ~ cosine similarity
+        sim = z_e_norm @ cb_norm.t()  # [N, K]
+
+        # nearest codebook entry = highest similarity
+        indices = sim.max(dim=1)[1]  # [N]
+
+        # lookup normalised codebook entry
+        z_q_norm = cb_norm[indices]  # [N, codebook_dim]
+
+        # losses on normalised vectors
+        commitment_loss = F.mse_loss(z_e_norm, z_q_norm.detach())  # push encoder direction → codebook
+        codebook_loss = F.mse_loss(z_e_norm.detach(), z_q_norm)    # push codebook → encoder direction
+
+        # STE in normalised space
+        z_q_st = z_e_norm + (z_q_norm - z_e_norm).detach()
+
+        # project back to full latent space
+        z_q_out = self.out_proj(z_q_st)  # [N, latent_ch]
+
+        return z_q_out, indices, commitment_loss, codebook_loss
+
+
+class RVQ(nn.Module):
+    def __init__(self, num_levels, latent_ch, K=1024, codebook_dim=8):
+        super().__init__()
+        self.num_levels = num_levels
+        self.levels = nn.ModuleList([
+            VQ(latent_ch, K=K, codebook_dim=codebook_dim) for _ in range(num_levels)
+        ])
+
+    def forward(self, z):
+        # z -> [N, C] 2d flat vector
+        r = z # initilise residual with z for the first level
+        quantised_sum = torch.zeros_like(z)
+        all_indices = []
+        total_commitment_loss = 0
+        total_codebook_loss = 0
+
+        for level in self.levels:
+            z_q, indices, commitment_loss, codebook_loss = level(r)
+            r = r - z_q.detach()  # next level quantizes the error
+            quantised_sum = quantised_sum + z_q  # accumulate: z ≈ q1 + q2 + q3 + ...
+            all_indices.append(indices)
+            total_commitment_loss = total_commitment_loss + commitment_loss
+            total_codebook_loss = total_codebook_loss + codebook_loss
+
+        return quantised_sum, all_indices, total_commitment_loss, total_codebook_loss
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, ch, dilation=1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(ch), # [B, C, T]
+            WNConv1d(ch, ch, kernel_size=7, dilation=dilation, padding=3 * dilation), # [B, C, T] sk=7, padding=3 to keep same shape
+            Snake1d(ch), # [B, C, T]
+            WNConv1d(ch, ch, kernel_size=1), # [B, C, T]
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, stride):
+        super().__init__()
+        self.res1 = ResidualUnit(in_ch, dilation=1)
+        self.res2 = ResidualUnit(in_ch, dilation=3)
+        self.res3 = ResidualUnit(in_ch, dilation=9)
+        self.downsample = nn.Sequential(
+            Snake1d(in_ch),
+            WNConv1d(in_ch, out_ch, kernel_size=2 * stride, stride=stride, padding=stride // 2),
+        )
+
+    def forward(self, x):
+        x = self.res1(x)
+        x = self.res2(x)
+        x = self.res3(x)
+        x = self.downsample(x)
+        return x
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, stride):
+        super().__init__()
+        self.upsample = nn.Sequential(
+            Snake1d(in_ch),
+            WNConvTranspose1d(in_ch, out_ch, kernel_size=2 * stride, stride=stride, padding=stride // 2),
+        )
+        self.res1 = ResidualUnit(out_ch, dilation=1)
+        self.res2 = ResidualUnit(out_ch, dilation=3)
+        self.res3 = ResidualUnit(out_ch, dilation=9)
+
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.res1(x)
+        x = self.res2(x)
+        x = self.res3(x)
+        return x
+
+
+class RVQCodec(nn.Module):
+    def __init__(self, in_ch=1, latent_ch=32, K=1024, num_rvq_levels=1, codebook_dim=8):
+        super().__init__()
+        # Encoder - [B, 1, T] → [B, D, T/128]
+        # strides - 2 × 4 × 4 × 4 = 128x downsample
+        self.encoder = nn.Sequential(
+            WNConv1d(in_ch, 64, kernel_size=7, padding=3),              # [B, 64, T]
+            EncoderBlock(64, 128, stride=2),                             # [B, 128, T/2]
+            EncoderBlock(128, 256, stride=4),                            # [B, 256, T/8]
+            EncoderBlock(256, 512, stride=4),                            # [B, 512, T/32]
+            EncoderBlock(512, 512, stride=4),                            # [B, 512, T/128]
+            Snake1d(512),
+            WNConv1d(512, latent_ch, kernel_size=3, padding=1),          # [B, D, T/128]
+        )
+        # Decoder - [B, D, T/128] → [B, 1, T]
+        # strides - 4 × 4 × 4 × 2 = 128x upsample
+        self.decoder = nn.Sequential(
+            WNConv1d(latent_ch, 512, kernel_size=7, padding=3),          # [B, 512, T/128]
+            DecoderBlock(512, 512, stride=4),                            # [B, 512, T/32]
+            DecoderBlock(512, 256, stride=4),                            # [B, 256, T/8]
+            DecoderBlock(256, 128, stride=4),                            # [B, 128, T/2]
+            DecoderBlock(128, 64, stride=2),                             # [B, 64, T]
+            Snake1d(64),
+            WNConv1d(64, in_ch, kernel_size=7, padding=3),              # [B, 1, T]
+            nn.Tanh(),
+        )
+        self.rvq = RVQ(num_levels=num_rvq_levels, latent_ch=latent_ch, K=K, codebook_dim=codebook_dim)
+
+    def forward(self, x: torch.tensor):
+        # x -> [B, C=1, T]
+        z = self.encoder(x) # [B, D, T/128]
+
+        # flatten to 2d vector for applying rvq on channel dim
+        B, C, T_128 = z.shape
+        z_flat = z.permute(0, 2, 1).contiguous().view(B * T_128, C)
+
+        # vector quantize
+        z_q, all_indices, commitment_loss, codebook_loss = self.rvq(z_flat)
+
+        # reshape back
+        z_q = z_q.view(B, T_128, C).permute(0, 2, 1) # [B, C, T_128]
+
+        x_recon = self.decoder(z_q) # [B, C=1, T]
+
+        return x_recon, all_indices, commitment_loss, codebook_loss
+
+
+if __name__ == "__main__":
+    device = "cuda"
+    x = torch.randn(1, 1, 8192)
+
+    model = RVQCodec()
+    print(model)
+    print(f"params: {sum(p.numel() for p in model.parameters()):,}")
+
+    x = x.to(device)
+    model = model.to(device)
+
+    out, _, _, _ = model(x)
+    print(f"in: {x.shape} → out: {out.shape}")