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