mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-19 07:42:59 +00:00
Brandon Thomas
72dd12b460
I copied this from my old and messy `tacotron-melgan` repo, where I had previously experimented with running the two models in a pipeline. I'll be moving over all the JIT pieces here and cleaning up when I can.
109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
from collections import OrderedDict
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .res_stack import ResStack
|
|
|
|
import numpy as np
|
|
|
|
from wav_images import render_histogram
|
|
from wav_images import rescale_mel
|
|
|
|
def load_melgan(model_checkpoint_path):
|
|
print('Loading melgan checkpoint: {}'.format(model_checkpoint_path))
|
|
|
|
checkpoint = torch.load(model_checkpoint_path, map_location=torch.device('cpu'))
|
|
melgan_state = checkpoint['model_g']
|
|
melgan_state = transform_melgate_state(melgan_state)
|
|
|
|
model_melgan = Generator()
|
|
model_melgan.load_state_dict(melgan_state)
|
|
return model_melgan
|
|
|
|
def transform_melgate_state(melgan_state):
|
|
import re
|
|
BLOCK_REGEX = re.compile('blocks\.(\d+)\.')
|
|
SHORTCUT_REGEX = re.compile('shortcuts\.(\d+)\.')
|
|
|
|
new_model_g = [] # Rebuild the ordered dict
|
|
for key, value in melgan_state.items():
|
|
if 'blocks.' in key:
|
|
new_key = BLOCK_REGEX.sub(r'block\1.', key)
|
|
#print('{} -> {}'.format(key, new_key))
|
|
key = new_key
|
|
elif 'shortcuts.' in key:
|
|
new_key = SHORTCUT_REGEX.sub(r'shortcut\1.', key)
|
|
#print('{} -> {}'.format(key, new_key))
|
|
key = new_key
|
|
|
|
new_model_g.append((key, value))
|
|
|
|
return OrderedDict(new_model_g)
|
|
|
|
class Generator(nn.Module):
|
|
def __init__(self):
|
|
super(Generator, self).__init__()
|
|
mel_channel = 80 # TODO: This is an hparam
|
|
self.mel_channel = mel_channel
|
|
|
|
self.generator = nn.Sequential(
|
|
nn.ReflectionPad1d(3),
|
|
nn.utils.weight_norm(nn.Conv1d(mel_channel, 512, kernel_size=7, stride=1)),
|
|
|
|
nn.LeakyReLU(0.2),
|
|
nn.utils.weight_norm(nn.ConvTranspose1d(512, 256, kernel_size=16, stride=8, padding=4)),
|
|
|
|
ResStack(256),
|
|
|
|
nn.LeakyReLU(0.2),
|
|
nn.utils.weight_norm(nn.ConvTranspose1d(256, 128, kernel_size=16, stride=8, padding=4)),
|
|
|
|
ResStack(128),
|
|
|
|
nn.LeakyReLU(0.2),
|
|
nn.utils.weight_norm(nn.ConvTranspose1d(128, 64, kernel_size=4, stride=2, padding=1)),
|
|
|
|
ResStack(64),
|
|
|
|
nn.LeakyReLU(0.2),
|
|
nn.utils.weight_norm(nn.ConvTranspose1d(64, 32, kernel_size=4, stride=2, padding=1)),
|
|
|
|
ResStack(32),
|
|
|
|
nn.LeakyReLU(0.2),
|
|
nn.ReflectionPad1d(3),
|
|
nn.utils.weight_norm(nn.Conv1d(32, 1, kernel_size=7, stride=1)),
|
|
nn.Tanh(),
|
|
)
|
|
|
|
def forward(self, mel):
|
|
mel = (mel + 5.0) / 5.0 # roughly normalize spectrogram
|
|
audio = self.generator(mel)
|
|
audio = audio.squeeze()
|
|
|
|
MAX_WAV_VALUE = 32768.0
|
|
audio = audio * MAX_WAV_VALUE
|
|
|
|
max_ = torch.max(torch.abs(audio))
|
|
#max_ = audio.abs().max()
|
|
#max_ = max(0.01, max_)
|
|
multiplier = 32767.0 / max_
|
|
audio = audio * multiplier
|
|
#max_ = np.max(np.abs(audio))
|
|
#max_ = max(0.01, audio)
|
|
#print("Max: {}".format(max_))
|
|
#print("Multiplier: {}".format(multiplier))
|
|
|
|
return audio
|
|
|
|
def eval(self):
|
|
super(Generator, self).eval()
|
|
|
|
def inference(self, mel):
|
|
audio = self.forward(mel)
|
|
#audio = audio.squeeze() # collapse all dimension except time axis
|
|
return audio
|
|
|