# kyutai-labs/pocket-tts architecture

> Architecture diagram of the GitHub repository kyutai-labs/pocket-tts, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-16). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/kyutai-labs/pocket-tts
- Repository: https://github.com/kyutai-labs/pocket-tts

## Overview

1. Pocket TTS is a single Python package centered on on-device text-to-speech generation. The architecture is built around a small CPU-friendly model, with a clean split between model logic, audio handling, voice conditioning, CLI/server entrypoints, and supporting utilities.

2. The main public entrypoints are:
   - `pocket_tts/__init__.py` and `pocket_tts/main.py` for the Python API
   - `pocket_tts/__main__.py` plus `pocket_tts/main.py` for the CLI
   - `pocket_tts/static/index.html` for the local web UI used by the server mode
   The same underlying model pipeline is reused by CLI generation, the HTTP server, and library calls.

3. The core orchestration lives in `pocket_tts/models/tts_model.py`. This is the high-level façade that loads weights, builds the model, prepares voice state from audio or cached state, and exposes audio generation. It is the main boundary between user-facing inputs and the neural inference stack.

4. The model itself is split into two major learned subsystems:
   - `pocket_tts/models/flow_lm.py`: the text-to-audio latent generation / language-model-like component
   - `pocket_tts/models/mimi.py`: the audio codec / representation side used to turn latent predictions into waveform output
   The naming and structure suggest a staged pipeline: text conditioning first, then latent/audio generation, then PCM output.

5. Voice cloning is a first-class path. The `conditioners/` and `data/` layers handle prompt audio and text conditioning:
   - `conditioners/text.py` and `conditioners/base.py` define conditioning interfaces
   - `data/audio.py` and `data/audio_utils.py` handle waveform loading, preprocessing, and sample-rate/audio-format concerns
   A voice state can come from a raw audio file or from a pre-exported serialized state, which is what enables fast reuse.

6. The repository explicitly supports exporting and reloading voice state for speed. The README’s `export_model_state` flow maps to the codebase’s stateful inference design and the presence of `modules/stateful_module.py`. This is important architecturally because long-form or repeated generation avoids recomputing prompt-derived caches.

7. `pocket_tts/config/*.yaml` holds model variants and language-specific settings. These configs are part of the loading boundary: pretrained language/model variants are selected by config, while custom weights are supported via local YAML paths. The included English, French, German, Italian, Portuguese, and Spanish configs reflect the multi-language deployment surface.

8. `pocket_tts/modules/` contains the reusable neural building blocks that make the model portable and CPU-efficient: transformer components, convolution/resampling layers, RoPE, MLPs, layer scaling, SEANet-style audio blocks, and a dummy quantizer. This layer is the implementation substrate for both the flow model and the audio codec.

9. `pocket_tts/quantization.py` and `scripts/evaluate_quantization.py` show that quantization is a distinct concern, not an afterthought. The repo treats quantized execution as a performance/size optimization path, which matters for the CPU-first runtime and for maintaining real-time or better-than-real-time generation.

10. `pocket_tts/utils/` handles cross-cutting concerns: config parsing, logging, debugging, weight loading, and general helpers. `utils/weights_loading.py` is especially important because the model is designed to load from packaged weights, local custom configs, and Hugging Face-style voice assets.

11. The project is packaged and deployable in multiple ways:
   - `Dockerfile`, `docker-compose.yaml`, `docker-bake.hcl`, and `deploy.sh` support containerized use
   - GitHub Actions workflows automate tests, package publishing, and documentation deployment
   - `mkdocs.yml` and the `docs/` tree define the generated documentation site
   This makes the repo both a library and an operational application.

12. The test suite covers the main architectural boundaries rather than only unit internals:
   - CLI generation behavior
   - Python API behavior
   - sentence splitting
   - quantization
   - documentation examples
   That tells you the stable contract is the end-to-end TTS pipeline, not just isolated model modules.

13. For an architecture diagram, the most important flow is:
   user text + voice prompt → conditioning/audio preprocessing → TTSModel orchestration → stateful latent generation model → audio codec / waveform reconstruction → PCM audio output or WAV file.
   The server and CLI are thin wrappers around this same pipeline.

14. The main external dependency boundary is Hugging Face-hosted model/voice assets, plus PyTorch as the runtime. Those assets are not just data files; they are part of the model-loading and voice-cloning workflow, so they should appear as an external model/asset source in any diagram.

