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:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user