refactor: load base model once for api v2 % v3

This commit is contained in:
kevin.zhang 2024-05-16 15:11:25 +08:00
parent b8b273ad0c
commit 95c761f492
7 changed files with 298 additions and 285 deletions

View File

@ -5,6 +5,7 @@ import random
import traceback
from tqdm import tqdm
now_dir = os.getcwd()
sys.path.append(now_dir)
import ffmpeg
@ -26,6 +27,7 @@ from my_utils import load_audio
from module.mel_processing import spectrogram_torch
from TTS_infer_pack.text_segmentation_method import splits
from TTS_infer_pack.TextPreprocessor import TextPreprocessor
i18n = I18nAuto()
# configs/tts_infer.yaml
@ -49,6 +51,7 @@ custom:
"""
def set_seed(seed: int):
seed = int(seed)
seed = seed if seed != -1 else random.randrange(1 << 32)
@ -71,6 +74,7 @@ def set_seed(seed:int):
pass
return seed
class TTS_Config:
default_configs = {
"device": "cpu",
@ -79,8 +83,10 @@ class TTS_Config:
"vits_weights_path": "GPT_SoVITS/pretrained_models/s2G488k.pth",
"cnhuhbert_base_path": "GPT_SoVITS/pretrained_models/chinese-hubert-base",
"bert_base_path": "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large",
"load_base": True,
}
configs: dict = None
def __init__(self, configs: Union[dict, str] = None):
# 设置默认配置文件路径
@ -105,14 +111,13 @@ class TTS_Config:
self.configs: dict = configs.get("custom", deepcopy(self.default_configs))
self.device = self.configs.get("device", torch.device("cpu"))
self.is_half = self.configs.get("is_half", False)
self.t2s_weights_path = self.configs.get("t2s_weights_path", None)
self.vits_weights_path = self.configs.get("vits_weights_path", None)
self.bert_base_path = self.configs.get("bert_base_path", None)
self.cnhuhbert_base_path = self.configs.get("cnhuhbert_base_path", None)
self.load_base = self.configs.get("load_base", True)
if (self.t2s_weights_path in [None, ""]) or (not os.path.exists(self.t2s_weights_path)):
self.t2s_weights_path = self.default_configs['t2s_weights_path']
@ -128,7 +133,6 @@ class TTS_Config:
print(f"fall back to default cnhuhbert_base_path: {self.cnhuhbert_base_path}")
self.update_configs()
self.max_sec = None
self.hz: int = 50
self.semantic_frame_rate: str = "25hz"
@ -141,7 +145,6 @@ class TTS_Config:
self.languages: list = ["auto", "en", "zh", "ja", "all_zh", "all_ja"]
def _load_configs(self, configs_path: str) -> dict:
with open(configs_path, 'r') as f:
configs = yaml.load(f, Loader=yaml.FullLoader)
@ -168,6 +171,7 @@ class TTS_Config:
"vits_weights_path": self.vits_weights_path,
"bert_base_path": self.bert_base_path,
"cnhuhbert_base_path": self.cnhuhbert_base_path,
"load_base": self.load_base,
}
return self.config
@ -190,6 +194,10 @@ class TTS_Config:
class TTS:
bert_tokenizer: AutoTokenizer = None
bert_model: AutoModelForMaskedLM = None
cnhuhbert_model: CNHubert = None
def __init__(self, configs: Union[dict, str, TTS_Config]):
if isinstance(configs, TTS_Config):
self.configs = configs
@ -198,18 +206,17 @@ class TTS:
self.t2s_model: Text2SemanticLightningModule = None
self.vits_model: SynthesizerTrn = None
self.bert_tokenizer:AutoTokenizer = None
self.bert_model:AutoModelForMaskedLM = None
self.cnhuhbert_model:CNHubert = None
# self.bert_tokenizer:AutoTokenizer = None
# self.bert_model:AutoModelForMaskedLM = None
# self.cnhuhbert_model:CNHubert = None
self._init_models()
self.text_preprocessor: TextPreprocessor = \
TextPreprocessor(self.bert_model,
self.bert_tokenizer,
TextPreprocessor(TTS.bert_model,
TTS.bert_tokenizer,
self.configs.device)
self.prompt_cache: dict = {
"ref_audio_path": None,
"prompt_semantic": None,
@ -221,37 +228,40 @@ class TTS:
"norm_text": None,
}
self.stop_flag: bool = False
self.precision: torch.dtype = torch.float16 if self.configs.is_half else torch.float32
def _init_models(self,):
def _init_models(self):
self.init_t2s_weights(self.configs.t2s_weights_path)
self.init_vits_weights(self.configs.vits_weights_path)
self.init_bert_weights(self.configs.bert_base_path)
self.init_cnhuhbert_weights(self.configs.cnhuhbert_base_path)
if self.configs.load_base:
TTS.init_bert_weights(self.configs)
TTS.init_cnhuhbert_weights(self.configs)
# self.enable_half_precision(self.configs.is_half)
@staticmethod
def init_base_models(configs: TTS_Config):
TTS.init_bert_weights(configs)
TTS.init_cnhuhbert_weights(configs)
@staticmethod
def init_cnhuhbert_weights(configs: TTS_Config):
print(f"Loading CNHuBERT weights from {configs.cnhuhbert_base_path}")
TTS.cnhuhbert_model = CNHubert(configs.cnhuhbert_base_path)
TTS.cnhuhbert_model = TTS.cnhuhbert_model.eval()
TTS.cnhuhbert_model = TTS.cnhuhbert_model.to(configs.device)
if configs.is_half and str(configs.device) != "cpu":
TTS.cnhuhbert_model = TTS.cnhuhbert_model.half()
def init_cnhuhbert_weights(self, base_path: str):
print(f"Loading CNHuBERT weights from {base_path}")
self.cnhuhbert_model = CNHubert(base_path)
self.cnhuhbert_model=self.cnhuhbert_model.eval()
self.cnhuhbert_model = self.cnhuhbert_model.to(self.configs.device)
if self.configs.is_half and str(self.configs.device)!="cpu":
self.cnhuhbert_model = self.cnhuhbert_model.half()
def init_bert_weights(self, base_path: str):
print(f"Loading BERT weights from {base_path}")
self.bert_tokenizer = AutoTokenizer.from_pretrained(base_path)
self.bert_model = AutoModelForMaskedLM.from_pretrained(base_path)
self.bert_model=self.bert_model.eval()
self.bert_model = self.bert_model.to(self.configs.device)
if self.configs.is_half and str(self.configs.device)!="cpu":
self.bert_model = self.bert_model.half()
@staticmethod
def init_bert_weights(configs: TTS_Config):
print(f"Loading BERT weights from {configs.bert_base_path}")
TTS.bert_tokenizer = AutoTokenizer.from_pretrained(configs.bert_base_path)
TTS.bert_model = AutoModelForMaskedLM.from_pretrained(configs.bert_base_path)
TTS.bert_model = TTS.bert_model.eval()
TTS.bert_model = TTS.bert_model.to(configs.device)
if configs.is_half and str(configs.device) != "cpu":
TTS.bert_model = TTS.bert_model.half()
def init_vits_weights(self, weights_path: str):
print(f"Loading VITS weights from {weights_path}")
@ -284,7 +294,6 @@ class TTS:
if self.configs.is_half and str(self.configs.device) != "cpu":
self.vits_model = self.vits_model.half()
def init_t2s_weights(self, weights_path: str):
print(f"Loading Text2Semantic weights from {weights_path}")
self.configs.t2s_weights_path = weights_path
@ -320,19 +329,19 @@ class TTS:
self.t2s_model = self.t2s_model.half()
if self.vits_model is not None:
self.vits_model = self.vits_model.half()
if self.bert_model is not None:
self.bert_model =self.bert_model.half()
if self.cnhuhbert_model is not None:
self.cnhuhbert_model = self.cnhuhbert_model.half()
if TTS.bert_model is not None:
TTS.bert_model = TTS.bert_model.half()
if TTS.cnhuhbert_model is not None:
TTS.cnhuhbert_model = TTS.cnhuhbert_model.half()
else:
if self.t2s_model is not None:
self.t2s_model = self.t2s_model.float()
if self.vits_model is not None:
self.vits_model = self.vits_model.float()
if self.bert_model is not None:
self.bert_model = self.bert_model.float()
if self.cnhuhbert_model is not None:
self.cnhuhbert_model = self.cnhuhbert_model.float()
if TTS.bert_model is not None:
TTS.bert_model = TTS.bert_model.float()
if TTS.cnhuhbert_model is not None:
TTS.cnhuhbert_model = TTS.cnhuhbert_model.float()
def set_device(self, device: torch.device):
'''
@ -346,10 +355,10 @@ class TTS:
self.t2s_model = self.t2s_model.to(device)
if self.vits_model is not None:
self.vits_model = self.vits_model.to(device)
if self.bert_model is not None:
self.bert_model = self.bert_model.to(device)
if self.cnhuhbert_model is not None:
self.cnhuhbert_model = self.cnhuhbert_model.to(device)
if TTS.bert_model is not None:
TTS.bert_model = TTS.bert_model.to(device)
if TTS.cnhuhbert_model is not None:
TTS.cnhuhbert_model = TTS.cnhuhbert_model.to(device)
def set_ref_audio(self, ref_audio_path: str):
'''
@ -380,7 +389,6 @@ class TTS:
# self.refer_spec = spec
self.prompt_cache["refer_spec"] = spec
def _set_prompt_semantic(self, ref_wav_path: str):
zero_wav = np.zeros(
int(self.configs.sampling_rate * 0.3),
@ -473,7 +481,6 @@ class TTS:
batch_index_list.append([])
batch_index_list[-1].append(i)
for batch_idx, index_list in enumerate(batch_index_list):
item_list = [data[idx] for idx in index_list]
phones_list = []
@ -513,7 +520,6 @@ class TTS:
all_phones_batch = all_phones_list
all_bert_features_batch = all_bert_features_list
max_len = max(bert_max_len, phones_max_len)
# phones_batch = self.batch_sequences(phones_list, axis=0, pad_value=0, max_length=max_len)
#### 直接对phones和bert_features进行pad。padding策略会影响T2S模型生成的结果但不直接影响复读概率。影响复读概率的主要因素是mask的策略
@ -652,7 +658,8 @@ class TTS:
if ref_audio_path in [None, ""] and \
((self.prompt_cache["prompt_semantic"] is None) or (self.prompt_cache["refer_spec"] is None)):
raise ValueError("ref_audio_path cannot be empty, when the reference audio is not set using set_ref_audio()")
raise ValueError(
"ref_audio_path cannot be empty, when the reference audio is not set using set_ref_audio()")
###### setting reference audio and prompt text preprocessing ########
t0 = ttime()
@ -706,7 +713,8 @@ class TTS:
batch_data = []
print(i18n("############ 提取文本Bert特征 ############"))
for text in tqdm(batch_texts):
phones, bert_features, norm_text = self.text_preprocessor.segment_and_extract_feature_for_text(text, text_lang)
phones, bert_features, norm_text = self.text_preprocessor.segment_and_extract_feature_for_text(text,
text_lang)
if phones is None:
continue
res = {
@ -727,7 +735,6 @@ class TTS:
)
return batch[0]
t2 = ttime()
try:
print("############ 推理 ############")
@ -755,8 +762,8 @@ class TTS:
if no_prompt_text:
prompt = None
else:
prompt = self.prompt_cache["prompt_semantic"].expand(len(all_phoneme_ids), -1).to(self.configs.device)
prompt = self.prompt_cache["prompt_semantic"].expand(len(all_phoneme_ids), -1).to(
self.configs.device)
pred_semantic_list, idx_list = self.t2s_model.model.infer_panel(
all_phoneme_ids,
@ -798,7 +805,8 @@ class TTS:
# ## vits并行推理 method 2
pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
upsample_rate = math.prod(self.vits_model.upsample_rates)
audio_frag_idx = [pred_semantic_list[i].shape[0]*2*upsample_rate for i in range(0, len(pred_semantic_list))]
audio_frag_idx = [pred_semantic_list[i].shape[0] * 2 * upsample_rate for i in
range(0, len(pred_semantic_list))]
audio_frag_end_idx = [sum(audio_frag_idx[:i + 1]) for i in range(0, len(audio_frag_idx))]
all_pred_semantic = torch.cat(pred_semantic_list).unsqueeze(0).unsqueeze(0).to(self.configs.device)
_batch_phones = torch.cat(batch_phones).unsqueeze(0).to(self.configs.device)
@ -806,7 +814,8 @@ class TTS:
all_pred_semantic, _batch_phones, refer_audio_spec
).detach()[0, 0, :])
audio_frag_end_idx.insert(0, 0)
batch_audio_fragment= [_batch_audio_fragment[audio_frag_end_idx[i-1]:audio_frag_end_idx[i]] for i in range(1, len(audio_frag_end_idx))]
batch_audio_fragment = [_batch_audio_fragment[audio_frag_end_idx[i - 1]:audio_frag_end_idx[i]] for i in
range(1, len(audio_frag_end_idx))]
# ## vits串行推理
# for i, idx in enumerate(idx_list):
@ -894,14 +903,12 @@ class TTS:
audio_fragment: torch.Tensor = torch.cat([audio_fragment, zero_wav], dim=0)
audio[i][j] = audio_fragment.cpu().numpy()
if split_bucket:
audio = self.recovery_order(audio, batch_index_list)
else:
# audio = [item for batch in audio for item in batch]
audio = sum(audio, [])
audio = np.concatenate(audio, 0)
audio = (audio * 32768).astype(np.int16)
@ -914,8 +921,6 @@ class TTS:
return sr, audio
def speed_change(input_audio: np.ndarray, speed: float, sr: int):
# 将 NumPy 数组转换为原始 PCM 流
raw_audio = input_audio.astype(np.int16).tobytes()

View File

@ -1,6 +1,7 @@
custom:
bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
load_base: true
device: cuda
is_half: true
t2s_weights_path: GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt

View File

@ -3,6 +3,7 @@ custom:
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
device: cpu
is_half: false
load_base: false
t2s_weights_path: GPT_weights/jackma-e10.ckpt
vits_weights_path: SoVITS_weights/jackma_e8_s192.pth
default:

View File

@ -3,6 +3,7 @@ custom:
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
device: cpu
is_half: false
load_base: false
t2s_weights_path: GPT_weights/liyunlong-e15.ckpt
vits_weights_path: SoVITS_weights/liyunlong_e8_s176.pth
default:

View File

@ -3,6 +3,7 @@ custom:
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
device: cpu
is_half: false
load_base: false
t2s_weights_path: GPT_weights/morgan-e15.ckpt
vits_weights_path: SoVITS_weights/morgan_e8_s120.pth
default:

View File

@ -3,6 +3,7 @@ custom:
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
device: cpu
is_half: false
load_base: false
t2s_weights_path: GPT_weights/stephenchow-e15.ckpt
vits_weights_path: SoVITS_weights/stephenchow_e8_s112.pth
default:

View File

@ -135,6 +135,9 @@ port = args.port
host = args.bind_addr
argv = sys.argv
default_tts_config = TTS_Config()
TTS.init_base_models(default_tts_config)
APP = FastAPI()