From 28b8244428cad108e0fc5634b81207d09168412d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=A1clav=20Volhejn?= <8401624+vvolhejn@users.noreply.github.com> Date: Tue, 27 Jan 2026 09:55:34 +0100 Subject: [PATCH] Improve consistency in variant -> config rename (#87) Small follow-up on #59 --- docs/generate.md | 3 +-- docs/python-api.md | 4 ++-- pocket_tts/main.py | 6 +++--- pocket_tts/models/tts_model.py | 13 +++++++------ tests/test_documentation_examples.py | 4 +--- 5 files changed, 14 insertions(+), 16 deletions(-) diff --git a/docs/generate.md b/docs/generate.md index a8504db..22af053 100644 --- a/docs/generate.md +++ b/docs/generate.md @@ -19,11 +19,10 @@ 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 -- `--variant VARIANT`: Model signature (default: "b6369a24") +- `--config CONFIG_PATH`: Path to custom config.yaml (for loading local model files) or model signature (default: "b6369a24") - `--lsd-decode-steps LSD_DECODE_STEPS`: Number of generation steps (default: 1) - `--temperature TEMPERATURE`: Temperature for generation (default: 0.7) - `--noise-clamp NOISE_CLAMP`: Noise clamp value (default: None) diff --git a/docs/python-api.md b/docs/python-api.md index 5fcf2dd..a473023 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -37,12 +37,12 @@ The main class for text-to-speech generation. #### Class Methods -##### `load_model(variant="b6369a24", temp=0.7, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0)` +##### `load_model(config="b6369a24", temp=0.7, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0)` Load and return a TTSModel instance with pre-trained weights. **Parameters:** -- `variant` (str): Model variant identifier (default: "b6369a24") +- `config` (str): Path to model config YAML file or a variant identifier (default: "b6369a24") - `temp` (float): Sampling temperature for generation (default: 0.7) - `lsd_decode_steps` (int): Number of generation steps (default: 1) - `noise_clamp` (float | None): Maximum value for noise sampling (default: None) diff --git a/pocket_tts/main.py b/pocket_tts/main.py index 356b7ad..f749b4b 100644 --- a/pocket_tts/main.py +++ b/pocket_tts/main.py @@ -184,7 +184,7 @@ def serve( config: Annotated[ str, typer.Option( - help="Model variant signature or path to locally-saved model config .yaml file" + help="Path to locally-saved model config .yaml file or model variant signature" ), ] = DEFAULT_VARIANT, ): @@ -286,7 +286,7 @@ def export_voice( 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, + 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, @@ -343,7 +343,7 @@ def export_voice( 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 9238fc8..f66f956 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -176,7 +176,7 @@ class TTSModel(nn.Module): @classmethod def load_model( cls, - variant: str = DEFAULT_VARIANT, + config: str | Path = DEFAULT_VARIANT, temp: float | int = DEFAULT_TEMPERATURE, lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS, noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP, @@ -189,8 +189,8 @@ class TTSModel(nn.Module): with the specified generation parameters and ready for inference. Args: - 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). + config: a path to a custom YAML config file saved locally (e.g., C://pocket_tts/pocket_tts_config.yaml) + or a model variant identifier (e.g., '610b0b2c'; must match a YAML file in the config directory). 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,12 +209,13 @@ class TTSModel(nn.Module): are not found. ValueError: If the configuration is invalid or incompatible. """ - if str(variant).endswith(".yaml"): - config_path = Path(variant) + if str(config).endswith(".yaml"): + config_path = Path(config) 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") + config = load_config(Path(__file__).parents[1] / f"config/{config}.yaml") + tts_model = TTSModel._from_pydantic_config_with_weights( config, temp, lsd_decode_steps, noise_clamp, eos_threshold ) diff --git a/tests/test_documentation_examples.py b/tests/test_documentation_examples.py index 06b918d..723d9ae 100644 --- a/tests/test_documentation_examples.py +++ b/tests/test_documentation_examples.py @@ -41,9 +41,7 @@ def test_load_model(): model = TTSModel.load_model() # Load with custom parameters - model = TTSModel.load_model( - variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0 - ) + model = TTSModel.load_model(config="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0) def test_device():