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

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>

docs/export_voice.md

@@ -0,0 +1,47 @@
+# Export Voice Command Documentation
+
+Kyutai Pocket TTS allows you to generate speech from an audio sample. However, processing an audio file each time is relatively slow and inefficient.
+
+The `export-voice` command allows you to convert an audio file to a voice embedding in safetensors format. The safetensors file can then be loaded very quickly whenever you generate speech.
+
+## Basic Usage
+
+```bash
+uvx pocket-tts export-voice audio-path export-path
+# or if installed manually:
+pocket-tts export-voice audio-path export-path
+```
+
+## Command Options
+
+### Required Parameters
+
+- `audio-path`: Path of the audio file or directory to convert. `audio-path` can point to an `http:` or `hf:` (hugging face) file. If `audio-path` is a local directory, all audio files found inside it will be batch converted. Supports popular audio file formats like wav and mp3.
+- `export-path`: Path of safetensors file or directory to export to. For batch conversion, export-path should be a directory. The directory will be created if it does not exist.
+
+### Options
+
+- `--truncate`: Automatically truncate long audio files down to 30 seconds.
+
+The other parameters such as `--lsd-decode-steps` and `--temperature` are the same as for the `generate` command. See the [generate documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/generate.md) for more details.
+
+## Examples
+
+```bash
+# export a single file
+pocket-tts export-voice voice_memo127762.mp3 jack.safetensors
+
+# export a single file to a different directory (output is embbeddings/mary.safetensors
+pocket-tts export-voice voices/mary.wav embeddings/
+
+# export an entire directory of audio files, truncate long audios
+pocket-tts export-voice voices/ embeddings/ --truncate
+
+# export an online file to current directory
+pocket-tts export-voice https://huggingface.co/kyutai/tts-voices/resolve/main/alba-mackenna/announcer.wav .
+
+# use the exported safetensors
+pocket-tts generate --text "Hello, welcome to today's game between the Bears and Cubs."  --voice announcer.safetensors
+```
+
+Note: to indicate a directory rather than a file, please be sure to include a trailing / (\ on Windows).

docs/generate.md

@@ -57,8 +57,12 @@ pocket-tts generate --voice "hf://kyutai/tts-voices/jessica-jian/casual.wav"
 
 # Use local voice file
 pocket-tts generate --voice "./my_voice.wav"
+
+# Use a safetensors file (such as one created using `pocket-tts export-voice`)
+pocket-tts generate --voice "./my_voice.safetensors"
 ```
 
+
 ### Quality Tuning
 
 ```bash

docs/python-api.md

@@ -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