WavTokenizer

用于音频语言建模的每秒四十个标记的最先进离散编解码器模型

🎉🎉 使用 WavTokenizer，你只需每秒40个标记就能表示语音、音乐和音频！

🎉🎉 使用 WavTokenizer，你可以获得强大的重建效果。

🎉🎉 WavTokenizer 拥有丰富的语义信息，专为 GPT-4o 等音频语言模型而设计。

🔥 新闻

• 2025年2月25日: 我们更新了 WavTokenizer 的 ICLR 2025 准备版本，并在 huggingface 上发布了 WavTokenizer-large-v2 检查点。
• 2024年10月22日: 我们在 arXiv 上更新了 WavTokenizer，并发布了 WavTokenizer-Large 检查点。
• 2024年9月9日: 我们在 huggingface 上发布了 WavTokenizer-medium 检查点。
• 2024年8月31日: 我们在 arXiv 上发布了 WavTokenizer。

安装

要使用 WavTokenizer，可以按照以下步骤安装：

conda create -n wavtokenizer python=3.9
conda activate wavtokenizer
pip install -r requirements.txt

推理

第一部分：从原始 WAV 文件重建音频

from encoder.utils import convert_audio
import torchaudio
import torch
from decoder.pretrained import WavTokenizer


device=torch.device('cpu')

config_path = "./configs/xxx.yaml"
model_path = "./xxx.ckpt"
audio_outpath = "xxx"

wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path)
wavtokenizer = wavtokenizer.to(device)


wav, sr = torchaudio.load(audio_path)
wav = convert_audio(wav, sr, 24000, 1) 
bandwidth_id = torch.tensor([0])
wav=wav.to(device)
features,discrete_code= wavtokenizer.encode_infer(wav, bandwidth_id=bandwidth_id)
audio_out = wavtokenizer.decode(features, bandwidth_id=bandwidth_id) 
torchaudio.save(audio_outpath, audio_out, sample_rate=24000, encoding='PCM_S', bits_per_sample=16)

第二部分：生成离散编解码器代码

from encoder.utils import convert_audio
import torchaudio
import torch
from decoder.pretrained import WavTokenizer

device=torch.device('cpu')

config_path = "./configs/xxx.yaml"
model_path = "./xxx.ckpt"

wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path)
wavtokenizer = wavtokenizer.to(device)

wav, sr = torchaudio.load(audio_path)
wav = convert_audio(wav, sr, 24000, 1) 
bandwidth_id = torch.tensor([0])
wav=wav.to(device)
_,discrete_code= wavtokenizer.encode_infer(wav, bandwidth_id=bandwidth_id)
print(discrete_code)

第三部分：通过编解码器重建音频

# audio_tokens [n_q,1,t]/[n_q,t]
features = wavtokenizer.codes_to_features(audio_tokens)
bandwidth_id = torch.tensor([0])  
audio_out = wavtokenizer.decode(features, bandwidth_id=bandwidth_id)

可用模型

🤗 Huggingface 模型库链接。

模型名称	HuggingFace	数据集	标记/秒	领域	开源
WavTokenizer-small-600-24k-4096	🤗	LibriTTS	40	语音	√
WavTokenizer-small-320-24k-4096	🤗	LibriTTS	75	语音	√
WavTokenizer-medium-320-24k-4096	🤗	10000 小时	75	语音、音频、音乐	√
WavTokenizer-large-600-24k-4096	🤗	80000 小时	40	语音、音频、音乐	√
WavTokenizer-large-320-24k-4096	🤗	80000 小时	75	语音、音频、音乐	√

训练

第一步：准备训练数据集

# 将数据处理成类似 ./data/demo.txt 的格式

第二步：修改配置文件

# ./configs/xxx.yaml
# 修改 batch_size、filelist_path、save_dir、device 等参数的值

第三步：开始训练过程

有关自定义训练流程的详细信息，请参阅 Pytorch Lightning 文档。

cd ./WavTokenizer
python train.py fit --config ./configs/xxx.yaml

引用

如果这段代码对您的研究有所帮助，请引用我们的工作——Language-Codec 和 WavTokenizer：

