MOSS-Audio-Tokenizer-Nano

MOSS-Audio-Tokenizer-Nano

This repository contains the Hugging Face remote-code implementation and weights for MOSS-Audio-Tokenizer-Nano, the lightweight audio tokenizer used by MOSS-TTS-Nano.

MOSS-Audio-Tokenizer-Nano is a compact discrete audio tokenizer based on the Cat (Causal Audio Tokenizer with Transformer) architecture from MOSS-Audio-Tokenizer: Scaling Audio Tokenizers for Future Audio Foundation Models. The checkpoint in this repository has 21,969,664 parameters (approximately 22M), making it much smaller than the full-size MOSS-Audio-Tokenizer while preserving the 48 kHz stereo tokenizer interface used by the MOSS-TTS family.

Key Features

• Small model size: approximately 22M parameters, including about 10.45M encoder parameters, 10.45M decoder parameters, and 1.07M quantizer parameters.
• Native high-resolution audio: supports 48 kHz input and output with 2-channel stereo audio, helping reduce compression loss and improve listening quality.
• Low-frame-rate discrete codes: compresses 48 kHz stereo audio into a 12.5 Hz token stream with a downsample rate of 7,680 samples.
• Variable bitrate reconstruction: uses a residual quantizer stack with 16 codebooks and 1,024 entries per codebook. Each codebook contributes about 0.125 kbps, for an inference range from 0.125 kbps to 2 kbps.
• Transformer-based tokenizer: uses causal Transformer blocks and supports low-latency streaming encode/decode.
• MOSS-TTS family interface: designed as the audio tokenizer backbone for MOSS-TTS-Nano and compatible MOSS-TTS-family workflows.

Summary: By combining a compact causal Transformer tokenizer with native 48 kHz stereo modeling, MOSS-Audio-Tokenizer-Nano reduces the deployment cost of the MOSS audio tokenizer interface while keeping high-fidelity reconstruction for speech, general audio, and music. It provides a lightweight, low-frame-rate, and streaming-friendly discrete audio representation for MOSS-TTS-Nano and other real-time speech generation workflows.

This repository contains a lightweight remote-code implementation that mirrors the current Hugging Face Transformers transformers.models.moss_audio_tokenizer module. Load it with trust_remote_code=True when needed.

Evaluation Metrics

The table below compares the reconstruction quality of MOSS-Audio-Tokenizer-Nano with open-source audio tokenizers with no more than 120M parameters on speech, audio, and music data. MOSS-Audio-Tokenizer-Nano keeps one of the smallest model sizes in the comparison while supporting 48 kHz stereo reconstruction.

• Speech metrics are evaluated on LibriSpeech test-clean (English) and AISHELL-2 (Chinese), reported as EN/ZH.
• Audio metrics are evaluated on the AudioSet evaluation subset, while music metrics are evaluated on MUSDB, reported as audio/music.
• STFT-Dist. denotes the STFT distance.
• Higher is better for speech metrics, while lower is better for audio/music metrics (Mel-Loss, STFT-Dist.).
• Ch. denotes the number of input/output channels supported by the audio tokenizer: ch=1 means mono audio, and ch=2 means stereo audio.
• Nvq denotes the number of quantizers.

Usage

Quickstart

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)

Attention Backend And Compute Dtype

config.attention_implementation controls whether Transformer layers prefer sdpa or flash_attention_2. config.compute_dtype controls the non-quantizer autocast dtype and supports fp32, bf16, and fp16.

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

The quantizer always runs in fp32.

Streaming

MossAudioTokenizerModel.encode, decode, batch_encode, and batch_decode all support streaming through a chunk_duration argument.

• chunk_duration is expressed in seconds.
• chunk_duration * MossAudioTokenizerConfig.sampling_rate must be divisible by MossAudioTokenizerConfig.downsample_rate.
• Streaming batch inference is supported.
• The public waveform interface expects stereo inputs shaped (2, T) or batched stereo inputs shaped (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)

Continuous Batch Streaming Decode

For decoder-side continuous batching, prefer batch_decode(..., streaming=True, ...).

• The first streaming call may pass max_batch_size=.... If it is omitted, the first batch size reserves the fixed-slot decoder budget for that public stream.
• Same-size calls continue the existing logical rows in order.
• If a later call is larger, the new rows are admitted by tail append.
• finalize_indices means "decode these rows one last time, then evict them". The indices are interpreted against the pre-call logical order.
• After a finalize call returns, the next streaming call may use the smaller survivor batch.
• reset_stream=True discards the hidden public streaming state and starts a fresh stream.

Milestone 1 boundaries:

• decode-only continuous batching
• one active streaming decode state per model instance
• fixed-slot decoder reservation from max_batch_size
• no encode-side continuous batching
• no physical compaction of surviving decode slots
• no multi-session concurrency on one model instance

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,
)

Repository Layout

• configuration_moss_audio_tokenizer.py
• modeling_moss_audio_tokenizer.py
• __init__.py
• config.json
• model weights

Citation

If you use this model or code in your work, please cite:

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