DOCSEARCHbyAlgolia
Submit Docs

NeuCodec

repository·main·Indexed 41 days ago

https://github.com/neuphonic/neucodec

A lightweight 0.8kbps neural audio codec designed for speech tokenization, optimized for text-to-speech models and Speech Language Models (SpeechLMs). It uses Finite Scalar Quantisation (FSQ) to encode 16kHz audio and reconstruct it at 24kHz. The package provides the NeuCodec class for standard operations, DistillNeuCodec for efficient encoding via distilhubert, and NeuCodecOnnxDecoder for high-performance CPU decoding using onnxruntime.

Tokens
2.6K
Snippets
8
Records
10
Agent score
50%

What's inside NeuCodec

  1. Install NeuCodec

    main

    You can install NeuCodec from PyPI. It is recommended to use a dedicated environment.

    Using conda + pip:

    conda create -n neucodec python>3.9
conda activate neucodec
pip install neucodec

    Using uv:

    uv venv neucodec --python 3.10
source neucodec/bin/activate  # On Windows: neucodec\Scripts\activate
uv pip install neucodec

    For ONNX decoder support: If you intend to use the ONNX decoder, you must also install onnxruntime:

    pip install onnxruntime
    pip install neucodec

  2. Encode and decode audio with NeuCodec

    main

    NeuCodec is a 0.8kbps neural audio codec designed for speech tokenization. It uses Finite Scalar Quantisation (FSQ) to produce a single vector for quantized output, making it ideal for Speech Language Models.

    Key technical specifications:

    • Input: 16kHz audio.
    • Output: 24kHz reconstructed audio (via an upsampling decoder).
    • Bit-rate: 0.8kbps (at 50 tokens/sec and 16 bits per token).
    • Architecture: Uses both audio (BigCodec) and semantic (Wav2Vec2-BERT) encoders.

    To use the codec, load the model using NeuCodec.from_pretrained, ensure your input audio is resampled to 16kHz, and use encode_code to generate FSQ codes and decode_code to reconstruct the audio.

    import librosa
import torch
import torchaudio
from torchaudio import transforms as T
from neucodec import NeuCodec
 
model = NeuCodec.from_pretrained("neuphonic/neucodec")
model.eval().cuda()   
 
y, sr = torchaudio.load(librosa.ex("libri1"))
if sr != 16_000:
    y = T.Resample(sr, 16_000)(y)[None, ...] # (B, 1, T_16)
 
with torch.no_grad():
    fsq_codes = model.encode_code(y)
    # fsq_codes = model.encode_code(librosa.ex("libri1")) # or directly pass your filepath!
    print(f"Codes shape: {fsq_codes.shape}")  
    recon = model.decode_code(fsq_codes).cpu() # (B, 1, T_24)
 
torchaudio.save("reconstructed.wav", recon[0, :, :], 24_000)

  3. NeuCodec input requirements and audio preparation

    main

    When passing raw tensors to encode_code, ensure they meet the following criteria:

    • Shape: Must be [B, 1, T] (Batch, Channels, Time).
    • Resampling: If loading from a file via a path, the model automatically resamples audio to 16,000Hz for the semantic encoder, though the final output is 24,000Hz.
    • Padding: The model automatically pads audio to ensure compatibility with the internal processing blocks (padding to a multiple of 320).

  4. Use NeuCodec for audio encoding and decoding

    main

    The NeuCodec class provides a high-level API for compressing audio into FSQ (Finite Scalar Quantization) codes and reconstructing it. It supports both standard and distilled versions.

    Encoding

    Use encode_code to convert audio into compressed codes. You can pass a file path (string or Path) or a PyTorch tensor of shape [B, 1, T].

    Decoding

    Use decode_code to reconstruct audio from the FSQ codes. The output is a reconstructed tensor of shape [B, 1, T] at 24kHz.

    Supported Models

    • neuphonic/neucodec: The standard model.
    • neuphonic/distill-neucodec: A distilled version using distilhubert.
    from neucodec.model import NeuCodec

# Initialize and load weights from Hugging Face
model = NeuCodec.from_pretrained("neuphonic/neucodec")
model.to("cuda")

# 1. Encode audio to FSQ codes
# audio_path can be a string, Path, or torch.Tensor [B, 1, T]
fsq_codes = model.encode_code("path/to/audio.wav")

