OpenMOSS/MOSS-Audio-Tokenizer-Nano
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

MOSS-Audio-Tokenizer-Nano

本仓库包含MOSS-Audio-Tokenizer-Nano的Hugging Face远程代码实现和权重，这是MOSS-TTS-Nano所使用的轻量级音频编码器。

MOSS-Audio-Tokenizer-Nano是一款紧凑的离散音频编码器，基于MOSS-Audio-Tokenizer: Scaling Audio Tokenizers for Future Audio Foundation Models中的Cat（Causal Audio Tokenizer with Transformer，因果音频Transformer编码器）架构。本仓库中的 checkpoint 拥有21,969,664个参数（约2200万），相比全尺寸的MOSS-Audio-Tokenizer大幅减小，同时保留了MOSS-TTS系列所使用的48 kHz立体声编码器接口。

主要特性

  • 模型体积小巧：约2200万参数，其中编码器参数约1045万，解码器参数约1045万，量化器参数约107万。
  • 原生高分辨率音频：支持48 kHz输入输出及双声道立体声，有助于减少压缩损失，提升聆听质量。
  • 低帧率离散编码：将48 kHz立体声音频压缩为12.5 Hz的令牌流，下采样率为7680样本。
  • 可变比特率重建：采用包含16个码本的残差量化器堆叠，每个码本含1024个条目。每个码本贡献约0.125 kbps，推理比特率范围为0.125 kbps至2 kbps。
  • 基于Transformer的编码器：使用因果Transformer块，支持低延迟流式编解码。
  • MOSS-TTS系列接口：设计作为MOSS-TTS-Nano的音频编码器核心，并兼容MOSS-TTS系列工作流。

总结： 通过将紧凑的因果Transformer编码器与原生48 kHz立体声建模相结合，MOSS-Audio-Tokenizer-Nano降低了MOSS音频编码器接口的部署成本，同时保持了对语音、通用音频和音乐的高保真重建能力。它为MOSS-TTS-Nano及其他实时语音生成工作流提供了轻量级、低帧率且流式友好的离散音频表示。

本仓库包含轻量级的远程代码实现，其镜像了当前Hugging Face Transformers的transformers.models.moss_audio_tokenizer模块。需要时，请使用trust_remote_code=True加载。

评估指标

下表对比了MOSS-Audio-Tokenizer-Nano与参数不超过120M的开源音频分词器在语音、音频和音乐数据上的重建质量。MOSS-Audio-Tokenizer-Nano在对比模型中保持了较小的模型尺寸，同时支持48 kHz立体声重建。

  • 语音指标在LibriSpeech test-clean（英语）和AISHELL-2（中文）上评估，结果以EN/ZH形式呈现。
  • 音频指标在AudioSet评估子集上评估，音乐指标在MUSDB上评估，结果以audio/music形式呈现。
  • STFT-Dist.表示STFT距离。
  • 语音指标数值越高越好，音频/音乐指标（Mel-Loss、STFT-Dist.）数值越低越好。
  • Ch.表示音频分词器支持的输入/输出通道数：ch=1表示单声道音频，ch=2表示立体声音频。
  • Nvq表示量化器数量。
模型参数（M）采样率通道数比特率量化器数量语音：SIM ↑（英/中）语音：STOI ↑（英/中）语音：PESQ-NB ↑（英/中）语音：PESQ-WB ↑（英/中）音频/音乐：Mel-Loss ↓音频/音乐：STFT-Dist. ↓
Mimi VAE2824k1----0.75 / 0.540.91 / 0.832.92 / 2.202.30 / 1.731.35 / 1.312.70 / 2.59
DAC7744.1k186110.30 / 0.200.76 / 0.681.55 / 1.361.24 / 1.151.25 / 1.182.71 / 2.54
SpeechTokenizer12016k1100020.36 / 0.250.77 / 0.681.59 / 1.381.25 / 1.17-- / ---- / --
Mimi9624k1110080.74 / 0.590.91 / 0.852.80 / 2.242.25 / 1.781.24 / 1.192.62 / 2.49
MOSS-Audio-Tokenizer-Nano2248k275060.64 / 0.610.90 / 0.852.65 / 2.282.11 / 1.871.04 / 1.012.42 / 2.27
MOSS-Audio-Tokenizer-Nano2248k2100080.75 / 0.690.92 / 0.872.92 / 2.482.36 / 2.041.00 / 0.972.37 / 2.22
EnCodec1948k2150010.35 / 0.300.76 / 0.751.54 / 1.601.25 / 1.321.25 / 1.052.73 / 2.30
SpeechTokenizer12016k1150030.52 / 0.380.84 / 0.752.00 / 1.601.57 / 1.33-- / ---- / --
Mimi9624k11512.5110.82 / 0.670.92 / 0.883.10 / 2.502.54 / 2.001.19 / 1.142.55 / 2.42
DAC7744.1k1172320.57 / 0.470.86 / 0.802.21 / 1.851.74 / 1.491.03 / 0.992.43 / 2.26
SpeechTokenizer12016k1200040.66 / 0.500.88 / 0.802.38 / 1.791.92 / 1.49-- / ---- / --
Mimi9624k12062.5150.87 / 0.730.94 / 0.903.36 / 2.762.81 / 2.221.14 / 1.092.49 / 2.36
MOSS-Audio-Tokenizer-Nano2248k21500120.84 / 0.770.94 / 0.903.25 / 2.772.71 / 2.310.95 / 0.912.31 / 2.14
MOSS-Audio-Tokenizer-Nano2248k22000160.88 / 0.810.95 / 0.913.40 / 2.932.89 / 2.470.93 / 0.892.28 / 2.11

