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>

README.md

@@ -77,6 +77,11 @@ Navigate to `http://localhost:8000` to try the web interface, it's faster than t
 
 You can check out the [serve documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/serve.md) for more details and examples.
 
+### The `export-voice` command
+
+Processing an audio file (e.g., a .wav or .mp3) for voice cloning is relatively slow, but loading a safetensors file -- a voice embedding converted from an audio file -- is very fast. You can use the `export-voice` command to do this conversion. See the [export-voice documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/export_voice.md) for more details and examples.
+
+
 ## Using it as a Python library
 
 You can try out the Python library on Colab [here](https://colab.research.google.com/github/kyutai-labs/pocket-tts/blob/main/docs/pocket-tts-example.ipynb).

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

pocket_tts/main.py

@@ -263,5 +263,115 @@ def generate(
         )
 
 
+# ----------------------------------------------
+# export audio to safetensors CLI implementation
+# ----------------------------------------------
+
+
+@cli_app.command()
+def export_voice(
+    audio_path: Annotated[
+        str, typer.Argument(help="Audio file or directory to convert and export")
+    ],
+    export_path: Annotated[str, typer.Argument(help="Output file or directory")],
+    truncate: Annotated[
+        bool, typer.Option("-tr", "--truncate", help="Truncate long audio")
+    ] = False,
+    quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False,
+    variant: Annotated[str, typer.Option(help="Model signature")] = DEFAULT_VARIANT,
+    lsd_decode_steps: Annotated[
+        int, typer.Option(help="Number of generation steps")
+    ] = DEFAULT_LSD_DECODE_STEPS,
+    temperature: Annotated[
+        float, typer.Option(help="Temperature for generation")
+    ] = DEFAULT_TEMPERATURE,
+    noise_clamp: Annotated[float, typer.Option(help="Noise clamp value")] = DEFAULT_NOISE_CLAMP,
+    eos_threshold: Annotated[float, typer.Option(help="EOS threshold")] = DEFAULT_EOS_THRESHOLD,
+    frames_after_eos: Annotated[
+        int, typer.Option(help="Number of frames to generate after EOS")
+    ] = DEFAULT_FRAMES_AFTER_EOS,
+    device: Annotated[str, typer.Option(help="Device to use")] = "cpu",
+):
+    """Convert and save audio to .safetensors file"""
+    import re
+
+    def url(path):
+        return path.startswith(("http:", "https:", "hf:"))
+
+    def normalize_url(url):
+        # utils.py expects urls to be xxx:// so normalize them
+        return re.sub(r"^(http|https|hf)\:\/*(.+)$", r"\1://\2", url)
+
+    def likely_file(path):
+        return not url(path) and not likely_dir(path)
+
+    def likely_dir(path):
+        return not url(path) and (path.endswith(("/", "\\")) or path == ".")
+
+    def convert_one(in_path, out_path, join_path):
+        """helper convert function"""
+        voice = in_path.stem
+        if url(str(in_path)):
+            in_path = normalize_url(str(in_path))
+        if join_path:
+            out_path = out_path / f"{voice}.safetensors"
+        else:
+            # ensure output file has correct extension
+            out_path = out_path.with_suffix(".safetensors")
+        try:
+            tts_model.save_audio_prompt(in_path, out_path, truncate)
+        except Exception as e:
+            logger.error(f"❌ Unable to export voice '{in_path}': {e}")
+            return False
+        logger.info(f"✅ Successfully exported voice '{voice}' to '{out_path}'")
+        return True
+
+    if "cuda" in device:
+        # Cuda graphs capturing does not play nice with multithreading.
+        os.environ["NO_CUDA_GRAPH"] = "1"
+
+    log_level = logging.ERROR if quiet else logging.INFO
+    success_count = 0
+
+    with enable_logging("pocket_tts", log_level):
+        tts_model = TTSModel.load_model(
+            variant, temperature, lsd_decode_steps, noise_clamp, eos_threshold
+        )
+        tts_model.to(device)
+
+        in_path = Path(audio_path)
+        out_path = Path(export_path)
+        if likely_dir(export_path):
+            # make sure output dir exists
+            out_path.mkdir(parents=True, exist_ok=True)
+
+        if likely_dir(audio_path):  # batch convert whole directory
+            if not in_path.is_dir():
+                logger.error(f"Input dir '{audio_path}' does not exists")
+                exit(1)
+            if not likely_dir(export_path):
+                # batch convert, output path must be directory, not file
+                out_path = Path("./")
+            for path in Path(in_path).iterdir():
+                if path.is_file() and path.suffix.lower() in [
+                    ".wav",
+                    ".mp3",
+                    ".flac",
+                    ".ogg",
+                    ".aiff",
+                ]:
+                    if convert_one(path, out_path, True):
+                        success_count += 1
+        else:  # convert single file
+            if likely_file(audio_path) and not in_path.exists():
+                logger.error(f"Input file '{in_path}'' does not exists")
+                exit(1)
+            if convert_one(in_path, out_path, likely_dir(export_path)):
+                success_count += 1
+
+        if success_count > 0:
+            logger.info(f"🎉 Successfully exported {success_count} voices.")
+
+
 if __name__ == "__main__":
     cli_app()

pocket_tts/models/tts_model.py

@@ -44,6 +44,15 @@ from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_st
 torch.set_num_threads(1)
 logger = logging.getLogger(__name__)
 
+VOICE_CLONING_UNSUPPORTED = (
+    f"We could not download the weights for the model with voice cloning, "
+    f"but you're trying to use voice cloning. "
+    f"Without voice cloning, you can use our catalog of voices {list(PREDEFINED_VOICES)}. "
+    f"If you want access to the model with voice cloning, go to "
+    f"https://huggingface.co/kyutai/pocket-tts and accept the terms, "
+    f"then make sure you're logged in locally with `uvx hf auth login`."
+)
+
 
 class TTSModel(nn.Module):
     def __init__(
@@ -637,9 +646,9 @@ class TTSModel(nn.Module):
         enables voice cloning and audio continuation with speaker consistency.
 
         Args:
-            audio_conditioning: Audio prompt to condition on. Can be:
-                - Path: Local file path to audio file
-                - str: URL to download audio file from
+            audio_conditioning: Audio prompt to condition (or .safetensors to load). Can be:
+                - Path: Local file path to audio file (or .safetensors)
+                - str: URL to download audio file (or .safetensors) from
                 - torch.Tensor: Pre-loaded audio tensor with shape [channels, samples]
             truncate: Whether to truncate long audio prompts to 30 seconds.
                 Helps prevent memory issues with very long inputs. Defaults to False.
@@ -662,19 +671,21 @@ class TTSModel(nn.Module):
             - Processing time is logged for performance monitoring
             - The state preserves speaker characteristics for voice cloning
         """
-        if isinstance(audio_conditioning, str) and audio_conditioning in PREDEFINED_VOICES:
+        if isinstance(audio_conditioning, (str, Path)) and str(audio_conditioning).endswith(
+            ".safetensors"
+        ):
+            if isinstance(audio_conditioning, str):
+                audio_conditioning = download_if_necessary(audio_conditioning)
+            import safetensors.torch
+
+            prompt = safetensors.torch.load_file(audio_conditioning)["audio_prompt"]
+        elif isinstance(audio_conditioning, str) and audio_conditioning in PREDEFINED_VOICES:
             # We get the audio conditioning directly from the safetensors file.
             prompt = load_predefined_voice(audio_conditioning)
         else:
             if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
-                raise ValueError(
-                    f"We could not download the weights for the model with voice cloning, "
-                    f"but you're trying to use voice cloning. "
-                    f"Without voice cloning, you can use our catalog of voices {list(PREDEFINED_VOICES)}. "
-                    f"If you want access to the model with voice cloning, go to "
-                    f"https://huggingface.co/kyutai/pocket-tts and accept the terms, "
-                    f"then make sure you're logged in locally with `uvx hf auth login`."
-                )
+                raise ValueError(VOICE_CLONING_UNSUPPORTED)
+
             if isinstance(audio_conditioning, str):
                 audio_conditioning = download_if_necessary(audio_conditioning)
 
@@ -693,11 +704,6 @@ class TTSModel(nn.Module):
 
             with display_execution_time("Encoding audio prompt"):
                 prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device))
-                # import safetensors.torch
-                # safetensors.torch.save_file(
-                #     {"audio_prompt": prompt},
-                #     "/projects/huggingface/pocket-tts/embeddings/cosette.safetensors"
-                # )
 
         model_state = init_states(self.flow_lm, batch_size=1, sequence_length=1000)
 
@@ -710,6 +716,72 @@ class TTSModel(nn.Module):
 
         return model_state
 
+    @torch.no_grad
+    def save_audio_prompt(
+        self,
+        audio_conditioning: Path | str | torch.Tensor,
+        export_path: Path | str,
+        truncate: bool = False,
+    ) -> torch.Tensor:
+        """Save audio prompt to .safetensors file
+
+        This method processes an audio prompt and exports it to a .safetensors file,
+        which can be loaded by get_state_for_audio_prompt in subsequent uses
+        without converting the audio again.
+
+        It also takes an already converted audio tensor and exports it as a .safetensors file
+
+        Args:
+            audio_conditioning: Audio to export
+                - Path: Local file path to audio file
+                - str: URL to download audio file
+                - torch.Tensor: Pre-loaded audio tensor with shape [channels, samples]
+            export_path: Path to output file
+            truncate: Whether to truncate long audio prompts to 30 seconds.
+
+        Returns:
+            Audio tensor of converted audio
+
+        Raises:
+            FileNotFoundError: If audio file path doesn't exist.
+            ValueError: If audio tensor export path is invalid or empty.
+            RuntimeError: If audio processing or encoding fails.
+
+        Note:
+            - Send resulting audio tensor to get_state_for_audio_prompt
+              in order to get the model state for generation.
+        """
+        if not export_path or not isinstance(export_path, (str, Path)):
+            raise ValueError("export_path must be of type str or Path")
+        export_path = Path(export_path).with_suffix(".safetensors")
+
+        if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
+            raise ValueError(VOICE_CLONING_UNSUPPORTED)
+
+        if isinstance(audio_conditioning, str):
+            audio_conditioning = download_if_necessary(audio_conditioning)
+
+        if isinstance(audio_conditioning, Path):
+            audio, conditioning_sample_rate = audio_read(audio_conditioning)
+
+            if truncate:
+                max_samples = int(30 * conditioning_sample_rate)  # 30 seconds of audio
+                if audio.shape[-1] > max_samples:
+                    audio = audio[..., :max_samples]
+                    logger.info(f"Audio truncated to first 30 seconds ({max_samples} samples)")
+
+            audio_conditioning = convert_audio(
+                audio, conditioning_sample_rate, self.config.mimi.sample_rate, 1
+            )
+
+        with display_execution_time("Exporting audio prompt"):
+            prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device))
+            import safetensors.torch
+
+            safetensors.torch.save_file({"audio_prompt": prompt}, export_path)
+
+        return audio_conditioning
+
 
 def prepare_text_prompt(text: str) -> tuple[str, int]:
     text = text.strip()