Training Pocket TTS

If you have a GPU and some speech data, you can train your own Pocket TTS! We try to make the training scripts work out of the box, meaning you don't need to know the ins-and-outs of the model to get a TTS. For best results, you might need to dig a bit deeper and tweak some hyperparameters. The training code follows the CALM paper, so start there to get an understanding of how the model works.

We're happy to feature community-trained models in the list of community-trained models (giving credit to the author) if they are novel, e.g. cover languages or voices that our official Pocket TTS models cannot do, or are better in other ways. Open an issue or PR to get started.

Quickstart

First, the minimal setup necessary to train a simple model – we'll talk about each of these steps in more detail below. This will give you a 24-layer model trained on 200h of English data. It won't be the best model in the world, but it'll produce intelligible speech, and demonstrate the training process. You can later re-run with 2000h of data to get results on par with the official model.

Downloading and preparing data (around 15 GB):

uv run training/scripts/prepare_data.py --hours 200

Training:

uv run training/train.py training/configs/scratch.yaml

Generate samples:

uv run pocket-tts generate --config pocket_tts/config/english.yaml \
    --checkpoint runs/scratch/checkpoint_00050000.pt \
    --voice some_speaker.wav --text "Hello there."

Congratulations, you've trained your first Pocket TTS!

With the 200 hours that we used, you should get intelligible speech, but to get a model that's on par with the production Pocket TTS, you'll want 2000 hours or more. You can try to rerun with --hours 2000 and edit scratch.yaml to point to the new manifest.

Now in detail:

Installation

Requirements:

• Linux - we don't provide official support for Windows/Mac training, as we don't have the hardware or time to support those platforms.
• One NVIDIA GPU (the default batch size wants ~56 GB; a consumer GPU runs batch_size: 16 with grad_accum_steps: 4 in ~16 GB); you can use more GPUs to train faster
• Python 3.10+ and uv
• ~60 GB disk per 1,000 hours of prepared audio

Training (and evaluation: WER / speaker similarity / UTMOS) dependencies live in the dev dependency group, so the regular project .venv covers everything — there's no separate extra or virtualenv to install:

uv sync

torch isn't pinned to a specific build, so this installs whatever CUDA-enabled build is current on PyPI. If your GPU driver is older than that build's CUDA version, torch.cuda.is_available() silently returns False — see the note below.

[tool.uv.sources]
torch = [
    { index = "pytorch-cu128" },
]

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

Data

Our example

Our example uses HiFiTTS-2 (nvidia/hifitts-2): 44.1 kHz LibriVox audiobook speech with curated transcripts. To download the dataset and alignments (see "Bring your own data" below on what that means), run:

uv run training/scripts/prepare_data.py --hours 2000

Use 200 hours of audio for a model that speaks, 2000 for one that actually sounds good.

The word alignments come precomputed from the kyutai/hifitts2-aligned dataset. If you want to run the aligner yourself, you can pass --alignments "" --align-shards 8 (--align-shards N splits it over N GPUs). This becomes relevant when you use your own data:

Bring your own data

The most common use case (we think) is training Pocket TTS on a new language: let's see how we can train a Pocket TTS model for Czech, as a concrete example.

We won't discuss the specific code you need to write - just use prepare_data.py as a reference implementation. Your coding agent should be able to adapt it to your custom dataset without difficulty. You just need to know what kind of data you need:

We'll need a dataset of speech paired with transcripts. You can get something working with 100 hours of data, but for a strong model, you'll want 1000+ hours. We'll be using ParCzech4Speech, a dataset of Czech parliamentary speech. Common Voice can be another good free source of data.

Next, we'll need alignments: it's not enough to have pairs of (speech audio, transcript), we also need to know when exactly each word was said. This is because for training, the audio is cut into smaller parts, and we need to know exactly what text corresponds to what part of the audio.

