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.

melgan-swpark/model_melgan/discriminator.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+
+        self.discriminator = nn.ModuleList([
+            nn.Sequential(
+                nn.ReflectionPad1d(7),
+                nn.utils.weight_norm(nn.Conv1d(1, 16, kernel_size=15, stride=1)),
+                nn.LeakyReLU(0.2, inplace=True),
+            ),
+            nn.Sequential(
+                nn.utils.weight_norm(nn.Conv1d(16, 64, kernel_size=41, stride=4, padding=20, groups=4)),
+                nn.LeakyReLU(0.2, inplace=True),
+            ),
+            nn.Sequential(
+                nn.utils.weight_norm(nn.Conv1d(64, 256, kernel_size=41, stride=4, padding=20, groups=16)),
+                nn.LeakyReLU(0.2, inplace=True),
+            ),
+            nn.Sequential(
+                nn.utils.weight_norm(nn.Conv1d(256, 1024, kernel_size=41, stride=4, padding=20, groups=64)),
+                nn.LeakyReLU(0.2, inplace=True),
+            ),
+            nn.Sequential(
+                nn.utils.weight_norm(nn.Conv1d(1024, 1024, kernel_size=41, stride=4, padding=20, groups=256)),
+                nn.LeakyReLU(0.2, inplace=True),
+            ),
+            nn.Sequential(
+                nn.utils.weight_norm(nn.Conv1d(1024, 1024, kernel_size=5, stride=1, padding=2)),
+                nn.LeakyReLU(0.2, inplace=True),
+            ),
+            nn.utils.weight_norm(nn.Conv1d(1024, 1, kernel_size=3, stride=1, padding=1)),
+        ])
+
+    def forward(self, x):
+        '''
+            returns: (list of 6 features, discriminator score)
+            we directly predict score without last sigmoid function
+            since we're using Least Squares GAN (https://arxiv.org/abs/1611.04076)
+        '''
+        features = list()
+        for module in self.discriminator:
+            x = module(x)
+            features.append(x)
+        return features[:-1], features[-1]
+
+
+if __name__ == '__main__':
+    model = Discriminator()
+
+    x = torch.randn(3, 1, 22050)
+    print(x.shape)
+
+    features, score = model(x)
+    for feat in features:
+        print(feat.shape)
+    print(score.shape)
+
+    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(pytorch_total_params)

melgan-swpark/model_melgan/generator.py

@@ -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
+

melgan-swpark/model_melgan/identity.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+

melgan-swpark/model_melgan/multiscale.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .discriminator import Discriminator
+from .identity import Identity
+
+
+class MultiScaleDiscriminator(nn.Module):
+    def __init__(self):
+        super(MultiScaleDiscriminator, self).__init__()
+
+        self.discriminators = nn.ModuleList(
+            [Discriminator() for _ in range(3)]
+        )
+        
+        self.pooling = nn.ModuleList(
+            [Identity()] +
+            [nn.AvgPool1d(kernel_size=4, stride=2, padding=1, count_include_pad=False) for _ in range(1, 3)]
+        )
+
+    def forward(self, x):
+        ret = list()
+
+        for pool, disc in zip(self.pooling, self.discriminators):
+            x = pool(x)
+            ret.append(disc(x))
+
+        return ret # [(feat, score), (feat, score), (feat, score)]

melgan-swpark/model_melgan/res_stack.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class ResStack(nn.Module):
+    def __init__(self, channel):
+        super(ResStack, self).__init__()
+
+        block_list = [
+            nn.Sequential(
+                nn.LeakyReLU(0.2),
+                nn.ReflectionPad1d(3**i),
+                nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=3, dilation=3**i)),
+                nn.LeakyReLU(0.2),
+                nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)),
+            )
+            for i in range(3)
+        ]
+
+        self.block0 = block_list[0]
+        self.block1 = block_list[1]
+        self.block2 = block_list[2]
+
+        shortcut_list = [
+            nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
+            for i in range(3)
+        ]
+
+        self.shortcut0 = shortcut_list[0]
+        self.shortcut1 = shortcut_list[1]
+        self.shortcut2 = shortcut_list[2]
+
+    def forward(self, x):
+        x = self.shortcut0(x) + self.block0(x)
+        x = self.shortcut1(x) + self.block1(x)
+        x = self.shortcut2(x) + self.block2(x)
+        return x
+