Add: Support for speech synthesis in English and Chinese for ONNX inference in Non-JP-Extra models

bert/bert_models.json

@@ -11,8 +11,16 @@
         "repo_id": "hfl/chinese-roberta-wwm-ext-large",
         "files": ["pytorch_model.bin"]
     },
+    "chinese-roberta-wwm-ext-large-onnx": {
+        "repo_id": "tsukumijima/chinese-roberta-wwm-ext-large-onnx",
+        "files": ["model.onnx"]
+    },
     "deberta-v3-large": {
         "repo_id": "microsoft/deberta-v3-large",
         "files": ["spm.model", "pytorch_model.bin"]
+    },
+    "deberta-v3-large-onnx": {
+        "repo_id": "tsukumijima/deberta-v3-large-onnx",
+        "files": ["model.onnx"]
     }
 }

bert/chinese-roberta-wwm-ext-large-onnx/added_tokens.json

@@ -0,0 +1 @@
+{}

bert/chinese-roberta-wwm-ext-large-onnx/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "directionality": "bidi",
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "pooler_fc_size": 768,
+  "pooler_num_attention_heads": 12,
+  "pooler_num_fc_layers": 3,
+  "pooler_size_per_head": 128,
+  "pooler_type": "first_token_transform",
+  "type_vocab_size": 2,
+  "vocab_size": 21128
+}

bert/chinese-roberta-wwm-ext-large-onnx/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

bert/chinese-roberta-wwm-ext-large-onnx/tokenizer_config.json

@@ -0,0 +1 @@
+{"init_inputs": []}

bert/deberta-v3-large-onnx/config.json

@@ -0,0 +1,22 @@
+{
+		"model_type": "deberta-v2",
+		"attention_probs_dropout_prob": 0.1,
+		"hidden_act": "gelu",
+		"hidden_dropout_prob": 0.1,
+		"hidden_size": 1024,
+		"initializer_range": 0.02,
+		"intermediate_size": 4096,
+		"max_position_embeddings": 512,
+		"relative_attention": true,
+		"position_buckets": 256,
+		"norm_rel_ebd": "layer_norm",
+		"share_att_key": true,
+		"pos_att_type": "p2c|c2p",
+		"layer_norm_eps": 1e-7,
+		"max_relative_positions": -1,
+		"position_biased_input": false,
+		"num_attention_heads": 16,
+		"num_hidden_layers": 24,
+		"type_vocab_size": 0,
+		"vocab_size": 128100
+}

bert/deberta-v3-large-onnx/generator_config.json

@@ -0,0 +1,22 @@
+{
+	"model_type": "deberta-v2",
+	"attention_probs_dropout_prob": 0.1,
+	"hidden_act": "gelu",
+	"hidden_dropout_prob": 0.1,
+	"hidden_size": 1024,
+	"initializer_range": 0.02,
+	"intermediate_size": 4096,
+	"max_position_embeddings": 512,
+	"relative_attention": true,
+	"position_buckets": 256,
+	"norm_rel_ebd": "layer_norm",
+	"share_att_key": true,
+	"pos_att_type": "p2c|c2p",
+	"layer_norm_eps": 1e-7,
+	"max_relative_positions": -1,
+	"position_biased_input": false,
+	"num_attention_heads": 16,
+	"num_hidden_layers": 12,
+	"type_vocab_size": 0,
+	"vocab_size": 128100
+}

bert/deberta-v3-large-onnx/tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "do_lower_case": false,
+  "vocab_type": "spm"
+}

convert_bert_onnx.py

@@ -47,7 +47,6 @@ if __name__ == "__main__":
     print(f"[bold green]Tokenizer JSON saved to {tokenizer_json_path}[/bold green]")
     print(Rule(characters="=", style=Style(color="blue")))
 
-    # TODO: JP, ZH は変換できるが、EN は途中で強制終了されてしまい変換できない
     class ONNXBert(nn.Module):
         def __init__(self):
             super(ONNXBert, self).__init__()
