# 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 ```bash # 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: ```bash # 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//