genie_tts/ModelManager.py

"""
不再新建 .bin 文件。
修改后内存: 6448 MB
修改前内存: 5952 MB
"""

import gc
import logging
import os
from dataclasses import dataclass
from typing import Optional, List, Dict

import numpy as np
import onnx
import onnxruntime
from onnxruntime import InferenceSession
from tokenizers import Tokenizer

from .Core.Resources import (HUBERT_MODEL_DIR, SV_MODEL, ROBERTA_MODEL_DIR)
from .Utils.Utils import LRUCacheDict

onnxruntime.set_default_logger_severity(3)
logger = logging.getLogger(__name__)


class GSVModelFile:
    T2S_ENCODER_FP32: str = 't2s_encoder_fp32.onnx'

    T2S_FIRST_STAGE_DECODER_FP32: str = 't2s_first_stage_decoder_fp32.onnx'
    T2S_FIRST_STAGE_DECODER_FP16: str = 't2s_first_stage_decoder_fp16.onnx'
    T2S_STAGE_DECODER_FP32: str = 't2s_stage_decoder_fp32.onnx'
    T2S_STAGE_DECODER_FP16: str = 't2s_stage_decoder_fp16.onnx'
    T2S_DECODER_WEIGHT_FP16: str = 't2s_shared_fp16.bin'

    VITS_FP32: str = 'vits_fp32.onnx'
    VITS_WEIGHT_FP16: str = 'vits_fp16.bin'

    PROMPT_ENCODER: str = 'prompt_encoder_fp32.onnx'
    PROMPT_ENCODER_WEIGHT_FP16: str = 'prompt_encoder_fp16.bin'

    HUBERT_MODEL = os.path.join(HUBERT_MODEL_DIR, "chinese-hubert-base.onnx")
    HUBERT_MODEL_WEIGHT_FP16 = os.path.join(HUBERT_MODEL_DIR, "chinese-hubert-base_weights_fp16.bin")

    ROBERTA_MODEL = os.path.join(ROBERTA_MODEL_DIR, 'RoBERTa.onnx')
    ROBERTA_TOKENIZER = os.path.join(ROBERTA_MODEL_DIR, 'roberta_tokenizer')


@dataclass
class GSVModel:
    LANGUAGE: str
    T2S_ENCODER: InferenceSession
    T2S_FIRST_STAGE_DECODER: InferenceSession
    T2S_STAGE_DECODER: InferenceSession
    VITS: InferenceSession
    PROMPT_ENCODER: Optional[InferenceSession] = None
    PROMPT_ENCODER_PATH: Optional[str] = None


def load_session_with_fp16_conversion(
        onnx_path: str,
        fp16_bin_path: str,
        providers: List[str],
        sess_options: Optional[onnxruntime.SessionOptions] = None
) -> InferenceSession:
    """
    通用函数：读取 ONNX 和 FP16 权重文件，在内存中将权重转换为 FP32，
    注入到 ONNX 模型中并加载 InferenceSession，不产生临时文件。
    """
    if not os.path.exists(onnx_path):
        raise FileNotFoundError(f"ONNX Model not found: {onnx_path}")
    if not os.path.exists(fp16_bin_path):
        raise FileNotFoundError(f"FP16 Weight file not found: {fp16_bin_path}")

    model_proto = onnx.load(onnx_path, load_external_data=False)
    fp16_data = np.fromfile(fp16_bin_path, dtype=np.float16)
    fp32_data = fp16_data.astype(np.float32)
    fp32_bytes = fp32_data.tobytes()

    # 遍历并修补模型中的 External Data Initializers
    for tensor in model_proto.graph.initializer:
        # 检查该 Tensor 是否使用外部数据
        if tensor.data_location == onnx.TensorProto.EXTERNAL:
            offset = 0
            length = 0
            # 解析外部数据信息
            for entry in tensor.external_data:
                if entry.key == 'offset':
                    offset = int(entry.value)
                elif entry.key == 'length':
                    length = int(entry.value)

            if offset + length > len(fp32_bytes):
                logger.warning(
                    f"Tensor {tensor.name} requested a data range that exceeds the size of the provided bin file. "
                    f"Offset: {offset}, Length: {length}, Buffer: {len(fp32_bytes)}"
                )
                continue

            tensor_data = fp32_bytes[offset: offset + length]
            tensor.raw_data = tensor_data

            del tensor.external_data[:]
            tensor.data_location = onnx.TensorProto.DEFAULT

    try:
        session = InferenceSession(
            model_proto.SerializeToString(),
            providers=providers,
            sess_options=sess_options
        )
        return session
    except Exception as e:
        logger.error(f"Failed to load in-memory model {os.path.basename(onnx_path)}: {e}")
        raise e


