valenti/spark-pocket-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add .safetensors save/load to get_state_for_audio_prompt (#69)

Browse Source 
Quality of life improvement to save user voice embedding as a
.safetensors file that can be loaded for subsequent use, achieving the
same efficiency as the predefined voices, all via the existing
get_state_for_audio_prompt function.

---------

Co-authored-by: Václav Volhejn <8401624+vvolhejn@users.noreply.github.com>
This commit is contained in:
thataboy
2026-01-26 10:36:27 +01:00
committed by GitHub
co-authored by Václav Volhejn
parent bfbaf80437
commit edec7f8790
6 changed files with 290 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/python-api.md
+35 -2
View File
@@ -91,10 +91,10 @@ print(f"Sample rate: {model.sample_rate} Hz")
##### `get_state_for_audio_prompt(audio_conditioning, truncate=False)`
Extract model state for a given audio file or URL (voice cloning).
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 file path, URL, or tensor
- `audio_conditioning` (Path | str | torch.Tensor): Audio or .safetensors file path, URL, or tensor
- `truncate` (bool): Whether to truncate the audio (default: False)
**Returns:**
@@ -111,6 +111,9 @@ voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mack
# 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"
@@ -169,6 +172,35 @@ for chunk in model.generate_audio_stream(voice_state, "Long text content..."):
# Could save chunks to file or play in real-time
```
##### `save_audio_prompt(audio_conditioning, export_path, truncate=False)`
Save audio prompt to a .safetensors file.
**Parameters:**
- `audio_conditioning` (Path | str | torch.Tensor): Audio file path, URL, or tensor
- `export_path` (Path | str): .safetensors file path
- `truncate` (bool): Whether to truncate the audio (default: False)
**Returns:**
- tensor of the converted audio.
**Example:**
```python
from pocket_tts import TTSModel
model = TTSModel.load_model()
# From HuggingFace URL
model.get_state_for_audio_prompt(
"hf://kyutai/tts-voices/alba-mackenna/casual.wav", "casual.safetensors"
)
# From local file (the .safetensors extension will be added automatically)
tensor = model.get_state_for_audio_prompt("./my_voice.wav", "my_voice")
# Use the tensor, Luke!
audio = model.generate_audio(tensor, "Hello world!")
```
## Advanced Usage
### Voice Management
@@ -190,6 +222,7 @@ casual_audio = model.generate_audio(voices["casual"], "Hey there!")
funny_audio = model.generate_audio(voices["funny"], "Good morning.")
```
### Batch Processing
```python