mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-18 07:33:20 +00:00
Brandon Thomas
79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
import matplotlib.pyplot as plt
|
|
|
|
import sys
|
|
sys.path.append('./waveglow/')
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import os
|
|
import glob
|
|
import json
|
|
|
|
import scipy.io.wavfile as wavfile
|
|
|
|
import torch
|
|
from text import text_to_sequence, cmudict
|
|
from text.symbols import symbols
|
|
import commons
|
|
import attentions
|
|
import modules
|
|
import models
|
|
import utils
|
|
|
|
def main(tst_stn):
|
|
|
|
# load waveglow
|
|
waveglow_path = '/home/bt/models/waveglow-nvidia/waveglow_256channels_universal_v5.pt'
|
|
waveglow = torch.load(waveglow_path)['model']
|
|
waveglow = waveglow.remove_weightnorm(waveglow)
|
|
_ = waveglow.cuda().eval()
|
|
from apex import amp
|
|
waveglow, _ = amp.initialize(waveglow, [], opt_level="O3") # Try if you want to boost up synthesis speed.
|
|
|
|
###############################################################################
|
|
|
|
model_dir = "./logs/base/"
|
|
hps = utils.get_hparams_from_dir(model_dir)
|
|
model = models.FlowGenerator(
|
|
len(symbols),
|
|
out_channels=hps.data.n_mel_channels,
|
|
**hps.model).to("cuda")
|
|
|
|
checkpoint_path = utils.latest_checkpoint_path(model_dir)
|
|
checkpoint_path = './logs/base/G_1007.pth'
|
|
print('Checkpoint path: {}'.format(checkpoint_path))
|
|
|
|
utils.load_checkpoint(checkpoint_path, model)
|
|
model.decoder.store_inverse() # do not calcuate jacobians for fast decoding
|
|
_ = model.eval()
|
|
|
|
cmu_dict = cmudict.CMUDict(hps.data.cmudict_path)
|
|
|
|
###############################################################################
|
|
|
|
# tst_stn = " Glow TTS is really awesome ! " # Adding spaces at the beginning and the end of utterance improves quality
|
|
sequence = np.array(text_to_sequence(tst_stn, ['english_cleaners'], cmu_dict))[None, :]
|
|
print("".join([symbols[c] for c in sequence[0]]))
|
|
x_tst = torch.autograd.Variable(torch.from_numpy(sequence)).cuda().long()
|
|
x_tst_lengths = torch.tensor([x_tst.shape[1]]).cuda()
|
|
|
|
###############################################################################
|
|
|
|
with torch.no_grad():
|
|
noise_scale = .667
|
|
length_scale = 1.0
|
|
(y_gen_tst, *r), attn_gen, *_ = model(x_tst, x_tst_lengths, gen=True, noise_scale=noise_scale, length_scale=length_scale)
|
|
try:
|
|
audio = waveglow.infer(y_gen_tst.half(), sigma=.666)
|
|
except:
|
|
audio = waveglow.infer(y_gen_tst, sigma=.666)
|
|
|
|
audio = audio.cpu().numpy()[0]
|
|
audio *= 2 ** 15
|
|
audio = audio.astype(np.int16)
|
|
|
|
wavfile.write('infer.wav', hps.data.sampling_rate, audio)
|
|
|
|
if __name__ == "__main__":
|
|
main(sys.argv[1])
|