CAMB-AI/MARS5-TTS

1dependencies = ['torch', 'torchaudio', 'numpy', 'vocos']2 3import logging4from pathlib import Path5 6import torch7from inference import Mars5TTS, InferenceConfig8 9ar_url = "https://github.com/Camb-ai/mars5-tts/releases/download/v0.1-checkpoints/mars5_en_checkpoints_ar-1680000.pt"10nar_url = "https://github.com/Camb-ai/mars5-tts/releases/download/v0.1-checkpoints/mars5_en_checkpoints_nar-1260000.pt"11 12def mars5_english(pretrained=True, progress=True, device=None, ar_path=None, nar_path=None) -> Mars5TTS:13    """ Load mars5 english model on `device`, optionally show `progress`. """14    if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu'15    logging.info(f"Using device: {device}")16    if pretrained == False: raise AssertionError('Only pretrained model currently supported.')17    logging.info("Loading AR checkpoint...")18    if ar_path is None:19        ar_ckpt = torch.hub.load_state_dict_from_url(20            ar_url, progress=progress, check_hash=False, map_location='cpu'21        )22    else: ar_ckpt = torch.load(str(ar_path), map_location='cpu')23 24    logging.info("Loading NAR checkpoint...")25    if nar_path is None:26        nar_ckpt = torch.hub.load_state_dict_from_url(27            nar_url, progress=progress, check_hash=False, map_location='cpu'28        )29    else: nar_ckpt = torch.load(str(nar_path), map_location='cpu')30    logging.info("Initializing modules...")31    mars5 = Mars5TTS(ar_ckpt, nar_ckpt, device=device)32    return mars5, InferenceConfig33 34