dependencies = ['torch', 'torchaudio', 'numpy', 'vocos']

import logging
from pathlib import Path

import torch
from inference import Mars5TTS, InferenceConfig

ar_url = "https://github.com/Camb-ai/mars5-tts/releases/download/v0.1-checkpoints/mars5_en_checkpoints_ar-1680000.pt"
nar_url = "https://github.com/Camb-ai/mars5-tts/releases/download/v0.1-checkpoints/mars5_en_checkpoints_nar-1260000.pt"

def mars5_english(pretrained=True, progress=True, device=None, ar_path=None, nar_path=None) -> Mars5TTS:
    """ Load mars5 english model on `device`, optionally show `progress`. """
    if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu'
    logging.info(f"Using device: {device}")
    if pretrained == False: raise AssertionError('Only pretrained model currently supported.')
    logging.info("Loading AR checkpoint...")
    if ar_path is None:
        ar_ckpt = torch.hub.load_state_dict_from_url(
            ar_url, progress=progress, check_hash=False, map_location='cpu'
        )
    else: ar_ckpt = torch.load(str(ar_path), map_location='cpu')

    logging.info("Loading NAR checkpoint...")
    if nar_path is None:
        nar_ckpt = torch.hub.load_state_dict_from_url(
            nar_url, progress=progress, check_hash=False, map_location='cpu'
        )
    else: nar_ckpt = torch.load(str(nar_path), map_location='cpu')
    logging.info("Initializing modules...")
    mars5 = Mars5TTS(ar_ckpt, nar_ckpt, device=device)
    return mars5, InferenceConfig