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

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

Browse Source
This commit is contained in:
Gabriel de Marmiesse
2026-02-16 15:37:00 +01:00
committed by GitHub
parent aca7dc8db6
commit 8443de8e06
11 changed files with 123 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+17
View File
@@ -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
+18 -16
View File
@@ -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
pocket_tts/__init__.py
+5 -2
View File
@@ -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
+5 -86
View File
@@ -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__":
pocket_tts/models/tts_model.py
+45 -103
View File
@@ -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
pocket_tts/modules/stateful_module.py
-1
View File
@@ -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
+1 -1
View File
@@ -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
+1
View File
@@ -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
Executable
+26
View File
@@ -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
+1 -1
View File
@@ -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
Generated
+2
View File
@@ -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]]