check this audio codec https://github.com/ZhangXInFD/SpeechTokenizer


The user wants me to check out the SpeechTokenizer repository. Let me fetch the GitHub page to learn about it.


ToolSearch

Question: select:WebFetch,Agent


Question: select:WebFetch,Agent


WebFetch

Prompt: Extract all key details about this project: what it does, architecture, number of quantizers/codebooks, bitrate, supported sample rate, how tokens are structured (semantic vs acoustic), training data, model sizes available, and any performance benchmarks. Also note installation/usage instructions and dependencies.


There's an issue with the selected model (claude-haiku-4-5@20251001). It may not exist or you may not have access to it. Run /model to pick a different model.


ToolSearch

Question: select:Bash


Question: select:Bash


Bash

gh api repos/ZhangXInFD/SpeechTokenizer --jq '.description' 2>/dev/null; echo "---"; gh api repos/ZhangXInFD/SpeechTokenizer/readme --jq '.content' 2>/dev/null | base64 -d 2>/dev/null || curl -sL https://raw.githubusercontent.com/ZhangXInFD/SpeechTokenizer/main/README.md 2>/dev/null

Fetch SpeechTokenizer README from GitHub


This is the code for the SpeechTokenizer presented in the SpeechTokenizer: Unified Speech Tokenizer for Speech Language Models. Samples are presented on
---
# SpeechTokenizer: Unified Speech Tokenizer for Speech Language Models

