WIP: Building JIT support for glow-tts (doesn't work yet)

glow-tts-jaywalnut/inference.py

@@ -0,0 +1,78 @@
+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])