9c233630ef
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.
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from style_bert_vits2.constants import Languages
|
|
from style_bert_vits2.nlp import bert_models
|
|
|
|
|
|
def extract_bert_feature(
|
|
text: str,
|
|
word2ph: list[int],
|
|
device: str,
|
|
assist_text: Optional[str] = None,
|
|
assist_text_weight: float = 0.7,
|
|
) -> torch.Tensor:
|
|
"""
|
|
英語のテキストから BERT の特徴量を抽出する
|
|
|
|
Args:
|
|
text (str): 英語のテキスト
|
|
word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
|
|
device (str): 推論に利用するデバイス
|
|
assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
|
|
assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
|
|
|
|
Returns:
|
|
torch.Tensor: BERT の特徴量
|
|
"""
|
|
|
|
if device == "cuda" and not torch.cuda.is_available():
|
|
device = "cpu"
|
|
model = bert_models.load_model(Languages.EN).to(device) # type: ignore
|
|
|
|
style_res_mean = None
|
|
with torch.no_grad():
|
|
tokenizer = bert_models.load_tokenizer(Languages.EN)
|
|
inputs = tokenizer(text, return_tensors="pt")
|
|
for i in inputs:
|
|
inputs[i] = inputs[i].to(device) # type: ignore
|
|
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 = 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)
|
|
|
|
assert len(word2ph) == res.shape[0], (text, res.shape[0], len(word2ph))
|
|
word2phone = word2ph
|
|
phone_level_feature = []
|
|
for i in range(len(word2phone)):
|
|
if assist_text:
|
|
assert style_res_mean is not None
|
|
repeat_feature = (
|
|
res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
|
|
+ style_res_mean.repeat(word2phone[i], 1) * assist_text_weight
|
|
)
|
|
else:
|
|
repeat_feature = res[i].repeat(word2phone[i], 1)
|
|
phone_level_feature.append(repeat_feature)
|
|
|
|
phone_level_feature = torch.cat(phone_level_feature, dim=0)
|
|
|
|
return phone_level_feature.T
|