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

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__":