import gc
import glob
import logging
import os
import traceback
import cpuinfo
import numpy as np
import psutil
import torch
# from contants import config
from contants import config
import utils
from bert_vits2 import Bert_VITS2
from contants import ModelType
from gpt_sovits.gpt_sovits import GPT_SoVITS
from logger import logger
from manager.observer import Subject
from utils.data_utils import HParams, check_is_none
from vits import VITS
from vits.hubert_vits import HuBert_VITS
from vits.w2v2_vits import W2V2_VITS
class ModelManager(Subject):
def __init__(self, device=config.system.device):
self.device = device
self.logger = logger
self.models = {
# ModelType:{model_id: {"model_path": model_path, "config_path": config_path, "model": model,
# "n_speakers": n_speakers}},
# model_id 类型为 int
ModelType.VITS: {},
ModelType.HUBERT_VITS: {},
ModelType.W2V2_VITS: {},
ModelType.BERT_VITS2: {},
ModelType.GPT_SOVITS: {},
}
self.sid2model = {
# ModelType:[{"real_id": real_id, "model": model, "model_id": model_id, "n_speakers": n_speakers}]
ModelType.VITS: [],
ModelType.HUBERT_VITS: [],
ModelType.W2V2_VITS: [],
ModelType.BERT_VITS2: [],
ModelType.GPT_SOVITS: [],
}
self.voice_speakers = {
ModelType.VITS.value: [],
ModelType.HUBERT_VITS.value: [],
ModelType.W2V2_VITS.value: [],
ModelType.BERT_VITS2.value: [],
ModelType.GPT_SOVITS.value: [],
}
self.emotion_reference = None
self.hubert = None
self.dimensional_emotion_model = None
self.tts_front = None
self.bert_models = {}
self.model_handler = None
self.emotion_model = None
self.processor = None
# self.sid2model = []
# self.name_mapping_id = []
self.voice_objs_count = 0
self._observers = []
self.model_class_map = {
ModelType.VITS: VITS,
ModelType.HUBERT_VITS: HuBert_VITS,
ModelType.W2V2_VITS: W2V2_VITS,
ModelType.BERT_VITS2: Bert_VITS2,
ModelType.GPT_SOVITS: GPT_SoVITS,
}
self.available_tts_model = set()
def model_init(self):
if config.tts_config.auto_load:
models = self.scan_path()
else:
models = config.tts_config.asdict().get("models")
for model in models:
self.load_model(model_path=model.get("model_path"),
config_path=model.get("config_path"),
sovits_path=model.get("sovits_path"),
gpt_path=model.get("gpt_path"))
dimensional_emotion_model_path = os.path.join(config.abs_path, config.system.data_path,
config.model_config.dimensional_emotion_model)
if os.path.isfile(dimensional_emotion_model_path):
if self.dimensional_emotion_model is None:
self.dimensional_emotion_model = self.load_dimensional_emotion_model(dimensional_emotion_model_path)
self.log_device_info()
if self.vits_speakers_count != 0:
self.logger.info(f"[{ModelType.VITS.value}] {self.vits_speakers_count} speakers")
if self.hubert_speakers_count != 0:
self.logger.info(f"[{ModelType.HUBERT_VITS.value}] {self.hubert_speakers_count} speakers")
if self.w2v2_speakers_count != 0:
self.logger.info(f"[{ModelType.W2V2_VITS.value}] {self.w2v2_speakers_count} speakers")
if self.bert_vits2_speakers_count != 0:
self.logger.info(f"[{ModelType.BERT_VITS2.value}] {self.bert_vits2_speakers_count} speakers")
if self.gpt_sovits_speakers_count != 0:
self.logger.info(f"[{ModelType.GPT_SOVITS.value}] {self.gpt_sovits_speakers_count} speakers")
self.logger.info(f"{self.speakers_count} speakers in total.")
if self.speakers_count == 0:
self.logger.warning(f"No model was loaded.")
@property
def vits_speakers(self):
return self.voice_speakers[ModelType.VITS]
@property
def speakers_count(self):
return self.vits_speakers_count + self.hubert_speakers_count + self.w2v2_speakers_count + self.bert_vits2_speakers_count + self.gpt_sovits_speakers_count
@property
def vits_speakers_count(self):
return len(self.voice_speakers[ModelType.VITS.value])
@property
def hubert_speakers_count(self):
return len(self.voice_speakers[ModelType.HUBERT_VITS.value])
@property
def w2v2_speakers_count(self):
return len(self.voice_speakers[ModelType.W2V2_VITS.value])
@property
def w2v2_emotion_count(self):
return len(self.emotion_reference) if self.emotion_reference is not None else 0
@property
def bert_vits2_speakers_count(self):
return len(self.voice_speakers[ModelType.BERT_VITS2.value])
@property
def gpt_sovits_speakers_count(self):
return len(self.voice_speakers[ModelType.GPT_SOVITS.value])
# 添加观察者
def attach(self, observer):
self._observers.append(observer)
# 移除观察者
def detach(self, observer):
self._observers.remove(observer)
# 通知所有观察者
def notify(self, event_type, **kwargs):
for observer in self._observers:
observer.update(event_type, **kwargs)
def log_device_info(self):
cuda_available = torch.cuda.is_available()
self.logger.info(
f"PyTorch Version: {torch.__version__} Cuda available:{cuda_available} Device type:{self.device.type}")
if self.device.type == 'cuda':
if cuda_available:
device_name = torch.cuda.get_device_name(self.device.index)
gpu_memory_info = round(torch.cuda.get_device_properties(self.device).total_memory / 1024 ** 3) # GB
self.logger.info(
f"Using GPU on {device_name} {gpu_memory_info}GB, GPU Device Index: {self.device.index}")
else:
self.logger.warning("GPU device specified, but CUDA is not available.")
else:
cpu_info = cpuinfo.get_cpu_info()
cpu_name = cpu_info.get("brand_raw")
cpu_count = psutil.cpu_count(logical=False)
thread_count = psutil.cpu_count(logical=True)
memory_info = psutil.virtual_memory()
total_memory = round(memory_info.total / (1024 ** 3))
self.logger.info(
f"Using CPU on {cpu_name} with {cpu_count} cores and {thread_count} threads. Total memory: {total_memory}GB")
def relative_to_absolute_path(self, *paths):
absolute_paths = []
for path in paths:
if path is None:
return None
path = os.path.normpath(path)
if path.startswith('models'):
path = os.path.join(config.abs_path, config.system.data_path, path)
else:
path = os.path.join(config.abs_path, config.system.data_path, config.tts_config.models_path,
path)
absolute_paths.append(path)
return absolute_paths
def absolute_to_relative_path(self, *paths):
relative_paths = []
for path in paths:
if path is None:
relative_paths.append(None)
continue
# 获取models目录下的相对路径
relative_path = os.path.relpath(path, os.path.join(config.abs_path, config.system.data_path,
config.tts_config.models_path))
relative_paths.append(relative_path)
return relative_paths
def _load_model_from_path(self, model_path, config_path, sovits_path, gpt_path):
if check_is_none(sovits_path, gpt_path):
hps = utils.get_hparams_from_file(config_path)
model_type = self.recognition_model_type(hps)
else:
hps = None
model_type = ModelType.GPT_SOVITS
model_args = {
"model_type": model_type,
"model_path": model_path,
"config_path": config_path,
"sovits_path": sovits_path,
"gpt_path": gpt_path,
"config": hps,
"device": self.device
}
model_class = self.model_class_map[model_type]
model = model_class(**model_args)
if model_type == ModelType.VITS:
bert_embedding = getattr(hps.data, 'bert_embedding', getattr(hps.model, 'bert_embedding', False))
if bert_embedding and self.tts_front is None:
self.load_VITS_PinYin_model(
os.path.join(config.abs_path, config.system.data_path, config.model_config.vits_chinese_bert))
if not config.vits_config.dynamic_loading:
model.load_model()
self.available_tts_model.add(ModelType.VITS.value)
elif model_type == ModelType.W2V2_VITS:
if self.emotion_reference is None:
self.emotion_reference = self.load_npy(
os.path.join(config.abs_path, config.system.data_path, config.model_config.dimensional_emotion_npy))
model.load_model(emotion_reference=self.emotion_reference,
dimensional_emotion_model=self.dimensional_emotion_model)
self.available_tts_model.add(ModelType.W2V2_VITS.value)
elif model_type == ModelType.HUBERT_VITS:
if self.hubert is None:
self.hubert = self.load_hubert_model(
os.path.join(config.abs_path, config.system.data_path, config.model_config.hubert_soft_0d54a1f4))
model.load_model(hubert=self.hubert)
elif model_type == ModelType.BERT_VITS2:
bert_model_names = model.bert_model_names
for bert_model_name in bert_model_names.values():
if self.model_handler is None:
from manager.model_handler import ModelHandler
self.model_handler = ModelHandler(self.device)
self.model_handler.load_bert(bert_model_name)
if model.hps_ms.model.emotion_embedding == 1:
self.model_handler.load_emotion()
elif model.hps_ms.model.emotion_embedding == 2:
self.model_handler.load_clap()
model.load_model(self.model_handler)
self.available_tts_model.add(ModelType.BERT_VITS2.value)
elif model_type == ModelType.GPT_SOVITS:
if self.model_handler is None:
from manager.model_handler import ModelHandler
self.model_handler = ModelHandler(self.device)
self.model_handler.load_ssl()
self.model_handler.load_bert("CHINESE_ROBERTA_WWM_EXT_LARGE")
model.load_model(self.model_handler)
sid2model = []
speakers = []
new_id = len(self.voice_speakers[model_type.value])
model_id = max([-1] + list(self.models[model_type].keys())) + 1
for real_id, name in enumerate(model.speakers):
sid2model.append({"real_id": real_id, "model": model, "model_id": model_id})
speakers.append({"id": new_id, "name": name, "lang": model.lang})
new_id += 1
model_data = {
"model": model,
"model_type": model_type,
"model_id": model_id,
"model_path": model_path,
"config": hps,
"sovits_path": sovits_path,
"gpt_path": gpt_path,
"sid2model": sid2model,
"speakers": speakers
}
if model_type == ModelType.GPT_SOVITS:
logging.info(
f"model_type:{model_type.value} model_id:{model_id} sovits_path:{sovits_path} gpt_path:{gpt_path}")
else:
logging.info(
f"model_type:{model_type.value} model_id:{model_id} n_speakers:{len(speakers)} model_path:{model_path}")
return model_data
def load_model(self, model_path: str, config_path: str, sovits_path: str, gpt_path: str):
try:
if not check_is_none(model_path, config_path):
model_path, config_path = self.relative_to_absolute_path(model_path, config_path)
else:
sovits_path, gpt_path = self.relative_to_absolute_path(sovits_path, gpt_path)
model_data = self._load_model_from_path(model_path, config_path, sovits_path, gpt_path)
model_id = model_data["model_id"]
sid2model = model_data["sid2model"]
model_type = model_data["model_type"]
self.models[model_type][model_id] = {
"model_type": model_data.get("model_type"),
"model_path": model_path,
"config_path": config_path,
"sovits_path": sovits_path,
"gpt_path": gpt_path,
"model": model_data.get("model"),
"n_speakers": len(model_data["speakers"])}
self.sid2model[model_type].extend(sid2model)
self.voice_speakers[model_type.value].extend(model_data["speakers"])
self.notify("model_loaded", model_manager=self)
state = True
except Exception as e:
self.logger.info(f"Loading failed. {e}")
self.logger.error(traceback.format_exc())
state = False
return state
def unload_model(self, model_type_value: str, model_id: str):
state = False
model_type = ModelType(model_type_value)
model_id = int(model_id)
try:
if model_id in self.models[model_type].keys():
model_data = self.models[model_type][model_id]
model = model_data.get("model")
n_speakers = model_data.get("n_speakers")
start = 0
for key, value in self.models[model_type].items():
if key == model_id:
break
start += value.get("n_speakers")
if model_type == ModelType.BERT_VITS2:
for bert_model_name in model.bert_model_names.values():
self.model_handler.release_bert(bert_model_name)
if model.version == "2.1":
self.model_handler.release_emotion()
elif model.version in ["2.2", "extra", "2.4"]:
self.model_handler.release_clap()
elif model_type == ModelType.GPT_SOVITS:
self.model_handler.release_bert("CHINESE_ROBERTA_WWM_EXT_LARGE")
self.model_handler.release_ssl_model()
del self.sid2model[model_type][start:start + n_speakers]
del self.voice_speakers[model_type.value][start:start + n_speakers]
del self.models[model_type][model_id]
for new_id, speaker in enumerate(self.voice_speakers[model_type.value]):
speaker["id"] = new_id
gc.collect()
torch.cuda.empty_cache()
state = True
self.notify("model_unloaded", model_manager=self)
self.logger.info(f"Unloading success.")
except Exception as e:
logging.error(traceback.print_exc())
logging.error(f"Unloading failed. {e}")
state = False
return state
def load_dimensional_emotion_model(self, model_path):
try:
import audonnx
root = os.path.dirname(model_path)
model_file = model_path
dimensional_emotion_model = audonnx.load(root=root, model_file=model_file)
self.notify("model_loaded", model_manager=self)
except Exception as e:
self.logger.warning(f"Load DIMENSIONAL_EMOTION_MODEL failed {e}")
return dimensional_emotion_model
def unload_dimensional_emotion_model(self):
self.dimensional_emotion_model = None
self.notify("model_unloaded", model_manager=self)
def load_hubert_model(self, model_path):
""""HuBERT-VITS"""
try:
from vits.hubert_model import hubert_soft
hubert = hubert_soft(model_path)
except Exception as e:
self.logger.warning(f"Load HUBERT_SOFT_MODEL failed {e}")
return hubert
def unload_hubert_model(self):
self.hubert = None
self.notify("model_unloaded", model_manager=self)
def load_VITS_PinYin_model(self, bert_path):
""""vits_chinese"""
from vits.text.vits_pinyin import VITS_PinYin
if self.tts_front is None:
self.tts_front = VITS_PinYin(bert_path, self.device)
def reorder_model(self, old_index, new_index):
"""重新排序模型，将old_index位置的模型移动到new_index位置"""
if 0 <= old_index < len(self.models) and 0 <= new_index < len(self.models):
model = self.models[old_index]
del self.models[old_index]
self.models.insert(new_index, model)
def get_models_path(self):
"""按返回模型路径列表，列表每一项为{"model_path": model_path, "config_path": config_path}"""
info = []
for models in self.models.values():
for model in models.values():
info.append({
"model_type": model.get("model_type"),
"model_path": model.get("model_path"),
"config_path": model.get("config_path"),
"sovits_path": model.get("sovits_path"),
"gpt_path": model.get("gpt_path"),
})
return info
def get_models_path_by_type(self):
"""按模型类型返回模型路径"""
info = {
ModelType.VITS.value: [],
ModelType.HUBERT_VITS.value: [],
ModelType.W2V2_VITS.value: [],
ModelType.BERT_VITS2.value: [],
ModelType.GPT_SOVITS.value: [],
}
for model_type, models in self.models.items():
for values in models.values():
info[model_type].append(values[0])
return info
def get_models_info(self):
"""按模型类型返回模型文件夹名以及模型文件名，speakers数量"""
info = {
ModelType.VITS.value: [],
ModelType.HUBERT_VITS.value: [],
ModelType.W2V2_VITS.value: [],
ModelType.BERT_VITS2.value: [],
ModelType.GPT_SOVITS.value: [],
}
for model_type, model_data in self.models.items():
if model_type != ModelType.GPT_SOVITS:
for model_id, model in model_data.items():
model_path = model.get("model_path")
config_path = model.get("config_path")
model_path = self.absolute_to_relative_path(model_path)[0].replace("\\", "/")
config_path = self.absolute_to_relative_path(config_path)[0].replace("\\", "/")
info[model_type.value].append(
{"model_id": model_id,
"model_path": model_path,
"config_path": config_path,
"n_speakers": model.get("n_speakers")})
else:
for model_id, model in model_data.items():
sovits_path = model.get("sovits_path")
gpt_path = model.get("gpt_path")
sovits_path = self.absolute_to_relative_path(sovits_path)[0].replace("\\", "/")
gpt_path = self.absolute_to_relative_path(gpt_path)[0].replace("\\", "/")
info[model_type.value].append(
{"model_id": model_id,
"sovits_path": sovits_path,
"gpt_path": gpt_path,
"n_speakers": model.get("n_speakers")})
return info
def get_model_by_index(self, model_type, model_id):
"""根据给定的索引返回模型"""
if 0 <= model_id < len(self.models):
_, model, _ = self.models[model_type][model_id]
return model
return None
# def get_bert_model(self, bert_model_name):
# if bert_model_name not in self.bert_models:
# raise ValueError(f"Model {bert_model_name} not loaded!")
# return self.bert_models[bert_model_name]
def clear_all(self):
"""清除所有模型"""
self.models.clear()
def recognition_model_type(self, hps: HParams) -> str:
# model_config = json.load(model_config_json)
symbols = getattr(hps, "symbols", None)
# symbols = model_config.get("symbols", None)
emotion_embedding = getattr(hps.data, "emotion_embedding", False)
if "use_spk_conditioned_encoder" in hps.model:
model_type = ModelType.BERT_VITS2
return model_type
if symbols != None:
if not emotion_embedding:
mode_type = ModelType.VITS
else:
mode_type = ModelType.W2V2_VITS
else:
mode_type = ModelType.HUBERT_VITS
return mode_type
def _load_npy_from_path(self, path):
model_extention = os.path.splitext(path)[1]
if model_extention != ".npy":
raise ValueError(f"Unsupported model type: {model_extention}")
return np.load(path).reshape(-1, 1024)
def load_npy(self, emotion_reference_npy):
emotion_reference = np.empty((0, 1024))
if isinstance(emotion_reference_npy, list):
for i in emotion_reference_npy:
emotion_reference = np.append(emotion_reference, self._load_npy_from_path(i), axis=0)
elif os.path.isdir(emotion_reference_npy):
for root, dirs, files in os.walk(emotion_reference_npy):
for file_name in files:
if file_name.endswith(".npy"):
file_path = os.path.join(root, file_name)
emotion_reference = np.append(emotion_reference, self._load_npy_from_path(file_path),
axis=0)
elif os.path.isfile(emotion_reference_npy):
emotion_reference = self._load_npy_from_path(emotion_reference_npy)
logging.info(f"Loaded emotional dimention npy range: {len(emotion_reference)}")
return emotion_reference
def scan_path(self):
folder_path = os.path.join(config.abs_path, config.system.data_path, config.tts_config.models_path)
model_paths = glob.glob(folder_path + "/**/*.pth", recursive=True)
all_paths = []
for id, pth_path in enumerate(model_paths):
pth_name = os.path.basename(pth_path)
if pth_name.startswith(("D_", "DUR_")):
continue
dir_name = os.path.dirname(pth_path)
config_paths = glob.glob(dir_name + "/*.json", recursive=True)
gpt_paths = glob.glob(dir_name + "/*.ckpt", recursive=True)
model_path, config_path, sovits_path, gpt_path, model_type = None, None, None, None, None
if len(config_paths) > 0:
model_path = pth_path
config_path = config_paths[0]
elif len(gpt_paths) > 0:
gpt_path = gpt_paths[0]
sovits_path = pth_path
model_type = ModelType.GPT_SOVITS
else:
continue
info = {
"model_id": id,
"model_type": model_type,
"model_path": model_path,
"config_path": config_path,
"sovits_path": sovits_path,
"gpt_path": gpt_path,
}
all_paths.append(info)
return all_paths
def scan_unload_path(self):
all_paths = self.scan_path()
unload_paths = []
loaded_paths = []
loaded_paths_2 = []
for model in self.get_models_path():
# 只取已加载的模型路径
if model.get("model_type") == ModelType.GPT_SOVITS:
sovits_path, gpt_path = self.absolute_to_relative_path(model.get("sovits_path"),
model.get("gpt_path"))
sovits_path, gpt_path = sovits_path.replace("\\", "/"), gpt_path.replace("\\", "/")
loaded_paths_2.append((sovits_path, gpt_path))
else:
model_path = self.absolute_to_relative_path(model.get("model_path"))[0].replace("\\", "/")
loaded_paths.append(model_path)
for info in all_paths:
# 将绝对路径修改为相对路径，并将分隔符格式化为'/'
if info.get("model_type") == ModelType.GPT_SOVITS:
sovits_path, gpt_path = self.absolute_to_relative_path(info.get("sovits_path"),
info.get("gpt_path"))
sovits_path, gpt_path = sovits_path.replace("\\", "/"), gpt_path.replace("\\", "/")
if not self.is_path_loaded((sovits_path, gpt_path), loaded_paths_2):
info.update(
{"model_type": info.get("model_type").value, "sovits_path": sovits_path, "gpt_path": gpt_path})
unload_paths.append(info)
else:
model_path, config_path = self.absolute_to_relative_path(info.get("model_path"),
info.get("config_path"))
model_path, config_path = model_path.replace("\\", "/"), config_path.replace("\\", "/")
if not self.is_path_loaded(model_path, loaded_paths):
info.update({"model_path": model_path, "config_path": config_path})
unload_paths.append(info)
return unload_paths
def is_path_loaded(self, paths, loaded_paths):
if len(paths) == 2:
sovits_path, gpt_path = paths
for loaded_path in loaded_paths:
if sovits_path == loaded_path[0] and gpt_path == loaded_path[1]:
return True
else:
path = paths
for loaded_path in loaded_paths:
if path == loaded_path:
return True
return False