@@ -88,8 +87,9 @@ if __name__ == "__main__":
         ],
         output_names=["output"],
         dynamic_axes={
-            "input_ids": {1: "batch_size"},
-            "attention_mask": {1: "batch_size"},
+            "input_ids": {0: "batch_size", 1: "sequence_length"},
+            "token_type_ids": {0: "batch_size", 1: "sequence_length"},
+            "attention_mask": {0: "batch_size", 1: "sequence_length"},
         },
     )
     print(

style_bert_vits2/constants.py

@@ -28,6 +28,8 @@ DEFAULT_BERT_MODEL_PATHS = {
 # 言語ごとのデフォルトの BERT モデル (ONNX 版) のパス
 DEFAULT_ONNX_BERT_MODEL_PATHS = {
     Languages.JP: BASE_DIR / "bert" / "deberta-v2-large-japanese-char-wwm-onnx",
+    Languages.EN: BASE_DIR / "bert" / "deberta-v3-large-onnx",
+    Languages.ZH: BASE_DIR / "bert" / "chinese-roberta-wwm-ext-large-onnx",
 }
 
 # デフォルトのユーザー辞書ディレクトリ

style_bert_vits2/nlp/__init__.py

@@ -81,6 +81,10 @@ def extract_bert_feature_onnx(
 
     if language == Languages.JP:
         from style_bert_vits2.nlp.japanese.bert_feature import extract_bert_feature_onnx
+    elif language == Languages.EN:
+        from style_bert_vits2.nlp.english.bert_feature import extract_bert_feature_onnx
+    elif language == Languages.ZH:
+        from style_bert_vits2.nlp.chinese.bert_feature import extract_bert_feature_onnx
     else:
         raise ValueError(f"Language {language} not supported")
 

style_bert_vits2/nlp/chinese/bert_feature.py

@@ -1,9 +1,16 @@
-from typing import Optional
+from __future__ import annotations
 
-import torch
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
+
+import numpy as np
+from numpy.typing import NDArray
 
 from style_bert_vits2.constants import Languages
-from style_bert_vits2.nlp import bert_models
+from style_bert_vits2.nlp import bert_models, onnx_bert_models
+
+
+if TYPE_CHECKING:
+    import torch
 
 
 def extract_bert_feature(
@@ -14,7 +21,7 @@ def extract_bert_feature(
     assist_text_weight: float = 0.7,
 ) -> torch.Tensor:
     """
-    中国語のテキストから BERT の特徴量を抽出する
+    中国語のテキストから BERT の特徴量を抽出する (PyTorch 推論)
 
     Args:
         text (str): 中国語のテキスト
@@ -27,6 +34,8 @@ def extract_bert_feature(
         torch.Tensor: BERT の特徴量
     """
 
+    import torch
+
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
     model = bert_models.load_model(Languages.ZH, device_map=device)
@@ -67,6 +76,76 @@ def extract_bert_feature(
     return phone_level_feature.T
 
 
+def extract_bert_feature_onnx(
+    text: str,
+    word2ph: list[int],
+    onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+) -> NDArray[Any]:
+    """
+    中国語のテキストから BERT の特徴量を抽出する (ONNX 推論)
+
+    Args:
+        text (str): 中国語のテキスト
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+        onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        NDArray[Any]: BERT の特徴量
+    """
+
+    tokenizer = onnx_bert_models.load_tokenizer(Languages.ZH)
+    inputs = tokenizer(text, return_tensors="pt")
+
+    session = onnx_bert_models.load_model(
+        language=Languages.ZH,
+        onnx_providers=onnx_providers,
+    )
+    output_name = session.get_outputs()[0].name
+    res = session.run(
+        [output_name],
+        {
+            "input_ids": inputs["input_ids"].detach().numpy(),  # type: ignore
+            "token_type_ids": inputs["token_type_ids"].detach().numpy(),  # type: ignore
+            "attention_mask": inputs["attention_mask"].detach().numpy(),  # type: ignore
+        },
+    )[0]
+
+    style_res_mean = None
+    if assist_text:
+        style_inputs = tokenizer(assist_text, return_tensors="pt")
+        style_res = session.run(
+            [output_name],
+            {
+                "input_ids": style_inputs["input_ids"].detach().numpy(),  # type: ignore
+                "token_type_ids": style_inputs["token_type_ids"].detach().numpy(),  # type: ignore
+                "attention_mask": style_inputs["attention_mask"].detach().numpy(),  # type: ignore
+            },
+        )[0]
+        style_res_mean = np.mean(style_res, axis=0)
+
+    assert len(word2ph) == len(text) + 2
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        if assist_text:
+            assert style_res_mean is not None
+            repeat_feature = (
+                np.tile(res[i], (word2phone[i], 1)) * (1 - assist_text_weight)
+                + np.tile(style_res_mean, (word2phone[i], 1)) * assist_text_weight
+            )
+        else:
+            repeat_feature = np.tile(res[i], (word2phone[i], 1))
+        phone_level_feature.append(repeat_feature)
+
+    phone_level_feature = np.concatenate(phone_level_feature, axis=0)
+
+    return phone_level_feature.T
+
+
 if __name__ == "__main__":
     word_level_feature = torch.rand(38, 1024)  # 12个词,每个词1024维特征
     word2phone = [

style_bert_vits2/nlp/english/bert_feature.py

@@ -1,9 +1,16 @@
-from typing import Optional
+from __future__ import annotations
 
-import torch
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
+
+import numpy as np
+from numpy.typing import NDArray
 
 from style_bert_vits2.constants import Languages
-from style_bert_vits2.nlp import bert_models
+from style_bert_vits2.nlp import bert_models, onnx_bert_models
+
+
+if TYPE_CHECKING:
+    import torch
 
 
 def extract_bert_feature(
@@ -14,7 +21,7 @@ def extract_bert_feature(
     assist_text_weight: float = 0.7,
 ) -> torch.Tensor:
     """
-    英語のテキストから BERT の特徴量を抽出する
+    英語のテキストから BERT の特徴量を抽出する (PyTorch 推論)
 
     Args:
         text (str): 英語のテキスト
@@ -27,6 +34,8 @@ def extract_bert_feature(
         torch.Tensor: BERT の特徴量
     """
 
+    import torch
+
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
     model = bert_models.load_model(Languages.EN, device_map=device)
@@ -65,3 +74,71 @@ def extract_bert_feature(
     phone_level_feature = torch.cat(phone_level_feature, dim=0)
 
     return phone_level_feature.T
+
+
+def extract_bert_feature_onnx(
+    text: str,
+    word2ph: list[int],
+    onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+) -> NDArray[Any]:
+    """
+    英語のテキストから BERT の特徴量を抽出する (ONNX 推論)
+
+    Args:
+        text (str): 英語のテキスト
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+        onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        NDArray[Any]: BERT の特徴量
+    """
+
+    tokenizer = onnx_bert_models.load_tokenizer(Languages.EN)
+    inputs = tokenizer(text, return_tensors="pt")
+
+    session = onnx_bert_models.load_model(
+        language=Languages.EN,
+        onnx_providers=onnx_providers,
+    )
+    output_name = session.get_outputs()[0].name
+    res = session.run(
+        [output_name],
+        {
+            "input_ids": inputs["input_ids"].detach().numpy(),  # type: ignore
+            "attention_mask": inputs["attention_mask"].detach().numpy(),  # type: ignore
+        },
+    )[0]
+
+    style_res_mean = None
+    if assist_text:
+        style_inputs = tokenizer(assist_text, return_tensors="pt")
+        style_res = session.run(
+            [output_name],
+            {
+                "input_ids": style_inputs["input_ids"].detach().numpy(),  # type: ignore
+                "attention_mask": style_inputs["attention_mask"].detach().numpy(),  # type: ignore
+            },
+        )[0]
+        style_res_mean = np.mean(style_res, axis=0)
+
+    assert len(word2ph) == res.shape[0], (text, res.shape[0], len(word2ph))
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        if assist_text:
+            assert style_res_mean is not None
+            repeat_feature = (
+                np.tile(res[i], (word2phone[i], 1)) * (1 - assist_text_weight)
+                + np.tile(style_res_mean, (word2phone[i], 1)) * assist_text_weight
+            )
+        else:
+            repeat_feature = np.tile(res[i], (word2phone[i], 1))
+        phone_level_feature.append(repeat_feature)
+
+    phone_level_feature = np.concatenate(phone_level_feature, axis=0)
+
+    return phone_level_feature.T

style_bert_vits2/nlp/japanese/bert_feature.py

@@ -121,8 +121,8 @@ def extract_bert_feature_onnx(
     res = session.run(
         [output_name],
         {
-            "input_ids": inputs["input_ids"].detach().numpy(),
-            "attention_mask": inputs["attention_mask"].detach().numpy(),
+            "input_ids": inputs["input_ids"].detach().numpy(),  # type: ignore
+            "attention_mask": inputs["attention_mask"].detach().numpy(),  # type: ignore
         },
     )[0]
 
@@ -132,8 +132,8 @@ def extract_bert_feature_onnx(
         style_res = session.run(
             [output_name],
             {
-                "input_ids": style_inputs["input_ids"].detach().numpy(),
-                "attention_mask": style_inputs["attention_mask"].detach().numpy(),
+                "input_ids": style_inputs["input_ids"].detach().numpy(),  # type: ignore
+                "attention_mask": style_inputs["attention_mask"].detach().numpy(),  # type: ignore
             },
         )[0]
         style_res_mean = np.mean(style_res, axis=0)