first commit; melgan support (perhaps bloated)

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.
2026-10-28 08:47:19 +00:00 · 2020-06-02 02:17:16 -04:00
commit 72dd12b460
10 changed files with 502 additions and 0 deletions
@@ -0,0 +1,108 @@
+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
+