S3Tokenizer

s3tokenizer is a reverse-engineered, pure PyTorch port of the Supervised Semantic Speech Tokenizer (S3Tokenizer) originally introduced in CosyVoice. CosyVoice’s authors released only an ONNX checkpoint and no training-friendly PyTorch code, which left users stuck extracting speech codes offline at batch size 1 — a slow, brittle workflow. This library rebuilds the model architecture in PyTorch, loads weights directly from the official ONNX files, and adds distributed, batched inference that the project reports as roughly 790x faster than the original single-sample extraction pipeline.

It ships four supported checkpoints spanning S3Tokenizer V1 through V3 at both 25Hz and 50Hz, a command-line tool for offline (multi-)GPU batch extraction via torchrun, and a Python API meant to be dropped directly into a SpeechLLM’s forward pass for online code extraction instead of a separate offline preprocessing step. Built-in sliding-window handling for audio longer than 30 seconds means callers use the same API regardless of clip length.

What You Get

• A pure PyTorch S3Tokenizer model (plus V2/V3 variants) that initializes its weights from the official CosyVoice ONNX checkpoints via a built-in onnx2torch() converter
• load_model() with automatic download-and-checksum-verify of any of the four published checkpoints (speech_tokenizer_v1, _v1_25hz, _v2_25hz, _v3_25hz)
• Distributed, batched offline inference through the s3tokenizer CLI, launchable under torchrun across multiple GPUs with a DistributedSampler-backed dataloader
• An online extraction path (model.quantize(mels, mels_lens)) meant to be called from inside another model’s forward(), replacing offline-only speech-code extraction
• Automatic long-audio handling: audio over 30 seconds is segmented with a sliding window (30s window, 4s overlap) and merged transparently, with no API change for the caller
• Audio preprocessing utilities (load_audio, log_mel_spectrogram, padding, mask helpers) matching the mel-spectrogram front end the tokenizer was trained on

Common Use Cases

• Extracting semantic speech tokens for CosyVoice-family voice-cloning or TTS pipelines without relying on the slow, batch-size-1 ONNX reference script
• Bulk-tokenizing large speech corpora across multiple GPUs ahead of SpeechLLM pretraining or fine-tuning
• Embedding the tokenizer directly inside a SpeechLLM’s forward() for online code extraction, removing a separate offline preprocessing stage from the training pipeline
• Benchmarking or reverse-engineering CosyVoice’s tokenizer architecture, since the ONNX-to-PyTorch weight mapping is implemented and readable in utils.py
• Processing long-form audio (podcasts, lectures) for speech-token extraction without manually chunking clips, thanks to the built-in sliding-window support

Under The Hood

Architecture The library exposes a single entry point, load_model() in s3tokenizer/__init__.py, which downloads and checksum-verifies an ONNX checkpoint and instantiates one of three model classes (S3Tokenizer, S3TokenizerV2, S3TokenizerV3 in model.py/model_v2.py/model_v3.py) — each a Whisper-derived transformer encoder paired with a vector-quantization codebook. utils.py bridges the gap between the ONNX release and native PyTorch: onnx2torch()/_rename_weights() remap CosyVoice’s ONNX weight names into the PyTorch module’s state dict, while load_audio, log_mel_spectrogram, padding, and merge_tokenized_segments handle preprocessing and long-audio stitching. The CLI (cli.py) is a separate, decoupled driver built on torch.distributed, DataLoader, and DistributedSampler for multi-GPU batch runs. The design is flat rather than layered — no dependency injection or plugin abstractions — so a change to the shared transformer block in model.py would require coordinated updates in both the CLI and any code calling model.quantize() directly.

Tech Stack Python 3.8+, built on torch and torchaudio for the model and audio pipeline, onnx for reading the source checkpoints, einops for tensor reshaping in the attention layers, and tqdm for progress reporting. Distributed inference uses torch.distributed with the NCCL backend, launched via torchrun. Packaging is classic setuptools/setup.py (no pyproject.toml), with a s3tokenizer console-script entry point. CI runs a CPU unit-test workflow on push and a separate release workflow that builds and publishes to PyPI via twine when a commit message matches a Release X.Y.Z pattern.

Code Quality Two test files (test/test_onnx.py, test/test_batch_efficiency.py) use pytest to compare the PyTorch model’s output against ONNX Runtime inference across a range of synthetic audio lengths, asserting numeric consistency rather than just smoke-testing imports — a real regression check for the ONNX-to-PyTorch weight port. A .flake8 config and .pre-commit-config.yaml enforce a linting baseline. Type hints are present in some signatures (typing.Optional/List/Union in cli.py and model.py) but not consistently applied throughout. Error handling leans on bare assert statements (e.g. CUDA availability checks, dimension checks) rather than typed exceptions, with RuntimeError reserved for checkpoint download/checksum failures.

What Makes It Unique The project’s value isn’t a novel model architecture — it’s a Whisper-style transformer encoder with a standard VQ codebook. The differentiator is reverse-engineering an artifact CosyVoice’s authors never open-sourced in trainable form: a working ONNX-to-PyTorch weight converter, plus production-oriented tooling CosyVoice’s own reference script lacked — distributed multi-GPU batch inference reported at roughly 790x the original offline extraction speed, and automatic sliding-window handling for audio over 30 seconds with no change to the calling API.