low-frame-rate-speech-codec-22khz

Model Architecture

Low Frame-rate Speech Codec model is composed of a fully convolutional generator neural network and three discriminators. The generator comprises an encoder, followed by vector quantization, and a HiFi-GAN-based decoder. The encoder consists of five residual blocks, each block containing three residual layers similar to the multi-receptive field fusion (MRF) module. For the vector quantization, we have used Finite Scalar Quantization (FSQ) with eight codebooks and four dimensions per code and 2016 codes per codebook. For the discriminators, we utilize three neural networks, all employing a squared-GAN and feature-matching loss. We adopt the multi-period discriminator and the multi-scale complex STFT discriminator. Additionally, we proposed the use of Speech Language Models (SLMs) as a discriminator. SLMs encode information ranging from acoustic to semantic aspects, which could benefit our model's training, especially in low frame rate settings where accurate pronunciation is difficult to achieve due to the high compression rate. We adopted the 12-layer WavLM as the SLM. During training, we resample the input audio to 16 kHz before feeding it into the WavLM model, extracting the intermediary layer features. These features are then fed to a discriminative head composed of four 1D convolutional layers.

For more details please check our paper.

Input

• Input Type: Audio
• Input Format(s): .wav files
• Input Parameters: One-Dimensional (1D)
• Other Properties Related to Input: 22050 Hz Mono-channel Audio

Output

• Output Type: Audio
• Output Format: .wav files
• Output Parameters: One Dimensional (1D)
• Other Properties Related to Output: 22050 Hz Mono-channel Audio

How to Use this Model

The model is available for use in the NVIDIA NeMo, and can be used as a pre-trained checkpoint for inference or for fine-tuning on another dataset.

Inference

For inference, you can refer to our Audio Codec Inference Tutorial, which automatically downloads the model checkpoint. Ensure that you set the model_name parameter to "nvidia/low-frame-rate-speech-codec-22khz".

Alternatively, you can use the code below, which also handles the automatic checkpoint download:

import librosa
import torch
import soundfile as sf
from nemo.collections.tts.models import AudioCodecModel

path_to_input_audio = ??? # path of the input audio
path_to_output_audio = ??? # path of the reconstructed output audio

# load audio codec model
nemo_codec_model = AudioCodecModel.from_pretrained("nvidia/low-frame-rate-speech-codec-22khz").eval()

# get discrete tokens from audio
audio, _ = librosa.load(path_to_input_audio, sr=nemo_codec_model.sample_rate)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
audio_tensor = torch.from_numpy(audio).unsqueeze(dim=0).to(device)
audio_len = torch.tensor([audio_tensor[0].shape[0]]).to(device)

encoded_tokens, encoded_len = nemo_codec_model.encode(audio=audio_tensor, audio_len=audio_len)

# Reconstruct audio from tokens
reconstructed_audio, _ = nemo_codec_model.decode(tokens=encoded_tokens, tokens_len=encoded_len)

# save reconstructed audio
output_audio = reconstructed_audio.cpu().numpy().squeeze()
sf.write(path_to_output_audio, output_audio, nemo_codec_model.sample_rate)

If preferred, you can manually download the checkpoint and use the provided code to run inference on the model:

import librosa
import torch
import soundfile as sf
from nemo.collections.tts.models import AudioCodecModel

codec_path = ??? # set here the model .nemo checkpoint path
path_to_input_audio = ??? # path of the input audio
path_to_output_audio = ??? # path of the reconstructed output audio

# load audio codec model
nemo_codec_model = AudioCodecModel.restore_from(restore_path=codec_path, map_location="cpu").eval()

# get discrete tokens from audio
audio, _ = librosa.load(path_to_input_audio, sr=nemo_codec_model.sample_rate)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
audio_tensor = torch.from_numpy(audio).unsqueeze(dim=0).to(device)
audio_len = torch.tensor([audio_tensor[0].shape[0]]).to(device)

encoded_tokens, encoded_len = nemo_codec_model.encode(audio=audio_tensor, audio_len=audio_len)

# Reconstruct audio from tokens
reconstructed_audio, _ = nemo_codec_model.decode(tokens=encoded_tokens, tokens_len=encoded_len)

# save reconstructed audio
output_audio = reconstructed_audio.cpu().numpy().squeeze()
sf.write(path_to_output_audio, output_audio, nemo_codec_model.sample_rate)

Training

For fine-tuning on another dataset please follow the steps available at our Audio Codec Training Tutorial. Note that you will