Custom checkpoint possibility

This commit is contained in:
Shivam Mehta
2023-09-17 07:50:04 +00:00
parent a9251ed984
commit bd55eb76f2
2 changed files with 19 additions and 0 deletions

View File

@@ -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)