Update inference

app.py

@@ -64,7 +64,9 @@ def tts_fn(
             wrong_tone_message = "アクセント指定は現在日本語のみ対応しています。"
         if line_split:
             logger.warning("Tone generation is not supported for line split.")
-            wrong_tone_message = "アクセント指定は改行で分けて生成を使わない場合のみ対応しています。"
+            wrong_tone_message = (
+                "アクセント指定は改行で分けて生成を使わない場合のみ対応しています。"
+            )
         try:
             kata_tone = []
             json_data = json.loads(kata_tone_json_str)
@@ -109,6 +111,9 @@ def tts_fn(
     except InvalidToneError as e:
         logger.error(f"Tone error: {e}")
         return f"Error: アクセント指定が不正です:\n{e}", None, kata_tone_json_str
+    except ValueError as e:
+        logger.error(f"Value error: {e}")
+        return f"Error: {e}", None, kata_tone_json_str
 
     end_time = datetime.datetime.now()
     duration = (end_time - start_time).total_seconds()
@@ -179,7 +184,10 @@ examples = [
         "Speech synthesis is the artificial production of human speech. A computer system used for this purpose is called a speech synthesizer, and can be implemented in software or hardware products.",
         "EN",
     ],
-    ["语音合成是人工制造人类语音。用于此目的的计算机系统称为语音合成器，可以通过软件或硬件产品实现。", "ZH"],
+    [
+        "语音合成是人工制造人类语音。用于此目的的计算机系统称为语音合成器，可以通过软件或硬件产品实现。",
+        "ZH",
+    ],
 ]
 
 initial_md = """
@@ -268,7 +276,9 @@ if __name__ == "__main__":
 
     model_names = model_holder.model_names
     if len(model_names) == 0:
-        logger.error(f"モデルが見つかりませんでした。{model_dir}にモデルを置いてください。")
+        logger.error(
+            f"モデルが見つかりませんでした。{model_dir}にモデルを置いてください。"
+        )
         sys.exit(1)
     initial_id = 0
     initial_pth_files = model_holder.model_files_dict[model_names[initial_id]]
@@ -295,7 +305,9 @@ if __name__ == "__main__":
                     load_button = gr.Button("ロード", scale=1, variant="primary")
                 text_input = gr.TextArea(label="テキスト", value=initial_text)
 
-                line_split = gr.Checkbox(label="改行で分けて生成", value=DEFAULT_LINE_SPLIT)
+                line_split = gr.Checkbox(
+                    label="改行で分けて生成", value=DEFAULT_LINE_SPLIT
+                )
                 split_interval = gr.Slider(
                     minimum=0.0,
                     maximum=2,
@@ -349,7 +361,9 @@ if __name__ == "__main__":
                         step=0.1,
                         label="Length",
                     )
-                    use_assist_text = gr.Checkbox(label="Assist textを使う", value=False)
+                    use_assist_text = gr.Checkbox(
+                        label="Assist textを使う", value=False
+                    )
                     assist_text = gr.Textbox(
                         label="Assist text",
                         placeholder="どうして私の意見を無視するの？許せない、ムカつく！死ねばいいのに。",
@@ -389,9 +403,13 @@ if __name__ == "__main__":
                     step=0.1,
                     label="スタイルの強さ",
                 )
-                ref_audio_path = gr.Audio(label="参照音声", type="filepath", visible=False)
+                ref_audio_path = gr.Audio(
+                    label="参照音声", type="filepath", visible=False
+                )
                 tts_button = gr.Button(
-                    "音声合成（モデルをロードしてください）", variant="primary", interactive=False
+                    "音声合成（モデルをロードしてください）",
+                    variant="primary",
+                    interactive=False,
                 )
                 text_output = gr.Textbox(label="情報")
                 audio_output = gr.Audio(label="結果")

common/tts_model.py

@@ -4,11 +4,12 @@ import torch
 import os
 import warnings
 from gradio.processing_utils import convert_to_16_bit_wav
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import utils
 from infer import get_net_g, infer
 from models import SynthesizerTrn
+from models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
 
 from .log import logger
 from .constants import (
@@ -52,7 +53,7 @@ class Model:
                 f"The number of styles ({self.num_styles}) does not match the number of style vectors ({self.style_vectors.shape[0]})"
             )
 
-        self.net_g: Optional[SynthesizerTrn] = None
+        self.net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None
 
     def load_net_g(self):
         self.net_g = get_net_g(
@@ -98,6 +99,10 @@ class Model:
         given_tone: Optional[list[int]] = None,
     ) -> tuple[int, np.ndarray]:
         logger.info(f"Start generating audio data from text:\n{text}")
+        if language != "JP" and self.hps.version.endswith("JP-Extra"):
+            raise ValueError(
+                "The model is trained with JP-Extra, but the language is not JP"
+            )
         if reference_audio_path == "":
             reference_audio_path = None
         if assist_text == "" or not use_assist_text:

infer.py

@@ -3,6 +3,7 @@ import torch
 import commons
 import utils
 from models import SynthesizerTrn
+from models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
 from text import cleaned_text_to_sequence, get_bert
 from text.cleaner import clean_text
 from text.symbols import symbols
@@ -16,13 +17,22 @@ class InvalidToneError(ValueError):
 
 
 def get_net_g(model_path: str, version: str, device: str, hps):
-    net_g = SynthesizerTrn(
-        len(symbols),
-        hps.data.filter_length // 2 + 1,
-        hps.train.segment_size // hps.data.hop_length,
-        n_speakers=hps.data.n_speakers,
-        **hps.model,
-    ).to(device)
+    if version.endswith("JP-Extra"):
+        net_g = SynthesizerTrnJPExtra(
+            len(symbols),
+            hps.data.filter_length // 2 + 1,
+            hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            **hps.model,
+        ).to(device)
+    else:
+        net_g = SynthesizerTrn(
+            len(symbols),
+            hps.data.filter_length // 2 + 1,
+            hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            **hps.model,
+        ).to(device)
     net_g.state_dict()
     _ = net_g.eval()
     if model_path.endswith(".pth") or model_path.endswith(".pt"):
@@ -109,6 +119,7 @@ def infer(
     assist_text_weight=0.7,
     given_tone=None,
 ):
+    is_jp_extra = hps.version.endswith("JP-Extra")
     bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
         text,
         language,
@@ -143,8 +154,22 @@ def infer(
         style_vec = torch.from_numpy(style_vec).to(device).unsqueeze(0)
         del phones
         sid_tensor = torch.LongTensor([sid]).to(device)
-        audio = (
-            net_g.infer(
+        if is_jp_extra:
+            output = net_g.infer(
+                x_tst,
+                x_tst_lengths,
+                sid_tensor,
+                tones,
+                lang_ids,
+                ja_bert,
+                style_vec=style_vec,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )
+        else:
+            output = net_g.infer(
                 x_tst,
                 x_tst_lengths,
                 sid_tensor,
@@ -158,11 +183,8 @@ def infer(
                 noise_scale=noise_scale,
                 noise_scale_w=noise_scale_w,
                 length_scale=length_scale,
-            )[0][0, 0]
-            .data.cpu()
-            .float()
-            .numpy()
-        )
+            )
+        audio = output[0][0, 0].data.cpu().float().numpy()
         del (
             x_tst,
             tones,