Improve consistency in variant -> config rename (#87)

Small follow-up on #59

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)

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)

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)
 

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)
+            config: a path to a custom YAML config file saved locally (e.g., C://pocket_tts/pocket_tts_config.yaml)
-                or 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"):
+        if str(config).endswith(".yaml"):
-            config_path = Path(variant)
+            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
         )

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(
+    model = TTSModel.load_model(config="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0)
-        variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0
-    )
 
 
 def test_device():