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

This commit is contained in:
tsukumi
2024-09-18 04:37:23 +09:00
parent e5a05a30cb
commit 7d721fdf2b
8 changed files with 65 additions and 26 deletions

View File

@@ -38,6 +38,7 @@ from style_bert_vits2.nlp import bert_models, onnx_bert_models
from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
from style_bert_vits2.nlp.japanese.user_dict import update_dict
from style_bert_vits2.tts_model import TTSModel, TTSModelHolder
from style_bert_vits2.utils import torch_device_to_onnx_providers
config = get_config()
@@ -103,15 +104,15 @@ if __name__ == "__main__":
bert_models.load_tokenizer(Languages.EN)
bert_models.load_model(Languages.ZH, device_map=device)
bert_models.load_tokenizer(Languages.ZH)
if device == "cpu":
onnx_provider = "CPUExecutionProvider"
else:
onnx_provider = ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"})
onnx_bert_models.load_model(Languages.JP, onnx_providers=[onnx_provider])
onnx_bert_models.load_model(
Languages.JP, onnx_providers=torch_device_to_onnx_providers(device)
)
onnx_bert_models.load_tokenizer(Languages.JP)
model_dir = Path(args.dir)
model_holder = TTSModelHolder(model_dir, device)
model_holder = TTSModelHolder(
model_dir, device, torch_device_to_onnx_providers(device)
)
if len(model_holder.model_names) == 0:
logger.error(f"Models not found in {model_dir}.")
sys.exit(1)