Update infer_onnx.py

infer_onnx.py

@@ -37,29 +37,36 @@ def process_text(i: int, text: str, device: torch.device, cleaner:str):
     return x.numpy(), x_lengths.numpy()
 
 # paths
-MODEL_PATH_MATCHA_MEL_BAL="matcha_multispeaker_cat_bal_opset_15_10_steps.onnx"
-MODEL_PATH_MATCHA_MEL_CAT="matcha_multispeaker_cat_cen_opset_15_10_steps.onnx"
-MODEL_PATH_MATCHA_MEL_OCC="matcha_multispeaker_cat_occ_opset_15_10_steps.onnx"
-MODEL_PATH_MATCHA_MEL_VAL="matcha_multispeaker_cat_val_opset_15_10_steps.onnx"
+# MODEL_PATH_MATCHA_MEL_BAL="matcha_multispeaker_cat_bal_opset_15_10_steps.onnx"
+# MODEL_PATH_MATCHA_MEL_CAT="matcha_multispeaker_cat_cen_opset_15_10_steps.onnx"
+# MODEL_PATH_MATCHA_MEL_OCC="matcha_multispeaker_cat_occ_opset_15_10_steps.onnx"
+# MODEL_PATH_MATCHA_MEL_VAL="matcha_multispeaker_cat_val_opset_15_10_steps.onnx"
+
+MODEL_PATH_MATCHA_MEL_ALL="matcha_multispeaker_cat_all_opset_15_10_steps.onnx"
+
 MODEL_PATH_VOCOS="mel_spec_22khz_cat.onnx"
 CONFIG_PATH="config.yaml"
-SPEAKER_ID_DICT="spk_to_id_2.json"
+SPEAKER_ID_DICT="spk_to_id_3.json"
 
 # Load models
 sess_options = onnxruntime.SessionOptions()
-model_matcha_mel_bal = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_BAL), sess_options=sess_options, providers=["CPUExecutionProvider"])
-model_matcha_mel_cat = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_CAT), sess_options=sess_options, providers=["CPUExecutionProvider"])
-model_matcha_mel_occ = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_OCC), sess_options=sess_options, providers=["CPUExecutionProvider"])
-model_matcha_mel_val = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_VAL), sess_options=sess_options, providers=["CPUExecutionProvider"])
+
+# model_matcha_mel_bal = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_BAL), sess_options=sess_options, providers=["CPUExecutionProvider"])
+# model_matcha_mel_cat = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_CAT), sess_options=sess_options, providers=["CPUExecutionProvider"])
+# model_matcha_mel_occ = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_OCC), sess_options=sess_options, providers=["CPUExecutionProvider"])
+# model_matcha_mel_val = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_VAL), sess_options=sess_options, providers=["CPUExecutionProvider"])
+
+model_matcha_mel_all = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA_MEL_ALL), sess_options=sess_options, providers=["CPUExecutionProvider"])
+
 model_vocos = onnxruntime.InferenceSession(str(MODEL_PATH_VOCOS), sess_options=sess_options, providers=["CPUExecutionProvider"])
 
 speaker_id_dict = json.load(open(SPEAKER_ID_DICT))
 accents = [e for e in speaker_id_dict.keys()]
 
-models={"balear":model_matcha_mel_bal,
-        "nord-occidental": model_matcha_mel_occ,
-        "valencia": model_matcha_mel_val,
-        "central": model_matcha_mel_cat}
+# models={"balear":model_matcha_mel_bal,
+#         "nord-occidental": model_matcha_mel_occ,
+#         "valencia": model_matcha_mel_val,
+#         "central": model_matcha_mel_cat}
 
 cleaners={"balear": "catalan_balear_cleaners",
         "nord-occidental": "catalan_occidental_cleaners",
@@ -161,7 +168,8 @@ def tts(text:str, accent:str, spk_name:str, temperature:float, length_scale:floa
         spk_id = speaker_id_dict[accent][spk_name]
         sid = np.array([int(spk_id)]) if spk_id is not None else None
         text_matcha , text_lengths = process_text(0,text,"cpu",cleaner=cleaners[accent])
-        model_matcha_mel = models[accent]
+        # model_matcha_mel = models[accent]
+        model_matcha_mel = model_matcha_mel_all
 
         # MATCHA VOCOS
         inputs = {