Improve: Support ONNX inference, add ONNX conversion script

This commit is contained in:
tsukumi
2024-09-17 16:55:03 +09:00
parent f42b9f0f89
commit 5e2c83c6c7
21 changed files with 44483 additions and 376 deletions

View File

@@ -34,7 +34,7 @@ from style_bert_vits2.constants import (
Languages,
)
from style_bert_vits2.logging import logger
from style_bert_vits2.nlp import bert_models
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
@@ -103,6 +103,12 @@ 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_tokenizer(Languages.JP)
model_dir = Path(args.dir)
model_holder = TTSModelHolder(model_dir, device)