From edec7f879089c217d40ce5de6bb70fbeca473ca9 Mon Sep 17 00:00:00 2001 From: thataboy Date: Mon, 26 Jan 2026 01:36:27 -0800 Subject: [PATCH] Add .safetensors save/load to get_state_for_audio_prompt (#69) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 5 ++ docs/export_voice.md | 47 ++++++++++++++ docs/generate.md | 4 ++ docs/python-api.md | 37 ++++++++++- pocket_tts/main.py | 110 +++++++++++++++++++++++++++++++++ pocket_tts/models/tts_model.py | 106 ++++++++++++++++++++++++++----- 6 files changed, 290 insertions(+), 19 deletions(-) create mode 100644 docs/export_voice.md diff --git a/README.md b/README.md index 8b0820c..3c52e36 100644 --- a/README.md +++ b/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). diff --git a/docs/export_voice.md b/docs/export_voice.md new file mode 100644 index 0000000..a19071b --- /dev/null +++ b/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). diff --git a/docs/generate.md b/docs/generate.md index 6fc614e..739ec78 100644 --- a/docs/generate.md +++ b/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 diff --git a/docs/python-api.md b/docs/python-api.md index 66b405a..5fcf2dd 100644 --- a/docs/python-api.md +++ b/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 diff --git a/pocket_tts/main.py b/pocket_tts/main.py index 865cb35..12db058 100644 --- a/pocket_tts/main.py +++ b/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() diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 8ca017b..d733f95 100644 --- a/pocket_tts/models/tts_model.py +++ b/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()