Improve: Support ONNX inference, add ONNX conversion script
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user