rename glow-tts-2 to glow-tts (removing git submodule)

glow-tts-jaywalnut/glow-tts/commons.py

@@ -0,0 +1,235 @@
+from typing import List
+
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+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
+
+
+@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
+
+
+# NB(bt): type hint is because of torchjit
+def convert_pad_shape(pad_shape: List[List[int]]):
+  l = pad_shape[::-1]
+  # NB(bt): torch jit
+  #pad_shape = [item for sublist in l for item in sublist]
+  pad_shape = l[0].copy()
+  for sublist in l[1:]:
+    pad_shape.extend(sublist)
+  #pad_shape = l.flatten()
+  return pad_shape
+
+
+def shift_1d(x):
+  x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+  return x
+
+
+def sequence_mask(length, max_length=None):
+  if max_length is None:
+    max_length = length.max()
+  x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+  return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def maximum_path(value, mask, max_neg_val=-np.inf):
+  """ Numpy-friendly version. It's about 4 times faster than torch version.
+  value: [b, t_x, t_y]
+  mask: [b, t_x, t_y]
+  """
+  value = value * mask
+
+  device = value.device
+  dtype = value.dtype
+  value = value.cpu().detach().numpy()
+  mask = mask.cpu().detach().numpy().astype(np.bool)
+  
+  b, t_x, t_y = value.shape
+  direction = np.zeros(value.shape, dtype=np.int64)
+  v = np.zeros((b, t_x), dtype=np.float32)
+  x_range = np.arange(t_x, dtype=np.float32).reshape(1,-1)
+  for j in range(t_y):
+    v0 = np.pad(v, [[0,0],[1,0]], mode="constant", constant_values=max_neg_val)[:, :-1]
+    v1 = v
+    max_mask = (v1 >= v0)
+    v_max = np.where(max_mask, v1, v0)
+    direction[:, :, j] = max_mask
+    
+    index_mask = (x_range <= j)
+    v = np.where(index_mask, v_max + value[:, :, j], max_neg_val)
+  direction = np.where(mask, direction, 1)
+    
+  path = np.zeros(value.shape, dtype=np.float32)
+  index = mask[:, :, 0].sum(1).astype(np.int64) - 1
+  index_range = np.arange(b)
+  for j in reversed(range(t_y)):
+    path[index_range, index, j] = 1
+    index = index + direction[index_range, index, j] - 1
+  path = path * mask.astype(np.float32)
+  path = torch.from_numpy(path).to(device=device, dtype=dtype)
+  return path
+
+
+def generate_path(duration, mask):
+  """
+  duration: [b, t_x]
+  mask: [b, t_x, t_y]
+  """
+  device = duration.device
+  
+  b, t_x, t_y = mask.shape
+  cum_duration = torch.cumsum(duration, 1)
+  path = torch.zeros(b, t_x, t_y, dtype=mask.dtype).to(device=device)
+  
+  cum_duration_flat = cum_duration.view(b * t_x)
+  path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+  path = path.view(b, t_x, t_y)
+  path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:,:-1]
+  path = path * mask
+  return path
+
+
+class Adam():
+  def __init__(self, params, scheduler, dim_model, warmup_steps=4000, lr=1e0, betas=(0.9, 0.98), eps=1e-9):
+    self.params = params
+    self.scheduler = scheduler
+    self.dim_model = dim_model
+    self.warmup_steps = warmup_steps
+    self.lr = lr
+    self.betas = betas
+    self.eps = eps
+
+    self.step_num = 1
+    self.cur_lr = lr * self._get_lr_scale()
+    
+    self._optim = torch.optim.Adam(params, lr=self.cur_lr, betas=betas, eps=eps)
+
+  def _get_lr_scale(self):
+    if self.scheduler == "noam":
+      return np.power(self.dim_model, -0.5) * np.min([np.power(self.step_num, -0.5), self.step_num * np.power(self.warmup_steps, -1.5)])
+    else:
+      return 1
+
+  def _update_learning_rate(self):
+    self.step_num += 1
+    if self.scheduler == "noam":
+      self.cur_lr = self.lr * self._get_lr_scale()
+      for param_group in self._optim.param_groups:
+        param_group['lr'] = self.cur_lr
+
+  def get_lr(self):
+    return self.cur_lr
+
+  def step(self):
+    self._optim.step()
+    self._update_learning_rate()
+
+  def zero_grad(self):
+    self._optim.zero_grad()
+
+  def load_state_dict(self, d):
+    self._optim.load_state_dict(d)
+
+  def state_dict(self):
+    return self._optim.state_dict()
+
+
+class TacotronSTFT(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=8000.0):
+    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)
+    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
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+  if isinstance(parameters, torch.Tensor):
+    parameters = [parameters]
+  parameters = list(filter(lambda p: p.grad is not None, parameters))
+  norm_type = float(norm_type)
+  clip_value = float(clip_value)
+
+  total_norm = 0
+  for p in parameters:
+    param_norm = p.grad.data.norm(norm_type)
+    total_norm += param_norm.item() ** norm_type
+
+    p.grad.data.clamp_(min=-clip_value, max=clip_value)
+  total_norm = total_norm ** (1. / norm_type)
+  return total_norm
+
+
+def squeeze(x, x_mask=None, n_sqz=2):
+  b, c, t = x.size()
+
+  t = (t // n_sqz) * n_sqz
+  x = x[:,:,:t]
+  x_sqz = x.view(b, c, t//n_sqz, n_sqz)
+  x_sqz = x_sqz.permute(0, 3, 1, 2).contiguous().view(b, c*n_sqz, t//n_sqz)
+  
+  if x_mask is not None:
+    x_mask = x_mask[:,:,n_sqz-1::n_sqz]
+  else:
+    x_mask = torch.ones(b, 1, t//n_sqz).to(device=x.device, dtype=x.dtype)
+  return x_sqz * x_mask, x_mask
+
+
+def unsqueeze(x, x_mask=None, n_sqz=2):
+  b, c, t = x.size()
+
+  x_unsqz = x.view(b, n_sqz, c//n_sqz, t)
+  x_unsqz = x_unsqz.permute(0, 2, 3, 1).contiguous().view(b, c//n_sqz, t*n_sqz)
+
+  if x_mask is not None:
+    x_mask = x_mask.unsqueeze(-1).repeat(1,1,1,n_sqz).view(b, 1, t*n_sqz)
+  else:
+    x_mask = torch.ones(b, 1, t*n_sqz).to(device=x.device, dtype=x.dtype)
+  return x_unsqz * x_mask, x_mask
+