zhangshuai/MeloTTS
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add HF hub compatibility

Browse Source
This commit is contained in:
Yoach Lacombe
2024-02-29 15:55:41 +00:00
parent 69983af7b1
commit b3ed38400f
3 changed files with 73 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
melo/api.py
View File

@@ -20,7 +20,8 @@ from .download_utils import load_or_download_config, load_or_download_model
class TTS(nn.Module):
def __init__(self,
language,
device='auto'):
device='auto',
use_hf=True):
super().__init__()
if device == 'auto':
device = 'cpu'
@@ -30,7 +31,7 @@ class TTS(nn.Module):
assert torch.cuda.is_available()
# config_path =
hps = load_or_download_config(language)
hps = load_or_download_config(language, use_hf=use_hf)
num_languages = hps.num_languages
num_tones = hps.num_tones
@@ -53,7 +54,7 @@ class TTS(nn.Module):
self.device = device
# load state_dict
checkpoint_dict = load_or_download_model(language, device)
checkpoint_dict = load_or_download_model(language, device, use_hf=use_hf)
self.model.load_state_dict(checkpoint_dict['model'], strict=True)
language = language.split('_')[0]

32
melo/download_utils.py
View File

@@ -2,6 +2,8 @@ import torch
import os
from . import utils
from cached_path import cached_path
from huggingface_hub import hf_hub_download
DOWNLOAD_CKPT_URLS = {
'EN': 'https://myshell-public-repo-hosting.s3.amazonaws.com/openvoice/basespeakers/EN/checkpoint.pth',
'EN_V2': 'https://myshell-public-repo-hosting.s3.amazonaws.com/openvoice/basespeakers/EN_V2/checkpoint.pth',
@@ -22,14 +24,32 @@ DOWNLOAD_CONFIG_URLS = {
'KR': 'https://myshell-public-repo-hosting.s3.amazonaws.com/openvoice/basespeakers/KR/config.json',
}
def load_or_download_config(locale):
LANG_TO_HF_REPO_ID = {
'EN': 'myshell-ai/MeloTTS-English',
'EN_V2': 'myshell-ai/MeloTTS-English-v2',
'FR': 'myshell-ai/MeloTTS-French',
'JP': 'myshell-ai/MeloTTS-Japanese',
'ES': 'myshell-ai/MeloTTS-Spanish',
'ZH': 'myshell-ai/MeloTTS-Chinese',
'KR': 'myshell-ai/MeloTTS-Korean',
}
def load_or_download_config(locale, use_hf=True):
language = locale.split('-')[0].upper()
assert language in DOWNLOAD_CONFIG_URLS
config_path = cached_path(DOWNLOAD_CONFIG_URLS[language])
if use_hf:
assert language in LANG_TO_HF_REPO_ID
config_path = hf_hub_download(repo_id=LANG_TO_HF_REPO_ID[language], filename="config.json")
else:
assert language in DOWNLOAD_CONFIG_URLS
config_path = cached_path(DOWNLOAD_CONFIG_URLS[language])
return utils.get_hparams_from_file(config_path)
def load_or_download_model(locale, device):
def load_or_download_model(locale, device, use_hf=True):
language = locale.split('-')[0].upper()
assert language in DOWNLOAD_CKPT_URLS
ckpt_path = cached_path(DOWNLOAD_CKPT_URLS[language])
if use_hf:
assert language in LANG_TO_HF_REPO_ID
ckpt_path = hf_hub_download(repo_id=LANG_TO_HF_REPO_ID[language], filename="checkpoint.pth")
else:
assert language in DOWNLOAD_CKPT_URLS
ckpt_path = cached_path(DOWNLOAD_CKPT_URLS[language])
return torch.load(ckpt_path, map_location=device)