Add neucodec package for self-contained inference

neucodec/__init__.py

@@ -0,0 +1 @@
+from .model import NeuCodec

neucodec/activations.py

@@ -0,0 +1,126 @@
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+#   LICENSE is in incl_licenses directory.
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+    """
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        """
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta ∶= x + 1/b * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x

neucodec/alias_free_torch/__init__.py

@@ -0,0 +1,6 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+from .filter import *
+from .resample import *
+from .act import *

neucodec/alias_free_torch/act.py

@@ -0,0 +1,30 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from .resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    # x: [B,C,T]
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+
+        return x

neucodec/alias_free_torch/filter.py

@@ -0,0 +1,99 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+    # This code is adopted from adefossez's julius.core.sinc under the MIT License
+    # https://adefossez.github.io/julius/julius/core.html
+    #   LICENSE is in incl_licenses directory.
+    def sinc(x: torch.Tensor):
+        """
+        Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+        __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+        """
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+#   LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # For kaiser window
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        # Normalize filter to have sum = 1, otherwise we will have a small leakage
+        # of the constant component in the input signal.
+        filter_ /= filter_.sum()
+        filter = filter_.view(1, 1, kernel_size)
+
+    return filter
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        # kernel_size should be even number for stylegan3 setup,
+        # in this implementation, odd number is also possible.
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    # input [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+        return out

neucodec/alias_free_torch/resample.py

@@ -0,0 +1,58 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .filter import LowPassFilter1d
+from .filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
+        self.register_buffer("filter", filter)
+
+    # x: [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
+
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        xx = self.lowpass(x)
+
+        return xx

neucodec/bs_roformer5.py

