From eb5d245903bd4a3f5d1824ba8f4aec1e951d13c3 Mon Sep 17 00:00:00 2001 From: tsukumi Date: Mon, 23 Sep 2024 10:04:02 +0900 Subject: [PATCH] Add: Test for ONNX inference code (without PyTorch dependency) --- convert_bert_onnx.py | 2 +- convert_onnx.py | 4 +- pyproject.toml | 26 +++++++++-- style_bert_vits2/nlp/bert_models.py | 13 +++++- style_bert_vits2/tts_model.py | 10 ++++ tests/test_main.py | 71 +++++++++++++++++++++++------ 6 files changed, 104 insertions(+), 22 deletions(-) diff --git a/convert_bert_onnx.py b/convert_bert_onnx.py index bfd0591..f30935c 100644 --- a/convert_bert_onnx.py +++ b/convert_bert_onnx.py @@ -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 diff --git a/convert_onnx.py b/convert_onnx.py index 1b1fe5d..b61f9b1 100644 --- a/convert_onnx.py +++ b/convert_onnx.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index ca77e23..0a10eb8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/style_bert_vits2/nlp/bert_models.py b/style_bert_vits2/nlp/bert_models.py index da74a4f..2775e29 100644 --- a/style_bert_vits2/nlp/bert_models.py +++ b/style_bert_vits2/nlp/bert_models.py @@ -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() diff --git a/style_bert_vits2/tts_model.py b/style_bert_vits2/tts_model.py index 2f9e592..28bf714 100644 --- a/style_bert_vits2/tts_model.py +++ b/style_bert_vits2/tts_model.py @@ -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]: """ スタイルベクトルを取得する。 diff --git a/tests/test_main.py b/tests/test_main.py index 2aa852f..74e2365 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -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"}) + ], + )