wavtokenizer-large-unify-40token

WavTokenizer large-unify 40 tokens/s

WavTokenizer is a discrete acoustic codec tokenizer that compresses 24 kHz mono audio into a single codebook of 4096 codes at an extremely low frame rate, making it well suited as an audio tokenizer for language models. The encoder and quantizer follow EnCodec's SEANet encoder with single-codebook vector quantization; the decoder is a Vocos-style backbone (ConvNeXt blocks and a positional conv/attention net) with an inverse STFT head.

It was proposed in WavTokenizer: an Efficient Acoustic Discrete Codec Tokenizer for Audio Language Modeling by Ji et al. (ICLR 2025). The original code (MIT license) is at jishengpeng/WavTokenizer.

This repository contains the large unified-domain checkpoint at 40 tokens/s (hop length 600), converted to the Transformers format from the original release novateur/WavTokenizer-large-unify-40token (wavtokenizer_large_unify_600_24k.ckpt). The conversion is verified bit-exact on encoding and within 1e-4 on decoded waveforms against the original implementation.

Usage

import torch
from datasets import Audio, load_dataset
from transformers import AutoFeatureExtractor, WavTokenizerModel

model = WavTokenizerModel.from_pretrained("swiss-ai/wavtokenizer-large-unify-40token")
feature_extractor = AutoFeatureExtractor.from_pretrained("swiss-ai/wavtokenizer-large-unify-40token")

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate))
audio = dataset[0]["audio"]["array"]

inputs = feature_extractor(audio=audio, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt")
with torch.no_grad():
    audio_codes = model.encode(inputs["input_values"]).audio_codes  # (batch, 1, ceil(samples / 600))
    reconstruction = model.decode(audio_codes).audio_values

[!WARNING] Run the model in float32, the default and the only precision used by the original implementation. As with any codec based on nearest-neighbour vector quantization, code assignment is an argmin over codebook distances: half precision (bfloat16/float16) perturbs the embeddings enough to flip codes near decision boundaries, so tokenization is not reproducible across precisions. Decoded audio is always returned in float32 (the inverse STFT head upcasts internally).

Available checkpoints

Citation

The original authors ask that work building on WavTokenizer cite it in exactly this format, together with their Language-Codec paper:

@article{ji2024wavtokenizer,
  title={Wavtokenizer: an efficient acoustic discrete codec tokenizer for audio language modeling},
  author={Ji, Shengpeng and Jiang, Ziyue and Wang, Wen and Chen, Yifu and Fang, Minghui and Zuo, Jialong and Yang, Qian and Cheng, Xize and Wang, Zehan and Li, Ruiqi and others},
  journal={arXiv preprint arXiv:2408.16532},
  year={2024}
}

@article{ji2024language,
  title={Language-codec: Reducing the gaps between discrete codec representation and speech language models},
  author={Ji, Shengpeng and Fang, Minghui and Jiang, Ziyue and Huang, Rongjie and Zuo, Jialung and Wang, Shulei and Zhao, Zhou},
  journal={arXiv preprint arXiv:2402.12208},
  year={2024}
}