From e401a724ac6d8f8b044b89036191025ef202aa7c Mon Sep 17 00:00:00 2001
From: jinzr <zengrui.jin0@gmail.com>
Date: Tue, 5 Dec 2023 08:53:20 +0800
Subject: [PATCH] Update export-onnx.py

---
 egs/vctk/TTS/vits/export-onnx.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/egs/vctk/TTS/vits/export-onnx.py b/egs/vctk/TTS/vits/export-onnx.py
index 2068adeea..94ec986f4 100755
--- a/egs/vctk/TTS/vits/export-onnx.py
+++ b/egs/vctk/TTS/vits/export-onnx.py
@@ -68,6 +68,12 @@ def get_parser():
         help="The experiment dir",
     )
 
+    parser.add_argument(
+        "--speakers",
+        type=Path,
+        default=Path("data/speakers.txt"),
+        help="Path to speakers.txt file.",
+    )
     parser.add_argument(
         "--tokens",
         type=str,
@@ -219,6 +225,10 @@ def main():
     params.oov_id = tokenizer.oov_id
     params.vocab_size = tokenizer.vocab_size
 
+    with open(args.speakers) as f:
+        speaker_map = {line.strip(): i for i, line in enumerate(f)}
+    params.num_spks = len(speaker_map)
+
     logging.info(params)
 
     logging.info("About to create model")