@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}
}

WavTokenizer 快速上手指南

WavTokenizer 是一款先进的离散编解码器（Discrete Codec）模型，专为音频语言模型设计。它能够以每秒仅 40 个 token 的高效速率表示语音、音乐和音频，同时保持强大的重建能力和丰富的语义信息，是构建类似 GPT-4o 音频模型的理想工具。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python 版本: 推荐 Python 3.9
• 依赖管理: 推荐使用 conda 进行环境隔离
• 硬件: 支持 CPU 推理；若需训练或加速推理，建议使用 NVIDIA GPU 并安装对应的 CUDA 版本

安装步骤

建议使用国内镜像源加速依赖下载。以下是完整的安装流程：

1. 创建并激活虚拟环境

   conda create -n wavtokenizer python=3.9
   conda activate wavtokenizer
   

2. 安装依赖 克隆项目代码后，使用 pip 安装所需库（已配置国内源加速）：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 获取模型文件 从 Hugging Face 下载所需的模型检查点（.ckpt）和配置文件（.yaml）。

   • 模型仓库地址：novateur/WavTokenizer
   • 推荐模型：WavTokenizer-large-unify-40token (40 tokens/s, 通用领域) 或 WavTokenizer-large-speech-75token (75 tokens/s, 语音优化)。

基本使用

以下示例展示如何加载预训练模型并将原始音频转换为离散 Token，再还原为音频。

1. 音频重建与编码示例

将以下代码保存为 infer_demo.py，并确保 config_path 和 model_path 指向你下载的实际文件路径。

from encoder.utils import convert_audio
import torchaudio
import torch
from decoder.pretrained import WavTokenizer

# 设置设备 (cpu 或 cuda)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 配置路径 (请替换为实际下载的路径)
config_path = "./configs/xxx.yaml"
model_path = "./xxx.ckpt"
audio_path = "./input.wav"       # 输入音频路径
audio_outpath = "./output.wav"   # 输出音频路径

# 加载预训练模型
wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path)
wavtokenizer = wavtokenizer.to(device)

# 加载并预处理音频
wav, sr = torchaudio.load(audio_path)
# 转换为单声道 24kHz
wav = convert_audio(wav, sr, 24000, 1) 
bandwidth_id = torch.tensor([0])
wav = wav.to(device)

# --- 核心功能 A: 编码并重建音频 ---
features, discrete_code = wavtokenizer.encode_infer(wav, bandwidth_id=bandwidth_id)
audio_out = wavtokenizer.decode(features, bandwidth_id=bandwidth_id) 
torchaudio.save(audio_outpath, audio_out, sample_rate=24000, encoding='PCM_S', bits_per_sample=16)
print(f"音频已重建并保存至：{audio_outpath}")

# --- 核心功能 B: 仅获取离散 Token (用于语言模型输入) ---
_, discrete_code = wavtokenizer.encode_infer(wav, bandwidth_id=bandwidth_id)
print("生成的离散 Token 序列:")
print(discrete_code)

2. 从 Token 还原音频

如果你已经拥有离散 Token 序列（例如来自大语言模型的输出），可以使用以下代码将其还原为音频：

# 假设 audio_tokens 形状为 [n_q, 1, t] 或 [n_q, t]
# features = wavtokenizer.codes_to_features(audio_tokens)
# bandwidth_id = torch.tensor([0])  
# audio_out = wavtokenizer.decode(features, bandwidth_id=bandwidth_id)

可用模型速查

模型名称	Token 速率	适用领域	特点
WavTokenizer-large-600-24k-4096	40	语音、音频、音乐	最高效，适合长上下文语言模型
WavTokenizer-large-320-24k-4096	75	语音、音频、音乐	平衡性能与效率
WavTokenizer-medium-320-24k-4096	75	语音、音频、音乐	中等规模，训练数据 1 万小时
WavTokenizer-small-320-24k-4096	75	语音	轻量级，仅限语音

提示: 对于大多数音频语言建模任务，推荐使用 40 tokens/s 的 Large 版本以获得最佳的压缩率和语义保留效果。