Add .safetensors save/load to get_state_for_audio_prompt (#69)

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>
2026-01-26 10:36:27 +01:00
parent bfbaf80437
commit edec7f8790
6 changed files with 290 additions and 19 deletions
@@ -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()