## Introduction
This is the code for the SpeechTokenizer presented in the [SpeechTokenizer: Unified Speech Tokenizer for Speech Language Models](https://arxiv.org/abs/2308.16692). SpeechTokenizer is a unified speech tokenizer for speech language models, which adopts the Encoder-Decoder architecture with residual vector quantization (RVQ). Unifying semantic and acoustic tokens, SpeechTokenizer disentangles different aspects of speech information hierarchically across different RVQ layers. Specifically, the code indices that the first quantizer of RVQ outputs can be considered as semantic tokens and the output of the remaining quantizers mainly contain timbre info, which serve as supplements for the information lost by the first quantizer. We provide our models:
* A model operated at 16khz on monophonic speech trained on Librispeech with average representation across all HuBERT layers as semantic teacher.
* A model with  [Snake activation](https://arxiv.org/abs/2306.06546) operated at 16khz on monophonic speech trained on Librispeech and Common Voice with average representation across all HuBERT layers as semantic teacher.

    Overview

    The SpeechTokenizer framework.

Welcome to try our [SLMTokBench](https://github.com/0nutation/SLMTokBench)
 and we will also open source our [USLM](https://github.com/0nutation/USLM)!

## Qick Link
* [Relase](#release)
* [Samples](#samples)
* [Installation](#installation)
* [Model List](#model-list)
* [Usage](#usage)
* [Train SpeechTokenizer](#train-speechtokenizer)
    * [Data Preprocess](#data-preprocess)
    * [Train](#train)
    * [Quick Start](#quick-start)
* [Citation](#citation)
* [License](#license)

## Release
- [2024/6/9] 🔥 We released the training code of SpeechTokenizer.
- [2024/3] 🔥 We released a checkpoint of SpeechTokenizer with [Snake activation](https://arxiv.org/abs/2306.06546) trained on LibriSpeech and Common Voice.
- [2023/9/11] 🔥 We released code of [soundstorm_speechtokenizer](https://github.com/ZhangXInFD/soundstorm-speechtokenizer).
- [2023/9/10] 🔥 We released code and checkpoints of [USLM](https://github.com/0nutation/USLM).
- [2023/9/1] 🔥 We released code and checkpoints of SpeechTokenizer. Checkout the [paper](https://arxiv.org/abs/2308.16692) and [demo](https://0nutation.github.io/SpeechTokenizer.github.io/).

## Samples

Samples are provided on [our demo page](https://0nutation.github.io/SpeechTokenizer.github.io/).

## Installation

SpeechTokenizer requires Python>=3.8, and a reasonly recent version of PyTorch.
To install SpeechTokenizer, you can run from this repository:
```bash
pip install -U speechtokenizer

# or you can clone the repo and install locally
git clone https://github.com/ZhangXInFD/SpeechTokenizer.git
cd SpeechTokenizer
pip install .
```
## Model List
| Model| Dataset |Discription|
|:----|:----:|:----|
|[speechtokenizer_hubert_avg](https://huggingface.co/fnlp/SpeechTokenizer/tree/main/speechtokenizer_hubert_avg)|LibriSpeech|Adopt average representation across all HuBERT layers as semantic teacher |
|[speechtokenizer_snake](https://huggingface.co/fnlp/AnyGPT-speech-modules/tree/main/speechtokenizer)|LibriSpeech + Common Voice|Snake activation, average representation across all HuBERT layers |
## Usage
### load model
```python
from speechtokenizer import SpeechTokenizer

config_path = '/path/config.json'
ckpt_path = '/path/SpeechTokenizer.pt'
model = SpeechTokenizer.load_from_checkpoint(config_path, ckpt_path)
model.eval()
```
### Extracting discrete representations
```python
import torchaudio
import torch

# Load and pre-process speech waveform
wav, sr = torchaudio.load('<SPEECH_FILE_PATH>')

# monophonic checking
if wav.shape(0) > 1:
    wav = wav[:1,:]

if sr != model.sample_rate:
    wav = torchaudio.functional.resample(wav, sr, model.sample_rate)

wav = wav.unsqueeze(0)

# Extract discrete codes from SpeechTokenizer
with torch.no_grad():
    codes = model.encode(wav) # codes: (n_q, B, T)

RVQ_1 = codes[:1, :, :] # Contain content info, can be considered as semantic tokens
RVQ_supplement = codes[1:, :, :] # Contain timbre info, complete info lost by the first quantizer
```

### Decoding discrete representations
```python
# Concatenating semantic tokens (RVQ_1) and supplementary timbre tokens and then decoding
wav = model.decode(torch.cat([RVQ_1, RVQ_supplement], axis=0))

# Decoding from RVQ-i:j tokens from the ith quantizers to the jth quantizers
wav = model.decode(codes[i: (j + 1)], st=i)
```

## Train SpeechTokenizer
In the following section, we describe how to train a SpeechTokenizer model by using our trainer.
### Data Preprocess
To train the SpeechTokenizer, the first step is to extract semantic teacher representations from raw audio waveforms. We provide an example of how to extract HuBERT representations in [scripts/hubert_rep_extract.sh](scripts/hubert_rep_extract.sh). We explain the arguments in the following:
* `--config`: Config file path. An example is provided in [config/spt_base_cfg.json](config/spt_base_cfg.json). You can modify the `semantic_model_path` and `semantic_model_layer` parameters in this file to change the Hubert model and the target layer.
* `--audio_dir`: The path to the folder containing all audio files.
* `--rep_dir`: The path to the folder storing all semantic representation files.
* `--exts`: The file extension of the audio files. Use ',' to separate multiple extensions if they exist.
* `--split_seed`: Random seed for splitting training set and validation set.
* `--valid_set_size`: The size of validation set. When this number is between 0 and 1, it represents the proportion of the total dataset used for the validation set.

### Train
You can use SpeechTokenizerTrainer to train a SpeechTokenizer as follows:
```python
from speechtokenizer import SpeechTokenizer, SpeechTokenizerTrainer
from speechtokenizer.discriminators import MultiPeriodDiscriminator, MultiScaleDiscriminator, MultiScaleSTFTDiscriminator
import json

# Load model and trainer config
with open('<CONFIG_FILE_PATH>') as f:
    cfg = json.load(f)

# Initialize SpeechTokenizer
generator = SpeechTokenizer(cfg)

# Initialize the discriminators. You can add any discriminator that is not yet implemented in this repository, as long as the output format remains consistent with the discriminators in `speechtokenizer.discriminators`.
discriminators = {'mpd':MultiPeriodDiscriminator(), 'msd':MultiScaleDiscriminator(), 'mstftd':MultiScaleSTFTDiscriminator(32)}

# Initialize Trainer
trainer = SpeechTokenizerTrainer(generator=generator,
                                discriminators=discriminators,
                                cfg=cfg)

# Start training
trainer.train()

# Continue training from checkpoints
trainer.continue_train()
```
We provide example training scripts in [scripts/train_example.sh](scripts/train_example.sh). All arguments for SpeechTokenizerTrainer are defined in [config/spt_base_cfg.json](config/spt_base_cfg.json). Below, we explain some of the important arguments:
* `train_files` and `valid_files`: Training file path and validation file path. These files should be text files listing the paths of all audio files and their corresponding semantic representation files in the training/validation set. Each line should follow the format: "<audio_file_path>\t<semantic_file_path>". If you use [scripts/hubert_rep_extract.sh](scripts/hubert_rep_extract.sh) to extract semantic representations, these two files will be genrated automantically.
* `distill_type`: Use "d_axis" for D-axis distillation loss and "t_axis" for T-axis distillation loss, as mentioned in the paper.

### Quick Start
If you want to fully follow our experimental setup, simply set `semantic_model_path` in [config/spt_base_cfg.json](config/spt_base_cfg.json), and `AUDIO_DIR`, `REP_DIR`, `EXTS` in [scripts/hubert_rep_extract.sh](scripts/hubert_rep_extract.sh), and other optional arguments , then execute the following code:
```shell
cd SpeechTokenizer

# Extact semantic representation
bash scripts/hubert_rep_extract.sh

# Train
bash scripts/train_example.sh
```
## Citation
If you use this code or result in your paper, please cite our work as:
```Tex
@misc{zhang2023speechtokenizer,
      title={SpeechTokenizer: Unified Speech Tokenizer for Speech Language Models},
      author={Xin Zhang and Dong Zhang and Shimin Li and Yaqian Zhou and Xipeng Qiu},
      year={2023},
      eprint={2308.16692},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
```
## License
The code in this repository is released under the Apache 2.0 license as found in the
[LICENSE](LICENSE) file.


Here's a summary of **SpeechTokenizer**:

## What it is
An **Encoder-Decoder speech codec with Residual Vector Quantization (RVQ)** that unifies semantic and acoustic tokens into a single tokenizer — designed specifically for Speech Language Models.

## Key Architecture Details
- **Encoder-Decoder** with RVQ (like EnCodec/DAC, but with a semantic distillation twist)
- **8 RVQ quantizers** total
- **16kHz mono** input
- **RVQ-1 (first quantizer)** = **semantic tokens** (content/linguistic info), distilled from HuBERT representations
- **RVQ 2-8** = **acoustic/timbre tokens**, supplementing info lost by the first quantizer
- Uses **Snake activation** in the newer variant for improved audio quality
- Discriminators: Multi-Period (MPD), Multi-Scale (MSD), Multi-Scale STFT (MSTFTD)

## Why it's interesting
The core innovation is **disentangling semantic and acoustic information across RVQ layers**. Most neural codecs (EnCodec, DAC) don't separate these — all quantizer layers mix content and timbre. SpeechTokenizer forces the first quantizer to align with HuBERT semantic features via a distillation loss, so you get a natural split:
- **RVQ-1 alone** → usable as input to an LLM for content understanding/generation
- **All RVQ layers** → full reconstruction with speaker identity

## Available Models
| Model | Training Data |
|-------|--------------|
| `speechtokenizer_hubert_avg` | LibriSpeech |
| `speechtokenizer_snake` | LibriSpeech + Common Voice |

Both on HuggingFace under `fnlp/SpeechTokenizer`.

## Relevance to your TTS work
This could be relevant for LeWM TTS in a few ways:
- **As a tokenizer/vocoder** — you could encode your 20h dataset into discrete tokens and train your JEPA model on those
- **Semantic-acoustic split** aligns well with two-stage TTS (predict semantic tokens from text, then predict acoustic tokens conditioned on semantic)
- However, at **16kHz** it's lower quality than some alternatives (DAC runs at 24/44.1kHz). For studio-quality output you might want to look at higher-sample-rate codecs like **DAC** or **Vocos**

Want me to clone it and dig into the architecture code in more detail, or compare it against other codecs like EnCodec/DAC/Vocos?


can u test a sample and check =qualirty of encode and decode


The user wants me to clone the repo, install it, encode a sample audio file, decode it back, and check the quality. Let me do that.


Let me clone the repo, install it, and test encode/decode on a sample.


git clone https://github.com/ZhangXInFD/SpeechTokenizer.git 2>&1 | tail -5

Clone SpeechTokenizer repo


ls /home/ubuntu/*.wav /home/ubuntu/*.flac /home/ubuntu/*.mp3 /home/ubuntu/data/*.wav /home/ubuntu/data/*.flac 2>/dev/null | head -5; echo "---"; find /home/ubuntu -maxdepth 3 -name "*.wav" -o -name "*.flac" -o -name "*.mp3" 2>/dev/null | head -10

Look for existing audio files to test with


Cloning into 'SpeechTokenizer'...


/home/ubuntu/bob5_BOB5HRg8tSE.wav
/home/ubuntu/bob5_lalal_16k.wav
/home/ubuntu/bob5_vocals_16k.wav
/home/ubuntu/clip_00.mp3
/home/ubuntu/cooking_input_16k.wav
---
/home/ubuntu/qwen3_hindi_samples/modi_medium.wav
/home/ubuntu/qwen3_hindi_samples/modi_speech.wav
/home/ubuntu/jalsa_mono.wav
/home/ubuntu/modi/094_XDh4q3zKYEY.wav
/home/ubuntu/modi/124_0Lc_11QFxMM.wav
/home/ubuntu/modi/110_3nZkxCsv-mY.wav
/home/ubuntu/modi/093_RUZK2yElCts.wav
/home/ubuntu/modi/079__1LkUMfJzzU.wav
/home/ubuntu/modi/010_hS7gTBmxMS8.wav
/home/ubuntu/modi/116_7xEF1hUdagU.wav


Good, there are audio files available. Let me install SpeechTokenizer and download the model.


cd /home/ubuntu/SpeechTokenizer && pip install -e . 2>&1 | tail -10

Install SpeechTokenizer


  Building editable for speechtokenizer (pyproject.toml): started
  Building editable for speechtokenizer (pyproject.toml): finished with status 'done'
  Created wheel for speechtokenizer: filename=speechtokenizer-1.0.0-0.editable-py3-none-any.whl size=10485 sha256=166457dbbf63886044f8d353852e83866de757322f6ad1cc7d5c7305f8a43901
  Stored in directory: /tmp/pip-ephem-wheel-cache-fxtqxmys/wheels/9c/a8/5f/d98413de9989f8d12eec7626d8600aa3f3ff3c72c0adafc80c
Successfully built speechtokenizer
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
Installing collected packages: lion_pytorch, speechtokenizer

WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
Successfully installed lion_pytorch-0.2.4 speechtokenizer-1.0.0


cd /home/ubuntu && huggingface-cli download fnlp/SpeechTokenizer speechtokenizer_hubert_avg/SpeechTokenizer.pt speechtokenizer_hubert_avg/config.json --local-dir speechtokenizer_models 2>&1 | tail -5

Download SpeechTokenizer checkpoint from HuggingFace


Downloading 'speechtokenizer_hubert_avg/config.json' to 'speechtokenizer_models/.cache/huggingface/download/speechtokenizer_hubert_avg/8_PA_wEVGiVa2goH2H4KQOQpvVY=.c8a8a0fe54262436174c168987313bd222af156d.incomplete'
Download complete. Moving file to speechtokenizer_models/speechtokenizer_hubert_avg/config.json
Download complete. Moving file to speechtokenizer_models/speechtokenizer_hubert_avg/SpeechTokenizer.pt

Fetching 2 files:  50%|█████     | 1/2 [00:03<00:03,  3.66s/it]
Fetching 2 files: 100%|██████████| 2/2 [00:03<00:00,  1.83s/it]
/home/ubuntu/speechtokenizer_models


Now let me run the encode/decode test.


python3 << 'PYEOF'
import torch
import torchaudio
from speechtokenizer import SpeechTokenizer
import time

# Load model
config_path = 'speechtokenizer_models/speechtokenizer_hubert_avg/config.json'
ckpt_path = 'speechtokenizer_models/speechtokenizer_hubert_avg/SpeechTokenizer.pt'
model = SpeechTokenizer.load_from_checkpoint(config_path, ckpt_path)
model.eval()
if torch.cuda.is_available():
    model = model.cuda()
print(f"Model loaded. Sample rate: {model.sample_rate}")

# Load audio
wav, sr = torchaudio.load('bob5_vocals_16k.wav')
print(f"Input: sr={sr}, shape={wav.shape}, duration={wav.shape[1]/sr:.2f}s")

# Mono
if wav.shape[0] > 1:
    wav = wav[:1, :]

# Resample if needed
if sr != model.sample_rate:
    wav = torchaudio.functional.resample(wav, sr, model.sample_rate)
    print(f"Resampled to {model.sample_rate}")

# Trim to 10s for test
max_samples = model.sample_rate * 10
if wav.shape[1] > max_samples:
    wav = wav[:, :max_samples]
    print(f"Trimmed to 10s")

wav_input = wav.unsqueeze(0)
if torch.cuda.is_available():
    wav_input = wav_input.cuda()

# Encode
t0 = time.time()
with torch.no_grad():
    codes = model.encode(wav_input)
t_enc = time.time() - t0
print(f"\nEncoded: codes shape={codes.shape} (n_q, B, T)")
print(f"  n_quantizers={codes.shape[0]}, n_frames={codes.shape[2]}")
print(f"  Encode time: {t_enc:.3f}s")

# Compute bitrate
n_q = codes.shape[0]
n_frames = codes.shape[2]
duration = wav.shape[1] / model.sample_rate
codebook_bits = 10  # 1024 codebook size = 10 bits
bitrate = n_q * n_frames * codebook_bits / duration
print(f"  Bitrate: {bitrate/1000:.2f} kbps")
print(f"  Compression ratio: {wav.shape[1] * 16 / (n_q * n_frames * codebook_bits):.1f}x")

# Decode full
t0 = time.time()
with torch.no_grad():
    wav_out = model.decode(codes)
t_dec = time.time() - t0
print(f"\nDecoded (all RVQ): shape={wav_out.shape}")
print(f"  Decode time: {t_dec:.3f}s")

# Save full reconstruction
wav_out_cpu = wav_out.squeeze(0).cpu()
torchaudio.save('test_speechtok_full.wav', wav_out_cpu, model.sample_rate)

# Decode from RVQ-1 only (semantic only)
with torch.no_grad():
    wav_semantic = model.decode(codes[:1, :, :])
wav_sem_cpu = wav_semantic.squeeze(0).cpu()
torchaudio.save('test_speechtok_rvq1.wav', wav_sem_cpu, model.sample_rate)

# Decode from RVQ 1-4 (half the quantizers)
with torch.no_grad():
    wav_half = model.decode(codes[:4, :, :])
wav_half_cpu = wav_half.squeeze(0).cpu()
torchaudio.save('test_speechtok_rvq1to4.wav', wav_half_cpu, model.sample_rate)

# Compute signal metrics
def snr(original, reconstructed):
    min_len = min(original.shape[1], reconstructed.shape[1])
    orig = original[:, :min_len]
    recon = reconstructed[:, :min_len]
    noise = orig - recon
    return 10 * torch.log10(torch.sum(orig**2) / (torch.sum(noise**2) + 1e-8))

def spectral_distance(original, reconstructed, n_fft=1024):
    min_len = min(original.shape[1], reconstructed.shape[1])
    orig_spec = torch.stft(original[:, :min_len].squeeze(), n_fft, return_complex=True).abs()
    recon_spec = torch.stft(reconstructed[:, :min_len].squeeze(), n_fft, return_complex=True).abs()
    return torch.mean(torch.abs(orig_spec - recon_spec) / (orig_spec + 1e-8)).item()

wav_cpu = wav.cpu()
snr_full = snr(wav_cpu, wav_out_cpu)
snr_sem = snr(wav_cpu, wav_sem_cpu)
snr_half = snr(wav_cpu, wav_half_cpu)

sd_full = spectral_distance(wav_cpu, wav_out_cpu)
sd_sem = spectral_distance(wav_cpu, wav_sem_cpu)
sd_half = spectral_distance(wav_cpu, wav_half_cpu)

print(f"\n=== Quality Metrics ===")
print(f"{'Config':<20} {'SNR (dB)':<12} {'Spectral Dist':<15}")
print(f"{'All RVQ (1-8)':<20} {snr_full.item():<12.2f} {sd_full:<15.4f}")
print(f"{'RVQ 1-4':<20} {snr_half.item():<12.2f} {sd_half:<15.4f}")
print(f"{'RVQ 1 only':<20} {snr_sem.item():<12.2f} {sd_sem:<15.4f}")

print(f"\nSaved: test_speechtok_full.wav, test_speechtok_rvq1to4.wav, test_speechtok_rvq1.wav")

# Show unique codes per quantizer (codebook utilization)
print(f"\n=== Codebook Utilization ===")
for i in range(codes.shape[0]):
    unique = codes[i].unique().numel()
    print(f"  RVQ-{i+1}: {unique}/1024 codes used ({unique/1024*100:.1f}%)")

PYEOF

Encode/decode audio with SpeechTokenizer and measure quality


Error: Exit code 1
Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
  File "/home/ubuntu/SpeechTokenizer/speechtokenizer/__init__.py", line 2, in <module>
    from .trainer import SpeechTokenizerTrainer
  File "/home/ubuntu/SpeechTokenizer/speechtokenizer/trainer/__init__.py", line 1, in <module>
    from .trainer import SpeechTokenizerTrainer
  File "/home/ubuntu/SpeechTokenizer/speechtokenizer/trainer/trainer.py", line 6, in <module>
    from beartype import beartype
ModuleNotFoundError: No module named 'beartype'


Install missing dependency


Installing collected packages: beartype
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
Successfully installed beartype-0.22.9


python3 << 'PYEOF'
import torch
import torchaudio
from speechtokenizer import SpeechTokenizer
import time

# Load model
config_path = 'speechtokenizer_models/speechtokenizer_hubert_avg/config.json'
ckpt_path = 'speechtokenizer_models/speechtokenizer_hubert_avg/SpeechTokenizer.pt'
model = SpeechTokenizer.load_from_checkpoint(config_path, ckpt_path)
model.eval()
if torch.cuda.is_available():
    model = model.cuda()
print(f"Model loaded. Sample rate: {model.sample_rate}")

# Load audio
wav, sr = torchaudio.load('bob5_vocals_16k.wav')
print(f"Input: sr={sr}, shape={wav.shape}, duration={wav.shape[1]/sr:.2f}s")

if wav.shape[0] > 1:
    wav = wav[:1, :]
if sr != model.sample_rate:
    wav = torchaudio.functional.resample(wav, sr, model.sample_rate)

# Trim to 10s
max_samples = model.sample_rate * 10
if wav.shape[1] > max_samples:
    wav = wav[:, :max_samples]
    print(f"Trimmed to 10s")

wav_input = wav.unsqueeze(0)
if torch.cuda.is_available():
    wav_input = wav_input.cuda()

# Encode
t0 = time.time()
with torch.no_grad():
    codes = model.encode(wav_input)
t_enc = time.time() - t0
print(f"\nEncoded: codes shape={codes.shape} (n_q, B, T)")
print(f"  n_quantizers={codes.shape[0]}, n_frames={codes.shape[2]}")
print(f"  Encode time: {t_enc:.3f}s")

n_q = codes.shape[0]
n_frames = codes.shape[2]
duration = wav.shape[1] / model.sample_rate
codebook_bits = 10
bitrate = n_q * n_frames * codebook_bits / duration
print(f"  Bitrate: {bitrate/1000:.2f} kbps")
print(f"  Compression ratio: {wav.shape[1] * 16 / (n_q * n_frames * codebook_bits):.1f}x")

# Decode full
t0 = time.time()
with torch.no_grad():
    wav_out = model.decode(codes)
t_dec = time.time() - t0
print(f"\nDecoded (all RVQ): shape={wav_out.shape}, time={t_dec:.3f}s")

wav_out_cpu = wav_out.squeeze(0).cpu()
torchaudio.save('test_speechtok_full.wav', wav_out_cpu, model.sample_rate)

# Decode RVQ-1 only
with torch.no_grad():
    wav_semantic = model.decode(codes[:1, :, :])
wav_sem_cpu = wav_semantic.squeeze(0).cpu()
torchaudio.save('test_speechtok_rvq1.wav', wav_sem_cpu, model.sample_rate)

# Decode RVQ 1-4
with torch.no_grad():
    wav_half = model.decode(codes[:4, :, :])
wav_half_cpu = wav_half.squeeze(0).cpu()
torchaudio.save('test_speechtok_rvq1to4.wav', wav_half_cpu, model.sample_rate)

# Metrics
def snr(original, reconstructed):
    min_len = min(original.shape[1], reconstructed.shape[1])
    orig = original[:, :min_len]
    recon = reconstructed[:, :min_len]
    noise = orig - recon
    return 10 * torch.log10(torch.sum(orig**2) / (torch.sum(noise**2) + 1e-8))

def spectral_distance(original, reconstructed, n_fft=1024):
    min_len = min(original.shape[1], reconstructed.shape[1])
    orig_spec = torch.stft(original[:, :min_len].squeeze(), n_fft, return_complex=True).abs()
    recon_spec = torch.stft(reconstructed[:, :min_len].squeeze(), n_fft, return_complex=True).abs()
    return torch.mean(torch.abs(orig_spec - recon_spec) / (orig_spec + 1e-8)).item()

wav_cpu = wav.cpu()
snr_full = snr(wav_cpu, wav_out_cpu)
snr_sem = snr(wav_cpu, wav_sem_cpu)
snr_half = snr(wav_cpu, wav_half_cpu)
sd_full = spectral_distance(wav_cpu, wav_out_cpu)
sd_sem = spectral_distance(wav_cpu, wav_sem_cpu)
sd_half = spectral_distance(wav_cpu, wav_half_cpu)

print(f"\n=== Quality Metrics ===")
print(f"{'Config':<20} {'SNR (dB)':<12} {'Spectral Dist':<15}")
print(f"{'All RVQ (1-8)':<20} {snr_full.item():<12.2f} {sd_full:<15.4f}")
print(f"{'RVQ 1-4':<20} {snr_half.item():<12.2f} {sd_half:<15.4f}")
print(f"{'RVQ 1 only':<20} {snr_sem.item():<12.2f} {sd_sem:<15.4f}")

print(f"\nSaved: test_speechtok_full.wav, test_speechtok_rvq1to4.wav, test_speechtok_rvq1.wav")

print(f"\n=== Codebook Utilization ===")
for i in range(codes.shape[0]):
    unique = codes[i].unique().numel()
    print(f"  RVQ-{i+1}: {unique}/1024 codes used ({unique/1024*100:.1f}%)")
PYEOF

Encode/decode audio with SpeechTokenizer and measure quality


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
/home/ubuntu/.local/lib/python3.10/site-packages/torch/functional.py:681: UserWarning: A window was not provided. A rectangular window will be applied,which is known to cause spectral leakage. Other windows such as torch.hann_window or torch.hamming_window are recommended to reduce spectral leakage.To suppress this warning and use a rectangular window, explicitly set `window=torch.ones(n_fft, device=<device>)`. (Triggered internally at /pytorch/aten/src/ATen/native/SpectralOps.cpp:836.)
  return _VF.stft(  # type: ignore[attr-defined]
Model loaded. Sample rate: 16000
Input: sr=16000, shape=torch.Size([1, 10973565]), duration=685.85s
Trimmed to 10s

Encoded: codes shape=torch.Size([8, 1, 500]) (n_q, B, T)
  n_quantizers=8, n_frames=500
  Encode time: 0.363s
  Bitrate: 4.00 kbps
  Compression ratio: 64.0x

Decoded (all RVQ): shape=torch.Size([1, 1, 160000]), time=0.087s

=== Quality Metrics ===
Config               SNR (dB)     Spectral Dist
All RVQ (1-8)        -0.47        2.4106
RVQ 1-4              -0.36        1.8364
RVQ 1 only           -0.60        3.9012

Saved: test_speechtok_full.wav, test_speechtok_rvq1to4.wav, test_speechtok_rvq1.wav

=== Codebook Utilization ===
  RVQ-1: 136/1024 codes used (13.3%)
  RVQ-2: 145/1024 codes used (14.2%)
  RVQ-3: 172/1024 codes used (16.8%)
  RVQ-4: 170/1024 codes used (16.6%)
  RVQ-5: 204/1024 codes used (19.9%)
  RVQ-6: 190/1024 codes used (18.6%)
  RVQ-7: 183/1024 codes used (17.9%)
  RVQ-8: 210/1024 codes used (20.5%)


Now let me also run a perceptual metric (PESQ) and check the spectrograms more carefully.


Install PESQ for perceptual quality measurement


Requirement already satisfied: pesq in ./.local/lib/python3.10/site-packages (0.0.4)
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)


python3 << 'PYEOF'
import torchaudio
import torch
from pesq import pesq
import numpy as np

# Load original and reconstructions
orig, sr = torchaudio.load('bob5_vocals_16k.wav')
orig = orig[0, :160000].numpy()  # first 10s

full, _ = torchaudio.load('test_speechtok_full.wav')
rvq4, _ = torchaudio.load('test_speechtok_rvq1to4.wav')
rvq1, _ = torchaudio.load('test_speechtok_rvq1.wav')

full = full[0, :160000].numpy()
rvq4 = rvq4[0, :160000].numpy()
rvq1 = rvq1[0, :160000].numpy()

# PESQ (wideband)
pesq_full = pesq(16000, orig, full, 'wb')
pesq_rvq4 = pesq(16000, orig, rvq4, 'wb')
pesq_rvq1 = pesq(16000, orig, rvq1, 'wb')

print("=== PESQ Scores (Wideband, range 1.0-4.5) ===")
print(f"  All RVQ (1-8): {pesq_full:.3f}")
print(f"  RVQ 1-4:       {pesq_rvq4:.3f}")
print(f"  RVQ 1 only:    {pesq_rvq1:.3f}")
print()
print("Reference: PESQ > 3.5 = good, > 3.0 = acceptable, < 2.5 = noticeable artifacts")

# Also check RMS levels
print(f"\n=== RMS Levels ===")
print(f"  Original:      {np.sqrt(np.mean(orig**2)):.4f}")
print(f"  Full recon:    {np.sqrt(np.mean(full**2)):.4f}")
print(f"  RVQ 1-4:       {np.sqrt(np.mean(rvq4**2)):.4f}")
print(f"  RVQ 1 only:    {np.sqrt(np.mean(rvq1**2)):.4f}")
PYEOF

Compute PESQ perceptual quality scores


=== PESQ Scores (Wideband, range 1.0-4.5) ===
  All RVQ (1-8): 2.155
  RVQ 1-4:       1.708
  RVQ 1 only:    1.205

Reference: PESQ > 3.5 = good, > 3.0 = acceptable, < 2.5 = noticeable artifacts

=== RMS Levels ===
  Original:      0.0297
  Full recon:    0.0276
  RVQ 1-4:       0.0217
  RVQ 1 only:    0.0125


Here are the results:

## SpeechTokenizer Test Results

**Input:** `bob5_vocals_16k.wav` (10s clip, 16kHz mono)

### Codec Stats
| Metric | Value |
|--------|-------|
| Quantizers | 8 |
| Frames per second | 50 fps |
| Bitrate | **4.0 kbps** (very low) |
| Compression ratio | 64x |
| Codebook size | 1024 per quantizer |

### Quality (PESQ, wideband 1.0-4.5 scale)
| Config | PESQ | Verdict |
|--------|------|---------|
| All RVQ (1-8) | **2.155** | Below acceptable, noticeable artifacts |
| RVQ 1-4 | 1.708 | Poor |
| RVQ 1 only (semantic) | 1.205 | Very poor (but expected — content only) |

### Codebook Utilization
Only **13-20%** of codebook entries are used per quantizer (136-210 out of 1024). This is quite low — suggests either the model is undertrained or the codebook is oversized for this data.

### Verdict

**The reconstruction quality is not great.** A PESQ of 2.15 with all 8 quantizers at 4 kbps means clearly audible artifacts. For comparison:
- **EnCodec** at 6 kbps typically scores PESQ ~3.0+
- **DAC** at 8 kbps scores PESQ ~3.5+

This is somewhat expected — SpeechTokenizer was designed to **prioritize semantic-acoustic disentanglement** over raw reconstruction quality. It was trained on LibriSpeech (clean read speech), and your test audio (vocals) is out-of-domain.

The output files are saved if you want to listen:
- `test_speechtok_full.wav` — all 8 quantizers
- `test_speechtok_rvq1to4.wav` — first 4 quantizers
- `test_speechtok_rvq1.wav` — semantic only

For your TTS project, I'd suggest looking at **DAC** or **Vocos** if reconstruction quality matters, or the newer **SpeechTokenizer with Snake activation** model which may perform better. Want me to test any of those?