Fix: CUDAExecutionProvider was not being used to infer Style-Bert-VITS2 ONNX models even though CUDA was available

style_bert_vits2/tts_model.py

@@ -131,6 +131,9 @@ class TTSModel:
                 device=self.device,
                 hps=self.hyper_parameters,
             )
+            logger.info(
+                f"Model loaded successfully from {self.model_path} to \"{self.device}\" device ({time.time() - start_time:.2f}s)"
+            )
 
         # ONNX 推論時
         else:
@@ -138,10 +141,9 @@ class TTSModel:
                 path_or_bytes=str(self.model_path),
                 providers=self.onnx_providers,
             )
-
-        logger.info(
-            f"Model loaded successfully from {self.model_path} ({time.time() - start_time:.2f}s)"
-        )
+            logger.info(
+                f"Model loaded successfully from {self.model_path} to {self.onnx_session.get_providers()[0]} ({time.time() - start_time:.2f}s)"
+            )
 
     def get_style_vector(self, style_id: int, weight: float = 1.0) -> NDArray[Any]:
         """
@@ -447,13 +449,18 @@ class TTSModelInfo(BaseModel):
 
 class TTSModelHolder:
     """
-    Style-Bert-Vits2 の音声合成モデルを管理するクラス。
+    Style-Bert-VITS2 の音声合成モデルを管理するクラス。
     model_holder.models_info から指定されたディレクトリ内にある音声合成モデルの一覧を取得できる。
     """
 
-    def __init__(self, model_root_dir: Path, device: str) -> None:
+    def __init__(
+        self,
+        model_root_dir: Path,
+        device: str,
+        onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
+    ) -> None:
         """
-        Style-Bert-Vits2 の音声合成モデルを管理するクラスを初期化する。
+        Style-Bert-VITS2 の音声合成モデルを管理するクラスを初期化する。
         音声合成モデルは下記のように配置されていることを前提とする (.safetensors / .onnx のファイル名は自由) 。
         ```
         model_root_dir
@@ -470,11 +477,13 @@ class TTSModelHolder:
 
         Args:
             model_root_dir (Path): 音声合成モデルが配置されているディレクトリのパス
-            device (str): 音声合成時に利用するデバイス (cpu, cuda, mps など)
+            device (str): PyTorch 推論での音声合成時に利用するデバイス (cpu, cuda, mps など)
+            onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
         """
 
         self.root_dir: Path = model_root_dir
         self.device: str = device
+        self.onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = onnx_providers  # fmt: skip
         self.model_files_dict: dict[str, list[Path]] = {}
         self.current_model: Optional[TTSModel] = None
         self.model_names: list[str] = []
@@ -551,6 +560,7 @@ class TTSModelHolder:
                 config_path=self.root_dir / model_name / "config.json",
                 style_vec_path=self.root_dir / model_name / "style_vectors.npy",
                 device=self.device,
+                onnx_providers=self.onnx_providers,
             )
 
         return self.current_model
@@ -580,6 +590,7 @@ class TTSModelHolder:
             config_path=self.root_dir / model_name / "config.json",
             style_vec_path=self.root_dir / model_name / "style_vectors.npy",
             device=self.device,
+            onnx_providers=self.onnx_providers,
         )
         speakers = list(self.current_model.spk2id.keys())
         styles = list(self.current_model.style2id.keys())

style_bert_vits2/utils/__init__.py

@@ -0,0 +1,15 @@
+from typing import Any, Sequence, Union
+
+
+def torch_device_to_onnx_providers(
+    device: str,
+) -> Sequence[Union[str, tuple[str, dict[str, Any]]]]:
+    if device.startswith("cuda"):
+        # cudnn_conv_algo_search を DEFAULT にすると推論速度が大幅に向上する
+        # ref: https://medium.com/neuml/debug-onnx-gpu-performance-c9290fe07459
+        return [
+            ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
+            ("CPUExecutionProvider", {}),
+        ]
+    else:
+        return ["CPUExecutionProvider"]