# 2. Decode FSQ codes back to audio
reconstructed_audio = model.decode_code(fsq_codes)

  5. NeuCodec API Reference

    main

    The NeuCodec class provides the primary interface for audio compression and reconstruction.

    NeuCodec.from_pretrained(model_id)

    Loads a pre-trained model from HuggingFace (e.g., "neuphonic/neucodec").

    model.encode_code(input)

    Encodes audio into Finite Scalar Quantisation (FSQ) codes.

    • Input: A torch Tensor of shape (B, 1, T_16) representing 16kHz audio, or a string representing a file path.
    • Returns: A tensor of quantized FSQ codes.

    model.decode_code(fsq_codes)

    Decodes FSQ codes back into audio.

    • Input: The output from encode_code.
    • Returns: A reconstructed audio tensor of shape (B, 1, T_24) representing 24kHz audio.

  6. Use NeuCodec, DistillNeuCodec, and NeuCodecOnnxDecoder

    main

    The neucodec package provides core classes for neural audio codec operations. You can import the following main classes directly from the package root:

    • NeuCodec: The primary model class for encoding and decoding audio.
    • DistillNeuCodec: A distilled version of the NeuCodec model.
    • NeuCodecOnnxDecoder: A specialized decoder implementation using the ONNX runtime.
    from neucodec import NeuCodec, DistillNeuCodec, NeuCodecOnnxDecoder

  7. Initialize and use the CodecEncoder class

    main

    The CodecEncoder is a PyTorch module used for encoding audio signals. It uses a series of convolutional blocks and encoder blocks with specified up-sampling ratios to process the input.

    Initialization Parameters

    • ngf (int, default: 48): Number of generator filters.
    • up_ratios (list[int], default: [2, 2, 4, 4, 5]): A list of up-sampling ratios used to determine the total hop_length.
    • dilations (tuple[int], default: (1, 3, 9)): Dilation factors for the encoder blocks.
    • hidden_dim (int, default: 1024): The dimension of the final hidden layer.
    • depth (int, default: 12): The depth of the encoder.
    • heads (int, default: 12): Number of attention heads.
    • pos_meb_dim (int, default: 64): Dimension for positional MEB.

    Key Methods

    • forward(x): Performs the forward pass on input tensor x.
    • inference(x): Performs inference on input tensor x.
    • apply_weight_norm(): Applies weight normalization to all nn.Conv1d layers in the module.
    • remove_weight_norm(): Removes weight normalization from all layers.
    • reset_parameters(): Resets the module parameters using truncated normal initialization for weights and zero for biases.
    from neucodec.codec_encoder import CodecEncoder
import torch

# Initialize the encoder
encoder = CodecEncoder(
    ngf=48,
    up_ratios=[2, 2, 4, 4, 5],
    hidden_dim=1024
)

# Example input (batch, channels, length)
x = torch.randn(1, 1, 16000)

# Forward pass
encoded = encoder(x)

# Inference mode
encoded_inf = encoder.inference(x)

  8. Use DistillNeuCodec for efficient encoding

    main

    The DistillNeuCodec class is a specialized version of NeuCodec that utilizes a distilled Hubert model (ntu-spml/distilhubert) for semantic encoding. It follows the same API pattern as NeuCodec for encode_code and decode_code.

    from neucodec.model import DistillNeuCodec

model = DistillNeuCodec.from_pretrained("neuphonic/distill-neucodec")
fsq_codes = model.encode_code("audio.wav")
reconstructed_audio = model.decode_code(fsq_codes)

  9. Manage weight normalization in CodecEncoder

    main

    The CodecEncoder provides methods to toggle weight normalization across its convolutional layers. This is useful when switching between training (where weight norm is often used) and inference (where it might be removed for speed or compatibility).

    • apply_weight_norm(): Iterates through the module and applies torch.nn.utils.weight_norm to all nn.Conv1d instances.
    • remove_weight_norm(): Iterates through the module and removes weight normalization from all layers using torch.nn.utils.remove_weight_norm.
    encoder = CodecEncoder()

# Apply weight normalization
encoder.apply_weight_norm()

# Remove weight normalization
encoder.remove_weight_norm()

  10. Use NeuCodecOnnxDecoder for high-performance CPU decoding

    main

    The NeuCodecOnnxDecoder is a specialized class for running only the decoding portion of the model using onnxruntime. This is optimized for CPU inference.

    Note: This class does not support encoding. It only contains the compiled decoder graph.

    Requirements

    • onnxruntime must be installed (pip install onnxruntime).
    • It currently only supports CPU runtimes.
    import numpy as np
from neucodec.model import NeuCodecOnnxDecoder

# Load the ONNX model from Hugging Face
decoder = NeuCodecOnnxDecoder.from_pretrained("neuphonic/neucodec")

# Prepare codes as a numpy array of shape [B, 1, F]
# Example: codes = np.random.randint(0, 100, (1, 1, 50)).astype(np.int32)
# Note: The exact shape/type depends on your FSQ code generation

# Decode using numpy arrays
reconstructed_audio = decoder.decode_code(fsq_codes_numpy)