## Motivation Human evaluations consistently prefer the English model at temperature **0.3** over the current default **0.7**, and the change is free on every objective axis we measured.
263 lines
7.8 KiB
Markdown
263 lines
7.8 KiB
Markdown
# Python API Documentation
|
|
|
|
Kyutai Pocket TTS provides a Python API for integrating text-to-speech capabilities into your applications.
|
|
|
|
## Installation
|
|
|
|
```bash
|
|
pip install pocket-tts
|
|
```
|
|
|
|
## Quick Start
|
|
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
import scipy.io.wavfile
|
|
|
|
# Load the model
|
|
tts_model = TTSModel.load_model()
|
|
|
|
# Get voice state from an audio file
|
|
voice_state = tts_model.get_state_for_audio_prompt(
|
|
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
|
)
|
|
|
|
# Generate audio
|
|
audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
|
|
|
|
# Save to file
|
|
scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
|
|
```
|
|
|
|
## Core Classes
|
|
|
|
### TTSModel
|
|
|
|
The main class for text-to-speech generation.
|
|
|
|
#### Class Methods
|
|
|
|
##### `load_model(language=None, config=None, temp=None, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0, quantize=False)`
|
|
|
|
Load and return a TTSModel instance with pre-trained weights.
|
|
|
|
**Parameters:**
|
|
- `language` (str | None): Name of built-in language config to load. Supported values: `"english_2026-01"`, `"english_2026-04"`, `"english"`, `"french_24l"`, `"german_24l"`, `"portuguese_24l"`, `"italian_24l"`, `"spanish_24l"`. If both `language` and `config` are omitted, defaults to `"english"`, which is the same model as `"english_2026-04"`. The "24l" variants are larger models that are not distilled yet and are here only as a preview.
|
|
- `config` (str | None): Path to model config YAML file. Incompatible with `language`.
|
|
- `temp` (float | None): Sampling temperature for generation. `None` uses the model's recommended default from its config file (`default_temperature`; 0.3 for the English model, 0.7 otherwise).
|
|
- `lsd_decode_steps` (int): Number of generation steps (default: 1)
|
|
- `noise_clamp` (float | None): Maximum value for noise sampling (default: None)
|
|
- `eos_threshold` (float): Threshold for end-of-sequence detection (default: -4.0)
|
|
- `quantize` (bool): Enable int8 quantization when loading the model (default: `False`)
|
|
|
|
**Returns:**
|
|
- `TTSModel`: Loaded model instance on CPU
|
|
|
|
**Example:**
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
# Load with default settings
|
|
model = TTSModel.load_model()
|
|
|
|
# Load with custom parameters
|
|
model = TTSModel.load_model(
|
|
language="english_2026-01", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0
|
|
)
|
|
```
|
|
|
|
#### Properties
|
|
|
|
##### `device` (str)
|
|
|
|
Returns the device type where the model is running ("cpu" or "cuda").
|
|
By default, the model runs on CPU.
|
|
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
model = TTSModel.load_model()
|
|
print(f"Model running on: {model.device}")
|
|
```
|
|
|
|
##### `sample_rate` (int)
|
|
|
|
Returns the generated audio sample rate (typically 24000 Hz).
|
|
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
model = TTSModel.load_model()
|
|
print(f"Sample rate: {model.sample_rate} Hz")
|
|
```
|
|
|
|
#### Methods
|
|
|
|
##### `get_state_for_audio_prompt(audio_conditioning, truncate=False)`
|
|
|
|
Extract model state for a given audio file or URL (voice cloning), or load from a .safetensors file.
|
|
|
|
**Parameters:**
|
|
- `audio_conditioning` (Path | str | torch.Tensor): Audio or .safetensors file path, URL, or tensor
|
|
- `truncate` (bool): Whether to truncate the audio (default: False)
|
|
|
|
**Returns:**
|
|
- `dict`: Model state dictionary containing hidden states and positional information
|
|
|
|
**Example:**
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
model = TTSModel.load_model()
|
|
# From HuggingFace URL
|
|
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
|
|
|
# From local file
|
|
voice_state = model.get_state_for_audio_prompt("./my_voice.wav")
|
|
|
|
# Reload state from a .safetensors file (much faster than extracting from an audio file)
|
|
voice_state = model.get_state_for_audio_prompt("./my_voices.safetensors")
|
|
|
|
# From HTTP URL
|
|
voice_state = model.get_state_for_audio_prompt(
|
|
"https://huggingface.co/kyutai/tts-voices/resolve"
|
|
"/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
|
)
|
|
```
|
|
|
|
##### `generate_audio(model_state, text_to_generate, frames_after_eos=None, copy_state=True)`
|
|
|
|
Generate complete audio tensor from text input.
|
|
|
|
**Parameters:**
|
|
- `model_state` (dict): Model state from `get_state_for_audio_prompt()`
|
|
- `text_to_generate` (str): Text to convert to speech
|
|
- `frames_after_eos` (int | None): Frames to generate after EOS detection (default: None)
|
|
- `copy_state` (bool): Whether to copy the state (default: True)
|
|
|
|
**Returns:**
|
|
- `torch.Tensor`: Audio 1D tensor with shape [samples]
|
|
|
|
**Example:**
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
model = TTSModel.load_model()
|
|
|
|
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
|
|
|
# Generate audio
|
|
audio = model.generate_audio(voice_state, "Hello world!", frames_after_eos=2, copy_state=True)
|
|
|
|
print(f"Generated audio shape: {audio.shape}")
|
|
print(f"Audio duration: {audio.shape[-1] / model.sample_rate:.2f} seconds")
|
|
```
|
|
|
|
##### `generate_audio_stream(model_state, text_to_generate, frames_after_eos=None, copy_state=True)`
|
|
|
|
Generate audio streaming chunks from text input.
|
|
|
|
**Parameters:** Same as `generate_audio()`
|
|
|
|
**Yields:**
|
|
- `torch.Tensor`: Audio chunks with shape [samples]
|
|
|
|
**Example:**
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
model = TTSModel.load_model()
|
|
|
|
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
|
# Stream generation
|
|
for chunk in model.generate_audio_stream(voice_state, "Long text content..."):
|
|
# Process each chunk as it's generated
|
|
print(f"Generated chunk: {chunk.shape[0]} samples")
|
|
# Could save chunks to file or play in real-time
|
|
```
|
|
|
|
|
|
## Functions
|
|
|
|
### export_model_state
|
|
|
|
Export a model state for a given voice conditioning to a safetensors file for fast loading later.
|
|
You can then load it again with the method `get_state_for_audio_prompt()`.
|
|
|
|
**Parameters:**
|
|
- `model_state` (dict): Model state dictionary from `get_state_for_audio_prompt()`
|
|
- `dest` (str | Path): Path to save the safetensors file
|
|
|
|
**Example:**
|
|
```python
|
|
from pocket_tts import TTSModel, export_model_state
|
|
|
|
model = TTSModel.load_model()
|
|
|
|
# Get voice state from an audio file
|
|
model_state_for_voice = model.get_state_for_audio_prompt(
|
|
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
|
)
|
|
|
|
# Export to safetensors for fast loading later
|
|
export_model_state(model_state_for_voice, "my_voice.safetensors")
|
|
|
|
# Quite fast, it's just loading the tensors without running any pytorch code
|
|
model_state_for_voice_copy = model.get_state_for_audio_prompt("my_voice.safetensors")
|
|
```
|
|
|
|
|
|
## Advanced Usage
|
|
|
|
### Voice Management
|
|
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
|
|
model = TTSModel.load_model()
|
|
# Preload multiple voices
|
|
voices = {
|
|
"casual": model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav"),
|
|
"funny": model.get_state_for_audio_prompt(
|
|
"https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
|
),
|
|
}
|
|
|
|
# Generate with different voices
|
|
casual_audio = model.generate_audio(voices["casual"], "Hey there!")
|
|
funny_audio = model.generate_audio(voices["funny"], "Good morning.")
|
|
```
|
|
|
|
|
|
### Batch Processing
|
|
|
|
```python
|
|
from pocket_tts import TTSModel
|
|
import scipy.io.wavfile
|
|
import torch
|
|
|
|
model = TTSModel.load_model()
|
|
|
|
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
|
# Process multiple texts efficiently by re-using the same voice state
|
|
texts = [
|
|
"First sentence to generate.",
|
|
"Second sentence to generate.",
|
|
"Third sentence to generate.",
|
|
]
|
|
|
|
audios = []
|
|
for text in texts:
|
|
audio = model.generate_audio(voice_state, text)
|
|
audios.append(audio)
|
|
|
|
# Concatenate all audio
|
|
full_audio = torch.cat(audios, dim=0)
|
|
scipy.io.wavfile.write("batch_output.wav", model.sample_rate, full_audio.numpy())
|
|
```
|
|
|
|
### Streaming to File
|
|
You can refer to our CLI implementation which can stream audio to a wav file.
|
|
|
|
For more information about the command-line interface, see the [Generate Documentation](generate.md) or [Serve Documentation](serve.md).
|