Most speech-text datasets don't include the alignment. To get those, we use forced alignment: we run a model that takes the audio and transcript and figures out where each word starts and ends. It uses a variant of CTC with the Wav2Vec2 model. Since Wav2Vec2 is language-specific, you will need a version fine-tuned for your language. For Czech, we found comodoro/wav2vec2-xls-r-300m-cs-250.

Alternatively, to get aligned transcripts, you could use an ASR with timestamps like whisper-timestamped, but this will come at the cost of noisier transcripts. (It's probably possible to force whisper-timestamped to output the transcripts you already have and only give you the timestamps, but it would require some plumbing.)

Finally, you'll need a text tokenizer to be able to feed the text to the model. This is just a quick training step - see train_tokenizer.py.

Just like ingredients matter in cooking, dataset quality is extremely important. Here are some common issues and solutions:

• As mentioned, you should aim for 1000+ hours of data, 100 hours is a minimum.
• If the TTS is good acoustically but doesn't follow the transcript: it means your transcripts are inaccurate. Prefer hand-transcribed data over automatically-transcribed if possible.
• If your outputs sound noisy/low-quality: This is the acoustic quality of your voice prompt, and the model learns to mimic it. This can be solved at inference time by using a "clean" voice prompt. (You could also train the model to clean up the voice automatically, but this is out of the scope of this README.)
• Your model doesn't do well on a certain kind of voice (e.g. high-pitched, or a specific accent): it means those voices are under-represented in the dataset.
• If your model is generating speech that cuts off part of the first or last word: the issue is probably with your alignment.
• If your model sounds like somebody is reading from a book: this can be due to the training dataset (audiobooks like LibriVox), the voice prompt, or both.

Train

# If you have a single GPU:
uv run training/train.py training/configs/scratch.yaml
# If you have multiple GPUs, launch using Torchrun:
uv run torchrun --nproc-per-node 8 training/train.py training/configs/scratch.yaml

The training is configured using a single YAML file. The official Pocket TTS training happens in two steps, corresponding to the two YAMLs we provide:

• scratch.yaml: trains a 24-layer teacher from scratch
• depth_distill.yaml: distills that teacher into a 6-layer student. This also bakes in classifier-free guidance (CFG), see paper or explanation.

If you want to finetune from the English teacher model, we provide two more configs:

• finetune.yaml: continues from the pretrained English teacher, useful if you want to finetune on more data in the same language (new voices, a new domain).
• finetune_language.yaml: continues from the pretrained English teacher, but initialises the text embedding fresh (the text tokenizer differs). On Czech, it reached the same WER as a scratch run about 2.5x sooner.

The two configs above will give you a teacher that you can then distil down to 6 layers. Training the model in two steps like this works better than training a 6-layer model from scratch.

drifting.yaml trains a 24-layer teacher whose sampler head uses drifting instead of LSD, as in english_drifting_26-09. Distill it with depth_distill.yaml and flow.type: drifting.

Reproducing our results

If you train on 2k hours of HiFiTTS-2 and batch size 64, this is the metrics progress you should expect to see:

• Loss optimizing: At ~2k steps, flow_loss should be ~0.35-0.4 and falling with grad norm ~1.
• Intelligibility (word error rate): Starts dropping around 15k steps, 1% by ~50k steps, then flat forever.
• Acoustic quality (UTMOS): Monotonically increasing to 3.7 at 150k steps, and then stable around 300k steps.

See Evaluate for more info about the metrics.

Regarding timing: the first run precomputes Mimi latents for the train manifest (once, using all GPUs; validation always encodes audio directly), then trains from them. On 2000 h of HiFiTTS-2, effective batch size 64:

Precompute stores ~6 GB of latents per 1000 h of audio next to the manifest. data.precompute: false skips the precomputing, which will start training immediately but at slower speeds.

And for the distillation step (depth_distill.yaml), here is the speed you can expect:

Training format

By default, the training saves:

|-- args.yaml
|-- checkpoint_00025000.pt    Last three checkpoints, for resuming training
|-- checkpoint_00027500.pt
|-- checkpoint_00030000.pt    
|-- progress.jsonl    
|-- model.safetensors         Checkpoint usable with the inference code
|-- optim_00030000.pt         Optimizer state, for resuming training
|-- samples
    |-- step00010000_0.wav
    `-- ...
`-- logs
    `-- ...

model.safetensors: different from the checkpoint_<N>.pt checkpoints in two ways. One, it uses the format that the inference code (in the repo root) expects. And two, it's an exponential moving average of the model weights, averaged over training timesteps. EMA'ing the weights is a common technique to squeeze a bit more performance out of models, see e.g. this.

progress.jsonl: We do not support experiment trackers like Tensorboard or Weights and Biases because it's easy to ask your coding agent to add support for whatever you like to use. We do provide a progress.jsonl that logs machine-readable events you can later parse if needed.

Evaluate

Being able to quantitatively evaluate models is crucial for experimentation. We use the following metrics:

• Word error rate (WER): measures intelligibility. Have the model speak out some sentences, transcribe the output using an ASR (speech-to-text) model, and compare to the reference text. Can be noisy if the ASR model or the eval dataset is bad.
• Speaker similarity: measures how similar the generated voice sounds to the voice prompt. Uses microsoft/wavlm-base-plus-sv to get speaker embeddings.
• UTMOS: runs a model that estimates audio quality, as measured by the "mean opinion score" (MOS) - a 1-5 scale commonly used in audio research.

Our example

In our example, we use the LibriSpeech dataset for evaluation. Run:

uv run training/eval/librispeech.py runs/scratch \
    --librispeech-root /data/LibriSpeech/test-clean --use-ema

Here are the numbers you should expect when training with our example setup:

LSD from scratch, EMA weights, evaluated on the full LibriSpeech test-clean cross-sentence list (1127 items, natural text prompts, Granite ASR, --temp 0.3 --cfg 2.0 --n-steps 1 --eos-threshold -1):

And for the distillation step (distilling a 24-layer teacher into a 6-layer student):

Bring your own data

When training on a new language, not all of these metrics transfer directly:

• Word error rate: you need an ASR that supports your language. For the Czech example above, we could just use Whisper, but for rarer languages, you might need to find a different ASR.
• Speaker similarity: Should work fine.
• UTMOS: Works, but might be less accurate for other languages.

Generate

To generate audio, you can use the usual pocket-tts generate command and pass in the checkpoint you obtained:

uv run pocket-tts generate --config my_config.yaml \
    --checkpoint runs/finetune/checkpoint_00124000.pt \
    --voice voice.wav --text "The quick brown fox jumps over the lazy dog."

Or, from Python:

model = TTSModel.load_model(
    config="my_config.yaml", checkpoint="runs/finetune/checkpoint_00124000.pt"
)
state = model.get_state_for_audio_prompt("voice.wav")
audio = model.generate_audio(state, "The quick brown fox jumps over the lazy dog.")

Tests

To run the unit tests (nothing to do with the models, just checks the correctness of the code):

uv run pytest training/tests -q

Notes on hyperparameters

Here are some messy Claude-written notes about some hyperparameter choices. We're leaving them here in case they help somebody with their experiments.

Distribution

To distribute your trained model, the easiest way is to upload it to Hugging Face. You need to upload the following files:

• the model weights (mimi + flow_lm model) as a .safetensors file
• the tokenizer file as a .model file.
• A .yaml file that has the same structure as the official Pocket TTS config files, but with the paths to your model weights and tokenizer. See this example file for reference.

Once those have been uploaded, try using your model with the official Pocket TTS wheel:

uvx pocket-tts generate --config hf://<your_hf_repo>/<your_yaml_path> \
    --voice <your_voice_prompt.wav> --text "Hello there."

If this works, then any user who has access to your Hugging Face repo can run the same command and get the same results!