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

README.md

@@ -118,6 +118,23 @@ 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

docs/API Reference/python-api.md

@@ -172,35 +172,37 @@ 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.
+## 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:**
-- `audio_conditioning` (Path | str | torch.Tensor): Audio file path, URL, or tensor
+- `model_state` (dict): Model state dictionary from `get_state_for_audio_prompt()`
-- `export_path` (Path | str): .safetensors file path
+- `dest` (str | Path): Path to save the safetensors file
-- `truncate` (bool): Whether to truncate the audio (default: False)
-
-**Returns:**
-- tensor of the converted audio.
 
 **Example:**
 ```python
-from pocket_tts import TTSModel
+from pocket_tts import TTSModel, export_model_state
 
 model = TTSModel.load_model()
-# From HuggingFace URL
+
-model.get_state_for_audio_prompt(
+# Get voice state from an audio file
-    "hf://kyutai/tts-voices/alba-mackenna/casual.wav", "casual.safetensors"
+model_state_for_voice = model.get_state_for_audio_prompt(
+    "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
 )
 
-# From local file (the .safetensors extension will be added automatically)
+# Export to safetensors for fast loading later
-tensor = model.get_state_for_audio_prompt("./my_voice.wav", "my_voice")
+export_model_state(model_state_for_voice, "my_voice.safetensors")
 
-# Use the tensor, Luke!
+# Quite fast, it's just loading the tensors without running any pytorch code
-audio = model.generate_audio(tensor, "Hello world!")
+model_state_for_voice_copy = model.get_state_for_audio_prompt("my_voice.safetensors")
 ```
 
+
 ## Advanced Usage
 
 ### Voice Management

pocket_tts/__init__.py

@@ -3,7 +3,10 @@ from beartype.claw import beartype_this_package
 
 beartype_this_package(conf=BeartypeConf(is_color=False))
 
-from pocket_tts.models.tts_model import TTSModel  # noqa: E402
+from pocket_tts.models.tts_model import (  # noqa: E402
+    TTSModel,
+    export_model_state,
+)
 
 # Public methods:
 # TTSModel.device
@@ -13,4 +16,4 @@ from pocket_tts.models.tts_model import TTSModel  # noqa: E402
 # TTSModel.generate_audio_stream
 # TTSModel.get_state_for_audio_prompt
 
-__all__ = ["TTSModel"]
+__all__ = ["TTSModel", "export_model_state"]

pocket_tts/main.py

@@ -25,7 +25,7 @@ from pocket_tts.default_parameters import (
     DEFAULT_VARIANT,
     MAX_TOKEN_PER_CHUNK,
 )
-from pocket_tts.models.tts_model import TTSModel
+from pocket_tts.models.tts_model import TTSModel, export_model_state
 from pocket_tts.utils.logging_utils import enable_logging
 from pocket_tts.utils.utils import PREDEFINED_VOICES, size_of_dict
 
