diff --git a/FastSpeech/.gitignore b/FastSpeech/.gitignore
new file mode 100644
index 0000000..84f4c08
--- /dev/null
+++ b/FastSpeech/.gitignore
@@ -0,0 +1,114 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+__pycache__
+.vscode
+.DS_Store
+
+
+data/train.txt
+model_new/
+mels/
+alignments/
diff --git a/FastSpeech/LICENSE b/FastSpeech/LICENSE
new file mode 100644
index 0000000..db81f11
--- /dev/null
+++ b/FastSpeech/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Zhengxi Liu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file
diff --git a/FastSpeech/README.md b/FastSpeech/README.md
new file mode 100644
index 0000000..f35f742
--- /dev/null
+++ b/FastSpeech/README.md
@@ -0,0 +1,68 @@
+# FastSpeech-Pytorch
+The Implementation of FastSpeech Based on Pytorch.
+
+## Update
+### 2019/10/23
+1. Fix bugs in alignment;
+2. Fix bugs in transformer;
+3. Fix bugs in LengthRegulator;
+4. Change the way to process audio;
+5. Use waveglow to synthesize.
+
+## Model
+
+
+
+
+
+
+
'
+EOS_WORD = ''
diff --git a/FastSpeech/transformer/Layers.py b/FastSpeech/transformer/Layers.py
new file mode 100644
index 0000000..a2db196
--- /dev/null
+++ b/FastSpeech/transformer/Layers.py
@@ -0,0 +1,230 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import numpy as np
+from collections import OrderedDict
+
+from transformer.SubLayers import MultiHeadAttention, PositionwiseFeedForward
+from text.symbols import symbols
+
+
+class Linear(nn.Module):
+ """
+ Linear Module
+ """
+
+ def __init__(self, in_dim, out_dim, bias=True, w_init='linear'):
+ """
+ :param in_dim: dimension of input
+ :param out_dim: dimension of output
+ :param bias: boolean. if True, bias is included.
+ :param w_init: str. weight inits with xavier initialization.
+ """
+ super(Linear, self).__init__()
+ self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)
+
+ nn.init.xavier_uniform_(
+ self.linear_layer.weight,
+ gain=nn.init.calculate_gain(w_init))
+
+ def forward(self, x):
+ return self.linear_layer(x)
+
+
+class PreNet(nn.Module):
+ """
+ Pre Net before passing through the network
+ """
+
+ def __init__(self, input_size, hidden_size, output_size, p=0.5):
+ """
+ :param input_size: dimension of input
+ :param hidden_size: dimension of hidden unit
+ :param output_size: dimension of output
+ """
+ super(PreNet, self).__init__()
+ self.input_size = input_size
+ self.output_size = output_size
+ self.hidden_size = hidden_size
+ self.layer = nn.Sequential(OrderedDict([
+ ('fc1', Linear(self.input_size, self.hidden_size)),
+ ('relu1', nn.ReLU()),
+ ('dropout1', nn.Dropout(p)),
+ ('fc2', Linear(self.hidden_size, self.output_size)),
+ ('relu2', nn.ReLU()),
+ ('dropout2', nn.Dropout(p)),
+ ]))
+
+ def forward(self, input_):
+
+ out = self.layer(input_)
+
+ return out
+
+
+class Conv(nn.Module):
+ """
+ Convolution Module
+ """
+
+ def __init__(self,
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ dilation=1,
+ bias=True,
+ w_init='linear'):
+ """
+ :param in_channels: dimension of input
+ :param out_channels: dimension of output
+ :param kernel_size: size of kernel
+ :param stride: size of stride
+ :param padding: size of padding
+ :param dilation: dilation rate
+ :param bias: boolean. if True, bias is included.
+ :param w_init: str. weight inits with xavier initialization.
+ """
+ super(Conv, self).__init__()
+
+ self.conv = nn.Conv1d(in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ bias=bias)
+
+ nn.init.xavier_uniform_(
+ self.conv.weight, gain=nn.init.calculate_gain(w_init))
+
+ def forward(self, x):
+ x = self.conv(x)
+ return x
+
+
+class FFTBlock(torch.nn.Module):
+ """FFT Block"""
+
+ def __init__(self,
+ d_model,
+ d_inner,
+ n_head,
+ d_k,
+ d_v,
+ dropout=0.1):
+ super(FFTBlock, self).__init__()
+ self.slf_attn = MultiHeadAttention(
+ n_head, d_model, d_k, d_v, dropout=dropout)
+ self.pos_ffn = PositionwiseFeedForward(
+ d_model, d_inner, dropout=dropout)
+
+ def forward(self, enc_input, non_pad_mask=None, slf_attn_mask=None):
+ enc_output, enc_slf_attn = self.slf_attn(
+ enc_input, enc_input, enc_input, mask=slf_attn_mask)
+ enc_output *= non_pad_mask
+
+ enc_output = self.pos_ffn(enc_output)
+ enc_output *= non_pad_mask
+
+ return enc_output, enc_slf_attn
+
+
+class ConvNorm(torch.nn.Module):
+ def __init__(self,
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=None,
+ dilation=1,
+ bias=True,
+ w_init_gain='linear'):
+ super(ConvNorm, self).__init__()
+
+ if padding is None:
+ assert(kernel_size % 2 == 1)
+ padding = int(dilation * (kernel_size - 1) / 2)
+
+ self.conv = torch.nn.Conv1d(in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ bias=bias)
+
+ torch.nn.init.xavier_uniform_(
+ self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain))
+
+ def forward(self, signal):
+ conv_signal = self.conv(signal)
+
+ return conv_signal
+
+
+class PostNet(nn.Module):
+ """
+ PostNet: Five 1-d convolution with 512 channels and kernel size 5
+ """
+
+ def __init__(self,
+ n_mel_channels=80,
+ postnet_embedding_dim=512,
+ postnet_kernel_size=5,
+ postnet_n_convolutions=5):
+
+ super(PostNet, self).__init__()
+ self.convolutions = nn.ModuleList()
+
+ self.convolutions.append(
+ nn.Sequential(
+ ConvNorm(n_mel_channels,
+ postnet_embedding_dim,
+ kernel_size=postnet_kernel_size,
+ stride=1,
+ padding=int((postnet_kernel_size - 1) / 2),
+ dilation=1,
+ w_init_gain='tanh'),
+
+ nn.BatchNorm1d(postnet_embedding_dim))
+ )
+
+ for i in range(1, postnet_n_convolutions - 1):
+ self.convolutions.append(
+ nn.Sequential(
+ ConvNorm(postnet_embedding_dim,
+ postnet_embedding_dim,
+ kernel_size=postnet_kernel_size,
+ stride=1,
+ padding=int((postnet_kernel_size - 1) / 2),
+ dilation=1,
+ w_init_gain='tanh'),
+
+ nn.BatchNorm1d(postnet_embedding_dim))
+ )
+
+ self.convolutions.append(
+ nn.Sequential(
+ ConvNorm(postnet_embedding_dim,
+ n_mel_channels,
+ kernel_size=postnet_kernel_size,
+ stride=1,
+ padding=int((postnet_kernel_size - 1) / 2),
+ dilation=1,
+ w_init_gain='linear'),
+
+ nn.BatchNorm1d(n_mel_channels))
+ )
+
+ def forward(self, x):
+ x = x.contiguous().transpose(1, 2)
+
+ for i in range(len(self.convolutions) - 1):
+ x = F.dropout(torch.tanh(
+ self.convolutions[i](x)), 0.5, self.training)
+ x = F.dropout(self.convolutions[-1](x), 0.5, self.training)
+
+ x = x.contiguous().transpose(1, 2)
+ return x
diff --git a/FastSpeech/transformer/Models.py b/FastSpeech/transformer/Models.py
new file mode 100644
index 0000000..37a2e42
--- /dev/null
+++ b/FastSpeech/transformer/Models.py
@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+import transformer.Constants as Constants
+from transformer.Layers import FFTBlock, PreNet, PostNet, Linear
+from text.symbols import symbols
+import hparams as hp
+
+
+def get_non_pad_mask(seq):
+ assert seq.dim() == 2
+ return seq.ne(Constants.PAD).type(torch.float).unsqueeze(-1)
+
+
+def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None):
+ ''' Sinusoid position encoding table '''
+
+ def cal_angle(position, hid_idx):
+ return position / np.power(10000, 2 * (hid_idx // 2) / d_hid)
+
+ def get_posi_angle_vec(position):
+ return [cal_angle(position, hid_j) for hid_j in range(d_hid)]
+
+ sinusoid_table = np.array([get_posi_angle_vec(pos_i)
+ for pos_i in range(n_position)])
+
+ sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i
+ sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1
+
+ if padding_idx is not None:
+ # zero vector for padding dimension
+ sinusoid_table[padding_idx] = 0.
+
+ return torch.FloatTensor(sinusoid_table)
+
+
+def get_attn_key_pad_mask(seq_k, seq_q):
+ ''' For masking out the padding part of key sequence. '''
+
+ # Expand to fit the shape of key query attention matrix.
+ len_q = seq_q.size(1)
+ padding_mask = seq_k.eq(Constants.PAD)
+ padding_mask = padding_mask.unsqueeze(
+ 1).expand(-1, len_q, -1) # b x lq x lk
+
+ return padding_mask
+
+
+class Encoder(nn.Module):
+ ''' Encoder '''
+
+ def __init__(self,
+ n_src_vocab=len(symbols)+1,
+ len_max_seq=hp.max_sep_len,
+ d_word_vec=hp.word_vec_dim,
+ n_layers=hp.encoder_n_layer,
+ n_head=hp.encoder_head,
+ d_k=64,
+ d_v=64,
+ d_model=hp.word_vec_dim,
+ d_inner=hp.encoder_conv1d_filter_size,
+ dropout=hp.dropout):
+
+ super(Encoder, self).__init__()
+
+ n_position = len_max_seq + 1
+
+ self.src_word_emb = nn.Embedding(
+ n_src_vocab, d_word_vec, padding_idx=Constants.PAD)
+
+ self.position_enc = nn.Embedding.from_pretrained(
+ get_sinusoid_encoding_table(n_position, d_word_vec, padding_idx=0),
+ freeze=True)
+
+ self.layer_stack = nn.ModuleList([FFTBlock(
+ d_model, d_inner, n_head, d_k, d_v, dropout=dropout) for _ in range(n_layers)])
+
+ def forward(self, src_seq, src_pos, return_attns=False):
+
+ enc_slf_attn_list = []
+
+ # -- Prepare masks
+ slf_attn_mask = get_attn_key_pad_mask(seq_k=src_seq, seq_q=src_seq)
+ non_pad_mask = get_non_pad_mask(src_seq)
+
+ # -- Forward
+ enc_output = self.src_word_emb(src_seq) + self.position_enc(src_pos)
+
+ for enc_layer in self.layer_stack:
+ enc_output, enc_slf_attn = enc_layer(
+ enc_output,
+ non_pad_mask=non_pad_mask,
+ slf_attn_mask=slf_attn_mask)
+ if return_attns:
+ enc_slf_attn_list += [enc_slf_attn]
+
+ return enc_output, non_pad_mask
+
+
+class Decoder(nn.Module):
+ """ Decoder """
+
+ def __init__(self,
+ len_max_seq=hp.max_sep_len,
+ d_word_vec=hp.word_vec_dim,
+ n_layers=hp.decoder_n_layer,
+ n_head=hp.decoder_head,
+ d_k=64,
+ d_v=64,
+ d_model=hp.word_vec_dim,
+ d_inner=hp.decoder_conv1d_filter_size,
+ dropout=hp.dropout):
+
+ super(Decoder, self).__init__()
+
+ n_position = len_max_seq + 1
+
+ self.position_enc = nn.Embedding.from_pretrained(
+ get_sinusoid_encoding_table(n_position, d_word_vec, padding_idx=0),
+ freeze=True)
+
+ self.layer_stack = nn.ModuleList([FFTBlock(
+ d_model, d_inner, n_head, d_k, d_v, dropout=dropout) for _ in range(n_layers)])
+
+ def forward(self, enc_seq, enc_pos, return_attns=False):
+
+ dec_slf_attn_list = []
+
+ # -- Prepare masks
+ slf_attn_mask = get_attn_key_pad_mask(seq_k=enc_pos, seq_q=enc_pos)
+ non_pad_mask = get_non_pad_mask(enc_pos)
+
+ # -- Forward
+ dec_output = enc_seq + self.position_enc(enc_pos)
+
+ for dec_layer in self.layer_stack:
+ dec_output, dec_slf_attn = dec_layer(
+ dec_output,
+ non_pad_mask=non_pad_mask,
+ slf_attn_mask=slf_attn_mask)
+ if return_attns:
+ dec_slf_attn_list += [dec_slf_attn]
+
+ return dec_output
diff --git a/FastSpeech/transformer/Modules.py b/FastSpeech/transformer/Modules.py
new file mode 100644
index 0000000..a9101a5
--- /dev/null
+++ b/FastSpeech/transformer/Modules.py
@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class ScaledDotProductAttention(nn.Module):
+ ''' Scaled Dot-Product Attention '''
+
+ def __init__(self, temperature, attn_dropout=0.1):
+ super().__init__()
+ self.temperature = temperature
+ self.dropout = nn.Dropout(attn_dropout)
+ self.softmax = nn.Softmax(dim=2)
+
+ def forward(self, q, k, v, mask=None):
+
+ attn = torch.bmm(q, k.transpose(1, 2))
+ attn = attn / self.temperature
+
+ if mask is not None:
+ attn = attn.masked_fill(mask, -np.inf)
+
+ attn = self.softmax(attn)
+ attn = self.dropout(attn)
+ output = torch.bmm(attn, v)
+
+ return output, attn
diff --git a/FastSpeech/transformer/SubLayers.py b/FastSpeech/transformer/SubLayers.py
new file mode 100644
index 0000000..5a92810
--- /dev/null
+++ b/FastSpeech/transformer/SubLayers.py
@@ -0,0 +1,97 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from transformer.Modules import ScaledDotProductAttention
+import hparams as hp
+
+
+class MultiHeadAttention(nn.Module):
+ ''' Multi-Head Attention module '''
+
+ def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
+ super().__init__()
+
+ self.n_head = n_head
+ self.d_k = d_k
+ self.d_v = d_v
+
+ self.w_qs = nn.Linear(d_model, n_head * d_k)
+ self.w_ks = nn.Linear(d_model, n_head * d_k)
+ self.w_vs = nn.Linear(d_model, n_head * d_v)
+ nn.init.normal_(self.w_qs.weight, mean=0,
+ std=np.sqrt(2.0 / (d_model + d_k)))
+ nn.init.normal_(self.w_ks.weight, mean=0,
+ std=np.sqrt(2.0 / (d_model + d_k)))
+ nn.init.normal_(self.w_vs.weight, mean=0,
+ std=np.sqrt(2.0 / (d_model + d_v)))
+
+ self.attention = ScaledDotProductAttention(
+ temperature=np.power(d_k, 0.5))
+ self.layer_norm = nn.LayerNorm(d_model)
+
+ self.fc = nn.Linear(n_head * d_v, d_model)
+ nn.init.xavier_normal_(self.fc.weight)
+
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, q, k, v, mask=None):
+
+ d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+
+ sz_b, len_q, _ = q.size()
+ sz_b, len_k, _ = k.size()
+ sz_b, len_v, _ = v.size()
+
+ residual = q
+
+ q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
+ k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
+ v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
+
+ q = q.permute(2, 0, 1, 3).contiguous().view(-1,
+ len_q, d_k) # (n*b) x lq x dk
+ k = k.permute(2, 0, 1, 3).contiguous().view(-1,
+ len_k, d_k) # (n*b) x lk x dk
+ v = v.permute(2, 0, 1, 3).contiguous().view(-1,
+ len_v, d_v) # (n*b) x lv x dv
+
+ mask = mask.repeat(n_head, 1, 1) # (n*b) x .. x ..
+ output, attn = self.attention(q, k, v, mask=mask)
+
+ output = output.view(n_head, sz_b, len_q, d_v)
+ output = output.permute(1, 2, 0, 3).contiguous().view(
+ sz_b, len_q, -1) # b x lq x (n*dv)
+
+ output = self.dropout(self.fc(output))
+ output = self.layer_norm(output + residual)
+
+ return output, attn
+
+
+class PositionwiseFeedForward(nn.Module):
+ ''' A two-feed-forward-layer module '''
+
+ def __init__(self, d_in, d_hid, dropout=0.1):
+ super().__init__()
+
+ # Use Conv1D
+ # position-wise
+ self.w_1 = nn.Conv1d(
+ d_in, d_hid, kernel_size=hp.fft_conv1d_kernel, padding=hp.fft_conv1d_padding)
+ # position-wise
+ self.w_2 = nn.Conv1d(
+ d_hid, d_in, kernel_size=hp.fft_conv1d_kernel, padding=hp.fft_conv1d_padding)
+
+ self.layer_norm = nn.LayerNorm(d_in)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x):
+ residual = x
+ output = x.transpose(1, 2)
+ output = self.w_2(F.relu(self.w_1(output)))
+ output = output.transpose(1, 2)
+ output = self.dropout(output)
+ output = self.layer_norm(output + residual)
+
+ return output
diff --git a/FastSpeech/transformer/__init__.py b/FastSpeech/transformer/__init__.py
new file mode 100644
index 0000000..068f707
--- /dev/null
+++ b/FastSpeech/transformer/__init__.py
@@ -0,0 +1,6 @@
+import transformer.Constants
+import transformer.Modules
+import transformer.Layers
+import transformer.SubLayers
+import transformer.Models
+import transformer.Beam
diff --git a/FastSpeech/utils.py b/FastSpeech/utils.py
new file mode 100644
index 0000000..25eada1
--- /dev/null
+++ b/FastSpeech/utils.py
@@ -0,0 +1,183 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+import os
+
+import tacotron2 as Tacotron2
+import text
+import hparams
+
+
+def process_text(train_text_path):
+ with open(train_text_path, "r", encoding="utf-8") as f:
+ txt = []
+ for line in f.readlines():
+ txt.append(line)
+
+ return txt
+
+
+def get_param_num(model):
+ num_param = sum(param.numel() for param in model.parameters())
+ return num_param
+
+
+def plot_data(data, figsize=(12, 4)):
+ _, axes = plt.subplots(1, len(data), figsize=figsize)
+ for i in range(len(data)):
+ axes[i].imshow(data[i], aspect='auto',
+ origin='bottom', interpolation='none')
+
+ if not os.path.exists("img"):
+ os.mkdir("img")
+ plt.savefig(os.path.join("img", "model_test.jpg"))
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+ if max_len == None:
+ max_len = torch.max(lengths).item()
+
+ ids = torch.arange(0, max_len, out=torch.cuda.LongTensor(max_len))
+ mask = (ids < lengths.unsqueeze(1)).byte()
+
+ return mask
+
+
+def get_WaveGlow():
+ waveglow_path = os.path.join("waveglow", "pretrained_model")
+ waveglow_path = os.path.join(waveglow_path, "waveglow_256channels.pt")
+ wave_glow = torch.load(waveglow_path)['model']
+ wave_glow = wave_glow.remove_weightnorm(wave_glow)
+ wave_glow.cuda().eval()
+ for m in wave_glow.modules():
+ if 'Conv' in str(type(m)):
+ setattr(m, 'padding_mode', 'zeros')
+
+ return wave_glow
+
+
+def get_Tacotron2():
+ checkpoint_path = "tacotron2_statedict.pt"
+ checkpoint_path = os.path.join(os.path.join(
+ "Tacotron2", "pretrained_model"), checkpoint_path)
+
+ model = Tacotron2.model.Tacotron2(
+ Tacotron2.hparams.create_hparams()).cuda()
+ model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
+ _ = model.cuda().eval()
+
+ return model
+
+
+def get_D(alignment):
+ D = np.array([0 for _ in range(np.shape(alignment)[1])])
+
+ for i in range(np.shape(alignment)[0]):
+ max_index = alignment[i].tolist().index(alignment[i].max())
+ D[max_index] = D[max_index] + 1
+
+ return D
+
+
+def pad_1D(inputs, PAD=0):
+
+ def pad_data(x, length, PAD):
+ x_padded = np.pad(x, (0, length - x.shape[0]),
+ mode='constant',
+ constant_values=PAD)
+ return x_padded
+
+ max_len = max((len(x) for x in inputs))
+ padded = np.stack([pad_data(x, max_len, PAD) for x in inputs])
+
+ return padded
+
+
+def pad_2D(inputs, maxlen=None):
+
+ def pad(x, max_len):
+ PAD = 0
+ if np.shape(x)[0] > max_len:
+ raise ValueError("not max_len")
+
+ s = np.shape(x)[1]
+ x_padded = np.pad(x, (0, max_len - np.shape(x)[0]),
+ mode='constant',
+ constant_values=PAD)
+ return x_padded[:, :s]
+
+ if maxlen:
+ output = np.stack([pad(x, maxlen) for x in inputs])
+ else:
+ max_len = max(np.shape(x)[0] for x in inputs)
+ output = np.stack([pad(x, max_len) for x in inputs])
+
+ return output
+
+
+def pad(input_ele, mel_max_length=None):
+ if mel_max_length:
+ out_list = list()
+ max_len = mel_max_length
+ for i, batch in enumerate(input_ele):
+ one_batch_padded = F.pad(
+ batch, (0, 0, 0, max_len-batch.size(0)), "constant", 0.0)
+ out_list.append(one_batch_padded)
+ out_padded = torch.stack(out_list)
+ return out_padded
+ else:
+ out_list = list()
+ max_len = max([input_ele[i].size(0)for i in range(len(input_ele))])
+
+ for i, batch in enumerate(input_ele):
+ one_batch_padded = F.pad(
+ batch, (0, 0, 0, max_len-batch.size(0)), "constant", 0.0)
+ out_list.append(one_batch_padded)
+ out_padded = torch.stack(out_list)
+ return out_padded
+
+
+def load_data(txt, mel, model):
+ character = text.text_to_sequence(txt, hparams.text_cleaners)
+ character = torch.from_numpy(np.stack([np.array(character)])).long().cuda()
+
+ text_length = torch.Tensor([character.size(1)]).long().cuda()
+ mel = torch.from_numpy(np.stack([mel.T])).float().cuda()
+ max_len = mel.size(2)
+ output_length = torch.Tensor([max_len]).long().cuda()
+
+ inputs = character, text_length, mel, max_len, output_length
+
+ with torch.no_grad():
+ [_, mel_tacotron2, _, alignment], cemb = model.forward(inputs)
+
+ alignment = alignment[0].cpu().numpy()
+ cemb = cemb[0].cpu().numpy()
+
+ D = get_D(alignment)
+ D = np.array(D)
+
+ mel_tacotron2 = mel_tacotron2[0].cpu().numpy()
+
+ return mel_tacotron2, cemb, D
+
+
+def load_data_from_tacotron2(txt, model):
+ character = text.text_to_sequence(txt, hparams.text_cleaners)
+ character = torch.from_numpy(np.stack([np.array(character)])).long().cuda()
+
+ with torch.no_grad():
+ [_, mel, _, alignment], cemb = model.inference(character)
+
+ alignment = alignment[0].cpu().numpy()
+ cemb = cemb[0].cpu().numpy()
+
+ D = get_D(alignment)
+ D = np.array(D)
+
+ mel = mel[0].cpu().numpy()
+
+ return mel, cemb, D
diff --git a/FastSpeech/waveglow/__init__.py b/FastSpeech/waveglow/__init__.py
new file mode 100644
index 0000000..2da378a
--- /dev/null
+++ b/FastSpeech/waveglow/__init__.py
@@ -0,0 +1,3 @@
+import waveglow.inference
+import waveglow.mel2samp
+import waveglow.glow
diff --git a/FastSpeech/waveglow/convert_model.py b/FastSpeech/waveglow/convert_model.py
new file mode 100644
index 0000000..2485e7e
--- /dev/null
+++ b/FastSpeech/waveglow/convert_model.py
@@ -0,0 +1,46 @@
+import sys
+import copy
+import torch
+
+def _check_model_old_version(model):
+ if hasattr(model.WN[0], 'res_layers'):
+ return True
+ else:
+ return False
+
+def update_model(old_model):
+ if not _check_model_old_version(old_model):
+ return old_model
+ new_model = copy.deepcopy(old_model)
+ for idx in range(0, len(new_model.WN)):
+ wavenet = new_model.WN[idx]
+ wavenet.res_skip_layers = torch.nn.ModuleList()
+ n_channels = wavenet.n_channels
+ n_layers = wavenet.n_layers
+ for i in range(0, n_layers):
+ if i < n_layers - 1:
+ res_skip_channels = 2*n_channels
+ else:
+ res_skip_channels = n_channels
+ res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
+ skip_layer = torch.nn.utils.remove_weight_norm(wavenet.skip_layers[i])
+ if i < n_layers - 1:
+ res_layer = torch.nn.utils.remove_weight_norm(wavenet.res_layers[i])
+ res_skip_layer.weight = torch.nn.Parameter(torch.cat([res_layer.weight, skip_layer.weight]))
+ res_skip_layer.bias = torch.nn.Parameter(torch.cat([res_layer.bias, skip_layer.bias]))
+ else:
+ res_skip_layer.weight = torch.nn.Parameter(skip_layer.weight)
+ res_skip_layer.bias = torch.nn.Parameter(skip_layer.bias)
+ res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
+ wavenet.res_skip_layers.append(res_skip_layer)
+ del wavenet.res_layers
+ del wavenet.skip_layers
+ return new_model
+
+if __name__ == '__main__':
+ old_model_path = sys.argv[1]
+ new_model_path = sys.argv[2]
+ model = torch.load(old_model_path)
+ model['model'] = update_model(model['model'])
+ torch.save(model, new_model_path)
+
diff --git a/FastSpeech/waveglow/glow.py b/FastSpeech/waveglow/glow.py
new file mode 100644
index 0000000..a04afed
--- /dev/null
+++ b/FastSpeech/waveglow/glow.py
@@ -0,0 +1,310 @@
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+import copy
+import torch
+from torch.autograd import Variable
+import torch.nn.functional as F
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+ n_channels_int = n_channels[0]
+ in_act = input_a+input_b
+ t_act = torch.tanh(in_act[:, :n_channels_int, :])
+ s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+ acts = t_act * s_act
+ return acts
+
+
+class WaveGlowLoss(torch.nn.Module):
+ def __init__(self, sigma=1.0):
+ super(WaveGlowLoss, self).__init__()
+ self.sigma = sigma
+
+ def forward(self, model_output):
+ z, log_s_list, log_det_W_list = model_output
+ for i, log_s in enumerate(log_s_list):
+ if i == 0:
+ log_s_total = torch.sum(log_s)
+ log_det_W_total = log_det_W_list[i]
+ else:
+ log_s_total = log_s_total + torch.sum(log_s)
+ log_det_W_total += log_det_W_list[i]
+
+ loss = torch.sum(z*z)/(2*self.sigma*self.sigma) - log_s_total - log_det_W_total
+ return loss/(z.size(0)*z.size(1)*z.size(2))
+
+
+class Invertible1x1Conv(torch.nn.Module):
+ """
+ The layer outputs both the convolution, and the log determinant
+ of its weight matrix. If reverse=True it does convolution with
+ inverse
+ """
+ def __init__(self, c):
+ super(Invertible1x1Conv, self).__init__()
+ self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
+ bias=False)
+
+ # Sample a random orthonormal matrix to initialize weights
+ W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
+
+ # Ensure determinant is 1.0 not -1.0
+ if torch.det(W) < 0:
+ W[:,0] = -1*W[:,0]
+ W = W.view(c, c, 1)
+ self.conv.weight.data = W
+
+ def forward(self, z, reverse=False):
+ # shape
+ batch_size, group_size, n_of_groups = z.size()
+
+ W = self.conv.weight.squeeze()
+
+ if reverse:
+ if not hasattr(self, 'W_inverse'):
+ # Reverse computation
+ W_inverse = W.float().inverse()
+ W_inverse = Variable(W_inverse[..., None])
+ if z.type() == 'torch.cuda.HalfTensor':
+ W_inverse = W_inverse.half()
+ self.W_inverse = W_inverse
+ z = F.conv1d(z, self.W_inverse, bias=None, stride=1, padding=0)
+ return z
+ else:
+ # Forward computation
+ log_det_W = batch_size * n_of_groups * torch.logdet(W)
+ z = self.conv(z)
+ return z, log_det_W
+
+
+class WN(torch.nn.Module):
+ """
+ This is the WaveNet like layer for the affine coupling. The primary difference
+ from WaveNet is the convolutions need not be causal. There is also no dilation
+ size reset. The dilation only doubles on each layer
+ """
+ def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels,
+ kernel_size):
+ super(WN, self).__init__()
+ assert(kernel_size % 2 == 1)
+ assert(n_channels % 2 == 0)
+ self.n_layers = n_layers
+ self.n_channels = n_channels
+ self.in_layers = torch.nn.ModuleList()
+ self.res_skip_layers = torch.nn.ModuleList()
+ self.cond_layers = torch.nn.ModuleList()
+
+ start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
+ start = torch.nn.utils.weight_norm(start, name='weight')
+ self.start = start
+
+ # Initializing last layer to 0 makes the affine coupling layers
+ # do nothing at first. This helps with training stability
+ end = torch.nn.Conv1d(n_channels, 2*n_in_channels, 1)
+ end.weight.data.zero_()
+ end.bias.data.zero_()
+ self.end = end
+
+ for i in range(n_layers):
+ dilation = 2 ** i
+ padding = int((kernel_size*dilation - dilation)/2)
+ in_layer = torch.nn.Conv1d(n_channels, 2*n_channels, kernel_size,
+ dilation=dilation, padding=padding)
+ in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
+ self.in_layers.append(in_layer)
+
+ cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels, 1)
+ cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+ self.cond_layers.append(cond_layer)
+
+ # last one is not necessary
+ if i < n_layers - 1:
+ res_skip_channels = 2*n_channels
+ else:
+ res_skip_channels = n_channels
+ res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
+ res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
+ self.res_skip_layers.append(res_skip_layer)
+
+ def forward(self, forward_input):
+ audio, spect = forward_input
+ audio = self.start(audio)
+
+ for i in range(self.n_layers):
+ acts = fused_add_tanh_sigmoid_multiply(
+ self.in_layers[i](audio),
+ self.cond_layers[i](spect),
+ torch.IntTensor([self.n_channels]))
+
+ res_skip_acts = self.res_skip_layers[i](acts)
+ if i < self.n_layers - 1:
+ audio = res_skip_acts[:,:self.n_channels,:] + audio
+ skip_acts = res_skip_acts[:,self.n_channels:,:]
+ else:
+ skip_acts = res_skip_acts
+
+ if i == 0:
+ output = skip_acts
+ else:
+ output = skip_acts + output
+ return self.end(output)
+
+
+class WaveGlow(torch.nn.Module):
+ def __init__(self, n_mel_channels, n_flows, n_group, n_early_every,
+ n_early_size, WN_config):
+ super(WaveGlow, self).__init__()
+
+ self.upsample = torch.nn.ConvTranspose1d(n_mel_channels,
+ n_mel_channels,
+ 1024, stride=256)
+ assert(n_group % 2 == 0)
+ self.n_flows = n_flows
+ self.n_group = n_group
+ self.n_early_every = n_early_every
+ self.n_early_size = n_early_size
+ self.WN = torch.nn.ModuleList()
+ self.convinv = torch.nn.ModuleList()
+
+ n_half = int(n_group/2)
+
+ # Set up layers with the right sizes based on how many dimensions
+ # have been output already
+ n_remaining_channels = n_group
+ for k in range(n_flows):
+ if k % self.n_early_every == 0 and k > 0:
+ n_half = n_half - int(self.n_early_size/2)
+ n_remaining_channels = n_remaining_channels - self.n_early_size
+ self.convinv.append(Invertible1x1Conv(n_remaining_channels))
+ self.WN.append(WN(n_half, n_mel_channels*n_group, **WN_config))
+ self.n_remaining_channels = n_remaining_channels # Useful during inference
+
+ def forward(self, forward_input):
+ """
+ forward_input[0] = mel_spectrogram: batch x n_mel_channels x frames
+ forward_input[1] = audio: batch x time
+ """
+ spect, audio = forward_input
+
+ # Upsample spectrogram to size of audio
+ spect = self.upsample(spect)
+ assert(spect.size(2) >= audio.size(1))
+ if spect.size(2) > audio.size(1):
+ spect = spect[:, :, :audio.size(1)]
+
+ spect = spect.unfold(2, self.n_group, self.n_group).permute(0, 2, 1, 3)
+ spect = spect.contiguous().view(spect.size(0), spect.size(1), -1).permute(0, 2, 1)
+
+ audio = audio.unfold(1, self.n_group, self.n_group).permute(0, 2, 1)
+ output_audio = []
+ log_s_list = []
+ log_det_W_list = []
+
+ for k in range(self.n_flows):
+ if k % self.n_early_every == 0 and k > 0:
+ output_audio.append(audio[:,:self.n_early_size,:])
+ audio = audio[:,self.n_early_size:,:]
+
+ audio, log_det_W = self.convinv[k](audio)
+ log_det_W_list.append(log_det_W)
+
+ n_half = int(audio.size(1)/2)
+ audio_0 = audio[:,:n_half,:]
+ audio_1 = audio[:,n_half:,:]
+
+ output = self.WN[k]((audio_0, spect))
+ log_s = output[:, n_half:, :]
+ b = output[:, :n_half, :]
+ audio_1 = torch.exp(log_s)*audio_1 + b
+ log_s_list.append(log_s)
+
+ audio = torch.cat([audio_0, audio_1],1)
+
+ output_audio.append(audio)
+ return torch.cat(output_audio,1), log_s_list, log_det_W_list
+
+ def infer(self, spect, sigma=1.0):
+ spect = self.upsample(spect)
+ # trim conv artifacts. maybe pad spec to kernel multiple
+ time_cutoff = self.upsample.kernel_size[0] - self.upsample.stride[0]
+ spect = spect[:, :, :-time_cutoff]
+
+ spect = spect.unfold(2, self.n_group, self.n_group).permute(0, 2, 1, 3)
+ spect = spect.contiguous().view(spect.size(0), spect.size(1), -1).permute(0, 2, 1)
+
+ if spect.type() == 'torch.cuda.HalfTensor':
+ audio = torch.cuda.HalfTensor(spect.size(0),
+ self.n_remaining_channels,
+ spect.size(2)).normal_()
+ else:
+ audio = torch.cuda.FloatTensor(spect.size(0),
+ self.n_remaining_channels,
+ spect.size(2)).normal_()
+
+ audio = torch.autograd.Variable(sigma*audio)
+
+ for k in reversed(range(self.n_flows)):
+ n_half = int(audio.size(1)/2)
+ audio_0 = audio[:,:n_half,:]
+ audio_1 = audio[:,n_half:,:]
+
+ output = self.WN[k]((audio_0, spect))
+ s = output[:, n_half:, :]
+ b = output[:, :n_half, :]
+ audio_1 = (audio_1 - b)/torch.exp(s)
+ audio = torch.cat([audio_0, audio_1],1)
+
+ audio = self.convinv[k](audio, reverse=True)
+
+ if k % self.n_early_every == 0 and k > 0:
+ if spect.type() == 'torch.cuda.HalfTensor':
+ z = torch.cuda.HalfTensor(spect.size(0), self.n_early_size, spect.size(2)).normal_()
+ else:
+ z = torch.cuda.FloatTensor(spect.size(0), self.n_early_size, spect.size(2)).normal_()
+ audio = torch.cat((sigma*z, audio),1)
+
+ audio = audio.permute(0,2,1).contiguous().view(audio.size(0), -1).data
+ return audio
+
+ @staticmethod
+ def remove_weightnorm(model):
+ waveglow = model
+ for WN in waveglow.WN:
+ WN.start = torch.nn.utils.remove_weight_norm(WN.start)
+ WN.in_layers = remove(WN.in_layers)
+ WN.cond_layers = remove(WN.cond_layers)
+ WN.res_skip_layers = remove(WN.res_skip_layers)
+ return waveglow
+
+
+def remove(conv_list):
+ new_conv_list = torch.nn.ModuleList()
+ for old_conv in conv_list:
+ old_conv = torch.nn.utils.remove_weight_norm(old_conv)
+ new_conv_list.append(old_conv)
+ return new_conv_list
diff --git a/FastSpeech/waveglow/inference.py b/FastSpeech/waveglow/inference.py
new file mode 100644
index 0000000..0c07ac4
--- /dev/null
+++ b/FastSpeech/waveglow/inference.py
@@ -0,0 +1,57 @@
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+import os
+from scipy.io.wavfile import write
+import torch
+from waveglow.mel2samp import files_to_list, MAX_WAV_VALUE
+# from denoiser import Denoiser
+
+
+def inference(mel, waveglow, audio_path, sigma=1.0, sampling_rate=22050):
+ with torch.no_grad():
+ audio = waveglow.infer(mel, sigma=sigma)
+ audio = audio * MAX_WAV_VALUE
+ audio = audio.squeeze()
+ audio = audio.cpu().numpy()
+ audio = audio.astype('int16')
+ write(audio_path, sampling_rate, audio)
+
+
+def test_speed(mel, waveglow, sigma=1.0, sampling_rate=22050):
+ with torch.no_grad():
+ audio = waveglow.infer(mel, sigma=sigma)
+ audio = audio * MAX_WAV_VALUE
+
+
+def get_wav(mel, waveglow, sigma=1.0, sampling_rate=22050):
+ with torch.no_grad():
+ audio = waveglow.infer(mel, sigma=sigma)
+ audio = audio * MAX_WAV_VALUE
+ audio = audio.squeeze()
+ audio = audio.cpu()
+
+ return audio
diff --git a/FastSpeech/waveglow/mel2samp.py b/FastSpeech/waveglow/mel2samp.py
new file mode 100644
index 0000000..df3e45a
--- /dev/null
+++ b/FastSpeech/waveglow/mel2samp.py
@@ -0,0 +1,147 @@
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************\
+# from tacotron2.layers import TacotronSTFT
+import os
+import random
+import argparse
+import json
+import torch
+import torch.utils.data
+import sys
+from scipy.io.wavfile import read
+
+# We're using the audio processing from TacoTron2 to make sure it matches
+sys.path.insert(0, 'tacotron2')
+
+MAX_WAV_VALUE = 32768.0
+
+
+def files_to_list(filename):
+ """
+ Takes a text file of filenames and makes a list of filenames
+ """
+ with open(filename, encoding='utf-8') as f:
+ files = f.readlines()
+
+ files = [f.rstrip() for f in files]
+ return files
+
+
+# def load_wav_to_torch(full_path):
+# """
+# Loads wavdata into torch array
+# """
+# sampling_rate, data = read(full_path)
+# return torch.from_numpy(data).float(), sampling_rate
+
+
+# class Mel2Samp(torch.utils.data.Dataset):
+# """
+# This is the main class that calculates the spectrogram and returns the
+# spectrogram, audio pair.
+# """
+
+# def __init__(self, training_files, segment_length, filter_length,
+# hop_length, win_length, sampling_rate, mel_fmin, mel_fmax):
+# self.audio_files = files_to_list(training_files)
+# random.seed(1234)
+# random.shuffle(self.audio_files)
+# self.stft = TacotronSTFT(filter_length=filter_length,
+# hop_length=hop_length,
+# win_length=win_length,
+# sampling_rate=sampling_rate,
+# mel_fmin=mel_fmin, mel_fmax=mel_fmax)
+# self.segment_length = segment_length
+# self.sampling_rate = sampling_rate
+
+# def get_mel(self, audio):
+# audio_norm = audio / MAX_WAV_VALUE
+# audio_norm = audio_norm.unsqueeze(0)
+# audio_norm = torch.autograd.Variable(audio_norm, requires_grad=False)
+# melspec = self.stft.mel_spectrogram(audio_norm)
+# melspec = torch.squeeze(melspec, 0)
+# return melspec
+
+# def __getitem__(self, index):
+# # Read audio
+# filename = self.audio_files[index]
+# audio, sampling_rate = load_wav_to_torch(filename)
+# if sampling_rate != self.sampling_rate:
+# raise ValueError("{} SR doesn't match target {} SR".format(
+# sampling_rate, self.sampling_rate))
+
+# # Take segment
+# if audio.size(0) >= self.segment_length:
+# max_audio_start = audio.size(0) - self.segment_length
+# audio_start = random.randint(0, max_audio_start)
+# audio = audio[audio_start:audio_start+self.segment_length]
+# else:
+# audio = torch.nn.functional.pad(
+# audio, (0, self.segment_length - audio.size(0)), 'constant').data
+
+# mel = self.get_mel(audio)
+# audio = audio / MAX_WAV_VALUE
+
+# return (mel, audio)
+
+# def __len__(self):
+# return len(self.audio_files)
+
+
+# # ===================================================================
+# # Takes directory of clean audio and makes directory of spectrograms
+# # Useful for making test sets
+# # ===================================================================
+# if __name__ == "__main__":
+# # Get defaults so it can work with no Sacred
+# parser = argparse.ArgumentParser()
+# parser.add_argument('-f', "--filelist_path", required=True)
+# parser.add_argument('-c', '--config', type=str,
+# help='JSON file for configuration')
+# parser.add_argument('-o', '--output_dir', type=str,
+# help='Output directory')
+# args = parser.parse_args()
+
+# with open(args.config) as f:
+# data = f.read()
+# data_config = json.loads(data)["data_config"]
+# mel2samp = Mel2Samp(**data_config)
+
+# filepaths = files_to_list(args.filelist_path)
+
+# # Make directory if it doesn't exist
+# if not os.path.isdir(args.output_dir):
+# os.makedirs(args.output_dir)
+# os.chmod(args.output_dir, 0o775)
+
+# for filepath in filepaths:
+# audio, sr = load_wav_to_torch(filepath)
+# melspectrogram = mel2samp.get_mel(audio)
+# filename = os.path.basename(filepath)
+# new_filepath = args.output_dir + '/' + filename + '.pt'
+# print(new_filepath)
+# torch.save(melspectrogram, new_filepath)
diff --git a/SqueezeWave/README.md b/SqueezeWave/README.md
new file mode 100644
index 0000000..a09bbc0
--- /dev/null
+++ b/SqueezeWave/README.md
@@ -0,0 +1,129 @@
+## SqueezeWave: Extremely Lightweight Vocoders for On-device Speech Synthesis
+By Bohan Zhai *, Tianren Gao *, Flora Xue, Daniel Rothchild, Bichen Wu, Joseph Gonzalez, and Kurt Keutzer (UC Berkeley)
+
+Automatic speech synthesis is a challenging task that is becoming increasingly important as edge devices begin to interact with users through speech. Typical text-to-speech pipelines include a vocoder, which translates intermediate audio representations into an audio waveform. Most existing vocoders are difficult to parallelize since each generated sample is conditioned on previous samples. WaveGlow is a flow-based feed-forward alternative to these auto-regressive models (Prenger et al., 2019). However, while WaveGlow can be easily parallelized, the model is too expensive for real-time speech synthesis on the edge. This paper presents SqueezeWave, a family of lightweight vocoders based on WaveGlow that can generate audio of similar quality to WaveGlow with 61x - 214x fewer MACs.
+
+Link to the paper: [paper]. If you find this work useful, please consider citing
+
+ ```
+ @inproceedings{squeezewave,
+ Author = {Bohan Zhai, Tianren Gao, Flora Xue, Daniel Rothchild, Bichen Wu, Joseph Gonzalez, Kurt Keutzer},
+ Title = {SqueezeWave: Extremely Lightweight Vocoders for On-device Speech Synthesis},
+ Journal = {arXiv:2001.05685},
+ Year = {2020}
+ }
+ ```
+
+### Audio samples generated by SqueezeWave
+Audio samples of SqueezeWave are here: https://tianrengao.github.io/SqueezeWaveDemo/
+
+### Results
+We introduce 4 variants of SqueezeWave in our paper. See the table below.
+
+
+ | Model | length | n_channels| MACs | Reduction | MOS |
+ | --------------- | ------ | --------- | ----- | --------- | --------- |
+ |WaveGlow | 2048 | 8 | 228.9 | 1x | 4.57±0.04 |
+ |SqueezeWave-128L | 128 | 256 | 3.78 | 60x | 4.07±0.06 |
+ |SqueezeWave-64L | 64 | 256 | 2.16 | 106x | 3.77±0.05 |
+ |SqueezeWave-128S | 128 | 128 | 1.06 | 214x | 3.79±0.05 |
+ |SqueezeWave-64S | 64 | 128 | 0.68 | 332x | 2.74±0.04 |
+
+### Model Complexity
+A detailed MAC calculation can be found from [here](https://github.com/tianrengao/SqueezeWave/blob/master/SqueezeWave_computational_complexity.ipynb)
+
+## Setup
+0. (Optional) Create a virtual environment
+
+ ```
+ virtualenv env
+ source env/bin/activate
+ ```
+
+1. Clone our repo and initialize submodule
+
+ ```command
+ git clone https://github.com/tianrengao/SqueezeWave.git
+ cd SqueezeWave
+ git submodule init
+ git submodule update
+ ```
+
+2. Install requirements
+```pip3 install -r requirements.txt```
+
+3. Install [Apex]
+ ```1
+ cd ../
+ git clone https://www.github.com/nvidia/apex
+ cd apex
+ python setup.py install
+ ```
+
+## Generate audio with our pretrained model
+
+1. Download our [pretrained models]. We provide 4 pretrained models as described in the paper.
+2. Download [mel-spectrograms]
+3. Generate audio. Please replace `SqueezeWave.pt` to the specific pretrained model's name.
+
+ ```python3 inference.py -f <(ls mel_spectrograms/*.pt) -w SqueezeWave.pt -o . --is_fp16 -s 0.6```
+
+
+## Train your own model
+
+1. Download [LJ Speech Data]. We assume all the waves are stored in the directory `^/data/`
+
+2. Make a list of the file names to use for training/testing
+
+ ```command
+ ls data/*.wav | tail -n+10 > train_files.txt
+ ls data/*.wav | head -n10 > test_files.txt
+ ```
+
+3. We provide 4 model configurations with audio channel and channel numbers specified in the table below. The configuration files are under ```/configs``` directory. To choose the model you want to train, select the corresponding configuration file.
+
+4. Train your SqueezeWave model
+
+ ```command
+ mkdir checkpoints
+ python train.py -c configs/config_a256_c128.json
+ ```
+
+ For multi-GPU training replace `train.py` with `distributed.py`. Only tested with single node and NCCL.
+
+ For mixed precision training set `"fp16_run": true` on `config.json`.
+
+5. Make test set mel-spectrograms
+
+ ```
+ mkdir -p eval/mels
+ python3 mel2samp.py -f test_files.txt -o eval/mels -c configs/config_a128_c256.json
+ ```
+
+6. Run inference on the test data.
+
+ ```command
+ ls eval/mels > eval/mel_files.txt
+ sed -i -e 's_.*_eval/mels/&_' eval/mel_files.txt
+ mkdir -p eval/output
+ python3 inference.py -f eval/mel_files.txt -w checkpoints/SqueezeWave_10000 -o eval/output --is_fp16 -s 0.6
+ ```
+ Replace `SqueezeWave_10000` with the checkpoint you want to test.
+
+## Credits
+The implementation of this work is based on WaveGlow: https://github.com/NVIDIA/waveglow
+
+
+[//]: # (TODO)
+[//]: # (PROVIDE INSTRUCTIONS FOR DOWNLOADING LJS)
+[pytorch 1.0]: https://github.com/pytorch/pytorch#installation
+[website]: https://nv-adlr.github.io/WaveGlow
+[paper]: https://arxiv.org/abs/2001.05685
+[WaveNet implementation]: https://github.com/r9y9/wavenet_vocoder
+[Glow]: https://blog.openai.com/glow/
+[WaveNet]: https://deepmind.com/blog/wavenet-generative-model-raw-audio/
+[PyTorch]: http://pytorch.org
+[pretrained models]: https://drive.google.com/file/d/1RyVMLY2l8JJGq_dCEAAd8rIRIn_k13UB/view?usp=sharing
+[mel-spectrograms]: https://drive.google.com/file/d/1g_VXK2lpP9J25dQFhQwx7doWl_p20fXA/view?usp=sharing
+[LJ Speech Data]: https://keithito.com/LJ-Speech-Dataset
+[Apex]: https://github.com/nvidia/apex
diff --git a/SqueezeWave/SqueezeWave_computational_complexity.ipynb b/SqueezeWave/SqueezeWave_computational_complexity.ipynb
new file mode 100644
index 0000000..15865c2
--- /dev/null
+++ b/SqueezeWave/SqueezeWave_computational_complexity.ipynb
@@ -0,0 +1,445 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+ "colab": {
+ "name": "SqueezeWave computational complexity.ipynb",
+ "provenance": []
+ },
+ "kernelspec": {
+ "name": "python2",
+ "display_name": "Python 2"
+ }
+ },
+ "cells": [
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "s8VYGy15fwqN",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "import numpy as np"
+ ],
+ "execution_count": 0,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "MDp5WalGf5Ji",
+ "colab_type": "text"
+ },
+ "source": [
+ "**WaveGlow**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "wrBBjKSYf89M",
+ "colab_type": "code",
+ "outputId": "4d77bc19-7a81-4f0b-bcad-65c42c4b2e9c",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 136
+ }
+ },
+ "source": [
+ "L = 2048 # audio length\n",
+ "n_audio_channel_init = 8 # initial audio channel \n",
+ "C_mel = 80 * 8 # After upsampling and unfolding \n",
+ "kernal_size = 3\n",
+ "C_wn = 256 # input channel size of in_layer\n",
+ "C_wn_middle = C_wn * 2 # output channel size of in_layer and cond_layer\n",
+ "n_flows = 12\n",
+ "n_layers = 8\n",
+ "n_early_output = 2\n",
+ "n_early_output_interval = 4\n",
+ "duration = 0.725\n",
+ "\n",
+ "n_audio_channels = []\n",
+ "n_audio = n_audio_channel_init\n",
+ "for i in range(n_flows):\n",
+ " if i % n_early_output_interval == 0 and i > 0:\n",
+ " n_audio -= n_early_output\n",
+ " n_audio_channels.append(n_audio) # audio channel after early output\n",
+ "\n",
+ "# in_layers\n",
+ "WN_in_layers = L * kernal_size * C_wn * C_wn_middle * n_layers * n_flows\n",
+ "print('MACs of in_layers', WN_in_layers / duration / 1e9)\n",
+ "# cond layers\n",
+ "WN_cond_layers = L * C_mel * C_wn_middle * n_layers * n_flows \n",
+ "print('MACs of cond_layers', WN_cond_layers / duration / 1e9)\n",
+ "# res skip layers\n",
+ "WN_res_layers = (L * C_wn * C_wn_middle * (n_layers - 1) + L * C_wn * C_wn) * n_flows\n",
+ "print('MACs of res_skip_layers', WN_res_layers / duration / 1e9)\n",
+ "# invertible convs\n",
+ "inv1x1 = np.sum([n**2 * L for n in n_audio_channels])\n",
+ "print('MACs of invertible conv layers', inv1x1 / duration / 1e9)\n",
+ "# start\n",
+ "starts = np.sum([n / 2 * C_wn * L for n in n_audio_channels])\n",
+ "print('MACs of start conv layers', starts / duration / 1e9)\n",
+ "# end\n",
+ "ends = np.sum([C_wn * n * L for n in n_audio_channels])\n",
+ "print('MACs of end conv layers', ends / duration / 1e9)\n",
+ "# total\n",
+ "WG_total = WN_in_layers + WN_cond_layers + WN_res_layers + inv1x1 + starts + ends\n",
+ "print('Total number of MACs is', WG_total / duration / 1e9)"
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "('MACs of in_layers', 106.63367079724138)\n",
+ "('MACs of cond_layers', 88.86139233103448)\n",
+ "('MACs of res_skip_layers', 33.32302212413793)\n",
+ "('MACs of invertible conv layers', 0.00131072)\n",
+ "('MACs of start conv layers', 0.02603361103448276)\n",
+ "('MACs of end conv layers', 0.05206722206896552)\n",
+ "('Total number of MACs is', 228.89749680551725)\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "QRQheCWjgC9D",
+ "colab_type": "text"
+ },
+ "source": [
+ "SqueezeWave L=64, C=128"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "zSlwPlvUgJue",
+ "colab_type": "code",
+ "outputId": "18e282ea-a071-4117-ba08-6e6abdc36c68",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 153
+ }
+ },
+ "source": [
+ "L = 64 # audio length\n",
+ "n_audio_channel_init = 256 # initial audio channel \n",
+ "L_mel = 64 # mel-spectrogram length\n",
+ "C_mel =80 # mel-spectrogram channel \n",
+ "kernal_size = 3\n",
+ "C_wn = 128 # input channel size of in_layer\n",
+ "C_wn_middle = C_wn * 2 # output channel size of in_layer and cond_layer\n",
+ "n_flows = 12\n",
+ "n_layers = 8\n",
+ "n_early_output = 16\n",
+ "n_early_output_interval = 2\n",
+ "duration = 0.725\n",
+ "\n",
+ "n_audio_channels = []\n",
+ "n_audio = n_audio_channel_init\n",
+ "for i in range(n_flows):\n",
+ " if i % n_early_output_interval == 0 and i > 0:\n",
+ " n_audio -= n_early_output\n",
+ " n_audio_channels.append(n_audio) # audio channel after early output\n",
+ "\n",
+ "# in_layers\n",
+ "WN_in_layers = L * kernal_size * C_wn * n_layers * n_flows # depthwise\n",
+ "WN_in_layers += L * C_wn * C_wn_middle * n_layers * n_flows # pointwise\n",
+ "print('MACs of in_layers', WN_in_layers / duration / 1e9)\n",
+ "# cond_layers\n",
+ "WN_cond_layers = L_mel * C_mel * C_wn_middle * n_layers * n_flows\n",
+ "print('MACs of cond_layers', WN_cond_layers / duration / 1e9)\n",
+ "# res_skip_layers\n",
+ "WN_res_layers = L * C_wn * C_wn * n_layers * n_flows\n",
+ "print('MACs of res_skip_layers', WN_res_layers / duration / 1e9)\n",
+ "# invertible convs\n",
+ "inv1x1 = np.sum([n**2 * L for n in n_audio_channels])\n",
+ "print('MACs of invertible conv layers', inv1x1 / duration / 1e9)\n",
+ "# start\n",
+ "starts = np.sum([n / 2 * C_wn * L for n in n_audio_channels])\n",
+ "print('MACs of start conv layers', starts / duration / 1e9)\n",
+ "#end\n",
+ "ends = np.sum([C_wn * n * L for n in n_audio_channels])\n",
+ "print('MACs of end conv layers', ends / duration / 1e9)\n",
+ "# total\n",
+ "total = WN_in_layers + WN_cond_layers + WN_res_layers + inv1x1 + starts + ends\n",
+ "print('Total number of MACs is', total / duration / 1e9)\n",
+ "print('Reduction compared with WaveGlow', WG_total / total)"
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "('MACs of in_layers', 0.2809460524137931)\n",
+ "('MACs of cond_layers', 0.17355740689655172)\n",
+ "('MACs of res_skip_layers', 0.1388459255172414)\n",
+ "('MACs of invertible conv layers', 0.0502141351724138)\n",
+ "('MACs of start conv layers', 0.014643906206896554)\n",
+ "('MACs of end conv layers', 0.029287812413793107)\n",
+ "('Total number of MACs is', 0.6874952386206896)\n",
+ "('Reduction compared with WaveGlow', 332)\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "M6K8zJ6cugYj",
+ "colab_type": "text"
+ },
+ "source": [
+ "**SqueezeWave L=64, C=256**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "ju5Xa4oAhScO",
+ "colab_type": "code",
+ "outputId": "c91361be-ff73-4113-a584-6dda74c3690e",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 153
+ }
+ },
+ "source": [
+ "L = 64 # audio length\n",
+ "n_audio_channel_init = 256 # initial audio channel \n",
+ "L_mel = 64 # mel-spectrogram length\n",
+ "C_mel =80 # mel-spectrogram channel \n",
+ "kernal_size = 3\n",
+ "C_wn = 256 # input channel size of in_layer\n",
+ "C_wn_middle = C_wn * 2 # output channel size of in_layer and cond_layer\n",
+ "n_flows = 12\n",
+ "n_layers = 8\n",
+ "n_early_output = 16\n",
+ "n_early_output_interval = 2\n",
+ "duration = 0.725\n",
+ "\n",
+ "n_audio_channels = []\n",
+ "n_audio = n_audio_channel_init\n",
+ "for i in range(n_flows):\n",
+ " if i % n_early_output_interval == 0 and i > 0:\n",
+ " n_audio -= n_early_output\n",
+ " n_audio_channels.append(n_audio) # audio channel after early output\n",
+ "\n",
+ "# in_layers\n",
+ "WN_in_layers = L * kernal_size * C_wn * n_layers * n_flows # depthwise\n",
+ "WN_in_layers += L * C_wn * C_wn_middle * n_layers * n_flows # pointwise\n",
+ "print('MACs of in_layers', WN_in_layers / duration / 1e9)\n",
+ "# cond_layers\n",
+ "WN_cond_layers = L_mel * C_mel * C_wn_middle * n_layers * n_flows\n",
+ "print('MACs of cond_layers', WN_cond_layers / duration / 1e9)\n",
+ "# res_skip_layers\n",
+ "WN_res_layers = L * C_wn * C_wn * n_layers * n_flows\n",
+ "print('MACs of res_skip_layers', WN_res_layers / duration / 1e9)\n",
+ "# invertible convs\n",
+ "inv1x1 = np.sum([n**2 * L for n in n_audio_channels])\n",
+ "print('MACs of invertible conv layers', inv1x1 / duration / 1e9)\n",
+ "# start\n",
+ "starts = np.sum([n / 2 * C_wn * L for n in n_audio_channels])\n",
+ "print('MACs of start conv layers', starts / duration / 1e9)\n",
+ "#end\n",
+ "ends = np.sum([C_wn * n * L for n in n_audio_channels])\n",
+ "print('MACs of end conv layers', ends / duration / 1e9)\n",
+ "# total\n",
+ "total = WN_in_layers + WN_cond_layers + WN_res_layers + inv1x1 + starts + ends\n",
+ "print('Total number of MACs is', total / duration / 1e9)\n",
+ "print('Reduction compared with WaveGlow', WG_total / total)"
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "('MACs of in_layers', 1.1172758068965518)\n",
+ "('MACs of cond_layers', 0.34711481379310344)\n",
+ "('MACs of res_skip_layers', 0.5553837020689656)\n",
+ "('MACs of invertible conv layers', 0.0502141351724138)\n",
+ "('MACs of start conv layers', 0.029287812413793107)\n",
+ "('MACs of end conv layers', 0.058575624827586215)\n",
+ "('Total number of MACs is', 2.157851895172414)\n",
+ "('Reduction compared with WaveGlow', 106)\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aIgnX6Yi4BFu",
+ "colab_type": "text"
+ },
+ "source": [
+ "**SqueezeWave L=128, C=128**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "W-3Q5jW84F_t",
+ "colab_type": "code",
+ "outputId": "436038c3-f3f8-4989-eeec-eb59c154b183",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 153
+ }
+ },
+ "source": [
+ "L = 128 # audio length\n",
+ "n_audio_channel_init = 128 # initial audio channel \n",
+ "L_mel = 64 # mel-spectrogram length\n",
+ "C_mel =80 # mel-spectrogram channel \n",
+ "kernal_size = 3\n",
+ "C_wn = 128 # input channel size of in_layer\n",
+ "C_wn_middle = C_wn * 2 # output channel size of in_layer and cond_layer\n",
+ "n_flows = 12\n",
+ "n_layers = 8\n",
+ "n_early_output = 16\n",
+ "n_early_output_interval = 2\n",
+ "duration = 0.725\n",
+ "\n",
+ "n_audio_channels = []\n",
+ "n_audio = n_audio_channel_init\n",
+ "for i in range(n_flows):\n",
+ " if i % n_early_output_interval == 0 and i > 0:\n",
+ " n_audio -= n_early_output\n",
+ " n_audio_channels.append(n_audio) # audio channel after early output\n",
+ "\n",
+ "# in_layers\n",
+ "WN_in_layers = L * kernal_size * C_wn * n_layers * n_flows # depthwise\n",
+ "WN_in_layers += L * C_wn * C_wn_middle * n_layers * n_flows # pointwise\n",
+ "print('MACs of in_layers', WN_in_layers / duration / 1e9)\n",
+ "# cond_layers\n",
+ "WN_cond_layers = L_mel * C_mel * C_wn_middle * n_layers * n_flows\n",
+ "print('MACs of cond_layers', WN_cond_layers / duration / 1e9)\n",
+ "# res_skip_layers\n",
+ "WN_res_layers = L * C_wn * C_wn * n_layers * n_flows\n",
+ "print('MACs of res_skip_layers', WN_res_layers / duration / 1e9)\n",
+ "# invertible convs\n",
+ "inv1x1 = np.sum([n**2 * L for n in n_audio_channels])\n",
+ "print('MACs of invertible conv layers', inv1x1 / duration / 1e9)\n",
+ "# start\n",
+ "starts = np.sum([n / 2 * C_wn * L for n in n_audio_channels])\n",
+ "print('MACs of start conv layers', starts / duration / 1e9)\n",
+ "#end\n",
+ "ends = np.sum([C_wn * n * L for n in n_audio_channels])\n",
+ "print('MACs of end conv layers', ends / duration / 1e9)\n",
+ "# total\n",
+ "total = WN_in_layers + WN_cond_layers + WN_res_layers + inv1x1 + starts + ends\n",
+ "print('Total number of MACs is', total / duration / 1e9)\n",
+ "print('Reduction compared with WaveGlow', WG_total / total)"
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "('MACs of in_layers', 0.5618921048275862)\n",
+ "('MACs of cond_layers', 0.17355740689655172)\n",
+ "('MACs of res_skip_layers', 0.2776918510344828)\n",
+ "('MACs of invertible conv layers', 0.017988502068965517)\n",
+ "('MACs of start conv layers', 0.011932071724137933)\n",
+ "('MACs of end conv layers', 0.023864143448275865)\n",
+ "('Total number of MACs is', 1.06692608)\n",
+ "('Reduction compared with WaveGlow', 214)\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "1kWvIBWU4Vwm",
+ "colab_type": "text"
+ },
+ "source": [
+ "**SqueezeWave L=128, C=256**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "6YM2bkC14WWc",
+ "colab_type": "code",
+ "outputId": "b1fd3d03-0135-400e-cfbc-28746c8d0cf0",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 153
+ }
+ },
+ "source": [
+ "L = 128 # audio length\n",
+ "n_audio_channel_init = 128 # initial audio channel \n",
+ "L_mel = 64 # mel-spectrogram length\n",
+ "C_mel =80 # mel-spectrogram channel \n",
+ "kernal_size = 3\n",
+ "C_wn = 256 # input channel size of in_layer\n",
+ "C_wn_middle = C_wn * 2 # output channel size of in_layer and cond_layer\n",
+ "n_flows = 12\n",
+ "n_layers = 8\n",
+ "n_early_output = 16\n",
+ "n_early_output_interval = 2\n",
+ "duration = 0.725\n",
+ "\n",
+ "n_audio_channels = []\n",
+ "n_audio = n_audio_channel_init\n",
+ "for i in range(n_flows):\n",
+ " if i % n_early_output_interval == 0 and i > 0:\n",
+ " n_audio -= n_early_output\n",
+ " n_audio_channels.append(n_audio) # audio channel after early output\n",
+ "\n",
+ "# in_layers\n",
+ "WN_in_layers = L * kernal_size * C_wn * n_layers * n_flows # depthwise\n",
+ "WN_in_layers += L * C_wn * C_wn_middle * n_layers * n_flows # pointwise\n",
+ "print('MACs of in_layers', WN_in_layers / duration / 1e9)\n",
+ "# cond_layers\n",
+ "WN_cond_layers = L_mel * C_mel * C_wn_middle * n_layers * n_flows\n",
+ "print('MACs of cond_layers', WN_cond_layers / duration / 1e9)\n",
+ "# res_skip_layers\n",
+ "WN_res_layers = L * C_wn * C_wn * n_layers * n_flows\n",
+ "print('MACs of res_skip_layers', WN_res_layers / duration / 1e9)\n",
+ "# invertible convs\n",
+ "inv1x1 = np.sum([n**2 * L for n in n_audio_channels])\n",
+ "print('MACs of invertible conv layers', inv1x1 / duration / 1e9)\n",
+ "# start\n",
+ "starts = np.sum([n / 2 * C_wn * L for n in n_audio_channels])\n",
+ "print('MACs of start conv layers', starts / duration / 1e9)\n",
+ "#end\n",
+ "ends = np.sum([C_wn * n * L for n in n_audio_channels])\n",
+ "print('MACs of end conv layers', ends / duration / 1e9)\n",
+ "# total\n",
+ "total = WN_in_layers + WN_cond_layers + WN_res_layers + inv1x1 + starts + ends\n",
+ "print('Total number of MACs is', total / duration / 1e9)\n",
+ "print('Reduction compared with WaveGlow', WG_total / total)"
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "('MACs of in_layers', 2.2345516137931036)\n",
+ "('MACs of cond_layers', 0.34711481379310344)\n",
+ "('MACs of res_skip_layers', 1.1107674041379312)\n",
+ "('MACs of invertible conv layers', 0.017988502068965517)\n",
+ "('MACs of start conv layers', 0.023864143448275865)\n",
+ "('MACs of end conv layers', 0.04772828689655173)\n",
+ "('Total number of MACs is', 3.7820147641379314)\n",
+ "('Reduction compared with WaveGlow', 60)\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ }
+ ]
+}
\ No newline at end of file
diff --git a/SqueezeWave/TacotronSTFT.py b/SqueezeWave/TacotronSTFT.py
new file mode 100644
index 0000000..ebfe6f8
--- /dev/null
+++ b/SqueezeWave/TacotronSTFT.py
@@ -0,0 +1,80 @@
+import torch
+from librosa.filters import mel as librosa_mel_fn
+from audio_processing import dynamic_range_compression
+from audio_processing import dynamic_range_decompression
+from stft import STFT
+
+
+class LinearNorm(torch.nn.Module):
+ def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+ super(LinearNorm, self).__init__()
+ self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
+
+ torch.nn.init.xavier_uniform(
+ self.linear_layer.weight,
+ gain=torch.nn.init.calculate_gain(w_init_gain))
+
+ def forward(self, x):
+ return self.linear_layer(x)
+
+
+class ConvNorm(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,
+ padding=None, dilation=1, bias=True, w_init_gain='linear'):
+ super(ConvNorm, self).__init__()
+ if padding is None:
+ assert(kernel_size % 2 == 1)
+ padding = int(dilation * (kernel_size - 1) / 2)
+
+ self.conv = torch.nn.Conv1d(in_channels, out_channels,
+ kernel_size=kernel_size, stride=stride,
+ padding=padding, dilation=dilation,
+ bias=bias)
+
+ torch.nn.init.xavier_uniform(
+ self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain))
+
+ def forward(self, signal):
+ conv_signal = self.conv(signal)
+ return conv_signal
+
+
+class TacotronSTFT(torch.nn.Module):
+ def __init__(self, filter_length=1024, hop_length=256, win_length=1024,
+ n_mel_channels=80, sampling_rate=22050, mel_fmin=0.0,
+ mel_fmax=None, n_group=256):
+ super(TacotronSTFT, self).__init__()
+ self.n_mel_channels = n_mel_channels
+ self.sampling_rate = sampling_rate
+ self.stft_fn = STFT(filter_length, hop_length, win_length, n_group=n_group)
+ mel_basis = librosa_mel_fn(
+ sampling_rate, filter_length, n_mel_channels, mel_fmin, mel_fmax)
+ mel_basis = torch.from_numpy(mel_basis).float()
+ self.register_buffer('mel_basis', mel_basis)
+
+ def spectral_normalize(self, magnitudes):
+ output = dynamic_range_compression(magnitudes)
+ return output
+
+ def spectral_de_normalize(self, magnitudes):
+ output = dynamic_range_decompression(magnitudes)
+ return output
+
+ def mel_spectrogram(self, y):
+ """Computes mel-spectrograms from a batch of waves
+ PARAMS
+ ------
+ y: Variable(torch.FloatTensor) with shape (B, T) in range [-1, 1]
+
+ RETURNS
+ -------
+ mel_output: torch.FloatTensor of shape (B, n_mel_channels, T)
+ """
+ assert(torch.min(y.data) >= -1)
+ assert(torch.max(y.data) <= 1)
+
+ magnitudes, phases = self.stft_fn.transform(y)
+ magnitudes = magnitudes.data
+ mel_output = torch.matmul(self.mel_basis, magnitudes)
+ mel_output = self.spectral_normalize(mel_output)
+ return mel_output
diff --git a/SqueezeWave/audio_processing.py b/SqueezeWave/audio_processing.py
new file mode 100644
index 0000000..b5af7f7
--- /dev/null
+++ b/SqueezeWave/audio_processing.py
@@ -0,0 +1,93 @@
+import torch
+import numpy as np
+from scipy.signal import get_window
+import librosa.util as librosa_util
+
+
+def window_sumsquare(window, n_frames, hop_length=200, win_length=800,
+ n_fft=800, dtype=np.float32, norm=None):
+ """
+ # from librosa 0.6
+ Compute the sum-square envelope of a window function at a given hop length.
+
+ This is used to estimate modulation effects induced by windowing
+ observations in short-time fourier transforms.
+
+ Parameters
+ ----------
+ window : string, tuple, number, callable, or list-like
+ Window specification, as in `get_window`
+
+ n_frames : int > 0
+ The number of analysis frames
+
+ hop_length : int > 0
+ The number of samples to advance between frames
+
+ win_length : [optional]
+ The length of the window function. By default, this matches `n_fft`.
+
+ n_fft : int > 0
+ The length of each analysis frame.
+
+ dtype : np.dtype
+ The data type of the output
+
+ Returns
+ -------
+ wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
+ The sum-squared envelope of the window function
+ """
+ if win_length is None:
+ win_length = n_fft
+
+ n = n_fft + hop_length * (n_frames - 1)
+ x = np.zeros(n, dtype=dtype)
+
+ # Compute the squared window at the desired length
+ win_sq = get_window(window, win_length, fftbins=True)
+ win_sq = librosa_util.normalize(win_sq, norm=norm)**2
+ win_sq = librosa_util.pad_center(win_sq, n_fft)
+
+ # Fill the envelope
+ for i in range(n_frames):
+ sample = i * hop_length
+ x[sample:min(n, sample + n_fft)] += win_sq[:max(0, min(n_fft, n - sample))]
+ return x
+
+
+def griffin_lim(magnitudes, stft_fn, n_iters=30):
+ """
+ PARAMS
+ ------
+ magnitudes: spectrogram magnitudes
+ stft_fn: STFT class with transform (STFT) and inverse (ISTFT) methods
+ """
+
+ angles = np.angle(np.exp(2j * np.pi * np.random.rand(*magnitudes.size())))
+ angles = angles.astype(np.float32)
+ angles = torch.autograd.Variable(torch.from_numpy(angles))
+ signal = stft_fn.inverse(magnitudes, angles).squeeze(1)
+
+ for i in range(n_iters):
+ _, angles = stft_fn.transform(signal)
+ signal = stft_fn.inverse(magnitudes, angles).squeeze(1)
+ return signal
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+ """
+ PARAMS
+ ------
+ C: compression factor
+ """
+ return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+ """
+ PARAMS
+ ------
+ C: compression factor used to compress
+ """
+ return torch.exp(x) / C
diff --git a/SqueezeWave/configs/config_a128_c128.json b/SqueezeWave/configs/config_a128_c128.json
new file mode 100644
index 0000000..1e3273c
--- /dev/null
+++ b/SqueezeWave/configs/config_a128_c128.json
@@ -0,0 +1,40 @@
+{
+ "train_config": {
+ "fp16_run": true,
+ "output_directory": "checkpoints",
+ "epochs": 100000,
+ "learning_rate": 4e-4,
+ "sigma": 1.0,
+ "iters_per_checkpoint": 2000,
+ "batch_size": 96,
+ "seed": 1234,
+ "checkpoint_path": "",
+ "with_tensorboard": true
+ },
+ "data_config": {
+ "training_files": "train_files.txt",
+ "segment_length": 16384,
+ "sampling_rate": 22050,
+ "filter_length": 1024,
+ "hop_length": 256,
+ "win_length": 1024,
+ "mel_fmin": 0.0,
+ "mel_fmax": 8000.0
+ },
+ "dist_config": {
+ "dist_backend": "nccl",
+ "dist_url": "tcp://localhost:54321"
+ },
+ "squeezewave_config": {
+ "n_mel_channels": 80,
+ "n_flows": 12,
+ "n_audio_channel": 128,
+ "n_early_every": 2,
+ "n_early_size": 16,
+ "WN_config": {
+ "n_layers": 8,
+ "n_channels": 128,
+ "kernel_size": 3
+ }
+ }
+}
diff --git a/SqueezeWave/configs/config_a128_c256.json b/SqueezeWave/configs/config_a128_c256.json
new file mode 100644
index 0000000..c342590
--- /dev/null
+++ b/SqueezeWave/configs/config_a128_c256.json
@@ -0,0 +1,40 @@
+{
+ "train_config": {
+ "fp16_run": true,
+ "output_directory": "checkpoints",
+ "epochs": 100000,
+ "learning_rate": 4e-4,
+ "sigma": 1.0,
+ "iters_per_checkpoint": 2000,
+ "batch_size": 96,
+ "seed": 1234,
+ "checkpoint_path": "checkpoints/Squeeze_244000",
+ "with_tensorboard": true
+ },
+ "data_config": {
+ "training_files": "train_files.txt",
+ "segment_length": 16384,
+ "sampling_rate": 22050,
+ "filter_length": 1024,
+ "hop_length": 256,
+ "win_length": 1024,
+ "mel_fmin": 0.0,
+ "mel_fmax": 8000.0
+ },
+ "dist_config": {
+ "dist_backend": "nccl",
+ "dist_url": "tcp://localhost:54321"
+ },
+ "squeezewave_config": {
+ "n_mel_channels": 80,
+ "n_flows": 12,
+ "n_audio_channel": 128,
+ "n_early_every": 2,
+ "n_early_size": 16,
+ "WN_config": {
+ "n_layers": 8,
+ "n_channels": 256,
+ "kernel_size": 3
+ }
+ }
+}
diff --git a/SqueezeWave/configs/config_a256_c128.json b/SqueezeWave/configs/config_a256_c128.json
new file mode 100644
index 0000000..29d8f26
--- /dev/null
+++ b/SqueezeWave/configs/config_a256_c128.json
@@ -0,0 +1,40 @@
+{
+ "train_config": {
+ "fp16_run": true,
+ "output_directory": "checkpoints",
+ "epochs": 100000,
+ "learning_rate": 4e-4,
+ "sigma": 1.0,
+ "iters_per_checkpoint": 2000,
+ "batch_size": 96,
+ "seed": 1234,
+ "checkpoint_path": "",
+ "with_tensorboard": true
+ },
+ "data_config": {
+ "training_files": "train_files.txt",
+ "segment_length": 16384,
+ "sampling_rate": 22050,
+ "filter_length": 1024,
+ "hop_length": 256,
+ "win_length": 1024,
+ "mel_fmin": 0.0,
+ "mel_fmax": 8000.0
+ },
+ "dist_config": {
+ "dist_backend": "nccl",
+ "dist_url": "tcp://localhost:54321"
+ },
+ "squeezewave_config": {
+ "n_mel_channels": 80,
+ "n_flows": 12,
+ "n_audio_channel": 256,
+ "n_early_every": 2,
+ "n_early_size": 16,
+ "WN_config": {
+ "n_layers": 8,
+ "n_channels": 128,
+ "kernel_size": 3
+ }
+ }
+}
diff --git a/SqueezeWave/configs/config_a256_c256.json b/SqueezeWave/configs/config_a256_c256.json
new file mode 100644
index 0000000..5ce633a
--- /dev/null
+++ b/SqueezeWave/configs/config_a256_c256.json
@@ -0,0 +1,40 @@
+{
+ "train_config": {
+ "fp16_run": true,
+ "output_directory": "checkpoints",
+ "epochs": 100000,
+ "learning_rate": 4e-4,
+ "sigma": 1.0,
+ "iters_per_checkpoint": 2000,
+ "batch_size": 96,
+ "seed": 1234,
+ "checkpoint_path": "",
+ "with_tensorboard": true
+ },
+ "data_config": {
+ "training_files": "train_files.txt",
+ "segment_length": 16384,
+ "sampling_rate": 22050,
+ "filter_length": 1024,
+ "hop_length": 256,
+ "win_length": 1024,
+ "mel_fmin": 0.0,
+ "mel_fmax": 8000.0
+ },
+ "dist_config": {
+ "dist_backend": "nccl",
+ "dist_url": "tcp://localhost:54321"
+ },
+ "squeezewave_config": {
+ "n_mel_channels": 80,
+ "n_flows": 12,
+ "n_audio_channel": 256,
+ "n_early_every": 2,
+ "n_early_size": 16,
+ "WN_config": {
+ "n_layers": 8,
+ "n_channels": 256,
+ "kernel_size": 3
+ }
+ }
+}
diff --git a/SqueezeWave/convert_model.py b/SqueezeWave/convert_model.py
new file mode 100644
index 0000000..b725a17
--- /dev/null
+++ b/SqueezeWave/convert_model.py
@@ -0,0 +1,70 @@
+import sys
+import copy
+import torch
+
+def _check_model_old_version(model):
+ if hasattr(model.WN[0], 'res_layers') or hasattr(model.WN[0], 'cond_layers'):
+ return True
+ else:
+ return False
+
+def _update_model_res_skip(old_model, new_model):
+ for idx in range(0, len(new_model.WN)):
+ wavenet = new_model.WN[idx]
+ n_channels = wavenet.n_channels
+ n_layers = wavenet.n_layers
+ wavenet.res_skip_layers = torch.nn.ModuleList()
+ for i in range(0, n_layers):
+ if i < n_layers - 1:
+ res_skip_channels = 2*n_channels
+ else:
+ res_skip_channels = n_channels
+ res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
+ skip_layer = torch.nn.utils.remove_weight_norm(wavenet.skip_layers[i])
+ if i < n_layers - 1:
+ res_layer = torch.nn.utils.remove_weight_norm(wavenet.res_layers[i])
+ res_skip_layer.weight = torch.nn.Parameter(torch.cat([res_layer.weight, skip_layer.weight]))
+ res_skip_layer.bias = torch.nn.Parameter(torch.cat([res_layer.bias, skip_layer.bias]))
+ else:
+ res_skip_layer.weight = torch.nn.Parameter(skip_layer.weight)
+ res_skip_layer.bias = torch.nn.Parameter(skip_layer.bias)
+ res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
+ wavenet.res_skip_layers.append(res_skip_layer)
+ del wavenet.res_layers
+ del wavenet.skip_layers
+
+def _update_model_cond(old_model, new_model):
+ for idx in range(0, len(new_model.WN)):
+ wavenet = new_model.WN[idx]
+ n_channels = wavenet.n_channels
+ n_layers = wavenet.n_layers
+ n_mel_channels = wavenet.cond_layers[0].weight.shape[1]
+ cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels*n_layers, 1)
+ cond_layer_weight = []
+ cond_layer_bias = []
+ for i in range(0, n_layers):
+ _cond_layer = torch.nn.utils.remove_weight_norm(wavenet.cond_layers[i])
+ cond_layer_weight.append(_cond_layer.weight)
+ cond_layer_bias.append(_cond_layer.bias)
+ cond_layer.weight = torch.nn.Parameter(torch.cat(cond_layer_weight))
+ cond_layer.bias = torch.nn.Parameter(torch.cat(cond_layer_bias))
+ cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+ wavenet.cond_layer = cond_layer
+ del wavenet.cond_layers
+
+def update_model(old_model):
+ if not _check_model_old_version(old_model):
+ return old_model
+ new_model = copy.deepcopy(old_model)
+ if hasattr(old_model.WN[0], 'res_layers'):
+ _update_model_res_skip(old_model, new_model)
+ if hasattr(old_model.WN[0], 'cond_layers'):
+ _update_model_cond(old_model, new_model)
+ return new_model
+
+if __name__ == '__main__':
+ old_model_path = sys.argv[1]
+ new_model_path = sys.argv[2]
+ model = torch.load(old_model_path)
+ model['model'] = update_model(model['model'])
+ torch.save(model, new_model_path)
diff --git a/SqueezeWave/denoiser.py b/SqueezeWave/denoiser.py
new file mode 100644
index 0000000..8f9ff57
--- /dev/null
+++ b/SqueezeWave/denoiser.py
@@ -0,0 +1,39 @@
+import sys
+import torch
+from stft import STFT
+
+
+class Denoiser(torch.nn.Module):
+ """ Removes model bias from audio produced with squeezewave"""
+
+ def __init__(self, squeezewave, filter_length=1024, n_overlap=4,
+ win_length=1024, mode='zeros'):
+ super(Denoiser, self).__init__()
+ self.stft = STFT(filter_length=filter_length,
+ hop_length=int(filter_length/n_overlap),
+ win_length=win_length).cuda()
+ if mode == 'zeros':
+ mel_input = torch.zeros(
+ (1, 80, 88),
+ dtype=squeezewave.upsample.weight.dtype,
+ device=squeezewave.upsample.weight.device)
+ elif mode == 'normal':
+ mel_input = torch.randn(
+ (1, 80, 88),
+ dtype=squeezewave.upsample.weight.dtype,
+ device=squeezewave.upsample.weight.device)
+ else:
+ raise Exception("Mode {} if not supported".format(mode))
+
+ with torch.no_grad():
+ bias_audio = squeezewave.infer(mel_input, sigma=0.0).float()
+ bias_spec, _ = self.stft.transform(bias_audio)
+
+ self.register_buffer('bias_spec', bias_spec[:, :, 0][:, :, None])
+
+ def forward(self, audio, strength=0.1):
+ audio_spec, audio_angles = self.stft.transform(audio.cuda().float())
+ audio_spec_denoised = audio_spec - self.bias_spec * strength
+ audio_spec_denoised = torch.clamp(audio_spec_denoised, 0.0)
+ audio_denoised = self.stft.inverse(audio_spec_denoised, audio_angles)
+ return audio_denoised
diff --git a/SqueezeWave/distributed.py b/SqueezeWave/distributed.py
new file mode 100644
index 0000000..e1b9f55
--- /dev/null
+++ b/SqueezeWave/distributed.py
@@ -0,0 +1,191 @@
+# We retain the copyright notice by NVIDIA from the original code. However, we
+# we reserve our rights on the modifications based on the original code.
+#
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+import os
+import sys
+import time
+import subprocess
+import argparse
+
+import torch
+import torch.distributed as dist
+from torch.autograd import Variable
+
+def reduce_tensor(tensor, num_gpus):
+ rt = tensor.clone()
+ dist.all_reduce(rt, op=dist.reduce_op.SUM)
+# rt /= (num_gpus*2)
+ rt /=num_gpus
+ return rt
+
+def init_distributed(rank, num_gpus, group_name, dist_backend, dist_url):
+ assert torch.cuda.is_available(), "Distributed mode requires CUDA."
+ print("Initializing Distributed")
+
+ # Set cuda device so everything is done on the right GPU.
+ torch.cuda.set_device(rank % torch.cuda.device_count())
+
+# os.environ['MASTER_ADDR'] = '172.31.44.232'
+# os.environ['MASTER_PORT'] = '58217'
+ # Initialize distributed communication
+ dist.init_process_group(dist_backend, init_method=dist_url,
+ world_size=num_gpus, rank=rank,
+ group_name=group_name)
+
+def _flatten_dense_tensors(tensors):
+ """Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of
+ same dense type.
+ Since inputs are dense, the resulting tensor will be a concatenated 1D
+ buffer. Element-wise operation on this buffer will be equivalent to
+ operating individually.
+ Arguments:
+ tensors (Iterable[Tensor]): dense tensors to flatten.
+ Returns:
+ A contiguous 1D buffer containing input tensors.
+ """
+ if len(tensors) == 1:
+ return tensors[0].contiguous().view(-1)
+ flat = torch.cat([t.contiguous().view(-1) for t in tensors], dim=0)
+ return flat
+
+def _unflatten_dense_tensors(flat, tensors):
+ """View a flat buffer using the sizes of tensors. Assume that tensors are of
+ same dense type, and that flat is given by _flatten_dense_tensors.
+ Arguments:
+ flat (Tensor): flattened dense tensors to unflatten.
+ tensors (Iterable[Tensor]): dense tensors whose sizes will be used to
+ unflatten flat.
+ Returns:
+ Unflattened dense tensors with sizes same as tensors and values from
+ flat.
+ """
+ outputs = []
+ offset = 0
+ for tensor in tensors:
+ numel = tensor.numel()
+ outputs.append(flat.narrow(0, offset, numel).view_as(tensor))
+ offset += numel
+ return tuple(outputs)
+
+def apply_gradient_allreduce(module):
+ """
+ Modifies existing model to do gradient allreduce, but doesn't change class
+ so you don't need "module"
+ """
+ if not hasattr(dist, '_backend'):
+ module.warn_on_half = True
+ else:
+ module.warn_on_half = True if dist._backend == dist.dist_backend.GLOO else False
+
+ for p in module.state_dict().values():
+ if not torch.is_tensor(p):
+ continue
+ dist.broadcast(p, 0)
+
+ def allreduce_params():
+ if(module.needs_reduction):
+ module.needs_reduction = False
+ buckets = {}
+ for param in module.parameters():
+ if param.requires_grad and param.grad is not None:
+ tp = type(param.data)
+ if tp not in buckets:
+ buckets[tp] = []
+ buckets[tp].append(param)
+ if module.warn_on_half:
+ if torch.cuda.HalfTensor in buckets:
+ print("WARNING: gloo dist backend for half parameters may be extremely slow." +
+ " It is recommended to use the NCCL backend in this case. This currently requires" +
+ "PyTorch built from top of tree master.")
+ module.warn_on_half = False
+
+ for tp in buckets:
+ bucket = buckets[tp]
+ grads = [param.grad.data for param in bucket]
+ coalesced = _flatten_dense_tensors(grads)
+ dist.all_reduce(coalesced)
+ coalesced /= dist.get_world_size()
+ for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):
+ buf.copy_(synced)
+
+ for param in list(module.parameters()):
+ def allreduce_hook(*unused):
+ Variable._execution_engine.queue_callback(allreduce_params)
+ if param.requires_grad:
+ param.register_hook(allreduce_hook)
+ dir(param)
+
+ def set_needs_reduction(self, input, output):
+ self.needs_reduction = True
+
+ module.register_forward_hook(set_needs_reduction)
+ return module
+
+
+def main(config, stdout_dir, args_str):
+ args_list = ['-u']
+ args_list.append('train.py')
+ args_list += args_str.split(' ') if len(args_str) > 0 else []
+
+ args_list.append('--config={}'.format(config))
+
+ num_gpus = torch.cuda.device_count()
+ args_list.append('--num_gpus={}'.format(num_gpus))
+ args_list.append("--group_name=group_{}".format(time.strftime("%Y_%m_%d-%H%M%S")))
+
+ if not os.path.isdir(stdout_dir):
+ os.makedirs(stdout_dir)
+ os.chmod(stdout_dir, 0o775)
+
+ workers = []
+
+ for i in range(num_gpus):
+ args_list[-2] = '--rank={}'.format(i)
+ stdout = None if i == 0 else open(
+ os.path.join(stdout_dir, "GPU_{}.log".format(i)), "w")
+ print(args_list)
+ p = subprocess.Popen([str(sys.executable)]+args_list, stdout=stdout)
+ workers.append(p)
+
+ for p in workers:
+ p.wait()
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-c', '--config', type=str, required=True,
+ help='JSON file for configuration')
+ parser.add_argument('-s', '--stdout_dir', type=str, default=".",
+ help='directory to save stoud logs')
+ parser.add_argument(
+ '-a', '--args_str', type=str, default='',
+ help='double quoted string with space separated key value pairs')
+
+ args = parser.parse_args()
+ main(args.config, args.stdout_dir, args.args_str)
diff --git a/SqueezeWave/glow.py b/SqueezeWave/glow.py
new file mode 100644
index 0000000..f692103
--- /dev/null
+++ b/SqueezeWave/glow.py
@@ -0,0 +1,328 @@
+# We retain the copyright notice by NVIDIA from the original code. However, we
+# we reserve our rights on the modifications based on the original code.
+#
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+import torch
+from torch.autograd import Variable
+import torch.nn.functional as F
+import numpy as np
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+ n_channels_int = n_channels[0]
+ in_act = input_a+input_b
+ t_act = torch.tanh(in_act[:, :n_channels_int, :])
+ s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+ acts = t_act * s_act
+ return acts
+
+
+class Upsample1d(torch.nn.Module):
+ def __init__(self, scale=2):
+ super(Upsample1d, self).__init__()
+ self.scale = scale
+
+ def forward(self, x):
+ y = F.interpolate(
+ x, scale_factor=self.scale, mode='nearest')
+ return y
+
+
+class SqueezeWaveLoss(torch.nn.Module):
+ def __init__(self, sigma=1.0):
+ super(SqueezeWaveLoss, self).__init__()
+ self.sigma = sigma
+
+ def forward(self, model_output):
+ z, log_s_list, log_det_W_list = model_output
+ for i, log_s in enumerate(log_s_list):
+ if i == 0:
+ log_s_total = torch.sum(log_s)
+ log_det_W_total = log_det_W_list[i]
+ else:
+ log_s_total = log_s_total + torch.sum(log_s)
+ log_det_W_total += log_det_W_list[i]
+
+ loss = torch.sum(z*z)/(2*self.sigma*self.sigma) - log_s_total - log_det_W_total
+ return loss/(z.size(0)*z.size(1)*z.size(2))
+
+
+class Invertible1x1Conv(torch.nn.Module):
+ """
+ The layer outputs both the convolution, and the log determinant
+ of its weight matrix. If reverse=True it does convolution with
+ inverse
+ """
+ def __init__(self, c):
+ super(Invertible1x1Conv, self).__init__()
+ self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
+ bias=False)
+
+ # Sample a random orthonormal matrix to initialize weights
+ W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
+
+ # Ensure determinant is 1.0 not -1.0
+ if torch.det(W) < 0:
+ W[:,0] = -1*W[:,0]
+ W = W.view(c, c, 1)
+ self.conv.weight.data = W
+
+ def forward(self, z, reverse=False):
+ # shape
+ batch_size, group_size, n_of_groups = z.size()
+ W = self.conv.weight.squeeze()
+
+ if reverse:
+ if not hasattr(self, 'W_inverse'):
+ # Reverse computation
+ W_inverse = W.float().inverse()
+ W_inverse = Variable(W_inverse[..., None])
+ if z.type() == 'torch.cuda.HalfTensor':
+ W_inverse = W_inverse.half()
+ self.W_inverse = W_inverse
+ z = F.conv1d(z, self.W_inverse, bias=None, stride=1, padding=0)
+ return z
+ else:
+ # Forward computation
+ log_det_W = batch_size * n_of_groups * torch.logdet(W)
+ z = self.conv(z)
+ return z, log_det_W
+
+
+class WN(torch.nn.Module):
+ """
+ This is the WaveNet like layer for the affine coupling. The primary difference
+ from WaveNet is the convolutions need not be causal. There is also no dilation
+ size reset. The dilation only doubles on each layer
+ """
+ def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels,
+ kernel_size):
+ super(WN, self).__init__()
+ assert(kernel_size % 2 == 1)
+ assert(n_channels % 2 == 0)
+ self.n_layers = n_layers
+ self.n_channels = n_channels
+ self.in_layers = torch.nn.ModuleList()
+ self.res_skip_layers = torch.nn.ModuleList()
+ self.upsample = Upsample1d(2)
+ start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
+ start = torch.nn.utils.weight_norm(start, name='weight')
+ self.start = start
+ end = torch.nn.Conv1d(n_channels, 2*n_in_channels, 1)
+ end.weight.data.zero_()
+ end.bias.data.zero_()
+ self.end = end
+
+ # cond_layer
+ cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels*n_layers, 1)
+ self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+ for i in range(n_layers):
+ dilation = 1
+ padding = int((kernel_size*dilation - dilation)/2)
+ # depthwise separable convolution
+ depthwise = torch.nn.Conv1d(n_channels, n_channels, 3,
+ dilation=dilation, padding=padding,
+ groups=n_channels).cuda()
+ pointwise = torch.nn.Conv1d(n_channels, 2*n_channels, 1).cuda()
+ bn = torch.nn.BatchNorm1d(n_channels)
+ self.in_layers.append(torch.nn.Sequential(bn, depthwise, pointwise))
+ # res_skip_layer
+ res_skip_layer = torch.nn.Conv1d(n_channels, n_channels, 1)
+ res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
+ self.res_skip_layers.append(res_skip_layer)
+
+ def forward(self, forward_input):
+ audio, spect = forward_input
+ audio = self.start(audio)
+ n_channels_tensor = torch.IntTensor([self.n_channels])
+ # pass all the mel_spectrograms to cond_layer
+ spect = self.cond_layer(spect)
+ for i in range(self.n_layers):
+ # split the corresponding mel_spectrogram
+ spect_offset = i*2*self.n_channels
+ spec = spect[:,spect_offset:spect_offset+2*self.n_channels,:]
+ if audio.size(2) > spec.size(2):
+ cond = self.upsample(spec)
+ else:
+ cond = spec
+ acts = fused_add_tanh_sigmoid_multiply(
+ self.in_layers[i](audio),
+ cond,
+ n_channels_tensor)
+ # res_skip
+ res_skip_acts = self.res_skip_layers[i](acts)
+ audio = audio + res_skip_acts
+ return self.end(audio)
+
+
+class SqueezeWave(torch.nn.Module):
+ def __init__(self, n_mel_channels, n_flows, n_audio_channel, n_early_every,
+ n_early_size, WN_config):
+ super(SqueezeWave, self).__init__()
+ assert(n_audio_channel % 2 == 0)
+ self.n_flows = n_flows
+ self.n_audio_channel = n_audio_channel
+ self.n_early_every = n_early_every
+ self.n_early_size = n_early_size
+ self.WN = torch.nn.ModuleList()
+ self.convinv = torch.nn.ModuleList()
+
+ n_half = int(n_audio_channel / 2)
+ # Set up layers with the right sizes based on how many dimensions
+ # have been output already
+ n_remaining_channels = n_audio_channel
+ for k in range(n_flows):
+ if k % self.n_early_every == 0 and k > 0:
+ n_half = n_half - int(self.n_early_size/2)
+ n_remaining_channels = n_remaining_channels - self.n_early_size
+ self.convinv.append(Invertible1x1Conv(n_remaining_channels))
+ self.WN.append(WN(n_half, n_mel_channels, **WN_config))
+
+ self.n_remaining_channels = n_remaining_channels # Useful during inference
+
+ def forward(self, forward_input):
+ """
+ forward_input[0] = mel_spectrogram: batch x n_mel_channels x frames
+ forward_input[1] = audio: batch x time
+ """
+ spect, audio = forward_input
+
+ audio = audio.unfold(
+ 1, self.n_audio_channel, self.n_audio_channel).permute(0, 2, 1)
+ output_audio = []
+ log_s_list = []
+ log_det_W_list = []
+
+ for k in range(self.n_flows):
+ if k % self.n_early_every == 0 and k > 0:
+ output_audio.append(audio[:,:self.n_early_size,:])
+ audio = audio[:,self.n_early_size:,:]
+
+ audio, log_det_W = self.convinv[k](audio)
+ log_det_W_list.append(log_det_W)
+
+ n_half = int(audio.size(1)/2)
+ audio_0 = audio[:,:n_half,:]
+ audio_1 = audio[:,n_half:,:]
+
+ output = self.WN[k]((audio_0, spect))
+ log_s = output[:, n_half:, :]
+ b = output[:, :n_half, :]
+
+ audio_1 = (torch.exp(log_s))*audio_1 + b
+ log_s_list.append(log_s)
+ audio = torch.cat([audio_0, audio_1], 1)
+
+ output_audio.append(audio)
+ return torch.cat(output_audio, 1), log_s_list, log_det_W_list
+
+ def infer(self, spect, sigma=1.0):
+ spect_size = spect.size()
+ l = spect.size(2)*(256 // self.n_audio_channel)
+ if spect.type() == 'torch.cuda.HalfTensor':
+ audio = torch.cuda.HalfTensor(spect.size(0),
+ self.n_remaining_channels,
+ l).normal_()
+ else:
+ audio = torch.cuda.FloatTensor(spect.size(0),
+ self.n_remaining_channels,
+ l).normal_()
+
+ for k in reversed(range(self.n_flows)):
+ n_half = int(audio.size(1)/2)
+ audio_0 = audio[:,:n_half,:]
+ audio_1 = audio[:,n_half:,:]
+ output = self.WN[k]((audio_0, spect))
+
+ s = output[:, n_half:, :]
+ b = output[:, :n_half, :]
+ audio_1 = (audio_1 - b)/torch.exp(s)
+ audio = torch.cat([audio_0, audio_1],1)
+
+ audio = self.convinv[k](audio, reverse=True)
+
+ if k % self.n_early_every == 0 and k > 0:
+ if spect.type() == 'torch.cuda.HalfTensor':
+ z = torch.cuda.HalfTensor(spect.size(0), self.n_early_size, l).normal_()
+ else:
+ z = torch.cuda.FloatTensor(spect.size(0), self.n_early_size, l).normal_()
+ audio = torch.cat((sigma*z, audio),1)
+
+ audio = audio.permute(0,2,1).contiguous().view(audio.size(0), -1).data
+ return audio
+
+
+ @staticmethod
+ def remove_weightnorm(model):
+ squeezewave = model
+ for WN in squeezewave.WN:
+ WN.start = torch.nn.utils.remove_weight_norm(WN.start)
+ WN.in_layers = remove_batch_norm(WN.in_layers)
+ WN.cond_layer = torch.nn.utils.remove_weight_norm(WN.cond_layer)
+ WN.res_skip_layers = remove(WN.res_skip_layers)
+ return squeezewave
+
+def fuse_conv_and_bn(conv, bn):
+ fusedconv = torch.nn.Conv1d(
+ conv.in_channels,
+ conv.out_channels,
+ kernel_size = conv.kernel_size,
+ padding=conv.padding,
+ bias=True,
+ groups=conv.groups)
+ w_conv = conv.weight.clone().view(conv.out_channels, -1)
+ w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps+bn.running_var)))
+ w_bn = w_bn.clone()
+ fusedconv.weight.data = torch.mm(w_bn, w_conv).view(fusedconv.weight.size())
+ if conv.bias is not None:
+ b_conv = conv.bias
+ else:
+ b_conv = torch.zeros( conv.weight.size(0) )
+ b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
+ b_bn = torch.unsqueeze(b_bn, 1)
+ bn_3 = b_bn.expand(-1, 3)
+ b = torch.matmul(w_conv, torch.transpose(bn_3, 0, 1))[range(b_bn.size()[0]), range(b_bn.size()[0])]
+ fusedconv.bias.data = ( b_conv + b )
+ return fusedconv
+
+def remove_batch_norm(conv_list):
+ new_conv_list = torch.nn.ModuleList()
+ for old_conv in conv_list:
+ depthwise = fuse_conv_and_bn(old_conv[1], old_conv[0])
+ pointwise = old_conv[2]
+ new_conv_list.append(torch.nn.Sequential(depthwise, pointwise))
+ return new_conv_list
+
+def remove(conv_list):
+ new_conv_list = torch.nn.ModuleList()
+ for old_conv in conv_list:
+ old_conv = torch.nn.utils.remove_weight_norm(old_conv)
+ new_conv_list.append(old_conv)
+ return new_conv_list
+
diff --git a/SqueezeWave/inference.py b/SqueezeWave/inference.py
new file mode 100644
index 0000000..568e6ce
--- /dev/null
+++ b/SqueezeWave/inference.py
@@ -0,0 +1,87 @@
+# We retain the copyright notice by NVIDIA from the original code. However, we
+# we reserve our rights on the modifications based on the original code.
+#
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+import os
+from scipy.io.wavfile import write
+import torch
+from mel2samp import files_to_list, MAX_WAV_VALUE
+from denoiser import Denoiser
+
+
+def main(mel_files, squeezewave_path, sigma, output_dir, sampling_rate, is_fp16,
+ denoiser_strength):
+ mel_files = files_to_list(mel_files)
+ squeezewave = torch.load(squeezewave_path)['model']
+ squeezewave = squeezewave.remove_weightnorm(squeezewave)
+ squeezewave.cuda().eval()
+ if is_fp16:
+ from apex import amp
+ squeezewave, _ = amp.initialize(squeezewave, [], opt_level="O3")
+
+ if denoiser_strength > 0:
+ denoiser = Denoiser(squeezewave).cuda()
+
+ for i, file_path in enumerate(mel_files):
+ file_name = os.path.splitext(os.path.basename(file_path))[0]
+ mel = torch.load(file_path)
+ mel = torch.autograd.Variable(mel.cuda())
+ mel = torch.unsqueeze(mel, 0)
+ mel = mel.half() if is_fp16 else mel
+ with torch.no_grad():
+ audio = squeezewave.infer(mel, sigma=sigma).float()
+ if denoiser_strength > 0:
+ audio = denoiser(audio, denoiser_strength)
+ audio = audio * MAX_WAV_VALUE
+ audio = audio.squeeze()
+ audio = audio.cpu().numpy()
+ audio = audio.astype('int16')
+ audio_path = os.path.join(
+ output_dir, "{}_synthesis.wav".format(file_name))
+ write(audio_path, sampling_rate, audio)
+ print(audio_path)
+
+
+if __name__ == "__main__":
+ import argparse
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-f', "--filelist_path", required=True)
+ parser.add_argument('-w', '--squeezewave_path',
+ help='Path to squeezewave decoder checkpoint with model')
+ parser.add_argument('-o', "--output_dir", required=True)
+ parser.add_argument("-s", "--sigma", default=1.0, type=float)
+ parser.add_argument("--sampling_rate", default=22050, type=int)
+ parser.add_argument("--is_fp16", action="store_true")
+ parser.add_argument("-d", "--denoiser_strength", default=0.0, type=float,
+ help='Removes model bias. Start with 0.1 and adjust')
+
+ args = parser.parse_args()
+
+ main(args.filelist_path, args.squeezewave_path, args.sigma, args.output_dir,
+ args.sampling_rate, args.is_fp16, args.denoiser_strength)
diff --git a/SqueezeWave/mel2samp.py b/SqueezeWave/mel2samp.py
new file mode 100644
index 0000000..0419202
--- /dev/null
+++ b/SqueezeWave/mel2samp.py
@@ -0,0 +1,150 @@
+# We retain the copyright notice by NVIDIA from the original code. However, we
+# we reserve our rights on the modifications based on the original code.
+#
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************\
+import os
+import random
+import argparse
+import json
+import torch
+import torch.utils.data
+import sys
+from scipy.io.wavfile import read
+
+# We're using the audio processing from TacoTron2 to make sure it matches
+from TacotronSTFT import TacotronSTFT
+
+MAX_WAV_VALUE = 32768.0
+
+def files_to_list(filename):
+ """
+ Takes a text file of filenames and makes a list of filenames
+ """
+ with open(filename, encoding='utf-8') as f:
+ files = f.readlines()
+
+ files = [f.rstrip() for f in files]
+ return files
+
+def load_wav_to_torch(full_path):
+ """
+ Loads wavdata into torch array
+ """
+ sampling_rate, data = read(full_path)
+ return torch.from_numpy(data).float(), sampling_rate
+
+
+class Mel2Samp(torch.utils.data.Dataset):
+ """
+ This is the main class that calculates the spectrogram and returns the
+ spectrogram, audio pair.
+ """
+ def __init__(self, n_audio_channel, training_files, segment_length,
+ filter_length, hop_length, win_length, sampling_rate, mel_fmin,
+ mel_fmax):
+ self.audio_files = files_to_list(training_files)
+ random.seed(1234)
+ random.shuffle(self.audio_files)
+ self.stft = TacotronSTFT(filter_length=filter_length,
+ hop_length=hop_length,
+ win_length=win_length,
+ sampling_rate=sampling_rate,
+ mel_fmin=mel_fmin, mel_fmax=mel_fmax,
+ n_group=n_audio_channel)
+ self.segment_length = segment_length
+ self.sampling_rate = sampling_rate
+
+ def get_mel(self, audio):
+ audio_norm = audio / MAX_WAV_VALUE
+ audio_norm = audio_norm.unsqueeze(0)
+ audio_norm = torch.autograd.Variable(audio_norm, requires_grad=False)
+ melspec = self.stft.mel_spectrogram(audio_norm)
+ melspec = torch.squeeze(melspec, 0)
+ return melspec
+
+ def __getitem__(self, index):
+ # Read audio
+ filename = self.audio_files[index]
+ audio, sampling_rate = load_wav_to_torch(filename)
+ if sampling_rate != self.sampling_rate:
+ raise ValueError("{} SR doesn't match target {} SR".format(
+ sampling_rate, self.sampling_rate))
+
+ # Take segment
+ if audio.size(0) >= self.segment_length:
+ max_audio_start = audio.size(0) - self.segment_length
+ audio_start = random.randint(0, max_audio_start)
+ audio = audio[audio_start:audio_start+self.segment_length]
+ else:
+ audio = torch.nn.functional.pad(
+ audio, (0, self.segment_length - audio.size(0)),
+ 'constant').data
+
+ mel = self.get_mel(audio)
+ audio = audio / MAX_WAV_VALUE
+
+ return (mel, audio)
+
+ def __len__(self):
+ return len(self.audio_files)
+
+# ===================================================================
+# Takes directory of clean audio and makes directory of spectrograms
+# Useful for making test sets
+# ===================================================================
+if __name__ == "__main__":
+ # Get defaults so it can work with no Sacred
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-f', "--filelist_path", required=True)
+ parser.add_argument('-c', '--config', type=str,
+ help='JSON file for configuration')
+ parser.add_argument('-o', '--output_dir', type=str,
+ help='Output directory')
+ args = parser.parse_args()
+
+ with open(args.config) as f:
+ data = f.read()
+ config = json.loads(data)
+ data_config = config["data_config"]
+ squeezewave_config = config["squeezewave_config"]
+ mel2samp = Mel2Samp(squeezewave_config['n_audio_channel'], **data_config)
+
+ filepaths = files_to_list(args.filelist_path)
+
+ # Make directory if it doesn't exist
+ if not os.path.isdir(args.output_dir):
+ os.makedirs(args.output_dir)
+ os.chmod(args.output_dir, 0o775)
+
+ for filepath in filepaths:
+ audio, sr = load_wav_to_torch(filepath)
+ melspectrogram = mel2samp.get_mel(audio)
+ filename = os.path.basename(filepath)
+ new_filepath = args.output_dir + '/' + filename + '.pt'
+ print(new_filepath)
+ torch.save(melspectrogram, new_filepath)
diff --git a/SqueezeWave/requirements.txt b/SqueezeWave/requirements.txt
new file mode 100644
index 0000000..cf54599
--- /dev/null
+++ b/SqueezeWave/requirements.txt
@@ -0,0 +1,8 @@
+torch==1.0
+matplotlib==2.1.0
+numpy==1.13.3
+inflect==0.2.5
+librosa==0.6.0
+scipy==1.0.0
+tensorboardX==1.1
+Unidecode==1.0.22
diff --git a/SqueezeWave/stft.py b/SqueezeWave/stft.py
new file mode 100644
index 0000000..98b56e5
--- /dev/null
+++ b/SqueezeWave/stft.py
@@ -0,0 +1,147 @@
+"""
+We retain the copyright notice from the original author. However, we reserve
+our rights on the modifications based on the original code.
+
+BSD 3-Clause License
+
+Copyright (c) 2017, Prem Seetharaman
+All rights reserved.
+
+* Redistribution and use in source and binary forms, with or without
+ modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice,
+ this list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice, this
+ list of conditions and the following disclaimer in the
+ documentation and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+ contributors may be used to endorse or promote products derived from this
+ software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+from torch.autograd import Variable
+from scipy.signal import get_window
+from librosa.util import pad_center, tiny
+from audio_processing import window_sumsquare
+
+
+class STFT(torch.nn.Module):
+ """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
+ def __init__(self, filter_length=800, hop_length=200, win_length=800,
+ window='hann', n_group=256):
+ super(STFT, self).__init__()
+ self.filter_length = filter_length
+ self.hop_length = hop_length
+ self.win_length = win_length
+ self.window = window
+ self.forward_transform = None
+ self.n_group = n_group
+ scale = self.filter_length / self.hop_length
+ fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+ cutoff = int((self.filter_length / 2 + 1))
+ fourier_basis = np.vstack([np.real(fourier_basis[:cutoff, :]),
+ np.imag(fourier_basis[:cutoff, :])])
+
+ forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+ inverse_basis = torch.FloatTensor(
+ np.linalg.pinv(scale * fourier_basis).T[:, None, :])
+
+ if window is not None:
+ assert(win_length >= filter_length)
+ # get window and zero center pad it to filter_length
+ fft_window = get_window(window, win_length, fftbins=True)
+ fft_window = pad_center(fft_window, filter_length)
+ fft_window = torch.from_numpy(fft_window).float()
+
+ # window the bases
+ forward_basis *= fft_window
+ inverse_basis *= fft_window
+
+ self.register_buffer('forward_basis', forward_basis.float())
+ self.register_buffer('inverse_basis', inverse_basis.float())
+
+ def transform(self, input_data):
+ num_batches = input_data.size(0)
+ num_samples = input_data.size(1)
+
+ self.num_samples = num_samples
+
+ # similar to librosa, reflect-pad the input
+ input_data = input_data.view(num_batches, 1, num_samples)
+ pad = ((64 - 1) * self.hop_length + self.filter_length - num_samples) // 2
+ if pad < 0:
+ pad = self.filter_length // 2
+ input_data = F.pad(
+ input_data.unsqueeze(1),
+ (int(pad), int(pad), 0, 0),
+ mode='reflect')
+ input_data = input_data.squeeze(1)
+
+ forward_transform = F.conv1d(
+ input_data,
+ Variable(self.forward_basis, requires_grad=False),
+ stride=self.hop_length,
+ padding=0)
+
+ cutoff = int((self.filter_length / 2) + 1)
+ real_part = forward_transform[:, :cutoff, :]
+ imag_part = forward_transform[:, cutoff:, :]
+
+ magnitude = torch.sqrt(real_part**2 + imag_part**2)
+ phase = torch.autograd.Variable(
+ torch.atan2(imag_part.data, real_part.data))
+
+ return magnitude, phase
+
+ def inverse(self, magnitude, phase):
+ recombine_magnitude_phase = torch.cat(
+ [magnitude*torch.cos(phase), magnitude*torch.sin(phase)], dim=1)
+
+ inverse_transform = F.conv_transpose1d(
+ recombine_magnitude_phase,
+ Variable(self.inverse_basis, requires_grad=False),
+ stride=self.hop_length,
+ padding=0)
+
+ if self.window is not None:
+ window_sum = window_sumsquare(
+ self.window, magnitude.size(-1), hop_length=self.hop_length,
+ win_length=self.win_length, n_fft=self.filter_length,
+ dtype=np.float32)
+ # remove modulation effects
+ approx_nonzero_indices = torch.from_numpy(
+ np.where(window_sum > tiny(window_sum))[0])
+ window_sum = torch.autograd.Variable(
+ torch.from_numpy(window_sum), requires_grad=False)
+ inverse_transform[:, :, approx_nonzero_indices] /= window_sum[approx_nonzero_indices]
+
+ # scale by hop ratio
+ inverse_transform *= float(self.filter_length) / self.hop_length
+
+ inverse_transform = inverse_transform[:, :, int(self.filter_length/2):]
+ inverse_transform = inverse_transform[:, :, :-int(self.filter_length/2):]
+
+ return inverse_transform
+
+ def forward(self, input_data):
+ self.magnitude, self.phase = self.transform(input_data)
+ reconstruction = self.inverse(self.magnitude, self.phase)
+ return reconstruction
diff --git a/SqueezeWave/train.py b/SqueezeWave/train.py
new file mode 100644
index 0000000..7cbc2e1
--- /dev/null
+++ b/SqueezeWave/train.py
@@ -0,0 +1,203 @@
+# We retain the copyright notice by NVIDIA from the original code. However, we
+# we reserve our rights on the modifications based on the original code.
+#
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+import argparse
+import json
+import os
+import torch
+
+#=====START: ADDED FOR DISTRIBUTED======
+from distributed import init_distributed, apply_gradient_allreduce, reduce_tensor
+from torch.utils.data.distributed import DistributedSampler
+#=====END: ADDED FOR DISTRIBUTED======
+
+from torch.utils.data import DataLoader
+from glow import SqueezeWave, SqueezeWaveLoss
+from mel2samp import Mel2Samp
+
+def load_checkpoint(
+ checkpoint_path, model, optimizer, n_flows, n_early_every,
+ n_early_size, n_mel_channels, n_audio_channel, WN_config):
+
+ assert os.path.isfile(checkpoint_path)
+
+ checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+ iteration = checkpoint_dict['iteration']
+ #iteration = 1
+ optimizer.load_state_dict(checkpoint_dict['optimizer'])
+ model_for_loading = checkpoint_dict['model']
+ state_dict = model_for_loading.state_dict()
+
+ model.load_state_dict(state_dict, strict = False)
+ print("Loaded checkpoint '{}' (iteration {})" .format(checkpoint_path, iteration))
+
+ return model, optimizer, iteration
+
+def save_checkpoint(model, optimizer, learning_rate, iteration, filepath):
+ print("Saving model and optimizer state at iteration {} to {}".format(
+ iteration, filepath))
+ model_for_saving = SqueezeWave(**squeezewave_config).cuda()
+ model_for_saving.load_state_dict(model.state_dict())
+ torch.save({'model': model_for_saving,
+ 'iteration': iteration,
+ 'optimizer': optimizer.state_dict(),
+ 'learning_rate': learning_rate}, filepath)
+
+def train(num_gpus, rank, group_name, output_directory, epochs, learning_rate,
+ sigma, iters_per_checkpoint, batch_size, seed, fp16_run,
+ checkpoint_path, with_tensorboard):
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+ #=====START: ADDED FOR DISTRIBUTED======
+ if num_gpus > 1:
+ init_distributed(rank, num_gpus, group_name, **dist_config)
+ #=====END: ADDED FOR DISTRIBUTED======
+
+ criterion = SqueezeWaveLoss(sigma)
+ model = SqueezeWave(**squeezewave_config).cuda()
+ print(model)
+ pytorch_total_params = sum(p.numel() for p in model.parameters())
+ pytorch_total_params_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
+ print("param", pytorch_total_params)
+ print("param trainable", pytorch_total_params_train)
+
+ #=====START: ADDED FOR DISTRIBUTED======
+ if num_gpus > 1:
+ model = apply_gradient_allreduce(model)
+ #=====END: ADDED FOR DISTRIBUTED======
+
+ optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+
+ if fp16_run:
+ from apex import amp
+ model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
+
+ # Load checkpoint if one exists
+ iteration = 0
+ if checkpoint_path != "":
+ model, optimizer, iteration = load_checkpoint(checkpoint_path, model,
+ optimizer, **squeezewave_config)
+ iteration += 1 # next iteration is iteration + 1
+
+ n_audio_channel = squeezewave_config["n_audio_channel"]
+ trainset = Mel2Samp(n_audio_channel, **data_config)
+ # =====START: ADDED FOR DISTRIBUTED======
+ train_sampler = DistributedSampler(trainset) if num_gpus > 1 else None
+ # =====END: ADDED FOR DISTRIBUTED======
+ train_loader = DataLoader(trainset, num_workers=0, shuffle=False,
+ sampler=train_sampler,
+ batch_size=batch_size,
+ pin_memory=False,
+ drop_last=True)
+
+ # Get shared output_directory ready
+ if rank == 0:
+ if not os.path.isdir(output_directory):
+ os.makedirs(output_directory)
+ os.chmod(output_directory, 0o775)
+ print("output directory", output_directory)
+
+ if with_tensorboard and rank == 0:
+ from tensorboardX import SummaryWriter
+ logger = SummaryWriter(os.path.join(output_directory, 'logs'))
+
+ model.train()
+ epoch_offset = max(0, int(iteration / len(train_loader)))
+ # ================ MAIN TRAINNIG LOOP! ===================
+ for epoch in range(epoch_offset, epochs):
+ print("Epoch: {}".format(epoch))
+ for i, batch in enumerate(train_loader):
+ model.zero_grad()
+
+ mel, audio = batch
+ mel = torch.autograd.Variable(mel.cuda())
+ audio = torch.autograd.Variable(audio.cuda())
+ outputs = model((mel, audio))
+
+ loss = criterion(outputs)
+ if num_gpus > 1:
+ reduced_loss = reduce_tensor(loss.data, num_gpus).item()
+ else:
+ reduced_loss = loss.item()
+
+ if fp16_run:
+ with amp.scale_loss(loss, optimizer) as scaled_loss:
+ scaled_loss.backward()
+ else:
+ loss.backward()
+
+ optimizer.step()
+
+ print("{}:\t{:.9f}\t".format(iteration, reduced_loss))
+ if with_tensorboard and rank == 0:
+ logger.add_scalar('training_loss', reduced_loss, i + len(train_loader) * epoch)
+ if (iteration % iters_per_checkpoint == 0):
+ if rank == 0:
+ checkpoint_path = "{}/SqueezeWave_{}".format(
+ output_directory, iteration)
+ save_checkpoint(model, optimizer, learning_rate, iteration,
+ checkpoint_path)
+
+ iteration += 1
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-c', '--config', type=str,
+ help='JSON file for configuration')
+ parser.add_argument('-r', '--rank', type=int, default=0,
+ help='rank of process for distributed')
+ parser.add_argument('-g', '--group_name', type=str, default='',
+ help='name of group for distributed')
+ args = parser.parse_args()
+
+ # Parse configs. Globals nicer in this case
+ with open(args.config) as f:
+ data = f.read()
+ config = json.loads(data)
+ train_config = config["train_config"]
+ global data_config
+ data_config = config["data_config"]
+ global dist_config
+ dist_config = config["dist_config"]
+ global squeezewave_config
+ squeezewave_config = config["squeezewave_config"]
+
+ num_gpus = torch.cuda.device_count()
+ if num_gpus > 1:
+ if args.group_name == '':
+ print("WARNING: Multiple GPUs detected but no distributed group set")
+ print("Only running 1 GPU. Use distributed.py for multiple GPUs")
+ num_gpus = 1
+
+ if num_gpus == 1 and args.rank != 0:
+ raise Exception("Doing single GPU training on rank > 0")
+
+ torch.backends.cudnn.enabled = True
+ torch.backends.cudnn.benchmark = False
+ train(num_gpus, args.rank, args.group_name, **train_config)