diff --git a/style_bert_vits2/nlp/chinese/bert_feature.py b/style_bert_vits2/nlp/chinese/bert_feature.py index f448b30..adc0788 100644 --- a/style_bert_vits2/nlp/chinese/bert_feature.py +++ b/style_bert_vits2/nlp/chinese/bert_feature.py @@ -1,16 +1,11 @@ -import sys from typing import Optional import torch -from transformers import PreTrainedModel from style_bert_vits2.constants import Languages from style_bert_vits2.nlp import bert_models -__models: dict[str, PreTrainedModel] = {} - - def extract_bert_feature( text: str, word2ph: list[int], @@ -32,18 +27,9 @@ def extract_bert_feature( torch.Tensor: BERT の特徴量 """ - if ( - sys.platform == "darwin" - and torch.backends.mps.is_available() - and device == "cpu" - ): - device = "mps" - if not device: - device = "cuda" if device == "cuda" and not torch.cuda.is_available(): device = "cpu" - if device not in __models.keys(): - __models[device] = bert_models.load_model(Languages.ZH).to(device) # type: ignore + model = bert_models.load_model(Languages.ZH).to(device) # type: ignore style_res_mean = None with torch.no_grad(): @@ -51,13 +37,13 @@ def extract_bert_feature( inputs = tokenizer(text, return_tensors="pt") for i in inputs: inputs[i] = inputs[i].to(device) # type: ignore - res = __models[device](**inputs, output_hidden_states=True) + res = model(**inputs, output_hidden_states=True) res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu() if assist_text: style_inputs = tokenizer(assist_text, return_tensors="pt") for i in style_inputs: style_inputs[i] = style_inputs[i].to(device) # type: ignore - style_res = __models[device](**style_inputs, output_hidden_states=True) + style_res = model(**style_inputs, output_hidden_states=True) style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu() style_res_mean = style_res.mean(0) diff --git a/style_bert_vits2/nlp/english/bert_feature.py b/style_bert_vits2/nlp/english/bert_feature.py index 27fd501..7692c9e 100644 --- a/style_bert_vits2/nlp/english/bert_feature.py +++ b/style_bert_vits2/nlp/english/bert_feature.py @@ -1,16 +1,11 @@ -import sys from typing import Optional import torch -from transformers import PreTrainedModel from style_bert_vits2.constants import Languages from style_bert_vits2.nlp import bert_models -__models: dict[str, PreTrainedModel] = {} - - def extract_bert_feature( text: str, word2ph: list[int], @@ -32,18 +27,9 @@ def extract_bert_feature( torch.Tensor: BERT の特徴量 """ - if ( - sys.platform == "darwin" - and torch.backends.mps.is_available() - and device == "cpu" - ): - device = "mps" - if not device: - device = "cuda" if device == "cuda" and not torch.cuda.is_available(): device = "cpu" - if device not in __models.keys(): - __models[device] = bert_models.load_model(Languages.EN).to(device) # type: ignore + model = bert_models.load_model(Languages.EN).to(device) # type: ignore style_res_mean = None with torch.no_grad(): @@ -51,13 +37,13 @@ def extract_bert_feature( inputs = tokenizer(text, return_tensors="pt") for i in inputs: inputs[i] = inputs[i].to(device) # type: ignore - res = __models[device](**inputs, output_hidden_states=True) + res = model(**inputs, output_hidden_states=True) res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu() if assist_text: style_inputs = tokenizer(assist_text, return_tensors="pt") for i in style_inputs: style_inputs[i] = style_inputs[i].to(device) # type: ignore - style_res = __models[device](**style_inputs, output_hidden_states=True) + style_res = model(**style_inputs, output_hidden_states=True) style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu() style_res_mean = style_res.mean(0) diff --git a/style_bert_vits2/nlp/japanese/bert_feature.py b/style_bert_vits2/nlp/japanese/bert_feature.py index 0d70014..d6e214b 100644 --- a/style_bert_vits2/nlp/japanese/bert_feature.py +++ b/style_bert_vits2/nlp/japanese/bert_feature.py @@ -1,17 +1,12 @@ -import sys from typing import Optional import torch -from transformers import PreTrainedModel from style_bert_vits2.constants import Languages from style_bert_vits2.nlp import bert_models from style_bert_vits2.nlp.japanese.g2p import text_to_sep_kata -__models: dict[str, PreTrainedModel] = {} - - def extract_bert_feature( text: str, word2ph: list[int], @@ -36,21 +31,12 @@ def extract_bert_feature( # 各単語が何文字かを作る `word2ph` を使う必要があるので、読めない文字は必ず無視する # でないと `word2ph` の結果とテキストの文字数結果が整合性が取れない text = "".join(text_to_sep_kata(text, raise_yomi_error=False)[0]) - if assist_text: assist_text = "".join(text_to_sep_kata(assist_text, raise_yomi_error=False)[0]) - if ( - sys.platform == "darwin" - and torch.backends.mps.is_available() - and device == "cpu" - ): - device = "mps" - if not device: - device = "cuda" + if device == "cuda" and not torch.cuda.is_available(): device = "cpu" - if device not in __models.keys(): - __models[device] = bert_models.load_model(Languages.JP).to(device) # type: ignore + model = bert_models.load_model(Languages.JP).to(device) # type: ignore style_res_mean = None with torch.no_grad(): @@ -58,13 +44,13 @@ def extract_bert_feature( inputs = tokenizer(text, return_tensors="pt") for i in inputs: inputs[i] = inputs[i].to(device) # type: ignore - res = __models[device](**inputs, output_hidden_states=True) + res = model(**inputs, output_hidden_states=True) res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu() if assist_text: style_inputs = tokenizer(assist_text, return_tensors="pt") for i in style_inputs: style_inputs[i] = style_inputs[i].to(device) # type: ignore - style_res = __models[device](**style_inputs, output_hidden_states=True) + style_res = model(**style_inputs, output_hidden_states=True) style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu() style_res_mean = style_res.mean(0)