commit 72dd12b4605e4a3549dac5c0416e4b91ecfaab20
Author: Brandon Thomas <echelon@gmail.com>
Date:   Tue Jun 2 02:17:16 2020 -0400

    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.

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bb8ff54
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+*~
+__pycache__/
+python/
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..fa27d4b
--- /dev/null
+++ b/README.md
@@ -0,0 +1,7 @@
+jit-models
+==========
+
+This project exists to turn PyTorch models into Torch JIT, which can be run in
+any language with libtorch bindings.
+
+
diff --git a/melgan-swpark/model_melgan/discriminator.py b/melgan-swpark/model_melgan/discriminator.py
new file mode 100644
index 0000000..63d2736
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/melgan-swpark/model_melgan/generator.py b/melgan-swpark/model_melgan/generator.py
new file mode 100644
index 0000000..9c45075
--- /dev/null
+++ b/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
+
diff --git a/melgan-swpark/model_melgan/identity.py b/melgan-swpark/model_melgan/identity.py
new file mode 100644
index 0000000..052b37e
--- /dev/null
+++ b/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
+
diff --git a/melgan-swpark/model_melgan/multiscale.py b/melgan-swpark/model_melgan/multiscale.py
new file mode 100644
index 0000000..640b59e
--- /dev/null
+++ b/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)]
diff --git a/melgan-swpark/model_melgan/res_stack.py b/melgan-swpark/model_melgan/res_stack.py
new file mode 100644
index 0000000..5b2658f
--- /dev/null
+++ b/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
+
diff --git a/melgan-swpark/requirements.txt b/melgan-swpark/requirements.txt
new file mode 100644
index 0000000..613ce4e
--- /dev/null
+++ b/melgan-swpark/requirements.txt
@@ -0,0 +1,16 @@
+# Requirements from Tacotron
+Unidecode==0.4.20
+#falcon==1.2.0
+falcon==2.0.0
+inflect==0.2.5
+librosa==0.5.1
+numpy==1.14.3
+tensorflow==1.14.0
+
+# Requirements from MelGan
+PyYAML==5.2
+torch==1.3.1
+tqdm==4.40.2
+
+# Personal Extensions
+Pillow==6.2.1
diff --git a/melgan-swpark/wav_images.py b/melgan-swpark/wav_images.py
new file mode 100644
index 0000000..118ac79
--- /dev/null
+++ b/melgan-swpark/wav_images.py
@@ -0,0 +1,181 @@
+from PIL import Image, ImageDraw
+
+def render_histogram(wav, filename='histograms/output.png', row_index=0, col_index=1):
+    dimensions = wav.shape
+    print('>>> Image input size: {}'.format(dimensions))
+    rows = dimensions[row_index]
+    cols = dimensions[col_index]
+
+    # Remove 1-sized dimensions
+    wav_signal = wav.squeeze()
+
+    print('>>> Generating image (rows: {} cols: {})'.format(rows, cols))
+    image = Image.new('RGB', (rows, cols))
+    pixels = image.load()
+
+    minimum = wav_signal[0,0]
+    maximum = wav_signal[0,0]
+
+    for x in range(0, rows):
+        for y in range(0, cols):
+            if wav_signal[x,y] > maximum:
+                maximum = wav_signal[x,y]
+            if wav_signal[x,y] < minimum:
+                minimum = wav_signal[x,y]
+
+    print('Minimum: ' + str(minimum))
+    print('Maximum: ' + str(maximum))
+
+    if maximum == minimum:
+        maximum = 1 # nb: to prevent division by zero below
+
+    for x in range(0, rows):
+        for y in range(0, cols):
+            v = wav_signal[x,y]
+            scaled = int((v - minimum) / (maximum - minimum) * 255)
+            pixels[x, y] = (scaled, scaled, scaled)
+
+    # Image Show only works on local X server:
+    #image.show()
+    image.save(filename)
+
+def rescale_mel(mel, rescaled_min=-1.0, rescaled_max=1.0):
+    # NB: Assumes 3 dimensions
+    dimensions = mel.shape
+    minimum = mel[0,0,0]
+    maximum = mel[0,0,0]
+
+    for x in range(0, dimensions[0]):
+        for y in range(0, dimensions[1]):
+            for z in range(0, dimensions[2]):
+                if mel[x,y,z] > maximum:
+                    maximum = mel[x,y,z]
+                if mel[x,y,z] < minimum:
+                    minimum = mel[x,y,z]
+
+    print('Actual Minimum: ' + str(minimum))
+    print('Actual Maximum: ' + str(maximum))
+
+    old_range = float(maximum) - minimum
+    rescaled_range = float(rescaled_max) - rescaled_min
+
+    print('Old Range: ' + str(old_range))
+    print('Rescaled Range: ' + str(rescaled_range))
+
+    for x in range(0, dimensions[0]):
+        for y in range(0, dimensions[1]):
+            for z in range(0, dimensions[2]):
+                d = mel[x,y,z]
+                rescaled = (d - minimum) * rescaled_range / old_range + rescaled_min
+                mel[x,y,z] = rescaled
+
+"""
+# The following is taken from SciPy Cookbook
+# https://scipy-cookbook.readthedocs.io/items/Rebinning.html
+import numpy as n
+import scipy.interpolate
+import scipy.ndimage
+
+def congrid(a, newdims, method='linear', centre=False, minusone=False):
+    '''Arbitrary resampling of source array to new dimension sizes.
+    Currently only supports maintaining the same number of dimensions.
+    To use 1-D arrays, first promote them to shape (x,1).
+
+    Uses the same parameters and creates the same co-ordinate lookup points
+    as IDL''s congrid routine, which apparently originally came from a VAX/VMS
+    routine of the same name.
+
+    method:
+    neighbour - closest value from original data
+    nearest and linear - uses n x 1-D interpolations using
+                         scipy.interpolate.interp1d
+    (see Numerical Recipes for validity of use of n 1-D interpolations)
+    spline - uses ndimage.map_coordinates
+
+    centre:
+    True - interpolation points are at the centres of the bins
+    False - points are at the front edge of the bin
+
+    minusone:
+    For example- inarray.shape = (i,j) & new dimensions = (x,y)
+    False - inarray is resampled by factors of (i/x) * (j/y)
+    True - inarray is resampled by(i-1)/(x-1) * (j-1)/(y-1)
+    This prevents extrapolation one element beyond bounds of input array.
+    '''
+    if not a.dtype in [n.float64, n.float32]:
+        a = n.cast[float](a)
+
+    m1 = n.cast[int](minusone)
+    ofs = n.cast[int](centre) * 0.5
+    old = n.array( a.shape )
+    ndims = len( a.shape )
+    if len( newdims ) != ndims:
+        print("[congrid] dimensions error. " \
+              "This routine currently only support " \
+              "rebinning to the same number of dimensions.")
+        return None
+    newdims = n.asarray( newdims, dtype=float )
+    dimlist = []
+
+    if method == 'neighbour':
+        for i in range( ndims ):
+            base = n.indices(newdims)[i]
+            dimlist.append( (old[i] - m1) / (newdims[i] - m1) \
+                            * (base + ofs) - ofs )
+        cd = n.array( dimlist ).round().astype(int)
+        newa = a[list( cd )]
+        return newa
+
+    elif method in ['nearest','linear']:
+        # calculate new dims
+        for i in range( ndims ):
+            base = n.arange( newdims[i] )
+            dimlist.append( (old[i] - m1) / (newdims[i] - m1) \
+                            * (base + ofs) - ofs )
+        # specify old dims
+        olddims = [n.arange(i, dtype = n.float) for i in list( a.shape )]
+
+        # first interpolation - for ndims = any
+        mint = scipy.interpolate.interp1d( olddims[-1], a, kind=method )
+        newa = mint( dimlist[-1] )
+
+        trorder = [ndims - 1] + list(range( ndims - 1 ))
+        for i in range( ndims - 2, -1, -1 ):
+            newa = newa.transpose( trorder )
+
+            mint = scipy.interpolate.interp1d( olddims[i], newa, kind=method )
+            newa = mint( dimlist[i] )
+
+        if ndims > 1:
+            # need one more transpose to return to original dimensions
+            newa = newa.transpose( trorder )
+
+        return newa
+    elif method in ['spline']:
+        oslices = [ slice(0,j) for j in old ]
+        oldcoords = n.ogrid[oslices]
+        nslices = [ slice(0,j) for j in list(newdims) ]
+        newcoords = n.mgrid[nslices]
+
+        newcoords_dims = list(range(n.rank(newcoords)))
+        #make first index last
+        newcoords_dims.append(newcoords_dims.pop(0))
+        newcoords_tr = newcoords.transpose(newcoords_dims)
+        # makes a view that affects newcoords
+
+        newcoords_tr += ofs
+
+        deltas = (n.asarray(old) - m1) / (newdims - m1)
+        newcoords_tr *= deltas
+
+        newcoords_tr -= ofs
+
+        newa = scipy.ndimage.map_coordinates(a, newcoords)
+        return newa
+    else:
+        print("Congrid error: Unrecognized interpolation type.\n", \
+              "Currently only \'neighbour\', \'nearest\',\'linear\',", \
+              "and \'spline\' are supported.")
+        return None
+"""
+
diff --git a/melgan-swpark/wrap_melgan_model.py b/melgan-swpark/wrap_melgan_model.py
new file mode 100755
index 0000000..8499f85
--- /dev/null
+++ b/melgan-swpark/wrap_melgan_model.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# From: https://github.com/pytorch/pytorch/issues/20356#issuecomment-545572400
+import torch
+
+from collections import OrderedDict
+import re
+import yaml
+
+from model_melgan.generator import load_melgan
+
+#torch.set_default_tensor_type('torch.DoubleTensor')
+
+#model_checkpoint_path = '/home/bt/models/melgan-swpark/trumptxlearn183_fa14347_7675.pt'
+model_checkpoint_path = '/home/bt/models/melgan-swpark/ljs-high-pitch_8741385_6975.pt'
+model_checkpoint_path = '/home/bt/models/melgan-swpark/trump-txlearn-2020.05.05_8741385_7525.pt'
+model_checkpoint_path = '/home/bt/models/melgan-swpark/trump-txlearn-2020.05.05_8741385_13675.pt'
+model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_trump_txlearn_2020-05-30_steno_chkpt9800.pt'
+
+melgan_model = load_melgan(model_checkpoint_path)
+
+#output_filename = 'shared_melgan_container2.pt'
+#output_filename = 'melgan_ljs-high-pitch_container.pt'
+output_filename = 'melgan_trump-txlearn-2020.05.05.jit'
+output_filename = 'melgan_trump-txlearn-2020.05.05_13675.jit'
+output_filename = 'melgan_trump_txlearn_2020-05-30_steno_chkpt9800.torchjit'
+
+# NB: Tracing evaluates the model on input and unrolls and hardcodes branching
+# and loops. Scripting allows these to remain by converting the entire program
+# to TorchScript. Scripting seems MUCH faster, is also harder get working.
+trace_model = True
+if trace_model:
+    print('Tracing model and saving as: {}'.format(output_filename))
+    mel_file = '/home/bt/dev/voder/data/mels/LJ002-0320.mel'
+    mel = torch.load(mel_file, map_location=torch.device('cpu'))
+    # NB: Getting `Tracing failed sanity checks! Graphs differed across invocations!`
+    traced_script_module = torch.jit.trace(melgan_model, mel, check_trace=True)
+    traced_script_module.save(output_filename)
+else:
+    print('Saving script model as: {}'.format(output_filename))
+    container = torch.jit.script(melgan_model)
+    container.save(output_filename)
+
+