manager/ModelManager.py

import gc
import glob
import logging
import os
import traceback

import cpuinfo
import numpy as np
import psutil
import torch

from config import config, BASE_DIR
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="cpu"):
        self.device = torch.device(device)
        self.logger = logger

        self.tts_models = {model_type: {} for model_type in ModelType}
        self.sid2model = {model_type: [] for model_type in ModelType}
        self.spk2model = {model_type: {} for model_type in ModelType}
        self.voice_speakers = {model_type: [] for model_type in ModelType}

        """
        self.tts_models: {
            ModelType: {model_id: {"vits_path": vits_path, "config_path": config_path, "model": model, "n_speakers": n_speakers}},
            ...
        }
        model_id 类型为 int

        self.sid2model: {
            ModelType: [{"real_id": real_id, "model": model, "model_id": model_id, "n_speakers": n_speakers}],
            ...
        }
        
        self.sid2model: {
            ModelType: {
                speaker: {"real_id": real_id, "model": model, "model_id": model_id},
                ...
            },
            ...
        }

        self.voice_speakers: {
            ModelType: [],
            ...
        }
        """

        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_model_config.auto_load:
            tts_models = self.scan_path()
            for tts_model in tts_models:
                self.load_model(tts_model)
        else:
            tts_models = config.tts_model_config.tts_models
            for tts_model in tts_models:
                self.load_model(tts_model.model_dump())

        dimensional_emotion_vits_path = os.path.join(BASE_DIR, config.system.data_path,
                                                     config.resource_paths_config.dimensional_emotion_model)
        if os.path.isfile(dimensional_emotion_vits_path):
            if self.dimensional_emotion_model is None:
                self.dimensional_emotion_model = self.load_dimensional_emotion_model(dimensional_emotion_vits_path)

        self.log_device_info()

        if self.vits_speakers_count != 0:
            self.logger.info(f"[{ModelType.VITS}] {self.vits_speakers_count} speakers")
        if self.hubert_speakers_count != 0:
            self.logger.info(f"[{ModelType.HUBERT_VITS}] {self.hubert_speakers_count} speakers")
        if self.w2v2_speakers_count != 0:
            self.logger.info(f"[{ModelType.W2V2_VITS}] {self.w2v2_speakers_count} speakers")
        if self.bert_vits2_speakers_count != 0:
            self.logger.info(f"[{ModelType.BERT_VITS2}] {self.bert_vits2_speakers_count} speakers")
        if self.gpt_sovits_speakers_count != 0:
            self.logger.info(f"[{ModelType.GPT_SOVITS}] {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])

    @property
    def hubert_speakers_count(self):
        return len(self.voice_speakers[ModelType.HUBERT_VITS])

    @property
    def w2v2_speakers_count(self):
        return len(self.voice_speakers[ModelType.W2V2_VITS])

    @property
    def w2v2_emotion_count(self):
        return len(self.emotion_reference) if self.emotion_reference is not None else 0

    @property
    def bert_vits2_spk2model(self):
        return self.spk2model[ModelType.BERT_VITS2]

    @property
    def bert_vits2_speakers_count(self):
        return len(self.voice_speakers[ModelType.BERT_VITS2])

    @property
    def gpt_sovits_speakers_count(self):
        return len(self.voice_speakers[ModelType.GPT_SOVITS])

    # 添加观察者
    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(BASE_DIR, config.system.data_path, path)
            else:
                path = os.path.join(BASE_DIR, config.system.data_path, config.tts_model_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(BASE_DIR, config.system.data_path,
                                                               config.tts_model_config.models_dir))

            relative_paths.append(relative_path)

        return relative_paths

    def _load_model_from_path(self, tts_model):
        if tts_model["model_type"] == ModelType.GPT_SOVITS:
            hps = None
            model_type = ModelType.GPT_SOVITS
            vits_path = tts_model["vits_path"]
            t2s_path = tts_model["t2s_path"]

            model_args = {
                "model_type": model_type,
                "vits_path": vits_path,
                "t2s_path": t2s_path,
                "config": hps,
                "device": self.device
            }

        else:
            hps = utils.get_hparams_from_file(tts_model["config_path"])
            model_type = self.recognition_model_type(hps)
            vits_path = tts_model["vits_path"]
            config_path = tts_model["config_path"]

            model_args = {
                "model_type": model_type,
                "vits_path": vits_path,
                "config_path": config_path,
                "config": hps,
                "device": self.device
            }
            if model_type == ModelType.VITS:
                model_args["dynamic_loading"] = tts_model.get("dynamic_loading", False)

        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(BASE_DIR, config.system.data_path, config.resource_paths_config.vits_chinese_bert))
            if not model.dynamic_loading:
                model.load_model()
            self.available_tts_model.add(ModelType.VITS)

        elif model_type == ModelType.W2V2_VITS:
            if self.emotion_reference is None:
                self.emotion_reference = self.load_npy(
                    os.path.join(BASE_DIR, config.system.data_path,
                                 config.resource_paths_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)

        elif model_type == ModelType.HUBERT_VITS:
            if self.hubert is None:
                self.hubert = self.load_hubert_model(
                    os.path.join(BASE_DIR, config.system.data_path, config.resource_paths_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)

        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)

            self.available_tts_model.add(ModelType.GPT_SOVITS)

        sid2model = []
        spk2model = {}
        speakers = []
        new_id = len(self.voice_speakers[model_type])
        model_id = max([-1] + list(self.tts_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})
            spk2model.update({name: {"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,
            "sid2model": sid2model,
            "spk2model": spk2model,
            "speakers": speakers
        }
        model_data.update(model_args)

        if model_type == ModelType.GPT_SOVITS:
            logging.info(
                f"model_type:{model_type} model_id:{model_id} vits_path:{vits_path} t2s_path:{t2s_path}")
        else:
            logging.info(
                f"model_type:{model_type} model_id:{model_id} n_speakers:{len(speakers)} vits_path:{vits_path}")

        return model_data

    def load_model(self, tts_model):
        try:
            model_data = self._load_model_from_path(tts_model)
            model_id = model_data["model_id"]
            sid2model = model_data["sid2model"]
            spk2model = model_data["spk2model"]
            model_type = model_data["model_type"]

            self.tts_models[model_type][model_id] = {
                "tts_model": tts_model,
                "model": model_data.get("model"),
                "n_speakers": len(model_data["speakers"]),
                "speakers": model_data["speakers"],
            }
            self.sid2model[model_type].extend(sid2model)
            self.spk2model[model_type].update(spk2model)
            self.voice_speakers[model_type].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: str, model_id: str):
        state = False
        model_id = int(model_id)
        try:
            if model_id in self.tts_models[model_type].keys():
                model_data = self.tts_models[model_type][model_id]
                model = model_data.get("model")
                n_speakers = model_data.get("n_speakers")
                speakers = model_data.get("speakers")

                start = 0

                for key, value in self.tts_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][start:start + n_speakers]
                del self.tts_models[model_type][model_id]

                for item in speakers:
                    self.spk2model[model_type].pop(item["name"], None)

                for new_id, speaker in enumerate(self.voice_speakers[model_type]):
                    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, vits_path):
        try:
            import audonnx
            root = os.path.dirname(vits_path)
            model_file = vits_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, vits_path):
        """"HuBERT-VITS"""
        try:
            from vits.hubert_model import hubert_soft
            hubert = hubert_soft(vits_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.tts_models) and 0 <= new_index < len(self.tts_models):
            model = self.tts_models[old_index]
            del self.tts_models[old_index]
            self.tts_models.insert(new_index, model)

    def get_models_path(self):
        """按返回模型路径列表，列表每一项为{"vits_path": vits_path, "config_path": config_path}"""
        info = []
        for models in self.tts_models.values():
            for model in models.values():
                info.append(model)

        return info

    def get_models_path_by_type(self):
        """按模型类型返回模型路径"""
        info = {
            ModelType.VITS: [],
            ModelType.HUBERT_VITS: [],
            ModelType.W2V2_VITS: [],
            ModelType.BERT_VITS2: [],
            ModelType.GPT_SOVITS: [],
        }
        for model_type, models in self.tts_models.items():
            for values in models.values():
                info[model_type].append(values[0])

        return info

    def get_models_info(self):
        """按模型类型返回模型文件夹名以及模型文件名，speakers数量"""
        info = {
            ModelType.VITS: [],
            ModelType.HUBERT_VITS: [],
            ModelType.W2V2_VITS: [],
            ModelType.BERT_VITS2: [],
            ModelType.GPT_SOVITS: [],
        }
        for model_type, model_data in self.tts_models.items():
            if model_type != ModelType.GPT_SOVITS:
                for model_id, model in model_data.items():
                    tts_model = model["tts_model"]

                    vits_path = tts_model["vits_path"]
                    config_path = tts_model["config_path"]

                    vits_path = self.absolute_to_relative_path(vits_path)[0].replace("\\", "/")
                    config_path = self.absolute_to_relative_path(config_path)[0].replace("\\", "/")

                    info[model_type].append(
                        {
                            "model_type": tts_model["model_type"],
                            "model_id": model_id,
                            "vits_path": vits_path,
                            "config_path": config_path,
                            "n_speakers": model["n_speakers"],
                        }
                    )
            else:
                for model_id, model in model_data.items():
                    tts_model = model["tts_model"]

                    vits_path = tts_model["vits_path"]
                    t2s_path = tts_model["t2s_path"]

                    vits_path = self.absolute_to_relative_path(vits_path)[0].replace("\\", "/")
                    t2s_path = self.absolute_to_relative_path(t2s_path)[0].replace("\\", "/")

                    info[model_type].append(
                        {
                            "model_type": tts_model["model_type"],
                            "model_id": model_id,
                            "vits_path": vits_path,
                            "t2s_path": t2s_path,
                            "n_speakers": model["n_speakers"],
                        }
                    )

        return info

    def get_model_by_index(self, model_type, model_id):
        """根据给定的索引返回模型"""
        if 0 <= model_id < len(self.tts_models):
            _, model, _ = self.tts_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.tts_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):
        models_dir = os.path.join(BASE_DIR, config.system.data_path, config.tts_model_config.models_dir)
        vits_paths = glob.glob(models_dir + "/**/*.pth", recursive=True)
        all_paths = []

        for id, pth_path in enumerate(vits_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)
            t2s_paths = glob.glob(dir_name + "/*.ckpt", recursive=True)

            if len(config_paths) > 0:
                vits_path = pth_path
                config_path = config_paths[0]
                hps = utils.get_hparams_from_file(config_path)
                model_type = self.recognition_model_type(hps)
                info = {
                    "vits_path": vits_path,
                    "config_path": config_path,
                }
            elif len(t2s_paths) > 0:
                t2s_path = t2s_paths[0]
                vits_path = pth_path
                model_type = ModelType.GPT_SOVITS
                info = {
                    "vits_path": vits_path,
                    "t2s_path": t2s_path,
                }
            else:
                continue
            info.update(
                {
                    "model_id": id,
                    "model_type": model_type,
                }
            )

            all_paths.append(info)

        return all_paths

    def scan_unload_path(self):
        all_paths = self.scan_path()
        unload_paths = []
        loaded_paths = self._get_loaded_paths()

        for info in all_paths:
            model_type = info["model_type"]
            if model_type == ModelType.GPT_SOVITS:
                vits_path, t2s_path = self._format_paths(
                    self.absolute_to_relative_path(
                        info["vits_path"],
                        info["t2s_path"]
                    )
                )
                if not self.is_path_loaded((vits_path, t2s_path), loaded_paths["GPT_SOVITS"]):
                    info.update(
                        {
                            "model_type": model_type,
                            "vits_path": vits_path,
                            "t2s_path": t2s_path
                        }
                    )
                    unload_paths.append(info)
            else:
                vits_path, config_path = self._format_paths(
                    self.absolute_to_relative_path(
                        info["vits_path"],
                        info["config_path"]
                    )
                )
                if not self.is_path_loaded(vits_path, loaded_paths["OTHER"]):
                    info.update(
                        {
                            "model_type": model_type,
                            "vits_path": vits_path,
                            "config_path": config_path
                        }
                    )
                    unload_paths.append(info)

        return unload_paths

    def _format_paths(self, paths):
        return tuple(path.replace("\\", "/") for path in paths)

    def _get_loaded_paths(self):
        loaded_paths = {
            "GPT_SOVITS": [],
            "OTHER": []
        }

        for model in self.get_models_path():
            tts_model = model["tts_model"]
            model_type = tts_model["model_type"]

            if model_type == ModelType.GPT_SOVITS:
                vits_path, t2s_path = self._format_paths(
                    self.absolute_to_relative_path(tts_model["vits_path"], tts_model["t2s_path"])
                )
                loaded_paths["GPT_SOVITS"].append((vits_path, t2s_path))
            else:
                vits_path = self._format_paths(self.absolute_to_relative_path(tts_model["vits_path"])[0])
                loaded_paths["OTHER"].append(vits_path)

        return loaded_paths

    def is_path_loaded(self, paths, loaded_paths):
        if isinstance(paths[0], tuple):
            return paths in loaded_paths
        return paths in loaded_paths