@@ -286,99 +286,18 @@ def export_voice(
         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,
     config: Annotated[str, typer.Option(help="Model config path or 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
 
     log_level = logging.ERROR if quiet else logging.INFO
-    success_count = 0
-
     with enable_logging("pocket_tts", log_level):
-        tts_model = TTSModel.load_model(
+        tts_model = TTSModel.load_model(config)
-            config, temperature, lsd_decode_steps, noise_clamp, eos_threshold
+        model_state = tts_model.get_state_for_audio_prompt(
+            audio_conditioning=audio_path, truncate=True
         )
-        tts_model.to(device)
+        export_model_state(model_state, export_path)
-
-        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__":

pocket_tts/models/tts_model.py

@@ -10,6 +10,7 @@ from functools import lru_cache
 from pathlib import Path
 
 import safetensors
+import safetensors.torch
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -31,13 +32,12 @@ from pocket_tts.models.mimi import MimiModel
 from pocket_tts.modules import mimi_transformer
 from pocket_tts.modules.dummy_quantizer import DummyQuantizer
 from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
-from pocket_tts.modules.stateful_module import increment_steps, init_states
+from pocket_tts.modules.stateful_module import StatefulModule, increment_steps, init_states
 from pocket_tts.utils.config import Config, load_config
 from pocket_tts.utils.utils import (
     PREDEFINED_VOICES,
     display_execution_time,
     download_if_necessary,
-    load_predefined_voice,
     size_of_dict,
 )
 from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_state_dict
@@ -175,6 +175,13 @@ class TTSModel(nn.Module):
         size_in_mb = size_of_dict(tts_model.state_dict()) // 1e6
         logging.info(f"TTS Model loaded successfully. Its size is {size_in_mb} MB")
 
+        # TODO: move this in the __init__ and make self.mimi in __init__
+        for top_module in (tts_model.flow_lm, tts_model.mimi):
+            for module_name, module in top_module.named_modules():
+                if not isinstance(module, StatefulModule):
+                    continue
+                module._module_absolute_name = module_name
+
         return tts_model
 
     @classmethod
@@ -733,13 +740,13 @@ class TTSModel(nn.Module):
         ):
             if isinstance(audio_conditioning, str):
                 audio_conditioning = download_if_necessary(audio_conditioning)
-            import safetensors.torch
 
-            prompt = safetensors.torch.load_file(audio_conditioning)["audio_prompt"]
+            return _import_model_state(audio_conditioning)
+
         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)
+            return _import_model_state(download_if_necessary(PREDEFINED_VOICES[audio_conditioning]))
-        else:
+
         if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
             raise ValueError(VOICE_CLONING_UNSUPPORTED)
 
@@ -778,89 +785,6 @@ class TTSModel(nn.Module):
         frame_rate = self.config.mimi.frame_rate
         return math.ceil(gen_len_sec * frame_rate)
 
-    @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.
-
-        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!")
-            ```
-
-        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()
@@ -940,3 +864,21 @@ def split_into_best_sentences(tokenizer, text_to_generate: str, max_tokens: int)
         chunks.append(current_chunk.strip())
 
     return chunks
+
+
+def export_model_state(model_state: dict[str, dict[str, torch.Tensor]], dest: str | Path):
+    dict_to_store = {}
+    for module_name, module_state in model_state.items():
+        for key, tensor_value in module_state.items():
+            dict_to_store[f"{module_name}/{key}"] = tensor_value
+    safetensors.torch.save_file(dict_to_store, dest)
+
+
+def _import_model_state(source: str | Path) -> dict[str, dict[str, torch.Tensor]]:
+    result = {}
+    with safetensors.safe_open(source, framework="pt") as f:
+        for key in f.keys():
+            module_name, tensor_key = key.split("/")
+            result.setdefault(module_name, {})
+            result[module_name][tensor_key] = f.get_tensor(key)
+    return result

pocket_tts/modules/stateful_module.py

@@ -11,7 +11,6 @@ def init_states(
     for module_name, module in model.named_modules():
         if not isinstance(module, StatefulModule):
             continue
-        module._module_absolute_name = module_name
         module_state = module.init_state(batch_size, sequence_length=sequence_length)
         result[module_name] = module_state
     return result

pocket_tts/utils/utils.py

@@ -14,7 +14,7 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent
 _voices_names = ["alba", "marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma"]
 PREDEFINED_VOICES = {
     # don't forget to change this
-    x: f"hf://kyutai/pocket-tts-without-voice-cloning/embeddings/{x}.safetensors@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3"
+    x: f"hf://kyutai/pocket-tts-without-voice-cloning/embeddings_v2/{x}.safetensors@2578fed2380333b621689eaed6fe144cf69dfeb3"
     for x in _voices_names
 }
 

pyproject.toml

@@ -37,6 +37,7 @@ dev = [
     "mkdocstrings-python>=2.0.1",
     "pytest>=9.0.2",
     "pytest-xdist>=3.8.0",
+    "soundfile",
 ]
 
 

scripts/generate_default_voices.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -ex
+# Generate safetensors files for all default voices
+
+OUTPUT_DIR="${1:-./default_voices}"
+mkdir -p "$OUTPUT_DIR"
+
+# Default voices with their Hugging Face paths
+# The marius voice can benefit from a bit of cleaning
+declare -A voices=(
+    ["alba"]="hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    ["marius"]="hf://kyutai/tts-voices/voice-donations/Selfie.wav"
+    ["javert"]="hf://kyutai/tts-voices/voice-donations/Butter.wav"
+    ["jean"]="hf://kyutai/tts-voices/ears/p010/freeform_speech_01.wav"
+    ["fantine"]="hf://kyutai/tts-voices/vctk/p244_023.wav"
+    ["cosette"]="hf://kyutai/tts-voices/expresso/ex04-ex02_confused_001_channel1_499s.wav"
+    ["eponine"]="hf://kyutai/tts-voices/vctk/p262_023.wav"
+    ["azelma"]="hf://kyutai/tts-voices/vctk/p303_023.wav"
+)
+
+for voice in "${!voices[@]}"; do
+    echo "Exporting voice: $voice"
+    uv run pocket-tts export-voice "${voices[$voice]}" "$OUTPUT_DIR/$voice.safetensors"
+done
+
+echo "Done! Voice files saved to $OUTPUT_DIR"

tests/test_python_api.py

@@ -6,7 +6,7 @@ from pocket_tts.models.tts_model import TTSModel as TTSModelImpl
 
 
 def test_public_api_exports_only_tts_model():
-    assert pocket_tts.__all__ == ["TTSModel"]
+    assert pocket_tts.__all__ == ["TTSModel", "export_model_state"]
 
 
 def test_public_api_tts_model_points_to_implementation():

uv.lock

@@ -1144,6 +1144,7 @@ dev = [
     { name = "mkdocstrings-python" },
     { name = "pytest" },
     { name = "pytest-xdist" },
+    { name = "soundfile" },
 ]
 
 [package.metadata]
@@ -1177,6 +1178,7 @@ dev = [
     { name = "mkdocstrings-python", specifier = ">=2.0.1" },
     { name = "pytest", specifier = ">=9.0.2" },
     { name = "pytest-xdist", specifier = ">=3.8.0" },
+    { name = "soundfile" },
 ]
 
 [[package]]