"""
放置公用模型
"""
import gc
import logging
import os
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM, BertTokenizer, MegatronBertModel
from contants import config
from utils.download import download_file
from bert_vits2.text.chinese_bert import get_bert_feature as zh_bert
from bert_vits2.text.english_bert_mock import get_bert_feature as en_bert
from bert_vits2.text.japanese_bert import get_bert_feature as ja_bert
from bert_vits2.text.japanese_bert_v111 import get_bert_feature as ja_bert_v111
from bert_vits2.text.japanese_bert_v200 import get_bert_feature as ja_bert_v200
from bert_vits2.text.english_bert_mock_v200 import get_bert_feature as en_bert_v200
from bert_vits2.text.chinese_bert_extra import get_bert_feature as zh_bert_extra
from bert_vits2.text.japanese_bert_extra import get_bert_feature as ja_bert_extra
class ModelHandler:
def __init__(self, device=config.system.device):
self.DOWNLOAD_PATHS = {
"CHINESE_ROBERTA_WWM_EXT_LARGE": [
"https://huggingface.co/hfl/chinese-roberta-wwm-ext-large/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/hfl/chinese-roberta-wwm-ext-large/resolve/main/pytorch_model.bin",
],
"BERT_BASE_JAPANESE_V3": [
"https://huggingface.co/cl-tohoku/bert-base-japanese-v3/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/cl-tohoku/bert-base-japanese-v3/resolve/main/pytorch_model.bin",
],
"BERT_LARGE_JAPANESE_V2": [
"https://huggingface.co/cl-tohoku/bert-large-japanese-v2/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/cl-tohoku/bert-large-japanese-v2/resolve/main/pytorch_model.bin",
],
"DEBERTA_V2_LARGE_JAPANESE": [
"https://huggingface.co/ku-nlp/deberta-v2-large-japanese/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/ku-nlp/deberta-v2-large-japanese/resolve/main/pytorch_model.bin",
],
"DEBERTA_V3_LARGE": [
"https://huggingface.co/microsoft/deberta-v3-large/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/microsoft/deberta-v3-large/resolve/main/pytorch_model.bin",
],
"SPM": [
"https://huggingface.co/microsoft/deberta-v3-large/resolve/main/spm.model",
"https://hf-mirror.com/microsoft/deberta-v3-large/resolve/main/spm.model",
],
"DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM": [
"https://huggingface.co/ku-nlp/deberta-v2-large-japanese-char-wwm/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/ku-nlp/deberta-v2-large-japanese-char-wwm/resolve/main/pytorch_model.bin",
],
"WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM": [
"https://huggingface.co/audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim/resolve/main/pytorch_model.bin",
],
"CLAP_HTSAT_FUSED": [
"https://huggingface.co/laion/clap-htsat-fused/resolve/main/pytorch_model.bin?download=true",
"https://hf-mirror.com/laion/clap-htsat-fused/resolve/main/pytorch_model.bin?download=true",
],
"Erlangshen_MegatronBert_1.3B_Chinese": [
"https://huggingface.co/IDEA-CCNL/Erlangshen-UniMC-MegatronBERT-1.3B-Chinese/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/IDEA-CCNL/Erlangshen-UniMC-MegatronBERT-1.3B-Chinese/resolve/main/pytorch_model.bin",
],
"G2PWModel": [
# "https://storage.googleapis.com/esun-ai/g2pW/G2PWModel-v2-onnx.zip",
"https://huggingface.co/ADT109119/G2PWModel-v2-onnx/resolve/main/g2pw.onnx",
"https://hf-mirror.com/ADT109119/G2PWModel-v2-onnx/resolve/main/g2pw.onnx",
],
"CHINESE_HUBERT_BASE": [
"https://huggingface.co/TencentGameMate/chinese-hubert-base/resolve/main/pytorch_model.bin",
"https://hf-mirror.com/TencentGameMate/chinese-hubert-base/resolve/main/pytorch_model.bin",
]
}
self.SHA256 = {
"CHINESE_ROBERTA_WWM_EXT_LARGE": "4ac62d49144d770c5ca9a5d1d3039c4995665a080febe63198189857c6bd11cd",
"BERT_BASE_JAPANESE_V3": "e172862e0674054d65e0ba40d67df2a4687982f589db44aa27091c386e5450a4",
"BERT_LARGE_JAPANESE_V2": "50212d714f79af45d3e47205faa356d0e5030e1c9a37138eadda544180f9e7c9",
"DEBERTA_V2_LARGE_JAPANESE": "a6c15feac0dea77ab8835c70e1befa4cf4c2137862c6fb2443b1553f70840047",
"DEBERTA_V3_LARGE": "dd5b5d93e2db101aaf281df0ea1216c07ad73620ff59c5b42dccac4bf2eef5b5",
"SPM": "c679fbf93643d19aab7ee10c0b99e460bdbc02fedf34b92b05af343b4af586fd",
"DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM": "bf0dab8ad87bd7c22e85ec71e04f2240804fda6d33196157d6b5923af6ea1201",
"WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM": "176d9d1ce29a8bddbab44068b9c1c194c51624c7f1812905e01355da58b18816",
"CLAP_HTSAT_FUSED": "1ed5d0215d887551ddd0a49ce7311b21429ebdf1e6a129d4e68f743357225253",
"Erlangshen_MegatronBert_1.3B_Chinese": "3456bb8f2c7157985688a4cb5cecdb9e229cb1dcf785b01545c611462ffe3579",
# "G2PWModel": "bb40c8c7b5baa755b2acd317c6bc5a65e4af7b80c40a569247fbd76989299999",
"G2PWModel": "",
"CHINESE_HUBERT_BASE": "2fefccd26c2794a583b80f6f7210c721873cb7ebae2c1cde3baf9b27855e24d8",
}
self.model_path = {
"CHINESE_ROBERTA_WWM_EXT_LARGE": os.path.join(config.abs_path, config.system.data_path,
config.model_config.chinese_roberta_wwm_ext_large),
"BERT_BASE_JAPANESE_V3": os.path.join(config.abs_path, config.system.data_path,
config.model_config.bert_base_japanese_v3),
"BERT_LARGE_JAPANESE_V2": os.path.join(config.abs_path, config.system.data_path,
config.model_config.bert_large_japanese_v2),
"DEBERTA_V2_LARGE_JAPANESE": os.path.join(config.abs_path, config.system.data_path,
config.model_config.deberta_v2_large_japanese),
"DEBERTA_V3_LARGE": os.path.join(config.abs_path, config.system.data_path,
config.model_config.deberta_v3_large),
"DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM": os.path.join(config.abs_path, config.system.data_path,
config.model_config.deberta_v2_large_japanese_char_wwm),
"WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM": os.path.join(config.abs_path, config.system.data_path,
config.model_config.wav2vec2_large_robust_12_ft_emotion_msp_dim),
"CLAP_HTSAT_FUSED": os.path.join(config.abs_path, config.system.data_path,
config.model_config.clap_htsat_fused),
"Erlangshen_MegatronBert_1.3B_Chinese": os.path.join(config.abs_path, config.system.data_path,
config.model_config.erlangshen_MegatronBert_1_3B_Chinese),
"G2PWModel": os.path.join(config.abs_path, config.system.data_path, config.model_config.g2pw_model),
"CHINESE_HUBERT_BASE": os.path.join(config.abs_path, config.system.data_path,
config.model_config.chinese_hubert_base),
}
self.lang_bert_func_map = {"zh": zh_bert, "en": en_bert, "ja": ja_bert, "ja_v111": ja_bert_v111,
"ja_v200": ja_bert_v200, "en_v200": en_bert_v200, "zh_extra": zh_bert_extra,
"ja_extra": ja_bert_extra}
self.bert_models = {} # Value: (tokenizer, model, reference_count)
self.emotion = None
self.clap = None
self.pinyinPlus = None
self.device = device
self.ssl_model = None
if config.bert_vits2_config.torch_data_type.lower() in ["float16", "fp16"]:
self.torch_dtype = torch.float16
else:
self.torch_dtype = None
@property
def emotion_model(self):
return self.emotion["model"]
@property
def emotion_processor(self):
return self.emotion["processor"]
@property
def clap_model(self):
return self.clap["model"]
@property
def clap_processor(self):
return self.clap["processor"]
def _download_model(self, model_name, target_path=None):
urls = self.DOWNLOAD_PATHS[model_name]
if target_path is None:
target_path = os.path.join(self.model_path[model_name], "pytorch_model.bin")
expected_sha256 = self.SHA256[model_name]
success, message = download_file(urls, target_path, expected_sha256=expected_sha256)
if not success:
logging.error(f"Failed to download {model_name}: {message}")
else:
logging.info(f"{message}")
def load_bert(self, bert_model_name, max_retries=3):
if bert_model_name not in self.bert_models:
retries = 0
model_path = ""
while retries < max_retries:
model_path = self.model_path[bert_model_name]
logging.info(f"Loading BERT model: {model_path}")
try:
if bert_model_name == "Erlangshen_MegatronBert_1.3B_Chinese":
tokenizer = BertTokenizer.from_pretrained(model_path, torch_dtype=self.torch_dtype)
model = MegatronBertModel.from_pretrained(model_path, torch_dtype=self.torch_dtype).to(
self.device)
else:
tokenizer = AutoTokenizer.from_pretrained(model_path, torch_dtype=self.torch_dtype)
model = AutoModelForMaskedLM.from_pretrained(model_path, torch_dtype=self.torch_dtype).to(
self.device)
self.bert_models[bert_model_name] = (tokenizer, model, 1) # 初始化引用计数为1
logging.info(f"Success loading: {model_path}")
break
except Exception as e:
logging.error(f"Failed loading {model_path}. {e}")
logging.info(f"Trying to download.")
if bert_model_name == "DEBERTA_V3_LARGE" and not os.path.exists(
os.path.join(model_path, "spm.model")):
self._download_model("SPM", os.path.join(model_path, "spm.model"))
self._download_model(bert_model_name)
retries += 1
if retries == max_retries:
logging.error(f"Failed to load {model_path} after {max_retries} retries.")
else:
tokenizer, model, count = self.bert_models[bert_model_name]
self.bert_models[bert_model_name] = (tokenizer, model, count + 1)
def load_emotion(self, max_retries=3):
"""Bert-VITS2 v2.1 EmotionModel"""
if self.emotion is None:
from transformers import Wav2Vec2Processor
from bert_vits2.get_emo import EmotionModel
retries = 0
model_path = self.model_path["WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM"]
while retries < max_retries:
logging.info(f"Loading WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM: {model_path}")
try:
self.emotion = {}
self.emotion["model"] = EmotionModel.from_pretrained(model_path).to(self.device)
self.emotion["processor"] = Wav2Vec2Processor.from_pretrained(model_path)
self.emotion["reference_count"] = 1
logging.info(f"Success loading: {model_path}")
break
except Exception as e:
logging.error(f"Failed loading {model_path}. {e}")
self._download_model("WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM")
retries += 1
if retries == max_retries:
logging.error(f"Failed to load {model_path} after {max_retries} retries.")
else:
self.emotion["reference_count"] += 1
def release_emotion(self):
if self.emotion is not None:
self.emotion["reference_count"] -= 1
if self.emotion["reference_count"] <= 0:
del self.emotion
self.emotion = None
gc.collect()
torch.cuda.empty_cache()
logging.info(f"Emotion model has been released.")
def load_clap(self, max_retries=3):
"""Bert-VITS2 v2.2 ClapModel"""
if self.clap is None:
from transformers import ClapModel, ClapProcessor
retries = 0
model_path = self.model_path["CLAP_HTSAT_FUSED"]
while retries < max_retries:
logging.info(f"Loading CLAP_HTSAT_FUSED: {model_path}")
try:
self.clap = {}
self.clap["model"] = ClapModel.from_pretrained(model_path, torch_dtype=self.torch_dtype).to(
self.device)
self.clap["processor"] = ClapProcessor.from_pretrained(model_path, torch_dtype=self.torch_dtype)
self.clap["reference_count"] = 1
logging.info(f"Success loading: {model_path}")
break
except Exception as e:
logging.error(f"Failed loading {model_path}. {e}")
self._download_model("CLAP_HTSAT_FUSED")
retries += 1
if retries == max_retries:
logging.error(f"Failed to load {model_path} after {max_retries} retries.")
else:
self.clap["reference_count"] += 1
def release_clap(self):
if self.clap is not None:
self.clap["reference_count"] -= 1
if self.clap["reference_count"] <= 0:
del self.clap
self.clap = None
gc.collect()
torch.cuda.empty_cache()
logging.info(f"Clap model has been released.")
def get_bert_model(self, bert_model_name):
if bert_model_name not in self.bert_models:
self.load_bert(bert_model_name)
tokenizer, model, _ = self.bert_models[bert_model_name]
return tokenizer, model
def get_bert_feature(self, norm_text, word2ph, language, bert_model_name, style_text=None, style_weight=0.7):
tokenizer, model = self.get_bert_model(bert_model_name)
bert_feature = self.lang_bert_func_map[language](norm_text, word2ph, tokenizer, model, self.device,
style_text=style_text, style_weight=style_weight)
return bert_feature
def get_pinyinPlus(self):
if self.pinyinPlus is None:
from bert_vits2.g2pW.pypinyin_G2pW_bv2 import G2PWPinyin
logging.info(f"Loading G2PWModel: {self.model_path['G2PWModel']}")
self.pinyinPlus = G2PWPinyin(
model_dir=self.model_path["G2PWModel"],
model_source=self.model_path["Erlangshen_MegatronBert_1.3B_Chinese"],
v_to_u=False,
neutral_tone_with_five=True,
)
logging.info("Success loading G2PWModel")
return self.pinyinPlus
def release_bert(self, bert_model_name):
if bert_model_name in self.bert_models:
_, _, count = self.bert_models[bert_model_name]
count -= 1
if count == 0:
# 当引用计数为0时，删除模型并释放其资源
del self.bert_models[bert_model_name]
gc.collect()
torch.cuda.empty_cache()
logging.info(f"BERT model {bert_model_name} has been released.")
else:
tokenizer, model = self.bert_models[bert_model_name][:2]
self.bert_models[bert_model_name] = (tokenizer, model, count)
def load_ssl(self, max_retries=3):
"""GPT-SoVITS"""
if self.ssl_model is None:
retries = 0
model_path = self.model_path["CHINESE_HUBERT_BASE"]
while retries < max_retries:
logging.info(f"Loading CHINESE_HUBERT_BASE: {model_path}")
try:
from gpt_sovits.feature_extractor.cnhubert import CNHubert
self.ssl_model = {}
model_path = self.model_path.get("CHINESE_HUBERT_BASE")
self.ssl_model["model"] = CNHubert(model_path)
self.ssl_model["model"].eval()
if config.gpt_sovits_config.is_half:
self.ssl_model["model"] = self.ssl_model["model"].half()
self.ssl_model["model"] = self.ssl_model["model"].to(self.device)
self.ssl_model["reference_count"] = 1
logging.info(f"Success loading: {model_path}")
break
except Exception as e:
logging.error(f"Failed loading {model_path}. {e}")
self._download_model("CHINESE_HUBERT_BASE")
retries += 1
if retries == max_retries:
logging.error(f"Failed to load {model_path} after {max_retries} retries.")
else:
self.ssl_model["reference_count"] += 1
def get_ssl_model(self):
if self.ssl_model is None:
self.load_ssl()
return self.ssl_model.get("model")
def release_ssl_model(self):
if self.ssl_model is not None:
self.ssl_model["reference_count"] -= 1
if self.ssl_model["reference_count"] <= 0:
del self.ssl_model
self.ssl_model = None
gc.collect()
torch.cuda.empty_cache()
logging.info(f"SSL model has been released.")
def is_model_loaded(self, bert_model_name):
return bert_model_name in self.bert_models
def reference_count(self, bert_model_name):
return self.bert_models[bert_model_name][2] if bert_model_name in self.bert_models else 0