Update infer_onnx.py

infer_onnx.py

@@ -6,6 +6,7 @@ from huggingface_hub import snapshot_download
 from num2words import num2words
 import re
 from transliterate import translit
+import json
 
 class TTS:
     def __init__(self, model_name: str, save_path: str = "./model", add_time_to_end: float = 0.8) -> None:
@@ -22,7 +23,9 @@ class TTS:
                             )
         
         self.model = onnxruntime.InferenceSession(os.path.join(model_dir, "exported/model.onnx"), providers=['CPUExecutionProvider'])
-        
+        with open(os.path.join(model_dir, "exported/config.json")) as config_file:
+            self.config = json.load(config_file)["model_config"]
+            
         if os.path.exists(os.path.join(model_dir, "exported/dictionary.txt")):
             from tokenizer import TokenizerG2P
             print("Use g2p")
@@ -43,9 +46,9 @@ class TTS:
         return audio_with_silence
 
 
-    def save_wav(self, audio, path:str):
+    def save_wav(self, audio, path:str, sample_rate: int = 22050):
         '''save audio to wav'''
-        scipy.io.wavfile.write(path, 22050, audio)
+        scipy.io.wavfile.write(path, sample_rate, audio)
     
     
     def _intersperse(self, lst, item):
@@ -83,5 +86,6 @@ class TTS:
                 "sid": None,
             },
         )[0][0,0][0]
-        audio = self._add_silent(audio, silence_duration = self.add_time_to_end)
+        
+        audio = self._add_silent(audio, silence_duration = self.add_time_to_end, sample_rate=self.config["samplerate"])
         return audio