spark-pocket-tts/AGENTS.md

AGENTS.md

This file provides guidance to AI agents when working with code in this repository.

Project Overview

pocket-tts is a CPU-based text-to-speech (TTS) model. The project uses a flow-based language model architecture with a neural audio codec (Mimi) for efficient speech synthesis.

Key Architecture Components:

• FlowLMModel: Transformer-based flow language model that generates latent representations from text using Lagrangian Self Distillation (LSD)
• MimiModel: Neural audio codec (from the moshi package) that compresses/decompresses audio to/from latent representations
• Conditioners: Text processing via SentencePiece tokenizer and lookup table embeddings
• Streaming Architecture: The entire pipeline supports streaming generation via stateful modules

Common Commands

Setup and Development

# Install pre-commit hooks
uvx pre-commit install

# Run tests (3 parallel workers)
uv run pytest -n 3 -v

# Run a single test
uv run pytest tests/test_python_api.py -v

# Run CLI locally (editable install)
uv run pocket-tts generate
uv run pocket-tts serve

Linting and Formatting

Pre-commit handles this automatically, but you can run manually:

# Ruff will run automatically on commit via pre-commit
# Includes: ruff-check, ruff-format (with --fix), and import sorting

Building (No Build Step)

This is a pure Python package with Rust extensions in training/rust_exts/audio_ds/ for training-time audio processing. The main package does not require building.

Code Structure

Main Package (`pocket_tts

Entry Points:

• main.py: CLI implementation with Typer (two commands: generate and serve)
• __init__.py: Public API exports only TTSModel

Core Models (models/):

• tts_model.py: Main TTSModel class - orchestrates the entire TTS pipeline
  • load_model(): Downloads weights from HuggingFace and initializes models
  • get_state_for_audio_prompt(): Encodes audio prompt (voice) into model state
  • generate_audio_stream(): Streaming generation that yields audio chunks
  • Uses LRU cache for voice prompts to avoid reprocessing
• flow_lm.py: FlowLMModel - transformer that generates latent audio codes from text

Modules (modules/):

• transformer.py: StreamingTransformer and StreamingMultiheadAttention with RoPE embeddings
• stateful_module.py: Base class for streaming support (maintains KV cache and state)
• rope.py: Rotary Position Embeddings
• mlp.py: SimpleMLPAdaLN (AdaLN-conditioned MLP for flow prediction)
• conv.py: Convolution utilities
• seanet.py: SEANet encoder/decoder (imported from moshi)

Conditioners (conditioners/):

• text.py: LUTConditioner - SentencePiece tokenizer + embedding lookup table for text

Data (data/):

• audio.py: Audio I/O utilities (reading, writing WAV, streaming)
• audio_utils.py: Audio processing (resampling, conversion)

Utils (utils/):

• config.py: Pydantic config models for FlowLM and Mimi
• utils.py: HuggingFace downloads, timing utilities

Configuration (config/):

• b6369a24.yaml: Model configuration (transformer dims, layers, vocab size, etc.)

Testing (tests/)

• test_python_api.py: Tests for public Python API
• test_cli_generate.py: Tests for CLI generate command
• test_documentation_examples.py: Ensures docs examples work

Development Workflow

Key Patterns

1. Streaming Generation: The model generates audio frame-by-frame (12.5 Hz frame rate, 80ms per frame). All modules inherit from StatefulModule to maintain internal state.

2. Voice Cloning: Audio prompts are encoded via Mimi encoder to create "voice state" (latent representations + speaker embeddings). This state is cached via lru_cache on _cached_get_state_for_audio_prompt().

3. Flow Matching: Uses Lagrangian Self Distillation (LSD) with configurable decode steps. Fewer steps = faster but lower quality.

4. EOS Detection: Model predicts end-of-speech via an EOS head. Generation continues for frames_after_eos frames after EOS is detected.

5. Config-Driven: Model architecture is defined in YAML configs. Weights are loaded from HuggingFace Hub via safetensors.

Important Implementation Details

• Thread Safety: The code is NOT thread-safe. Server mode does not support concurrent requests.
• Batching: Batch size is always 1. No batching support currently.
• Device: Defaults to CPU. GPU does not provide speedup for this small model.
• Torch Threads: torch.set_num_threads(1) in tts_model.py for optimal CPU performance
• dtype: Models use float32 by default (configurable in YAML)
• Beartype: Runtime type checking is enabled via beartype claw in __init__.py

Adding Features

When adding features, be aware of:

• The streaming architecture: any changes to model forward passes need to maintain state correctly
• The config system: new model parameters must be added to config classes in utils/config.py
• The public API: only TTSModel is exported; keep implementation details internal
• Ruff formatting: line length 100, LF line endings, skip magic trailing comma

Model Weights

Weights are downloaded from HuggingFace Hub on first use:

• Model weights: hf://kyutai/pocket-tts/tts_b6369a24.safetensors
• Tokenizer: hf://kyutai/pocket-tts/tokenizer.model
• Voice prompts: hf://kyutai/tts-voices/<speaker>/<style>.wav

The download_if_necessary() utility handles hf:// URLs and caches locally.

Common Gotchas

1. PyTorch Version: Requires PyTorch 2.5+. Version 2.4.0 produces incorrect audio.
2. Python Version: Supports 3.11, 3.12, 3.13 only.
3. uv Python Preference: Set to "only-managed" in pyproject.toml because system Python may lack headers.
4. Moshi Dependency: The package depends on moshi for the MimiModel audio codec.
5. CPU-Only PyTorch: Uses PyTorch CPU index from download.pytorch.org/whl/cpu in uv config.