class ModelManager:
    def __init__(self):
        capacity_str = os.getenv('Max_Cached_Character_Models', '3')
        self.character_to_model: Dict[str, Dict[str, Optional[InferenceSession]]] = LRUCacheDict(
            capacity=int(capacity_str)
        )
        self.character_to_language: Dict[str, str] = {}
        self.character_model_paths: Dict[str, str] = {}
        self.providers = ["CPUExecutionProvider"]

        self.cn_hubert: Optional[InferenceSession] = None
        self.speaker_verification_model: Optional[InferenceSession] = None
        self.roberta_model: Optional[InferenceSession] = None
        self.roberta_tokenizer: Optional[Tokenizer] = None

    def load_roberta_model(self, model_path: str = GSVModelFile.ROBERTA_MODEL) -> bool:
        if self.roberta_model is not None:
            return True
        if not os.path.exists(model_path):
            # logger.warning(f'RoBERTa model does not exist: {model_path}. BERT features will not be used.')
            return False
        try:
            self.roberta_model = onnxruntime.InferenceSession(
                model_path,
                providers=self.providers,
            )
            self.roberta_tokenizer = Tokenizer.from_file(
                os.path.join(GSVModelFile.ROBERTA_TOKENIZER, 'tokenizer.json')
            )
            logger.info(f"Successfully loaded RoBERTa model.")
            return True
        except Exception as e:
            logger.error(
                f"Error: Failed to load ONNX model '{GSVModelFile.ROBERTA_MODEL}'.\n"
                f"Details: {e}"
            )
        return False

    def load_sv_model(self, model_path: str = SV_MODEL) -> bool:
        if self.speaker_verification_model is not None:
            return True
        try:
            self.speaker_verification_model = onnxruntime.InferenceSession(
                model_path,
                providers=self.providers,
            )
            logger.info(f"Successfully loaded Speaker Verification model.")
            return True
        except Exception as e:
            logger.error(
                f"Error: Failed to load ONNX model '{SV_MODEL}'.\n"
                f"Details: {e}"
            )
        return False

    def load_cn_hubert(self, model_path: str = GSVModelFile.HUBERT_MODEL) -> bool:
        if self.cn_hubert is not None:
            return True
        try:
            # Hubert 也应用内存转换逻辑
            if model_path == GSVModelFile.HUBERT_MODEL and os.path.exists(GSVModelFile.HUBERT_MODEL_WEIGHT_FP16):
                self.cn_hubert = load_session_with_fp16_conversion(
                    model_path,
                    GSVModelFile.HUBERT_MODEL_WEIGHT_FP16,
                    self.providers
                )
            else:
                self.cn_hubert = onnxruntime.InferenceSession(
                    model_path,
                    providers=self.providers,
                )
            logger.info("Successfully loaded CN_HuBERT model.")
            return True
        except Exception as e:
            logger.error(
                f"Error: Failed to load ONNX model '{GSVModelFile.HUBERT_MODEL}'.\n"
                f"Details: {e}"
            )
        return False

    def get(self, character_name: str) -> Optional[GSVModel]:
        character_name = character_name.lower()
        language = self.character_to_language.get(character_name, 'Japanese')
        if character_name in self.character_to_model:
            model_map: dict = self.character_to_model[character_name]
            # 简化获取逻辑
            t2s_first_stage_decoder = model_map.get(GSVModelFile.T2S_FIRST_STAGE_DECODER_FP32) or \
                                      model_map.get(GSVModelFile.T2S_FIRST_STAGE_DECODER_FP16)
            t2s_stage_decoder = model_map.get(GSVModelFile.T2S_STAGE_DECODER_FP32) or \
                                model_map.get(GSVModelFile.T2S_STAGE_DECODER_FP16)
            prompt_encoder_path = os.path.join(self.character_model_paths[character_name], GSVModelFile.PROMPT_ENCODER)

            return GSVModel(
                LANGUAGE=language,
                T2S_ENCODER=model_map[GSVModelFile.T2S_ENCODER_FP32],
                T2S_FIRST_STAGE_DECODER=t2s_first_stage_decoder,
                T2S_STAGE_DECODER=t2s_stage_decoder,
                VITS=model_map[GSVModelFile.VITS_FP32],
                PROMPT_ENCODER=model_map[GSVModelFile.PROMPT_ENCODER],
                PROMPT_ENCODER_PATH=prompt_encoder_path,
            )
        if character_name in self.character_model_paths:
            model_dir = self.character_model_paths[character_name]
            if self.load_character(character_name, model_dir, language=language):
                return self.get(character_name)
            else:
                del self.character_model_paths[character_name]
                return None
        return None

    def has_character(self, character_name: str) -> bool:
        character_name = character_name.lower()
        return character_name in self.character_model_paths

    def load_character(
            self,
            character_name: str,
            model_dir: str,
            language: str,
    ) -> bool:
        """
        加载角色模型，如果需要，在内存中动态转换 FP16 权重。
        """
        character_name = character_name.lower()
        if character_name in self.character_to_model:
            _ = self.character_to_model[character_name]
            return True

        model_dict: Dict[str, Optional[InferenceSession]] = {}

        # 定义 ONNX 文件到 FP16 Bin 文件的映射关系
        onnx_to_fp16_map = {
            GSVModelFile.T2S_FIRST_STAGE_DECODER_FP32: GSVModelFile.T2S_DECODER_WEIGHT_FP16,
            GSVModelFile.T2S_STAGE_DECODER_FP32: GSVModelFile.T2S_DECODER_WEIGHT_FP16,
            GSVModelFile.VITS_FP32: GSVModelFile.VITS_WEIGHT_FP16,
            GSVModelFile.PROMPT_ENCODER: GSVModelFile.PROMPT_ENCODER_WEIGHT_FP16
        }

        # 确定需要加载的模型列表
        model_files_to_load = [
            GSVModelFile.T2S_ENCODER_FP32,
            GSVModelFile.VITS_FP32,
            GSVModelFile.PROMPT_ENCODER,
        ]

        fp32_decoders = [GSVModelFile.T2S_FIRST_STAGE_DECODER_FP32, GSVModelFile.T2S_STAGE_DECODER_FP32]
        model_files_to_load.extend(fp32_decoders)

        try:
            for model_file in model_files_to_load:
                model_path = os.path.normpath(os.path.join(model_dir, model_file))

                # 设置 Session Options
                sess_options = onnxruntime.SessionOptions()
                sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL

                if os.path.exists(model_path):
                    fp16_bin_name = onnx_to_fp16_map.get(model_file)
                    fp16_bin_path = os.path.join(model_dir, fp16_bin_name) if fp16_bin_name else None

                    if fp16_bin_path and os.path.exists(fp16_bin_path):
                        model_dict[model_file] = load_session_with_fp16_conversion(
                            model_path, fp16_bin_path, self.providers, sess_options
                        )
                    else:
                        model_dict[model_file] = onnxruntime.InferenceSession(
                            model_path,
                            providers=self.providers,
                            sess_options=sess_options,
                        )
                elif model_file == GSVModelFile.PROMPT_ENCODER:
                    model_dict[model_file] = None
                else:
                    raise FileNotFoundError(f'文件 {model_path} 不存在！')

            # 日志信息
            is_v2pp = model_dict[GSVModelFile.PROMPT_ENCODER] is not None
            logger.info(
                f"Character {character_name.capitalize()} loaded successfully.\n"
                f"- Model Path: {model_dir}\n"
                f"- Model Type: {'V2ProPlus' if is_v2pp else 'V2'}"
            )

            self.character_to_model[character_name] = model_dict
            self.character_to_language[character_name] = language
            self.character_model_paths[character_name] = model_dir
            return True

        except Exception as e:
            logger.error(
                f"Error: Failed to load ONNX model '{model_dir}'.\n"
                f"Details: {e}"
            )
            return False

    def remove_all_character(self) -> None:
        self.character_to_model.clear()
        gc.collect()

    def remove_character(self, character_name: str) -> None:
        character_name = character_name.lower()
        if character_name in self.character_to_model:
            del self.character_to_model[character_name]
            gc.collect()
            logger.info(f"Character {character_name.capitalize()} removed successfully.")


model_manager: ModelManager = ModelManager()