$ show HeCheng0625/Diffusion-Speech-Tokenizer
This repository provides the official PyTorch implementation of TaDiCodec, a text-aware diffusion speech tokenizer accepted at NeurIPS 2025. The core contribution is a novel speech codec that achieves ultra-low bitrate (0.0875 kbps at 6.25 Hz) by combining a non-autoregressive Llama encoder, binary spherical quantization (BSQ), and a flow-matching diffusion decoder conditioned on text and prompt speech. The repo also includes token-based zero-shot TTS models: an autoregressive (AR) model using LLMs (Qwen2.5, Phi-3.5) and a masked diffusion (MGM) model. Inference pipelines with automatic model downloading from Hugging Face are provided. Training scripts for the tokenizer are partially implemented (dataloader, trainer class), but TTS model training scripts are incomplete. The project has migrated to a new repository (AmphionTeam/Diffusion-Speech-Tokenizer), and this repo is no longer the primary development branch, though recent commits indicate some ongoing activity.
The LLM identified 3 files worth highlighting for reuse:
A clean, standalone implementation of Binary Spherical Quantization (BSQ) with support for both hard quantization and soft entropy loss via differentiable codebook entropy. Includes codebook index mapping, group-wise entropy computation, and L2-normalized code vectors. Can be extracted and reused in any VQ-VAE or discrete representation learning project.
A Llama-based non-autoregressive transformer (DiffLlamaPrefix) with adaptive RMSNorm conditioned on diffusion timesteps. Supports text prefix conditioning and can project input/output dimensions. Demonstrates how to modify Hugging Face Llama layers for bidirectional attention and conditional normalization. Reusable as a backbone for any flow-matching or diffusion generation task that requires sequence-level conditioning.
A neural vocoder implementation (Vocos) using a ConvNeXt backbone and ISTFT head with support for both center and same padding. Includes optional adaptive layer normalization for speaker or bandwidth conditioning. The module is fully self-contained (only depends on torch) and can be reused as a mel-to-waveform decoder in any audio generation pipeline. Also provides MDCT-based alternatives (IMDCTSymExpHead, IMDCTCosHead).