mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-12 23:48:57 +00:00
Brandon Thomas
73 lines
1.7 KiB
Python
73 lines
1.7 KiB
Python
import os
|
|
import json
|
|
# import re
|
|
import argparse
|
|
# from string import punctuation
|
|
|
|
import torch
|
|
import numpy as np
|
|
from torch.utils.data import DataLoader
|
|
# from g2p_en import G2p
|
|
# from pypinyin import pinyin, Style
|
|
|
|
from utils.model import get_model, get_vocoder
|
|
from utils.tools import get_configs_of, to_device #, read_lexicon
|
|
from dataset import TextDataset
|
|
from text import text_to_sequence, sequence_to_text
|
|
from text.cleaners import to_arpa
|
|
import sys
|
|
device = "cpu"
|
|
from model import Tacotron2
|
|
|
|
def get_tac2(configs, device):
|
|
(preprocess_config, model_config, train_config) = configs
|
|
|
|
model = Tacotron2(preprocess_config, model_config, train_config).to(device)
|
|
|
|
model.eval()
|
|
model.requires_grad_ = False
|
|
return model
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--dataset",
|
|
type=str,
|
|
required=True,
|
|
help="name of dataset",
|
|
)
|
|
parser.add_argument(
|
|
"--checkpoint",
|
|
type=str,
|
|
required=True,
|
|
help="Path to model file.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.isfile(args.checkpoint):
|
|
print(f"Model path {args.checkpoint} incorrect, file not found.")
|
|
sys.exit(2)
|
|
|
|
|
|
|
|
# Read Config
|
|
preprocess_config, model_config, train_config = get_configs_of(args.dataset)
|
|
configs = (preprocess_config, model_config, train_config)
|
|
stats = None
|
|
mel_stats = None
|
|
|
|
# Get model
|
|
model = get_tac2(configs, device)
|
|
try:
|
|
ckpt = torch.load(args.checkpoint)
|
|
model.load_state_dict(ckpt["model"])
|
|
except:
|
|
print("Error loading model")
|
|
sys.exit(1) # Error
|
|
|
|
print("Model is valid.")
|
|
sys.exit(0) # Success
|
|
|
|
|