From 62919e904e643d0d72ef18fd5424260a38157cb8 Mon Sep 17 00:00:00 2001 From: tsukumi Date: Thu, 7 Mar 2024 03:32:07 +0000 Subject: [PATCH] Refactor: moved the module for extracting BERT features from text in each language to style_bert_vits2/text_processing/(language)/bert_feature.py --- app.py | 3 +- bert_gen.py | 10 +-- data_utils.py | 6 +- server_fastapi.py | 4 +- style_bert_vits2/models/infer.py | 4 +- style_bert_vits2/text_processing/__init__.py | 68 +++++++++++++++++++ .../text_processing/chinese/bert_feature.py | 36 +++++++--- .../text_processing/english/bert_feature.py | 36 +++++++--- .../text_processing/japanese/bert_feature.py | 39 ++++++++--- text/__init__.py | 43 ------------ text/chinese.py | 10 +-- text/english.py | 6 -- webui_dataset.py | 1 + webui_merge.py | 3 +- webui_style_vectors.py | 3 +- webui_train.py | 1 + 16 files changed, 172 insertions(+), 101 deletions(-) create mode 100644 style_bert_vits2/text_processing/__init__.py rename text/chinese_bert.py => style_bert_vits2/text_processing/chinese/bert_feature.py (74%) rename text/english_bert_mock.py => style_bert_vits2/text_processing/english/bert_feature.py (65%) rename text/japanese_bert.py => style_bert_vits2/text_processing/japanese/bert_feature.py (63%) delete mode 100644 text/__init__.py diff --git a/app.py b/app.py index a0a215c..02c01cd 100644 --- a/app.py +++ b/app.py @@ -10,6 +10,7 @@ import gradio as gr import torch import yaml +from common.tts_model import ModelHolder from style_bert_vits2.constants import ( DEFAULT_ASSIST_TEXT_WEIGHT, DEFAULT_LENGTH, @@ -25,11 +26,11 @@ from style_bert_vits2.constants import ( Languages, ) from style_bert_vits2.logging import logger -from common.tts_model import ModelHolder from style_bert_vits2.models.infer import InvalidToneError from style_bert_vits2.text_processing.japanese.g2p_utils import g2kata_tone, kata_tone2phone_tone from style_bert_vits2.text_processing.japanese.normalizer import normalize_text + # Get path settings with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f: path_config: dict[str, str] = yaml.safe_load(f.read()) diff --git a/bert_gen.py b/bert_gen.py index fd0b54e..2b512d2 100644 --- a/bert_gen.py +++ b/bert_gen.py @@ -5,12 +5,12 @@ import torch import torch.multiprocessing as mp from tqdm import tqdm -from style_bert_vits2.models import commons import utils -from style_bert_vits2.logging import logger -from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT from config import config -from text import cleaned_text_to_sequence, get_bert +from style_bert_vits2.logging import logger +from style_bert_vits2.models import commons +from style_bert_vits2.text_processing import cleaned_text_to_sequence, extract_bert_feature +from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT def process_line(x): @@ -45,7 +45,7 @@ def process_line(x): bert = torch.load(bert_path) assert bert.shape[-1] == len(phone) except Exception: - bert = get_bert(text, word2ph, language_str, device) + bert = extract_bert_feature(text, word2ph, language_str, device) assert bert.shape[-1] == len(phone) torch.save(bert, bert_path) diff --git a/data_utils.py b/data_utils.py index 1118102..7738247 100644 --- a/data_utils.py +++ b/data_utils.py @@ -7,12 +7,12 @@ import torch import torch.utils.data from tqdm import tqdm -from style_bert_vits2.models import commons from config import config from mel_processing import mel_spectrogram_torch, spectrogram_torch -from text import cleaned_text_to_sequence -from style_bert_vits2.logging import logger from utils import load_filepaths_and_text, load_wav_to_torch +from style_bert_vits2.logging import logger +from style_bert_vits2.models import commons +from style_bert_vits2.text_processing import cleaned_text_to_sequence """Multi speaker version""" diff --git a/server_fastapi.py b/server_fastapi.py index 132cc17..ca9520c 100644 --- a/server_fastapi.py +++ b/server_fastapi.py @@ -20,6 +20,8 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, Response from scipy.io import wavfile +from common.tts_model import Model, ModelHolder +from config import config from style_bert_vits2.constants import ( DEFAULT_ASSIST_TEXT_WEIGHT, DEFAULT_LENGTH, @@ -33,8 +35,6 @@ from style_bert_vits2.constants import ( Languages, ) from style_bert_vits2.logging import logger -from common.tts_model import Model, ModelHolder -from config import config ln = config.server_config.language diff --git a/style_bert_vits2/models/infer.py b/style_bert_vits2/models/infer.py index 99999a9..c4f6aed 100644 --- a/style_bert_vits2/models/infer.py +++ b/style_bert_vits2/models/infer.py @@ -1,12 +1,12 @@ import torch import utils -from text import cleaned_text_to_sequence, get_bert from style_bert_vits2.constants import Languages from style_bert_vits2.logging import logger from style_bert_vits2.models import commons from style_bert_vits2.models.models import SynthesizerTrn from style_bert_vits2.models.models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra +from style_bert_vits2.text_processing import cleaned_text_to_sequence, extract_bert_feature from style_bert_vits2.text_processing.cleaner import clean_text from style_bert_vits2.text_processing.symbols import SYMBOLS @@ -77,7 +77,7 @@ def get_text( for i in range(len(word2ph)): word2ph[i] = word2ph[i] * 2 word2ph[0] += 1 - bert_ori = get_bert( + bert_ori = extract_bert_feature( norm_text, word2ph, language_str, diff --git a/style_bert_vits2/text_processing/__init__.py b/style_bert_vits2/text_processing/__init__.py new file mode 100644 index 0000000..42bb3a5 --- /dev/null +++ b/style_bert_vits2/text_processing/__init__.py @@ -0,0 +1,68 @@ +import torch + +from style_bert_vits2.constants import Languages +from style_bert_vits2.text_processing.symbols import ( + LANGUAGE_ID_MAP, + LANGUAGE_TONE_START_MAP, + SYMBOLS, +) + + +_symbol_to_id = {s: i for i, s in enumerate(SYMBOLS)} + + +def cleaned_text_to_sequence(cleaned_text: str, tones: list[int], language: Languages) -> tuple[list[int], list[int], list[int]]: + """ + Converts a string of text to a sequence of IDs corresponding to the symbols in the text. + + Args: + cleaned_text (str): string to convert to a sequence + tones (list[int]): List of tones + language (Languages): Language of the text + + Returns: + tuple[list[int], list[int], list[int]]: List of integers corresponding to the symbols in the text + """ + + phones = [_symbol_to_id[symbol] for symbol in cleaned_text] + tone_start = LANGUAGE_TONE_START_MAP[language] + tones = [i + tone_start for i in tones] + lang_id = LANGUAGE_ID_MAP[language] + lang_ids = [lang_id for i in phones] + + return phones, tones, lang_ids + + +def extract_bert_feature( + text: str, + word2ph: list[int], + language: Languages, + device: torch.device | str, + assist_text: str | None = None, + assist_text_weight: float = 0.7, +) -> torch.Tensor: + """ + テキストから BERT の特徴量を抽出する + + Args: + text (str): テキスト + word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト + language (Languages): テキストの言語 + device (torch.device | str): 推論に利用するデバイス + assist_text (str | None, optional): 補助テキスト (デフォルト: None) + assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7) + + Returns: + torch.Tensor: BERT の特徴量 + """ + + if language == Languages.JP: + from style_bert_vits2.text_processing.japanese.bert_feature import extract_bert_feature + elif language == Languages.EN: + from style_bert_vits2.text_processing.english.bert_feature import extract_bert_feature + elif language == Languages.ZH: + from style_bert_vits2.text_processing.chinese.bert_feature import extract_bert_feature + else: + raise ValueError(f"Language {language} not supported") + + return extract_bert_feature(text, word2ph, device, assist_text, assist_text_weight) diff --git a/text/chinese_bert.py b/style_bert_vits2/text_processing/chinese/bert_feature.py similarity index 74% rename from text/chinese_bert.py rename to style_bert_vits2/text_processing/chinese/bert_feature.py index 2ee6e9d..c2085a4 100644 --- a/text/chinese_bert.py +++ b/style_bert_vits2/text_processing/chinese/bert_feature.py @@ -1,22 +1,36 @@ import sys import torch +from transformers import PreTrainedModel -from config import config from style_bert_vits2.constants import Languages from style_bert_vits2.text_processing import bert_models -models = dict() +models: dict[str, PreTrainedModel] = {} -def get_bert_feature( +def extract_bert_feature( text: str, - word2ph, - device = config.bert_gen_config.device, + word2ph: list[int], + device: torch.device | str, assist_text: str | None = None, assist_text_weight: float = 0.7, -): +) -> torch.Tensor: + """ + 中国語のテキストから BERT の特徴量を抽出する + + Args: + text (str): 中国語のテキスト + word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト + device (torch.device | str): 推論に利用するデバイス + assist_text (str | None, optional): 補助テキスト (デフォルト: None) + assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7) + + Returns: + torch.Tensor: BERT の特徴量 + """ + if ( sys.platform == "darwin" and torch.backends.mps.is_available() @@ -28,26 +42,30 @@ def get_bert_feature( 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) + models[device] = bert_models.load_model(Languages.ZH).to(device) # type: ignore + + style_res_mean = None with torch.no_grad(): tokenizer = bert_models.load_tokenizer(Languages.ZH) inputs = tokenizer(text, return_tensors="pt") for i in inputs: - inputs[i] = inputs[i].to(device) + inputs[i] = inputs[i].to(device) # type: ignore res = models[device](**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) + style_inputs[i] = style_inputs[i].to(device) # type: ignore style_res = models[device](**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) == len(text) + 2 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 diff --git a/text/english_bert_mock.py b/style_bert_vits2/text_processing/english/bert_feature.py similarity index 65% rename from text/english_bert_mock.py rename to style_bert_vits2/text_processing/english/bert_feature.py index 8e57df1..d9c1025 100644 --- a/text/english_bert_mock.py +++ b/style_bert_vits2/text_processing/english/bert_feature.py @@ -1,22 +1,36 @@ import sys import torch +from transformers import PreTrainedModel -from config import config from style_bert_vits2.constants import Languages from style_bert_vits2.text_processing import bert_models -models = dict() +models: dict[str, PreTrainedModel] = {} -def get_bert_feature( +def extract_bert_feature( text: str, - word2ph, - device = config.bert_gen_config.device, + word2ph: list[int], + device: torch.device | str, assist_text: str | None = None, assist_text_weight: float = 0.7, -): +) -> torch.Tensor: + """ + 英語のテキストから BERT の特徴量を抽出する + + Args: + text (str): 英語のテキスト + word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト + device (torch.device | str): 推論に利用するデバイス + assist_text (str | None, optional): 補助テキスト (デフォルト: None) + assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7) + + Returns: + torch.Tensor: BERT の特徴量 + """ + if ( sys.platform == "darwin" and torch.backends.mps.is_available() @@ -28,26 +42,30 @@ def get_bert_feature( 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) + models[device] = 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) + inputs[i] = inputs[i].to(device) # type: ignore res = models[device](**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) + style_inputs[i] = style_inputs[i].to(device) # type: ignore style_res = models[device](**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 diff --git a/text/japanese_bert.py b/style_bert_vits2/text_processing/japanese/bert_feature.py similarity index 63% rename from text/japanese_bert.py rename to style_bert_vits2/text_processing/japanese/bert_feature.py index 4c1bccc..078bb5c 100644 --- a/text/japanese_bert.py +++ b/style_bert_vits2/text_processing/japanese/bert_feature.py @@ -1,25 +1,39 @@ import sys import torch +from transformers import PreTrainedModel -from config import config from style_bert_vits2.constants import Languages from style_bert_vits2.text_processing import bert_models from style_bert_vits2.text_processing.japanese.g2p import text_to_sep_kata -models = dict() +models: dict[str, PreTrainedModel] = {} -def get_bert_feature( +def extract_bert_feature( text: str, - word2ph, - device = config.bert_gen_config.device, + word2ph: list[int], + device: torch.device | str, assist_text: str | None = None, assist_text_weight: float = 0.7, -): - # 各単語が何文字かを作る`word2ph`を使う必要があるので、読めない文字は必ず無視する - # でないと`word2ph`の結果とテキストの文字数結果が整合性が取れない +) -> torch.Tensor: + """ + 日本語のテキストから BERT の特徴量を抽出する + + Args: + text (str): 日本語のテキスト + word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト + device (torch.device | str): 推論に利用するデバイス + assist_text (str | None, optional): 補助テキスト (デフォルト: None) + assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7) + + Returns: + torch.Tensor: BERT の特徴量 + """ + + # 各単語が何文字かを作る `word2ph` を使う必要があるので、読めない文字は必ず無視する + # でないと `word2ph` の結果とテキストの文字数結果が整合性が取れない text = "".join(text_to_sep_kata(text, raise_yomi_error=False)[0]) if assist_text: @@ -35,18 +49,20 @@ def get_bert_feature( 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) + models[device] = bert_models.load_model(Languages.JP).to(device) # type: ignore + + style_res_mean = None with torch.no_grad(): tokenizer = bert_models.load_tokenizer(Languages.JP) inputs = tokenizer(text, return_tensors="pt") for i in inputs: - inputs[i] = inputs[i].to(device) + inputs[i] = inputs[i].to(device) # type: ignore res = models[device](**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) + style_inputs[i] = style_inputs[i].to(device) # type: ignore style_res = models[device](**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) @@ -56,6 +72,7 @@ def get_bert_feature( 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 diff --git a/text/__init__.py b/text/__init__.py deleted file mode 100644 index efff830..0000000 --- a/text/__init__.py +++ /dev/null @@ -1,43 +0,0 @@ -from style_bert_vits2.constants import Languages -from style_bert_vits2.text_processing.symbols import * - - -_symbol_to_id = {s: i for i, s in enumerate(SYMBOLS)} - - -def cleaned_text_to_sequence(cleaned_text: str, tones: list[int], language: Languages): - """ - Converts a string of text to a sequence of IDs corresponding to the symbols in the text. - - Args: - text: string to convert to a sequence - - Returns: - List of integers corresponding to the symbols in the text - """ - phones = [_symbol_to_id[symbol] for symbol in cleaned_text] - tone_start = LANGUAGE_TONE_START_MAP[language] - tones = [i + tone_start for i in tones] - lang_id = LANGUAGE_ID_MAP[language] - lang_ids = [lang_id for i in phones] - return phones, tones, lang_ids - - -def get_bert( - text: str, - word2ph, - language: Languages, - device: str, - assist_text: str | None = None, - assist_text_weight: float = 0.7, -): - if language == Languages.ZH: - from .chinese_bert import get_bert_feature - elif language == Languages.EN: - from .english_bert_mock import get_bert_feature - elif language == Languages.JP: - from .japanese_bert import get_bert_feature - else: - raise ValueError(f"Language {language} not supported") - - return get_bert_feature(text, word2ph, device, assist_text, assist_text_weight) diff --git a/text/chinese.py b/text/chinese.py index 3d9c392..94266c2 100644 --- a/text/chinese.py +++ b/text/chinese.py @@ -176,20 +176,14 @@ def normalize_text(text): return text -def get_bert_feature(text, word2ph): - from text import chinese_bert - - return chinese_bert.get_bert_feature(text, word2ph) - - if __name__ == "__main__": - from text.chinese_bert import get_bert_feature + from style_bert_vits2.text_processing.chinese.bert_feature import extract_bert_feature text = "啊！但是《原神》是由,米哈\游自主， [研发]的一款全.新开放世界.冒险游戏" text = normalize_text(text) print(text) phones, tones, word2ph = g2p(text) - bert = get_bert_feature(text, word2ph) + bert = extract_bert_feature(text, word2ph, 'cuda') print(phones, tones, word2ph, bert.shape) diff --git a/text/english.py b/text/english.py index a6d71e9..419be3b 100644 --- a/text/english.py +++ b/text/english.py @@ -477,12 +477,6 @@ def g2p(text): return phones, tones, word2ph -def get_bert_feature(text, word2ph): - from text import english_bert_mock - - return english_bert_mock.get_bert_feature(text, word2ph) - - if __name__ == "__main__": # print(get_dict()) # print(eng_word_to_phoneme("hello")) diff --git a/webui_dataset.py b/webui_dataset.py index 3ad63c1..1796864 100644 --- a/webui_dataset.py +++ b/webui_dataset.py @@ -8,6 +8,7 @@ from style_bert_vits2.constants import GRADIO_THEME from style_bert_vits2.logging import logger from style_bert_vits2.utils.subprocess import run_script_with_log + # Get path settings with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f: path_config: dict[str, str] = yaml.safe_load(f.read()) diff --git a/webui_merge.py b/webui_merge.py index 0a39902..5bf0568 100644 --- a/webui_merge.py +++ b/webui_merge.py @@ -11,9 +11,10 @@ import yaml from safetensors import safe_open from safetensors.torch import save_file +from common.tts_model import Model, ModelHolder from style_bert_vits2.constants import DEFAULT_STYLE, GRADIO_THEME from style_bert_vits2.logging import logger -from common.tts_model import Model, ModelHolder + voice_keys = ["dec"] voice_pitch_keys = ["flow"] diff --git a/webui_style_vectors.py b/webui_style_vectors.py index cf53ca2..1cbabbd 100644 --- a/webui_style_vectors.py +++ b/webui_style_vectors.py @@ -12,9 +12,10 @@ from sklearn.cluster import DBSCAN, AgglomerativeClustering, KMeans from sklearn.manifold import TSNE from umap import UMAP +from config import config from style_bert_vits2.constants import DEFAULT_STYLE, GRADIO_THEME from style_bert_vits2.logging import logger -from config import config + # Get path settings with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f: diff --git a/webui_train.py b/webui_train.py index 6c89b26..fda31ca 100644 --- a/webui_train.py +++ b/webui_train.py @@ -19,6 +19,7 @@ from style_bert_vits2.logging import logger from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT from style_bert_vits2.utils.subprocess import run_script_with_log, second_elem_of + logger_handler = None tensorboard_executed = False