tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Improve: Support ONNX inference, add ONNX conversion script

Browse Source
This commit is contained in:
tsukumi
2024-09-17 16:55:03 +09:00
parent f42b9f0f89
commit 5e2c83c6c7
21 changed files with 44483 additions and 376 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
style_bert_vits2/models/infer_onnx.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Optional, Sequence
from typing import Any, Optional, Sequence, Union
import numpy as np
import onnxruntime
@@ -35,8 +35,7 @@ def get_text_onnx(
text: str,
language_str: Languages,
hps: HyperParameters,
onnx_providers: list[str],
onnx_provider_options: Optional[Sequence[dict[str, Any]]],
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
assist_text: Optional[str] = None,
assist_text_weight: float = 0.7,
given_phone: Optional[list[str]] = None,
@@ -68,7 +67,6 @@ def get_text_onnx(
word2ph,
language_str,
onnx_providers,
onnx_provider_options,
assist_text,
assist_text_weight,
)
@@ -111,8 +109,7 @@ def infer_onnx(
language: Languages,
hps: HyperParameters,
onnx_session: onnxruntime.InferenceSession,
onnx_providers: list[str],
onnx_provider_options: Optional[Sequence[dict[str, Any]]],
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
skip_start: bool = False,
skip_end: bool = False,
assist_text: Optional[str] = None,
@@ -126,7 +123,6 @@ def infer_onnx(
language,
hps,
onnx_providers=onnx_providers,
onnx_provider_options=onnx_provider_options,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
given_phone=given_phone,
@@ -171,8 +167,8 @@ def infer_onnx(
input_names[4]: lang_ids,
input_names[5]: ja_bert,
input_names[6]: style_vec_tensor,
input_names[7]: length_scale,
input_names[8]: sdp_ratio,
input_names[7]: np.array([length_scale], dtype=np.float32),
input_names[8]: np.array([sdp_ratio], dtype=np.float32),
},
)
else:

7
style_bert_vits2/nlp/__init__.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Optional, Sequence
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
from numpy.typing import NDArray
@@ -60,8 +60,7 @@ def extract_bert_feature_onnx(
text: str,
word2ph: list[int],
language: Languages,
onnx_providers: list[str],
onnx_provider_options: Optional[Sequence[dict[str, Any]]],
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
assist_text: Optional[str] = None,
assist_text_weight: float = 0.7,
) -> NDArray[Any]:
@@ -73,7 +72,6 @@ def extract_bert_feature_onnx(
word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
language (Languages): テキストの言語
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)
@@ -90,7 +88,6 @@ def extract_bert_feature_onnx(
text,
word2ph,
onnx_providers,
onnx_provider_options,
assist_text,
assist_text_weight,
)

7
style_bert_vits2/nlp/japanese/bert_feature.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Optional, Sequence
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
import numpy as np
from numpy.typing import NDArray
@@ -86,8 +86,7 @@ def extract_bert_feature(
def extract_bert_feature_onnx(
text: str,
word2ph: list[int],
onnx_providers: list[str],
onnx_provider_options: Optional[Sequence[dict[str, Any]]],
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
assist_text: Optional[str] = None,
assist_text_weight: float = 0.7,
) -> NDArray[Any]:
@@ -98,7 +97,6 @@ def extract_bert_feature_onnx(
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)
@@ -118,7 +116,6 @@ def extract_bert_feature_onnx(
session = onnx_bert_models.load_model(
language=Languages.JP,
onnx_providers=onnx_providers,
onnx_provider_options=onnx_provider_options,
)
output_name = session.get_outputs()[0].name
res = session.run(

7
style_bert_vits2/nlp/onnx_bert_models.py
View File

@@ -39,11 +39,10 @@ __loaded_tokenizers: dict[
def load_model(
language: Languages,
pretrained_model_name_or_path: Optional[str] = None,
onnx_providers: list[str] = ["CPUExecutionProvider"],
onnx_provider_options: Optional[Sequence[dict[str, Any]]] = None,
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = ["CPUExecutionProvider"],
cache_dir: Optional[str] = None,
revision: str = "main",
) -> onnxruntime.InferenceSession:
) -> onnxruntime.InferenceSession: # fmt: skip
"""
指定された言語の ONNX 版 BERT モデルをロードし、ロード済みの ONNX 版 BERT モデルを返す。
一度ロードされていれば、ロード済みの ONNX 版 BERT モデルを即座に返す。
@@ -59,7 +58,6 @@ def load_model(
language (Languages): ロードする学習済みモデルの対象言語
pretrained_model_name_or_path (Optional[str]): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
onnx_provider_options (Optional[dict[str, Any]]): ONNX 推論で利用する ExecutionProvider のオプション
cache_dir (Optional[str]): モデルのキャッシュディレクトリ。指定しない場合はデフォルトのキャッシュディレクトリが利用される (デフォルト: None)
revision (str): モデルの Hugging Face 上の Git リビジョン。指定しない場合は最新の main ブランチの内容が利用される (デフォルト: None)
@@ -98,7 +96,6 @@ def load_model(
__loaded_models[language] = onnxruntime.InferenceSession(
model_path,
providers=onnx_providers,
provider_options=onnx_provider_options,
)
logger.info(
f"Loaded the {language} ONNX BERT model from {pretrained_model_name_or_path}"

82
style_bert_vits2/tts_model.py
View File

@@ -44,9 +44,8 @@ class TTSModel:
config_path: Union[Path, HyperParameters],
style_vec_path: Union[Path, NDArray[Any]],
device: str = "cpu",
onnx_providers: list[str] = ["CPUExecutionProvider"],
onnx_provider_options: Optional[Sequence[dict[str, Any]]] = None,
) -> None:
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = ["CPUExecutionProvider"],
) -> None: # fmt: skip
"""
Style-Bert-VITS2 の音声合成モデルを初期化する。
この時点ではモデルはロードされていない (明示的にロードしたい場合は model.load() を呼び出す)。
@@ -57,13 +56,11 @@ class TTSModel:
style_vec_path (Union[Path, NDArray[Any]]): スタイルベクトル (style_vectors.npy) のパス (直接 NDArray を指定することも可能)
device (str): PyTorch 推論での音声合成時に利用するデバイス (cpu, cuda, mps など)
onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
onnx_provider_options (Optional[dict[str, Any]]): ONNX 推論で利用する ExecutionProvider のオプション
"""
self.model_path: Path = model_path
self.device: str = device
self.onnx_providers: list[str] = onnx_providers
self.onnx_provider_options: Optional[Sequence[dict[str, Any]]] = onnx_provider_options # fmt: skip
self.onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = onnx_providers # fmt: skip
# ONNX 形式のモデルかどうか
if self.model_path.suffix == ".onnx":
@@ -85,11 +82,11 @@ class TTSModel:
# スタイルベクトルの NDArray が直接指定された
if isinstance(style_vec_path, np.ndarray):
self.style_vec_path: Path = Path("") # 互換性のため空の Path を設定
self.__style_vectors: NDArray[Any] = style_vec_path
self.style_vectors: NDArray[Any] = style_vec_path
# スタイルベクトルのパスが指定された
else:
self.style_vec_path: Path = style_vec_path
self.__style_vectors: NDArray[Any] = np.load(self.style_vec_path)
self.style_vectors: NDArray[Any] = np.load(self.style_vec_path)
self.spk2id: dict[str, int] = self.hyper_parameters.data.spk2id
self.id2spk: dict[int, str] = {v: k for k, v in self.spk2id.items()}
@@ -104,17 +101,17 @@ class TTSModel:
f"Number of styles ({num_styles}) does not match the number of style2id ({len(self.style2id)})"
)
if self.__style_vectors.shape[0] != num_styles:
if self.style_vectors.shape[0] != num_styles:
raise ValueError(
f"The number of styles ({num_styles}) does not match the number of style vectors ({self.__style_vectors.shape[0]})"
f"The number of styles ({num_styles}) does not match the number of style vectors ({self.style_vectors.shape[0]})"
)
self.__style_vector_inference: Optional[Any] = None
self.style_vector_inference: Optional[Any] = None
# __net_g は PyTorch 推論時のみ遅延初期化される
self.__net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None
# net_g は PyTorch 推論時のみ遅延初期化される
self.net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None
# __onnx_session は ONNX 推論時のみ遅延初期化される
self.__onnx_session: Optional[onnxruntime.InferenceSession] = None
# onnx_session は ONNX 推論時のみ遅延初期化される
self.onnx_session: Optional[onnxruntime.InferenceSession] = None
def load(self) -> None:
"""
@@ -125,7 +122,7 @@ class TTSModel:
if not self.is_onnx_model:
from style_bert_vits2.models.infer import get_net_g
self.__net_g = get_net_g(
self.net_g = get_net_g(
model_path=str(self.model_path),
version=self.hyper_parameters.version,
device=self.device,
@@ -134,12 +131,13 @@ class TTSModel:
# ONNX 推論時
else:
self.__onnx_session = onnxruntime.InferenceSession(
self.onnx_session = onnxruntime.InferenceSession(
path_or_bytes=str(self.model_path),
providers=self.onnx_providers,
provider_options=self.onnx_provider_options,
)
logger.info(f"Model loaded successfully from {self.model_path}")
def get_style_vector(self, style_id: int, weight: float = 1.0) -> NDArray[Any]:
"""
スタイルベクトルを取得する。
@@ -151,8 +149,8 @@ class TTSModel:
Returns:
NDArray[Any]: スタイルベクトル
"""
mean = self.__style_vectors[0]
style_vec = self.__style_vectors[style_id]
mean = self.style_vectors[0]
style_vec = self.style_vectors[style_id]
style_vec = mean + (style_vec - mean) * weight
return style_vec
@@ -169,7 +167,7 @@ class TTSModel:
NDArray[Any]: スタイルベクトル
"""
if self.__style_vector_inference is None:
if self.style_vector_inference is None:
# pyannote.audio は scikit-learn などの大量の重量級ライブラリに依存しているため、
# TTSModel.infer() に reference_audio_path を指定し音声からスタイルベクトルを推論する場合のみ遅延 import する
@@ -183,17 +181,17 @@ class TTSModel:
# スタイルベクトルを取得するための推論モデルを初期化
import torch
self.__style_vector_inference = pyannote.audio.Inference(
self.style_vector_inference = pyannote.audio.Inference(
model=pyannote.audio.Model.from_pretrained(
"pyannote/wespeaker-voxceleb-resnet34-LM"
),
window="whole",
)
self.__style_vector_inference.to(torch.device(self.device))
self.style_vector_inference.to(torch.device(self.device))
# 音声からスタイルベクトルを推論
xvec = self.__style_vector_inference(audio_path)
mean = self.__style_vectors[0]
xvec = self.style_vector_inference(audio_path)
mean = self.style_vectors[0]
xvec = mean + (xvec - mean) * weight
return xvec
@@ -310,9 +308,9 @@ class TTSModel:
from style_bert_vits2.models.infer import infer
# モデルがロードされていない場合はロードする
if self.__net_g is None:
if self.net_g is None:
self.load()
assert self.__net_g is not None
assert self.net_g is not None
# 通常のテキストから音声を生成
if not line_split:
@@ -326,7 +324,7 @@ class TTSModel:
sid=speaker_id,
language=language,
hps=self.hyper_parameters,
net_g=self.__net_g,
net_g=self.net_g,
device=self.device,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
@@ -352,7 +350,7 @@ class TTSModel:
sid=speaker_id,
language=language,
hps=self.hyper_parameters,
net_g=self.__net_g,
net_g=self.net_g,
device=self.device,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
@@ -368,9 +366,9 @@ class TTSModel:
from style_bert_vits2.models.infer_onnx import infer_onnx
# モデルがロードされていない場合はロードする
if self.__onnx_session is None:
if self.onnx_session is None:
self.load()
assert self.__onnx_session is not None
assert self.onnx_session is not None
# 通常のテキストから音声を生成
if not line_split:
@@ -383,9 +381,8 @@ class TTSModel:
sid=speaker_id,
language=language,
hps=self.hyper_parameters,
onnx_session=self.__onnx_session,
onnx_session=self.onnx_session,
onnx_providers=self.onnx_providers,
onnx_provider_options=self.onnx_provider_options,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
style_vec=style_vector,
@@ -409,9 +406,8 @@ class TTSModel:
sid=speaker_id,
language=language,
hps=self.hyper_parameters,
onnx_session=self.__onnx_session,
onnx_session=self.onnx_session,
onnx_providers=self.onnx_providers,
onnx_provider_options=self.onnx_provider_options,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
style_vec=style_vector,
@@ -487,13 +483,17 @@ class TTSModelHolder:
self.current_model = None
self.models_info = []
model_dirs = [d for d in self.root_dir.iterdir() if d.is_dir()]
model_dirs = sorted([d for d in self.root_dir.iterdir() if d.is_dir()])
for model_dir in model_dirs:
model_files = [
f
for f in model_dir.iterdir()
if f.suffix in [".pth", ".pt", ".safetensors", ".onnx"]
]
if model_dir.name.startswith("."):
continue
model_files = sorted(
[
f
for f in model_dir.iterdir()
if f.suffix in [".pth", ".pt", ".safetensors", ".onnx"]
]
)
if len(model_files) == 0:
logger.warning(f"No model files found in {model_dir}, so skip it")
continue