使用方法

快速开始

import torchaudio
from transformers import AutoModel

repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano"
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval()

wav, sr = torchaudio.load("demo/demo_gt.wav")
if sr != model.sampling_rate:
    wav = torchaudio.functional.resample(wav, sr, model.sampling_rate)

# The public waveform interface expects stereo audio.
if wav.shape[0] == 1:
    wav = wav.repeat(model.config.number_channels, 1)
else:
    wav = wav[: model.config.number_channels]

wav = wav.unsqueeze(0)
enc = model.encode(wav, return_dict=True)
print(f"enc.audio_codes.shape: {enc.audio_codes.shape}")

dec = model.decode(enc.audio_codes, return_dict=True)
print(f"dec.audio.shape: {dec.audio.shape}")

wav = dec.audio.squeeze(0)
torchaudio.save("demo/demo_rec.wav", wav, sample_rate=model.sampling_rate)

# Decode with the first 8 codebooks, roughly 1 kbps.
dec_rvq8 = model.decode(enc.audio_codes[:8], return_dict=True)
wav_rvq8 = dec_rvq8.audio.squeeze(0)
torchaudio.save("demo/demo_rec_rvq8.wav", wav_rvq8, sample_rate=model.sampling_rate)

注意力后端与计算数据类型

config.attention_implementation 控制 Transformer 层是优先使用 sdpa 还是 flash_attention_2。 config.compute_dtype 控制非量化器的自动转换数据类型，支持 fp32、bf16 和 fp16。

model.set_attention_implementation("flash_attention_2")
model.set_compute_dtype("fp16")

量化器始终以 fp32 模式运行。

流式处理

MossAudioTokenizerModel.encode、decode、batch_encode 和 batch_decode 均通过 chunk_duration 参数支持流式处理。

  • chunk_duration 以秒为单位。
  • chunk_duration * MossAudioTokenizerConfig.sampling_rate 必须能被 MossAudioTokenizerConfig.downsample_rate 整除。
  • 支持流式批量推理。
  • 公共波形接口要求立体声输入的形状为 (2, T)，或批量立体声输入的形状为 (B, 2, T)。
import torch
from transformers import AutoModel

repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano"
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval()
audio = torch.randn(2, 48000 * 6)  # dummy stereo waveform

# 6.0s @ 48kHz = 288000 samples, divisible by downsample_rate=3840
enc = model.encode(audio.unsqueeze(0), return_dict=True, chunk_duration=0.08)
dec = model.decode(enc.audio_codes, return_dict=True, chunk_duration=0.08)

batch_enc = model.batch_encode([audio, audio[:, : 48000 * 3]], chunk_duration=0.08)
codes_list = [
    batch_enc.audio_codes[:, i, : batch_enc.audio_codes_lengths[i]]
    for i in range(batch_enc.audio_codes.shape[1])
]
batch_dec = model.batch_decode(codes_list, chunk_duration=0.08)

连续批处理流式解码

