diff --git a/docs/generate.md b/docs/generate.md index 739ec78..a8504db 100644 --- a/docs/generate.md +++ b/docs/generate.md @@ -19,6 +19,7 @@ This will generate a WAV file `./tts_output.wav` with the default text and voice - `--text TEXT`: Text to generate (default: "Hello world! I am Kyutai Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.") - `--voice VOICE`: Path to audio conditioning file (voice to clone) (default: "hf://kyutai/tts-voices/alba-mackenna/casual.wav"). Urls and local paths are supported. - `--output-path OUTPUT_PATH`: Output path for generated audio (default: "./tts_output.wav") +- `--confg CONFIG_PATH`: Path to custom config.yaml (for loading local model files) ### Generation Parameters @@ -76,6 +77,19 @@ pocket-tts generate --temperature 1.0 pocket-tts generate --eos-threshold -3.0 ``` +### Custom Model Config + +If you'd like to override the paths from which the models are loaded, you can provide a custom YAML configuration. + +Copy pocket_tts/config/b6369a24.yaml and change weights_path:, weights_path_without_voice_cloning: and tokenizer_path: to the paths of the models you want to load. + +Then, use the --config option to point to your newly created config. + +```bash +# Use a different config +pocket-tts generate --config "C://pocket-tts/my_config.yaml" +``` + ## Output Format The generate command always outputs WAV files in the following format: diff --git a/docs/serve.md b/docs/serve.md index 72c999f..0f4a2ce 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -18,6 +18,7 @@ This starts a server on `http://localhost:8000` with the default voice model. - `--host HOST`: Host to bind to (default: "localhost") - `--port PORT`: Port to bind to (default: 8000) - `--reload`: Enable auto-reload for development +- `--config`: Path to a custom config .yaml ## Examples @@ -40,6 +41,17 @@ pocket-tts serve --voice "hf://kyutai/tts-voices/jessica-jian/casual.wav" # Use local voice file pocket-tts serve --voice "./my_voice.wav" ``` +### Custom Model Config +If you'd like to override the paths from which the models are loaded, you can provide a custom YAML configuration. + +Copy pocket_tts/config/b6369a24.yaml and change weights_path:, weights_path_without_voice_cloning: and tokenizer_path: to the paths of the models you want to load. + +Then, use the --config option to point to your newly created config. + +```bash +# Use a different config +pocket-tts serve --config "C://pocket-tts/my_config.yaml" +``` ## Web Interface diff --git a/pocket_tts/main.py b/pocket_tts/main.py index 12db058..356b7ad 100644 --- a/pocket_tts/main.py +++ b/pocket_tts/main.py @@ -181,11 +181,17 @@ def serve( host: Annotated[str, typer.Option(help="Host to bind to")] = "localhost", port: Annotated[int, typer.Option(help="Port to bind to")] = 8000, reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False, + config: Annotated[ + str, + typer.Option( + help="Model variant signature or path to locally-saved model config .yaml file" + ), + ] = DEFAULT_VARIANT, ): """Start the FastAPI server.""" global tts_model, global_model_state - tts_model = TTSModel.load_model(DEFAULT_VARIANT) + tts_model = TTSModel.load_model(config) # Pre-load the voice prompt global_model_state = tts_model.get_state_for_audio_prompt(voice) @@ -208,7 +214,9 @@ def generate( str, typer.Option(help="Path to audio conditioning file (voice to clone)") ] = DEFAULT_AUDIO_PROMPT, quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False, - variant: Annotated[str, typer.Option(help="Model signature")] = DEFAULT_VARIANT, + config: Annotated[ + str, typer.Option(help="Model signature or path to config .yaml file") + ] = DEFAULT_VARIANT, lsd_decode_steps: Annotated[ int, typer.Option(help="Number of generation steps") ] = DEFAULT_LSD_DECODE_STEPS, @@ -236,7 +244,7 @@ def generate( log_level = logging.ERROR if quiet else logging.INFO with enable_logging("pocket_tts", log_level): tts_model = TTSModel.load_model( - variant, temperature, lsd_decode_steps, noise_clamp, eos_threshold + config, temperature, lsd_decode_steps, noise_clamp, eos_threshold ) tts_model.to(device) diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 4a95431..9238fc8 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -189,8 +189,8 @@ class TTSModel(nn.Module): with the specified generation parameters and ready for inference. Args: - variant: Model variant identifier corresponding to a config file name - (e.g., '610b0b2c'). Must match a YAML file in the config directory. + variant: Model variant identifier (e.g., '610b0b2c'; must match a YAML file in the config directory) + or a path to a custom YAML config file saved locally (e.g., C://pocket_tts/pocket_tts_config.yaml). temp: Sampling temperature for generation. Higher values produce more diverse but potentially lower quality output. lsd_decode_steps: Number of steps for Lagrangian Self Distillation @@ -209,7 +209,12 @@ class TTSModel(nn.Module): are not found. ValueError: If the configuration is invalid or incompatible. """ - config = load_config(Path(__file__).parents[1] / f"config/{variant}.yaml") + if str(variant).endswith(".yaml"): + config_path = Path(variant) + config = load_config(config_path) + logger.info(f"Loading model from config at {config_path}...") + else: + config = load_config(Path(__file__).parents[1] / f"config/{variant}.yaml") tts_model = TTSModel._from_pydantic_config_with_weights( config, temp, lsd_decode_steps, noise_clamp, eos_threshold ) diff --git a/tests/test_cli_generate.py b/tests/test_cli_generate.py index d6d70b1..bc0bc68 100644 --- a/tests/test_cli_generate.py +++ b/tests/test_cli_generate.py @@ -76,7 +76,7 @@ def test_generate_with_custom_parameters(tmp_path): "generate", "--text", "Testing custom parameters.", - "--variant", + "--config", DEFAULT_VARIANT, "--temperature", "0.8",