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

Add: Test for ONNX inference code (without PyTorch dependency)

Browse Source
This commit is contained in:
tsukumi
2024-09-23 10:04:02 +09:00
parent 5a82105df3
commit eb5d245903
6 changed files with 104 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
convert_bert_onnx.py
View File

@@ -1,4 +1,4 @@
# usage: .venv/bin/python convert_bert_onnx.py --language JP
# Usage: .venv/bin/python convert_bert_onnx.py --language JP
# ref: https://github.com/tuna2134/sbv2-api/blob/main/convert/convert_deberta.py
import time

4
convert_onnx.py
View File

@@ -1,5 +1,5 @@
# usage: .venv/bin/python convert_onnx.py --model model_assets/koharune-ami/koharune-ami.safetensors
# usage: .venv/bin/python convert_onnx.py --model model_assets/ (All models in the directory will be converted)
# Usage: .venv/bin/python convert_onnx.py --model model_assets/koharune-ami/koharune-ami.safetensors
# Usage: .venv/bin/python convert_onnx.py --model model_assets/ (All models in the directory will be converted)
# ref: https://github.com/tuna2134/sbv2-api/blob/main/convert/convert_model.py
import time

26
pyproject.toml
View File

@@ -30,7 +30,7 @@ dependencies = [
"nltk<=3.8.1",
"num2words",
"numba",
"numpy",
"numpy<2",
"onnxruntime",
"pydantic>=2.0",
"pyopenjtalk-dict",
@@ -71,14 +71,17 @@ exclude = [".git", ".gitignore", ".gitattributes"]
[tool.hatch.build.targets.wheel]
packages = ["style_bert_vits2"]
# for PyTorch inference
[tool.hatch.envs.test]
dependencies = ["coverage[toml]>=6.5", "pytest", "scipy"]
features = ["torch"]
[tool.hatch.envs.test.scripts]
# Usage: `hatch run test:test`
test = "pytest {args:tests}"
test = "pytest tests/test_main.py::test_synthesize_cpu"
# Usage: `hatch run test:test-cuda`
test-cuda = "pytest tests/test_main.py::test_synthesize_cuda"
# Usage: `hatch run test:coverage`
test-cov = "coverage run -m pytest {args:tests}"
test-cov = "coverage run -m pytest tests/test_main.py::test_synthesize_cpu"
# Usage: `hatch run test:cov-report`
cov-report = ["- coverage combine", "coverage report"]
# Usage: `hatch run test:cov`
@@ -86,6 +89,23 @@ cov = ["test-cov", "cov-report"]
[[tool.hatch.envs.test.matrix]]
python = ["3.9", "3.10", "3.11"]
# for ONNX inference (without PyTorch dependency)
[tool.hatch.envs.test-onnx]
dependencies = ["coverage[toml]>=6.5", "pytest", "scipy", "onnxruntime-gpu"]
[tool.hatch.envs.test-onnx.scripts]
# Usage: `hatch run test-onnx:test`
test = "pytest tests/test_main.py::test_synthesize_onnx_cpu"
# Usage: `hatch run test-onnx:test-cuda`
test-cuda = "pytest tests/test_main.py::test_synthesize_onnx_cuda"
# Usage: `hatch run test-onnx:coverage`
test-cov = "coverage run -m pytest tests/test_main.py::test_synthesize_onnx_cpu"
# Usage: `hatch run test-onnx:cov-report`
cov-report = ["- coverage combine", "coverage report"]
# Usage: `hatch run test-onnx:cov`
cov = ["test-cov", "cov-report"]
[[tool.hatch.envs.test-onnx.matrix]]
python = ["3.9", "3.10", "3.11"]
[tool.hatch.envs.style]
detached = true
dependencies = ["black[jupyter]", "isort"]

13
style_bert_vits2/nlp/bert_models.py
View File

@@ -8,11 +8,12 @@ Style-Bert-VITS2 の学習・推論に必要な各言語ごとの BERT モデル
一度 load_model/tokenizer() で当該言語の BERT モデルがロードされていれば、ライブラリ内部のどこからでもロード済みのモデル/トークナイザーを取得できる。
"""
from __future__ import annotations
import gc
import time
from typing import Optional, Union, cast
from typing import TYPE_CHECKING, Optional, Union, cast
import torch
from transformers import (
AutoModelForMaskedLM,
AutoTokenizer,
@@ -27,6 +28,10 @@ from style_bert_vits2.constants import DEFAULT_BERT_MODEL_PATHS, Languages
from style_bert_vits2.logging import logger
if TYPE_CHECKING:
import torch
# 各言語ごとのロード済みの BERT モデルを格納する辞書
__loaded_models: dict[Languages, Union[PreTrainedModel, DebertaV2Model]] = {}
@@ -208,6 +213,8 @@ def unload_model(language: Languages) -> None:
language (Languages): アンロードする BERT モデルの言語
"""
import torch
if language in __loaded_models:
del __loaded_models[language]
gc.collect()
@@ -224,6 +231,8 @@ def unload_tokenizer(language: Languages) -> None:
language (Languages): アンロードする BERT トークナイザーの言語
"""
import torch
if language in __loaded_tokenizers:
del __loaded_tokenizers[language]
gc.collect()

10
style_bert_vits2/tts_model.py
View File

@@ -154,6 +154,16 @@ class TTSModel:
f"Model loaded successfully from {self.model_path} to {self.onnx_session.get_providers()[0]} ({time.time() - start_time:.2f}s)"
)
def unload(self) -> None:
"""
音声合成モデルをデバイスからアンロードする。
"""
if self.net_g is not None:
self.net_g = None
if self.onnx_session is not None:
self.onnx_session = None
def get_style_vector(self, style_id: int, weight: float = 1.0) -> NDArray[Any]:
"""
スタイルベクトルを取得する。

71
tests/test_main.py
View File

@@ -1,3 +1,5 @@
from typing import Any, Literal, Sequence, Union
import pytest
from scipy.io import wavfile
@@ -6,54 +8,95 @@ from style_bert_vits2.tts_model import TTSModelHolder
def synthesize(
inference_type: Literal["torch", "onnx"] = "torch",
device: str = "cpu",
onnx_providers: list[str] = ["CPUExecutionProvider"],
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = [
"CPUExecutionProvider"
],
):
# 音声合成モデルが配置されていれば、音声合成を実行
model_holder = TTSModelHolder(BASE_DIR / "model_assets", device, onnx_providers)
if len(model_holder.models_info) > 0:
# jvnv-F2-jp モデルを探す
# "koharune-ami" または "amitaro" モデルを探す
for model_info in model_holder.models_info:
if model_info.name == "jvnv-F2-jp":
if model_info.name == "koharune-ami" or model_info.name == "amitaro":
# Safetensors 形式または ONNX 形式のモデルファイルに絞り込む
if inference_type == "torch":
model_files = [
f
for f in model_info.files
if f.endswith(".safetensors") and not f.startswith(".")
]
else:
model_files = [
f
for f in model_info.files
if f.endswith(".onnx") and not f.startswith(".")
]
if len(model_files) == 0:
pytest.skip(f"音声合成モデル \"{model_info.name}\" のモデルファイルが見つかりませんでした。")
# モデルをロード
model = model_holder.get_model(model_info.name, model_files[0])
model.load()
# すべてのスタイルに対して音声合成を実行
for style in model_info.styles:
# 音声合成を実行
model = model_holder.get_model(model_info.name, model_info.files[0])
model.load()
sample_rate, audio_data = model.infer(
"あらゆる現実を、すべて自分のほうへねじ曲げたのだ。",
# 言語 (JP, EN, ZH / JP-Extra モデルの場合は JP のみ)
language=Languages.JP,
# 話者 ID (音声合成モデルに複数の話者が含まれる場合のみ必須、単一話者のみの場合は 0)
speaker_id=0,
# 感情表現の強さ (0.0 〜 1.0)
# テンポの緩急 (0.0 〜 1.0)
sdp_ratio=0.4,
# スタイル (Neutral, Happy など)
style=style,
# スタイルの強さ (0.0 〜 100.0)
style_weight=6.0,
style_weight=2.0,
)
# 音声データを保存
(BASE_DIR / "tests/wavs").mkdir(exist_ok=True, parents=True)
wav_file_path = BASE_DIR / f"tests/wavs/{style}.wav"
(BASE_DIR / f"tests/wavs/{model_info.name}").mkdir(
exist_ok=True, parents=True
)
wav_file_path = (
BASE_DIR / f"tests/wavs/{model_info.name}/{style}.wav"
)
with open(wav_file_path, "wb") as f:
wavfile.write(f, sample_rate, audio_data)
# 音声データが保存されたことを確認
assert wav_file_path.exists()
# wav_file_path.unlink()
# モデルをアンロード
model.unload()
else:
pytest.skip("音声合成モデルが見つかりませんでした。")
def test_synthesize_cpu():
synthesize(device="cpu")
synthesize(inference_type="torch", device="cpu")
# Windows環境ではtorchのcudaが簡単に入らないため、テストをスキップ
# def test_synthesize_cuda():
# synthesize(device="cuda")
def test_synthesize_cuda():
pytest.importorskip("torch.cuda")
synthesize(inference_type="torch", device="cuda")
def test_synthesize_onnx_cpu():
synthesize(inference_type="onnx", onnx_providers=["CPUExecutionProvider"])
def test_synthesize_onnx_cuda():
synthesize(
inference_type="onnx",
onnx_providers=[
("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"})
],
)