Add functions to export and import the model state for a given voice (#129)

This commit is contained in:
Gabriel de Marmiesse
2026-02-16 15:37:00 +01:00
committed by GitHub
parent aca7dc8db6
commit 8443de8e06
11 changed files with 123 additions and 212 deletions
+19 -2
View File
@@ -113,11 +113,28 @@ audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
```
You can have multiple voice states around if
you have multiple voices you want to use. `load_model()`
You can have multiple voice states around if
you have multiple voices you want to use. `load_model()`
and `get_state_for_audio_prompt()` are relatively slow operations,
so we recommend to keep the model and voice states in memory if you can.
For faster voice loading, you can export voice states to safetensors files:
```python
from pocket_tts import TTSModel, export_model_state
model = TTSModel.load_model()
# Export a voice state for fast loading later
model_state = model.get_state_for_audio_prompt("some_voice.wav")
export_model_state(model_state, "./some_voice.safetensors")
# Later, load it quickly, this is quite fast as it's just reading the kvcache
# from disk and doesn't do any others computations.
model_state_copy = model.get_state_for_audio_prompt("./some_voice.safetensors")
audio = model.generate_audio(model_state_copy, "Hello world!")
```
You can check out the [Python API documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/python-api.md) for more details and examples.
## Unsupported features