Main entry points: [`pocket_tts/models/tts_model.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/tts_model.py) (TTS Model), [`pocket_tts/main.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/main.py) (Main API), [`pocket_tts/modules`](https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/modules) (Modules), [`pocket_tts/__main__.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/__main__.py) (CLI entrypoint)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g1["Entrypoints"]
  node_n1["Package API<br/>python package<br/>[__init__.py]"]
  node_n2["Main API<br/>api entrypoint<br/>[main.py]"]
  node_n3["CLI entrypoint<br/>[__main__.py]"]
  node_n4["Web UI<br/>static ui<br/>[index.html]"]
end

subgraph group_g2["Core model"]
  node_n5["TTS Model<br/>orchestrator<br/>[tts_model.py]"]
  node_n6["Flow LM<br/>latent generator<br/>[flow_lm.py]"]
  node_n7["Mimi Codec<br/>audio codec<br/>[mimi.py]"]
  node_n8["Stateful Module<br/>cache/state<br/>[stateful_module.py]"]
end

subgraph group_g3["Conditioning"]
  node_n9["Text Conditioner<br/>conditioning<br/>[text.py]"]
  node_n10["Conditioning Base<br/>interfaces<br/>[base.py]"]
  node_n11["Audio I/O<br/>audio preprocess<br/>[audio.py]"]
  node_n12["Audio Utils<br/>[audio_utils.py]"]
end

subgraph group_g4["Infrastructure"]
  node_n13["Modules<br/>neural blocks"]
  node_n14["Config YAMLs<br/>model config"]
  node_n15["Weights Loader<br/>loading utils<br/>[weights_loading.py]"]
  node_n16["Core Utils<br/>utilities"]
  node_n17["Quantization<br/>optimization<br/>[quantization.py]"]
end

subgraph group_g5["External runtime"]
  node_n18[("Model Assets<br/>external assets")]
  node_n19["PyTorch<br/>runtime"]
end

node_n1 -->|"exports"| node_n2
node_n3 -->|"invokes"| node_n2
node_n2 -->|"drives"| node_n5
node_n5 -->|"runs"| node_n6
node_n5 -->|"decodes"| node_n7
node_n5 -->|"caches"| node_n8
node_n5 -->|"conditions"| node_n9
node_n5 -->|"loads audio"| node_n11
node_n11 -->|"uses"| node_n12
node_n9 -->|"implements"| node_n10
node_n6 -->|"built from"| node_n13
node_n7 -->|"built from"| node_n13
node_n8 -->|"restores"| node_n15
node_n5 -->|"configured by"| node_n14
node_n15 -->|"fetches"| node_n18
node_n5 -->|"executes on"| node_n19
node_n13 -->|"depends on"| node_n19
node_n17 -->|"optimizes"| node_n19
node_n2 -->|"exposes"| node_n17
node_n3 -->|"serves"| node_n4
node_n4 -->|"calls"| node_n2

click node_n1 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/__init__.py"
click node_n2 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/main.py"
click node_n3 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/__main__.py"
click node_n4 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/static/index.html"
click node_n5 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/tts_model.py"
click node_n6 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/flow_lm.py"
click node_n7 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/mimi.py"
click node_n8 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/stateful_module.py"
click node_n9 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/conditioners/text.py"
click node_n10 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/conditioners/base.py"
click node_n11 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/data/audio.py"
click node_n12 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/data/audio_utils.py"
click node_n13 "https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/modules"
click node_n14 "https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/config"
click node_n15 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/utils/weights_loading.py"
click node_n16 "https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/utils"
click node_n17 "https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/quantization.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_n1,node_n2,node_n3,node_n4 toneBlue
class node_n5,node_n6,node_n7,node_n8 toneAmber
class node_n9,node_n10,node_n11,node_n12 toneMint
class node_n13,node_n14,node_n15,node_n16,node_n17 toneRose
class node_n18,node_n19 toneIndigo
```

## Components

### Entrypoints

Public API, CLI, and local server UI

- [Package API](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/__init__.py): python package · Public library entrypoint for importing Pocket TTS · `pocket_tts/__init__.py`
- [Main API](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/main.py): api entrypoint · Shared runtime surface for library calls and generation · `pocket_tts/main.py`
- [CLI entrypoint](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/__main__.py): cli · Module entrypoint for command-line generation and serving · `pocket_tts/__main__.py`
- [Web UI](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/static/index.html): static ui · Local browser interface used by server mode · `pocket_tts/static/index.html`

### Core model

High-level orchestration and learned inference stack

- [TTS Model](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/tts_model.py): orchestrator · Loads weights, builds the model, prepares voice state, and runs generation · `pocket_tts/models/tts_model.py`
- [Flow LM](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/flow_lm.py): latent generator · Text-conditioned latent prediction stack · `pocket_tts/models/flow_lm.py`
- [Mimi Codec](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/mimi.py): audio codec · Audio representation and waveform reconstruction path · `pocket_tts/models/mimi.py`
- [Stateful Module](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/stateful_module.py): cache/state · Reusable state and prompt-derived cache handling · `pocket_tts/modules/stateful_module.py`

### Conditioning

Prompt text and audio preparation for voice cloning

- [Text Conditioner](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/conditioners/text.py): conditioning · Text conditioning interfaces and implementations · `pocket_tts/conditioners/text.py`
- [Conditioning Base](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/conditioners/base.py): interfaces · Shared conditioning abstractions · `pocket_tts/conditioners/base.py`
- [Audio I/O](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/data/audio.py): audio preprocess · Waveform loading, preprocessing, and sample-rate handling · `pocket_tts/data/audio.py`
- [Audio Utils](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/data/audio_utils.py): audio utils · Low-level audio helper routines · `pocket_tts/data/audio_utils.py`

### Infrastructure

Reusable building blocks, config, utils, and quantization

- [Modules](https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/modules): neural blocks · Shared transformer, convolution, resampling, RoPE, MLP, and SEANet-style blocks · `pocket_tts/modules`
- [Config YAMLs](https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/config): model config · Language and model variant selection files · `pocket_tts/config`
- [Weights Loader](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/utils/weights_loading.py): loading utils · Loading packaged weights, custom configs, and voice assets · `pocket_tts/utils/weights_loading.py`
- [Core Utils](https://github.com/kyutai-labs/pocket-tts/tree/main/pocket_tts/utils): utilities · Config parsing, logging, debugging, and helper functions · `pocket_tts/utils`
- [Quantization](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/quantization.py): optimization · Quantized execution path for CPU-friendly deployment · `pocket_tts/quantization.py`

### External runtime

Model assets and compute backend outside the repository

- **Model Assets**: external assets · Hugging Face-hosted pretrained models and voice files
- **PyTorch**: runtime · Neural network execution backend

## Connections

- Package API → Main API: exports · Package import surface forwards to shared API
- CLI entrypoint → Main API: invokes · CLI uses the same generation API
- Main API → TTS Model: drives · Shared API delegates synthesis to the orchestrator
- TTS Model → Flow LM: runs · High-level model calls the latent generation stack
- TTS Model → Mimi Codec: decodes · Waveform reconstruction goes through the audio codec
- TTS Model → Stateful Module: caches · Prompt-derived voice state is managed as reusable state
- TTS Model → Text Conditioner: conditions · Text and prompt conditioning are prepared before inference
- TTS Model → Audio I/O: loads audio · Voice prompts and input audio are preprocessed from files
- Audio I/O → Audio Utils: uses · Audio preprocessing relies on helper routines
- Text Conditioner → Conditioning Base: implements · Text conditioning follows shared interfaces
- Flow LM → Modules: built from · Flow model is assembled from reusable neural blocks
- Mimi Codec → Modules: built from · Codec uses the same portable building blocks
- Stateful Module → Weights Loader: restores · State and weights are loaded through dedicated loaders
- TTS Model → Config YAMLs: configured by · Model variants and language settings come from YAML configs
- Weights Loader → Model Assets: fetches · Loader can resolve pretrained and voice assets from external storage
- TTS Model → PyTorch: executes on · Inference runs on the PyTorch backend
- Modules → PyTorch: depends on · Neural building blocks are implemented for the runtime backend
- Quantization → PyTorch: optimizes · Quantization shapes the deployment runtime
- Main API → Quantization: exposes · API surface includes quantized execution support
- CLI entrypoint → Web UI: serves · Server mode presents the local web UI
- Web UI → Main API: calls · Web UI talks to the shared generation API