tabs/audios/load_models.py

import os

import torch
from df.enhance import init_df
from modelscope import snapshot_download

from config import Config
from .modules.CosyVoice.cosyvoice.cli.cosyvoice import CosyVoice


def init_sys():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Load DeepFilterNet2
    df_model, df_state, _ = init_df()
    
    # Download CosyVoice models
    snapshot_download('iic/CosyVoice-300M', local_dir=f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-300M')
    snapshot_download('iic/CosyVoice-300M-SFT', local_dir=f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-300M-SFT')
    snapshot_download('iic/CosyVoice-300M-Instruct', local_dir=f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-300M-Instruct')
    snapshot_download('iic/CosyVoice-ttsfrd', local_dir=f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-ttsfrd')
    
    # Add `tabs/audios/modules/CosyVoice/third_party/Matcha-TTS` to your `PYTHONPATH`
    os.environ['PYTHONPATH'] = f'{os.path.dirname(__file__)}/modules/CosyVoice/third_party/Matcha-TTS:{os.environ.get("PYTHONPATH", "")}'
    
    # Load CosyVoice SFT
    cv_sft = CosyVoice(f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-300M-SFT')
    sft_speakers = cv_sft.list_avaliable_spks()
    
    # Load CosyVoice TTS
    cv_vc = CosyVoice(f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-300M')
    
    # Load CosyVoice Instruct
    cv_instruct = CosyVoice(f'{Config.MODEL_DOWNLOAD_DIR}/audios/CosyVoice-300M-Instruct')
    
    return device, df_model, df_state, cv_vc, cv_sft, sft_speakers, cv_instruct

device, df_model, df_state, cv_vc, cv_sft, sft_speakers, cv_instruct = init_sys()