From 8443de8e067f71d2da565eceddcb8212dc366eaf Mon Sep 17 00:00:00 2001 From: Gabriel de Marmiesse Date: Mon, 16 Feb 2026 15:37:00 +0100 Subject: [PATCH] Add functions to export and import the model state for a given voice (#129) --- README.md | 21 +++- docs/API Reference/python-api.md | 34 +++--- pocket_tts/__init__.py | 7 +- pocket_tts/main.py | 91 +--------------- pocket_tts/models/tts_model.py | 148 ++++++++------------------ pocket_tts/modules/stateful_module.py | 1 - pocket_tts/utils/utils.py | 2 +- pyproject.toml | 1 + scripts/generate_default_voices.sh | 26 +++++ tests/test_python_api.py | 2 +- uv.lock | 2 + 11 files changed, 123 insertions(+), 212 deletions(-) create mode 100755 scripts/generate_default_voices.sh diff --git a/README.md b/README.md index 8ad7754..b71a208 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docs/API Reference/python-api.md b/docs/API Reference/python-api.md index a473023..866e291 100644 --- a/docs/API Reference/python-api.md +++ b/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 -- `export_path` (Path | str): .safetensors file path -- `truncate` (bool): Whether to truncate the audio (default: False) - -**Returns:** -- tensor of the converted audio. +- `model_state` (dict): Model state dictionary from `get_state_for_audio_prompt()` +- `dest` (str | Path): Path to save the safetensors file **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( - "hf://kyutai/tts-voices/alba-mackenna/casual.wav", "casual.safetensors" + +# Get voice state from an audio file +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) -tensor = model.get_state_for_audio_prompt("./my_voice.wav", "my_voice") +# Export to safetensors for fast loading later +export_model_state(model_state_for_voice, "my_voice.safetensors") -# Use the tensor, Luke! -audio = model.generate_audio(tensor, "Hello world!") +# Quite fast, it's just loading the tensors without running any pytorch code +model_state_for_voice_copy = model.get_state_for_audio_prompt("my_voice.safetensors") ``` + ## Advanced Usage ### Voice Management diff --git a/pocket_tts/__init__.py b/pocket_tts/__init__.py index 4869107..0667888 100644 --- a/pocket_tts/__init__.py +++ b/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"] diff --git a/pocket_tts/main.py b/pocket_tts/main.py index 6b4de4f..a693776 100644 --- a/pocket_tts/main.py +++ b/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( - config, temperature, lsd_decode_steps, noise_clamp, eos_threshold + tts_model = TTSModel.load_model(config) + model_state = tts_model.get_state_for_audio_prompt( + audio_conditioning=audio_path, truncate=True ) - 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.") + export_model_state(model_state, export_path) if __name__ == "__main__": diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index c2e91a2..63a0444 100644 --- a/pocket_tts/models/tts_model.py +++ b/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,106 +740,12 @@ 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) - else: - 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("Encoding audio prompt"): - prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device)) - - model_state = init_states(self.flow_lm, batch_size=1, sequence_length=prompt.shape[1]) - - with display_execution_time("Prompting audio"): - self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt) - - logger.info( - "Size of the model state for audio prompt: %d MB", size_of_dict(model_state) // 1e6 - ) - - return model_state - - def _estimate_max_gen_len(self, token_count: int) -> int: - gen_len_sec = token_count / self._TOKENS_PER_SECOND_ESTIMATE + self._GEN_SECONDS_PADDING - 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") + return _import_model_state(download_if_necessary(PREDEFINED_VOICES[audio_conditioning])) if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)): raise ValueError(VOICE_CLONING_UNSUPPORTED) @@ -853,13 +766,24 @@ class TTSModel(nn.Module): audio, conditioning_sample_rate, self.config.mimi.sample_rate, 1 ) - with display_execution_time("Exporting audio prompt"): + 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}, export_path) + model_state = init_states(self.flow_lm, batch_size=1, sequence_length=prompt.shape[1]) - return audio_conditioning + with display_execution_time("Prompting audio"): + self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt) + + logger.info( + "Size of the model state for audio prompt: %d MB", size_of_dict(model_state) // 1e6 + ) + + return model_state + + def _estimate_max_gen_len(self, token_count: int) -> int: + gen_len_sec = token_count / self._TOKENS_PER_SECOND_ESTIMATE + self._GEN_SECONDS_PADDING + frame_rate = self.config.mimi.frame_rate + return math.ceil(gen_len_sec * frame_rate) def prepare_text_prompt(text: str) -> tuple[str, int]: @@ -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 diff --git a/pocket_tts/modules/stateful_module.py b/pocket_tts/modules/stateful_module.py index ff9edaf..3d030aa 100644 --- a/pocket_tts/modules/stateful_module.py +++ b/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 diff --git a/pocket_tts/utils/utils.py b/pocket_tts/utils/utils.py index 83ba38d..125f395 100644 --- a/pocket_tts/utils/utils.py +++ b/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 } diff --git a/pyproject.toml b/pyproject.toml index fb1d516..53b86ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dev = [ "mkdocstrings-python>=2.0.1", "pytest>=9.0.2", "pytest-xdist>=3.8.0", + "soundfile", ] diff --git a/scripts/generate_default_voices.sh b/scripts/generate_default_voices.sh new file mode 100755 index 0000000..3bdd6ac --- /dev/null +++ b/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" diff --git a/tests/test_python_api.py b/tests/test_python_api.py index 5740d7d..9abf056 100644 --- a/tests/test_python_api.py +++ b/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(): diff --git a/uv.lock b/uv.lock index 9655b73..7483802 100644 --- a/uv.lock +++ b/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]]