utils/load_model.py

import os
import json
import logging
import config
import numpy as np

import utils
from utils.data_utils import check_is_none, HParams
from vits import VITS
from voice import TTS
from config import DEVICE as device
from utils.lang_dict import lang_dict
from contants import ModelType


def recognition_model_type(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(emotion_reference_npy):
    if isinstance(emotion_reference_npy, list):
        # check if emotion_reference_npy is endwith .npy
        for i in emotion_reference_npy:
            model_extention = os.path.splitext(i)[1]
            if model_extention != ".npy":
                raise ValueError(f"Unsupported model type: {model_extention}")

        # merge npy files
        emotion_reference = np.empty((0, 1024))
        for i in emotion_reference_npy:
            tmp = np.load(i).reshape(-1, 1024)
            emotion_reference = np.append(emotion_reference, tmp, axis=0)

    elif os.path.isdir(emotion_reference_npy):
        emotion_reference = np.empty((0, 1024))
        for root, dirs, files in os.walk(emotion_reference_npy):
            for file_name in files:
                # check if emotion_reference_npy is endwith .npy
                model_extention = os.path.splitext(file_name)[1]
                if model_extention != ".npy":
                    continue
                file_path = os.path.join(root, file_name)

                # merge npy files
                tmp = np.load(file_path).reshape(-1, 1024)
                emotion_reference = np.append(emotion_reference, tmp, axis=0)

    elif os.path.isfile(emotion_reference_npy):
        # check if emotion_reference_npy is endwith .npy
        model_extention = os.path.splitext(emotion_reference_npy)[1]
        if model_extention != ".npy":
            raise ValueError(f"Unsupported model type: {model_extention}")

        emotion_reference = np.load(emotion_reference_npy)
    logging.info(f"Loaded emotional dimention npy range:{len(emotion_reference)}")
    return emotion_reference


def parse_models(model_list):
    categorized_models = {
        ModelType.VITS: [],
        ModelType.HUBERT_VITS: [],
        ModelType.W2V2_VITS: [],
        ModelType.BERT_VITS2: []
    }

    for model_info in model_list:
        config_path = model_info[1]
        hps = utils.get_hparams_from_file(config_path)
        model_info.append(hps)
        model_type = recognition_model_type(hps)
        # with open(config_path, 'r', encoding='utf-8') as model_config:
        #     model_type = recognition_model_type(model_config)
        if model_type in categorized_models:
            categorized_models[model_type].append(model_info)

    return categorized_models


def merge_models(model_list, model_class, model_type, additional_arg=None):
    id_mapping_objs = []
    speakers = []
    new_id = 0

    for obj_id, (model_path, config_path, hps) in enumerate(model_list):
        obj_args = {
            "model": model_path,
            "config": hps,
            "model_type": model_type,
            "device": device
        }

        if model_type == ModelType.BERT_VITS2:
            from bert_vits2.utils import process_legacy_versions
            legacy_versions = process_legacy_versions(hps)
            key = f"{model_type.value}_v{legacy_versions}" if legacy_versions else model_type.value         
        else:
            key = getattr(hps.data, "text_cleaners", ["none"])[0]

        if additional_arg:
            obj_args.update(additional_arg)

        obj = model_class(**obj_args)

        lang = lang_dict.get(key, ["unknown"])

        for real_id, name in enumerate(obj.get_speakers()):
            id_mapping_objs.append([real_id, obj, obj_id])
            speakers.append({"id": new_id, "name": name, "lang": lang})
            new_id += 1

    return id_mapping_objs, speakers


def load_model(model_list) -> TTS:
    categorized_models = parse_models(model_list)

    # Handle VITS
    vits_objs, vits_speakers = merge_models(categorized_models[ModelType.VITS], VITS, ModelType.VITS)

    # Handle HUBERT-VITS
    hubert_vits_objs, hubert_vits_speakers = [], []
    if len(categorized_models[ModelType.HUBERT_VITS]) != 0:
        if getattr(config, "HUBERT_SOFT_MODEL", None) is None or check_is_none(config.HUBERT_SOFT_MODEL):
            raise ValueError(f"Please configure HUBERT_SOFT_MODEL path in config.py")
        try:
            from vits.hubert_model import hubert_soft
            hubert = hubert_soft(config.HUBERT_SOFT_MODEL)
        except Exception as e:
            raise ValueError(f"Load HUBERT_SOFT_MODEL failed {e}")

        hubert_vits_objs, hubert_vits_speakers = merge_models(categorized_models[ModelType.HUBERT_VITS], VITS, ModelType.HUBERT_VITS,
                                                              additional_arg={"additional_model": hubert})

    # Handle W2V2-VITS
    w2v2_vits_objs, w2v2_vits_speakers = [], []
    w2v2_emotion_count = 0
    if len(categorized_models[ModelType.W2V2_VITS]) != 0:
        if getattr(config, "DIMENSIONAL_EMOTION_NPY", None) is None or check_is_none(
                config.DIMENSIONAL_EMOTION_NPY):
            raise ValueError(f"Please configure DIMENSIONAL_EMOTION_NPY path in config.py")
        try:
            emotion_reference = load_npy(config.DIMENSIONAL_EMOTION_NPY)
        except Exception as e:
            emotion_reference = None
            raise ValueError(f"Load DIMENSIONAL_EMOTION_NPY failed {e}")

        w2v2_vits_objs, w2v2_vits_speakers = merge_models(categorized_models[ModelType.W2V2_VITS], VITS, ModelType.W2V2_VITS,
                                                          additional_arg={"additional_model": emotion_reference})
        w2v2_emotion_count = len(emotion_reference) if emotion_reference is not None else 0

    # Handle BERT-VITS2
    bert_vits2_objs, bert_vits2_speakers = [], []
    if len(categorized_models[ModelType.BERT_VITS2]) != 0:
        from bert_vits2 import Bert_VITS2
        bert_vits2_objs, bert_vits2_speakers = merge_models(categorized_models[ModelType.BERT_VITS2], Bert_VITS2, ModelType.BERT_VITS2)

    voice_obj = {ModelType.VITS: vits_objs,
                 ModelType.HUBERT_VITS: hubert_vits_objs,
                 ModelType.W2V2_VITS: w2v2_vits_objs,
                 ModelType.BERT_VITS2: bert_vits2_objs}
    voice_speakers = {ModelType.VITS.value: vits_speakers,
                      ModelType.HUBERT_VITS.value: hubert_vits_speakers,
                      ModelType.W2V2_VITS.value: w2v2_vits_speakers,
                      ModelType.BERT_VITS2.value: bert_vits2_speakers}

    tts = TTS(voice_obj, voice_speakers, device=device, w2v2_emotion_count=w2v2_emotion_count)
    return tts