Basic changes to allow loading of local model files (#59)

-modify main,py to not hardcode using DEFAULT_VARIANT

-allow submission of path string to path of variant.yaml file saved in
custom directory in lieu of model variant identifier in CLI to use when
loading model

-addresses #57

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:

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
 

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)
 

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
+            variant: Model variant identifier (e.g., '610b0b2c'; must match a YAML file in the config directory)
-                (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,6 +209,11 @@ class TTSModel(nn.Module):
                 are not found.
             ValueError: If the configuration is invalid or incompatible.
         """
+        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

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",