对于解码器端的连续批处理，建议使用 batch_decode(..., streaming=True, ...)。

  • 首次流式调用可传入 max_batch_size=...。若省略该参数，首个批处理大小会为该公共流预留固定槽位的解码器资源。
  • 相同大小的调用会按顺序延续现有的逻辑行。
  • 若后续调用的批处理更大，新行将通过尾部追加的方式加入。
  • finalize_indices 表示“对这些行进行最后一次解码，然后将其逐出”。索引是相对于调用前的逻辑顺序进行解释的。
  • 完成 finalize 调用并返回后，下一次流式调用可使用规模更小的幸存批处理。
  • reset_stream=True 会丢弃隐藏的公共流式状态并启动新的流。

里程碑 1 边界：

  • 仅解码的连续批处理
  • 每个模型实例一个活跃的流式解码状态
  • 由 max_batch_size 确定的固定槽位解码器预留
  • 无编码器端连续批处理
  • 不对幸存的解码槽位进行物理压缩
  • 一个模型实例上不支持多会话并发
import torch
from transformers import AutoModel

repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano"
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval()
num_quantizers = model.config.quantizer_kwargs["num_quantizers"]
codebook_size = model.config.quantizer_kwargs["codebook_size"]

codes_a0 = torch.randint(0, codebook_size, (num_quantizers, 2))
codes_b0 = torch.randint(0, codebook_size, (num_quantizers, 3))
codes_a1 = torch.randint(0, codebook_size, (num_quantizers, 2))
codes_b1 = torch.randint(0, codebook_size, (num_quantizers, 2))
codes_c0 = torch.randint(0, codebook_size, (num_quantizers, 1))
codes_a2 = torch.randint(0, codebook_size, (num_quantizers, 1))
codes_b2 = torch.randint(0, codebook_size, (num_quantizers, 2))
codes_c1 = torch.randint(0, codebook_size, (num_quantizers, 2))
codes_b3 = torch.randint(0, codebook_size, (num_quantizers, 1))
codes_c2 = torch.randint(0, codebook_size, (num_quantizers, 1))

# First call reserves 3 fixed decoder slots for A and B.
out_ab0 = model.batch_decode(
    [codes_a0, codes_b0],
    streaming=True,
    max_batch_size=3,
    reset_stream=True,
)

# Same logical rows continue in order; C is a tail append.
out_abc1 = model.batch_decode(
    [codes_a1, codes_b1, codes_c0],
    streaming=True,
)

# Finalize A against the pre-call logical order. A still decodes in this call,
# then is evicted immediately afterward.
out_abc2 = model.batch_decode(
    [codes_a2, codes_b2, codes_c1],
    streaming=True,
    finalize_indices=[0],
)

# The next call can shrink to the surviving logical rows only.
out_bc3 = model.batch_decode(
    [codes_b3, codes_c2],
    streaming=True,
)

仓库结构

  • configuration_moss_audio_tokenizer.py
  • modeling_moss_audio_tokenizer.py
  • __init__.py
  • config.json
  • 模型权重

引用说明

如果您在研究工作中使用了本模型或代码，请引用：

@misc{gong2026mossttstechnicalreport,
  title={MOSS-TTS Technical Report},
  author={Yitian Gong and Botian Jiang and Yiwei Zhao and Yucheng Yuan and Kuangwei Chen and Yaozhou Jiang and Cheng Chang and Dong Hong and Mingshu Chen and Ruixiao Li and Yiyang Zhang and Yang Gao and Hanfu Chen and Ke Chen and Songlin Wang and Xiaogui Yang and Yuqian Zhang and Kexin Huang and ZhengYuan Lin and Kang Yu and Ziqi Chen and Jin Wang and Zhaoye Fei and Qinyuan Cheng and Shimin Li and Xipeng Qiu},
  year={2026},
  eprint={2603.18090},
  archivePrefix={arXiv},
  primaryClass={cs.SD},
  url={https://arxiv.org/abs/2603.18090}
}
@misc{gong2026mossaudiotokenizerscalingaudiotokenizers,
  title={MOSS-Audio-Tokenizer: Scaling Audio Tokenizers for Future Audio Foundation Models},
  author={Yitian Gong and Kuangwei Chen and Zhaoye Fei and Xiaogui Yang and Ke Chen and Yang Wang and Kexin Huang and Mingshu Chen and Ruixiao Li and Qingyuan Cheng and Shimin Li and Xipeng Qiu},
  year={2026},
  eprint={2602.10934},
  archivePrefix={arXiv},
  primaryClass={cs.SD},
  url={https://arxiv.org/abs/2602.10934}
}