manager/model_handler.py

"""
放置公用模型
"""

import gc
import logging
import os

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM, BertTokenizer, MegatronBertModel

from config import config, BASE_DIR
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="cpu"):
        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-MegatronBert-1.3B/resolve/main/pytorch_model.bin",
                "https://hf-mirror.com/IDEA-CCNL/Erlangshen-MegatronBert-1.3B/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": "d89257b2b59ac562f50227082127b077167c0ff8fb2379cdec0b0692c7b1dbdd",
            # "G2PWModel": "bb40c8c7b5baa755b2acd317c6bc5a65e4af7b80c40a569247fbd76989299999",
            "G2PWModel": "",
            "CHINESE_HUBERT_BASE": "2fefccd26c2794a583b80f6f7210c721873cb7ebae2c1cde3baf9b27855e24d8",
        }
        self.model_path = {
            "CHINESE_ROBERTA_WWM_EXT_LARGE": os.path.join(BASE_DIR, config.system.data_path,
                                                          config.resource_paths_config.chinese_roberta_wwm_ext_large),
            "BERT_BASE_JAPANESE_V3": os.path.join(BASE_DIR, config.system.data_path,
                                                  config.resource_paths_config.bert_base_japanese_v3),
            "BERT_LARGE_JAPANESE_V2": os.path.join(BASE_DIR, config.system.data_path,
                                                   config.resource_paths_config.bert_large_japanese_v2),
            "DEBERTA_V2_LARGE_JAPANESE": os.path.join(BASE_DIR, config.system.data_path,
                                                      config.resource_paths_config.deberta_v2_large_japanese),
            "DEBERTA_V3_LARGE": os.path.join(BASE_DIR, config.system.data_path,
                                             config.resource_paths_config.deberta_v3_large),
            "DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM": os.path.join(BASE_DIR, config.system.data_path,
                                                               config.resource_paths_config.deberta_v2_large_japanese_char_wwm),
            "WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM": os.path.join(BASE_DIR, config.system.data_path,
                                                                        config.resource_paths_config.wav2vec2_large_robust_12_ft_emotion_msp_dim),
            "CLAP_HTSAT_FUSED": os.path.join(BASE_DIR, config.system.data_path,
                                             config.resource_paths_config.clap_htsat_fused),
            "Erlangshen_MegatronBert_1.3B_Chinese": os.path.join(BASE_DIR, config.system.data_path,
                                                                 config.resource_paths_config.erlangshen_MegatronBert_1_3B_Chinese),
            "G2PWModel": os.path.join(BASE_DIR, config.system.data_path, config.resource_paths_config.g2pw_model),
            "CHINESE_HUBERT_BASE": os.path.join(BASE_DIR, config.system.data_path,
                                                config.resource_paths_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.pinyin_g2pw = None
        self.device = torch.device(device)
        self.ssl_model = None

        torch_data_type = config.bert_vits2_config.torch_data_type
        if torch_data_type and 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_pinyin_g2pw(self,model_source="Erlangshen_MegatronBert_1.3B_Chinese"):
        if self.pinyin_g2pw is None:
            from module.g2pW.g2pw import G2PWPinyin

            logging.info(f"Loading G2PWModel: {self.model_path['G2PWModel']}")
            self.pinyin_g2pw = G2PWPinyin(
                model_dir=self.model_path["G2PWModel"],
                model_source=self.model_path[model_source],
                v_to_u=False,
                neutral_tone_with_five=True,
            )
            logging.info("Success loading G2PWModel")

        return self.pinyin_g2pw

    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