From 9c233630efbe48840dbff1e4d4774cf37413de42 Mon Sep 17 00:00:00 2001 From: tsukumi Date: Sun, 10 Mar 2024 14:18:57 +0000 Subject: [PATCH] Refactor: don't keep models in bert_feature module for each language BERT models and tokenizers are already stored and managed in the bert_models module and should not be stored here. In addition, since there may be situations where the user would like to use cpu instead of mps for inference when using it as a library, the automatic switching process to mps was removed. --- style_bert_vits2/nlp/chinese/bert_feature.py | 20 +++-------------- style_bert_vits2/nlp/english/bert_feature.py | 20 +++-------------- style_bert_vits2/nlp/japanese/bert_feature.py | 22 ++++--------------- 3 files changed, 10 insertions(+), 52 deletions(-) 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)