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

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"})
+        ],
+    )