From bd55eb76f2d6225e1d5002e7b6677a50f39f3d6e Mon Sep 17 00:00:00 2001
From: Shivam Mehta <shivam.mehta25@gmail.com>
Date: Sun, 17 Sep 2023 07:50:04 +0000
Subject: [PATCH] Custom checkpoint possibility

---
 README.md     |  6 ++++++
 matcha/cli.py | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/README.md b/README.md
index fafdc89..ae5463d 100644
--- a/README.md
+++ b/README.md
@@ -182,6 +182,12 @@ python matcha/train.py experiment=ljspeech_min_memory
 python matcha/train.py experiment=ljspeech trainer.devices=[0,1]
 ```
 
+6. Synthesise from the custom trained model
+
+```bash
+matcha_tts --text "<INPUT TEXT>" --checkpoint_path <PATH TO CHECKPOINT>
+```
+
 ## Acknowledgements
 
 Since this code uses: [Lightning-Hydra-Template](https://github.com/ashleve/lightning-hydra-template), you have all the powers that comes with it.
diff --git a/matcha/cli.py b/matcha/cli.py
index 49004fb..1e7f8a9 100644
--- a/matcha/cli.py
+++ b/matcha/cli.py
@@ -155,6 +155,14 @@ def cli():
         help="Model to use",
         choices=MATCHA_URLS.keys(),
     )
+
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        default=None,
+        help="Path to the custom model checkpoint",
+    )
+
     parser.add_argument(
         "--vocoder",
         type=str,
@@ -201,6 +209,11 @@ def cli():
     print_config(args)
     paths = assert_required_models_available(args)
 
+    if args.checkpoint_path is not None:
+        print(f"[🍵] Loading custom model from {args.checkpoint_path}")
+        paths["matcha"] = args.checkpoint_path
+        args.model = "custom_model"
+
     model = load_matcha(args.model, paths["matcha"], device)
     vocoder, denoiser = load_vocoder(args.vocoder, paths["vocoder"], device)