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.
This commit is contained in:
tsukumi
2024-03-10 14:18:57 +00:00
parent 00bf496325
commit 9c233630ef
3 changed files with 10 additions and 52 deletions

View File

@@ -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)