From 860957006be81500067f28cd747bf4efba9e7fea Mon Sep 17 00:00:00 2001 From: litagin02 Date: Wed, 13 Mar 2024 16:27:06 +0900 Subject: [PATCH] Add typing and load bert models on the top --- bert_gen.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/bert_gen.py b/bert_gen.py index 068afd6..933fa21 100644 --- a/bert_gen.py +++ b/bert_gen.py @@ -13,7 +13,15 @@ from style_bert_vits2.nlp import cleaned_text_to_sequence, extract_bert_feature from style_bert_vits2.nlp.japanese import pyopenjtalk_worker from style_bert_vits2.nlp.japanese.user_dict import update_dict from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT +from style_bert_vits2.nlp import bert_models +from style_bert_vits2.constants import Languages +bert_models.load_model(Languages.JP) +bert_models.load_tokenizer(Languages.JP) +bert_models.load_model(Languages.EN) +bert_models.load_tokenizer(Languages.EN) +bert_models.load_model(Languages.ZH) +bert_models.load_tokenizer(Languages.ZH) # このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化 pyopenjtalk_worker.initialize_worker() @@ -22,7 +30,7 @@ pyopenjtalk_worker.initialize_worker() update_dict() -def process_line(x): +def process_line(x: tuple[str, bool]): line, add_blank = x device = config.bert_gen_config.device if config.bert_gen_config.use_multi_device: @@ -30,9 +38,9 @@ def process_line(x): rank = rank[0] if len(rank) > 0 else 0 if torch.cuda.is_available(): gpu_id = rank % torch.cuda.device_count() - device = torch.device(f"cuda:{gpu_id}") + device = f"cuda:{gpu_id}" else: - device = torch.device("cpu") + device = "cpu" wav_path, _, language_str, text, phones, tone, word2ph = line.strip().split("|") phone = phones.split(" ") tone = [int(i) for i in tone.split(" ")] @@ -72,7 +80,7 @@ if __name__ == "__main__": args, _ = parser.parse_known_args() config_path = args.config hps = HyperParameters.load_from_json(config_path) - lines = [] + lines: list[str] = [] with open(hps.data.training_files, "r", encoding="utf-8") as f: lines.extend(f.readlines())