@@ -0,0 +1,118 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torchtune.modules import RotaryPositionalEmbeddings
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        r"""https://github.com/meta-llama/llama/blob/main/llama/model.py"""
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm_x = torch.mean(x**2, dim=-1, keepdim=True)
+        output = x * torch.rsqrt(norm_x + self.eps) * self.weight
+        return output
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+
+        self.fc1 = nn.Linear(dim, 4 * dim, bias=False)
+        self.silu = nn.SiLU()
+        self.fc2 = nn.Linear(4 * dim, dim, bias=False)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.silu(x)
+        x = self.fc2(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings
+    ):
+        super().__init__()
+
+        assert dim % n_heads == 0
+
+        self.n_heads = n_heads
+        self.dim = dim
+        self.rotary_embed = rotary_embed
+
+        self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
+        assert self.flash, "Must have flash attention."
+
+        self.c_attn = nn.Linear(dim, 3 * dim, bias=False)
+        self.c_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x: (b, t, h*d)
+
+        Constants:
+            b: batch_size
+            t: time steps
+            r: 3
+            h: heads_num
+            d: heads_dim
+        """
+        B, T, C = x.size()
+
+        q, k, v = rearrange(
+            self.c_attn(x), "b t (r h d) -> r b h t d", r=3, h=self.n_heads
+        )
+        # q, k, v: (b, h, t, d)
+
+        q = self.rotary_embed(q)
+        k = self.rotary_embed(k)
+
+        if self.flash:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=0, is_causal=False
+            )
+
+        y = rearrange(y, "b h t d -> b t (h d)")
+
+        y = self.c_proj(y)
+        # shape: (b, t, h*d)
+
+        return y
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_heads = n_heads
+
+        self.att_norm = RMSNorm(dim)
+        self.ffn_norm = RMSNorm(dim)
+        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed)
+        self.mlp = MLP(dim=dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        x = x + self.att(self.att_norm(x))
+        x = x + self.mlp(self.ffn_norm(x))
+        return x
+
+
+if __name__ == "__main__":
+    rotary_embed_128 = RotaryPositionalEmbeddings(dim=128)
+    transformer_block = TransformerBlock(
+        dim=1024, n_heads=8, rotary_embed=rotary_embed_128
+    )
+    x = torch.randn(2, 128, 1024)
+    y = transformer_block(x)
+    print(y.shape)
+    c = 1

neucodec/codec_decoder_vocos.py

@@ -0,0 +1,431 @@
+import torch
+import torch.nn as nn
+
+from typing import List
+from torchtune.modules import RotaryPositionalEmbeddings
+from vector_quantize_pytorch import ResidualFSQ
+
+from .bs_roformer5 import TransformerBlock
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(
+        self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"
+    ):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+
+        Returns:
+            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(
+                spec,
+                self.n_fft,
+                self.hop_length,
+                self.win_length,
+                self.window,
+                center=True,
+            )
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        assert spec.dim() == 3, "Expected a 3D tensor as input"
+        B, N, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        # Normalize
+        assert (window_envelope > 1e-11).all()
+        y = y / window_envelope
+
+        return y
+
+
+class FourierHead(nn.Module):
+    """Base class for inverse fourier modules."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(
+            n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x_pred = self.out(x)
+        # x_pred = x
+        x_pred = x_pred.transpose(1, 2)
+        mag, p = x_pred.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(
+            mag, max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value
+        S = mag * (x + 1j * y)
+        audio = self.istft(S)
+        return audio.unsqueeze(1), x_pred
+
+
+def nonlinearity(x):
+    # swish
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+    return torch.nn.GroupNorm(
+        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout,
+        temb_channels=512,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv1d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv1d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv1d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = torch.nn.Conv1d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    def forward(self, x, temb=None):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Backbone(nn.Module):
+    """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
+                        C denotes output features, and L is the sequence length.
+
+        Returns:
+            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
+                    and H denotes the model dimension.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class VocosBackbone(Backbone):
+    """
+    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+                                                None means non-conditional model. Defaults to None.
+    """
+
+    def __init__(self, hidden_dim=1024, depth=12, heads=16, pos_meb_dim=64):
+        super().__init__()
+
+        self.embed = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=7, padding=3)
+
+        self.temb_ch = 0
+        block_in = hidden_dim
+        dropout = 0.1
+
+        prior_net: List[nn.Module] = [
+            ResnetBlock(
+                in_channels=block_in,
+                out_channels=block_in,
+                temb_channels=self.temb_ch,
+                dropout=dropout,
+            ),
+            ResnetBlock(
+                in_channels=block_in,
+                out_channels=block_in,
+                temb_channels=self.temb_ch,
+                dropout=dropout,
+            ),
+        ]
+        self.prior_net = nn.Sequential(*prior_net)
+
+        depth = depth
+        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
+
+        transformer_blocks = [
+            TransformerBlock(
+                dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed
+            )
+            for _ in range(depth)
+        ]
+
+        self.transformers = nn.Sequential(*transformer_blocks)
+        self.final_layer_norm = nn.LayerNorm(hidden_dim, eps=1e-6)
+        post_net: List[nn.Module] = [
+            ResnetBlock(
+                in_channels=block_in,
+                out_channels=block_in,
+                temb_channels=self.temb_ch,
+                dropout=dropout,
+            ),
+            ResnetBlock(
+                in_channels=block_in,
+                out_channels=block_in,
+                temb_channels=self.temb_ch,
+                dropout=dropout,
+            ),
+        ]
+        self.post_net = nn.Sequential(*post_net)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        x = self.embed(x)
+        x = self.prior_net(x)
+        x = x.transpose(1, 2)
+        x = self.transformers(x)
+        x = x.transpose(1, 2)
+        x = self.post_net(x)
+        x = x.transpose(1, 2)
+        x = self.final_layer_norm(x)
+        return x
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class CodecDecoderVocos(nn.Module):
+    def __init__(
+        self,
+        hidden_dim=1024,
+        depth=12,
+        heads=16,
+        pos_meb_dim=64,
+        hop_length=320,
+        vq_num_quantizers=1,
+        vq_dim=2048,  # 1024 2048
+        vq_commit_weight=0.25,
+        vq_weight_init=False,
+        vq_full_commit_loss=False,
+        codebook_size=16384,
+        codebook_dim=16,
+    ):
+        super().__init__()
+        self.hop_length = hop_length
+
+        self.quantizer = ResidualFSQ(
+            dim=vq_dim, levels=[4, 4, 4, 4, 4, 4, 4, 4], num_quantizers=1
+        )
+
+        self.backbone = VocosBackbone(
+            hidden_dim=hidden_dim, depth=depth, heads=heads, pos_meb_dim=pos_meb_dim
+        )
+
+        self.head = ISTFTHead(
+            dim=hidden_dim,
+            n_fft=self.hop_length * 4,
+            hop_length=self.hop_length,
+            padding="same",
+        )
+
+        self.reset_parameters()
+
+    def forward(self, x, vq=True):
+        if vq is True:
+            # x, q, commit_loss = self.quantizer(x)
+            x = x.permute(0, 2, 1)
+            x, q = self.quantizer(x)
+            x = x.permute(0, 2, 1)
+            q = q.permute(0, 2, 1)
+            return x, q, None
+        x = self.backbone(x)
+        x, _ = self.head(x)
+
+        return x, _
+
+    def vq2emb(self, vq):
+        self.quantizer = self.quantizer.eval()
+        x = self.quantizer.vq2emb(vq)
+        return x
+
+    def get_emb(self):
+        self.quantizer = self.quantizer.eval()
+        embs = self.quantizer.get_emb()
+        return embs
+
+    def inference_vq(self, vq):
+        x = vq[None, :, :]
+        x = self.model(x)
+        return x
+
+    def inference_0(self, x):
+        x, q, loss, perp = self.quantizer(x)
+        x = self.model(x)
+        return x, None
+
+    def inference(self, x):
+        x = self.model(x)
+        return x, None
+
+    def remove_weight_norm(self):
+        """Remove weight normalization module from all of the layers."""
+
+        def _remove_weight_norm(m):
+            try:
+                torch.nn.utils.remove_weight_norm(m)
+            except ValueError:  # this module didn't have weight norm
+                return
+
+        self.apply(_remove_weight_norm)
+
+    def apply_weight_norm(self):
+        """Apply weight normalization module from all of the layers."""
+
+        def _apply_weight_norm(m):
+            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
+                torch.nn.utils.weight_norm(m)
+
+        self.apply(_apply_weight_norm)
+
+    def reset_parameters(self):
+        self.apply(init_weights)

neucodec/codec_encoder.py

@@ -0,0 +1,84 @@
+import torch
+import numpy as np
+
+from torch import nn
+
+from .module import WNConv1d, EncoderBlock
+from .alias_free_torch import Activation1d
+from . import activations
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class CodecEncoder(nn.Module):
+    def __init__(
+        self,
+        ngf=48,
+        up_ratios=[2, 2, 4, 4, 5],
+        dilations=(1, 3, 9),
+        hidden_dim=1024,
+        depth=12,
+        heads=12,
+        pos_meb_dim=64,
+    ):
+        super().__init__()
+        self.hop_length = np.prod(up_ratios)
+        self.ngf = ngf
+        self.up_ratios = up_ratios
+
+        d_model = ngf
+        self.conv_blocks = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        for i, stride in enumerate(up_ratios):
+            d_model *= 2
+            self.conv_blocks += [
+                EncoderBlock(d_model, stride=stride, dilations=dilations)
+            ]
+
+        self.conv_blocks = nn.Sequential(*self.conv_blocks)
+
+        self.conv_final_block = [
+            Activation1d(
+                activation=activations.SnakeBeta(d_model, alpha_logscale=True)
+            ),
+            WNConv1d(d_model, hidden_dim, kernel_size=3, padding=1),
+        ]
+        self.conv_final_block = nn.Sequential(*self.conv_final_block)
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        x = self.conv_blocks(x)
+        x = self.conv_final_block(x)
+        x = x.permute(0, 2, 1)
+        return x
+
+    def inference(self, x):
+        return self.block(x)
+
+    def remove_weight_norm(self):
+        """Remove weight normalization module from all of the layers."""
+
+        def _remove_weight_norm(m):
+            try:
+                torch.nn.utils.remove_weight_norm(m)
+            except ValueError:  # this module didn't have weight norm
+                return
+
+        self.apply(_remove_weight_norm)
+
+    def apply_weight_norm(self):
+        """Apply weight normalization module from all of the layers."""
+
+        def _apply_weight_norm(m):
+            if isinstance(m, nn.Conv1d):
+                torch.nn.utils.weight_norm(m)
+
+        self.apply(_apply_weight_norm)
+
+    def reset_parameters(self):
+        self.apply(init_weights)

neucodec/codec_encoder_distill.py

@@ -0,0 +1,388 @@
+import torch
+from torch import Tensor
+from torch import nn
+from local_attention.transformer import DynamicPositionBias, LocalMHA, FeedForward
+from .distill_layers import ChannelNorm, Conv1d, Linear, GRN, Snake1d
+from .tconv.t_first import FirstBlock
+
+
+class LocalTrans(nn.Module):
+    def __init__(
+        self,
+        dim=512,
+        depth=6,
+        causal=True,
+        local_attn_window_size=512,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        use_dynamic_pos_bias=False,
+        qk_rmsnorm=False,
+    ):
+        super().__init__()
+
+        self.layers = nn.ModuleList([])
+
+        self.window_size = local_attn_window_size
+        self.use_rotary_pos_emb = not use_dynamic_pos_bias
+        self.dynamic_pos_bias = (
+            None
+            if self.use_rotary_pos_emb
+            else DynamicPositionBias(dim=dim // 2, heads=heads)
+        )
+
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        LocalMHA(
+                            dim=dim,
+                            dim_head=dim_head,
+                            heads=heads,
+                            dropout=attn_dropout,
+                            causal=causal,
+                            window_size=self.window_size,
+                            use_xpos=False,
+                            xpos_scale_base=None,
+                            use_rotary_pos_emb=self.use_rotary_pos_emb,
+                            prenorm=True,
+                            qk_rmsnorm=qk_rmsnorm,
+                            exact_windowsize=False,
+                        ),
+                        FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout),
+                    ]
+                )
+            )
+
+    def forward(self, x, mask=None):
+        attn_bias = (
+            None
+            if self.use_rotary_pos_emb
+            else self.dynamic_pos_bias(self.window_size, self.window_size * 2)
+        )
+        for attn, ff in self.layers:
+            x = attn(x, mask=mask, attn_bias=attn_bias) + x
+            x = ff(x) + x
+
+        return x
+
+    @classmethod
+    def builder(
+        cls, feature_dim=128, depth=2, local_window_size=200, use_dynamic_pos_bias=False
+    ):
+        return cls(
+            dim=feature_dim,
+            depth=depth,
+            dim_head=feature_dim // 4,
+            heads=6,
+            ff_mult=4,
+            causal=True,
+            local_attn_window_size=local_window_size,
+            use_dynamic_pos_bias=use_dynamic_pos_bias,
+        )
+
+
+class LocalEncoder(nn.Module):
+    def __init__(
+        self,
+        feature_dim=128,
+        depth=2,
+        local_window_size=200,
+        use_dynamic_pos_bias=False,
+    ):
+        super().__init__()
+        self.local_trans = LocalTrans.builder(
+            feature_dim=feature_dim,
+            depth=depth,
+            local_window_size=local_window_size,
+            use_dynamic_pos_bias=use_dynamic_pos_bias,
+        )
+
+    def forward(self, feature):
+        """
+        Args:
+            feature: (B, C, T)
+        Returns:
+            local_feature: (B, T, C)
+        """
+        feature = feature.permute(0, 2, 1)
+        feature = self.local_trans(feature)
+        return feature
+
+
+class DownTrans(nn.Module):
+    def __init__(
+        self, feature_dim=128, window_size=200, compress_rate=2, depth=2, **kwargs
+    ):
+        super().__init__()
+        assert window_size % compress_rate == 0
+        self.feature_dim = feature_dim
+        self.compress_rate = compress_rate
+        self.trans = LocalTrans.builder(
+            feature_dim, local_window_size=window_size, depth=depth, **kwargs
+        )
+        self.down_layer = Conv1d(
+            feature_dim, feature_dim, kernel_size=compress_rate, stride=compress_rate
+        )
+
+    def forward(self, x):
+        x = self.trans(x)
+        # x = x[:, ::self.compress_rate, :]  # v1
+        x = self.down_layer(x.permute(0, 2, 1)).permute(0, 2, 1)  # v2
+        return x
+
+
+class CompressedLocalEncoderWithCache(nn.Module):
+    def __init__(
+        self,
+        feature_dim=128,
+        local_window_size=200,
+        compress_rate=2,
+        cache_size=3,
+        depth=4,
+        **kwargs,
+    ):
+        super().__init__()
+        self.local_window_size = local_window_size
+        self.cache_size = cache_size
+        self.compress_rate = compress_rate
+        self.trans_window_size = local_window_size + cache_size
+
+        self.cache_token = nn.Parameter(
+            torch.randn(1, self.cache_size * self.compress_rate, feature_dim)
+        )
+
+        self.down_trans = DownTrans(
+            feature_dim,
+            window_size=self.trans_window_size * compress_rate,
+            compress_rate=compress_rate,
+            depth=2,
+            **kwargs,
+        )
+
+        self.local_trans = LocalTrans.builder(
+            feature_dim,
+            local_window_size=self.trans_window_size,
+            depth=depth - 2,
+            **kwargs,
+        )
+
+    def forward(self, feature):
+        feature = feature.permute(0, 2, 1)
+        split_feature = torch.split(
+            feature, self.local_window_size * self.compress_rate, dim=1
+        )
+        cache_token = self.cache_token.expand(feature.shape[0], -1, -1)
+        feature = torch.cat(
+            [
+                f
+                for fs in split_feature
+                for f in (
+                    cache_token,
+                    fs,
+                )
+            ],
+            dim=1,
+        )
+        # assert feature[:, self.down_trans_window_size: 2*self.down_trans_window_size, :].equal(
+        #     feature.reshape(B, -1, self.down_trans_window_size, C)[:, 1, :, :])
+        feature = self.down_trans(feature)
+        feature = self.local_trans(feature)
+        return feature
+
+
+class ConvUnit(nn.Module):
+    """
+    Args:
+        dim (int): Number of input channels.
+    """
+
+    def __init__(self, dim, snake_act=True, norm=False, dilation=1, kernel_size=7):
+        super().__init__()
+        total_pad = (kernel_size - 1) * dilation
+        self.dw_conv = Conv1d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            padding=total_pad // 2,
+            groups=dim,
+        )  # depth-wise conv
+
+        self.norm = (
+            ChannelNorm(dim, data_format="channels_last") if norm else nn.Identity()
+        )
+        self.pw_conv1 = Linear(
+            dim, 4 * dim
+        )  # point-wise/1x1 conv, implemented with linear layer
+
+        if snake_act:
+            self.act = Snake1d(4 * dim, data_format="channels_last")
+        else:
+            self.act = nn.GELU()
+        self.grn = GRN(4 * dim)
+        self.pw_conv2 = Linear(4 * dim, dim)
+
+    def forward(self, x):
+        x = self.dw_conv(x)
+        x = x.permute(0, 2, 1)  # (N, C, T) -> (N, T, C)
+        x = self.norm(x)
+        x = self.pw_conv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pw_conv2(x)
+        x = x.permute(0, 2, 1)  # (N, T, C) -> (N, C, T)
+        return x
+
+
+class Residual(nn.Module):
+    def __init__(
+        self, module: nn.Module, drop_prob: float = 0.0, scale_by_keep: bool = True
+    ):
+        super().__init__()
+        assert 0 <= drop_prob < 1
+        self.module = module
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def drop_path(self, x_side: Tensor):
+        if self.drop_prob == 0.0 or not self.training:
+            return x_side
+        keep_prob = 1 - self.drop_prob
+        shape = (x_side.shape[0],) + (1,) * (x_side.ndim - 1)
+        keep_mask = x_side.new_empty(shape).bernoulli_(keep_prob)
+        if self.scale_by_keep:
+            keep_mask.div_(keep_prob)
+        return x_side * keep_mask
+
+    def forward(self, x: Tensor):
+        x_side = self.module(x)
+        x_side = self.drop_path(x_side)
+        return x + x_side
+
+
+ResidualUnit = lambda *args, drop_rate=0.0, **kwargs: Residual(
+    ConvUnit(*args, **kwargs), drop_prob=drop_rate
+)
+
+
+class LegacyUnit(nn.Module):
+    def __init__(self, dim, snake_act=True, norm=False, dilation=1, kernel_size=7):
+        super().__init__()
+        assert snake_act, "LegacyUnit only supports snake_act=True"
+        assert norm == False, "LegacyUnit only supports norm=False"
+        total_pad = (kernel_size - 1) * dilation
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            Conv1d(
+                dim,
+                dim,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                padding=total_pad // 2,
+            ),
+            Snake1d(dim),
+            Conv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+ResidualLegacyUnit = lambda *args, **kwargs: Residual(
+    LegacyUnit(*args, **kwargs), drop_prob=0.0
+)
+
+BaseUnit = ResidualUnit
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        feature_dim: int = 512,
+        strides: tuple = (2, 2, 2, 2),
+        depths: tuple = (1, 1, 1, 1, 1),
+        dims: tuple = (32, 64, 128, 256, 512),
+        drop_path_rate: float = 0.0,
+        use_norm=False,
+        use_snake_act=True,
+    ):
+        super().__init__()
+        # Create first convolution
+        blocks = [
+            # Conv1d(1, dims[0], kernel_size=7, padding=3),
+            FirstBlock(dims[0]),
+        ]
+
+        drop_path_rates = [
+            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+        ]
+        cur = 0
+        for i_d, o_d, stride, depth in zip(dims[:-1], dims[1:], strides, depths):
+            stage = nn.Sequential(
+                *[
+                    BaseUnit(
+                        dim=i_d,
+                        drop_rate=drop_path_rates[cur + j],
+                        snake_act=use_snake_act,
+                        norm=use_norm,
+                    )
+                    for j in range(depth)
+                ]
+            )
+            down_layer = nn.Sequential(
+                Conv1d(i_d, o_d, kernel_size=stride, stride=stride),
+                ChannelNorm(o_d, data_format="channels_first")
+                if use_norm
+                else nn.Identity(),
+            )
+            blocks += [stage, down_layer]
+            cur += depth
+
+        # Create last convolution
+        blocks += [
+            nn.Sequential(
+                *[
+                    BaseUnit(
+                        dim=dims[-1],
+                        drop_rate=drop_path_rates[cur + j],
+                        snake_act=use_snake_act,
+                        norm=use_norm,
+                    )
+                    for j in range(depths[-1])
+                ]
+            ),
+            # Snake1d(dims[-1]),
+            Conv1d(dims[-1], feature_dim, kernel_size=3, padding=1),
+        ]
+
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.blocks(x)
+
+
+class DistillCodecEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = Encoder(
+            feature_dim=512,
+            strides=(4, 4, 4, 4),
+            depths=(1, 1, 1, 2),
+            dims=(32, 64, 128, 256),
+        )
+        self.en_encoder = CompressedLocalEncoderWithCache(
+            feature_dim=512,
+            local_window_size=300,
+            compress_rate=5,
+            cache_size=0,
+            depth=5,
+            use_dynamic_pos_bias=True,
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.en_encoder(x)
+        return x

neucodec/distill_layers.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.utils.parametrizations import weight_norm
+
+
+def get_eps(data_type):
+    return torch.finfo(data_type).eps
+
+
+EPS = get_eps(torch.float32)
+
+
+def nn_wrapper(nn_class, norm_weight=True, init_weight=True):
+    def nn_builder(*args, **kwargs):
+        nn_instance = nn_class(*args, **kwargs)
+        if init_weight:
+            nn.init.trunc_normal_(nn_instance.weight, std=0.02)
+            nn.init.constant_(nn_instance.bias, 0)
+        if norm_weight:
+            nn_instance = weight_norm(nn_instance)
+        return nn_instance
+
+    return nn_builder
+
+
+Conv1d = nn_wrapper(nn.Conv1d, norm_weight=True, init_weight=True)
+Linear = nn_wrapper(nn.Linear, norm_weight=True, init_weight=True)
+
+
+class Residual(nn.Module):
+    def __init__(
+        self, module: nn.Module, drop_prob: float = 0.0, scale_by_keep: bool = True
+    ):
+        super().__init__()
+        assert 0 <= drop_prob < 1
+        self.module = module
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def drop_path(self, x_side: Tensor):
+        if self.drop_prob == 0.0 or not self.training:
+            return x_side
+        keep_prob = 1 - self.drop_prob
+        shape = (x_side.shape[0],) + (1,) * (x_side.ndim - 1)
+        keep_mask = x_side.new_empty(shape).bernoulli_(keep_prob)
+        if self.scale_by_keep:
+            keep_mask.div_(keep_prob)
+        return x_side * keep_mask
+
+    def forward(self, x: Tensor):
+        x_side = self.module(x)
+        x_side = self.drop_path(x_side)
+        return x + x_side
+
+
+class GRN(nn.Module):
+    """GRN (Global Response Normalization) layer
+    Which supports two data formats: channels_last (default) or channels_first.
+    Channels_last corresponds to inputs with shape (batch_size, Sequence, channels)
+    while channels_first corresponds to inputs with shape (batch_size, channels, Sequence).
+    """
+
+    def __init__(self, n_channels, eps=EPS, data_format="channels_last"):
+        super().__init__()
+        self.n_channels = n_channels
+        self.data_format = data_format
+        if data_format == "channels_last":
+            self.gamma = nn.Parameter(torch.zeros(1, n_channels))
+            self.beta = nn.Parameter(torch.zeros(1, n_channels))
+            self.channel_dim = -1
+        elif data_format == "channels_first":
+            self.gamma = nn.Parameter(torch.zeros(n_channels, 1))
+            self.beta = nn.Parameter(torch.zeros(n_channels, 1))
+            self.channel_dim = 1
+        else:
+            raise ValueError(f"Unsupported data_format: {data_format}")
+        self.eps = torch.tensor(eps)
+
+    def forward(self, x):
+        g_x = torch.norm(x, p=2, dim=[1, 2], keepdim=True)
+        n_x = g_x / (g_x.mean(dim=self.channel_dim, keepdim=True) + self.eps)
+        return self.gamma * (x * n_x) + self.beta + x
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(n_channels={self.n_channels}, {self.data_format})"
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    # torch.clamp_(alpha, 0.05, 50.)
+    eps = 1.1920928955078125e-07
+    x = x + (alpha + eps).reciprocal() * torch.sin(alpha * x).pow(2)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels, data_format="channels_first"):
+        super().__init__()
+        if data_format == "channels_first":
+            self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+        elif data_format == "channels_last":
+            self.alpha = nn.Parameter(torch.ones(1, 1, channels))
+        else:
+            raise NotImplementedError
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+@torch.jit.script
+def channel_norm(x, weight, bias, eps):
+    u = x.mean(1, keepdim=True)
+    s = (x - u).pow(2).mean(1, keepdim=True)
+    x = (x - u) / torch.sqrt(s + eps)
+    x = weight * x + bias
+    return x
+
+
+class ChannelNorm(nn.Module):
+    """ChannelNorm that supports two data formats: channels_last (default) or channels_first.
+    Channels_last corresponds to inputs with shape (batch_size, ..., channels)
+    while channels_first corresponds to inputs with shape (batch_size, channels, ...).
+    """
+
+    def __init__(self, n_channels, eps=EPS, data_format="channels_last"):
+        super().__init__()
+        self.n_channels = n_channels
+        self.data_format = data_format
+        self.weight = nn.Parameter(torch.ones(n_channels))
+        self.bias = nn.Parameter(torch.zeros(n_channels))
+        self.eps = torch.tensor(eps)
+
+    def forward(self, x):
+        if self.data_format == "channels_first":
+            extend_dims = (1,) * len(x.shape[2:])
+            return channel_norm(
+                x,
+                self.weight.view(-1, *extend_dims),
+                self.bias.view(-1, *extend_dims),
+                self.eps,
+            )
+
+        elif self.data_format == "channels_last":
+            return F.layer_norm(
+                x, (self.n_channels,), self.weight, self.bias, self.eps.item()
+            )
+
+        else:
+            raise NotImplementedError
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(n_channels={self.n_channels}, {self.data_format})"

neucodec/model.py

@@ -0,0 +1,218 @@
+from typing import Optional, Dict
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from torchaudio import transforms as T
+from huggingface_hub import PyTorchModelHubMixin, ModelHubMixin, hf_hub_download
+from transformers import AutoFeatureExtractor, HubertModel, Wav2Vec2BertModel
+
+from .codec_encoder import CodecEncoder
+from .codec_encoder_distill import DistillCodecEncoder
+from .codec_decoder_vocos import CodecDecoderVocos
+from .module import SemanticEncoder
+
+
+class NeuCodec(
+    nn.Module,
+    PyTorchModelHubMixin,
+    repo_url="https://github.com/neuphonic/neucodec",
+    license="apache-2.0",
+):
+
+    def __init__(self, sample_rate: int, hop_length: int, decoder_depth: int = 12):
+        super().__init__()
+        self.sample_rate = sample_rate
+        self.hop_length = hop_length
+        self.semantic_model = Wav2Vec2BertModel.from_pretrained(
+            "facebook/w2v-bert-2.0", output_hidden_states=True
+        )
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
+            "facebook/w2v-bert-2.0"
+        )
+        self.SemanticEncoder_module = SemanticEncoder(1024, 1024, 1024)
+        self.CodecEnc = CodecEncoder()
+        self.generator = CodecDecoderVocos(hop_length=hop_length, depth=decoder_depth)
+        self.fc_prior = nn.Linear(2048, 2048)
+        self.fc_post_a = nn.Linear(2048, 1024)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @classmethod
+    def _from_pretrained(
+        cls,
+        *,
+        model_id: str = None,
+        revision: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        force_download: bool = False,
+        proxies: Optional[Dict] = None,
+        resume_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[str] = None,
+        map_location: str = "cpu",
+        strict: bool = False,
+        local_ckpt_path: str = None,
+        **model_kwargs,
+    ):
+        if model_id == "neuphonic/neucodec":
+            ignore_keys = ["fc_post_s", "SemanticDecoder"]
+        elif model_id == "neuphonic/distill-neucodec":
+            ignore_keys = []
+        else:
+            ignore_keys = []
+
+        if model_id is not None:
+            ckpt_path = hf_hub_download(
+                repo_id=model_id,
+                filename="pytorch_model.bin",
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                local_files_only=local_files_only,
+                token=token,
+            )
+        else:
+            # incase we interpolate the weight to become 960 instead train from scratch 
+            ckpt_path = local_ckpt_path
+
+        # initialize model
+        decoder_depth = model_kwargs.pop('decoder_depth', 12)
+        model = cls(44_100, 882, decoder_depth=decoder_depth)
+
+        # load weights
+        state_dict = torch.load(ckpt_path, map_location)
+        contains_list = lambda s, l: any(i in s for i in l)
+        state_dict = {
+            k:v for k, v in state_dict.items() 
+            if not contains_list(k, ignore_keys)
+        }
+
+        # Filter out keys with shape mismatches (e.g. 48k model vs 24k checkpoint)
+        model_state = model.state_dict()
+        state_dict = {
+            k: v for k, v in state_dict.items()
+            if k in model_state and v.shape == model_state[k].shape
+        }
+
+        model.load_state_dict(state_dict, strict=False)
+
+        return model
+    
+    def _prepare_audio(self, audio_or_path: torch.Tensor | Path | str):
+        
+        # load from file
+        if isinstance(audio_or_path, (Path, str)):
+            y, sr = torchaudio.load(audio_or_path)
+            if sr != 16_000:
+                y, sr = (T.Resample(sr, 16_000)(y), 16_000)
+                y = y[None, :]  # [1, T] -> [B, 1, T]
+
+        # ensure input tensor is of correct shape
+        elif isinstance(audio_or_path, torch.Tensor):
+            y = audio_or_path
+            if len(y.shape) == 3:
+                y = audio_or_path
+            else:
+                raise ValueError(
+                    f"NeuCodec expects tensor audio input to be of shape [B, 1, T] -- received shape: {y.shape}"
+                )
+
+        # pad audio
+        pad_for_wav = 320 - (y.shape[-1] % 320)
+        y = torch.nn.functional.pad(y, (0, pad_for_wav))
+        
+        return y
+        
+    def encode_code(self, audio_or_path: torch.Tensor | Path | str) -> torch.Tensor:
+        """
+        Args:
+            audio_or_path: torch.Tensor [B, 1, T] | Path | str, input audio
+
+        Returns:
+            fsq_codes: torch.Tensor [B, 1, F], 50hz FSQ codes
+        """
+         
+        # prepare inputs
+        y = self._prepare_audio(audio_or_path)
+        semantic_features = self.feature_extractor(
+            [w for w in y.squeeze(1).cpu()], sampling_rate=16_000, return_tensors="pt"
+        ).input_features.to(self.device)
+
+        # acoustic encoding
+        acoustic_emb = self.CodecEnc(y.to(self.device))
+        acoustic_emb = acoustic_emb.transpose(1, 2)
+
+        # semantic encoding
+        semantic_output = (
+            self.semantic_model(semantic_features).hidden_states[16].transpose(1, 2)
+        )
+        semantic_encoded = self.SemanticEncoder_module(semantic_output)
+
+        # concatenate embeddings
+        if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
+            min_len = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
+            acoustic_emb = acoustic_emb[:, :, :min_len]
+            semantic_encoded = semantic_encoded[:, :, :min_len]        
+        concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
+        concat_emb = self.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
+
+        # quantize
+        _, fsq_codes, _ = self.generator(concat_emb, vq=True)
+        return fsq_codes
+
+    def encode_code_from_features(self, audio: torch.Tensor, semantic_features: torch.Tensor) -> torch.Tensor:
+        """Encode using pre-computed semantic features, avoiding CPU feature extraction.
+
+        Args:
+            audio: torch.Tensor [B, 1, T], 16kHz input audio
+            semantic_features: torch.Tensor [B, seq_len, feat_dim], pre-computed features
+
+        Returns:
+            fsq_codes: torch.Tensor [B, 1, F], 50hz FSQ codes
+        """
+        y = self._prepare_audio(audio)
+        semantic_features = semantic_features.to(self.device)
+
+        # acoustic encoding
+        acoustic_emb = self.CodecEnc(y.to(self.device))
+        acoustic_emb = acoustic_emb.transpose(1, 2)
+
+        # semantic encoding
+        semantic_output = (
+            self.semantic_model(semantic_features).hidden_states[16].transpose(1, 2)
+        )
+        semantic_encoded = self.SemanticEncoder_module(semantic_output)
+
+        # concatenate embeddings
+        if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
+            min_len = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
+            acoustic_emb = acoustic_emb[:, :, :min_len]
+            semantic_encoded = semantic_encoded[:, :, :min_len]
+        concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
+        concat_emb = self.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
+
+        # quantize
+        _, fsq_codes, _ = self.generator(concat_emb, vq=True)
+        return fsq_codes
+
+    def decode_code(self, fsq_codes: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            fsq_codes: torch.Tensor [B, 1, F], 50hz FSQ codes
+
+        Returns:
+            recon: torch.Tensor [B, 1, T], reconstructed 48kHz audio
+        """
+
+        fsq_post_emb = self.generator.quantizer.get_output_from_indices(fsq_codes.transpose(1, 2))
+        fsq_post_emb = fsq_post_emb.transpose(1, 2)
+        fsq_post_emb = self.fc_post_a(fsq_post_emb.transpose(1, 2)).transpose(1, 2) 
+        recon = self.generator(fsq_post_emb.transpose(1, 2), vq=False)[0]
+        return recon

neucodec/module.py

@@ -0,0 +1,102 @@
+import torch.nn as nn
+
+from torch.nn.utils import weight_norm
+
+from .activations import SnakeBeta
+from .alias_free_torch import Activation1d
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1, dilations=(1, 3, 9)):
+        super().__init__()
+        runits = [ResidualUnit(dim // 2, dilation=d) for d in dilations]
+        self.block = nn.Sequential(
+            *runits,
+            Activation1d(activation=SnakeBeta(dim // 2, alpha_logscale=True)),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=stride // 2 + stride % 2,
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class SemanticEncoder(nn.Module):
+    def __init__(
+        self,
+        input_channels: int,
+        code_dim: int,
+        encode_channels: int,
+        kernel_size: int = 3,
+        bias: bool = True,
+    ):
+        super(SemanticEncoder, self).__init__()
+
+        self.initial_conv = nn.Conv1d(
+            in_channels=input_channels,
+            out_channels=encode_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+            bias=False,
+        )
+
+        self.residual_blocks = nn.Sequential(
+            nn.ReLU(inplace=True),
+            nn.Conv1d(
+                encode_channels,
+                encode_channels,
+                kernel_size=kernel_size,
+                stride=1,
+                padding=(kernel_size - 1) // 2,
+                bias=bias,
+            ),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(
+                encode_channels,
+                encode_channels,
+                kernel_size=kernel_size,
+                stride=1,
+                padding=(kernel_size - 1) // 2,
+                bias=bias,
+            ),
+        )
+
+        self.final_conv = nn.Conv1d(
+            in_channels=encode_channels,
+            out_channels=code_dim,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+            bias=False,
+        )
+
+    def forward(self, x):
+        x = self.initial_conv(x)  # (Batch, Encode_channels, Length)
+        x = self.residual_blocks(x) + x  # 残差连接
+        x = self.final_conv(x)  # (Batch, Code_dim, Length)
+        return x

neucodec/tconv/base.py

@@ -0,0 +1,92 @@
+import einops
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from ..distill_layers import Conv1d
+
+
+def trend_pool(x, kernel_size):
+    if kernel_size > 1:
+        pool_args = dict(kernel_size=kernel_size, stride=1, padding=kernel_size // 2)
+        return F.avg_pool1d(F.max_pool1d(x.abs(), **pool_args), **pool_args)
+        # return F.avg_pool1d(F.max_pool1d(x, **pool_args), **pool_args)  # woabs
+    else:
+        return x
+
+
+class TrendPool(nn.Module):
+    def __init__(self, kernel_size=5):
+        super().__init__()
+        self.kernel_size = kernel_size
+
+    def forward(self, x):
+        return trend_pool(x, self.kernel_size)
+
+
+class FirstBlock(nn.Module):
+    def __init__(
+        self,
+        target_dim,
+        conv_kernels=(7, 7, 7, 7),
+        pool_kernels=(1, 3, 5, 9),
+        dilation_rate=2,
+    ):
+        super().__init__()
+        assert target_dim % len(pool_kernels) == 0
+        each_dim = target_dim // len(pool_kernels)
+        blocks = []
+        for conv_kernel, pool_kernel in zip(conv_kernels, pool_kernels):
+            conv_dilation = pool_kernel // dilation_rate + 1
+            conv_padding = (conv_kernel - 1) * conv_dilation // 2
+            blocks.append(
+                nn.Sequential(
+                    TrendPool(pool_kernel),
+                    Conv1d(
+                        1,
+                        each_dim,
+                        kernel_size=conv_kernel,
+                        dilation=conv_dilation,
+                        padding=conv_padding,
+                    ),
+                )
+            )
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, x):
+        return torch.cat([block(x) for block in self.blocks], dim=1)
+
+
+class EnhanceBlock(FirstBlock):
+    def __init__(self, dim):
+        super().__init__(4, conv_kernels=(7, 7, 7, 7), pool_kernels=(1, 3, 5, 9))
+        self.dim = dim
+        self.merge_layer = nn.Sequential(
+            # nn.LeakyReLU(),  # ! if active or use InstanceNorm1d
+            nn.InstanceNorm1d(4, affine=True),
+            nn.Conv1d(4, 1, kernel_size=1),
+        )
+
+    def forward(self, x):
+        x = einops.rearrange(x, "b c t -> (b c) 1 t", c=self.dim)
+        y = super().forward(x)
+        y = self.merge_layer(y)
+        y = einops.rearrange(y, "(b c) 1 t -> b c t", c=self.dim)
+        return y  # ! x + y or x + y * x
+
+
+class SimpleEnhanceBlock(FirstBlock):
+    def __init__(self, dim):
+        super().__init__(4, conv_kernels=(7, 7, 7, 7), pool_kernels=(1, 3, 5, 9))
+        self.dim = dim
+        self.merge_layer = nn.Sequential(
+            # nn.LeakyReLU(),  # ! if active or use InstanceNorm1d
+            nn.InstanceNorm1d(4, affine=True),
+            nn.Conv1d(4, self.dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        xi = x[:, :1, :]
+        yi = super().forward(xi)
+        y = self.merge_layer(yi)
+        return x + y * x

neucodec/tconv/t_first.py

@@ -0,0 +1,38 @@
+import torch
+from torch import nn
+
+from ..distill_layers import Conv1d
+from . import base
+
+
+class V3FirstBlock(base.FirstBlock):  # (1, 5, 11, 21, 45)
+    def __init__(
+        self,
+        target_dim,
+        conv_kernels=(7, 7, 7, 7, 7),
+        pool_kernels=(1, 5, 11, 21, 45),
+        dilation_rate=7,
+    ):
+        h_dim = len(pool_kernels) * 4
+        super().__init__(h_dim, conv_kernels, pool_kernels, dilation_rate=dilation_rate)
+        self.conv_1 = Conv1d(h_dim, h_dim * 4, kernel_size=1)
+        self.act = nn.GELU()
+        self.conv_2 = Conv1d(h_dim * 4 + 1, target_dim, kernel_size=1)
+
+    def forward(self, x):
+        h = super().forward(x)
+        h = self.conv_1(h)
+        h = self.act(h)
+        y = torch.cat([h, x], dim=1)
+        y = self.conv_2(y)
+        return y
+
+
+FirstBlock = lambda dim: (
+    V3FirstBlock(
+        dim,
+        conv_kernels=(7, 7, 7, 7, 7),
+        pool_kernels=(1, 5, 11, 21, 45),
+        dilation_rate=99,
+    )  # fv36
+)

neucodec/token_interpolator.py

@@ -0,0 +1,112 @@
+"""
+TokenInterpolator: Upsample low-rate (25/12 TPS) token embeddings back to 50 TPS.
+
+The original 50 TPS NeuCodec codebook is completely frozen and unchanged.
+This module operates purely in embedding space (after quantizer lookup).
+
+Encode at 25 TPS:
+    audio -> NeuCodec.encode_code() -> 50 TPS codes [B, 1, T]
+          -> take every 2nd token   -> 25 TPS codes [B, 1, T//2]
+
+Decode from 25 TPS:
+    25 TPS codes -> quantizer.get_output_from_indices -> 25 TPS embeddings [B, T//2, 1024]
+                -> TokenInterpolator(factor=2) -> 50 TPS embeddings [B, T, 1024]
+                -> NeuCodec decoder backbone + ISTFT -> audio
+
+Training:
+    Freeze entire NeuCodec. Only train TokenInterpolator.
+    Loss: MSE on the predicted (odd-position) embeddings vs the true 50 TPS embeddings.
+    Optional: reconstruction loss via frozen decoder.
+"""
+
+import torch
+import torch.nn as nn
+from torchtune.modules import RotaryPositionalEmbeddings
+from .bs_roformer5 import TransformerBlock
+
+
+class TokenInterpolator(nn.Module):
+    """
+    Upsamples from low-rate token embeddings to 50 TPS embeddings.
+
+    Args:
+        dim:    embedding dimension (1024, matching fc_post_a output)
+        factor: upsample factor — 2 for 25->50 TPS, 4 for 12->50 TPS
+        depth:  number of transformer layers
+        heads:  attention heads
+    """
+
+    def __init__(self, dim: int = 1024, factor: int = 2, depth: int = 4, heads: int = 8):
+        super().__init__()
+        assert factor in (2, 4), "factor must be 2 (25 TPS) or 4 (12 TPS)"
+        self.factor = factor
+        self.dim = dim
+
+        # Learned sub-position embeddings to distinguish slots within each group.
+        # e.g. factor=2: slot 0 = known token, slot 1 = to be predicted.
+        self.sub_pos_embed = nn.Embedding(factor, dim)
+
+        rotary_embed = RotaryPositionalEmbeddings(dim=64)
+        self.transformer = nn.Sequential(*[
+            TransformerBlock(dim=dim, n_heads=heads, rotary_embed=rotary_embed)
+            for _ in range(depth)
+        ])
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T_low, dim]  — embeddings at 25 or 12 TPS
+
+        Returns:
+            out: [B, T_low * factor, dim]  — embeddings at 50 TPS
+        """
+        B, T, D = x.shape
+
+        # Repeat each embedding `factor` times along time axis
+        # [B, T, D] -> [B, T, factor, D] -> [B, T*factor, D]
+        x = x.unsqueeze(2).expand(B, T, self.factor, D).reshape(B, T * self.factor, D)
+
+        # Add sub-position embedding so the model knows which slot it's filling.
+        # sub_idx: [0,1,0,1,...] for factor=2; [0,1,2,3,0,1,2,3,...] for factor=4
+        sub_idx = torch.arange(self.factor, device=x.device).repeat(T)  # [T*factor]
+        x = x + self.sub_pos_embed(sub_idx)  # broadcast over batch
+
+        x = self.transformer(x)
+        x = self.norm(x)
+        return x  # [B, T*factor, D]
+
+
+def encode_low_rate(neucodec, audio, factor: int = 2) -> torch.Tensor:
+    """
+    Encode audio to low-rate codes.
+
+    Returns:
+        codes: [B, 1, T//factor]  integer token indices
+    """
+    codes = neucodec.encode_code(audio)          # [B, 1, T] at 50 TPS
+    codes = codes[:, :, ::factor]                # [B, 1, T//factor]
+    return codes
+
+
+def decode_low_rate(neucodec, interpolator: TokenInterpolator, codes: torch.Tensor) -> torch.Tensor:
+    """
+    Decode low-rate codes back to 48kHz audio via interpolation.
+
+    Args:
+        codes: [B, 1, T_low]  — 25 or 12 TPS codes
+
+    Returns:
+        audio: [B, 1, T_audio]  — 48kHz audio
+    """
+    # 1. Lookup embeddings for the known tokens [B, T_low, 2048]
+    emb = neucodec.generator.quantizer.get_output_from_indices(codes.transpose(1, 2))
+    # 2. Project to 1024-dim space [B, T_low, 1024]
+    emb = neucodec.fc_post_a(emb)
+
+    # 3. Interpolate to 50 TPS [B, T_high, 1024]
+    emb = interpolator(emb)
+
+    # 4. Decode with existing frozen backbone + ISTFT
+    audio, _ = neucodec.generator(emb, vq=False)
+    return audio