Add: Preparation for ONNX inference support ②

style_bert_vits2/nlp/japanese/bert_feature.py

@@ -1,12 +1,19 @@
-from typing import Optional
+from __future__ import annotations
 
-import torch
+from typing import TYPE_CHECKING, Any, Optional, Sequence
+
+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.japanese.g2p import text_to_sep_kata
 
 
+if TYPE_CHECKING:
+    import torch
+
+
 def extract_bert_feature(
     text: str,
     word2ph: list[int],
@@ -15,7 +22,7 @@ def extract_bert_feature(
     assist_text_weight: float = 0.7,
 ) -> torch.Tensor:
     """
-    日本語のテキストから BERT の特徴量を抽出する
+    日本語のテキストから BERT の特徴量を抽出する (PyTorch 推論)
 
     Args:
         text (str): 日本語のテキスト
@@ -28,6 +35,8 @@ def extract_bert_feature(
         torch.Tensor: BERT の特徴量
     """
 
+    import torch
+
     # 各単語が何文字かを作る `word2ph` を使う必要があるので、読めない文字は必ず無視する
     # でないと `word2ph` の結果とテキストの文字数結果が整合性が取れない
     text = "".join(text_to_sep_kata(text, raise_yomi_error=False)[0])
@@ -72,3 +81,64 @@ 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: list[str],
+    onnx_provider_options: Optional[Sequence[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 など)
+        onnx_provider_options (Optional[dict[str, Any]]): ONNX 推論で利用する ExecutionProvider のオプション
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        NDArray[Any]: BERT の特徴量
+    """
+
+    # 各単語が何文字かを作る `word2ph` を使う必要があるので、読めない文字は必ず無視する
+    # でないと `word2ph` の結果とテキストの文字数結果が整合性が取れない
+    text = "".join(text_to_sep_kata(text, raise_yomi_error=False)[0])
+    if assist_text:
+        assist_text = "".join(text_to_sep_kata(assist_text, raise_yomi_error=False)[0])
+
+    tokenizer = Tokenizer.from_file("tokenizer.json")
+    token_ids = [1]
+    attention_mask = [1]
+    for word in text:
+        encoded = tokenizer.encode(word)
+        token_ids.extend(encoded.ids[1:-1])
+        attention_mask.extend(encoded.attention_mask[1:-1])
+
+    token_ids.append(2)
+    attention_mask.append(1)
+
+    bert_output_name = bert_session.get_outputs()[0].name
+    res = bert_session.run(
+        [bert_output_name],
+        {
+            "input_ids": np.array(token_ids).reshape(1, -1),
+            "attention_mask": np.array(attention_mask).reshape(1, -1),
+        },
+    )[0]
+
+    assert len(word2ph) == len(text) + 2, text
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        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