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