DDPM-based Pitch Generation and Pitch Controllable Voice Synthesis.

import math
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn.utils import remove_weight_norm, weight_norm

import commons
from modules import LayerNorm
   
class Encoder_VITS2(nn.Module): #backward compatible vits2 encoder
  def __init__(
      self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs
      ):
    super().__init__()
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.window_size = window_size

    self.drop = nn.Dropout(p_dropout)
    self.attn_layers = nn.ModuleList()
    self.norm_layers_1 = nn.ModuleList()
    self.ffn_layers = nn.ModuleList()
    self.norm_layers_2 = nn.ModuleList()
    # if kwargs has spk_emb_dim, then add a linear layer to project spk_emb_dim to hidden_channels
    self.cond_layer_idx = self.n_layers
    if 'gin_channels' in kwargs:
      self.gin_channels = kwargs['gin_channels']
      if self.gin_channels != 0:
        self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
        # vits2 says 3rd block, so idx is 2 by default
        self.cond_layer_idx = kwargs['cond_layer_idx'] if 'cond_layer_idx' in kwargs else 2
        assert self.cond_layer_idx < self.n_layers, 'cond_layer_idx should be less than n_layers'
      
    for i in range(self.n_layers):
      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
      self.norm_layers_1.append(LayerNorm(hidden_channels))
      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
      self.norm_layers_2.append(LayerNorm(hidden_channels))

  def forward(self, x, x_mask, g=None):
    attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
    x = x * x_mask
    for i in range(self.n_layers):
      if i == self.cond_layer_idx and g is not None:
        g = self.spk_emb_linear(g.transpose(1, 2))
        g = g.transpose(1, 2)
        x = x + g
        x = x * x_mask
      y = self.attn_layers[i](x, x, attn_mask)
      y = self.drop(y)
      x = self.norm_layers_1[i](x + y)

      y = self.ffn_layers[i](x, x_mask)
      y = self.drop(y)
      x = self.norm_layers_2[i](x + y)
    x = x * x_mask
    return x

class Encoder(nn.Module):
  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
    super().__init__()
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.window_size = window_size

    self.drop = nn.Dropout(p_dropout)
    self.attn_layers = nn.ModuleList()
    self.norm_layers_1 = nn.ModuleList()
    self.ffn_layers = nn.ModuleList()
    self.norm_layers_2 = nn.ModuleList()
    for i in range(self.n_layers):
      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
      self.norm_layers_1.append(LayerNorm(hidden_channels))
      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
      self.norm_layers_2.append(LayerNorm(hidden_channels))

  def forward(self, x, x_mask):
    attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
    x = x * x_mask
    for i in range(self.n_layers):
      y = self.attn_layers[i](x, x, attn_mask)
      y = self.drop(y)
      x = self.norm_layers_1[i](x + y)

      y = self.ffn_layers[i](x, x_mask)
      y = self.drop(y)
      x = self.norm_layers_2[i](x + y)
    x = x * x_mask
    return x

class Decoder(nn.Module):
  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs):
    super().__init__()
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.proximal_bias = proximal_bias
    self.proximal_init = proximal_init

    self.drop = nn.Dropout(p_dropout)
    self.self_attn_layers = nn.ModuleList()
    self.norm_layers_0 = nn.ModuleList()
    self.encdec_attn_layers = nn.ModuleList()
    self.norm_layers_1 = nn.ModuleList()
    self.ffn_layers = nn.ModuleList()
    self.norm_layers_2 = nn.ModuleList()
    for i in range(self.n_layers):
      self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init))
      self.norm_layers_0.append(LayerNorm(hidden_channels))
      self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
      self.norm_layers_1.append(LayerNorm(hidden_channels))
      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
      self.norm_layers_2.append(LayerNorm(hidden_channels))

  def forward(self, x, x_mask, h, h_mask):
    """
    x: decoder input
    h: encoder output
    """
    self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
    encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
    x = x * x_mask
    for i in range(self.n_layers):
      y = self.self_attn_layers[i](x, x, self_attn_mask)
      y = self.drop(y)
      x = self.norm_layers_0[i](x + y)

      y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
      y = self.drop(y)
      x = self.norm_layers_1[i](x + y)
      
      y = self.ffn_layers[i](x, x_mask)
      y = self.drop(y)
      x = self.norm_layers_2[i](x + y)
    x = x * x_mask
    return x


class MultiHeadAttention(nn.Module):
  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
    super().__init__()
    assert channels % n_heads == 0

    self.channels = channels
    self.out_channels = out_channels
    self.n_heads = n_heads
    self.p_dropout = p_dropout
    self.window_size = window_size
    self.heads_share = heads_share
    self.block_length = block_length
    self.proximal_bias = proximal_bias
    self.proximal_init = proximal_init
    self.attn = None

    self.k_channels = channels // n_heads
    self.conv_q = nn.Conv1d(channels, channels, 1)
    self.conv_k = nn.Conv1d(channels, channels, 1)
    self.conv_v = nn.Conv1d(channels, channels, 1)
    self.conv_o = nn.Conv1d(channels, out_channels, 1)
    self.drop = nn.Dropout(p_dropout)

    if window_size is not None:
      n_heads_rel = 1 if heads_share else n_heads
      rel_stddev = self.k_channels**-0.5
      self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
      self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)

    nn.init.xavier_uniform_(self.conv_q.weight)
    nn.init.xavier_uniform_(self.conv_k.weight)
    nn.init.xavier_uniform_(self.conv_v.weight)
    if proximal_init:
      with torch.no_grad():
        self.conv_k.weight.copy_(self.conv_q.weight)
        self.conv_k.bias.copy_(self.conv_q.bias)
      
  def forward(self, x, c, attn_mask=None):
    q = self.conv_q(x)
    k = self.conv_k(c)
    v = self.conv_v(c)
    
    x, self.attn = self.attention(q, k, v, mask=attn_mask)

    x = self.conv_o(x)
    return x

  def attention(self, query, key, value, mask=None):
    # reshape [b, d, t] -> [b, n_h, t, d_k]
    b, d, t_s, t_t = (*key.size(), query.size(2))
    query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
    key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
    value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)

    scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
    if self.window_size is not None:
      assert t_s == t_t, "Relative attention is only available for self-attention."
      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
      rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings)
      scores_local = self._relative_position_to_absolute_position(rel_logits)
      scores = scores + scores_local
    if self.proximal_bias:
      assert t_s == t_t, "Proximal bias is only available for self-attention."
      scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
    if mask is not None:
      scores = scores.masked_fill(mask == 0, -1e4)
      if self.block_length is not None:
        assert t_s == t_t, "Local attention is only available for self-attention."
        block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
        scores = scores.masked_fill(block_mask == 0, -1e4)
    p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
    p_attn = self.drop(p_attn)
    output = torch.matmul(p_attn, value)
    if self.window_size is not None:
      relative_weights = self._absolute_position_to_relative_position(p_attn)
      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
    output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
    return output, p_attn

  def _matmul_with_relative_values(self, x, y):
    """
    x: [b, h, l, m]
    y: [h or 1, m, d]
    ret: [b, h, l, d]
    """
    ret = torch.matmul(x, y.unsqueeze(0))
    return ret

  def _matmul_with_relative_keys(self, x, y):
    """
    x: [b, h, l, d]
    y: [h or 1, m, d]
    ret: [b, h, l, m]
    """
    ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
    return ret

  def _get_relative_embeddings(self, relative_embeddings, length):
    max_relative_position = 2 * self.window_size + 1
    # Pad first before slice to avoid using cond ops.
    pad_length = max(length - (self.window_size + 1), 0)
    slice_start_position = max((self.window_size + 1) - length, 0)
    slice_end_position = slice_start_position + 2 * length - 1
    if pad_length > 0:
      padded_relative_embeddings = F.pad(
          relative_embeddings,
          commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
    else:
      padded_relative_embeddings = relative_embeddings
    used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position]
    return used_relative_embeddings

  def _relative_position_to_absolute_position(self, x):
    """
    x: [b, h, l, 2*l-1]
    ret: [b, h, l, l]
    """
    batch, heads, length, _ = x.size()
    # Concat columns of pad to shift from relative to absolute indexing.
    x = F.pad(x, commons.convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))

    # Concat extra elements so to add up to shape (len+1, 2*len-1).
    x_flat = x.view([batch, heads, length * 2 * length])
    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0,0],[0,0],[0,length-1]]))

    # Reshape and slice out the padded elements.
    x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
    return x_final

  def _absolute_position_to_relative_position(self, x):
    """
    x: [b, h, l, l]
    ret: [b, h, l, 2*l-1]
    """
    batch, heads, length, _ = x.size()
    # padd along column
    x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
    x_flat = x.view([batch, heads, length**2 + length*(length -1)])
    # add 0's in the beginning that will skew the elements after reshape
    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
    x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:]
    return x_final

  def _attention_bias_proximal(self, length):
    """Bias for self-attention to encourage attention to close positions.
    Args:
      length: an integer scalar.
    Returns:
      a Tensor with shape [1, 1, length, length]
    """
    r = torch.arange(length, dtype=torch.float32)
    diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
    return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)


class FFN(nn.Module):
  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
    super().__init__()
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.filter_channels = filter_channels
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.activation = activation
    self.causal = causal

    if causal:
      self.padding = self._causal_padding
    else:
      self.padding = self._same_padding

    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
    self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
    self.drop = nn.Dropout(p_dropout)

  def forward(self, x, x_mask):
    x = self.conv_1(self.padding(x * x_mask))
    if self.activation == "gelu":
      x = x * torch.sigmoid(1.702 * x)
    else:
      x = torch.relu(x)
    x = self.drop(x)
    x = self.conv_2(self.padding(x * x_mask))
    return x * x_mask
  
  def _causal_padding(self, x):
    if self.kernel_size == 1:
      return x
    pad_l = self.kernel_size - 1
    pad_r = 0
    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
    x = F.pad(x, commons.convert_pad_shape(padding))
    return x

  def _same_padding(self, x):
    if self.kernel_size == 1:
      return x
    pad_l = (self.kernel_size - 1) // 2
    pad_r = self.kernel_size // 2
    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
    x = F.pad(x, commons.convert_pad_shape(padding))
    return x

def weight_norm_modules(module, name = 'weight', dim = 0):
    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
      module.weight_norm()
      return module
    else:
      return weight_norm(module,name,dim)

def remove_weight_norm_modules(module, name = 'weight'):
    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
      module.remove_weight_norm()
    else:
      remove_weight_norm(module,name)
    

class FFT(nn.Module):
  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers=1, kernel_size=1, p_dropout=0.,
               proximal_bias=False, proximal_init=True, isflow = False, **kwargs):
    super().__init__()
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.proximal_bias = proximal_bias
    self.proximal_init = proximal_init
    if isflow and 'gin_channels' in kwargs and kwargs["gin_channels"] > 0:
      cond_layer = torch.nn.Conv1d(kwargs["gin_channels"], 2*hidden_channels*n_layers, 1)
      self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
      self.cond_layer = weight_norm_modules(cond_layer, name='weight')
      self.gin_channels = kwargs["gin_channels"]
    self.drop = nn.Dropout(p_dropout)
    self.self_attn_layers = nn.ModuleList()
    self.norm_layers_0 = nn.ModuleList()
    self.ffn_layers = nn.ModuleList()
    self.norm_layers_1 = nn.ModuleList()
    for i in range(self.n_layers):
      self.self_attn_layers.append(
        MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias,
                           proximal_init=proximal_init))
      self.norm_layers_0.append(LayerNorm(hidden_channels))
      self.ffn_layers.append(
        FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
      self.norm_layers_1.append(LayerNorm(hidden_channels))

  def forward(self, x, x_mask, g = None):
    """
    x: decoder input
    h: encoder output
    """
    if g is not None:
      g = self.cond_layer(g)

    self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
    x = x * x_mask
    for i in range(self.n_layers):
      if g is not None:
        x = self.cond_pre(x)
        cond_offset = i * 2 * self.hidden_channels
        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
        x = commons.fused_add_tanh_sigmoid_multiply(
          x,
          g_l,
          torch.IntTensor([self.hidden_channels]))
      y = self.self_attn_layers[i](x, x, self_attn_mask)
      y = self.drop(y)
      x = self.norm_layers_0[i](x + y)

      y = self.ffn_layers[i](x, x_mask)
      y = self.drop(y)
      x = self.norm_layers_1[i](x + y)
    x = x * x_mask
    return x
  

class Depthwise_Separable_Conv1D(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride = 1,
        padding = 0,
        dilation = 1,
        bias = True,
        padding_mode = 'zeros',  # TODO: refine this type
        device=None,
        dtype=None
    ):
      super().__init__()
      self.depth_conv = nn.Conv1d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, groups=in_channels,stride = stride,padding=padding,dilation=dilation,bias=bias,padding_mode=padding_mode,device=device,dtype=dtype)
      self.point_conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias, device=device,dtype=dtype)
    
    def forward(self, input):
      return self.point_conv(self.depth_conv(input))

    def weight_norm(self):
      self.depth_conv = weight_norm(self.depth_conv, name = 'weight')
      self.point_conv = weight_norm(self.point_conv, name = 'weight')

    def remove_weight_norm(self):
      self.depth_conv = remove_weight_norm(self.depth_conv, name = 'weight')
      self.point_conv = remove_weight_norm(self.point_conv, name = 'weight')

class Depthwise_Separable_TransposeConv1D(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride = 1,
        padding = 0, 
        output_padding = 0,
        bias = True,
        dilation = 1,
        padding_mode = 'zeros',  # TODO: refine this type
        device=None,
        dtype=None
    ):
      super().__init__()
      self.depth_conv = nn.ConvTranspose1d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, groups=in_channels,stride = stride,output_padding=output_padding,padding=padding,dilation=dilation,bias=bias,padding_mode=padding_mode,device=device,dtype=dtype)
      self.point_conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias, device=device,dtype=dtype)
    
    def forward(self, input):
      return self.point_conv(self.depth_conv(input))

    def weight_norm(self):
      self.depth_conv = weight_norm(self.depth_conv, name = 'weight')
      self.point_conv = weight_norm(self.point_conv, name = 'weight')

    def remove_weight_norm(self):
      remove_weight_norm(self.depth_conv, name = 'weight')
      remove_weight_norm(self.point_conv, name = 'weight')
attentions.py
tonnetonne814/SiFi-VITS2-44100-Ja
import math
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F


def init_weights(m, mean=0.0, std=0.01):
  classname = m.__class__.__name__
  if classname.find("Conv") != -1:
    m.weight.data.normal_(mean, std)


def get_padding(kernel_size, dilation=1):
  return int((kernel_size*dilation - dilation)/2)


def convert_pad_shape(pad_shape):
  l = pad_shape[::-1]
  pad_shape = [item for sublist in l for item in sublist]
  return pad_shape


def intersperse(lst, item):
  result = [item] * (len(lst) * 2 + 1)
  result[1::2] = lst
  return result


def kl_divergence(m_p, logs_p, m_q, logs_q):
  """KL(P||Q)"""
  kl = (logs_q - logs_p) - 0.5
  kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q)
  return kl


def rand_gumbel(shape):
  """Sample from the Gumbel distribution, protect from overflows."""
  uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
  return -torch.log(-torch.log(uniform_samples))


def rand_gumbel_like(x):
  g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
  return g
"""

def slice_segments(x, ids_str, segment_size=4):
  ret = torch.zeros_like(x[:, :, :segment_size])
  for i in range(x.size(0)):
    try:
      idx_str = ids_str[i]
      idx_end = idx_str + segment_size
      ret[i] = x[i, :, idx_str:idx_end]
    except:
      idx_str = ids_str[i]
      idx_end = idx_str + segment_size
      ret[i] =  torch.cat((x[i, :, idx_str:] , torch.zeros), dim=１)
      
  return ret

"""

def slice_pitch_segments(x, ids_str, segment_size=4):
  ret = torch.zeros_like(x[:, :segment_size])
  for i in range(x.size(0)):
    idx_str = ids_str[i]
    idx_end = idx_str + segment_size
    ret[i] = x[i, idx_str:idx_end]
  return ret

def rand_slice_segments_with_pitch(x, pitch, x_lengths=None, segment_size=4):
  b, d, t = x.size()
  if x_lengths is None:
    x_lengths = t
  ids_str_max = x_lengths - segment_size + 1
  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
  ret = slice_segments(x, ids_str, segment_size)
  ret_pitch = slice_pitch_segments(pitch, ids_str, segment_size)
  return ret, ret_pitch, ids_str

def slice_segments_with_pitch(x, pitch, ids_str, segment_size=4):
  b, d, t = x.size()
  ret = slice_segments(x, ids_str, segment_size)
  ret_pitch = slice_pitch_segments(pitch, ids_str, segment_size)
  return ret, ret_pitch, ids_str

def slice_segments(x, ids_str, segment_size=4):
  ret = torch.zeros_like(x[:, :, :segment_size])
  for i in range(x.size(0)):
    idx_str = ids_str[i]
    idx_end = idx_str + segment_size
    ret[i] = x[i, :, int(idx_str):int(idx_end)]
  return ret


def rand_slice_segments(x, x_lengths=None, segment_size=4):
  b, d, t = x.size()
  if x_lengths is None:
    x_lengths = t
  ids_str_max = torch.floor((x_lengths - segment_size + 1))
  #ids_str_max = torch.floor((x_lengths - segment_size + 1) * 0.7)
  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
  ret = slice_segments(x, ids_str, segment_size)
  return ret, ids_str

def rand_slice_segments_with_sinewave(x, pitch, x_lengths=None, hop_size=512, segment_size=4):
  b, d, t = x.size()
  if x_lengths is None:
    x_lengths = t
  ids_str_max = x_lengths - segment_size + 1
  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
  ret = slice_segments(x, ids_str, segment_size)
  ret_pitch = slice_pitch_segments(pitch, ids_str*hop_size, segment_size*hop_size)
  return ret, ret_pitch.unsqueeze(1), ids_str

def dfs_slice_segment(dfs, ids_str, upscales, segment_size=4):
  dfs_out = []

  for x, scale in zip(dfs, upscales):
    ids_str     *= scale
    segment_size*= scale
    ret = slice_segments(x, ids_str, segment_size)
    dfs_out += [ret]
  return dfs_out


def last_slice_segments(x, x_lengths=None, segment_size=4):
  b, d, t = x.size()
  if x_lengths is None:
    x_lengths = t
  ids_str_max = torch.floor((x_lengths - segment_size + 1))
  #ids_str_max = torch.floor((x_lengths - segment_size + 1) * 0.7)
  ids_str = ids_str_max - 1
  ret = slice_segments(x, ids_str, segment_size)
  return ret, ids_str



def get_timing_signal_1d(
    length, channels, min_timescale=1.0, max_timescale=1.0e4):
  position = torch.arange(length, dtype=torch.float)
  num_timescales = channels // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (num_timescales - 1))
  inv_timescales = min_timescale * torch.exp(
      torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment)
  scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
  signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
  signal = F.pad(signal, [0, 0, 0, channels % 2])
  signal = signal.view(1, channels, length)
  return signal


def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
  b, channels, length = x.size()
  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
  return x + signal.to(dtype=x.dtype, device=x.device)


def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
  b, channels, length = x.size()
  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
  return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)


def subsequent_mask(length):
  mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
  return mask


@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


def convert_pad_shape(pad_shape):
  l = pad_shape[::-1]
  pad_shape = [item for sublist in l for item in sublist]
  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 generate_path(duration, mask):
  """
  duration: [b, 1, t_x]
  mask: [b, 1, t_y, t_x]
  """
  device = duration.device
  
  b, _, t_y, t_x = mask.shape
  cum_duration = torch.cumsum(duration, -1)
  
  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.unsqueeze(1).transpose(2,3) * mask
  return path


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)
  if clip_value is not None:
    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
    if clip_value is not None:
      p.grad.data.clamp_(min=-clip_value, max=clip_value)
  total_norm = total_norm ** (1. / norm_type)
  return total_norm


def convert_logdur_to_intdur(log_duration,d_control=1):
    duration_rounded = torch.clamp(
        (torch.round(torch.exp(log_duration) - 1) * d_control),
        min=0,
    )
    return duration_rounded

def convert_intdur_to_logdur(int_dur):
    return torch.log(int_dur.float() + 1)
commons.py
tonnetonne814/SiFi-VITS2-44100-Ja
import math
import torch
from torch import nn
from torch.nn import Parameter
import torch.onnx.operators
import torch.nn.functional as F
from utils.tts_utils import make_positions, softmax, get_incremental_state, set_incremental_state


class Reshape(nn.Module):
    def __init__(self, *args):
        super(Reshape, self).__init__()
        self.shape = args

    def forward(self, x):
        return x.view(self.shape)


class Permute(nn.Module):
    def __init__(self, *args):
        super(Permute, self).__init__()
        self.args = args

    def forward(self, x):
        return x.permute(self.args)


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


def Embedding(num_embeddings, embedding_dim, padding_idx=None):
    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
    nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
    if padding_idx is not None:
        nn.init.constant_(m.weight[padding_idx], 0)
    return m


class GroupNorm1DTBC(nn.GroupNorm):
    def forward(self, input):
        return super(GroupNorm1DTBC, self).forward(input.permute(1, 2, 0)).permute(2, 0, 1)


def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
    if not export and torch.cuda.is_available():
        try:
            from apex.normalization import FusedLayerNorm
            return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
        except ImportError:
            pass
    return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)


def Linear(in_features, out_features, bias=True):
    m = nn.Linear(in_features, out_features, bias)
    nn.init.xavier_uniform_(m.weight)
    if bias:
        nn.init.constant_(m.bias, 0.)
    return m


class SinusoidalPositionalEmbedding(nn.Module):
    """This module produces sinusoidal positional embeddings of any length.

    Padding symbols are ignored.
    """

    def __init__(self, embedding_dim, padding_idx, init_size=1024):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx
        self.weights = SinusoidalPositionalEmbedding.get_embedding(
            init_size,
            embedding_dim,
            padding_idx,
        )
        self.register_buffer('_float_tensor', torch.FloatTensor(1))

    @staticmethod
    def get_embedding(num_embeddings, embedding_dim, padding_idx=None):
        """Build sinusoidal embeddings.

        This matches the implementation in tensor2tensor, but differs slightly
        from the description in Section 3.5 of "Attention Is All You Need".
        """
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if padding_idx is not None:
            emb[padding_idx, :] = 0
        return emb

    def forward(self, input, incremental_state=None, timestep=None, positions=None, **kwargs):
        """Input is expected to be of size [bsz x seqlen]."""
        bsz, seq_len = input.shape[:2]
        max_pos = self.padding_idx + 1 + seq_len
        if self.weights is None or max_pos > self.weights.size(0):
            # recompute/expand embeddings if needed
            self.weights = SinusoidalPositionalEmbedding.get_embedding(
                max_pos,
                self.embedding_dim,
                self.padding_idx,
            )
        self.weights = self.weights.to(self._float_tensor)

        if incremental_state is not None:
            # positions is the same for every token when decoding a single step
            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
            return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)

        positions = make_positions(input, self.padding_idx) if positions is None else positions
        return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()

    def max_positions(self):
        """Maximum number of supported positions."""
        return int(1e5)  # an arbitrary large number


class ConvTBC(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding=0):
        super(ConvTBC, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.padding = padding

        self.weight = torch.nn.Parameter(torch.Tensor(
            self.kernel_size, in_channels, out_channels))
        self.bias = torch.nn.Parameter(torch.Tensor(out_channels))

    def forward(self, input):
        return torch.conv_tbc(input.contiguous(), self.weight, self.bias, self.padding)


class MultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0., bias=True,
                 add_bias_kv=False, add_zero_attn=False, self_attention=False,
                 encoder_decoder_attention=False):
        super().__init__()
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim

        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
        self.scaling = self.head_dim ** -0.5

        self.self_attention = self_attention
        self.encoder_decoder_attention = encoder_decoder_attention

        assert not self.self_attention or self.qkv_same_dim, 'Self-attention requires query, key and ' \
                                                             'value to be of the same size'

        if self.qkv_same_dim:
            self.in_proj_weight = Parameter(torch.Tensor(3 * embed_dim, embed_dim))
        else:
            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))

        if bias:
            self.in_proj_bias = Parameter(torch.Tensor(3 * embed_dim))
        else:
            self.register_parameter('in_proj_bias', None)

        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        if add_bias_kv:
            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
        else:
            self.bias_k = self.bias_v = None

        self.add_zero_attn = add_zero_attn

        self.reset_parameters()

        self.enable_torch_version = False
        if hasattr(F, "multi_head_attention_forward"):
            self.enable_torch_version = True
        else:
            self.enable_torch_version = False
        self.last_attn_probs = None

    def reset_parameters(self):
        if self.qkv_same_dim:
            nn.init.xavier_uniform_(self.in_proj_weight)
        else:
            nn.init.xavier_uniform_(self.k_proj_weight)
            nn.init.xavier_uniform_(self.v_proj_weight)
            nn.init.xavier_uniform_(self.q_proj_weight)

        nn.init.xavier_uniform_(self.out_proj.weight)
        if self.in_proj_bias is not None:
            nn.init.constant_(self.in_proj_bias, 0.)
            nn.init.constant_(self.out_proj.bias, 0.)
        if self.bias_k is not None:
            nn.init.xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            nn.init.xavier_normal_(self.bias_v)

    def forward(
            self,
            query, key, value,
            key_padding_mask=None,
            incremental_state=None,
            need_weights=True,
            static_kv=False,
            attn_mask=None,
            before_softmax=False,
            need_head_weights=False,
            enc_dec_attn_constraint_mask=None,
            reset_attn_weight=None
    ):
        """Input shape: Time x Batch x Channel

        Args:
            key_padding_mask (ByteTensor, optional): mask to exclude
                keys that are pads, of shape `(batch, src_len)`, where
                padding elements are indicated by 1s.
            need_weights (bool, optional): return the attention weights,
                averaged over heads (default: False).
            attn_mask (ByteTensor, optional): typically used to
                implement causal attention, where the mask prevents the
                attention from looking forward in time (default: None).
            before_softmax (bool, optional): return the raw attention
                weights and values before the attention softmax.
            need_head_weights (bool, optional): return the attention
                weights for each head. Implies *need_weights*. Default:
                return the average attention weights over all heads.
        """
        if need_head_weights:
            need_weights = True

        tgt_len, bsz, embed_dim = query.size()
        assert embed_dim == self.embed_dim
        assert list(query.size()) == [tgt_len, bsz, embed_dim]
        if self.enable_torch_version and incremental_state is None and not static_kv and reset_attn_weight is None:
            if self.qkv_same_dim:
                return F.multi_head_attention_forward(query, key, value,
                                                      self.embed_dim, self.num_heads,
                                                      self.in_proj_weight,
                                                      self.in_proj_bias, self.bias_k, self.bias_v,
                                                      self.add_zero_attn, self.dropout,
                                                      self.out_proj.weight, self.out_proj.bias,
                                                      self.training, key_padding_mask, need_weights,
                                                      attn_mask)
            else:
                return F.multi_head_attention_forward(query, key, value,
                                                      self.embed_dim, self.num_heads,
                                                      torch.empty([0]),
                                                      self.in_proj_bias, self.bias_k, self.bias_v,
                                                      self.add_zero_attn, self.dropout,
                                                      self.out_proj.weight, self.out_proj.bias,
                                                      self.training, key_padding_mask, need_weights,
                                                      attn_mask, use_separate_proj_weight=True,
                                                      q_proj_weight=self.q_proj_weight,
                                                      k_proj_weight=self.k_proj_weight,
                                                      v_proj_weight=self.v_proj_weight)

        if incremental_state is not None:
            saved_state = self._get_input_buffer(incremental_state)
            if 'prev_key' in saved_state:
                # previous time steps are cached - no need to recompute
                # key and value if they are static
                if static_kv:
                    assert self.encoder_decoder_attention and not self.self_attention
                    key = value = None
        else:
            saved_state = None

        if self.self_attention:
            # self-attention
            q, k, v = self.in_proj_qkv(query)
        elif self.encoder_decoder_attention:
            # encoder-decoder attention
            q = self.in_proj_q(query)
            if key is None:
                assert value is None
                k = v = None
            else:
                k = self.in_proj_k(key)
                v = self.in_proj_v(key)

        else:
            q = self.in_proj_q(query)
            k = self.in_proj_k(key)
            v = self.in_proj_v(value)
        q *= self.scaling

        if self.bias_k is not None:
            assert self.bias_v is not None
            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
            if attn_mask is not None:
                attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
            if key_padding_mask is not None:
                key_padding_mask = torch.cat(
                    [key_padding_mask, key_padding_mask.new_zeros(key_padding_mask.size(0), 1)], dim=1)

        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if k is not None:
            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if v is not None:
            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)

        if saved_state is not None:
            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
            if 'prev_key' in saved_state:
                prev_key = saved_state['prev_key'].view(bsz * self.num_heads, -1, self.head_dim)
                if static_kv:
                    k = prev_key
                else:
                    k = torch.cat((prev_key, k), dim=1)
            if 'prev_value' in saved_state:
                prev_value = saved_state['prev_value'].view(bsz * self.num_heads, -1, self.head_dim)
                if static_kv:
                    v = prev_value
                else:
                    v = torch.cat((prev_value, v), dim=1)
            if 'prev_key_padding_mask' in saved_state and saved_state['prev_key_padding_mask'] is not None:
                prev_key_padding_mask = saved_state['prev_key_padding_mask']
                if static_kv:
                    key_padding_mask = prev_key_padding_mask
                else:
                    key_padding_mask = torch.cat((prev_key_padding_mask, key_padding_mask), dim=1)

            saved_state['prev_key'] = k.view(bsz, self.num_heads, -1, self.head_dim)
            saved_state['prev_value'] = v.view(bsz, self.num_heads, -1, self.head_dim)
            saved_state['prev_key_padding_mask'] = key_padding_mask

            self._set_input_buffer(incremental_state, saved_state)

        src_len = k.size(1)

        # This is part of a workaround to get around fork/join parallelism
        # not supporting Optional types.
        if key_padding_mask is not None and key_padding_mask.shape == torch.Size([]):
            key_padding_mask = None

        if key_padding_mask is not None:
            assert key_padding_mask.size(0) == bsz
            assert key_padding_mask.size(1) == src_len

        if self.add_zero_attn:
            src_len += 1
            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
            if attn_mask is not None:
                attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
            if key_padding_mask is not None:
                key_padding_mask = torch.cat(
                    [key_padding_mask, torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask)], dim=1)

        attn_weights = torch.bmm(q, k.transpose(1, 2))
        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)

        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]

        if attn_mask is not None:
            if len(attn_mask.shape) == 2:
                attn_mask = attn_mask.unsqueeze(0)
            elif len(attn_mask.shape) == 3:
                attn_mask = attn_mask[:, None].repeat([1, self.num_heads, 1, 1]).reshape(
                    bsz * self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights + attn_mask

        if enc_dec_attn_constraint_mask is not None:  # bs x head x L_kv
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.masked_fill(
                enc_dec_attn_constraint_mask.unsqueeze(2).bool(),
                -1e8,
            )
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if key_padding_mask is not None:
            # don't attend to padding symbols
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2),
                -1e8,
            )
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        attn_logits = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)

        if before_softmax:
            return attn_weights, v

        attn_weights_float = softmax(attn_weights, dim=-1)
        attn_weights = attn_weights_float.type_as(attn_weights)
        attn_probs = F.dropout(attn_weights_float.type_as(attn_weights), p=self.dropout, training=self.training)

        if reset_attn_weight is not None:
            if reset_attn_weight:
                self.last_attn_probs = attn_probs.detach()
            else:
                assert self.last_attn_probs is not None
                attn_probs = self.last_attn_probs
        attn = torch.bmm(attn_probs, v)
        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
        attn = self.out_proj(attn)

        if need_weights:
            attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0)
            if not need_head_weights:
                # average attention weights over heads
                attn_weights = attn_weights.mean(dim=0)
        else:
            attn_weights = None

        return attn, (attn_weights, attn_logits)

    def in_proj_qkv(self, query):
        return self._in_proj(query).chunk(3, dim=-1)

    def in_proj_q(self, query):
        if self.qkv_same_dim:
            return self._in_proj(query, end=self.embed_dim)
        else:
            bias = self.in_proj_bias
            if bias is not None:
                bias = bias[:self.embed_dim]
            return F.linear(query, self.q_proj_weight, bias)

    def in_proj_k(self, key):
        if self.qkv_same_dim:
            return self._in_proj(key, start=self.embed_dim, end=2 * self.embed_dim)
        else:
            weight = self.k_proj_weight
            bias = self.in_proj_bias
            if bias is not None:
                bias = bias[self.embed_dim:2 * self.embed_dim]
            return F.linear(key, weight, bias)

    def in_proj_v(self, value):
        if self.qkv_same_dim:
            return self._in_proj(value, start=2 * self.embed_dim)
        else:
            weight = self.v_proj_weight
            bias = self.in_proj_bias
            if bias is not None:
                bias = bias[2 * self.embed_dim:]
            return F.linear(value, weight, bias)

    def _in_proj(self, input, start=0, end=None):
        weight = self.in_proj_weight
        bias = self.in_proj_bias
        weight = weight[start:end, :]
        if bias is not None:
            bias = bias[start:end]
        return F.linear(input, weight, bias)

    def _get_input_buffer(self, incremental_state):
        return get_incremental_state(
            self,
            incremental_state,
            'attn_state',
        ) or {}

    def _set_input_buffer(self, incremental_state, buffer):
        set_incremental_state(
            self,
            incremental_state,
            'attn_state',
            buffer,
        )

    def apply_sparse_mask(self, attn_weights, tgt_len, src_len, bsz):
        return attn_weights

    def clear_buffer(self, incremental_state=None):
        if incremental_state is not None:
            saved_state = self._get_input_buffer(incremental_state)
            if 'prev_key' in saved_state:
                del saved_state['prev_key']
            if 'prev_value' in saved_state:
                del saved_state['prev_value']
            self._set_input_buffer(incremental_state, saved_state)


class Swish(torch.autograd.Function):
    @staticmethod
    def forward(ctx, i):
        result = i * torch.sigmoid(i)
        ctx.save_for_backward(i)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        i = ctx.saved_variables[0]
        sigmoid_i = torch.sigmoid(i)
        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))


class CustomSwish(nn.Module):
    def forward(self, input_tensor):
        return Swish.apply(input_tensor)


class TransformerFFNLayer(nn.Module):
    def __init__(self, hidden_size, filter_size, padding="SAME", kernel_size=1, dropout=0., act='gelu'):
        super().__init__()
        self.kernel_size = kernel_size
        self.dropout = dropout
        self.act = act
        if padding == 'SAME':
            self.ffn_1 = nn.Conv1d(hidden_size, filter_size, kernel_size, padding=kernel_size // 2)
        elif padding == 'LEFT':
            self.ffn_1 = nn.Sequential(
                nn.ConstantPad1d((kernel_size - 1, 0), 0.0),
                nn.Conv1d(hidden_size, filter_size, kernel_size)
            )
        self.ffn_2 = Linear(filter_size, hidden_size)
        if self.act == 'swish':
            self.swish_fn = CustomSwish()

    def forward(self, x, incremental_state=None):
        # x: T x B x C
        if incremental_state is not None:
            saved_state = self._get_input_buffer(incremental_state)
            if 'prev_input' in saved_state:
                prev_input = saved_state['prev_input']
                x = torch.cat((prev_input, x), dim=0)
            x = x[-self.kernel_size:]
            saved_state['prev_input'] = x
            self._set_input_buffer(incremental_state, saved_state)

        x = self.ffn_1(x.permute(1, 2, 0)).permute(2, 0, 1)
        x = x * self.kernel_size ** -0.5

        if incremental_state is not None:
            x = x[-1:]
        if self.act == 'gelu':
            x = F.gelu(x)
        if self.act == 'relu':
            x = F.relu(x)
        if self.act == 'swish':
            x = self.swish_fn(x)
        x = F.dropout(x, self.dropout, training=self.training)
        x = self.ffn_2(x)
        return x

    def _get_input_buffer(self, incremental_state):
        return get_incremental_state(
            self,
            incremental_state,
            'f',
        ) or {}

    def _set_input_buffer(self, incremental_state, buffer):
        set_incremental_state(
            self,
            incremental_state,
            'f',
            buffer,
        )

    def clear_buffer(self, incremental_state):
        if incremental_state is not None:
            saved_state = self._get_input_buffer(incremental_state)
            if 'prev_input' in saved_state:
                del saved_state['prev_input']
            self._set_input_buffer(incremental_state, saved_state)


class BatchNorm1dTBC(nn.Module):
    def __init__(self, c):
        super(BatchNorm1dTBC, self).__init__()
        self.bn = nn.BatchNorm1d(c)

    def forward(self, x):
        """

        :param x: [T, B, C]
        :return: [T, B, C]
        """
        x = x.permute(1, 2, 0)  # [B, C, T]
        x = self.bn(x)  # [B, C, T]
        x = x.permute(2, 0, 1)  # [T, B, C]
        return x


class EncSALayer(nn.Module):
    def __init__(self, c, num_heads, dropout, attention_dropout=0.1,
                 relu_dropout=0.1, kernel_size=9, padding='SAME', norm='ln', act='gelu'):
        super().__init__()
        self.c = c
        self.dropout = dropout
        self.num_heads = num_heads
        if num_heads > 0:
            if norm == 'ln':
                self.layer_norm1 = LayerNorm(c)
            elif norm == 'bn':
                self.layer_norm1 = BatchNorm1dTBC(c)
            elif norm == 'gn':
                self.layer_norm1 = GroupNorm1DTBC(8, c)
            self.self_attn = MultiheadAttention(
                self.c, num_heads, self_attention=True, dropout=attention_dropout, bias=False)
        if norm == 'ln':
            self.layer_norm2 = LayerNorm(c)
        elif norm == 'bn':
            self.layer_norm2 = BatchNorm1dTBC(c)
        elif norm == 'gn':
            self.layer_norm2 = GroupNorm1DTBC(8, c)
        self.ffn = TransformerFFNLayer(
            c, 4 * c, kernel_size=kernel_size, dropout=relu_dropout, padding=padding, act=act)

    def forward(self, x, encoder_padding_mask=None, **kwargs):
        layer_norm_training = kwargs.get('layer_norm_training', None)
        if layer_norm_training is not None:
            self.layer_norm1.training = layer_norm_training
            self.layer_norm2.training = layer_norm_training
        if self.num_heads > 0:
            residual = x
            x = self.layer_norm1(x)
            x, _, = self.self_attn(
                query=x,
                key=x,
                value=x,
                key_padding_mask=encoder_padding_mask
            )
            x = F.dropout(x, self.dropout, training=self.training)
            x = residual + x
            x = x * (1 - encoder_padding_mask.float()).transpose(0, 1)[..., None]

        residual = x
        x = self.layer_norm2(x)
        x = self.ffn(x)
        x = F.dropout(x, self.dropout, training=self.training)
        x = residual + x
        x = x * (1 - encoder_padding_mask.float()).transpose(0, 1)[..., None]
        return x


class DecSALayer(nn.Module):
    def __init__(self, c, num_heads, dropout, attention_dropout=0.1, relu_dropout=0.1,
                 kernel_size=9, act='gelu', norm='ln'):
        super().__init__()
        self.c = c
        self.dropout = dropout
        if norm == 'ln':
            self.layer_norm1 = LayerNorm(c)
        elif norm == 'gn':
            self.layer_norm1 = GroupNorm1DTBC(8, c)
        self.self_attn = MultiheadAttention(
            c, num_heads, self_attention=True, dropout=attention_dropout, bias=False
        )
        if norm == 'ln':
            self.layer_norm2 = LayerNorm(c)
        elif norm == 'gn':
            self.layer_norm2 = GroupNorm1DTBC(8, c)
        self.encoder_attn = MultiheadAttention(
            c, num_heads, encoder_decoder_attention=True, dropout=attention_dropout, bias=False,
        )
        if norm == 'ln':
            self.layer_norm3 = LayerNorm(c)
        elif norm == 'gn':
            self.layer_norm3 = GroupNorm1DTBC(8, c)
        self.ffn = TransformerFFNLayer(
            c, 4 * c, padding='LEFT', kernel_size=kernel_size, dropout=relu_dropout, act=act)

    def forward(
            self,
            x,
            encoder_out=None,
            encoder_padding_mask=None,
            incremental_state=None,
            self_attn_mask=None,
            self_attn_padding_mask=None,
            attn_out=None,
            reset_attn_weight=None,
            **kwargs,
    ):
        layer_norm_training = kwargs.get('layer_norm_training', None)
        if layer_norm_training is not None:
            self.layer_norm1.training = layer_norm_training
            self.layer_norm2.training = layer_norm_training
            self.layer_norm3.training = layer_norm_training
        residual = x
        x = self.layer_norm1(x)
        x, _ = self.self_attn(
            query=x,
            key=x,
            value=x,
            key_padding_mask=self_attn_padding_mask,
            incremental_state=incremental_state,
            attn_mask=self_attn_mask
        )
        x = F.dropout(x, self.dropout, training=self.training)
        x = residual + x

        attn_logits = None
        if encoder_out is not None or attn_out is not None:
            residual = x
            x = self.layer_norm2(x)
        if encoder_out is not None:
            x, attn = self.encoder_attn(
                query=x,
                key=encoder_out,
                value=encoder_out,
                key_padding_mask=encoder_padding_mask,
                incremental_state=incremental_state,
                static_kv=True,
                enc_dec_attn_constraint_mask=get_incremental_state(self, incremental_state,
                                                                   'enc_dec_attn_constraint_mask'),
                reset_attn_weight=reset_attn_weight
            )
            attn_logits = attn[1]
        elif attn_out is not None:
            x = self.encoder_attn.in_proj_v(attn_out)
        if encoder_out is not None or attn_out is not None:
            x = F.dropout(x, self.dropout, training=self.training)
            x = residual + x

        residual = x
        x = self.layer_norm3(x)
        x = self.ffn(x, incremental_state=incremental_state)
        x = F.dropout(x, self.dropout, training=self.training)
        x = residual + x
        return x, attn_logits

    def clear_buffer(self, input, encoder_out=None, encoder_padding_mask=None, incremental_state=None):
        self.encoder_attn.clear_buffer(incremental_state)
        self.ffn.clear_buffer(incremental_state)

    def set_buffer(self, name, tensor, incremental_state):
        return set_incremental_state(self, incremental_state, name, tensor)


class ConvBlock(nn.Module):
    def __init__(self, idim=80, n_chans=256, kernel_size=3, stride=1, norm='gn', dropout=0):
        super().__init__()
        self.conv = ConvNorm(idim, n_chans, kernel_size, stride=stride)
        self.norm = norm
        if self.norm == 'bn':
            self.norm = nn.BatchNorm1d(n_chans)
        elif self.norm == 'in':
            self.norm = nn.InstanceNorm1d(n_chans, affine=True)
        elif self.norm == 'gn':
            self.norm = nn.GroupNorm(n_chans // 16, n_chans)
        elif self.norm == 'ln':
            self.norm = LayerNorm(n_chans // 16, n_chans)
        elif self.norm == 'wn':
            self.conv = torch.nn.utils.weight_norm(self.conv.conv)
        self.dropout = nn.Dropout(dropout)
        self.relu = nn.ReLU()

    def forward(self, x):
        """

        :param x: [B, C, T]
        :return: [B, C, T]
        """
        x = self.conv(x)
        if not isinstance(self.norm, str):
            if self.norm == 'none':
                pass
            elif self.norm == 'ln':
                x = self.norm(x.transpose(1, 2)).transpose(1, 2)
            else:
                x = self.norm(x)
        x = self.relu(x)
        x = self.dropout(x)
        return x


class ConvStacks(nn.Module):
    def __init__(self, idim=80, n_layers=5, n_chans=256, odim=32, kernel_size=5, norm='gn',
                 dropout=0, strides=None, res=True):
        super().__init__()
        self.conv = torch.nn.ModuleList()
        self.kernel_size = kernel_size
        self.res = res
        self.in_proj = Linear(idim, n_chans)
        if strides is None:
            strides = [1] * n_layers
        else:
            assert len(strides) == n_layers
        for idx in range(n_layers):
            self.conv.append(ConvBlock(
                n_chans, n_chans, kernel_size, stride=strides[idx], norm=norm, dropout=dropout))
        self.out_proj = Linear(n_chans, odim)

    def forward(self, x, return_hiddens=False):
        """

        :param x: [B, T, H]
        :return: [B, T, H]
        """
        x = self.in_proj(x)
        x = x.transpose(1, -1)  # (B, idim, Tmax)
        hiddens = []
        for f in self.conv:
            x_ = f(x)
            x = x + x_ if self.res else x_  # (B, C, Tmax)
            hiddens.append(x)
        x = x.transpose(1, -1)
        x = self.out_proj(x)  # (B, Tmax, H)
        if return_hiddens:
            hiddens = torch.stack(hiddens, 1)  # [B, L, C, T]
            return x, hiddens
        return x


class ConvGlobalStacks(nn.Module):
    def __init__(self, idim=80, n_layers=5, n_chans=256, odim=32, kernel_size=5, norm='gn', dropout=0,
                 strides=[2, 2, 2, 2, 2]):
        super().__init__()
        self.conv = torch.nn.ModuleList()
        self.pooling = torch.nn.ModuleList()
        self.kernel_size = kernel_size
        self.in_proj = Linear(idim, n_chans)
        for idx in range(n_layers):
            self.conv.append(ConvBlock(n_chans, n_chans, kernel_size, stride=strides[idx],
                                       norm=norm, dropout=dropout))
            self.pooling.append(nn.MaxPool1d(strides[idx]))
        self.out_proj = Linear(n_chans, odim)

    def forward(self, x):
        """

        :param x: [B, T, H]
        :return: [B, T, H]
        """
        x = self.in_proj(x)
        x = x.transpose(1, -1)  # (B, idim, Tmax)
        for f, p in zip(self.conv, self.pooling):
            x = f(x)  # (B, C, T)
        x = x.transpose(1, -1)
        x = self.out_proj(x.mean(1))  # (B, H)
        return x


class ConvDecoder(nn.Module):
    def __init__(self, c, dropout, kernel_size=9, act='gelu'):
        super().__init__()
        self.c = c
        self.dropout = dropout

        self.pre_convs = nn.ModuleList()
        self.pre_lns = nn.ModuleList()
        for i in range(2):
            self.pre_convs.append(TransformerFFNLayer(
                c, c * 2, padding='LEFT', kernel_size=kernel_size, dropout=dropout, act=act))
            self.pre_lns.append(LayerNorm(c))

        self.layer_norm_attn = LayerNorm(c)
        self.encoder_attn = MultiheadAttention(c, 1, encoder_decoder_attention=True, bias=False)

        self.post_convs = nn.ModuleList()
        self.post_lns = nn.ModuleList()
        for i in range(8):
            self.post_convs.append(TransformerFFNLayer(
                c, c * 2, padding='LEFT', kernel_size=kernel_size, dropout=dropout, act=act))
            self.post_lns.append(LayerNorm(c))

    def forward(
            self,
            x,
            encoder_out=None,
            encoder_padding_mask=None,
            incremental_state=None,
            **kwargs,
    ):
        attn_logits = None
        for conv, ln in zip(self.pre_convs, self.pre_lns):
            residual = x
            x = ln(x)
            x = conv(x) + residual
        if encoder_out is not None:
            residual = x
            x = self.layer_norm_attn(x)
            x, attn = self.encoder_attn(
                query=x,
                key=encoder_out,
                value=encoder_out,
                key_padding_mask=encoder_padding_mask,
                incremental_state=incremental_state,
                static_kv=True,
                enc_dec_attn_constraint_mask=get_incremental_state(self, incremental_state,
                                                                   'enc_dec_attn_constraint_mask'),
            )
            attn_logits = attn[1]
            x = F.dropout(x, self.dropout, training=self.training)
            x = residual + x
        for conv, ln in zip(self.post_convs, self.post_lns):
            residual = x
            x = ln(x)
            x = conv(x) + residual
        return x, attn_logits

    def clear_buffer(self, input, encoder_out=None, encoder_padding_mask=None, incremental_state=None):
        self.encoder_attn.clear_buffer(incremental_state)
        self.ffn.clear_buffer(incremental_state)

    def set_buffer(self, name, tensor, incremental_state):
        return set_incremental_state(self, incremental_state, name, tensor)

commons/common_layers.py
tonnetonne814/SiFi-VITS2-44100-Ja
import time
import os
import random
import numpy as np
import torch
import torch.utils.data
from tqdm import tqdm
import torchaudio
from singDB_loader import get_g2p_dict_from_tabledata, get_g2p_dict_from_training_data
import logger, copy
import commons 
from mel_processing import spectrogram_torch
from utils import load_wav_to_torch, load_filepaths_and_text
from scipy.interpolate import interp1d
import commons 
from sifigan.utils import dilated_factor        
from sifigan.utils.features import SignalGenerator
import pyopenjtalk
import jaconv
import re

# データセット読み込み君本体
class TextAudioLoader_TTS_f0diff(torch.utils.data.Dataset):
    def __init__(self, audiopaths_and_text, hparams):

        self.small_letters = ["ぁ", "ぃ", "ぅ", "ぇ", "ぉ", 
                              "ァ", "ィ", "ゥ", "ェ", "ォ",
                              "ゃ", "ゅ", "ょ", "っ", "ゎ",
                              "ャ", "ュ", "ョ", "ッ", "ヮ"]
        self.find_desc = "|".join(self.small_letters)

        self.basepath_list = load_filepaths_and_text(audiopaths_and_text, split="|")
        
        self.sampling_rate  = hparams["sampling_rate"]
        self.hop_length     = hparams["hop_length"] 
        self.filter_length  = hparams["filter_length"]
        self.win_length     = hparams["win_length"]
        self.wav_max_ms     = hparams["wav_max_ms"]
        self.wav_min_ms     = hparams["wav_min_ms"]
        self.f0_max         = hparams["f0_max"]

        self.oto2lab, self.ph_symbol_to_id,   self.id_to_ph_symbol, \
                      self.word_symbol_to_id,self.id_to_word_symbol = get_g2p_dict_from_tabledata(table_path=hparams["oto2lab_path"], include_converter=True)

        random.seed(hparams["seed"])
        random.shuffle(self.basepath_list) # ここでシャッフルしている
        self._filter()

    # 最小以下及び最大以上のデータを弾くフィルター関数
    def _filter(self):
        filtered_list = []
        lengths = []
        for basepath in tqdm(self.basepath_list, desc="Dataset Filtering..."):
            wav, sr = torchaudio.load(basepath[0]+".wav")
            ch, wav_len = wav.shape
            ms = wav_len / sr * 1000
            if self.wav_min_ms <= ms  and ms  <= self.wav_max_ms:
                filtered_list.append(basepath)
                lengths.append(wav_len // (2 * self.hop_length))
            else:
                print(f"EXCEEDED LENGTH : {basepath[0]}")
        self.basepath_list = filtered_list
        self.lengths = lengths
        
    def get_ph_vocab_size(self):
        return len(self.ph_symbol_to_id) + 1 # mask用

    def get_ph_ID(self, text, oto2lab, ph_symbol_to_id):
        small_letters = ["ぁ", "ぃ", "ぅ", "ぇ", "ぉ", 
                         "ァ", "ィ", "ゥ", "ェ", "ォ",
                         "ゃ", "ゅ", "ょ", "っ", "ゎ",
                         "ャ", "ュ", "ョ", "ッ", "ヮ"]
        find_desc = "|".join(small_letters)
    
        kana = pyopenjtalk.g2p(text, kana=True)
        hira = " ".join(jaconv.kata2hira(kana))
        hira_back= hira
    
        skip = 0
        for idx, m in enumerate(re.finditer(find_desc, hira)):
            pos = m.span()[0]-(idx-skip)-1 
            if pos < 0:
                skip += 1
                pass
            else:
                hira = hira[:pos] + hira[pos+1:]
        hira = hira.split(" ")
        ph_list = []
        for word in hira:
            try:
                # jsut base covert rule
                if word == "・" or word == "-" or word == "−" or word == "？" :
                    continue
                if word == '．':
                    word = "。"
                if word == '乃':
                    word = "の"
                elif word == '珂':
                    ph_list += [oto2lab["か"]]
                    continue
                elif word == '是':
                    ph_list += [oto2lab["ぜ"]]
                    continue
                
                elif word == '哉':
                    ph_list += [oto2lab["や"]]
                    continue
                elif word == '孺':
                    ph_list += [oto2lab["じゅ"]]
                    continue
                
                if word == "吐":
                    ph_list += [oto2lab["ば"]]
                    ph_list += [oto2lab["き"]]
                    continue
                elif word == '熏':
                    ph_list += [oto2lab["く"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '撃':
                    ph_list += [oto2lab["げ"]]
                    ph_list += [oto2lab["き"]]
                    continue
                elif word == '哨':
                    ph_list += [oto2lab["しょ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '逓':
                    ph_list += [oto2lab["て"]]
                    ph_list += [oto2lab["い"]]
                    continue
                elif word == '騎':
                    ph_list += [oto2lab["き"]]
                    continue
                elif word == '閃':
                    ph_list += [oto2lab["せ"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '抽':
                    ph_list += [oto2lab["ちゅ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '分':
                    ph_list += [oto2lab["ぶ"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '斬':
                    ph_list += [oto2lab["ざ"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '醸':
                    ph_list += [oto2lab["じょ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '勃':
                    ph_list += [oto2lab["ぼ"]]
                    ph_list += [oto2lab["つ"]]
                    continue
                elif word == '煖':
                    ph_list += [oto2lab["え"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '四':
                    ph_list += [oto2lab["し"]]
                    continue
                elif word == '川':
                    ph_list += [oto2lab["せ"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '省':
                    ph_list += [oto2lab["しょ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '雅':
                    ph_list += [oto2lab["が"]]
                    continue
                elif word == '安':
                    ph_list += [oto2lab["あ"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '県':
                    ph_list += [oto2lab["け"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '高':
                    ph_list += [oto2lab["こ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '頤':
                    ph_list += [oto2lab["い"]]
                    continue
                elif word == '墓':
                    ph_list += [oto2lab["ぼ"]]
                    continue
                elif word == '闕':
                    ph_list += [oto2lab["け"]]
                    ph_list += [oto2lab["つ"]]
                    continue
                elif word == '窒':
                    ph_list += [oto2lab["ち"]]
                    ph_list += [oto2lab["つ"]]
                    continue
                elif word == '菌':
                    ph_list += [oto2lab["き"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '聚':
                    ph_list += [oto2lab["じゅ"]]
                    continue
                elif word == '慎':
                    ph_list += [oto2lab["し"]]
                    ph_list += [oto2lab["ん"]]
                    continue
                elif word == '果':
                    ph_list += [oto2lab["か"]]
                    continue
                elif word == '胞':
                    ph_list += [oto2lab["ほ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '溶':
                    ph_list += [oto2lab["よ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '結':
                    ph_list += [oto2lab["け"]]
                    ph_list += [oto2lab["つ"]]
                    continue
                elif word == '腔':
                    ph_list += [oto2lab["こ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '療':
                    ph_list += [oto2lab["りょ"]]
                    ph_list += [oto2lab["う"]]
                    continue
                elif word == '々': 
                    ph_list += [oto2lab["し"]]
                    continue
                elif word == '禕': 
                    ph_list += [oto2lab["い"]]
                    continue
                
                if word == 'ー':
                    c_ph = ph_list[-1]
                    if len(c_ph) != 1:
                        ph_list += c_ph[-1]
                    else:
                        ph_list += c_ph
                    continue
                
                elif word == "ーっ":
                    c_ph = ph_list[-1]
                    if len(c_ph) != 1:
                        ph_list += [c_ph[-1]+" "+ oto2lab["っ"]]
                    else:
                        ph_list += [c_ph     + " " + oto2lab["っ"]]
                    continue
                
                try:
                    word_ph = oto2lab[word]
                except:
                    txt = ""   
                    word_ph = ""   
                    # check cl
                    for w in word:
                        if w == "っ":
                            word_ph += oto2lab[txt] + " " + oto2lab[w] # pull from dict
                            txt = "" #reset
                        else:
                            txt += w # stack
                    # last pull
                    if txt != "":
                        word_ph += oto2lab[txt]
                ph_list += [word_ph]
            except:
                tmp = []
                for w in word:
                    ph_list += [oto2lab[w]]
                    tmp += [oto2lab[w]]
                print(f"[WARNING] {word} is not existed in oto2lab.")
                print(f"[WARNING] Full text is {hira_back}.")
                print(f"[WARNING] {word} is converted to {tmp}.")
    
        sequence = []
        for symbols in ph_list:
            for symbol in symbols.split(" "):
                symbol_id = ph_symbol_to_id[symbol]
                sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)

    def get_audio(self, filename):
        #print(filename)
        audio_norm, 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))

        audio_norm = audio_norm.unsqueeze(0)
        spec_filename = filename.replace(".wav", "_spec.pt")
        if os.path.exists(spec_filename):
            spec = torch.load(spec_filename)
        else:
            spec = spectrogram_torch(audio_norm, self.filter_length,
                self.sampling_rate, self.hop_length, self.win_length,
                center=False)
            spec = torch.squeeze(spec, 0)
            if spec.size(1) == -1:
                print("ERROR SPEC")
            torch.save(spec, spec_filename)
        return spec, audio_norm

    def get_item(self, basepath):

        filepath = basepath[0]
        text = basepath[1]

        # labのデータは推論時存在しない
        f0           = torch.load(filepath + "_f0.pt"          ) 
        ph_frame_dur = torch.load(filepath + "_ph_frame_dur.pt")
        speakerID    = 0 # 未実装 

        # tokenize and get duration
        ph_IDs              = self.get_ph_ID(text=text, oto2lab=self.oto2lab, ph_symbol_to_id=self.ph_symbol_to_id)
        f0_len              = len(f0)

        # padの影響で1ずれる。無問題と妄想
        if f0_len != int(torch.sum(torch.from_numpy(ph_frame_dur))):
            ph_frame_dur[-1] += 1 
            
        # 保障

        return (torch.tensor(f0          ,            dtype=torch.float32),          
                torch.tensor(ph_IDs,                    dtype=torch.int64)+1,        # maskを0とする。
                torch.tensor(ph_frame_dur,              dtype=torch.int64),   
                torch.tensor(speakerID,                 dtype=torch.int64)  )

    def __getitem__(self, index):
        return self.get_item(self.basepath_list[index])

    def __len__(self):
        return len(self.basepath_list)


class TextAudioCollate_TTS_f0diff():
    def __init__(self, return_ids=False):
        self.return_ids = return_ids

    def __call__(self, batch):
        _, ids_sorted_decreasing = torch.sort( torch.LongTensor( [x[0].size(0) for x in batch] ),dim=0, descending=True )

        max_f0_len              = max([len(x[0]) for x in batch])
        #max_vuv_len            = max([len(x[1]) for x in batch])
        max_ph_IDs_len          = max([len(x[1]) for x in batch])
        max_ph_frame_dur_len    = max([len(x[2]) for x in batch])

        f0_lengths              = torch.LongTensor(len(batch))
        #vuv_lengths             = torch.LongTensor(len(batch))
        ph_IDs_lengths          = torch.LongTensor(len(batch))
        ph_frame_dur_lengths    = torch.LongTensor(len(batch))
        spkids                  = torch.LongTensor(len(batch))

        f0_padded               = torch.FloatTensor(len(batch), max_f0_len)
        #vuv_padded              = torch.LongTensor(len(batch),  max_vuv_len) 
        ph_IDs_padded           = torch.LongTensor(len(batch),   max_ph_IDs_len)
        ph_frame_dur_padded     = torch.LongTensor(len(batch),   max_ph_frame_dur_len)

        f0_padded.zero_()
        ph_IDs_padded.zero_()
        ph_frame_dur_padded.zero_()
        spkids.zero_()

        for i in range(len(ids_sorted_decreasing)):
            row = batch[ids_sorted_decreasing[i]]
            
            f0 = row[0]
            f0_padded[i, :f0.size(0)] = f0
            f0_lengths[i] = f0.size(0)

            #vuv = row[1]
            #vuv_padded[i, :vuv.size(0)] = vuv
            #vuv_lengths[i] = vuv.size(0)
            
            ph_IDs = row[1]
            ph_IDs_padded[i, :ph_IDs.size(0)] = ph_IDs
            ph_IDs_lengths[i] = ph_IDs.size(0) 
            
            ph_frame_dur = row[2]
            ph_frame_dur_padded[i,     :ph_frame_dur.size(0)] = ph_frame_dur
            ph_frame_dur_lengths[i] =   ph_frame_dur.size(0)
            
            spkids[i] = row[3]
            

        # 次元調整
        f0_padded = torch.unsqueeze(f0_padded, dim=1)  
           
        if self.return_ids:
            return  f0_padded,              f0_lengths,             \
                    ph_IDs_padded,          ph_IDs_lengths,         \
                    ph_frame_dur_padded,                            \
                    spkids,                                   \
                    ids_sorted_decreasing    
        
        return  f0_padded,              f0_lengths,             \
                ph_IDs_padded,          ph_IDs_lengths,         \
                ph_frame_dur_padded,                            \
                spkids


# データセット読み込み君本体
class TextAudioLoader_TTS_synth(torch.utils.data.Dataset):
    def __init__(self, audiopaths_and_text, hparams):
        self.small_letters = ["ぁ", "ぃ", "ぅ", "ぇ", "ぉ", 
                              "ァ", "ィ", "ゥ", "ェ", "ォ",
                              "ゃ", "ゅ", "ょ", "っ", "ゎ",
                              "ャ", "ュ", "ョ", "ッ", "ヮ"]
        self.find_desc = "|".join(self.small_letters)

        self.basepath_list = load_filepaths_and_text(audiopaths_and_text, split="|")
        
        self.sampling_rate  = hparams["sampling_rate"]
        self.hop_length     = hparams["hop_length"] 
        self.filter_length  = hparams["filter_length"]
        self.win_length     = hparams["win_length"]
        self.wav_max_ms     = hparams["wav_max_ms"]
        self.wav_min_ms     = hparams["wav_min_ms"]
        self.f0_max         = hparams["f0_max"]

        self.df_f0_type = hparams["SiFiGAN_utils"]["df_f0_type"]
        self.dense_factors = hparams["SiFiGAN_utils"]["dense_factors"]
        self.upsample_scales= hparams["SiFiGANGenerator"]["upsample_scales"]

        # for SiFiGAN
        self.signal_generator = SignalGenerator(sample_rate=hparams["sampling_rate"],
                                                hop_size=hparams["hop_length"] ,
                                                sine_amp=hparams["SiFiGAN_utils"]["sine_amp"],
                                                noise_amp=hparams["SiFiGAN_utils"]["noise_amp"],
                                                signal_types=hparams["SiFiGAN_utils"]["signal_types"],)
        
        self.oto2lab, self.ph_symbol_to_id,   self.id_to_ph_symbol, \
                      self.word_symbol_to_id,self.id_to_word_symbol = get_g2p_dict_from_tabledata(table_path=hparams["oto2lab_path"], include_converter=True)
        random.seed(hparams["seed"])
        random.shuffle(self.basepath_list) # ここでシャッフルしている
        self._filter()

    # 最小以下及び最大以上のデータを弾くフィルター関数
    def _filter(self):
        filtered_list = []
        lengths = []
        for basepath in tqdm(self.basepath_list, desc="Dataset Filtering..."):
            wav, sr = torchaudio.load(basepath[0]+".wav")
            ch, wav_len = wav.shape
            ms = wav_len / sr * 1000
            if self.wav_min_ms <= ms  and ms  <= self.wav_max_ms:
                filtered_list.append(basepath)
                lengths.append(wav_len // (2 * self.hop_length))
            else:
                print(f"EXCEEDED LENGTH : {basepath}")
        self.basepath_list = filtered_list
        self.lengths = lengths
        
    def get_ph_vocab_size(self):
        return len(self.ph_symbol_to_id) + 1 # mask用

    def get_ph_ID(self, text):
        kana = pyopenjtalk.g2p(text, kana=True)
        hira = " ".join(jaconv.kata2hira(kana))
        for idx, m in enumerate(re.finditer(self.find_desc, hira)):
            pos = m.span()[0]-idx-1
            hira = hira[:pos] + hira[pos+1:]
        hira = hira.split(" ")

        ph_list = []
        for word in hira:
            if word == 'ー':
                c_ph = ph_list[-1]
                if len(c_ph) != 1:
                    ph_list += c_ph[-1]
                else:
                    ph_list += c_ph
                continue

            try:
                word_ph = self.oto2lab[word]
            except:
                txt = ""   
                word_ph = ""   
                # check cl
                for w in word:
                    if w == "っ":
                        word_ph += self.oto2lab[txt] + " " + self.oto2lab[w] # pull from dict
                        txt = "" #reset
                    else:
                        txt += w # stack
                # last pull
                if txt != "":
                    word_ph += self.oto2lab[txt]
            ph_list += [word_ph]
        sequence = []
        for symbols in ph_list:
            for symbol in symbols.split(" "):
                symbol_id = self.ph_symbol_to_id[symbol]
                sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)

    def get_audio(self, filename):
        #print(filename)
        audio_norm, 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))

        audio_norm = audio_norm.unsqueeze(0)
        spec_filename = filename.replace(".wav", "_spec.pt")
        if os.path.exists(spec_filename):
            spec = torch.load(spec_filename)
        else:
            spec = spectrogram_torch(audio_norm, self.filter_length,
                self.sampling_rate, self.hop_length, self.win_length,
                center=False)
            spec = torch.squeeze(spec, 0)
            if spec.size(1) == -1:
                print("ERROR SPEC")
            torch.save(spec, spec_filename)
        return spec, audio_norm

    def get_continuos_f0(self, f0):
        """Convert F0 to continuous F0
        Args:
            f0 (ndarray): original f0 sequence with the shape (T)
        Return:
            (ndarray): continuous f0 with the shape (T)
        """
        # get uv information as binary
        uv = np.float32(f0 != 0)
        # get start and end of f0
        if (f0 == 0).all():
            logger.warn("all of the f0 values are 0.")
            return uv, f0, False
        start_f0 = f0[f0 != 0][0]
        end_f0 = f0[f0 != 0][-1]
        # padding start and end of f0 sequence
        cf0 = copy.deepcopy(f0)
        start_idx = np.where(cf0 == start_f0)[0][0]
        end_idx = np.where(cf0 == end_f0)[0][-1]
        cf0[:start_idx] = start_f0
        cf0[end_idx:] = end_f0
        # get non-zero frame index
        nz_frames = np.where(cf0 != 0)[0]
        # perform linear interpolation
        f = interp1d(nz_frames, cf0[nz_frames])
        cf0 = f(np.arange(0, cf0.shape[0]))

        return uv, cf0, True

    def get_item(self, basepath):
        filepath = basepath[0]
        text = basepath[1]

        # labのデータは推論時存在しない
        spec, wav   = self.get_audio(filename= filepath+".wav") 
        f0          = torch.load(filepath + "_f0.pt"          ) 
        vuv         = torch.load(filepath + "_vuv.pt"         ) 
        speakerID   = 0 # 未実装 
        
        ph_IDs = self.get_ph_ID(text)
        _, spec_len         = spec.shape

        # paddingの影響で、1長くなることがあるので、その対策
        if  spec_len != len(vuv) or  spec_len != len(f0):
            f0 = f0[:spec_len]
            vuv = vuv[:spec_len]

        _, c_f0, _ = self.get_continuos_f0(f0) # for SiFiGAN
        Sinewave = self.signal_generator(torch.tensor(c_f0, dtype=torch.float32).view(1,1,-1))# for SiFiGAN

        prod_upsample_scales = np.cumprod(self.upsample_scales)
        df_sample_rates = [self.sampling_rate / self.hop_length * s for s in prod_upsample_scales]
        dfs = []
        for df, us in zip(self.dense_factors, prod_upsample_scales):
            dfs += [
                np.repeat(dilated_factor(c_f0, self.sampling_rate, df), us)
                if self.df_f0_type == "cf0"
                else np.repeat(dilated_factor(c_f0, self.sampling_rate, df), us)
            ]

        # 保障
        assert spec_len == len(f0)
        assert spec_len == len(vuv)
        for i in range(len(self.dense_factors)):
            assert Sinewave.size(2) * df_sample_rates[i] == len(dfs[i]) * self.sampling_rate

        return (torch.tensor(wav,                       dtype=torch.float32),    
                torch.tensor(spec,                      dtype=torch.float32),    
                torch.tensor(f0,                        dtype=torch.float32),  # diffusion用に正規化  
                torch.tensor(ph_IDs,                    dtype=torch.int64)+1,    # maskを0とする。
                torch.tensor(c_f0,                      dtype=torch.float32),
                dfs ,
                torch.tensor(speakerID,                 dtype=torch.int64))

    def __getitem__(self, index):
        return self.get_item(self.basepath_list[index])

    def __len__(self):
        return len(self.basepath_list)


class TextAudioCollate_TTS_synth():
    def __init__(self, hparams, return_ids=False):
        self.return_ids = return_ids
        self.dense_factors = hparams["SiFiGAN_utils"]["dense_factors"]
        self.upsample_scales = hparams["SiFiGANGenerator"]["upsample_scales"]
        
        # for SiFiGAN
        self.signal_generator = SignalGenerator(sample_rate=hparams["sampling_rate"],
                                                hop_size=hparams["hop_length"] ,
                                                sine_amp=hparams["SiFiGAN_utils"]["sine_amp"],
                                                noise_amp=hparams["SiFiGAN_utils"]["noise_amp"],
                                                signal_types=hparams["SiFiGAN_utils"]["signal_types"],)

    def __call__(self, batch):
        _, ids_sorted_decreasing = torch.sort( torch.LongTensor( [x[1].size(1) for x in batch] ),dim=0, descending=True )

        max_wav_len                 = max([    x[0].size(1) for x in batch])
        max_spec_len                = max([    x[1].size(1) for x in batch])
        max_f0_len                  = max([len(x[2]) for x in batch])
        max_ph_IDs_len              = max([len(x[3]) for x in batch])
        max_c_f0_len                = max([len(x[4]) for x in batch])

        spec_lengths    = torch.LongTensor(len(batch))
        ph_IDs_lengths          = torch.LongTensor(len(batch))
        spkids                  = torch.LongTensor(len(batch))

        wav_padded  = torch.FloatTensor(len(batch), 1, max_wav_len)
        spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
        f0_padded   = torch.FloatTensor(len(batch), max_f0_len)
        ph_IDs_padded       = torch.LongTensor(len(batch),   max_ph_IDs_len)
        c_f0_padded   = torch.FloatTensor(len(batch), max_c_f0_len)

        wav_padded.zero_()
        spec_padded.zero_()
        f0_padded.zero_()
        ph_IDs_padded.zero_()
        c_f0_padded   .zero_()

        dfs_padded = []
        dfs_lengths = torch.LongTensor(len(batch), len(self.upsample_scales))
        lengths=max_c_f0_len
        for idx, scales in enumerate(self.upsample_scales):
            lengths *= scales
            dfs_padded.append(torch.FloatTensor(len(batch), 1, int(lengths)).zero_())

        for i in range(len(ids_sorted_decreasing)):
            row = batch[ids_sorted_decreasing[i]]

            wav = row[0]
            wav_padded[i, :, :wav.size(1)] = wav

            spec = row[1]
            spec_padded[i, :, :spec.size(1)] = spec
            spec_lengths[i] = spec.size(1)

            f0 = row[2]
            f0_padded[i, :f0.size(0)] = f0

            ph_IDs = row[3]
            ph_IDs_padded[i, :ph_IDs.size(0)] = ph_IDs
            ph_IDs_lengths[i] = ph_IDs.size(0) 
            
            c_f0 = row[4]
            c_f0_padded[i, :c_f0.size(0)] = c_f0
            
            dfs = row[5]
            for idx in range(len(self.upsample_scales)):
                df = dfs[idx]
                dfs_padded[idx][i, 0, :len(df)] = torch.tensor(df, dtype=torch.float32)
                dfs_lengths[i, idx] = len(df)
            
            spkids[i] = row[6]

        # 次元調整
        f0_padded = torch.unsqueeze(f0_padded, dim=1)  
            
        # バッチ単位で作成
        Sinewaves = self.signal_generator(c_f0_padded.unsqueeze(1)).float()

        if self.return_ids:
            return  wav_padded,\
                    spec_padded,            spec_lengths,           \
                    f0_padded,\
                    ph_IDs_padded,          ph_IDs_lengths,         \
                    dfs_padded, \
                    Sinewaves, \
                    spkids, \
                    ids_sorted_decreasing    
        
        return  wav_padded,\
                spec_padded,            spec_lengths,           \
                f0_padded,\
                ph_IDs_padded,          ph_IDs_lengths,         \
                dfs_padded, \
                Sinewaves, \
                spkids


# データセット読み込み君本体
class TextAudioLoader_SVS_synth(torch.utils.data.Dataset):
    def __init__(self, audiopaths_and_text, hparams):

        self.basepath_list = load_filepaths_and_text(audiopaths_and_text, split="|")
        
        self.sampling_rate  = hparams["sampling_rate"]
        self.hop_length     = hparams["hop_length"] 
        self.filter_length  = hparams["filter_length"]
        self.win_length     = hparams["win_length"]
        self.wav_max_ms     = hparams["wav_max_ms"]
        self.wav_min_ms     = hparams["wav_min_ms"]
        self.f0_max         = hparams["f0_max"]

        self.df_f0_type = hparams["SiFiGAN_utils"]["df_f0_type"]
        self.dense_factors = hparams["SiFiGAN_utils"]["dense_factors"]
        self.upsample_scales= hparams["SiFiGANGenerator"]["upsample_scales"]

        # for SiFiGAN
        self.signal_generator = SignalGenerator(sample_rate=hparams["sampling_rate"],
                                                hop_size=hparams["hop_length"] ,
                                                sine_amp=hparams["SiFiGAN_utils"]["sine_amp"],
                                                noise_amp=hparams["SiFiGAN_utils"]["noise_amp"],
                                                signal_types=hparams["SiFiGAN_utils"]["signal_types"],)
        
        self.oto2lab, self.ph_symbol_to_id,   self.id_to_ph_symbol, \
                      self.word_symbol_to_id,self.id_to_word_symbol = get_g2p_dict_from_tabledata(table_path=hparams["oto2lab_path"], include_converter=True)
        random.seed(hparams["seed"])
        random.shuffle(self.basepath_list) # ここでシャッフルしている
        self._filter()

    # 最小以下及び最大以上のデータを弾くフィルター関数
    def _filter(self):
        filtered_list = []
        lengths = []
        for basepath in tqdm(self.basepath_list, desc="Dataset Filtering..."):
            wav, sr = torchaudio.load(basepath[0]+".wav")
            ch, wav_len = wav.shape
            ms = wav_len / sr * 1000
            if self.wav_min_ms <= ms  and ms  <= self.wav_max_ms:
                filtered_list.append(basepath[0])
                lengths.append(wav_len // (2 * self.hop_length))
            else:
                print(f"EXCEEDED LENGTH : {basepath[0]}")
        self.basepath_list = filtered_list
        self.lengths = lengths
        
    def get_ph_vocab_size(self):
        return len(self.ph_symbol_to_id) + 1 # mask用

    def get_ph_ID(self, ph_list):
        sequence = []
        for symbol in ph_list:
            symbol_id = self.ph_symbol_to_id[symbol]
            sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)

    def get_word_ID(self, word_list):
        sequence = []
        for symbol in word_list:
            symbol_id = self.word_symbol_to_id[symbol]
            sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)
    
    def expand_note_info(self, ph_IDs, noteID, note_dur, n_ph_pooling):
        ph_IDs_lengths = torch.tensor(ph_IDs.size(1), dtype=torch.int64)
        ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths.view(1), ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
        noteID_lengths = torch.tensor(noteID.size(1), dtype=torch.int64)
        noteID_mask = torch.unsqueeze(commons.sequence_mask(noteID_lengths.view(1), noteID.size(1)), 1).to(noteID.dtype) # [B, 1, ph_len]

        attn_mask     = torch.unsqueeze(noteID_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
        attn          = commons.generate_path(duration=torch.unsqueeze(n_ph_pooling,dim=1), mask=attn_mask)
        attn          = torch.squeeze(attn, dim=1).permute(0,2,1).float()                             # attn=[Batch, note_len,] 
        # expand
        noteID        = torch.matmul(noteID.float().unsqueeze(1), attn)                                            # to [Batch, inner_channel, ph_len] 
        note_dur      = torch.matmul(note_dur.float().unsqueeze(1), attn)                     # to [Batch, inner_channel, ph_len] 

        return noteID.view(-1), note_dur.view(-1)

    def get_audio(self, filename):
        #print(filename)
        audio_norm, 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))

        audio_norm = audio_norm.unsqueeze(0)
        spec_filename = filename.replace(".wav", "_spec.pt")
        if os.path.exists(spec_filename):
            spec = torch.load(spec_filename)
        else:
            spec = spectrogram_torch(audio_norm, self.filter_length,
                self.sampling_rate, self.hop_length, self.win_length,
                center=False)
            spec = torch.squeeze(spec, 0)
            if spec.size(1) == -1:
                print("ERROR SPEC")
            torch.save(spec, spec_filename)
        return spec, audio_norm
    
    def get_dur_frame_from_e_ms(self, e_ms):
        e_ms = torch.tensor(e_ms,dtype=torch.float32)
        #frames = torch.ceil(  (e_ms/1000)*self.sampling_rate / self.hop_length )   # 切り上げ
        frames = torch.floor(  (e_ms / 1000)*self.sampling_rate / self.hop_length )   # 切り捨て
        frames = torch.diff(frames, dim=0, prepend=frames.new_zeros(1))
        #for idx in reversed(range(len(frames))):
        #    if idx == 0:
        #        continue
        #    frames[idx] = frames[idx] - frames[idx-1]
        return torch.tensor(frames, dtype=torch.int64)
    
    def get_linspace_ph_id(self, ph_idx_in_a_word):
        for idx, n_ph in enumerate(ph_idx_in_a_word):
            if idx == 0:
                data = torch.linspace(start=int(n_ph-1), end=0, steps=int(n_ph), dtype=torch.int64)
            else:
                data = torch.concat([data, torch.linspace(start=int(n_ph-1), end=0, steps=int(n_ph), dtype=torch.int64)], dim=0)
        return data
    
    def get_frame_level_word_id(self, wordID, word_dur):
        #for idx, (id, dur) in enumerate(zip(wordID, word_dur)):
        #    if idx == 0:
        #        data = torch.ones(size=(dur,),dtype=torch.int64) * id
        #    else:
        #        data = torch.concat([data, torch.ones(size=(dur,),dtype=torch.int64) * id], dim=0)
        for idx, (id, dur) in enumerate(zip(wordID, word_dur)):
            if idx == 0:
                data = torch.ones(size=(dur,),dtype=torch.int64) * idx
            else:
                data = torch.concat([data, torch.ones(size=(dur,),dtype=torch.int64) * idx], dim=0)
        return data

    def get_continuos_f0(self, f0):
        """Convert F0 to continuous F0
        Args:
            f0 (ndarray): original f0 sequence with the shape (T)
        Return:
            (ndarray): continuous f0 with the shape (T)
        """
        # get uv information as binary
        uv = np.float32(f0 != 0)
        # get start and end of f0
        if (f0 == 0).all():
            logger.warn("all of the f0 values are 0.")
            return uv, f0, False
        start_f0 = f0[f0 != 0][0]
        end_f0 = f0[f0 != 0][-1]
        # padding start and end of f0 sequence
        cf0 = copy.deepcopy(f0)
        start_idx = np.where(cf0 == start_f0)[0][0]
        end_idx = np.where(cf0 == end_f0)[0][-1]
        cf0[:start_idx] = start_f0
        cf0[end_idx:] = end_f0
        # get non-zero frame index
        nz_frames = np.where(cf0 != 0)[0]
        # perform linear interpolation
        f = interp1d(nz_frames, cf0[nz_frames])
        cf0 = f(np.arange(0, cf0.shape[0]))

        return uv, cf0, True

    def get_ph_pooling_dur(self, ph_e, word_e):
        out = list()
        z_t = 0
        word_idx = 0
        for idx, e in enumerate(ph_e):
            idx += 1
            if word_e[word_idx] == e:
                out.append(int(idx - z_t))
                z_t = idx
                word_idx += 1
        return out

    def get_item(self, basepath):

        # labのデータは推論時存在しない
        spec, wav   = self.get_audio(filename= basepath+".wav") 
        f0          = torch.load(basepath + "_f0.pt"          )   
        word_dur_ms = torch.load(basepath + "_word_dur_ms.pt")
        speakerID   = 0 # 未実装 

        # tokenize and get duration
        ph_ids_path             = basepath + "_ph_ids.pt"
        ph_frame_dur_path       = basepath + "_ph_frame_dur.pt"
        word_frame_dur_path     = basepath + "_word_frame_dur.pt"
        ph_idx_in_a_word_path   = basepath + "_ph_idx_in_a_word.pt"
        n_ph_pooling_path       = basepath + "_n_ph_pooling.pt"
        note_ID_dur_path        = basepath + "_noteID_dur.pt"

        if os.path.exists(ph_ids_path) is True:
            ph_IDs = torch.load(ph_ids_path)
        else:
            ph = torch.load(basepath + "_ph.pt"          ) # ust or lab
            ph_IDs              = self.get_ph_ID(ph_list=ph)
            torch.save(ph_IDs, ph_ids_path)
        
        if os.path.exists(ph_frame_dur_path) is True:
            ph_frame_dur = torch.load(ph_frame_dur_path)
        else:
            ph_e_ms     = torch.load(basepath + "_ph_e_ms.pt"     )     # lab
            ph_frame_dur = self.get_dur_frame_from_e_ms(e_ms=ph_e_ms)
            torch.save(ph_frame_dur, ph_frame_dur_path)

        if os.path.exists(word_frame_dur_path) is True:
            word_frame_dur = torch.load(word_frame_dur_path)
        else:
            word_e_ms   = torch.load(basepath + "_word_e_ms.pt"   )     # lab
            word_frame_dur = self.get_dur_frame_from_e_ms(e_ms=word_e_ms)
            torch.save(word_frame_dur, word_frame_dur_path)
            
        if os.path.exists(ph_idx_in_a_word_path) is True:
            ph_idx_in_a_word = torch.load(ph_idx_in_a_word_path)
        else:
            ph_idx_in_a_word = torch.load(basepath + "_n_ph_in_a_word.pt" )   # lab
            ph_idx_in_a_word = self.get_linspace_ph_id(ph_idx_in_a_word=ph_idx_in_a_word)
            torch.save(ph_idx_in_a_word, ph_idx_in_a_word_path)

        if os.path.exists(n_ph_pooling_path) is True:
            n_ph_pooling = torch.load(n_ph_pooling_path)
        else:
            word_e_ms   = torch.load(basepath + "_word_e_ms.pt"   )     # lab
            ph_e_ms     = torch.load(basepath + "_ph_e_ms.pt"     )     # lab
            n_ph_pooling = self.get_ph_pooling_dur(ph_e=ph_e_ms, word_e=word_e_ms)
            torch.save(n_ph_pooling, n_ph_pooling_path)

        if os.path.exists(note_ID_dur_path) is True:
            note_ID_dur = torch.load(note_ID_dur_path)
            noteID, notedur = note_ID_dur
        else:
            noteID      = torch.load(basepath + "_noteID.pt"      ) # ust
            notedur     = torch.load(basepath + "_notedur.pt"     ) # ust    
            noteID, notedur     = self.expand_note_info(ph_IDs=ph_IDs.view(1,-1), 
                                                    noteID      =torch.tensor(noteID,dtype=torch.int64).view(1,-1),
                                                    note_dur    =torch.tensor(notedur,dtype=torch.int64).view(1,-1),
                                                    n_ph_pooling=torch.tensor(n_ph_pooling, dtype=torch.int64).view(1,-1))
            note_ID_dur = [noteID, notedur]
            torch.save(note_ID_dur, note_ID_dur_path)

        _, spec_len         = spec.shape
        
        # paddingの影響で、1長くなることがあるので、その対策
        if  spec_len != len(f0):
            #back = f0 # for debug
            f0 = f0[:spec_len]
        if spec_len != int(torch.sum(word_frame_dur)):
            word_frame_dur[-1] += 1 
        if spec_len != int(torch.sum(ph_frame_dur)):
            ph_frame_dur[-1] += 1 
        
        c_f0_path             = basepath+"_c_f0.pt"
        Sinewave_path             = basepath+"_Sinewave.pt"
        dfs_path             = basepath+"_dfs.pt"
        
        if os.path.exists(c_f0_path) is True:
            c_f0 = torch.load(c_f0_path)
        else:
            _, c_f0, _ = self.get_continuos_f0(f0) # for SiFiGAN
            torch.save(c_f0, c_f0_path)

        if os.path.exists(Sinewave_path) is True:
            Sinewave = torch.load(Sinewave_path)
        else:
            Sinewave = self.signal_generator(torch.tensor(c_f0, dtype=torch.float32).view(1,1,-1))# for SiFiGAN
            torch.save(Sinewave, Sinewave_path)

        if os.path.exists(Sinewave_path) is True:
            Sinewave = torch.load(Sinewave_path)
        else:
            Sinewave = self.signal_generator(torch.tensor(c_f0, dtype=torch.float32).view(1,1,-1))# for SiFiGAN
            torch.save(Sinewave, Sinewave_path)

        if os.path.exists(dfs_path) is True:
            dfs = torch.load(dfs_path)
        else:
            prod_upsample_scales = np.cumprod(self.upsample_scales)
            df_sample_rates = [self.sampling_rate / self.hop_length * s for s in prod_upsample_scales]
            dfs = []
            for df, us in zip(self.dense_factors, prod_upsample_scales):
                dfs += [
                    np.repeat(dilated_factor(c_f0, self.sampling_rate, df), us)
                    if self.df_f0_type == "cf0"
                    else np.repeat(dilated_factor(c_f0, self.sampling_rate, df), us)
                ]
            torch.save(dfs, dfs_path)
        
        # 保障 (Debug時のみ)
        # assert sum(n_ph_pooling) == len(ph_IDs)
        # assert int(torch.sum(word_frame_dur)) == int(torch.sum(ph_frame_dur))
        # assert len(ph_idx_in_a_word) == len(ph_IDs)
        # assert spec_len == int(torch.sum(word_frame_dur))
        # assert spec_len == len(f0)
        # assert len(ph_IDs) == len(noteID) # WORD ID が多い。おそらく"っ"の影響
        # for i in range(len(self.dense_factors)):
        #     assert Sinewave.size(2) * df_sample_rates[i] == len(dfs[i]) * self.sampling_rate

        return (torch.tensor(wav,                       dtype=torch.float32),    
                torch.tensor(spec,                      dtype=torch.float32),    
                torch.tensor(f0,                        dtype=torch.float32),  # diffusion用に正規化  
                torch.tensor(ph_IDs,                    dtype=torch.int64)+1,    # maskを0とする。
                torch.tensor(ph_frame_dur,              dtype=torch.int64),  
                torch.tensor(word_dur_ms,               dtype=torch.float32)/1000,  
                torch.tensor(word_frame_dur,            dtype=torch.int64),    
                torch.tensor(ph_idx_in_a_word,            dtype=torch.int64)+1,     # maskを0とする。
                torch.tensor(n_ph_pooling,              dtype=torch.int64),    # maskを0とする。
                torch.tensor(c_f0,                      dtype=torch.float32),
                dfs ,
                torch.tensor(speakerID,                 dtype=torch.int64))

    def __getitem__(self, index):
        return self.get_item(self.basepath_list[index])

    def __len__(self):
        return len(self.basepath_list)


class TextAudioCollate_SVS_synth():
    def __init__(self, hparams, return_ids=False):
        self.return_ids = return_ids
        self.dense_factors = hparams["SiFiGAN_utils"]["dense_factors"]
        self.upsample_scales = hparams["SiFiGANGenerator"]["upsample_scales"]
        
        # for SiFiGAN
        self.signal_generator = SignalGenerator(sample_rate=hparams["sampling_rate"],
                                                hop_size=hparams["hop_length"] ,
                                                sine_amp=hparams["SiFiGAN_utils"]["sine_amp"],
                                                noise_amp=hparams["SiFiGAN_utils"]["noise_amp"],
                                                signal_types=hparams["SiFiGAN_utils"]["signal_types"],)

    def __call__(self, batch):
        _, ids_sorted_decreasing = torch.sort( torch.LongTensor( [x[1].size(1) for x in batch] ),dim=0, descending=True )

        max_wav_len                 = max([    x[0].size(1) for x in batch])
        max_spec_len                = max([    x[1].size(1) for x in batch])
        max_f0_len                  = max([len(x[2]) for x in batch])
        max_ph_IDs_len              = max([len(x[3]) for x in batch])
        max_ph_frame_dur_len        = max([len(x[4]) for x in batch])
        max_word_dur_ms_len         = max([len(x[5]) for x in batch])
        max_word_frame_dur_len      = max([len(x[6]) for x in batch])
        max_ph_idx_in_a_word_len    = max([len(x[7]) for x in batch])
        max_n_ph_pooling_len        = max([len(x[8]) for x in batch])
        max_c_f0_len                = max([len(x[9]) for x in batch])

        spec_lengths    = torch.LongTensor(len(batch))
        ph_IDs_lengths          = torch.LongTensor(len(batch))
        word_frame_dur_lengths      = torch.LongTensor(len(batch))
        spkids                  = torch.LongTensor(len(batch))

        wav_padded  = torch.FloatTensor(len(batch), 1, max_wav_len)
        spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
        f0_padded   = torch.FloatTensor(len(batch), max_f0_len)
        ph_IDs_padded       = torch.LongTensor(len(batch),   max_ph_IDs_len)
        ph_frame_dur_padded = torch.LongTensor(len(batch),   max_ph_frame_dur_len)
        word_dur_ms_padded      = torch.FloatTensor(len(batch), max_word_dur_ms_len)
        word_frame_dur_padded   = torch.LongTensor(len(batch),  max_word_frame_dur_len)
        ph_idx_in_a_word_padded        = torch.LongTensor(len(batch), max_ph_idx_in_a_word_len)
        n_ph_pooling_padded          = torch.LongTensor(len(batch), max_n_ph_pooling_len)
        c_f0_padded   = torch.FloatTensor(len(batch), max_c_f0_len)

        wav_padded.zero_()
        spec_padded.zero_()
        f0_padded.zero_()
        ph_IDs_padded.zero_()
        ph_frame_dur_padded.zero_()
        word_dur_ms_padded.zero_()
        word_frame_dur_padded.zero_()
        ph_idx_in_a_word_padded.zero_()
        n_ph_pooling_padded.zero_()
        c_f0_padded   .zero_()

        dfs_padded = []
        dfs_lengths = torch.LongTensor(len(batch), len(self.upsample_scales))
        lengths=max_c_f0_len
        for idx, scales in enumerate(self.upsample_scales):
            lengths *= scales
            dfs_padded.append(torch.FloatTensor(len(batch), 1, int(lengths)).zero_())

        for i in range(len(ids_sorted_decreasing)):
            row = batch[ids_sorted_decreasing[i]]

            wav = row[0]
            wav_padded[i, :, :wav.size(1)] = wav

            spec = row[1]
            spec_padded[i, :, :spec.size(1)] = spec
            spec_lengths[i] = spec.size(1)

            f0 = row[2]
            f0_padded[i, :f0.size(0)] = f0

            ph_IDs = row[3]
            ph_IDs_padded[i, :ph_IDs.size(0)] = ph_IDs
            ph_IDs_lengths[i] = ph_IDs.size(0) 

            ph_frame_dur = row[4]
            ph_frame_dur_padded[i,     :ph_frame_dur.size(0)] = ph_frame_dur
            
            word_dur_ms = row[5]
            word_dur_ms_padded[i,     :word_dur_ms.size(0)] = word_dur_ms
            
            word_frame_dur = row[6]
            word_frame_dur_padded[i,     :word_frame_dur.size(0)] = word_frame_dur
            word_frame_dur_lengths[i] =   word_frame_dur.size(0)

            ph_idx_in_a_word = row[7]
            ph_idx_in_a_word_padded[i, :ph_idx_in_a_word.size(0)] = ph_idx_in_a_word
            
            n_ph_pooling = row[8]
            n_ph_pooling_padded[i, :n_ph_pooling.size(0)] = n_ph_pooling
            
            c_f0 = row[9]
            c_f0_padded[i, :c_f0.size(0)] = c_f0
            
            dfs = row[10]
            for idx in range(len(self.upsample_scales)):
                df = dfs[idx]
                dfs_padded[idx][i, 0, :len(df)] = torch.tensor(df, dtype=torch.float32)
                dfs_lengths[i, idx] = len(df)
            
            spkids[i] = row[11]

        # 次元調整
        f0_padded = torch.unsqueeze(f0_padded, dim=1)  
            
        # バッチ単位で作成
        Sinewaves = self.signal_generator(c_f0_padded.unsqueeze(1)).float()

        if self.return_ids:
            return  wav_padded,\
                    spec_padded,            spec_lengths,           \
                    f0_padded,\
                    ph_IDs_padded,          ph_IDs_lengths,         \
                    ph_frame_dur_padded,\
                    word_dur_ms_padded,     \
                    word_frame_dur_padded,  word_frame_dur_lengths, \
                    ph_idx_in_a_word_padded,  \
                    n_ph_pooling_padded, \
                    dfs_padded, \
                    Sinewaves, \
                    spkids, \
                    ids_sorted_decreasing    
        
        return  wav_padded,\
                spec_padded,            spec_lengths,           \
                f0_padded,\
                ph_IDs_padded,          ph_IDs_lengths,         \
                ph_frame_dur_padded,\
                word_dur_ms_padded,     \
                word_frame_dur_padded,  word_frame_dur_lengths, \
                ph_idx_in_a_word_padded,  \
                n_ph_pooling_padded, \
                dfs_padded, \
                Sinewaves, \
                spkids

# データセット読み込み君本体
class TextAudioLoader_SVS_f0diff(torch.utils.data.Dataset):
    def __init__(self, audiopaths_and_text, hparams):

        self.basepath_list = load_filepaths_and_text(audiopaths_and_text, split="|")
        
        self.sampling_rate  = hparams["sampling_rate"]
        self.hop_length     = hparams["hop_length"] 
        self.filter_length  = hparams["filter_length"]
        self.win_length     = hparams["win_length"]
        self.wav_max_ms     = hparams["wav_max_ms"]
        self.wav_min_ms     = hparams["wav_min_ms"]
        self.f0_max         = hparams["f0_max"]

        self.oto2lab, self.ph_symbol_to_id,   self.id_to_ph_symbol, \
                      self.word_symbol_to_id,self.id_to_word_symbol = get_g2p_dict_from_tabledata(table_path=hparams["oto2lab_path"], include_converter=True)
        with open(hparams["noteid2hz_txt_path"], mode="r", encoding="utf-8") as f:
            lines = f.readlines()
        self.id_to_hz = {}
        for idx, line in enumerate(lines):
            id, hz = line.split(",")

            self.id_to_hz[idx-1] = float(hz)

        random.seed(hparams["seed"])
        random.shuffle(self.basepath_list) # ここでシャッフルしている
        self._filter()

    # 最小以下及び最大以上のデータを弾くフィルター関数
    def _filter(self):
        filtered_list = []
        lengths = []
        for basepath in tqdm(self.basepath_list, desc="Dataset Filtering..."):
            wav, sr = torchaudio.load(basepath[0]+".wav")
            ch, wav_len = wav.shape
            ms = wav_len / sr * 1000
            if self.wav_min_ms <= ms  and ms  <= self.wav_max_ms:
                filtered_list.append(basepath[0])
                lengths.append(wav_len // (2 * self.hop_length))
            else:
                print(f"EXCEEDED LENGTH : {basepath[0]}")
        self.basepath_list = filtered_list
        self.lengths = lengths
        
    def get_ph_vocab_size(self):
        return len(self.ph_symbol_to_id) + 1 # mask用

    def get_ph_ID(self, ph_list):
        sequence = []
        for symbol in ph_list:
            symbol_id = self.ph_symbol_to_id[symbol]
            sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)

    def get_word_ID(self, word_list):
        sequence = []
        for symbol in word_list:
            symbol_id = self.word_symbol_to_id[symbol]
            sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)

    def expand_note_info(self, ph_IDs, noteID, note_dur, n_ph_pooling):
        ph_IDs_lengths = torch.tensor(ph_IDs.size(1), dtype=torch.int64)
        ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths.view(1), ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
        noteID_lengths = torch.tensor(noteID.size(1), dtype=torch.int64)
        noteID_mask = torch.unsqueeze(commons.sequence_mask(noteID_lengths.view(1), noteID.size(1)), 1).to(noteID.dtype) # [B, 1, ph_len]

        attn_mask     = torch.unsqueeze(noteID_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
        attn          = commons.generate_path(duration=torch.unsqueeze(n_ph_pooling,dim=1), mask=attn_mask)
        attn          = torch.squeeze(attn, dim=1).permute(0,2,1).float()                             # attn=[Batch, note_len,] 
        # expand
        noteID        = torch.matmul(noteID.float().unsqueeze(1), attn)                                            # to [Batch, inner_channel, ph_len] 
        note_dur      = torch.matmul(note_dur.float().unsqueeze(1), attn)                     # to [Batch, inner_channel, ph_len] 

        return noteID.view(-1), note_dur.view(-1)

    def get_audio(self, filename):
        #print(filename)
        audio_norm, 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))

        audio_norm = audio_norm.unsqueeze(0)
        spec_filename = filename.replace(".wav", "_spec.pt")
        if os.path.exists(spec_filename):
            spec = torch.load(spec_filename)
        else:
            spec = spectrogram_torch(audio_norm, self.filter_length,
                self.sampling_rate, self.hop_length, self.win_length,
                center=False)
            spec = torch.squeeze(spec, 0)
            if spec.size(1) == -1:
                print("ERROR SPEC")
            torch.save(spec, spec_filename)
        return spec, audio_norm
    
    def get_dur_frame_from_e_ms(self, e_ms):
        e_ms = torch.tensor(e_ms,dtype=torch.float32)
        #frames = torch.ceil(  (e_ms/1000)*self.sampling_rate / self.hop_length )   # 切り上げ
        frames = torch.floor(  (e_ms / 1000)*self.sampling_rate / self.hop_length )   # 切り捨て
        frames = torch.diff(frames, dim=0, prepend=frames.new_zeros(1))
        #for idx in reversed(range(len(frames))):
        #    if idx == 0:
        #        continue
        #    frames[idx] = frames[idx] - frames[idx-1]
        return torch.tensor(frames, dtype=torch.int64)

    def get_ph_pooling_dur(self, ph_e, word_e):
        out = list()
        z_t = 0
        word_idx = 0
        for idx, e in enumerate(ph_e):
            idx += 1
            if word_e[word_idx] == e:
                out.append(int(idx - z_t))
                z_t = idx
                word_idx += 1
        return out

    def get_item(self, basepath):

        # labのデータは推論時存在しない
        f0          = torch.load(basepath + "_f0.pt"          ) 
        ph          = torch.load(basepath + "_ph.pt"          ) # ust or lab
        ph_e_ms     = torch.load(basepath + "_ph_e_ms.pt"     )     # lab
        #word        = torch.load(basepath + "_word.pt"        ) # ust
        #word_dur_ms = torch.load(basepath + "_word_dur_ms.pt" )     # lab
        word_e_ms   = torch.load(basepath + "_word_e_ms.pt"   )     # lab
        noteID      = torch.load(basepath + "_noteID.pt"      ) # ust
        notedur     = torch.load(basepath + "_notedur.pt"     ) # ust   
        speakerID   = 0 # 未実装 

        # tokenize and get duration
        ph_IDs              = self.get_ph_ID(ph_list=ph)
        #word_IDs            = self.get_word_ID(word_list=word)
        ph_frame_dur        = self.get_dur_frame_from_e_ms(e_ms=ph_e_ms)
        n_ph_pooling        = self.get_ph_pooling_dur(ph_e=ph_e_ms, word_e=word_e_ms)
        f0_len            = len(f0)
        notes = []
        for id in noteID:
            notes += [self.id_to_hz[int(id)]]
        noteID, notedur     = self.expand_note_info(ph_IDs=ph_IDs.view(1,-1), 
                                                    noteID      =torch.tensor(notes, dtype=torch.float32).view(1,-1),
                                                    note_dur    =torch.tensor(notedur,dtype=torch.int64).view(1,-1),
                                                    n_ph_pooling=torch.tensor(n_ph_pooling, dtype=torch.int64).view(1,-1))
        
        # padの影響で1ずれる。無問題と妄想
        if f0_len != int(torch.sum(ph_frame_dur)):
            ph_frame_dur[-1] += 1 
            
        # 保障
        assert sum(n_ph_pooling) == len(ph_IDs)
        #assert spec_len == len(vuv)
        assert len(ph_IDs) == len(noteID) 
        #assert len(word_IDs) == len(word_dur_ms)
        #assert len(word_IDs) == len(word_e_ms)

        return (torch.tensor(f0          ,            dtype=torch.float32),          
                #torch.tensor(vuv,                       dtype=torch.int64)+1,       # maskを0とする。
                torch.tensor(ph_IDs,                    dtype=torch.int64)+1,        # maskを0とする。
                torch.tensor(ph_frame_dur,              dtype=torch.int64),   
                #torch.tensor(word_dur_ms,               dtype=torch.float32) / 1000, # ここで秒になる
                torch.tensor(noteID,                    dtype=torch.float32) ,       # maskを0とする。z
                torch.tensor(speakerID,                 dtype=torch.int64)  )

    def __getitem__(self, index):
        return self.get_item(self.basepath_list[index])

    def __len__(self):
        return len(self.basepath_list)


class TextAudioCollate_SVS_f0diff():
    def __init__(self, return_ids=False):
        self.return_ids = return_ids

    def __call__(self, batch):
        _, ids_sorted_decreasing = torch.sort( torch.LongTensor( [x[0].size(0) for x in batch] ),dim=0, descending=True )

        max_f0_len              = max([len(x[0]) for x in batch])
        #max_vuv_len            = max([len(x[1]) for x in batch])
        max_ph_IDs_len          = max([len(x[1]) for x in batch])
        max_ph_frame_dur_len    = max([len(x[2]) for x in batch])
        max_noteID_len          = max([len(x[3]) for x in batch])

        f0_lengths              = torch.LongTensor(len(batch))
        #vuv_lengths             = torch.LongTensor(len(batch))
        ph_IDs_lengths          = torch.LongTensor(len(batch))
        ph_frame_dur_lengths    = torch.LongTensor(len(batch))
        noteID_lengths          = torch.LongTensor(len(batch))
        spkids                  = torch.LongTensor(len(batch))

        f0_padded               = torch.FloatTensor(len(batch), max_f0_len)
        #vuv_padded              = torch.LongTensor(len(batch),  max_vuv_len) 
        ph_IDs_padded           = torch.LongTensor(len(batch),   max_ph_IDs_len)
        ph_frame_dur_padded     = torch.LongTensor(len(batch),   max_ph_frame_dur_len)
        noteID_padded           = torch.FloatTensor(len(batch), max_noteID_len)

        f0_padded.zero_()
        ph_IDs_padded.zero_()
        ph_frame_dur_padded.zero_()
        noteID_padded.zero_()
        spkids.zero_()

        for i in range(len(ids_sorted_decreasing)):
            row = batch[ids_sorted_decreasing[i]]
            
            f0 = row[0]
            f0_padded[i, :f0.size(0)] = f0
            f0_lengths[i] = f0.size(0)

            #vuv = row[1]
            #vuv_padded[i, :vuv.size(0)] = vuv
            #vuv_lengths[i] = vuv.size(0)
            
            ph_IDs = row[1]
            ph_IDs_padded[i, :ph_IDs.size(0)] = ph_IDs
            ph_IDs_lengths[i] = ph_IDs.size(0) 
            
            ph_frame_dur = row[2]
            ph_frame_dur_padded[i,     :ph_frame_dur.size(0)] = ph_frame_dur
            ph_frame_dur_lengths[i] =   ph_frame_dur.size(0)
            
            noteID = row[3]
            noteID_padded[i, :noteID.size(0)] = noteID
            noteID_lengths[i] = noteID.size(0)

            spkids[i] = row[4]
            

        # 次元調整
        f0_padded = torch.unsqueeze(f0_padded, dim=1)  
           
        if self.return_ids:
            return  f0_padded,              f0_lengths,             \
                    ph_IDs_padded,          ph_IDs_lengths,         \
                    ph_frame_dur_padded,                            \
                    noteID_padded,          noteID_lengths,         \
                    spkids,                                   \
                    ids_sorted_decreasing    
        
        return  f0_padded,              f0_lengths,             \
                ph_IDs_padded,          ph_IDs_lengths,         \
                ph_frame_dur_padded,                            \
                noteID_padded,          noteID_lengths,         \
                spkids

# データセット読み込み君本体
class TextAudioLoader_ALL(torch.utils.data.Dataset):
    def __init__(self, audiopaths_and_text, hparams):

        self.basepath_list = load_filepaths_and_text(audiopaths_and_text, split="|")
        
        self.sampling_rate  = hparams["sampling_rate"]
        self.hop_length     = hparams["hop_length"] 
        self.filter_length  = hparams["filter_length"]
        self.win_length     = hparams["win_length"]
        self.wav_max_ms     = hparams["wav_max_ms"]
        self.wav_min_ms     = hparams["wav_min_ms"]
        self.f0_max         = hparams["f0_max"]

        self.df_f0_type = hparams["SiFiGAN_utils"]["df_f0_type"]
        self.dense_factors = hparams["SiFiGAN_utils"]["dense_factors"]
        self.upsample_scales= hparams["SiFiGANGenerator"]["upsample_scales"]

        # for SiFiGAN
        self.signal_generator = SignalGenerator(sample_rate=hparams["sampling_rate"],
                                                hop_size=hparams["hop_length"] ,
                                                sine_amp=hparams["SiFiGAN_utils"]["sine_amp"],
                                                noise_amp=hparams["SiFiGAN_utils"]["noise_amp"],
                                                signal_types=hparams["SiFiGAN_utils"]["signal_types"],)
        
        self.oto2lab, self.ph_symbol_to_id,   self.id_to_ph_symbol, \
                      self.word_symbol_to_id,self.id_to_word_symbol = get_g2p_dict_from_tabledata(table_path=hparams["oto2lab_path"], include_converter=True)
        random.seed(hparams["seed"])
        random.shuffle(self.basepath_list) # ここでシャッフルしている
        self._filter()

    # 最小以下及び最大以上のデータを弾くフィルター関数
    def _filter(self):
        filtered_list = []
        lengths = []
        for basepath in tqdm(self.basepath_list, desc="Dataset Filtering..."):
            wav, sr = torchaudio.load(basepath[0]+".wav")
            ch, wav_len = wav.shape
            ms = wav_len / sr * 1000
            if self.wav_min_ms <= ms  and ms  <= self.wav_max_ms:
                filtered_list.append(basepath[0])
                lengths.append(wav_len // (2 * self.hop_length))
            else:
                print(f"EXCEEDED LENGTH : {basepath[0]}")
        self.basepath_list = filtered_list
        self.lengths = lengths
        
    def get_ph_vocab_size(self):
        return len(self.ph_symbol_to_id) + 1 # mask用

    def get_ph_ID(self, ph_list):
        sequence = []
        for symbol in ph_list:
            symbol_id = self.ph_symbol_to_id[symbol]
            sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)
    
    def get_word_ID(self, word_list):
        sequence = []
        for symbol in word_list:
            symbol_id = self.word_symbol_to_id[symbol]
            sequence += [symbol_id]
        return torch.tensor(sequence, dtype=torch.int64)
    
    def expand_note_info(self, ph_IDs, noteID, note_dur, n_ph_pooling):
        ph_IDs_lengths = torch.tensor(ph_IDs.size(1), dtype=torch.int64)
        ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths.view(1), ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
        noteID_lengths = torch.tensor(noteID.size(1), dtype=torch.int64)
        noteID_mask = torch.unsqueeze(commons.sequence_mask(noteID_lengths.view(1), noteID.size(1)), 1).to(noteID.dtype) # [B, 1, ph_len]

        attn_mask     = torch.unsqueeze(noteID_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
        attn          = commons.generate_path(duration=torch.unsqueeze(n_ph_pooling,dim=1), mask=attn_mask)
        attn          = torch.squeeze(attn, dim=1).permute(0,2,1).float()                             # attn=[Batch, note_len,] 
        # expand
        noteID        = torch.matmul(noteID.float().unsqueeze(1), attn)                                            # to [Batch, inner_channel, ph_len] 
        note_dur      = torch.matmul(note_dur.float().unsqueeze(1), attn)                     # to [Batch, inner_channel, ph_len] 

        return noteID.view(-1), note_dur.view(-1)

    def get_audio(self, filename):
        #print(filename)
        audio_norm, 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))

        audio_norm = audio_norm.unsqueeze(0)
        spec_filename = filename.replace(".wav", "_spec.pt")
        if os.path.exists(spec_filename):
            spec = torch.load(spec_filename)
        else:
            spec = spectrogram_torch(audio_norm, self.filter_length,
                self.sampling_rate, self.hop_length, self.win_length,
                center=False)
            spec = torch.squeeze(spec, 0)
            if spec.size(1) == -1:
                print("ERROR SPEC")
            torch.save(spec, spec_filename)
        return spec, audio_norm
    
    def get_dur_frame_from_e_ms(self, e_ms):
        e_ms = torch.tensor(e_ms,dtype=torch.float32)
        #frames = torch.ceil(  (e_ms/1000)*self.sampling_rate / self.hop_length )   # 切り上げ
        frames = torch.floor(  (e_ms / 1000)*self.sampling_rate / self.hop_length )   # 切り捨て
        frames = torch.diff(frames, dim=0, prepend=frames.new_zeros(1))
        #for idx in reversed(range(len(frames))):
        #    if idx == 0:
        #        continue
        #    frames[idx] = frames[idx] - frames[idx-1]
        return torch.tensor(frames, dtype=torch.int64)
    
    def get_linspace_ph_id(self, ph_idx_in_a_word):
        for idx, n_ph in enumerate(ph_idx_in_a_word):
            if idx == 0:
                data = torch.linspace(start=int(n_ph-1), end=0, steps=int(n_ph), dtype=torch.int64)
            else:
                data = torch.concat([data, torch.linspace(start=int(n_ph-1), end=0, steps=int(n_ph), dtype=torch.int64)], dim=0)
        return data
    
    def get_frame_level_word_id(self, wordID, word_dur):
        #for idx, (id, dur) in enumerate(zip(wordID, word_dur)):
        #    if idx == 0:
        #        data = torch.ones(size=(dur,),dtype=torch.int64) * id
        #    else:
        #        data = torch.concat([data, torch.ones(size=(dur,),dtype=torch.int64) * id], dim=0)
        for idx, (id, dur) in enumerate(zip(wordID, word_dur)):
            if idx == 0:
                data = torch.ones(size=(dur,),dtype=torch.int64) * idx
            else:
                data = torch.concat([data, torch.ones(size=(dur,),dtype=torch.int64) * idx], dim=0)
        return data

    def get_continuos_f0(self, f0):
        """Convert F0 to continuous F0
        Args:
            f0 (ndarray): original f0 sequence with the shape (T)
        Return:
            (ndarray): continuous f0 with the shape (T)
        """
        # get uv information as binary
        uv = np.float32(f0 != 0)
        # get start and end of f0
        if (f0 == 0).all():
            logger.warn("all of the f0 values are 0.")
            return uv, f0, False
        start_f0 = f0[f0 != 0][0]
        end_f0 = f0[f0 != 0][-1]
        # padding start and end of f0 sequence
        cf0 = copy.deepcopy(f0)
        start_idx = np.where(cf0 == start_f0)[0][0]
        end_idx = np.where(cf0 == end_f0)[0][-1]
        cf0[:start_idx] = start_f0
        cf0[end_idx:] = end_f0
        # get non-zero frame index
        nz_frames = np.where(cf0 != 0)[0]
        # perform linear interpolation
        f = interp1d(nz_frames, cf0[nz_frames])
        cf0 = f(np.arange(0, cf0.shape[0]))

        return uv, cf0, True

    def get_ph_pooling_dur(self, ph_e, word_e):
        out = list()
        z_t = 0
        word_idx = 0
        for idx, e in enumerate(ph_e):
            idx += 1
            if word_e[word_idx] == e:
                out.append(int(idx - z_t))
                z_t = idx
                word_idx += 1
        return out

    def get_item(self, basepath):

        # labのデータは推論時存在しない
        spec, wav   = self.get_audio(filename= basepath+".wav") 
        f0          = torch.load(basepath + "_f0.pt"          ) 
        vuv         = torch.load(basepath + "_vuv.pt"         ) 

        ph          = torch.load(basepath + "_ph.pt"          ) # ust or lab
        ph_dur_ms   = torch.load(basepath + "_ph_dur_ms.pt"   )     # lab
        ph_s_ms     = torch.load(basepath + "_ph_s_ms.pt"     )     # lab
        ph_e_ms     = torch.load(basepath + "_ph_e_ms.pt"     )     # lab

        word        = torch.load(basepath + "_word.pt"        ) # ust
        word_dur_ms = torch.load(basepath + "_word_dur_ms.pt" )     # lab
        word_e_ms   = torch.load(basepath + "_word_e_ms.pt"   )     # lab
        word_s_ms   = torch.load(basepath + "_word_s_ms.pt"   )     # lab

        noteID      = torch.load(basepath + "_noteID.pt"      ) # ust
        notename    = torch.load(basepath + "_notename.pt"    ) # ust    
        notedur     = torch.load(basepath + "_notedur.pt"     ) # ust    
        
        ph_idx_in_a_word = torch.load(basepath + "_n_ph_in_a_word.pt" )   # lab

        # tokenize and get duration
        ph_IDs              = self.get_ph_ID(ph_list=ph)
        word_IDs            = self.get_word_ID(word_list=word)
        ph_frame_dur        = self.get_dur_frame_from_e_ms(e_ms=ph_e_ms)
        word_frame_dur      = self.get_dur_frame_from_e_ms(e_ms=word_e_ms)
        frame_level_word_id = self.get_frame_level_word_id(wordID=word_IDs, word_dur=word_frame_dur)
        note_frame_dur      = self.get_dur_frame_from_e_ms(e_ms=torch.cumsum(torch.tensor(notedur,dtype=torch.float32),dim=0))
        ph_idx_in_a_word    = self.get_linspace_ph_id(ph_idx_in_a_word=ph_idx_in_a_word)
        n_ph_pooling        = self.get_ph_pooling_dur(ph_e=ph_e_ms, word_e=word_e_ms)
        _, spec_len         = spec.shape
        noteID, notedur     = self.expand_note_info(ph_IDs=ph_IDs.view(1,-1), 
                                                    noteID      =torch.tensor(noteID,dtype=torch.int64).view(1,-1),
                                                    note_dur    =torch.tensor(notedur,dtype=torch.int64).view(1,-1),
                                                    n_ph_pooling=torch.tensor(n_ph_pooling, dtype=torch.int64).view(1,-1))
        
        # paddingの影響で、1長くなることがあるので、その対策
        if  spec_len != len(vuv) or  spec_len != len(f0):
            #back = f0 # for debug
            f0 = f0[:spec_len]
            vuv = vuv[:spec_len]
        if spec_len != int(torch.sum(word_frame_dur)):
            word_frame_dur[-1] += 1 
        if spec_len != int(torch.sum(ph_frame_dur)):
            ph_frame_dur[-1] += 1 
        if spec_len != len(frame_level_word_id):
            frame_level_word_id = torch.tensor(frame_level_word_id)
            frame_level_word_id = torch.concat( [frame_level_word_id, torch.unsqueeze(frame_level_word_id[-1],dim=0)],  dim=0 )

        _, c_f0, _ = self.get_continuos_f0(f0) # for SiFiGAN
        Sinewave = self.signal_generator(torch.tensor(c_f0, dtype=torch.float32).view(1,1,-1))# for SiFiGAN

        prod_upsample_scales = np.cumprod(self.upsample_scales)
        df_sample_rates = [self.sampling_rate / self.hop_length * s for s in prod_upsample_scales]
        dfs = []
        for df, us in zip(self.dense_factors, prod_upsample_scales):
            dfs += [
                np.repeat(dilated_factor(c_f0, self.sampling_rate, df), us)
                if self.df_f0_type == "cf0"
                else np.repeat(dilated_factor(c_f0, self.sampling_rate, df), us)
            ]

        # 保障
        assert sum(n_ph_pooling) == len(ph_IDs)
        assert int(torch.sum(word_frame_dur)) == int(torch.sum(ph_frame_dur))
        assert len(ph_idx_in_a_word) == len(ph_IDs)
        assert spec_len == int(torch.sum(word_frame_dur))
        assert spec_len == len(frame_level_word_id)
        assert spec_len == len(f0)
        assert spec_len == len(vuv)
        assert len(ph_IDs) == len(noteID) # WORD ID が多い。おそらく"っ"の影響
        #assert len(word_IDs) == len(noteID) # WORD ID が多い。おそらく"っ"の影響
        assert len(word_IDs) == len(notename)
        assert len(word_IDs) == len(word_dur_ms)
        assert len(word_IDs) == len(word_e_ms)
        assert len(word_IDs) == len(word_s_ms)
        for i in range(len(self.dense_factors)):
            assert Sinewave.size(2) * df_sample_rates[i] == len(dfs[i]) * self.sampling_rate

        return (torch.tensor(wav,                       dtype=torch.float32),    
                torch.tensor(spec,                      dtype=torch.float32),    
                torch.tensor(f0,                        dtype=torch.float32),  
                torch.tensor(vuv,                       dtype=torch.int64)+1,    # maskを0とする。

                torch.tensor(ph_IDs,                    dtype=torch.int64)+1,    # maskを0とする。
                torch.tensor(ph_dur_ms,                 dtype=torch.float32),   
                torch.tensor(ph_frame_dur,              dtype=torch.int64),  
                torch.tensor(ph_s_ms,                   dtype=torch.float32),  
                torch.tensor(ph_e_ms,                   dtype=torch.float32),   
            
                torch.tensor(word_IDs,                  dtype=torch.int64)+1,   # maskを0とする。 
                torch.tensor(word_dur_ms,               dtype=torch.float32)/1000,  
                torch.tensor(word_frame_dur,            dtype=torch.int64),    
                torch.tensor(word_s_ms,                 dtype=torch.float32),  
                torch.tensor(word_e_ms,                 dtype=torch.float32),   

                torch.tensor(noteID,                    dtype=torch.int64)+1 ,   # maskを0とする。
                #torch.tensor(note_frame_dur,            dtype=torch.int64),     
                torch.tensor(notedur,                   dtype=torch.float32) / 1000,     
                torch.tensor(ph_idx_in_a_word,            dtype=torch.int64)+1,     # maskを0とする。
                torch.tensor(frame_level_word_id,       dtype=torch.int64)+1,    # maskを0とする。
                torch.tensor(n_ph_pooling,              dtype=torch.int64),    # maskを0とする。
                torch.tensor(c_f0,                      dtype=torch.float32),
                #torch.tensor(Sinewave,                  dtype=torch.float32),
                dfs )

    def __getitem__(self, index):
        return self.get_item(self.basepath_list[index])

    def __len__(self):
        return len(self.basepath_list)


class TextAudioCollate_ALL():
    def __init__(self, hparams, return_ids=False):
        self.return_ids = return_ids
        self.dense_factors = hparams["SiFiGAN_utils"]["dense_factors"]
        self.upsample_scales = hparams["SiFiGANGenerator"]["upsample_scales"]
        
        # for SiFiGAN
        self.signal_generator = SignalGenerator(sample_rate=hparams["sampling_rate"],
                                                hop_size=hparams["hop_length"] ,
                                                sine_amp=hparams["SiFiGAN_utils"]["sine_amp"],
                                                noise_amp=hparams["SiFiGAN_utils"]["noise_amp"],
                                                signal_types=hparams["SiFiGAN_utils"]["signal_types"],)

    def __call__(self, batch):
        _, ids_sorted_decreasing = torch.sort( torch.LongTensor( [x[1].size(1) for x in batch] ),dim=0, descending=True )

        max_wav_len     = max([    x[0].size(1) for x in batch])
        max_spec_len    = max([    x[1].size(1) for x in batch])
        max_f0_len      = max([len(x[2]) for x in batch])
        max_vuv_len     = max([len(x[3]) for x in batch])

        max_ph_IDs_len          = max([len(x[4]) for x in batch])
        max_ph_dur_ms_len       = max([len(x[5]) for x in batch])
        max_ph_frame_dur_len    = max([len(x[6]) for x in batch])
        max_ph_s_ms_len         = max([len(x[7]) for x in batch])
        max_ph_e_ms_len         = max([len(x[8]) for x in batch])

        max_word_IDs_len        = max([len(x[9]) for x in batch])
        max_word_dur_ms_len     = max([len(x[10]) for x in batch])
        max_word_frame_dur_len  = max([len(x[11]) for x in batch])
        max_word_s_ms_len       = max([len(x[12]) for x in batch])
        max_word_e_ms_len       = max([len(x[13]) for x in batch])

        max_noteID_len          = max([len(x[14]) for x in batch])
        max_note_frame_dur_len  = max([len(x[15]) for x in batch])
        
        max_ph_idx_in_a_word_len  = max([len(x[16]) for x in batch])
        max_frame_level_word_id_len  = max([len(x[17]) for x in batch])
        max_n_ph_pooling_len  = max([len(x[18]) for x in batch])
        
        max_c_f0_len  = max([len(x[19]) for x in batch])

        #max_Sinewave_len  = max([len(x[20]) for x in batch])

        wav_lengths     = torch.LongTensor(len(batch))
        spec_lengths    = torch.LongTensor(len(batch))
        f0_lengths      = torch.LongTensor(len(batch))
        vuv_lengths     = torch.LongTensor(len(batch))
        
        ph_IDs_lengths          = torch.LongTensor(len(batch))
        ph_dur_ms_lengths       = torch.LongTensor(len(batch))
        ph_frame_dur_lengths    = torch.LongTensor(len(batch))
        ph_s_ms_lengths         = torch.LongTensor(len(batch))
        ph_e_ms_lengths         = torch.LongTensor(len(batch))

        word_IDs_lengths            = torch.LongTensor(len(batch))
        word_dur_ms_lengths         = torch.LongTensor(len(batch))
        word_frame_dur_lengths      = torch.LongTensor(len(batch))
        word_s_ms_lengths           = torch.LongTensor(len(batch))
        word_e_ms_lengths           = torch.LongTensor(len(batch))

        noteID_lengths          = torch.LongTensor(len(batch))
        note_frame_dur_lengths  = torch.LongTensor(len(batch))

        ph_idx_in_a_word_lengths      = torch.LongTensor(len(batch))
        frame_level_word_id_lengths = torch.LongTensor(len(batch))
        n_ph_pooling_lengths        = torch.LongTensor(len(batch))
        
        c_f0_lengths        = torch.LongTensor(len(batch))
        #Sinewave_lengths        = torch.LongTensor(len(batch))

        wav_padded  = torch.FloatTensor(len(batch), 1, max_wav_len)
        spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
        f0_padded   = torch.FloatTensor(len(batch), max_f0_len)
        vuv_padded  = torch.LongTensor(len(batch),  max_vuv_len) 

        ph_IDs_padded       = torch.LongTensor(len(batch),   max_ph_IDs_len)
        ph_dur_ms_padded    = torch.FloatTensor(len(batch),  max_ph_dur_ms_len)
        ph_frame_dur_padded = torch.LongTensor(len(batch),   max_ph_frame_dur_len)
        ph_s_ms_padded      = torch.FloatTensor(len(batch),  max_ph_s_ms_len)
        ph_e_ms_padded      = torch.FloatTensor(len(batch),  max_ph_e_ms_len)
        
        word_IDs_padded         = torch.LongTensor(len(batch),  max_word_IDs_len)
        word_dur_ms_padded      = torch.FloatTensor(len(batch), max_word_dur_ms_len)
        word_frame_dur_padded   = torch.LongTensor(len(batch),  max_word_frame_dur_len)
        word_s_ms_padded        = torch.FloatTensor(len(batch), max_word_s_ms_len)
        word_e_ms_padded        = torch.FloatTensor(len(batch), max_word_e_ms_len)

        noteID_padded           = torch.LongTensor(len(batch), max_noteID_len)
        note_frame_dur_padded   = torch.FloatTensor(len(batch), max_note_frame_dur_len)

        ph_idx_in_a_word_padded        = torch.LongTensor(len(batch), max_ph_idx_in_a_word_len)
        frame_level_word_id_padded   = torch.LongTensor(len(batch), max_frame_level_word_id_len)
        n_ph_pooling_padded          = torch.LongTensor(len(batch), max_n_ph_pooling_len)
        
        c_f0_padded   = torch.FloatTensor(len(batch), max_c_f0_len)
        #Sinewave_padded   = torch.FloatTensor(len(batch), max_Sinewave_len)

        wav_padded.zero_()
        spec_padded.zero_()
        f0_padded.zero_()
        vuv_padded.zero_()

        ph_IDs_padded.zero_()
        ph_dur_ms_padded.zero_()
        ph_frame_dur_padded.zero_()
        ph_s_ms_padded.zero_()
        ph_e_ms_padded.zero_()

        word_IDs_padded.zero_()
        word_dur_ms_padded.zero_()
        word_frame_dur_padded.zero_()
        word_s_ms_padded.zero_()
        word_e_ms_padded.zero_()
        
        noteID_padded.zero_()
        note_frame_dur_padded.zero_()
        
        ph_idx_in_a_word_padded.zero_()
        frame_level_word_id_padded.zero_()
        n_ph_pooling_padded.zero_()
        
        c_f0_padded   .zero_()
        #Sinewave_padded   .zero_()

        dfs_padded = []
        dfs_lengths = torch.LongTensor(len(batch), len(self.upsample_scales))
        lengths=max_c_f0_len
        for idx, scales in enumerate(self.upsample_scales):
            lengths *= scales
            dfs_padded.append(torch.FloatTensor(len(batch), 1, int(lengths)).zero_())

        #dfs_batch_np = [[] for _ in range(len(self.dense_factors))]
        #dfs_batch_tensor = [[] for _ in range(len(self.dense_factors))]
        for i in range(len(ids_sorted_decreasing)):
            row = batch[ids_sorted_decreasing[i]]

            wav = row[0]
            wav_padded[i, :, :wav.size(1)] = wav
            wav_lengths[i] = wav.size(1)

            spec = row[1]
            spec_padded[i, :, :spec.size(1)] = spec
            spec_lengths[i] = spec.size(1)
            
            f0 = row[2]
            f0_padded[i, :f0.size(0)] = f0
            f0_lengths[i] = f0.size(0)

            vuv = row[3]
            vuv_padded[i, :vuv.size(0)] = vuv
            vuv_lengths[i] = vuv.size(0)
            
            ph_IDs = row[4]
            ph_IDs_padded[i, :ph_IDs.size(0)] = ph_IDs
            ph_IDs_lengths[i] = ph_IDs.size(0) 
            
            ph_dur_ms = row[5]
            ph_dur_ms_padded[i,     :ph_dur_ms.size(0)] = ph_dur_ms
            ph_dur_ms_lengths[i] =   ph_dur_ms.size(0)

            ph_frame_dur = row[6]
            ph_frame_dur_padded[i,     :ph_frame_dur.size(0)] = ph_frame_dur
            ph_frame_dur_lengths[i] =   ph_frame_dur.size(0)
            
            ph_s_ms = row[7]
            ph_s_ms_padded[i,     :ph_s_ms.size(0)] = ph_s_ms
            ph_s_ms_lengths[i] =   ph_s_ms.size(0)
            
            ph_e_ms = row[8]
            ph_e_ms_padded[i,     :ph_e_ms.size(0)] = ph_e_ms
            ph_e_ms_lengths[i] =   ph_e_ms.size(0)
            

            word_IDs = row[9]
            word_IDs_padded[i,  :word_IDs.size(0)] = word_IDs
            word_IDs_lengths[i] = word_IDs.size(0) 

            word_dur_ms = row[10]
            word_dur_ms_padded[i,     :word_dur_ms.size(0)] = word_dur_ms
            word_dur_ms_lengths[i] =   word_dur_ms.size(0)
            
            word_frame_dur = row[11]
            word_frame_dur_padded[i,     :word_frame_dur.size(0)] = word_frame_dur
            word_frame_dur_lengths[i] =   word_frame_dur.size(0)

            word_s_ms = row[12]
            word_s_ms_padded[i,     :word_s_ms.size(0)] = word_s_ms
            word_s_ms_lengths[i] =   word_s_ms.size(0)    

            word_e_ms = row[13]
            word_e_ms_padded[i,     :word_e_ms.size(0)] = word_e_ms
            word_e_ms_lengths[i] =   word_e_ms.size(0)
            
            noteID = row[14]
            noteID_padded[i, :noteID.size(0)] = noteID
            noteID_lengths[i] = noteID.size(0)

            note_frame_dur = row[15]
            note_frame_dur_padded[i, :note_frame_dur.size(0)] = note_frame_dur
            note_frame_dur_lengths[i] = note_frame_dur.size(0)

            ph_idx_in_a_word = row[16]
            ph_idx_in_a_word_padded[i, :ph_idx_in_a_word.size(0)] = ph_idx_in_a_word
            ph_idx_in_a_word_lengths[i] = ph_idx_in_a_word.size(0)

            frame_level_word_id = row[17]
            frame_level_word_id_padded[i, :frame_level_word_id.size(0)] = frame_level_word_id
            frame_level_word_id_lengths[i] = frame_level_word_id.size(0)
        
            n_ph_pooling = row[18]
            n_ph_pooling_padded[i, :n_ph_pooling.size(0)] = n_ph_pooling
            n_ph_pooling_lengths[i] = n_ph_pooling.size(0)

            c_f0 = row[19]
            c_f0_padded[i, :c_f0.size(0)] = c_f0
            c_f0_lengths[i] = c_f0.size(0)

            dfs = row[20]
            for idx in range(len(self.upsample_scales)):
                df = dfs[idx]
                dfs_padded[idx][i, 0, :len(df)] = torch.tensor(df, dtype=torch.float32)
                dfs_lengths[i, idx] = len(df)

            """
            for i in range(len(self.dense_factors)):
                dfs_batch_np[i] += [
                    dfs[i].astype(np.float32).reshape(-1, 1)
                ]  # [(T', 1), ...]

        for i in range(len(self.dense_factors)):
            dfs_batch_tensor[i] = torch.FloatTensor(np.array(dfs_batch_np[i])).transpose(
                    2, 1
                )  # (B, 1, T')
                
            """
        # 次元調整
        f0_padded = torch.unsqueeze(f0_padded, dim=1)  
            
        # バッチ単位で作成
        Sinewaves = self.signal_generator(c_f0_padded.unsqueeze(1)).float()

        if self.return_ids:
            return  wav_padded,             wav_lengths,            \
                    spec_padded,            spec_lengths,           \
                    f0_padded,              f0_lengths,             \
                    vuv_padded,             vuv_lengths,            \
                    ph_IDs_padded,          ph_IDs_lengths,         \
                    ph_dur_ms_padded,       ph_dur_ms_lengths,      \
                    ph_frame_dur_padded,    ph_frame_dur_lengths,   \
                    ph_s_ms_padded,         ph_s_ms_lengths,        \
                    ph_e_ms_padded,         ph_e_ms_lengths,        \
                    word_IDs_padded,        word_IDs_lengths,       \
                    word_dur_ms_padded,     word_dur_ms_lengths,    \
                    word_frame_dur_padded,  word_frame_dur_lengths, \
                    word_s_ms_padded,       word_s_ms_lengths,      \
                    word_e_ms_padded,       word_e_ms_lengths,      \
                    noteID_padded,          noteID_lengths,         \
                    note_frame_dur_padded,  note_frame_dur_lengths, \
                    ph_idx_in_a_word_padded,  ph_idx_in_a_word_lengths, \
                    frame_level_word_id_padded, frame_level_word_id_lengths, \
                    n_ph_pooling_padded, n_ph_pooling_lengths, \
                    c_f0_padded, c_f0_lengths, \
                    dfs_padded, dfs_lengths, \
                    Sinewaves, \
                    ids_sorted_decreasing    
        
        return  wav_padded,             wav_lengths,            \
                spec_padded,            spec_lengths,           \
                f0_padded,              f0_lengths,             \
                vuv_padded,             vuv_lengths,            \
                ph_IDs_padded,          ph_IDs_lengths,         \
                ph_dur_ms_padded,       ph_dur_ms_lengths,      \
                ph_frame_dur_padded,    ph_frame_dur_lengths,   \
                ph_s_ms_padded,         ph_s_ms_lengths,        \
                ph_e_ms_padded,         ph_e_ms_lengths,        \
                word_IDs_padded,        word_IDs_lengths,       \
                word_dur_ms_padded,     word_dur_ms_lengths,    \
                word_frame_dur_padded,  word_frame_dur_lengths, \
                word_s_ms_padded,       word_s_ms_lengths,      \
                word_e_ms_padded,       word_e_ms_lengths,      \
                noteID_padded,          noteID_lengths,         \
                note_frame_dur_padded,  note_frame_dur_lengths, \
                ph_idx_in_a_word_padded,  ph_idx_in_a_word_lengths, \
                frame_level_word_id_padded, frame_level_word_id_lengths, \
                n_ph_pooling_padded, n_ph_pooling_lengths, \
                c_f0_padded, c_f0_lengths, \
                dfs_padded, dfs_lengths, \
                Sinewaves


"""Multi speaker version"""

class TextAudioSpeakerLoader(torch.utils.data.Dataset):
    """
        1) loads audio, speaker_id, text pairs
        2) normalizes text and converts them to sequences of integers
        3) computes spectrograms from audio files.
    """
    def __init__(self, audiopaths_sid_text, hparams):
        self.audiopaths_sid_text = load_filepaths_and_text(audiopaths_sid_text)
        self.text_cleaners = hparams.text_cleaners
        self.max_wav_value = hparams.max_wav_value
        self.sampling_rate = hparams.sampling_rate
        self.filter_length  = hparams.filter_length
        self.hop_length     = hparams.hop_length
        self.win_length     = hparams.win_length
        self.sampling_rate  = hparams.sampling_rate

        self.cleaned_text = getattr(hparams, "cleaned_text", False)

        self.add_blank = hparams.add_blank
        self.min_text_len = getattr(hparams, "min_text_len", 1)
        self.max_text_len = getattr(hparams, "max_text_len", 190)

        random.seed(1234)
        random.shuffle(self.audiopaths_sid_text)
        self._filter()

    def _filter(self):
        """
        Filter text & store spec lengths
        """
        # Store spectrogram lengths for Bucketing
        # wav_length ~= file_size / (wav_channels * Bytes per dim) = file_size / (1 * 2)
        # spec_length = wav_length // hop_length

        audiopaths_sid_text_new = []
        lengths = []
        for audiopath, sid, text in self.audiopaths_sid_text:
            if self.min_text_len <= len(text) and len(text) <= self.max_text_len:
                audiopaths_sid_text_new.append([audiopath, sid, text])
                lengths.append(os.path.getsize(audiopath) // (2 * self.hop_length))
        self.audiopaths_sid_text = audiopaths_sid_text_new
        self.lengths = lengths

    def get_audio_text_speaker_pair(self, audiopath_sid_text):
        # separate filename, speaker_id and text
        audiopath, sid, text = audiopath_sid_text[0], audiopath_sid_text[1], audiopath_sid_text[2]
        text = self.get_text(text)
        spec, wav = self.get_audio(audiopath)
        sid = self.get_sid(sid)
        return (text, spec, wav, sid)

    def get_audio(self, filename):
        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))
        #audio_norm = audio / self.max_wav_value
        audio_norm = audio_norm.unsqueeze(0)
        spec_filename = filename.replace(".wav", ".spec.pt")
        if os.path.exists(spec_filename):
            spec = torch.load(spec_filename)
        else:
            spec = spectrogram_torch(audio_norm, self.filter_length,
                self.sampling_rate, self.hop_length, self.win_length,
                center=False)
            spec = torch.squeeze(spec, 0)
            torch.save(spec, spec_filename)

        return spec, audio_norm

    def get_text(self, text):
        if self.cleaned_text:
            text_norm = cleaned_text_to_sequence(text)
        else:
            text_norm = text_to_sequence(text, self.text_cleaners)
        if self.add_blank:
            text_norm = commons.intersperse(text_norm, 0)
        text_norm = torch.LongTensor(text_norm)
        return text_norm

    def get_sid(self, sid):
        sid = torch.LongTensor([int(sid)])
        return sid

    def __getitem__(self, index):
        return self.get_audio_text_speaker_pair(self.audiopaths_sid_text[index])

    def __len__(self):
        return len(self.audiopaths_sid_text)


class TextAudioSpeakerCollate():
    """ Zero-pads model inputs and targets
    """
    def __init__(self, return_ids=False):
        self.return_ids = return_ids

    def __call__(self, batch):
        """Collate's training batch from normalized text, audio and speaker identities
        PARAMS
        ------
        batch: [text_normalized, spec_normalized, wav_normalized, sid]
        """
        # Right zero-pad all one-hot text sequences to max input length
        _, ids_sorted_decreasing = torch.sort(
            torch.LongTensor([x[1].size(1) for x in batch]),
            dim=0, descending=True)

        max_text_len = max([len(x[0]) for x in batch])
        max_spec_len = max([x[1].size(1) for x in batch])
        max_wav_len = max([x[2].size(1) for x in batch])

        text_lengths = torch.LongTensor(len(batch))
        spec_lengths = torch.LongTensor(len(batch))
        wav_lengths = torch.LongTensor(len(batch))
        sid = torch.LongTensor(len(batch))

        text_padded = torch.LongTensor(len(batch), max_text_len)
        spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
        wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)
        text_padded.zero_()
        spec_padded.zero_()
        wav_padded.zero_()
        for i in range(len(ids_sorted_decreasing)):
            row = batch[ids_sorted_decreasing[i]]

            text = row[0]
            text_padded[i, :text.size(0)] = text
            text_lengths[i] = text.size(0)

            spec = row[1]
            spec_padded[i, :, :spec.size(1)] = spec
            spec_lengths[i] = spec.size(1)

            wav = row[2]
            wav_padded[i, :, :wav.size(1)] = wav
            wav_lengths[i] = wav.size(1)

            sid[i] = row[3]

        if self.return_ids:
            return text_padded, text_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, sid, ids_sorted_decreasing
        return text_padded, text_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, sid


class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
    """
    Maintain similar input lengths in a batch.
    Length groups are specified by boundaries.
    Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.
  
    It removes samples which are not included in the boundaries.
    Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
    """
    def __init__(self, dataset, batch_size, boundaries, num_replicas=None, rank=None, shuffle=True):
        super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
        self.lengths = dataset.lengths
        self.batch_size = batch_size
        self.boundaries = boundaries
  
        self.buckets, self.num_samples_per_bucket = self._create_buckets()
        self.total_size = sum(self.num_samples_per_bucket)
        self.num_samples = self.total_size // self.num_replicas
  
    def _create_buckets(self):
        buckets = [[] for _ in range(len(self.boundaries) - 1)]
        for i in range(len(self.lengths)):
            length = self.lengths[i]
            idx_bucket = self._bisect(length)
            if idx_bucket != -1:
                buckets[idx_bucket].append(i)
  
        for i in range(len(buckets) - 1, 0, -1):
            if len(buckets[i]) == 0:
                buckets.pop(i)
                self.boundaries.pop(i+1)
  
        num_samples_per_bucket = []
        for i in range(len(buckets)):
            len_bucket = len(buckets[i])
            total_batch_size = self.num_replicas * self.batch_size
            rem = (total_batch_size - (len_bucket % total_batch_size)) % total_batch_size
            num_samples_per_bucket.append(len_bucket + rem)
        return buckets, num_samples_per_bucket
  
    def __iter__(self):
      # deterministically shuffle based on epoch
      g = torch.Generator()
      g.manual_seed(self.epoch)
  
      indices = []
      if self.shuffle:
          for bucket in self.buckets:
              indices.append(torch.randperm(len(bucket), generator=g).tolist())
      else:
          for bucket in self.buckets:
              indices.append(list(range(len(bucket))))
  
      batches = []
      for i in range(len(self.buckets)):
          bucket = self.buckets[i]
          len_bucket = len(bucket)
          ids_bucket = indices[i]
          num_samples_bucket = self.num_samples_per_bucket[i]
  
          # add extra samples to make it evenly divisible
          rem = num_samples_bucket - len_bucket
          ids_bucket = ids_bucket + ids_bucket * (rem // len_bucket) + ids_bucket[:(rem % len_bucket)]
  
          # subsample
          ids_bucket = ids_bucket[self.rank::self.num_replicas]
  
          # batching
          for j in range(len(ids_bucket) // self.batch_size):
              batch = [bucket[idx] for idx in ids_bucket[j*self.batch_size:(j+1)*self.batch_size]]
              batches.append(batch)
  
      if self.shuffle:
          batch_ids = torch.randperm(len(batches), generator=g).tolist()
          batches = [batches[i] for i in batch_ids]
      self.batches = batches
  
      assert len(self.batches) * self.batch_size == self.num_samples
      return iter(self.batches)
  
    def _bisect(self, x, lo=0, hi=None):
      if hi is None:
          hi = len(self.boundaries) - 1
  
      if hi > lo:
          mid = (hi + lo) // 2
          if self.boundaries[mid] < x and x <= self.boundaries[mid+1]:
              return mid
          elif x <= self.boundaries[mid]:
              return self._bisect(x, lo, mid)
          else:
              return self._bisect(x, mid + 1, hi)
      else:
          return -1

    def __len__(self):
        return self.num_samples // self.batch_size

if __name__ == "__main__":
    print(1 * False)

    f0_length = 678
    IDs_length = 55
    Batch = 3
    n_speaker = 10
    vocab_size= 20
    max_step = 1000

    dummy_f0        = torch.rand   (                size=( f0_length, ))
    dummy_f0_len    = torch.randint(f0_length ,     size=(1,))
    dummy_vuv       = torch.randint(1,              size=(f0_length,))
    dummy_vuv_len   = torch.randint(f0_length,      size=(1,))
    dummy_IDs       = torch.randint(vocab_size,     size=(IDs_length,))
    dummy_IDs_len   = torch.randint(IDs_length,     size=(1,))
    dummy_attn      = torch.randint(1,              size=(f0_length, IDs_length))
    dummy_g         = torch.randint(n_speaker,      size=(1,))
    dummy_diff_steps= torch.randint(max_step,       size=(1,))

    test = TextAudioCollate()
    batch = [[dummy_f0, dummy_attn, dummy_IDs, dummy_vuv] for x in range(5)]
    _ = test.__call__(batch=batch)
data_utils.py
tonnetonne814/SiFi-VITS2-44100-Ja

import torch.nn as nn
import torch.nn.functional as F

# For Downsampling 
class Downsample1D(nn.Module):
    """A 1D downsampling layer with an optional convolution.

    Parameters:
        channels (`int`):
            number of channels in the inputs and outputs.
        use_conv (`bool`, default `False`):
            option to use a convolution.
        out_channels (`int`, optional):
            number of output channels. Defaults to `channels`.
        padding (`int`, default `1`):
            padding for the convolution.
    """

    def __init__(self, channels, use_conv=False, out_channels=None, padding=1): # , name="conv"):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_conv = use_conv
        self.padding = padding
        stride = 2
        #self.name = name

        if use_conv:
            self.conv = nn.Conv1d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
        else:
            assert self.channels == self.out_channels
            self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride)

    def forward(self, inputs):
        assert inputs.shape[1] == self.channels
        return self.conv(inputs)

# For Upsampling 
class Upsample1D(nn.Module):
    """A 1D upsampling layer with an optional convolution.

    Parameters:
        channels (`int`):
            number of channels in the inputs and outputs.
        use_conv (`bool`, default `False`):
            option to use a convolution.
        use_conv_transpose (`bool`, default `False`):
            option to use a convolution transpose.
        out_channels (`int`, optional):
            number of output channels. Defaults to `channels`.
    """

    def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None): #, name="conv"):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_conv = use_conv
        self.use_conv_transpose = use_conv_transpose
        #self.name = name

        self.conv = None
        if use_conv_transpose:
            self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1)
        elif use_conv:
            self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1)

    def forward(self, inputs):
        assert inputs.shape[1] == self.channels
        if self.use_conv_transpose:
            return self.conv(inputs)

        outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest")

        if self.use_conv:
            outputs = self.conv(outputs)

        return outputs

# No dimensional change ResBlock 
class ResBlock1D(nn.Module):
    def __init__(self, inp_channels, out_channels, time_embed_dim, kernel_size=5):
        super().__init__()
        self.conv_in = Conv1dBlock(inp_channels, out_channels, kernel_size)
        self.conv_out = Conv1dBlock(out_channels, out_channels, kernel_size)

        self.time_emb_act = nn.Mish()
        self.time_emb = nn.Linear(time_embed_dim, out_channels)

        self.residual_conv = (
            nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity()
        )

    def forward(self, inputs, t):
        """
        Args:
            inputs : [ batch_size x inp_channels x horizon ]
            t : [ batch_size x embed_dim ]

        returns:
            out : [ batch_size x out_channels x horizon ]
        """
        t = self.time_emb_act(t)
        t = self.time_emb(t)
        out = self.conv_in(inputs) + rearrange_dims(t)
        out = self.conv_out(out)
        return out + self.residual_conv(inputs)

# for ResidualTemporalBlock1D
class Conv1dBlock(nn.Module):
    """
    Conv1d --> GroupNorm --> Mish
    """

    def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):
        super().__init__()
        
        self.conv1d = nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2)
        self.group_norm = nn.GroupNorm(n_groups, out_channels)
        self.mish = nn.Mish()

    def forward(self, inputs):
        intermediate_repr = self.conv1d(inputs)
        intermediate_repr = rearrange_dims(intermediate_repr)
        intermediate_repr = self.group_norm(intermediate_repr)
        intermediate_repr = rearrange_dims(intermediate_repr)
        output = self.mish(intermediate_repr)
        return output

# For Final output
class OutConv1DBlock(nn.Module):
    def __init__(self, num_groups_out, out_channels, embed_dim, act_fn):
        super().__init__()
        self.final_conv1d_1 = nn.Conv1d(embed_dim, embed_dim, 5, padding=2)
        self.final_conv1d_gn = nn.GroupNorm(num_groups_out, embed_dim)
        self.final_conv1d_act = get_activation(act_fn)
        self.final_conv1d_2 = nn.Conv1d(embed_dim, out_channels, 1)

    def forward(self, hidden_states, temb=None):
        hidden_states = self.final_conv1d_1(hidden_states)
        hidden_states = rearrange_dims(hidden_states)
        hidden_states = self.final_conv1d_gn(hidden_states)
        hidden_states = rearrange_dims(hidden_states)
        hidden_states = self.final_conv1d_act(hidden_states)
        hidden_states = self.final_conv1d_2(hidden_states)
        return hidden_states

def get_activation(act_fn):
    if act_fn in ["swish", "silu"]:
        return nn.SiLU()
    elif act_fn == "mish":
        return nn.Mish()
    elif act_fn == "gelu":
        return nn.GELU()
    elif act_fn == "relu":
        return nn.ReLU()
    else:
        raise ValueError(f"Unsupported activation function: {act_fn}")

# unet_rl.py
def rearrange_dims(tensor):
    if len(tensor.shape) == 2:
        return tensor[:, :, None]
    if len(tensor.shape) == 3:
        return tensor[:, :, None, :]
    elif len(tensor.shape) == 4:
        return tensor[:, :, 0, :]
    else:
        raise ValueError(f"`len(tensor)`: {len(tensor)} has to be 2, 3 or 4.")

diffusion/diffusers_modules.py
tonnetonne814/SiFi-VITS2-44100-Ja
import torch.nn as nn
from diffusion.diffusion_utils import  compute_hyperparams_given_schedule, std_normal, map_noise_scale_to_time_step
import  torch
from tqdm import tqdm

class F0_Diffusion(nn.Module):
    def __init__(self, hps ):
        super(F0_Diffusion, self).__init__()

        self.hps = hps
        self.diffusion_hyperparameters = self.make_diffusion_hyperparams()
        self.alpha = self.diffusion_hyperparameters["alpha"]

        # ロス関数定義はここ
        self.loss_fn = nn.MSELoss()

    def get_diff_hyperparameters(self):
        return self.diffusion_hyperparameters
    
    def make_diffusion_hyperparams(self, noise_schedule="cosine_beta", training=True):
        
        beta = self.get_beta(noise_schedule=noise_schedule,training=training)

        diffusion_hyperparams = compute_hyperparams_given_schedule(beta)
        for key in diffusion_hyperparams:
            if key in ["beta", "alpha", "sigma"]:
                diffusion_hyperparams[key] = diffusion_hyperparams[key].cuda()
        return diffusion_hyperparams
    
    def get_beta(self, noise_schedule, timesteps=None, training=True):
        
        if timesteps is None:   
            timesteps = int(self.hps["Diffusion"]["T"])

        if training is True:
            beta_0 = 0.0001
            beta_T = 0.9999
        else:
            beta_0 = float(self.hps["Diffusion"]["beta_0"])
            beta_T = float(self.hps["Diffusion"]["beta_T"])
                       
        if noise_schedule == "cosine_beta":
            beta = self.get_cosine_beta_schedule(beta_start=beta_0, beta_end=beta_T, timesteps=timesteps)
        elif noise_schedule == "quadratic_beta":
            beta = self.get_quadratic_beta_schedule(beta_start=beta_0, beta_end=beta_T, timesteps=timesteps)
        elif noise_schedule == "sigmoid_beta":
            beta = self.get_sigmoid_beta_schedule(beta_start=beta_0, beta_end=beta_T, timesteps=timesteps)
        else: # Liner schedule
            noise_schedule = "linear"
            beta = torch.linspace(beta_0, beta_T,timesteps).cuda()
        
        print(f"# Noise Schedule # : {noise_schedule}  beta_0:{beta_0}, beta_T:{beta_T}")

        return beta

    def get_cosine_beta_schedule(self, beta_start, beta_end, timesteps, s=0.008): # timesteps = 1000
        """
        cosine schedule as proposed in https://arxiv.org/abs/2102.09672
        """
        steps = timesteps + 1
        x = torch.linspace(0, timesteps, steps)
        alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
        return torch.clip(betas, beta_start, beta_end)

    def get_quadratic_beta_schedule(self, beta_start, beta_end, timesteps):
        return torch.linspace(beta_start**0.5, beta_end**0.5, timesteps) ** 2

    def get_sigmoid_beta_schedule(self, beta_start, beta_end, timesteps):
        betas = torch.linspace(-6, 6, timesteps)
        return torch.sigmoid(betas) * (beta_end - beta_start) + beta_start

    def get_noisy_f0(self, f0, ts):
        noise_gt = std_normal(f0.shape)                  
        ts=ts.view(f0.size(0), 1, 1)
        self.delta = (1 - self.alpha[ts] ** 2.).sqrt() 
        self.alpha_cur = self.alpha[ts]
        noisy_f0 = self.alpha_cur * f0 + self.delta * noise_gt  # q(x_t|x_0)
        return noisy_f0, noise_gt
    
    def get_loss(self, f0_noise_pd, noisy_f0, noise_gt , reverse=False):

        if reverse:
            x0 = (noisy_f0 - self.delta * f0_noise_pd) / self.alpha_cur
            return self.loss_fn(f0_noise_pd, noise_gt), x0

        loss_f0 = self.loss_fn(f0_noise_pd, noise_gt)
        return loss_f0
    
    @torch.no_grad()
    def get_noise_schedule(self, timesteps=None, training=False):
        
        # Specify a specific noise schedule
        if self.hps["Diffusion"]['noise_schedule'] != '':
            noise_schedule = self.hps["Diffusion"]['noise_schedule']
            if isinstance(noise_schedule, list):
                beta = torch.FloatTensor(noise_schedule).cuda()
            else:
                beta = self.get_beta(noise_schedule=noise_schedule, timesteps=timesteps, training=training)

        # Select Schedule
        else:
            try:
                reverse_step = int(self.hps["Diffusion"]['N'])
            except:
                print('Please specify $N (the number of revere iterations) in config file. Now denoise with 4 iterations.')
                reverse_step = 4
            if reverse_step == 1000:
                beta = torch.linspace(0.000001, 0.01, 1000).cuda()
            elif reverse_step == 200:
                beta = torch.linspace(0.0001, 0.02, 200).cuda()
            # Below are schedules derived by Noise Predictor.
            elif reverse_step == 8:
                beta = [6.689325005027058e-07, 1.0033881153503899e-05, 0.00015496854030061513,
                                 0.002387222135439515, 0.035597629845142365, 0.3681158423423767, 0.4735414385795593, 0.5]
            elif reverse_step == 6:
                beta = [1.7838445955931093e-06, 2.7984189728158526e-05, 0.00043231004383414984,
                                0.006634317338466644, 0.09357017278671265, 0.6000000238418579]
            elif reverse_step == 4:
                beta = [3.2176e-04, 2.5743e-03, 2.5376e-02, 7.0414e-01]
            elif reverse_step == 3:
                beta = [9.0000e-05, 9.0000e-03, 6.0000e-01]
            else:
                raise NotImplementedError

        # put schedule on cuda
        #if isinstance(noise_schedule, list):
        #  noise_schedule = torch.FloatTensor(noise_schedule).cuda()

        return torch.FloatTensor(beta).cuda()
    
    @torch.no_grad()
    def generate_x(self, size):
        return std_normal(size)
    
    def get_step_infer(self, size, inference_noise_schedule):
        self.size=size
        T, alpha = self.diffusion_hyperparameters["T"], self.diffusion_hyperparameters["alpha"]
        assert len(alpha) == T
        assert len(size) == 3

        N = len(inference_noise_schedule)
        self.beta_infer = inference_noise_schedule
        alpha_infer = 1 - self.beta_infer
        sigma_infer = self.beta_infer + 0
        for n in range(1, N):
            alpha_infer[n] *= alpha_infer[n - 1]
            sigma_infer[n] *= (1 - alpha_infer[n - 1]) / (1 - alpha_infer[n])
        self.alpha_infer = torch.sqrt(alpha_infer) # self
        self.sigma_infer = torch.sqrt(sigma_infer) # self

        # Mapping noise scales to time steps
        steps_infer = []
        for n in tqdm(range(N), desc="Mapping noise scales"):
            step = map_noise_scale_to_time_step(self.alpha_infer[n], alpha)
            if step >= 0:
                steps_infer.append(step)
        #print(steps_infer, flush=True)
        steps_infer = torch.FloatTensor(steps_infer)

        # N may change since alpha_infer can be out of the range of alpha
        N = len(steps_infer)

        return N, steps_infer
    
    @torch.no_grad()
    def denoising(self, x, noise_pd, ddim, n):
        
        if ddim:
            alpha_next = self.alpha_infer[n] / (1 - self.beta_infer[n]).sqrt()
            c1 = alpha_next / self.alpha_infer[n]
            c2 = -(1 - self.alpha_infer[n] ** 2.).sqrt() * c1
            c3 = (1 - alpha_next ** 2.).sqrt()
            x = c1 * x + c2 * noise_pd + c3 * noise_pd  # std_normal(size)
        else:
            x -= self.beta_infer[n] / torch.sqrt(1 - self.alpha_infer[n] ** 2.) * noise_pd
            x /= torch.sqrt(1 - self.beta_infer[n])
            if n > 0:
                x = x + self.sigma_infer[n] * self.generate_x(self.size)
            else:
                #print()
                pass
        return x
    
diffusion/diffusion_model.py
tonnetonne814/SiFi-VITS2-44100-Ja
import torch
import torch.nn.functional as F
import numpy as np
from inspect import isfunction


"""
Based in part on: https://github.com/lucidrains/denoising-diffusion-pytorch/blob/5989f4c77eafcdc6be0fb4739f0f277a6dd7f7d8/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py#L281
"""
eps = 1e-8

class MultinomialDiffusion(torch.nn.Module):
    def __init__(self, num_classes, timesteps=1000, loss_type='vb_stochastic', parametrization='x0'):
        super(MultinomialDiffusion, self).__init__()
        assert loss_type in ('vb_stochastic', 'vb_all')
        assert parametrization in ('x0', 'direct')

        if loss_type == 'vb_all':
            print('Computing the loss using the bound on _all_ timesteps.'
                  ' This is expensive both in terms of memory and computation.')

        self.num_classes = num_classes
        #self._denoise_fn = denoise_fn       # ノイズ予測器、実装はTransformerになってる。PE有
        self.loss_type = loss_type
        self.shape = None
        self.num_timesteps = timesteps
        self.parametrization = parametrization

        alphas = cosine_beta_schedule(timesteps)

        alphas = torch.tensor(alphas.astype('float64'))
        log_alpha = np.log(alphas)
        log_cumprod_alpha = np.cumsum(log_alpha)

        log_1_min_alpha = log_1_min_a(log_alpha)
        log_1_min_cumprod_alpha = log_1_min_a(log_cumprod_alpha)

        assert log_add_exp(log_alpha, log_1_min_alpha).abs().sum().item() < 1.e-5
        assert log_add_exp(log_cumprod_alpha, log_1_min_cumprod_alpha).abs().sum().item() < 1e-5
        assert (np.cumsum(log_alpha) - log_cumprod_alpha).abs().sum().item() < 1.e-5

        # Convert to float32 and register buffers.
        self.register_buffer('log_alpha', log_alpha.float())
        self.register_buffer('log_1_min_alpha', log_1_min_alpha.float())
        self.register_buffer('log_cumprod_alpha', log_cumprod_alpha.float())
        self.register_buffer('log_1_min_cumprod_alpha', log_1_min_cumprod_alpha.float())

        self.register_buffer('Lt_history', torch.zeros(timesteps))
        self.register_buffer('Lt_count', torch.zeros(timesteps))


    # 学習時のロス計算コード xはIDs、出力はロス
    def preprocess(self, x, t_int, detach_mean=False):
        if self.loss_type == 'vb_stochastic':
            #x_start = x # x= text_IDs =[Batch, seq_len]
            
            # これは使わない。通常は線形のステップになる
            #t, pt = self.sample_time(b, device, 'importance') # t=学習用整数ステップ数、pt=0~1(importance or uniform)。

            log_x_start = index_to_log_onehot(x, self.num_classes) # x_start

            log_sample = self.q_sample(log_x_start=log_x_start, t=t_int)

            ###############
            log_true_prob = self.q_posterior(                   ## t=tのノイズ環境下での正解分布を出力
                log_x_start=log_x_start, log_x_t=log_sample, t=t_int)

            model_in = log_onehot_to_index(log_sample) # ここでlog_x_t(float)がx_t(int)に変換される。    
            return model_in, [log_sample, log_x_start, log_true_prob]

    def postprocess(self, model_out, parameters, t_int, t_float, detach_mean=False):

        log_sample, log_x_start, log_true_prob = parameters

        if self.loss_type == 'vb_stochastic':

            log_pred = F.log_softmax(model_out, dim=1)    # model_out = [Batch, class, seq_len]

            if self.parametrization == 'x0':
                log_model_prob = self.q_posterior(log_x_start=log_pred, log_x_t=log_sample, t=t_int)
            elif self.parametrization == 'direct':
                log_model_prob = log_pred
            else:
                raise ValueError

            if detach_mean:
                log_model_prob = log_model_prob.detach()

            kl = self.multinomial_kl(log_true_prob, log_model_prob)
            kl = sum_except_batch(kl)

            decoder_nll = -log_categorical(log_x_start, log_model_prob)
            decoder_nll = sum_except_batch(decoder_nll)

            mask = (t_int == torch.zeros_like(t_int)).float()
            kl = mask * decoder_nll + (1. - mask) * kl

            # 重みづけ有無判定用計算
            Lt2 = kl.pow(2)
            Lt2_prev = self.Lt_history.gather(dim=0, index=t_int)
            new_Lt_history = (0.1 * Lt2 + 0.9 * Lt2_prev).detach()
            self.Lt_history.scatter_(dim=0, index=t_int, src=new_Lt_history)
            self.Lt_count.scatter_add_(dim=0, index=t_int, src=torch.ones_like(Lt2))

            kl_prior = self.kl_prior(log_x_start)

            # 重みづけロス計算
            vb_loss = kl / t_float + kl_prior

            return torch.sum(vb_loss)   # マイナスを掛けて渡すのが数式的に正解だが、最終的には符号は＋でlossに加算するので、＋のまま

        elif self.loss_type == 'vb_all':
            # Expensive, dont do it ;).
            return -self.nll(model_out)
        else:
            raise ValueError()

    def q_posterior(self, log_x_start, log_x_t, t):
        # ベイズの定理によりqの事後確率を計算する
        # q(xt-1 | xt, x0) = q(xt | xt-1, x0) * q(xt-1 | x0) / q(xt | x0)
        # マルコフ連鎖より→となる。 q(xt | xt-1, x0) = q(xt | xt-1).

        t_minus_1 = t - 1
        # マイナス除去。最終的にデコーダーでは使わない。
        t_minus_1 = torch.where(t_minus_1 < 0, torch.zeros_like(t_minus_1), t_minus_1)
        log_EV_qxtmin_x0 = self.q_pred(log_x_start, t_minus_1) # log_probs=[Batch, class, seq_len]

        num_axes = (1,) * (len(log_x_start.size()) - 1) # Batch軸を無視した次元の数
        t_broadcast = t.view(-1, *num_axes) * torch.ones_like(log_x_start) # t_broadcast=[Batch, class, seq_len]
        log_EV_qxtmin_x0 = torch.where(t_broadcast == 0, log_x_start, log_EV_qxtmin_x0)

        # x_tmin1 ではない。これが一般的な公式の使い方らしい
        # 理解するのは簡単ではないが、事実こうなっている（マジで謎
        unnormed_logprobs = log_EV_qxtmin_x0 + self.q_pred_one_timestep(log_x_t, t)

        log_EV_xtmin_given_xt_given_xstart = \
            unnormed_logprobs \
            - torch.logsumexp(unnormed_logprobs, dim=1, keepdim=True)

        return log_EV_xtmin_given_xt_given_xstart


    def multinomial_kl(self, log_prob1, log_prob2):
        kl = (log_prob1.exp() * (log_prob1 - log_prob2)).sum(dim=1)
        return kl

    def q_pred_one_timestep(self, log_x_t, t):
        log_alpha_t = extract(self.log_alpha, t, log_x_t.shape)
        log_1_min_alpha_t = extract(self.log_1_min_alpha, t, log_x_t.shape)

        # alpha_t * E[xt] + (1 - alpha_t) 1 / K
        log_probs = log_add_exp(
            log_x_t + log_alpha_t,
            log_1_min_alpha_t - np.log(self.num_classes)
        )

        return log_probs

    def q_pred(self, log_x_start, t):
        log_cumprod_alpha_t = extract(self.log_cumprod_alpha, t, log_x_start.shape) # cumprod(a,t,b)=累積積計算 tはint型
        log_1_min_cumprod_alpha = extract(self.log_1_min_cumprod_alpha, t, log_x_start.shape)

        log_probs = log_add_exp(
            log_x_start + log_cumprod_alpha_t,
            log_1_min_cumprod_alpha - np.log(self.num_classes)      # log_probs=[Batch, class, seq_len]
        )

        return log_probs

    
    def predict_start(self, log_x_t, t):
        x_t = log_onehot_to_index(log_x_t) # ここでlog_x_t(float)がx_t(int)に変換される

        out = self._denoise_fn(t, x_t)  # ノイズ予測モデル本体。

        assert out.size(0) == x_t.size(0)
        assert out.size(1) == self.num_classes
        assert out.size()[2:] == x_t.size()[1:]
        log_pred = F.log_softmax(out, dim=1)    # [Batch, class, seq_len]
        return log_pred

    def p_pred(self, log_x, t):
        if self.parametrization == 'x0':
            log_x_recon = self.predict_start(log_x, t=t)
            log_model_pred = self.q_posterior(
                log_x_start=log_x_recon, log_x_t=log_x, t=t)
        elif self.parametrization == 'direct':
            log_model_pred = self.predict_start(log_x, t=t)
        else:
            raise ValueError
        return log_model_pred

    @torch.no_grad()
    def p_sample(self, log_x, t):
        model_log_prob = self.p_pred(log_x=log_x, t=t)
        out = self.log_sample_categorical(model_log_prob)
        return out

    @torch.no_grad()
    def p_sample_loop(self, shape):
        device = self.log_alpha.device

        b = shape[0]
        # start with random normal image.
        img = torch.randn(shape, device=device)

        for i in reversed(range(1, self.num_timesteps)):
            img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
        return img

    @torch.no_grad()
    def _sample(self, image_size, batch_size = 16):
        return self.p_sample_loop((batch_size, 3, image_size, image_size))

    @torch.no_grad()
    def interpolate(self, x1, x2, t = None, lam = 0.5):
        b, *_, device = *x1.shape, x1.device
        t = default(t, self.num_timesteps - 1)

        assert x1.shape == x2.shape

        t_batched = torch.stack([torch.tensor(t, device=device)] * b)
        xt1, xt2 = map(lambda x: self.q_sample(x, t=t_batched), (x1, x2))

        img = (1 - lam) * xt1 + lam * xt2
        for i in reversed(range(0, t)):
            img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))

        return img

    # log(textID)から、q(xt|x0)の分散期待値を計算する　log (E[V])
    def q_sample(self, log_x_start, t):
        log_EV_qxt_x0 = self.q_pred(log_x_start, t)

        log_sample = self.log_sample_categorical(log_EV_qxt_x0)

        return log_sample

    def nll(self, log_x_start):
        b = log_x_start.size(0)
        device = log_x_start.device
        loss = 0
        for t in range(0, self.num_timesteps):
            t_array = (torch.ones(b, device=device) * t).long()

            kl = self.compute_Lt(
                log_x_start=log_x_start,
                log_x_t=self.q_sample(log_x_start=log_x_start, t=t_array),
                t=t_array)

            loss += kl

        loss += self.kl_prior(log_x_start)

        return loss

    def kl_prior(self, log_x_start):
        b = log_x_start.size(0)
        device = log_x_start.device
        ones = torch.ones(b, device=device).long()

        log_qxT_prob = self.q_pred(log_x_start, t=(self.num_timesteps - 1) * ones)
        log_half_prob = -torch.log(self.num_classes * torch.ones_like(log_qxT_prob))

        kl_prior = self.multinomial_kl(log_qxT_prob, log_half_prob)
        return sum_except_batch(kl_prior)

    def compute_Lt(self, log_x_start, log_x_t, t, detach_mean=False):
        log_true_prob = self.q_posterior(                   ## t=tのノイズ環境下での正解を出力
            log_x_start=log_x_start, log_x_t=log_x_t, t=t)

        log_model_prob = self.p_pred(log_x=log_x_t, t=t)    ## t=tでの正解を予測

        if detach_mean:
            log_model_prob = log_model_prob.detach()

        kl = self.multinomial_kl(log_true_prob, log_model_prob)
        kl = sum_except_batch(kl)

        decoder_nll = -log_categorical(log_x_start, log_model_prob)
        decoder_nll = sum_except_batch(decoder_nll)

        mask = (t == torch.zeros_like(t)).float()
        loss = mask * decoder_nll + (1. - mask) * kl

        return loss

    # タイムステップ作成　uniform=線形的
    def sample_time(self, b, device, method='uniform'):
        if method == 'importance':
            if not (self.Lt_count > 10).all():
                return self.sample_time(b, device, method='uniform')

            Lt_sqrt = torch.sqrt(self.Lt_history + 1e-10) + 0.0001
            Lt_sqrt[0] = Lt_sqrt[1]  # Overwrite decoder term with L1.
            pt_all = Lt_sqrt / Lt_sqrt.sum()

            t = torch.multinomial(pt_all, num_samples=b, replacement=True)

            pt = pt_all.gather(dim=0, index=t)
            return t, pt

        elif method == 'uniform':
            t= torch.randint(self.num_timesteps, size=(b, 1, 1)).cuda()
            pt = torch.ones_like(t).float() / self.num_timesteps
            return t, pt
        else:
            raise ValueError

    def log_prob(self, x):
        b, device = x.size(0), x.device
        if self.training:
            return self._train_loss(x)

        else: # 評価用
            log_x_start = index_to_log_onehot(x, self.num_classes)

            t, pt = self.sample_time(b, device, 'importance')

            kl = self.compute_Lt(
                log_x_start, self.q_sample(log_x_start=log_x_start, t=t), t)

            kl_prior = self.kl_prior(log_x_start)

            # Upweigh loss term of the kl
            loss = kl / pt + kl_prior

            return -loss

    # サンプリング開始時の初期log_zを生成する。
    @torch.no_grad()
    def generate_log_z(self, size):
        device = "cuda"
        uniform_logits = torch.zeros(size=size, device=device)
        log_z = self.log_sample_categorical(uniform_logits)
        return log_z

    # for文最初のやつ
    @torch.no_grad()
    def sample_preprocess(self, log_z, t, num_samples=1):
        device = "cuda"
        

        print(f'Sample timestep {t:4d}', end='\r')
        t = torch.full((num_samples,), t, device=device, dtype=torch.long)
        model_in = log_onehot_to_index(log_z) # ここでlog_x_t(float)がx_t(int)に変換される
        return model_in, log_z

    # モデルの出力を入れて、for文内の最後を出し、これがまたfor文最初にループする。
    @torch.no_grad()
    def sample_postprocess(self, model_out, log_z, t_int):
        t_int = torch.tensor(t_int, dtype=torch.int64, device="cuda")
        t_int = torch.unsqueeze(t_int, dim=0)
        model_log_prob = F.log_softmax(model_out, dim=1)    # [Batch, class, seq_len]

        if self.parametrization == 'x0':
            model_log_prob = self.q_posterior(
                log_x_start=model_log_prob, log_x_t=log_z, t=t_int)
        elif self.parametrization == 'direct':
            pass
        else:
            raise ValueError
        
        log_z = self.log_sample_categorical(model_log_prob)

        return log_z

    @torch.no_grad()
    def decode_log_z(self, log_z):
        return log_onehot_to_index(log_z) 

    @torch.no_grad()
    def p_sample(self, log_x, t):
        model_log_prob = self.p_pred(log_x=log_x, t=t)
        out = self.log_sample_categorical(model_log_prob)
        return out

    def log_sample_categorical(self, logits):
        uniform = torch.rand_like(logits)
        gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
        sample = (gumbel_noise + logits).argmax(dim=1)
        log_sample = index_to_log_onehot(sample, self.num_classes)
        return log_sample
    
    def sample_chain(self, num_samples):
        b = num_samples
        device = self.log_alpha.device
        uniform_logits = torch.zeros(
            (b, self.num_classes) + self.shape, device=device)

        zs = torch.zeros((self.num_timesteps, b) + self.shape).long()

        log_z = self.log_sample_categorical(uniform_logits)
        for i in reversed(range(0, self.num_timesteps)):
            print(f'Chain timestep {i:4d}', end='\r')
            t = torch.full((b,), i, device=device, dtype=torch.long)
            log_z = self.p_sample(log_z, t)

            zs[i] = log_onehot_to_index(log_z)
        print()
        return zs


def sum_except_batch(x, num_dims=1):
    '''
    Sums all dimensions except the first.

    Args:
        x: Tensor, shape (batch_size, ...)
        num_dims: int, number of batch dims (default=1)

    Returns:
        x_sum: Tensor, shape (batch_size,)
    '''
    return x.reshape(*x.shape[:num_dims], -1).sum(-1)


def log_1_min_a(a):
    return torch.log(1 - a.exp() + 1e-40)


def log_add_exp(a, b):
    maximum = torch.max(a, b)
    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))


def exists(x):
    return x is not None


def extract(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def default(val, d):
    if exists(val):
        return val
    return d() if isfunction(d) else d


def log_categorical(log_x_start, log_prob):
    return (log_x_start.exp() * log_prob).sum(dim=1)


def index_to_log_onehot(x, num_classes):
    assert x.max().item() < num_classes, \
        f'Error: {x.max().item()} >= {num_classes}'
    x_onehot = F.one_hot(x, num_classes) # x_onehot = [Batch, seq_len, classes ]

    permute_order = (0, -1) + tuple(range(1, len(x.size())))

    x_onehot = x_onehot.permute(permute_order) # x_onehot = [Batch, classes, seq_len ]

    log_x = torch.log(x_onehot.float().clamp(min=1e-30)) # logを取る

    return log_x


def log_onehot_to_index(log_x):
    return log_x.argmax(1)


def cosine_beta_schedule(timesteps, s = 0.008):
    """
    cosine schedule
    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
    """
    steps = timesteps + 1
    x = np.linspace(0, steps, steps)
    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])

    alphas = np.clip(alphas, a_min=0.001, a_max=1.)

    # Use sqrt of this, so the alpha in our paper is the alpha_sqrt from the
    # Gaussian diffusion in Ho et al.
    alphas = np.sqrt(alphas)
    return alphas


diffusion/diffusion_multinomial.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import numpy as np
import torch
import torch.nn as nn
import copy
import math 
from tqdm import tqdm

def flatten(v):
    """
    Flatten a list of lists/tuples
    """

    return [x for y in v for x in y]


def rescale(x):
    """
    Rescale a tensor to 0-1
    """

    return (x - x.min()) / (x.max() - x.min())


def find_max_epoch(path):
    """
    Find maximum epoch/iteration in path, formatted ${n_iter}.pkl
    E.g. 100000.pkl

    Parameters:
    path (str): checkpoint path
    
    Returns:
    maximum iteration, -1 if there is no (valid) checkpoint
    """

    files = os.listdir(path)
    epoch = -1
    for f in files:
        if len(f) <= 4:
            continue
        if f[-4:]  == '.pkl':
            try:
                epoch = max(epoch, int(f[:-4]))
            except:
                continue
    #print(path, epoch, flush=True)
    return epoch


def print_size(net):
    """
    Print the number of parameters of a network
    """

    if net is not None and isinstance(net, torch.nn.Module):
        module_parameters = filter(lambda p: p.requires_grad, net.parameters())
        params = sum([np.prod(p.size()) for p in module_parameters])
        print("{} Parameters: {:.6f}M".format(
            net.__class__.__name__, params / 1e6), flush=True)


# Utilities for diffusion models

def std_normal(size):
    """
    Generate the standard Gaussian variable of a certain size
    """

    return torch.normal(0, 1, size=size).cuda()


def calc_noise_scale_embedding(noise_scales, noise_scale_embed_dim_in):
    """
    Embed a noise scale $t$ into a higher dimensional space
    E.g. the embedding vector in the 128-dimensional space is
    [sin(t * 10^(0*4/63)), ... , sin(t * 10^(63*4/63)), cos(t * 10^(0*4/63)), ... , cos(t * 10^(63*4/63))]

    Parameters:
    noise_scales (torch.long tensor, shape=(batchsize, 1)):     
                                noise scales for batch data
    noise_scale_embed_dim_in (int, default=128):  
                                dimensionality of the embedding space for discrete noise scales
    
    Returns:
    the embedding vectors (torch.tensor, shape=(batchsize, noise_scale_embed_dim_in)):
    """

    assert noise_scale_embed_dim_in % 2 == 0

    half_dim = noise_scale_embed_dim_in // 2
    _embed = np.log(10000) / (half_dim - 1)
    _embed = torch.exp(torch.arange(half_dim) * -_embed).cuda()
    _embed = noise_scales * _embed
    noise_scale_embed = torch.cat((torch.sin(_embed), 
                                      torch.cos(_embed)), 1)
    
    return noise_scale_embed


def calc_diffusion_hyperparams_given_beta(beta):
    """
    Compute diffusion process hyperparameters

    Parameters:
    beta (tensor):  beta schedule 
    
    Returns:
    a dictionary of diffusion hyperparameters including:
        T (int), beta/alpha/sigma (torch.tensor on cpu, shape=(T, ))
        These cpu tensors are changed to cuda tensors on each individual gpu
    """

    T = len(beta)
    alpha = 1 - beta
    sigma = beta + 0
    for t in range(1, T):
        alpha[t] *= alpha[t-1]  # \alpha^2_t = \prod_{s=1}^t (1-\beta_s)
        sigma[t] *= (1-alpha[t-1]) / (1-alpha[t])  # \sigma^2_t = \beta_t * (1-\alpha_{t-1}) / (1-\alpha_t)
    alpha = torch.sqrt(alpha)
    sigma = torch.sqrt(sigma)
    
    _dh = {}
    _dh["T"], _dh["beta"], _dh["alpha"], _dh["sigma"] = T, beta, alpha, sigma
    diffusion_hyperparams = _dh
    return diffusion_hyperparams


def calc_diffusion_hyperparams(T, beta_0, beta_T, tau, N, beta_N, alpha_N, rho):
    """
    Compute diffusion process hyperparameters

    Parameters:
    T (int):                    number of noise scales
    beta_0 and beta_T (float):  beta schedule start/end value, 
                                where any beta_t in the middle is linearly interpolated
    
    Returns:
    a dictionary of diffusion hyperparameters including:
        T (int), beta/alpha/sigma (torch.tensor on cpu, shape=(T, ))
        These cpu tensors are changed to cuda tensors on each individual gpu
    """

    beta = torch.linspace(beta_0, beta_T, T)
    alpha = 1 - beta
    sigma = beta + 0
    for t in range(1, T):
        alpha[t] *= alpha[t-1]  # \alpha^2_t = \prod_{s=1}^t (1-\beta_s)
        sigma[t] *= (1-alpha[t-1]) / (1-alpha[t])  # \sigma^2_t = \beta_t * (1-\alpha_{t-1}) / (1-\alpha_t)
    alpha = torch.sqrt(alpha)
    sigma = torch.sqrt(sigma)
    
    _dh = {}
    _dh["T"], _dh["beta"], _dh["alpha"], _dh["sigma"] = T, beta, alpha, sigma
    _dh["tau"], _dh["N"], _dh["betaN"], _dh["alphaN"], _dh["rho"] = tau, N, beta_N, alpha_N, rho
    diffusion_hyperparams = _dh
    return diffusion_hyperparams


def sampling_given_noise_schedule(
        net,
        size,
        diffusion_hyperparams,
        inference_noise_schedule,
        condition=None,
        ddim=False,
        return_sequence=False):
    """
    Perform the complete sampling step according to p(x_0|x_T) = \prod_{t=1}^T p_{\theta}(x_{t-1}|x_t)

    Parameters:
    net (torch network):            the wavenet models
    size (tuple):                   size of tensor to be generated,
                                    usually is (number of audios to generate, channels=1, length of audio)
    diffusion_hyperparams (dict):   dictionary of diffusion hyperparameters returned by calc_diffusion_hyperparams
                                    note, the tensors need to be cuda tensors
    condition (torch.tensor):       ground truth mel spectrogram read from disk
                                    None if used for unconditional generation

    Returns:
    the generated audio(s) in torch.tensor, shape=size
    """

    IDs_gt, IDs_gt_len, dur_gt, speaker_ID = condition
    _dh = diffusion_hyperparams
    T, alpha = _dh["T"], _dh["alpha"]
    assert len(alpha) == T
    assert len(size) == 3

    N = len(inference_noise_schedule)
    beta_infer = inference_noise_schedule
    alpha_infer = 1 - beta_infer
    sigma_infer = beta_infer + 0
    for n in range(1, N):
        alpha_infer[n] *= alpha_infer[n - 1]
        sigma_infer[n] *= (1 - alpha_infer[n - 1]) / (1 - alpha_infer[n])
    alpha_infer = torch.sqrt(alpha_infer)
    sigma_infer = torch.sqrt(sigma_infer)

    # Mapping noise scales to time steps
    steps_infer = []
    for n in tqdm(range(N), desc="Mapping noise scales"):
        step = map_noise_scale_to_time_step(alpha_infer[n], alpha)
        if step >= 0:
            steps_infer.append(step)
    #print(steps_infer, flush=True)
    steps_infer = torch.FloatTensor(steps_infer)

    # N may change since alpha_infer can be out of the range of alpha
    N = len(steps_infer)

    print('begin sampling, total number of reverse steps = %s' % N)

    x = std_normal(size).cuda()
    log_z = net.multinomial_diff.generate_log_z(size)

    if return_sequence:
        x_ = copy.deepcopy(x)
        xs = [x_]

    f0_length = torch.tensor([dur_gt.size(1)], dtype=torch.int64, device="cuda:0")
    with torch.no_grad():
        for n in range(N - 1, -1, -1):
            diffusion_steps = (steps_infer[n] * torch.ones((size[0], 1))).cuda()

            # preprocess
            model_in, log_z = net.multinomial_diff.sample_preprocess(log_z=log_z, t=n)

            # 出力も要修正
            epsilon_theta, model_out =  net(f0=x,  
                                f0_len=f0_length,
                                IDs=IDs_gt, 
                                IDs_len=IDs_gt_len, 
                                vuv=model_in,                  
                                vuv_len=f0_length,               
                                attn=dur_gt, 
                                timesteps=diffusion_steps,
                                g=speaker_ID,
                                sampling=True)
            
            # postprocess
            log_z = net.multinomial_diff.sample_postprocess(model_out=model_out, log_z=log_z, t_int=n)

            if ddim:
                alpha_next = alpha_infer[n] / (1 - beta_infer[n]).sqrt()
                c1 = alpha_next / alpha_infer[n]
                c2 = -(1 - alpha_infer[n] ** 2.).sqrt() * c1
                c3 = (1 - alpha_next ** 2.).sqrt()
                x = c1 * x + c2 * epsilon_theta + c3 * epsilon_theta  # std_normal(size)
            else:
                x -= beta_infer[n] / torch.sqrt(1 - alpha_infer[n] ** 2.) * epsilon_theta
                x /= torch.sqrt(1 - beta_infer[n])
                if n > 0:
                    x = x + sigma_infer[n] * std_normal(size)
            if return_sequence:
                x_ = copy.deepcopy(x)
                xs.append(x_)
    if return_sequence:
        return xs
    return x, net.multinomial_diff.decode_log_z(log_z)

def noise_scheduling(net, size, diffusion_hyperparams, condition=None, ddim=False):
    """
    Perform the complete sampling step according to p(x_0|x_T) = \prod_{t=1}^T p_{\theta}(x_{t-1}|x_t)

    Parameters:
    net (torch network):            the wavenet models
    size (tuple):                   size of tensor to be generated,
                                    usually is (number of audios to generate, channels=1, length of audio)
    diffusion_hyperparams (dict):   dictionary of diffusion hyperparameters returned by calc_diffusion_hyperparams
                                    note, the tensors need to be cuda tensors
    condition (torch.tensor):       ground truth mel spectrogram read from disk
                                    None if used for unconditional generation

    Returns:
    noise schedule:                 a list of noise scales in torch.tensor, length <= N
    """

    _dh = diffusion_hyperparams
    N, betaN, alphaN, rho, alpha = _dh["N"], _dh["betaN"], _dh["alphaN"], _dh["rho"], _dh["alpha"]

    print('begin noise scheduling, maximum number of reverse steps = %d' % (N))

    betas = []
    x = std_normal(size)
    with torch.no_grad():
        beta_cur = torch.ones(1, 1, 1).cuda() * betaN
        alpha_cur = torch.ones(1, 1, 1).cuda() * alphaN
        for n in range(N - 1, -1, -1):
            # print(n, beta_cur.squeeze().item(), alpha_cur.squeeze().item())
            step = map_noise_scale_to_time_step(alpha_cur.squeeze().item(), alpha)
            if step >= 0:
                betas.append(beta_cur.squeeze().item())
            diffusion_steps = (step * torch.ones((size[0], 1))).cuda()
            epsilon_theta = net((x, condition, diffusion_steps,))
            if ddim:
                alpha_nxt = alpha_cur / (1 - beta_cur).sqrt()
                c1 = alpha_nxt / alpha_cur
                c2 = -(1 - alpha_cur ** 2.).sqrt() * c1
                c3 = (1 - alpha_nxt ** 2.).sqrt()
                x = c1 * x + c2 * epsilon_theta + c3 * epsilon_theta  # std_normal(size)
            else:
                x -= beta_cur / torch.sqrt(1 - alpha_cur ** 2.) * epsilon_theta
                x /= torch.sqrt(1 - beta_cur)
            alpha_nxt, beta_nxt = alpha_cur, beta_cur
            alpha_cur = alpha_nxt / (1 - beta_nxt).sqrt()
            if alpha_cur > 1:
                break
            beta_cur = net.noise_pred(
                x.squeeze(1), (beta_nxt.view(-1, 1), (1 - alpha_cur ** 2.).view(-1, 1)))
            if beta_cur.squeeze().item() < rho:
                break
    return torch.FloatTensor(betas[::-1]).cuda()


def theta_timestep_loss(net, X, diffusion_hyperparams, reverse=False):
    """
    Compute the training loss for learning theta

    Parameters:
    net (torch network):            the wavenet models
    X (tuple, shape=(2,)):          training data in tuple form (mel_spectrograms, audios)
                                    mel_spectrograms: torch.tensor, shape is batchsize followed by each mel_spectrogram shape
                                    audios: torch.tensor, shape=(batchsize, 1, length of audio)
    diffusion_hyperparams (dict):   dictionary of diffusion hyperparameters returned by calc_diffusion_hyperparams
                                    note, the tensors need to be cuda tensors

    Returns:
    theta loss
    """
    # assert type(X) == tuple and len(X) == 2

    ##############
    # ロス関数定義# 
    loss_fn = nn.MSELoss()
    ##############
    
    f0, f0_lengths, dur, IDs, IDs_lengths, vuv, vuv_lengths, speakerID = X
    _dh = diffusion_hyperparams
    T, alpha = _dh["T"], _dh["alpha"]

    B, C, L = f0.shape  
    #ts = torch.randint(T, size=(B, 1, 1)).cuda()  # ランダムなステップを生成 t=1~T → cudaへ # F0 Diffusionのもの。
    ts, pt = net.multinomial_diff.sample_time(B, device="cuda", method='importance')
    noise_gt = std_normal(f0.shape)                  # ノイズ生成
    ts=ts.view(B, 1, 1)
    
    # q(x_t|x_0)からx_tを計算
    delta = (1 - alpha[ts] ** 2.).sqrt() 
    alpha_cur = alpha[ts]
    noisy_f0 = alpha_cur * f0 + delta * noise_gt  

    # ノイズを予測する。
    f0_noise_pd, vuv_loss_minus = net(f0=noisy_f0,              # f0→noizy_f0へ修正済み。
                              f0_len=f0_lengths,
                              IDs=IDs, 
                              IDs_len=IDs_lengths, 
                              vuv=vuv,                  # 用修正、本来はnoisy_vuv ?
                              vuv_len=vuv_lengths,
                              attn=dur,
                              g=speakerID,
                              pt=pt,
                              timesteps=ts.view(B, 1), 
                              sampling=False)
    
    if reverse:
        x0 = (noisy_f0 - delta * f0_noise_pd) / alpha_cur
        return loss_fn(f0_noise_pd, noise_gt), x0
    
    num_elem = torch.sum(vuv_lengths)
    vuv_loss_minus = vuv_loss_minus.sum() / (math.log(2) * num_elem)
    loss_f0 = loss_fn(f0_noise_pd, noise_gt)
    return loss_f0 - vuv_loss_minus


def phi_loss(net, X, diffusion_hyperparams):
    """
    Compute the training loss for learning phi
    Parameters:
    net (torch network):            the wavenet models
    X (tuple, shape=(2,)):          training data in tuple form (mel_spectrograms, audios)
                                    mel_spectrograms: torch.tensor, shape is batchsize followed by each mel_spectrogram shape
                                    audios: torch.tensor, shape=(batchsize, 1, length of audio)
    diffusion_hyperparams (dict):   dictionary of diffusion hyperparameters returned by calc_diffusion_hyperparams
                                    note, the tensors need to be cuda tensors

    Returns:
    phi loss
    """
    assert type(X) == tuple and len(X) == 2
    _dh = diffusion_hyperparams
    T, alpha, tau = _dh["T"], _dh["alpha"], _dh["tau"]

    mel_spectrogram, audio = X
    B, C, L = audio.shape  # B is batchsize, C=1, L is audio length
    ts = torch.randint(tau, T - tau, size=(B,)).cuda()  # randomly sample steps from 1~T
    alpha_cur = alpha.index_select(0, ts).view(B, 1, 1)
    alpha_nxt = alpha.index_select(0, ts + tau).view(B, 1, 1)
    beta_nxt = 1 - (alpha_nxt / alpha_cur) ** 2.
    delta = (1 - alpha_cur ** 2.).sqrt()
    z = std_normal(audio.shape)
    noisy_audio = alpha_cur * audio + delta * z  # compute x_t from q(x_t|x_0)
    epsilon_theta = net((noisy_audio, mel_spectrogram, ts.view(B, 1),))
    beta_est = net.noise_pred(noisy_audio.squeeze(1), (beta_nxt.view(B, 1), delta.view(B, 1) ** 2.))
    phi_loss = 1 / (2. * (delta ** 2. - beta_est)) * (
            delta * z - beta_est / delta * epsilon_theta) ** 2.
    phi_loss += torch.log(1e-8 + delta ** 2. / (beta_est + 1e-8)) / 4.
    phi_loss = (torch.mean(phi_loss, -1, keepdim=True) + beta_est / delta ** 2 / 2.).mean()

    return phi_loss


def compute_hyperparams_given_schedule(beta):
    """
    Compute diffusion process hyperparameters

    Parameters:
    beta (tensor):  beta schedule

    Returns:
    a dictionary of diffusion hyperparameters including:
        T (int), beta/alpha/sigma (torch.tensor on cpu, shape=(T, ))
        These cpu tensors are changed to cuda tensors on each individual gpu
    """

    T = len(beta)
    alpha = 1 - beta
    sigma = beta + 0
    for t in range(1, T):
        alpha[t] *= alpha[t - 1]  # \alpha^2_t = \prod_{s=1}^t (1-\beta_s)
        sigma[t] *= (1 - alpha[t - 1]) / (1 - alpha[t])  # \sigma^2_t = \beta_t * (1-\alpha_{t-1}) / (1-\alpha_t)
    alpha = torch.sqrt(alpha)
    sigma = torch.sqrt(sigma)

    _dh = {}
    _dh["T"], _dh["beta"], _dh["alpha"], _dh["sigma"] = T, beta, alpha, sigma
    diffusion_hyperparams = _dh
    return diffusion_hyperparams



def map_noise_scale_to_time_step(alpha_infer, alpha):
    if alpha_infer < alpha[-1]:
        return len(alpha) - 1
    if alpha_infer > alpha[0]:
        return 0
    for t in range(len(alpha) - 1):
        if alpha[t+1] <= alpha_infer <= alpha[t]:
             step_diff = alpha[t] - alpha_infer
             step_diff /= alpha[t] - alpha[t+1]
             return t + step_diff.item()
    return -1


def calc_diffusion_step_embedding(diffusion_steps, diffusion_step_embed_dim_in):
    """
    Embed a diffusion step $t$ into a higher dimensional space
    E.g. the embedding vector in the 128-dimensional space is
    [sin(t * 10^(0*4/63)), ... , sin(t * 10^(63*4/63)), cos(t * 10^(0*4/63)), ... , cos(t * 10^(63*4/63))]

    Parameters:
    diffusion_steps (torch.long tensor, shape=(batchsize, 1)):
                                diffusion steps for batch data
    diffusion_step_embed_dim_in (int, default=128):
                                dimensionality of the embedding space for discrete diffusion steps

    Returns:
    the embedding vectors (torch.tensor, shape=(batchsize, diffusion_step_embed_dim_in)):
    """

    assert diffusion_step_embed_dim_in % 2 == 0

    half_dim = diffusion_step_embed_dim_in // 2
    _embed = np.log(10000) / (half_dim - 1)
    _embed = torch.exp(torch.arange(half_dim) * -_embed).cuda()
    _embed = diffusion_steps * _embed
    diffusion_step_embed = torch.cat((torch.sin(_embed),
                                      torch.cos(_embed)), 1)

    return diffusion_step_embed

if __name__ == "__main__":
    noise_gt = std_normal((1,1,800))
    print(torch.max(noise_gt))
    print(torch.min(noise_gt))
diffusion/diffusion_utils.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import soundfile as sf
import librosa 
import numpy as np
from typing import Optional,Union
try:
    from typing import Literal
except Exception as e:
    from typing_extensions import Literal
import numpy as np
import torch
import torchcrepe
from torch import nn
from torch.nn import functional as F
import scipy
import numpy 
import pyworld
import os
import torch.utils.data
import numpy as np
import torchaudio
import librosa
from librosa.filters import mel as librosa_mel_fn
import matplotlib.pyplot as plt

import os
import soundfile as sf
import librosa 
import numpy as np
from torchaudio.functional import resample
    
MAX_WAV_VALUE = 32768.0

f0_bin = 256
f0_max = 1100.0
f0_min = 50.0
f0_mel_min = 1127 * np.log(1 + f0_min / 700)
f0_mel_max = 1127 * np.log(1 + f0_max / 700)

class F0_extractor(object):
    def __init__(self, samplerate=44100, hop_size=512, device="cuda"):
        self.device = device
        self.fs = samplerate
        self.hop = hop_size
        
        self.f0_bin = f0_bin
        self.f0_max = f0_max
        self.f0_min = f0_min
        self.f0_mel_min = f0_mel_min
        self.f0_mel_max = f0_mel_max

        self.alert = False

    def compute_f0(self, path, f0_method="rmvpe"):
        vuv_vector = False

        x, sr = torchaudio.load(path)
        if sr != self.fs:
            x = resample(x, sr, self.fs)
            if self.alert is False:
                print(f"[WARNING] Detected a discrepancy between the input audio sampling rate and the target sampling rate.")
                self.alert = True
        x = x[0].to('cpu').detach().numpy().copy()
        
        if f0_method == "rmvpe":
            if hasattr(self, "model_rmvpe") == False:
                self.model_dir = "model_data"
                os.makedirs(self.model_dir, exist_ok=True)
                self.rmvpe_filename = "rmvpe.pt"
                self.rmvpe_modelpath = os.path.join(self.model_dir, self.rmvpe_filename)
                if os.path.isfile(self.rmvpe_modelpath) is False :
                    import urllib.request
                    url='https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/rmvpe.pt'
                    print("[INFO] Rmvpe model is not found, download from HuggingFace.")
                    urllib.request.urlretrieve(url, self.rmvpe_modelpath)
                    print(f"[INFO] Done. Rmvpe model path : {self.rmvpe_modelpath}")
                from rmvpe import RMVPE
                print("[INFO] Load rmvpe model")
                self.model_rmvpe = RMVPE(model_path=self.rmvpe_modelpath, 
                                         is_half=True,
                                         sampling_rate=self.fs, 
                                         hop_length=self.hop,
                                         win_size=2048,
                                         f0_mel_min = self.f0_mel_min,
                                         f0_mel_max = self.f0_mel_max, 
                                         device=self.device)
            f0 = self.model_rmvpe.infer_from_audio(x, thred=0.03)

        
        elif f0_method == "harvest":
            f0 = compute_f0_harvest(x, sampling_rate=self.fs, hop_length=512)

        elif f0_method == "parselmouth":
            f0 = compute_f0_parselmouth(x, sampling_rate=self.fs, hop_length=512)

        elif f0_method == "crepe":
            f0, vuv_vector = compute_f0_torchcrepe(x, sampling_rate=self.fs, hop_length=512)

        elif f0_method == "dio":
            f0 = compute_f0_dio(x, sampling_rate=self.fs, hop_length=512)


        elif f0_method == "check_f0_method":
            f0_crepe, _= compute_f0_torchcrepe(x, sampling_rate=self.fs, hop_length=512)
            f0_harvest = compute_f0_harvest(x, sampling_rate=self.fs, hop_length=512)
            f0_dio = compute_f0_dio(x, sampling_rate=self.fs, hop_length=512)
            f0_parselmouth = compute_f0_parselmouth(x, sampling_rate=self.fs, hop_length=512)

            if hasattr(self, "model_rmvpe") == False:
                self.rmvpe_filename = "rmvpe.pt"
                self.rmvpe_modelpath = os.path.join(self.model_dir, self.rmvpe_filename)
                if os.path.isfile(self.rmvpe_modelpath) is False :
                    import urllib.request
                    url='https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/rmvpe.pt'
                    print("[INFO] Rmvpe model is not found, download from HuggingFace.")
                    urllib.request.urlretrieve(url, self.rmvpe_modelpath)
                    print(f"[INFO] Done. Rmvpe model path : {self.rmvpe_modelpath}")
                from lib.rmvpe import RMVPE
                print("[INFO] Load rmvpe model")
                self.model_rmvpe = RMVPE(self.rmvpe_modelpath, 
                                         sampling_rate=self.fs, 
                                         hop_length=self.hop,
                                         is_half=True, 
                                         device=self.device)
            f0_rmvpe = self.model_rmvpe.infer_from_audio(x, thred=0.03)

            x = numpy.linspace(0, 1, len(f0_parselmouth))
            plt.plot(x, f0_crepe, label="crepe")
            plt.plot(x, f0_harvest, label="harvest")
            plt.plot(x, f0_dio, label="dio")
            plt.plot(x, f0_parselmouth, label="pm")
            f0_rmvpe = f0_rmvpe[:len(f0_dio)]
            plt.plot(x, f0_parselmouth, label="rmvpe")
            plt.legend()
            plt.show()
            plt.close()
            f0 = None

        if vuv_vector is False:
            f0, vuv_vector = interpolate_f0(f0) # vuvは0が発話無、1が発話アリ
        return f0, vuv_vector

    def coarse_f0(self, path, f0_method="rmvpe"):
        f0 = self.compute_f0(path, f0_method=f0_method)
        f0_mel = 1127 * np.log(1 + f0 / 700)
        f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - self.f0_mel_min) * (
            self.f0_bin - 2
        ) / (self.f0_mel_max - self.f0_mel_min) + 1

        # use 0 or 1
        f0_mel[f0_mel <= 1] = 1
        f0_mel[f0_mel > self.f0_bin - 1] = self.f0_bin - 1
        f0_coarse = np.rint(f0_mel).astype(int)
        assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (
            f0_coarse.max(),
            f0_coarse.min(),
        )
        return f0_coarse

def normalize_f0(f0, x_mask, uv, random_scale=True):
    # calculate means based on x_mask
    uv_sum = torch.sum(uv, dim=1, keepdim=True)
    uv_sum[uv_sum == 0] = 9999
    means = torch.sum(f0[:, 0, :] * uv, dim=1, keepdim=True) / uv_sum

    if random_scale:
        factor = torch.Tensor(f0.shape[0], 1).uniform_(0.8, 1.2).to(f0.device)
    else:
        factor = torch.ones(f0.shape[0], 1).to(f0.device)
    # normalize f0 based on means and factor
    f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)
    if torch.isnan(f0_norm).any():
        exit(0)
    return f0_norm * x_mask

def compute_f0_uv_torchcrepe(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512,device=None,cr_threshold=0.05):
    
    x = wav_numpy
    if p_len is None:
        p_len = x.shape[0]//hop_length
    else:
        assert abs(p_len-x.shape[0]//hop_length) < 4, "pad length error"
    
    F0Creper = CrepePitchExtractor(hop_length=hop_length,f0_min=f0_min,f0_max=f0_max,device=device,threshold=cr_threshold)
    f0,uv = F0Creper(x[None,:].float(),sampling_rate,pad_to=p_len)
    return f0,uv


def compute_f0_parselmouth(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
    import parselmouth
    x = wav_numpy
    if p_len is None:
        p_len = x.shape[0]//hop_length
    else:
        assert abs(p_len-x.shape[0]//hop_length) < 4, "pad length error"
    time_step = hop_length / sampling_rate * 1000
    f0 = parselmouth.Sound(x, sampling_rate).to_pitch_ac(
        time_step=time_step / 1000, voicing_threshold=0.6,
        pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']

    pad_size=(p_len - len(f0) + 1) // 2
    if(pad_size>0 or p_len - len(f0) - pad_size>0):
        f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
    return f0

def resize_f0(x, target_len):
    source = np.array(x)
    source[source<0.001] = np.nan
    target = np.interp(np.arange(0, len(source)*target_len, len(source))/ target_len, np.arange(0, len(source)), source)
    res = np.nan_to_num(target)
    return res

def compute_f0_dio(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
    import pyworld
    if p_len is None:
        p_len = wav_numpy.shape[0]//hop_length
    f0, t = pyworld.dio(
        wav_numpy.astype(np.double),
        fs=sampling_rate,
        f0_ceil=800,
        frame_period=1000 * hop_length / sampling_rate,
    )
    f0 = pyworld.stonemask(wav_numpy.astype(np.double), f0, t, sampling_rate)
    for index, pitch in enumerate(f0):
        f0[index] = round(pitch, 1)
    return resize_f0(f0, p_len)

def f0_to_coarse(f0):
  is_torch = isinstance(f0, torch.Tensor)
  f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
  f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1

  f0_mel[f0_mel <= 1] = 1
  f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
  f0_coarse = (f0_mel + 0.5).int() if is_torch else np.rint(f0_mel).astype(np.int)
  assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())
  return f0_coarse


def interpolate_f0(f0):

    data = np.reshape(f0, (f0.size, 1))

    vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
    vuv_vector[data > 0.0] = 1.0
    vuv_vector[data <= 0.0] = 0.0

    ip_data = data

    frame_number = data.size
    last_value = 0.0
    for i in range(frame_number):
        if data[i] <= 0.0:
            j = i + 1
            for j in range(i + 1, frame_number):
                if data[j] > 0.0:
                    break
            if j < frame_number - 1:
                if last_value > 0.0:
                    step = (data[j] - data[i - 1]) / float(j - i)
                    for k in range(i, j):
                        ip_data[k] = data[i - 1] + step * (k - i + 1)
                else:
                    for k in range(i, j):
                        ip_data[k] = data[j]
            else:
                for k in range(i, frame_number):
                    ip_data[k] = last_value
        else:
            ip_data[i] = data[i] # this may not be necessary
            last_value = data[i]

    return ip_data[:,0], vuv_vector[:,0]


def repeat_expand_2d(content, target_len):
    # content : [h, t]

    src_len = content.shape[-1]
    target = torch.zeros([content.shape[0], target_len], dtype=torch.float).to(content.device)
    temp = torch.arange(src_len+1) * target_len / src_len
    current_pos = 0
    for i in range(target_len):
        if i < temp[current_pos+1]:
            target[:, i] = content[:, current_pos]
        else:
            current_pos += 1
            target[:, i] = content[:, current_pos]

    return target

def dynamic_range_compression_torch(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_torch(x, C=1):
    """
    PARAMS
    ------
    C: compression factor used to compress
    """
    return torch.exp(x) / C


def spectral_normalize_torch(magnitudes):
    output = dynamic_range_compression_torch(magnitudes)
    return output


def spectral_de_normalize_torch(magnitudes):
    output = dynamic_range_decompression_torch(magnitudes)
    return output


mel_basis = {}
hann_window = {}


def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
    if torch.min(y) < -1.:
        print('min value is ', torch.min(y))
    if torch.max(y) > 1.:
        print('max value is ', torch.max(y))

    global hann_window
    dtype_device = str(y.dtype) + '_' + str(y.device)
    wnsize_dtype_device = str(win_size) + '_' + dtype_device
    if wnsize_dtype_device not in hann_window:
        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)

    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
    y = y.squeeze(1)

    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)

    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
    return spec


def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
    global mel_basis
    dtype_device = str(spec.dtype) + '_' + str(spec.device)
    fmax_dtype_device = str(fmax) + '_' + dtype_device
    if fmax_dtype_device not in mel_basis:
        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=spec.dtype, device=spec.device)
    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
    spec = spectral_normalize_torch(spec)
    return spec


def mel_spectrogram_torch(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
    if torch.min(y) < -1.:
        print('min value is ', torch.min(y))
    if torch.max(y) > 1.:
        print('max value is ', torch.max(y))

    global mel_basis, hann_window
    dtype_device = str(y.dtype) + '_' + str(y.device)
    fmax_dtype_device = str(fmax) + '_' + dtype_device
    wnsize_dtype_device = str(win_size) + '_' + dtype_device
    if fmax_dtype_device not in mel_basis:
        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=y.dtype, device=y.device)
    if wnsize_dtype_device not in hann_window:
        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)

    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
    y = y.squeeze(1)

    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)

    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)

    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
    spec = spectral_normalize_torch(spec)

    return spec


def plot_spectrogram_to_numpy(spectrogram, count=0):
  global MATPLOTLIB_FLAG
  if not MATPLOTLIB_FLAG:
    import matplotlib
    matplotlib.use("Agg")
    MATPLOTLIB_FLAG = True
    mpl_logger = logging.getLogger('matplotlib')
    mpl_logger.setLevel(logging.WARNING)
  import matplotlib.pylab as plt
  import numpy as np
  
  fig, ax = plt.subplots(figsize=(10,2))
  im = ax.imshow(spectrogram, aspect="auto", origin="lower",
                  interpolation='none')
  plt.colorbar(im, ax=ax)
  plt.xlabel("Frames")
  plt.ylabel("Channels")
  plt.tight_layout()

  fig.canvas.draw()
  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
  plt.savefig(f"{count}_fft.png")
  plt.close()
  return data

f0_bin = 256
f0_max = 1100.0
f0_min = 50.0
f0_mel_min = 1127 * numpy.log(1 + f0_min / 700)
f0_mel_max = 1127 * numpy.log(1 + f0_max / 700)

def normalize_f0(f0, x_mask, uv, random_scale=True):
    # calculate means based on x_mask
    uv_sum = torch.sum(uv, dim=1, keepdim=True)
    uv_sum[uv_sum == 0] = 9999
    means = torch.sum(f0[:, 0, :] * uv, dim=1, keepdim=True) / uv_sum

    if random_scale:
        factor = torch.Tensor(f0.shape[0], 1).uniform_(0.8, 1.2).to(f0.device)
    else:
        factor = torch.ones(f0.shape[0], 1).to(f0.device)
    # normalize f0 based on means and factor
    f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)
    if torch.isnan(f0_norm).any():
        exit(0)
    return f0_norm * x_mask

def compute_f0_torchcrepe(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512,device=None,cr_threshold=0.05):
    x = wav_numpy
    if p_len is None:
        p_len = x.shape[0]//hop_length
    else:
        assert abs(p_len-x.shape[0]//hop_length) < 4, "pad length error"
    
    x = torch.from_numpy(x.astype(numpy.float32)).clone()
    F0Creper = CrepePitchExtractor(hop_length=hop_length,f0_min=f0_min,f0_max=f0_max,device=device,threshold=cr_threshold)
    f0,uv = F0Creper(x[None,:].float(),sampling_rate,pad_to=p_len)
    f0[uv<0.5] = 0
    return f0,uv

def compute_f0_harvest(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
    x = wav_numpy
    if p_len is None:
        p_len = x.shape[0]//hop_length
    else:
        assert abs(p_len-x.shape[0]//hop_length) < 4, "pad length error"
    
    f0, t = pyworld.harvest(
        x.astype(numpy.double),
        fs=sampling_rate,
        f0_ceil=f0_max,
        f0_floor=f0_min,
        frame_period=1000 * hop_length / sampling_rate,
    )
    f0 = pyworld.stonemask(x.astype(numpy.double), f0, t, sampling_rate)
    return resize_f0(f0, p_len)

def compute_f0_parselmouth(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
    import parselmouth
    x = wav_numpy
    if p_len is None:
        p_len = x.shape[0]//hop_length
    else:
        assert abs(p_len-x.shape[0]//hop_length) < 4, "pad length error"
    time_step = hop_length / sampling_rate * 1000
    f0 = parselmouth.Sound(x, sampling_rate).to_pitch_ac(
        time_step=time_step / 1000, voicing_threshold=0.6,
        pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']

    pad_size=(p_len - len(f0) + 1) // 2
    if(pad_size>0 or p_len - len(f0) - pad_size>0):
        f0 = numpy.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
    return f0

def resize_f0(x, target_len):
    source = numpy.array(x)
    source[source<0.001] = numpy.nan
    target = numpy.interp(numpy.arange(0, len(source)*target_len, len(source))/ target_len, numpy.arange(0, len(source)), source)
    res = numpy.nan_to_num(target)
    return res

def compute_f0_dio(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
    import pyworld
    if p_len is None:
        p_len = wav_numpy.shape[0]//hop_length
    f0, t = pyworld.dio(
        wav_numpy.astype(numpy.double),
        fs=sampling_rate,
        f0_ceil=f0_max,
        f0_floor=f0_min,
        frame_period=1000 * hop_length / sampling_rate,
    )
    f0 = pyworld.stonemask(wav_numpy.astype(numpy.double), f0, t, sampling_rate)
    for index, pitch in enumerate(f0):
        f0[index] = round(pitch, 1)
    return resize_f0(f0, p_len)

def f0_to_coarse(f0):
  is_torch = isinstance(f0, torch.Tensor)
  f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * numpy.log(1 + f0 / 700)
  f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1

  f0_mel[f0_mel <= 1] = 1
  f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
  f0_coarse = (f0_mel + 0.5).int() if is_torch else numpy.rint(f0_mel).astype(numpy.int)
  assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())
  return f0_coarse


def interpolate_f0(f0):

    data = numpy.reshape(f0, (f0.size, 1))

    vuv_vector = numpy.zeros((data.size, 1), dtype=numpy.float32)
    vuv_vector[data > 0.0] = 1.0
    vuv_vector[data <= 0.0] = 0.0

    ip_data = data

    frame_number = data.size
    last_value = 0.0
    for i in range(frame_number):
        if data[i] <= 0.0:
            j = i + 1
            for j in range(i + 1, frame_number):
                if data[j] > 0.0:
                    break
            if j < frame_number - 1:
                if last_value > 0.0:
                    step = (data[j] - data[i - 1]) / float(j - i)
                    for k in range(i, j):
                        ip_data[k] = data[i - 1] + step * (k - i + 1)
                else:
                    for k in range(i, j):
                        ip_data[k] = data[j]
            else:
                for k in range(i, frame_number):
                    ip_data[k] = last_value
        else:
            ip_data[i] = data[i] # this may not be necessary
            last_value = data[i]

    return ip_data[:,0], vuv_vector[:,0]

#from:https://github.com/fishaudio/fish-diffusion

def repeat_expand(
    content: Union[torch.Tensor, numpy.ndarray], target_len: int, mode: str = "nearest"
):
    """Repeat content to target length.
    This is a wrapper of torch.nn.functional.interpolate.

    Args:
        content (torch.Tensor): tensor
        target_len (int): target length
        mode (str, optional): interpolation mode. Defaults to "nearest".

    Returns:
        torch.Tensor: tensor
    """

    ndim = content.ndim

    if content.ndim == 1:
        content = content[None, None]
    elif content.ndim == 2:
        content = content[None]

    assert content.ndim == 3

    is_np = isinstance(content, numpy.ndarray)
    if is_np:
        content = torch.from_numpy(content)

    results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)

    if is_np:
        results = results.numpy()

    if ndim == 1:
        return results[0, 0]
    elif ndim == 2:
        return results[0]


class BasePitchExtractor:
    def __init__(
        self,
        hop_length: int = 512,
        f0_min: float = 50.0,
        f0_max: float = 1100.0,
        keep_zeros: bool = True,
    ):
        """Base pitch extractor.

        Args:
            hop_length (int, optional): Hop length. Defaults to 512.
            f0_min (float, optional): Minimum f0. Defaults to 50.0.
            f0_max (float, optional): Maximum f0. Defaults to 1100.0.
            keep_zeros (bool, optional): Whether keep zeros in pitch. Defaults to True.
        """

        self.hop_length = hop_length
        self.f0_min = f0_min
        self.f0_max = f0_max
        self.keep_zeros = keep_zeros

    def __call__(self, x, sampling_rate=44100, pad_to=None):
        raise NotImplementedError("BasePitchExtractor is not callable.")

    def post_process(self, x, sampling_rate, f0, pad_to):
        if isinstance(f0, numpy.ndarray):
            f0 = torch.from_numpy(f0).float().to(x.device)

        if pad_to is None:
            return f0

        f0 = repeat_expand(f0, pad_to)

        if self.keep_zeros:
            return f0
        
        vuv_vector = torch.zeros_like(f0)
        vuv_vector[f0 > 0.0] = 1.0
        vuv_vector[f0 <= 0.0] = 0.0
        
        # Remove 0 frequency and apply linear interpolation
        nzindex = torch.nonzero(f0).squeeze()
        f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
        time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
        time_frame = numpy.arange(pad_to) * self.hop_length / sampling_rate

        if f0.shape[0] <= 0:
            return torch.zeros(pad_to, dtype=torch.float, device=x.device),torch.zeros(pad_to, dtype=torch.float, device=x.device)

        if f0.shape[0] == 1:
            return torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0],torch.ones(pad_to, dtype=torch.float, device=x.device)
    
        # Probably can be rewritten with torch?
        f0 = numpy.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
        vuv_vector = vuv_vector.cpu().numpy()
        vuv_vector = numpy.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
        
        return f0,vuv_vector


class MaskedAvgPool1d(nn.Module):
    def __init__(
        self, kernel_size: int, stride: Optional[int] = None, padding: Optional[int] = 0
    ):
        """An implementation of mean pooling that supports masked values.

        Args:
            kernel_size (int): The size of the median pooling window.
            stride (int, optional): The stride of the median pooling window. Defaults to None.
            padding (int, optional): The padding of the median pooling window. Defaults to 0.
        """

        super(MaskedAvgPool1d, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride or kernel_size
        self.padding = padding

    def forward(self, x, mask=None):
        ndim = x.dim()
        if ndim == 2:
            x = x.unsqueeze(1)

        assert (
            x.dim() == 3
        ), "Input tensor must have 2 or 3 dimensions (batch_size, channels, width)"

        # Apply the mask by setting masked elements to zero, or make NaNs zero
        if mask is None:
            mask = ~torch.isnan(x)

        # Ensure mask has the same shape as the input tensor
        assert x.shape == mask.shape, "Input tensor and mask must have the same shape"

        masked_x = torch.where(mask, x, torch.zeros_like(x))
        # Create a ones kernel with the same number of channels as the input tensor
        ones_kernel = torch.ones(x.size(1), 1, self.kernel_size, device=x.device)

        # Perform sum pooling
        sum_pooled = nn.functional.conv1d(
            masked_x,
            ones_kernel,
            stride=self.stride,
            padding=self.padding,
            groups=x.size(1),
        )

        # Count the non-masked (valid) elements in each pooling window
        valid_count = nn.functional.conv1d(
            mask.float(),
            ones_kernel,
            stride=self.stride,
            padding=self.padding,
            groups=x.size(1),
        )
        valid_count = valid_count.clamp(min=1)  # Avoid division by zero

        # Perform masked average pooling
        avg_pooled = sum_pooled / valid_count

        # Fill zero values with NaNs
        avg_pooled[avg_pooled == 0] = float("nan")

        if ndim == 2:
            return avg_pooled.squeeze(1)

        return avg_pooled


class MaskedMedianPool1d(nn.Module):
    def __init__(
        self, kernel_size: int, stride: Optional[int] = None, padding: Optional[int] = 0
    ):
        """An implementation of median pooling that supports masked values.

        This implementation is inspired by the median pooling implementation in
        https://gist.github.com/rwightman/f2d3849281624be7c0f11c85c87c1598

        Args:
            kernel_size (int): The size of the median pooling window.
            stride (int, optional): The stride of the median pooling window. Defaults to None.
            padding (int, optional): The padding of the median pooling window. Defaults to 0.
        """

        super(MaskedMedianPool1d, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride or kernel_size
        self.padding = padding

    def forward(self, x, mask=None):
        ndim = x.dim()
        if ndim == 2:
            x = x.unsqueeze(1)

        assert (
            x.dim() == 3
        ), "Input tensor must have 2 or 3 dimensions (batch_size, channels, width)"

        if mask is None:
            mask = ~torch.isnan(x)

        assert x.shape == mask.shape, "Input tensor and mask must have the same shape"

        masked_x = torch.where(mask, x, torch.zeros_like(x))

        x = F.pad(masked_x, (self.padding, self.padding), mode="reflect")
        mask = F.pad(
            mask.float(), (self.padding, self.padding), mode="constant", value=0
        )

        x = x.unfold(2, self.kernel_size, self.stride)
        mask = mask.unfold(2, self.kernel_size, self.stride)

        x = x.contiguous().view(x.size()[:3] + (-1,))
        mask = mask.contiguous().view(mask.size()[:3] + (-1,)).to(x.device)

        # Combine the mask with the input tensor
        #x_masked = torch.where(mask.bool(), x, torch.fill_(torch.zeros_like(x),float("inf")))
        x_masked = torch.where(mask.bool(), x, torch.FloatTensor([float("inf")]).to(x.device))

        # Sort the masked tensor along the last dimension
        x_sorted, _ = torch.sort(x_masked, dim=-1)

        # Compute the count of non-masked (valid) values
        valid_count = mask.sum(dim=-1)

        # Calculate the index of the median value for each pooling window
        median_idx = (torch.div((valid_count - 1), 2, rounding_mode='trunc')).clamp(min=0)

        # Gather the median values using the calculated indices
        median_pooled = x_sorted.gather(-1, median_idx.unsqueeze(-1).long()).squeeze(-1)

        # Fill infinite values with NaNs
        median_pooled[torch.isinf(median_pooled)] = float("nan")
        
        if ndim == 2:
            return median_pooled.squeeze(1)

        return median_pooled


class CrepePitchExtractor(BasePitchExtractor):
    def __init__(
        self,
        hop_length: int = 512,
        f0_min: float = 50.0,
        f0_max: float = 1100.0,
        threshold: float = 0.05,
        keep_zeros: bool = False,
        device = None,
        model: Literal["full", "tiny"] = "full",
        use_fast_filters: bool = True,
    ):
        super().__init__(hop_length, f0_min, f0_max, keep_zeros)

        self.threshold = threshold
        self.model = model
        self.use_fast_filters = use_fast_filters
        self.hop_length = hop_length
        if device is None:
            self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.dev = torch.device(device)
        if self.use_fast_filters:
            self.median_filter = MaskedMedianPool1d(3, 1, 1).to(device)
            self.mean_filter = MaskedAvgPool1d(3, 1, 1).to(device)

    def __call__(self, x, sampling_rate=44100, pad_to=None):
        """Extract pitch using crepe.


        Args:
            x (torch.Tensor): Audio signal, shape (1, T).
            sampling_rate (int, optional): Sampling rate. Defaults to 44100.
            pad_to (int, optional): Pad to length. Defaults to None.

        Returns:
            torch.Tensor: Pitch, shape (T // hop_length,).
        """

        assert x.ndim == 2, f"Expected 2D tensor, got {x.ndim}D tensor."
        assert x.shape[0] == 1, f"Expected 1 channel, got {x.shape[0]} channels."

        x = x.to(self.dev)
        f0, pd = torchcrepe.predict(
            x,
            sampling_rate,
            self.hop_length,
            self.f0_min,
            self.f0_max,
            pad=True,
            model=self.model,
            batch_size=1024,
            device=x.device,
            return_periodicity=True,
        )

        # Filter, remove silence, set uv threshold, refer to the original warehouse readme
        if self.use_fast_filters:
            pd = self.median_filter(pd)
        else:
            pd = torchcrepe.filter.median(pd, 3)

        pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, 512)
        f0 = torchcrepe.threshold.At(self.threshold)(f0, pd)
        
        if self.use_fast_filters:
            f0 = self.mean_filter(f0)
        else:
            f0 = torchcrepe.filter.mean(f0, 3)

        f0 = torch.where(torch.isnan(f0), torch.full_like(f0, 0), f0)[0]

        if torch.all(f0 == 0):
            rtn = f0.cpu().numpy() if pad_to==None else numpy.zeros(pad_to)
            return rtn,rtn
        
        return self.post_process(x, sampling_rate, f0, pad_to)



if __name__ == "__main__":
   testwav = "./basic5000\wav\BASIC5000_0001.wav"

   model = F0_extractor()
   s = model.compute_f0(path = testwav)
f0_extractor.py
tonnetonne814/SiFi-VITS2-44100-Ja
import warnings
warnings.filterwarnings(action='ignore')

import os
import time
import torch
import utils
import argparse
import utils
from models import VITS2_based_SiFiSinger, DiffusionModels
import soundcard as sc
import soundfile as sf
from singDB_loader import get_ust_info, get_g2p_dict_from_tabledata
import commons
from sifigan.utils import dilated_factor        
from sifigan.utils.features import SignalGenerator
import numpy as np
import logger
import copy
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import interp1d

def inference(args):

    ###### CPU Inference is not allowed ######
    assert torch.cuda.is_available(), "CPU Inference is not allowed."
    ##########################################

    config_path = args.config
    F0_model_path = args.F0diff_model_path
    Synthesizer_model_path = args.Synthesizer_model_path
    ust_path = args.UST_path
    
    ask_retake = args.ask_retake

    # load config.json
    hps = utils.get_hparams(config_path)
    oto2lab_tablepath = hps["oto2lab_path"]

    # save directory process
    if args.is_save is True:
        save_dir = os.path.join("./infer_logs/")
        os.makedirs(save_dir, exist_ok=True)

    # preprocess
    g2p, ph_symbol_to_id, _, _, _ = get_g2p_dict_from_tabledata(oto2lab_tablepath, include_converter=True)
    inference_loader = preprocess_dataset(get_ust_info(ust_path), g2p, ph_symbol_to_id, hps)

    # load checkpoint
    SVS = VITS2_based_SiFiSinger(
      hps = hps,
      n_vocab=len(ph_symbol_to_id) + 1, # +1 means masked symbol id
      spec_channels=hps["spec_encoder"]["spec_channels"],
      segment_size=int(hps["segments_size"] // hps["hop_length"]),
      n_speakers=hps["common"]["n_speaker"],
      gin_channels=hps["common"]["gin_channels"],
      **hps["hifi_gan"]).cuda()
    _ = SVS.eval()
    utils.load_model_diffsize(Synthesizer_model_path, SVS, hps, None)

    # モデルを作成
    F0DIFF = DiffusionModels(hps=hps, vocab_size=len(ph_symbol_to_id)+1, training=False).cuda()
    _ = F0DIFF.eval()
    utils.load_model_diffsize(F0_model_path, F0DIFF, hps, None)

    # for SiFiGAN
    df_f0_type = hps["SiFiGAN_utils"]["df_f0_type"]
    dense_factors = hps["SiFiGAN_utils"]["dense_factors"]
    upsample_scales= hps["SiFiGANGenerator"]["upsample_scales"]
    prod_upsample_scales = np.cumprod(upsample_scales)
    df_sample_rates = [hps["sampling_rate"] / hps["hop_length"] * s for s in prod_upsample_scales]
    signal_generator = SignalGenerator(sample_rate=hps["sampling_rate"],
                                            hop_size=hps["hop_length"] ,
                                            sine_amp=hps["SiFiGAN_utils"]["sine_amp"],
                                            noise_amp=hps["SiFiGAN_utils"]["noise_amp"],
                                            signal_types=hps["SiFiGAN_utils"]["signal_types"],)

    # play audio by system default
    speaker = sc.get_speaker(sc.default_speaker().name)

    out_voices = list()
    diff_lists = list()
    total_lengths = list()

    # required_grad is False
    with torch.inference_mode():
        for idx, (ph_seq, ph_length, word_frame_dur, word_lengths, word_dur_ms, ph_idx_in_a_word, n_ph_seq,\
             noteid_seq, noteid_length) in enumerate(inference_loader):
            speakerID = torch.zeros(size=(1,), dtype=torch.int64) # speaker id is fixed

            # for retake flag
            retake_flag = True
            while retake_flag:

                # measure the execution time 
                torch.cuda.synchronize()
                start = time.time()

                # 長い休符の時は、無音出力で終わらせる。
                if ph_seq[0][0] == 35 and word_dur_ms[0][0]>3:
                    length=int(word_dur_ms[0][0] * hps["sampling_rate"])
                    data = np.zeros(shape=(length,))
                    save_path = os.path.join(save_dir, str(idx).zfill(3)+".wav")
                    sf.write(
                        file=save_path,
                        data=data,
                        samplerate=hps["sampling_rate"],
                        format="WAV")
                    out_voices.append(data)
                    diff_lists.append(0)
                    total_lengths.append(length)
                    retake_flag = False
                    print(f"[INFO] Silent wavs is saved at {save_path}")
                    continue

                # duration predict and encoding
                z_spec, _, w_ceil, g, diff = SVS.encode_and_dp(ph_IDs=ph_seq.cuda(),
                                                    ph_IDs_lengths=ph_length.cuda(),
                                                    speakerID=speakerID.cuda(),
                                                    word_frame_dur=word_frame_dur.cuda(),
                                                    word_frame_dur_lenngths=word_lengths.cuda(),
                                                    word_dur_ms=word_dur_ms.cuda(),
                                                    ph_word_flag=ph_idx_in_a_word.cuda(),
                                                    n_ph_pool=n_ph_seq.cuda(),
                                                    noise_scale=0.5)

                # f0 diffusion [1,1,T']
                f0_pd = F0DIFF.sampling(condition=[ph_seq.cuda(), ph_length.cuda(), w_ceil.cuda(), 
                                      noteid_seq.cuda(), noteid_length.cuda(), None, speakerID.cuda()])

                # Below threshold, truncated
                f0_pd [f0_pd<int(hps["f0_min"])] = 0

                # calc continuousf0 and dilated factors
                _, continuous_f0, _ = get_continuos_f0(f0_pd[0][0].to('cpu').detach().numpy().copy())
                Sinewave = signal_generator(torch.tensor(continuous_f0, dtype=torch.float32).view(1,1,-1))# for SiFiGAN
                dfs = []
                for df, us in zip(dense_factors, prod_upsample_scales):
                    dfs += [
                        torch.from_numpy(np.repeat(dilated_factor(continuous_f0, hps["sampling_rate"], df), us).astype(np.float32)).clone()
                        if df_f0_type == "cf0"
                        else torch.from_numpy(np.repeat(dilated_factor(continuous_f0, hps["sampling_rate"], df), us).astype(np.float32)).clone()
                    ]

                # check length 
                for i in range(len(dense_factors)):
                    assert Sinewave.size(2) * df_sample_rates[i] == len(dfs[i]) * hps["sampling_rate"]

                # synthesize voice=[1,1,T]
                voice = SVS.synthesize(Sinewave.float().cuda(), z_spec, dfs, g)
                
                # measure the execution time 
                torch.cuda.synchronize()
                elapsed_time = time.time() - start
                print(f"Gen Time : {elapsed_time}")

                if ask_retake is True:
                    speaker.play(data=voice[0][0].to('cpu').detach().numpy().copy(),
                                 samplerate=hps["sampling_rate"],
                                 channels=1)
                    x = input("OK:0, Retake:1  select==>")
                    if x == "":
                        retake_flag = False
                    elif int(x) == 1:
                        pass
                    else:
                        retake_flag = False

                else:
                    retake_flag = False
            
            if ph_seq[0][0] == 35 and word_dur_ms[0][0]>3:
                continue

            diff_lists.append(diff) # first ph duration difference
            total_lengths.append(int(torch.sum(word_dur_ms) * hps["sampling_rate"])) # total note length 
            out_voices.append(voice[0][0].to('cpu').detach().numpy().copy()) # synthesized voice

            if args.is_save is True:
                data = voice[0][0].to('cpu').detach().numpy().copy()
                save_path = os.path.join(save_dir, str(idx).zfill(3)+".wav")
                sf.write(
                    file=save_path,
                    data=data,
                    samplerate=hps["sampling_rate"],
                    format="WAV")
                print(f"[INFO] Save synthesized voice at {save_path}")

                f0_pd = f0_pd[0][0].to('cpu').detach().numpy().copy()
                generate_graph(vector=f0_pd , 
                            label="F0 gt",
                            color="red", 
                            x_label = 'Frames',
                            y_label = "f0",
                            savename = os.path.join(save_dir, str(idx).zfill(3)+".png"))

    if len(out_voices) != 1:
        concat_voices(out_voices, save_dir, hps, diff_lists, total_lengths)
    return 0

def concat_voices(out_voices,save_dir, hps, diff_lists, total_lengths):
    lengths = 0
    for idx, voice in enumerate(out_voices):
        lengths += len(voice)

    wavs = np.zeros(shape=(int(lengths,)))

    z_idx = 0
    for idx, (voice, diff, total) in enumerate(zip(out_voices, diff_lists, total_lengths)):
        start_idx = z_idx-diff
        wavs[start_idx:start_idx+len(voice)] += voice
        z_idx += total

    sf.write(file=os.path.join(save_dir, "full_chorus.wav"),
             data=wavs[:z_idx],
             samplerate=hps["sampling_rate"],
             format="WAV")

def generate_graph(vector, 
                   label="NoName",
                   color="blue", 
                   title="Title",
                   x_label = 'Frames',
                   y_label = "y_labels",
                   figsize=(20,5),
                   savename=None):
   
    fig = plt.figure(figsize=figsize) 
    x = np.arange(0, len(vector))
    plt.plot(x, vector, label=label, color=color)
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    #plt.ylim(y_val_min, y_val_max)
    fig.canvas.draw() 
    plt.savefig(savename)
    plt.clf()
    plt.close()
    
def preprocess_dataset(UST_data, g2p, ph2id, hps):

    with open(hps["noteid2hz_txt_path"], mode="r", encoding="utf-8") as f:
        lines = f.readlines()
    id_to_hz = {}
    for idx, line in enumerate(lines):
        id, hz = line.split(",")
        id_to_hz[idx-1] = float(hz)

    word_seq = list()
    for idx, ust in enumerate(UST_data):
        #print(ust)
        ms = ust["length_ms"]
        if ms == 0:
            continue
        noteid = float(id_to_hz[ust["note_MIDI"]])
        lyric = ust["lyric"].split(" ")[-1]
        if "っ" in lyric and len(lyric)>1:
            txt = ""   
            word_ph = ""   
            # check cl
            for w in lyric:
                if w == "っ":
                    word_ph += g2p[txt] + " " + g2p[w] # pull from dict
                    txt = "" #reset
                else:
                    txt += w # stack
            # last pull
            if txt != "":
                word_ph += g2p[txt]
        elif lyric == "ー":
            ph = z_ust["lyric"]
            if len(ph) == 1:
                word_ph = g2p[ph]
            else:
                word_ph = g2p[ph[-1]]
        else:      
            word_ph = g2p[lyric]
        word_ph = word_ph.split(" ")
        n_ph = len(word_ph)
        id_seq = []
        for ph in word_ph:
            #print(ph)
            id = ph2id[ph]
            id_seq += [id]
        word_seq .append( [ms, id_seq, n_ph, noteid])
        z_ust = ust
    batch_process = list()
    total_time = 0
    ph_seq = []
    n_ph_seq = []
    noteid_seq = []
    ms_seq =[]
    for idx, (ms, ph_ids, n_ph, noteid) in enumerate(word_seq):
        ms_seq      += [ms]
        ph_seq      += [ph_ids]
        n_ph_seq    += [n_ph ]
        noteid_seq  += [noteid]

        # 初手長い休憩の時
        if ph_ids[0] == ph2id[g2p["R"]] and ms > 3000 and total_time == 0: 
            batch_process.append([ms_seq, ph_seq, n_ph_seq, noteid_seq])
            ph_seq = []
            n_ph_seq = []
            noteid_seq = []
            ms_seq =[]
            total_time = 0

        # 途中で来たR時
        if ph_ids[0] == ph2id[g2p["R"]] and total_time!=0:
            
            # 長い時は分割
            if ms > 3000:
                if len(ms_seq) > 1:
                    ms_last = float(ms_seq[-1]) - 1000

                    batch_process.append([ms_seq[:-1] + [1000], 
                                          ph_seq, 
                                          n_ph_seq, 
                                          noteid_seq])
                    batch_process.append([ms_last, ph_seq[-1], n_ph_seq[-1], noteid_seq[-1]]) 
                elif len(ms_seq) == 1:
                    batch_process.append([ms_seq, ph_seq, n_ph_seq, noteid_seq])
            # それ以外はそのまま
            else:
                batch_process.append([ms_seq, ph_seq, n_ph_seq, noteid_seq])

            ph_seq = []
            n_ph_seq = []
            noteid_seq = []
            ms_seq =[]
            total_time = 0

        total_time += ms
    if ms_seq != []:
        batch_process.append([ms_seq, ph_seq, n_ph_seq, noteid_seq])

    inference_loader = []
    for ms_seq, ph_seq, n_ph_seq, noteid_seq in batch_process:
        ph_seq_t = []
        for p in ph_seq:
            if type(p) == int:
                ph_seq_t.extend([p])
            else:
                ph_seq_t.extend(p)

        word_dur_ms     = torch.tensor(ms_seq,  dtype=torch.float32).view(-1)
        ph_seq          = torch.tensor(ph_seq_t,  dtype=torch.int64).view(-1)
        ph_length       = torch.tensor(len(ph_seq),  dtype=torch.int64).view(-1)
        n_ph_seq        = torch.tensor(n_ph_seq,  dtype=torch.int64).view(-1)
        noteid_seq      = torch.tensor(noteid_seq,  dtype=torch.float32).view(-1)

        e_ms=torch.cumsum(torch.tensor(word_dur_ms,dtype=torch.float32),dim=0).view(-1)
        if e_ms.size(0) == 1:
            word_frame_dur = e_ms
        else:
            word_frame_dur = get_dur_frame_from_e_ms(e_ms, sr=hps["sampling_rate"], hop_len=hps["hop_length"])
        word_lengths       = torch.tensor(word_frame_dur.size(0), dtype=torch.int64).view(-1)
        ph_idx_in_a_word   = get_linspace_ph_id(n_ph_seq)
        noteid_seq         = expand_note_info(ph_seq.unsqueeze(0), noteid_seq.unsqueeze(0), n_ph_seq.unsqueeze(0))
        noteid_length      = torch.tensor(len(noteid_seq), dtype=torch.int64)
        inference_loader.append([ph_seq.view(1,-1)+1 , 
                                 ph_length.view(-1), 
                                 word_frame_dur.view(1,-1), 
                                 word_lengths.view(-1), 
                                 word_dur_ms.view(1,-1) / 1000, 
                                 ph_idx_in_a_word.view(1,-1)+1, 
                                 n_ph_seq.view(1,-1),
                                 noteid_seq.view(1,-1)+1,
                                 noteid_length.view(-1)])
    return inference_loader

def get_continuos_f0(f0):
    """Convert F0 to continuous F0
    Args:
        f0 (ndarray): original f0 sequence with the shape (T)
    Return:
        (ndarray): continuous f0 with the shape (T)
    """
    # get uv information as binary
    uv = np.float32(f0 != 0)
    # get start and end of f0
    if (f0 == 0).all():
        logger.warn("all of the f0 values are 0.")
        return uv, f0, False
    start_f0 = f0[f0 != 0][0]
    end_f0 = f0[f0 != 0][-1]
    # padding start and end of f0 sequence
    cf0 = copy.deepcopy(f0)
    start_idx = np.where(cf0 == start_f0)[0][0]
    end_idx = np.where(cf0 == end_f0)[0][-1]
    cf0[:start_idx] = start_f0
    cf0[end_idx:] = end_f0
    # get non-zero frame index
    nz_frames = np.where(cf0 != 0)[0]
    # perform linear interpolation
    f = interp1d(nz_frames, cf0[nz_frames])
    cf0 = f(np.arange(0, cf0.shape[0]))
    return uv, cf0, True

def expand_note_info( ph_IDs, noteID, n_ph_pooling):
    ph_IDs_lengths = torch.tensor(ph_IDs.size(1), dtype=torch.int64)
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths.view(1), ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    noteID_lengths = torch.tensor(noteID.size(1), dtype=torch.int64)
    noteID_mask = torch.unsqueeze(commons.sequence_mask(noteID_lengths.view(1), noteID.size(1)), 1).to(noteID.dtype) # [B, 1, ph_len]
    attn_mask     = torch.unsqueeze(noteID_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn          = commons.generate_path(duration=torch.unsqueeze(n_ph_pooling,dim=1), mask=attn_mask)
    attn          = torch.squeeze(attn, dim=1).permute(0,2,1).float()                             # attn=[Batch, note_len,] 
    # expand
    noteID        = torch.matmul(noteID.float().unsqueeze(1), attn).float()                                         # to [Batch, inner_channel, ph_len] 
    return noteID.view(-1)

def get_linspace_ph_id( ph_idx_in_a_word):
    for idx, n_ph in enumerate(ph_idx_in_a_word):
        if idx == 0:
            data = torch.linspace(start=int(n_ph-1), end=0, steps=int(n_ph), dtype=torch.int64)
        else:
            data = torch.concat([data, torch.linspace(start=int(n_ph-1), end=0, steps=int(n_ph), dtype=torch.int64)], dim=0)
    return data

def get_dur_frame_from_e_ms(e_ms, sr, hop_len):
    e_ms = torch.tensor(e_ms,dtype=torch.float32)
    frames = torch.floor(  (e_ms / 1000)*sr / hop_len )   # 切り捨て
    frames = torch.diff(frames, dim=0, prepend=frames.new_zeros(1))
    return torch.tensor(frames, dtype=torch.int64)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--config',
                        type=str,
                        #required=True,
                        default="./configs/config.yaml" ,    
                        help='Path to configuration file')
    parser.add_argument('--F0diff_model_path',
                        type=str,
                        #required=True,
                        default="./inference_models/F0DIFF.pth",
                        help='Path to checkpoint')
    parser.add_argument('--Synthesizer_model_path',
                        type=str,
                        #required=True,
                        default="./inference_models/SVS.pth",
                        help='Path to checkpoint')
    parser.add_argument('--UST_path',
                        type=str,
                        #required=True,
                        default="./inference_models/ﾄﾞﾚﾐﾌｧｿﾗｼﾄﾞ.ust",
                        help='Path to checkpoint')
    parser.add_argument('--ask_retake',
                        type=bool,
                        default=True,
                        help='Whether to save output or not')
    parser.add_argument('--is_save',
                        type=str,
                        default=True,
                        help='Whether to save output or not')
    args = parser.parse_args()

    print(f"[INFO] AI Retake function is {args.ask_retake}")
    inference(args)

inference_svs.py
tonnetonne814/SiFi-VITS2-44100-Ja
import torch 
from torch.nn import functional as F

def feature_loss(fmap_r, fmap_g):
  loss = 0
  for dr, dg in zip(fmap_r, fmap_g):
    for rl, gl in zip(dr, dg):
      rl = rl.float().detach()
      gl = gl.float()
      loss += torch.mean(torch.abs(rl - gl))

  return loss * 2 

def discriminator_loss(disc_real_outputs, disc_generated_outputs):
  loss = 0
  r_losses = []
  g_losses = []
  for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
    dr = dr.float()
    dg = dg.float()
    r_loss = torch.mean((1-dr)**2)
    g_loss = torch.mean(dg**2)
    loss += (r_loss + g_loss)
    r_losses.append(r_loss.item())
    g_losses.append(g_loss.item())

  return loss, r_losses, g_losses


def generator_loss(disc_outputs):
  loss = 0
  gen_losses = []
  for dg in disc_outputs:
    dg = dg.float()
    l = torch.mean((1-dg)**2)
    gen_losses.append(l)
    loss += l

  return loss, gen_losses


def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
  """
  z_p, logs_q: [b, h, t_t]
  m_p, logs_p: [b, h, t_t]
  """
  z_p = z_p.float()
  logs_q = logs_q.float()
  m_p = m_p.float()
  logs_p = logs_p.float()
  z_mask = z_mask.float()

  kl = logs_p - logs_q - 0.5
  kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p)
  kl = torch.sum(kl * z_mask)
  l = kl / torch.sum(z_mask)
  return l



losses.py
tonnetonne814/SiFi-VITS2-44100-Ja
import torch.nn as nn
from diffusion.diffusion_utils import calc_diffusion_step_embedding
#from diffusion.diffusion_multinomial import MultinomialDiffusion
from diffusion.diffusion_model import F0_Diffusion
from diffusion.diffusers_modules import Downsample1D, Upsample1D, ResBlock1D, OutConv1DBlock
from attentions import Decoder as CrossAttn
from attentions import Encoder as SelfAttn
import commons, modules, attentions, math, torch, copy, logging
import math
from commons import generate_path
from tqdm import tqdm 
import torch
import torch.nn as nn
from collections import OrderedDict
#from vdecoder.hifigan.models import Generator # from so-vits-svc
from torch.nn.utils import weight_norm, spectral_norm
from torch.nn import Conv1d,   Conv2d
from commons import get_padding
from commons import convert_logdur_to_intdur
import copy
from sifigan.models.generator import SiFiGANGenerator
from singDB_loader import get_g2p_dict_from_tabledata
import monotonic_align


class VITS2_based_SiFiTTS(nn.Module):
  """
  Synthesizer for Training
  """

  def __init__(self, 
    hps,
    n_vocab,
    spec_channels,
    segment_size,
    inter_channels,
    hidden_channels,
    filter_channels,
    n_heads,
    n_layers,
    kernel_size,
    p_dropout,
    resblock, 
    resblock_kernel_sizes, 
    resblock_dilation_sizes, 
    upsample_rates, 
    upsample_initial_channel, 
    upsample_kernel_sizes,
    n_speakers,
    gin_channels,
    **kwargs):

    super().__init__()
    self.n_vocab = n_vocab
    self.spec_channels = spec_channels
    self.inter_channels = inter_channels
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.resblock = resblock
    self.resblock_kernel_sizes = resblock_kernel_sizes
    self.resblock_dilation_sizes = resblock_dilation_sizes
    self.upsample_rates = upsample_rates
    self.upsample_initial_channel = upsample_initial_channel
    self.upsample_kernel_sizes = upsample_kernel_sizes
    self.segment_size = segment_size
    self.n_speakers = hps["common"]["n_speaker"]
    self.gin_channels = hps["common"]["gin_channels"]
    self.transformer_flow_type = "fft" # "fft" / "mono_layer" / "pre_conv" ### When mono_layer and pre_conv, kl div loss went negative.
    self.current_mas_noise_scale = float(hps["VITS2_config"]["mas_noise_scale"])
    self.enc_gin_channels = gin_channels

    # VITS2 Text Encoder
    self.enc_p = TextEncoder_VITS2(n_vocab,hps["note_encoder"]["n_note"]+1,
        inter_channels,
        hidden_channels,
        filter_channels,
        n_heads,
        n_layers,
        kernel_size,
        p_dropout,
        gin_channels=self.enc_gin_channels)

    # SiFi Decoder
    self.dec = SiFiGANGenerator(**hps["SiFiGANGenerator"])
    self.upsample_scales= hps["SiFiGANGenerator"]["upsample_scales"]

    # VITS1 Encoder
    self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)

    # VITS2 Flow
    self.flow = ResidualCouplingTransformersBlock(
       inter_channels, 
       hidden_channels, 
       5, 
       1, 
       4, 
       gin_channels=gin_channels, 
       use_transformer_flows=True,
       transformer_flow_type=self.transformer_flow_type
       )

    # FastSpeech2 DurationPredictor
    self.dp = VariancePredictor(hps=hps, 
                                input_size      =hps["dur_predictor"]["input_size"],
                                filter_size     =hps["dur_predictor"]["filter_size"],
                                kernel          =hps["dur_predictor"]["kernel_size"],
                                conv_output_size=hps["dur_predictor"]["filter_size"],
                                dropout         =hps["dur_predictor"]["dropout"],
                                n_speaker=hps["common"]["n_speaker"])

    if n_speakers >= 1:
      self.emb_g = nn.Embedding(n_speakers, gin_channels)

    self.ms_per_frame = hps["sampling_rate"] / (hps["hop_length"] * 1000)
    self.hop_length = hps["hop_length"]
    self.oto2lab, self.ph_to_id,   self.id_to_ph, _,_ = get_g2p_dict_from_tabledata(table_path=hps["oto2lab_path"], 
                                                                                    include_converter=True)
    try:
      self.ph_statistics = torch.load(hps["ph_statistics_path"]) #{ph:[mean,var]} 
      print(f"[INFO] Loaded :", hps["ph_statistics_path"])
    except:
      self.ph_statistics = False

  def forward(self,
                spec, spec_lengths,
                ph_IDs, ph_IDs_lengths,
                dfs, 
                sinewave,
                speakerID):
    
    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [B, hidden, 1]
    else:
      g = None

    # posterior encoder
    z_spec, z_spec_m_q, z_spec_logs_q, spec_mask    = self.enc_q(spec, spec_lengths.float(), g=g)           # z_spec=[B, hidden, spec_len]

    # Flow 
    z_spec_text = self.flow(z_spec, spec_mask, g=g) # z_spec_text=[B, hidden, spec_len]

    # prior encoder 
    H_ph,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                              w_dur_ms=None,
                                                              ph_w_idx=None,
                                                              g=g)       # H_ph=[B, hidden, ph_len]

    with torch.no_grad():
      # negative cross-entropy
      s_p_sq_r = torch.exp(-2 * H_ph_logs_p) # [b, d, t]
      neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - H_ph_logs_p, [1], keepdim=True) # [b, 1, t_s]
      neg_cent2 = torch.matmul(-0.5 * (z_spec_text ** 2).transpose(1, 2), s_p_sq_r) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
      neg_cent3 = torch.matmul(z_spec_text.transpose(1, 2), (H_ph_m_p * s_p_sq_r)) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
      neg_cent4 = torch.sum(-0.5 * (H_ph_m_p ** 2) * s_p_sq_r, [1], keepdim=True) # [b, 1, t_s]
      neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4

      epsilon = torch.std(neg_cent) * torch.randn_like(neg_cent) * self.current_mas_noise_scale
      neg_cent = neg_cent + epsilon

      attn_mask = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(spec_mask, -1)
      attn = monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()  # pathが出来れば誤差伝播カット

    ph_dur = attn.sum(2)
    # duration predict
    logw_ = torch.log(ph_dur + 1e-6) * H_ph_mask
    logw = self.dp(H_ph, H_ph_mask, g=g)  # Note Normalization
    l2_dur_loss = torch.sum((logw - logw_)**2) / torch.sum(H_ph_mask)      # phoneme dur loss
    dp_H_ph = H_ph # for dur discriminator H_ph=[B, hidden, ph_len]
    dp_H_ph_mask = H_ph_mask

    # path between ph and spec
    spec_mask   = torch.unsqueeze(commons.sequence_mask(spec_lengths, spec.size(2)), 1).to(spec.dtype) # [B, 1, spec_len]
    attn_mask   = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(spec_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_gt     = generate_path(duration=torch.unsqueeze(ph_dur,dim=1), mask=attn_mask)
    attn_gt     = torch.squeeze(attn_gt, dim=1).permute(0,2,1).float()                             # attn=[Batch, note_len,] 

    # expand prior(from ph_len to spec_len) [B, hidden, spec_len]
    H_ph_mask    = torch.matmul(H_ph_mask, attn_gt  )
    H_ph         = torch.matmul(H_ph, attn_gt       ) * spec_mask 
    H_ph_m_p     = torch.matmul(H_ph_m_p , attn_gt  ) * spec_mask 
    H_ph_logs_p  = torch.matmul(H_ph_logs_p, attn_gt) * spec_mask 

    # slice process 
    z_slice, sinewave_slice, ids_slice = commons.rand_slice_segments_with_sinewave(x=z_spec, 
                                                                                pitch=torch.squeeze(sinewave, dim=1),  # ここをSineWaveへ
                                                                                x_lengths=z_spec.size(2), 
                                                                                hop_size= self.hop_length, 
                                                                                segment_size=self.segment_size)  # frame level
    dfs_slice = commons.dfs_slice_segment(dfs=dfs,
                                          ids_str=copy.deepcopy(ids_slice),
                                          upscales=self.upsample_scales,
                                          segment_size=self.segment_size)  # frame level

    # SiFi Decoder
    voice, excitation = self.dec(x=sinewave_slice.cuda(), 
                                 c=z_slice, 
                                 d=tuple([d.to("cuda:0") for d in dfs_slice]), 
                                 g=g)

    return voice, excitation, l2_dur_loss, attn_gt, ids_slice, dp_H_ph_mask, spec_mask, \
          (z_spec, z_spec_text, H_ph_m_p, H_ph_logs_p, z_spec_m_q, z_spec_logs_q), \
          (dp_H_ph, logw, logw_)

  # batch 1 only
  def eval_infer(self, ph_IDs, ph_IDs_lengths,
                  speakerID,
                  dfs, 
                  sinewave,
                  f0_lengths,
                  noise_scale=1,
                  length_scale=1, 
                  noise_scale_w=1.):
    
    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [b, h, 1]
    else:
      g = None

    # text 
    H_ph,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                                w_dur_ms=None,
                                                                ph_w_idx=None,
                                                                g=g)       # H_ph=[B, hidden, ph_len]

    logw = self.dp(H_ph, H_ph_mask, g=g)  # Note Normalization
    w = torch.exp(logw) * H_ph_mask * length_scale
    w = (w/torch.sum(w)) * f0_lengths  # length regurator
    w_ceil = torch.ceil(w)
    w_sum = torch.sum(w_ceil,dim=2).view(-1)
    if w_sum != f0_lengths:
       w_ceil[:,:,-1] -= w_sum - f0_lengths
    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
    y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(H_ph_mask.dtype)
    attn_mask = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(y_mask, -1)
    attn = commons.generate_path(w_ceil, attn_mask).squeeze(1).permute(0,2,1)

    # expand prior(from ph_len to spec_len) [B, hidden, spec_len]
    H_ph_mask    = torch.matmul(H_ph_mask  , attn)
    H_ph_m_p     = torch.matmul(H_ph_m_p   , attn) * H_ph_mask 
    H_ph_logs_p  = torch.matmul(H_ph_logs_p, attn) * H_ph_mask 

    z_spec_text = H_ph_m_p + torch.randn_like(H_ph_m_p) * torch.exp(H_ph_logs_p) * noise_scale
    z_spec = self.flow(z_spec_text, y_mask, g=g, reverse=True)

    # SiFi Decoder
    voice, _ = self.dec(x=sinewave.cuda(), 
                                 c=z_spec, 
                                 d=tuple([d.to("cuda:0") for d in dfs]), 
                                 g=g)

    return voice, attn, y_mask
  
  def get_mas_output(self,
                spec, spec_lengths,
                ph_IDs, ph_IDs_lengths,
                speakerID):
    
    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [B, hidden, 1]
    else:
      g = None

    # posterior encoder
    z_spec, _, _, spec_mask    = self.enc_q(spec, spec_lengths.float(), g=g)           # z_spec=[B, hidden, spec_len]

    # Flow 
    z_spec_text = self.flow(z_spec, spec_mask, g=g) # z_spec_text=[B, hidden, spec_len]

    # prior encoder 
    _,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                              w_dur_ms=None,
                                                              ph_w_idx=None,
                                                              g=g)       # H_ph=[B, hidden, ph_len]

    with torch.no_grad():
      # negative cross-entropy
      s_p_sq_r = torch.exp(-2 * H_ph_logs_p) # [b, d, t]
      neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - H_ph_logs_p, [1], keepdim=True) # [b, 1, t_s]
      neg_cent2 = torch.matmul(-0.5 * (z_spec_text ** 2).transpose(1, 2), s_p_sq_r) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
      neg_cent3 = torch.matmul(z_spec_text.transpose(1, 2), (H_ph_m_p * s_p_sq_r)) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
      neg_cent4 = torch.sum(-0.5 * (H_ph_m_p ** 2) * s_p_sq_r, [1], keepdim=True) # [b, 1, t_s]
      neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4

      epsilon = torch.std(neg_cent) * torch.randn_like(neg_cent) * self.current_mas_noise_scale
      neg_cent = neg_cent + epsilon

      attn_mask = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(spec_mask, -1)
      attn = monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()  # pathが出来れば誤差伝播カット
    
    ph_dur = attn.sum(2)

    return ph_dur, attn

  def voice_conversion(self, y, y_lengths, sid_src, sid_tgt):
    assert self.n_speakers > 0, "n_speakers have to be larger than 0."
    g_src = self.emb_g(sid_src).unsqueeze(-1)
    g_tgt = self.emb_g(sid_tgt).unsqueeze(-1)
    z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src)
    z_p = self.flow(z, y_mask, g=g_src)
    z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True)
    o_hat = self.dec(z_hat * y_mask, g=g_tgt)
    return o_hat, y_mask, (z, z_p, z_hat)

   # inference only and batch size=1 only 
  def adjust_duration(self, ph_ids, ph_dur_pd, word_dur, n_ph_in_word, noise_scale=0.33):
    ph_ids      = ph_ids[0, :]
    ph_dur_pd   = ph_dur_pd[0][0]
    word_dur    = word_dur[0]
    n_ph_in_word=n_ph_in_word[0]
    sum_duration = torch.sum(word_dur)
    total_diff = 0

    out_ph_dur  = torch.zeros_like(ph_dur_pd) # for duration output
    ph_idx = len(ph_dur_pd) 
    for idx in reversed(range(len(word_dur))):
        n_ph = n_ph_in_word[idx]
        # first
        if idx == len(word_dur)-1: 
            out_ph_dur[-1] = word_dur[idx]
            ph_idx -= 1
        # other
        else:
            target_ph_dur = ph_dur_pd[int(ph_idx-(1+z_n_ph)):int(ph_idx)] # Vowels this time + previous consonants
            adjusted_pd_dur = torch.ceil( (target_ph_dur/torch.sum(target_ph_dur) ) * word_dur[idx]) 
            if torch.sum(adjusted_pd_dur) != word_dur[idx]:
                undur = torch.sum(adjusted_pd_dur) - word_dur[idx]
                adjusted_pd_dur[0] = adjusted_pd_dur[0] - undur
            out_ph_dur[int(ph_idx-(1+z_n_ph)):int(ph_idx)] = adjusted_pd_dur
            ph_idx -= 1+z_n_ph
        z_n_ph = n_ph - 1 
    # last consonants
    if z_n_ph != 0: 
      out_ph_dur[0:z_n_ph] = ph_dur_pd[0:z_n_ph] # Vowels this time + previous consonants
    #assert torch.sum(out_ph_dur) == sum_duration  # check duration

    # adjust ph duration by statistics
    if self.ph_statistics is not False:
      diff = 0
      for idx in reversed(range(len(ph_ids))):
        ph = self.id_to_ph[int(ph_ids[idx]-1)] # minus 1 is for mask
        # fix ph_dur and calc diff
        if ph in self.ph_statistics:
          dur = out_ph_dur[idx]
          mean_ms = torch.tensor(self.ph_statistics[ph][0]).float().cuda()
          std_ms  = torch.tensor(self.ph_statistics[ph][1]).float().cuda()
          corr_ms = mean_ms + std_ms*torch.randn_like(mean_ms)*noise_scale
          corr_dur = torch.ceil(self.ms_per_frame*corr_ms)
          diff += dur - corr_dur
          out_ph_dur[idx] = corr_dur
        # shifting diff
        else:  
          out_ph_dur[idx] += diff
          diff=0

      # shifting diff last (Perhaps not necessary)
      #if diff != 0:
      #   out_ph_dur[0] += diff

      if z_n_ph != 0: 
          out_ph_dur[0:z_n_ph] = ph_dur_pd[0:z_n_ph] # Vowels this time + previous consonants
          
          #out_ph_dur[-1] -= diff # 長さ補正はconcatでやる
      #assert torch.sum(out_ph_dur) == sum_duration  # check duration
    
    # 最後に子音があれば、それだけずらす。
    if int(z_n_ph) == 0:
      total_diff = 0
    else:
      total_diff = int(torch.sum(out_ph_dur[0:int(z_n_ph)]))

    return out_ph_dur.view(1,1,-1), total_diff
    
  def encode_and_dp(self, 
                    ph_IDs, ph_IDs_lengths,
                    speakerID,
                    word_frame_dur, word_frame_dur_lenngths,
                    word_dur_ms, ph_word_flag,
                    n_ph_pool,  
                    noise_scale=1,
                    length_scale=1,
                    get_mas_output=False):

    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [b, h, 1]
    else:
      g = None

    # path between ph and word
    word_mask   = torch.unsqueeze(commons.sequence_mask(word_frame_dur_lenngths, word_frame_dur.size(1)), 1).to(word_frame_dur.dtype) # [B, 1, spec_len]
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    attn_mask   = torch.unsqueeze(word_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_ph_word= generate_path(duration=torch.unsqueeze(n_ph_pool,dim=1), mask=attn_mask)
    attn_ph_word= torch.squeeze(attn_ph_word, dim=1).float()                          # attn=[Batch, note_len,] 

    word_dur_ms= torch.matmul(attn_ph_word, word_dur_ms.unsqueeze(2).float()).squeeze(2)

    # text 
    H_ph,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                                w_dur_ms=word_dur_ms,
                                                                ph_w_idx=ph_word_flag,
                                                                g=g)       # H_ph=[B, hidden, ph_len]
    logw = self.dp(H_ph, H_ph_mask, g=g)  
    w = torch.exp(logw) * H_ph_mask * length_scale
    w_ceil = torch.ceil(w)

    w_ceil, diff = self.adjust_duration(ph_IDs, w_ceil, word_frame_dur, n_ph_pool)

    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
    y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(H_ph_mask.dtype)
    attn_mask = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(y_mask, -1)
    attn = commons.generate_path(w_ceil, attn_mask).squeeze(1).permute(0,2,1)

    # expand prior(from ph_len to spec_len) [B, hidden, spec_len]
    H_ph_mask    = torch.matmul(H_ph_mask  , attn)
    H_ph_m_p     = torch.matmul(H_ph_m_p   , attn) * H_ph_mask 
    H_ph_logs_p  = torch.matmul(H_ph_logs_p, attn) * H_ph_mask 

    z_spec_text = H_ph_m_p + torch.randn_like(H_ph_m_p) * torch.exp(H_ph_logs_p) * noise_scale
    z_spec = self.flow(z_spec_text, y_mask, g=g, reverse=True)

    return z_spec, attn, w_ceil, g, int(diff)
  
  def synthesize(self,sinewave, z_spec, dfs, g):
    # SiFi Decoder
    voice, _ = self.dec(x=sinewave.cuda(), 
                                 c=z_spec, 
                                 d=tuple([d.to("cuda:0") for d in dfs]), 
                                 g=g)
    return voice

class VITS2_based_SiFiSinger(nn.Module):
  """
  Synthesizer for Training
  """

  def __init__(self, 
    hps,
    n_vocab,
    spec_channels,
    segment_size,
    inter_channels,
    hidden_channels,
    filter_channels,
    n_heads,
    n_layers,
    kernel_size,
    p_dropout,
    resblock, 
    resblock_kernel_sizes, 
    resblock_dilation_sizes, 
    upsample_rates, 
    upsample_initial_channel, 
    upsample_kernel_sizes,
    n_speakers,
    gin_channels,
    use_sdp=False,
    **kwargs):

    super().__init__()
    self.n_vocab = n_vocab
    self.spec_channels = spec_channels
    self.inter_channels = inter_channels
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.resblock = resblock
    self.resblock_kernel_sizes = resblock_kernel_sizes
    self.resblock_dilation_sizes = resblock_dilation_sizes
    self.upsample_rates = upsample_rates
    self.upsample_initial_channel = upsample_initial_channel
    self.upsample_kernel_sizes = upsample_kernel_sizes
    self.segment_size = segment_size
    self.n_speakers = hps["common"]["n_speaker"]
    self.gin_channels = hps["common"]["gin_channels"]
    self.transformer_flow_type = "fft" # "fft" / "mono_layer" / "pre_conv" ### When mono_layer and pre_conv, kl div loss went negative.
    self.current_mas_noise_scale = float(hps["VITS2_config"]["mas_noise_scale"])
    self.enc_gin_channels = gin_channels
  
    # VITS2 Text Encoder (Speaker Embedding)
    self.enc_p = TextEncoder_VITS2(n_vocab,hps["note_encoder"]["n_note"]+1,
        inter_channels,
        hidden_channels,
        filter_channels,
        n_heads,
        n_layers,
        kernel_size,
        p_dropout,
        gin_channels=self.enc_gin_channels)

    # SiFi Decoder
    self.dec = SiFiGANGenerator(**hps["SiFiGANGenerator"])
    self.upsample_scales= hps["SiFiGANGenerator"]["upsample_scales"]

    # VITS1 Encoder
    self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)

    # VITS2 Flow
    self.flow = ResidualCouplingTransformersBlock(
       inter_channels, 
       hidden_channels, 
       5, 
       1, 
       4, 
       gin_channels=gin_channels, 
       use_transformer_flows=True, # Trueにすると、kl_divergense がマイナスになる。
       transformer_flow_type=self.transformer_flow_type
       )
    
    # FastSpeech2 DurationPredictor
    self.dp = VariancePredictor(hps=hps, 
                                input_size      =hps["dur_predictor"]["input_size"],
                                filter_size     =hps["dur_predictor"]["filter_size"],
                                kernel          =hps["dur_predictor"]["kernel_size"],
                                conv_output_size=hps["dur_predictor"]["filter_size"],
                                dropout         =hps["dur_predictor"]["dropout"],
                                n_speaker=hps["common"]["n_speaker"])

    if n_speakers >= 1:
      self.emb_g = nn.Embedding(n_speakers, gin_channels)

    self.ms_per_frame = hps["sampling_rate"] / (hps["hop_length"] * 1000)
    self.hop_length = hps["hop_length"]
    self.oto2lab, self.ph_to_id,   self.id_to_ph, _,_ = get_g2p_dict_from_tabledata(table_path=hps["oto2lab_path"], 
                                                                                    include_converter=True)
    try:
      self.ph_statistics = torch.load(hps["ph_statistics_path"]) #{ph:[mean,var]} 
      print(f"[INFO] Loaded :", hps["ph_statistics_path"])
    except:
      self.ph_statistics = False

  def forward(self,
                spec, spec_lengths,
                ph_IDs, ph_IDs_lengths,
                ph_dur, 
                word_frame_dur, word_frame_dur_lenngths,
                word_dur_ms, ph_word_flag,
                n_ph_pool,  
                dfs, 
                sinewave,
                speakerID):
    
    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [B, hidden, 1]
    else:
      g = None
      
    # generate path between ph and word
    word_mask   = torch.unsqueeze(commons.sequence_mask(word_frame_dur_lenngths, word_frame_dur.size(1)), 1).to(word_frame_dur.dtype) # [B, 1, spec_len]
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    attn_mask   = torch.unsqueeze(word_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_ph_word= generate_path(duration=torch.unsqueeze(n_ph_pool,dim=1), mask=attn_mask)
    attn_ph_word= torch.squeeze(attn_ph_word, dim=1).float()                             # attn=[Batch, note_len,] 

    # expand 
    word_dur_ms= torch.matmul(attn_ph_word, word_dur_ms.unsqueeze(2).float()).squeeze(2)
    
    # posterior encoder
    z_spec, z_spec_m_q, z_spec_logs_q, spec_mask    = self.enc_q(spec, spec_lengths.float(), g=g)           # z_spec=[B, hidden, spec_len]

    # Flow 
    z_spec_text = self.flow(z_spec, spec_mask, g=g) # z_spec_text=[B, hidden, spec_len]

    # prior encoder 
    H_ph,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                              w_dur_ms=word_dur_ms,
                                                              ph_w_idx=ph_word_flag,
                                                              g=g)       # H_ph=[B, hidden, ph_len]
    
    # duration predict
    logw_ = torch.log(torch.unsqueeze(ph_dur,dim=1) + 1e-6) * ph_IDs_mask
    logw = self.dp(H_ph, ph_IDs_mask, g=g)  # Note Normalization
    l2_ph_dur_loss = torch.sum((logw - logw_)**2) / torch.sum(ph_IDs_mask)      # phoneme dur loss
    l2_word_dur_loss = torch.sum((torch.matmul(logw, attn_ph_word) - torch.matmul(logw_, attn_ph_word))**2) / torch.sum(ph_IDs_mask)     # word dur loss
    l2_full_dur_loss = (torch.sum(logw_) - torch.sum(logw) )**2 / torch.sum(ph_IDs_mask) # lengths dur loss
    l2_dur_loss = [l2_ph_dur_loss, l2_word_dur_loss, l2_full_dur_loss]

    dp_H_ph = H_ph # for dur discriminator H_ph=[B, hidden, ph_len]

    # path between ph and spec
    spec_mask   = torch.unsqueeze(commons.sequence_mask(spec_lengths, spec.size(2)), 1).to(spec.dtype) # [B, 1, spec_len]
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    attn_mask   = torch.unsqueeze(ph_IDs_mask, 2) * torch.unsqueeze(spec_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_gt     = generate_path(duration=torch.unsqueeze(ph_dur,dim=1), mask=attn_mask)
    attn_gt     = torch.squeeze(attn_gt, dim=1).permute(0,2,1).float()                             # attn=[Batch, note_len,] 

    # expand prior(from ph_len to spec_len) [B, hidden, spec_len]
    H_ph_mask    = torch.matmul(H_ph_mask, attn_gt  )
    H_ph         = torch.matmul(H_ph, attn_gt       ) * spec_mask 
    H_ph_m_p     = torch.matmul(H_ph_m_p , attn_gt  ) * spec_mask 
    H_ph_logs_p  = torch.matmul(H_ph_logs_p, attn_gt) * spec_mask 

    # slice process 
    z_slice, sinewave_slice, ids_slice = commons.rand_slice_segments_with_sinewave(x=z_spec, 
                                                                                pitch=torch.squeeze(sinewave, dim=1),  # ここをSineWaveへ
                                                                                x_lengths=z_spec.size(2), 
                                                                                hop_size= self.hop_length, 
                                                                                segment_size=self.segment_size)  # frame level
    dfs_slice = commons.dfs_slice_segment(dfs=dfs,
                                          ids_str=copy.deepcopy(ids_slice),
                                          upscales=self.upsample_scales,
                                          segment_size=self.segment_size)  # frame level
    
    # SiFi Decoder
    voice, excitation = self.dec(x=sinewave_slice.cuda(), 
                                 c=z_slice, 
                                 d=tuple([d.to("cuda:0") for d in dfs_slice]), 
                                 g=g)
    
    return voice, excitation, l2_dur_loss, attn_gt, ids_slice, ph_IDs_mask, spec_mask, \
          (z_spec, z_spec_text, H_ph_m_p, H_ph_logs_p, z_spec_m_q, z_spec_logs_q), \
          (dp_H_ph, logw, logw_)

  # batch 1 only
  def eval_infer(self, ph_IDs, ph_IDs_lengths,
                  speakerID,
                  word_frame_dur, word_frame_dur_lenngths,
                  word_dur_ms, ph_word_flag,
                  n_ph_pool,  
                  dfs, 
                  sinewave,
                  noise_scale=1,
                  length_scale=1, 
                  noise_scale_w=1.):
    
    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [b, h, 1]
    else:
      g = None

    # path between ph and word
    word_mask   = torch.unsqueeze(commons.sequence_mask(word_frame_dur_lenngths, word_frame_dur.size(1)), 1).to(word_frame_dur.dtype) # [B, 1, spec_len]
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    attn_mask   = torch.unsqueeze(word_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_ph_word= generate_path(duration=torch.unsqueeze(n_ph_pool,dim=1), mask=attn_mask)
    attn_ph_word= torch.squeeze(attn_ph_word, dim=1).float()                             # attn=[Batch, note_len,] 

    word_dur_ms= torch.matmul(attn_ph_word, word_dur_ms.unsqueeze(2).float()).squeeze(2)

    # text 
    H_ph,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                                w_dur_ms=word_dur_ms,
                                                                ph_w_idx=ph_word_flag,
                                                                g=g)       # H_ph=[B, hidden, ph_len]

    logw = self.dp(H_ph, H_ph_mask, g=g)  # Note Normalization
    w = torch.exp(logw) * H_ph_mask * length_scale
    w_ceil = torch.ceil(w)

    w_ceil, diff = self.adjust_duration(ph_IDs, w_ceil, word_frame_dur, n_ph_pool)
    w_ceil[0,0,-1] -= diff # eval時だとph分余分なframeが作られるので、末尾でつじつま合わせ。

    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
    y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(H_ph_mask.dtype)
    attn_mask = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(y_mask, -1)
    attn = commons.generate_path(w_ceil, attn_mask).squeeze(1).permute(0,2,1)

    # expand prior(from ph_len to spec_len) [B, hidden, spec_len]
    H_ph_mask    = torch.matmul(H_ph_mask  , attn)
    H_ph_m_p     = torch.matmul(H_ph_m_p   , attn) * H_ph_mask 
    H_ph_logs_p  = torch.matmul(H_ph_logs_p, attn) * H_ph_mask 

    z_spec_text = H_ph_m_p + torch.randn_like(H_ph_m_p) * torch.exp(H_ph_logs_p) * noise_scale
    z_spec = self.flow(z_spec_text, y_mask, g=g, reverse=True)
      
    # SiFi Decoder
    voice, _ = self.dec(x=sinewave.cuda(), 
                                 c=z_spec, 
                                 d=tuple([d.to("cuda:0") for d in dfs]), 
                                 g=g)

    return voice, attn, y_mask
    
  def voice_conversion(self, y, y_lengths, sid_src, sid_tgt):
    assert self.n_speakers > 0, "n_speakers have to be larger than 0."
    g_src = self.emb_g(sid_src).unsqueeze(-1)
    g_tgt = self.emb_g(sid_tgt).unsqueeze(-1)
    z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src)
    z_p = self.flow(z, y_mask, g=g_src)
    z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True)
    o_hat = self.dec(z_hat * y_mask, g=g_tgt)
    return o_hat, y_mask, (z, z_p, z_hat)

  # inference only and batch size=1 only 
  def adjust_duration(self, ph_ids, ph_dur_pd, word_dur, n_ph_in_word, noise_scale=0.33):
    ph_ids      = ph_ids[0, :]
    ph_dur_pd   = ph_dur_pd[0][0]
    word_dur    = word_dur[0]
    n_ph_in_word=n_ph_in_word[0]
    sum_duration = torch.sum(word_dur)
    total_diff = 0

    out_ph_dur  = torch.zeros_like(ph_dur_pd) # for duration output
    ph_idx = len(ph_dur_pd) 
    for idx in reversed(range(len(word_dur))):
        n_ph = n_ph_in_word[idx]
        # first
        if idx == len(word_dur)-1: 
            out_ph_dur[-1] = word_dur[idx]
            ph_idx -= 1
        # other
        else:
            target_ph_dur = ph_dur_pd[int(ph_idx-(1+z_n_ph)):int(ph_idx)] # Vowels this time + previous consonants
            adjusted_pd_dur = torch.ceil( (target_ph_dur/torch.sum(target_ph_dur) ) * word_dur[idx]) 
            if torch.sum(adjusted_pd_dur) != word_dur[idx]:
                undur = torch.sum(adjusted_pd_dur) - word_dur[idx]
                adjusted_pd_dur[0] = adjusted_pd_dur[0] - undur
            out_ph_dur[int(ph_idx-(1+z_n_ph)):int(ph_idx)] = adjusted_pd_dur
            ph_idx -= 1+z_n_ph
        z_n_ph = n_ph - 1 
    # last consonants
    if z_n_ph != 0: 
      out_ph_dur[0:z_n_ph] = ph_dur_pd[0:z_n_ph] # Vowels this time + previous consonants
    #assert torch.sum(out_ph_dur) == sum_duration  # check duration

    # adjust ph duration by statistics
    if self.ph_statistics is not False:
      diff = 0
      for idx in reversed(range(len(ph_ids))):
        ph = self.id_to_ph[int(ph_ids[idx]-1)] # minus 1 is for mask
        # fix ph_dur and calc diff
        if ph in self.ph_statistics:
          dur = out_ph_dur[idx]
          mean_ms = torch.tensor(self.ph_statistics[ph][0]).float().cuda()
          std_ms  = torch.tensor(self.ph_statistics[ph][1]).float().cuda()
          corr_ms = mean_ms + std_ms*torch.randn_like(mean_ms)*noise_scale
          corr_dur = torch.ceil(self.ms_per_frame*corr_ms)
          diff += dur - corr_dur
          out_ph_dur[idx] = corr_dur
        # shifting diff
        else:  
          out_ph_dur[idx] += diff
          diff=0

      # shifting diff last (Perhaps not necessary)
      #if diff != 0:
      #   out_ph_dur[0] += diff

      if z_n_ph != 0: 
          out_ph_dur[0:z_n_ph] = ph_dur_pd[0:z_n_ph] # Vowels this time + previous consonants
          
          #out_ph_dur[-1] -= diff # 長さ補正はconcatでやる
      #assert torch.sum(out_ph_dur) == sum_duration  # check duration
    
    # 最後に子音があれば、それだけずらす。
    if int(z_n_ph) == 0:
      total_diff = 0
    else:
      total_diff = int(torch.sum(out_ph_dur[0:int(z_n_ph)]))

    return out_ph_dur.view(1,1,-1), total_diff
    
  def encode_and_dp(self, 
                    ph_IDs, ph_IDs_lengths,
                    speakerID,
                    word_frame_dur, word_frame_dur_lenngths,
                    word_dur_ms, ph_word_flag,
                    n_ph_pool,  
                    noise_scale=1,
                    length_scale=1):

    if self.n_speakers > 0:
      g = self.emb_g(speakerID).unsqueeze(-1) # [b, h, 1]
    else:
      g = None

    # path between ph and word
    word_mask   = torch.unsqueeze(commons.sequence_mask(word_frame_dur_lenngths, word_frame_dur.size(1)), 1).to(word_frame_dur.dtype) # [B, 1, spec_len]
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    attn_mask   = torch.unsqueeze(word_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_ph_word= generate_path(duration=torch.unsqueeze(n_ph_pool,dim=1), mask=attn_mask)
    attn_ph_word= torch.squeeze(attn_ph_word, dim=1).float()                          # attn=[Batch, note_len,] 

    word_dur_ms= torch.matmul(attn_ph_word, word_dur_ms.unsqueeze(2).float()).squeeze(2)

    # text 
    H_ph,   H_ph_m_p,   H_ph_logs_p,  H_ph_mask  = self.enc_p(ph_IDs, ph_IDs_lengths, 
                                                                w_dur_ms=word_dur_ms,
                                                                ph_w_idx=ph_word_flag,
                                                                g=g)       # H_ph=[B, hidden, ph_len]

    logw = self.dp(H_ph, H_ph_mask, g=g)  
    w = torch.exp(logw) * H_ph_mask * length_scale
    w_ceil = torch.ceil(w)

    w_ceil, diff = self.adjust_duration(ph_IDs, w_ceil, word_frame_dur, n_ph_pool)

    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
    y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(H_ph_mask.dtype)
    attn_mask = torch.unsqueeze(H_ph_mask, 2) * torch.unsqueeze(y_mask, -1)
    attn = commons.generate_path(w_ceil, attn_mask).squeeze(1).permute(0,2,1)

    # expand prior(from ph_len to spec_len) [B, hidden, spec_len]
    H_ph_mask    = torch.matmul(H_ph_mask  , attn)
    H_ph_m_p     = torch.matmul(H_ph_m_p   , attn) * H_ph_mask 
    H_ph_logs_p  = torch.matmul(H_ph_logs_p, attn) * H_ph_mask 

    z_spec_text = H_ph_m_p + torch.randn_like(H_ph_m_p) * torch.exp(H_ph_logs_p) * noise_scale
    z_spec = self.flow(z_spec_text, y_mask, g=g, reverse=True)

    return z_spec, attn, w_ceil, g, int(diff)
  
  def synthesize(self,sinewave, z_spec, dfs, g):
    # SiFi Decoder
    voice, _ = self.dec(x=sinewave.cuda(), 
                                 c=z_spec, 
                                 d=tuple([d.to("cuda:0") for d in dfs]), 
                                 g=g)
    return voice
  
class Rezero(torch.nn.Module):
    def __init__(self):
        super(Rezero, self).__init__()
        self.alpha = torch.nn.Parameter(torch.zeros(size=(1,)))

    def forward(self, x):
        return self.alpha * x

class StochasticDurationPredictor(nn.Module):
  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
    super().__init__()
    filter_channels = in_channels # it needs to be removed from future version.
    self.in_channels = in_channels
    self.filter_channels = filter_channels
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.n_flows = n_flows
    self.gin_channels = gin_channels

    self.log_flow = modules.Log()
    self.flows = nn.ModuleList()
    self.flows.append(modules.ElementwiseAffine(2))
    for i in range(n_flows):
      self.flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
      self.flows.append(modules.Flip())

    self.post_pre = nn.Conv1d(1, filter_channels, 1)
    self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
    self.post_convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
    self.post_flows = nn.ModuleList()
    self.post_flows.append(modules.ElementwiseAffine(2))
    for i in range(4):
      self.post_flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
      self.post_flows.append(modules.Flip())

    self.pre = nn.Conv1d(in_channels, filter_channels, 1)
    self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
    self.convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
    if gin_channels != 0:
      self.cond = nn.Conv1d(gin_channels, filter_channels, 1)

  def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
    x = torch.detach(x)
    x = self.pre(x)
    if g is not None:
      g = torch.detach(g)
      x = x + self.cond(g)
    x = self.convs(x, x_mask)
    x = self.proj(x) * x_mask

    if not reverse:
      flows = self.flows
      assert w is not None

      logdet_tot_q = 0 
      h_w = self.post_pre(w)
      h_w = self.post_convs(h_w, x_mask)
      h_w = self.post_proj(h_w) * x_mask
      e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask
      z_q = e_q
      for flow in self.post_flows:
        z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
        logdet_tot_q += logdet_q
      z_u, z1 = torch.split(z_q, [1, 1], 1) 
      u = torch.sigmoid(z_u) * x_mask
      z0 = (w - u) * x_mask
      logdet_tot_q += torch.sum((nn.functional.logsigmoid(z_u) + nn.functional.logsigmoid(-z_u)) * x_mask, [1,2])
      logq = torch.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - logdet_tot_q

      logdet_tot = 0
      z0, logdet = self.log_flow(z0, x_mask)
      logdet_tot += logdet
      z = torch.cat([z0, z1], 1)
      for flow in flows:
        z, logdet = flow(z, x_mask, g=x, reverse=reverse)
        logdet_tot = logdet_tot + logdet
      nll = torch.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - logdet_tot
      return nll + logq # [b]
    else:
      flows = list(reversed(self.flows))
      flows = flows[:-2] + [flows[-1]] # remove a useless vflow
      z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale
      for flow in flows:
        z = flow(z, x_mask, g=x, reverse=reverse)
      z0, z1 = torch.split(z, [1, 1], 1)
      logw = z0
      return logw

class DurationPredictor(nn.Module):
  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
    super().__init__()

    self.in_channels = in_channels
    self.filter_channels = filter_channels
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.gin_channels = gin_channels

    self.drop = nn.Dropout(p_dropout)
    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2)
    self.norm_1 = modules.LayerNorm(filter_channels)
    self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2)
    self.norm_2 = modules.LayerNorm(filter_channels)
    self.proj = nn.Conv1d(filter_channels, 1, 1)

    if gin_channels != 0:
      self.cond = nn.Conv1d(gin_channels, in_channels, 1)

  def forward(self, x, x_mask, g=None):
    x = torch.detach(x)
    if g is not None:
      g = torch.detach(g)
      x = x + self.cond(g)
    x = self.conv_1(x * x_mask)
    x = torch.relu(x)
    x = self.norm_1(x)
    x = self.drop(x)
    x = self.conv_2(x * x_mask)
    x = torch.relu(x)
    x = self.norm_2(x)
    x = self.drop(x)
    x = self.proj(x * x_mask)
    return x * x_mask

class DurationDiscriminator(nn.Module): #vits2
  # TODO : not using "spk conditioning" for now according to the paper.
  # Can be a better discriminator if we use it.
  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
    super().__init__()

    self.in_channels = in_channels
    self.filter_channels = filter_channels
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.gin_channels = gin_channels

    self.drop = nn.Dropout(p_dropout)
    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2)
    # self.norm_1 = modules.LayerNorm(filter_channels)
    self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2)
    # self.norm_2 = modules.LayerNorm(filter_channels)
    self.dur_proj = nn.Conv1d(1, filter_channels, 1)

    self.pre_out_conv_1 = nn.Conv1d(2*filter_channels, filter_channels, kernel_size, padding=kernel_size//2)
    self.pre_out_norm_1 = modules.LayerNorm(filter_channels)
    self.pre_out_conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2)
    self.pre_out_norm_2 = modules.LayerNorm(filter_channels)

    # if gin_channels != 0:
    #   self.cond = nn.Conv1d(gin_channels, in_channels, 1)

    self.output_layer = nn.Sequential(
        nn.Linear(filter_channels, 1), 
        nn.Sigmoid() 
    )

  def forward_probability(self, x, x_mask, dur, g=None):
    dur = self.dur_proj(dur)
    x = torch.cat([x, dur], dim=1)
    x = self.pre_out_conv_1(x * x_mask)
    # x = torch.relu(x)
    # x = self.pre_out_norm_1(x)
    # x = self.drop(x)
    x = self.pre_out_conv_2(x * x_mask)
    # x = torch.relu(x)
    # x = self.pre_out_norm_2(x)
    # x = self.drop(x)
    x = x * x_mask
    x = x.transpose(1, 2)
    output_prob = self.output_layer(x)
    return output_prob

  # x=hidden_input, dur=dur
  def forward(self, x, x_mask, dur_r, dur_hat, g=None):
    x = torch.detach(x)
    # if g is not None:
    #   g = torch.detach(g)
    #   x = x + self.cond(g)
    x = self.conv_1(x * x_mask)
    # x = torch.relu(x)
    # x = self.norm_1(x)
    # x = self.drop(x)
    x = self.conv_2(x * x_mask)
    # x = torch.relu(x)
    # x = self.norm_2(x)
    # x = self.drop(x)
    
    output_probs = []
    for dur in [dur_r, dur_hat]:
      output_prob = self.forward_probability(x, x_mask, dur, g)
      output_probs.append(output_prob)
    
    return output_probs
  
class ResidualCouplingTransformersLayer(nn.Module): #vits2
  def __init__(
      self,
      channels,
      hidden_channels,
      kernel_size,
      dilation_rate,
      n_layers,
      p_dropout=0,
      gin_channels=0,
      mean_only=False,
  ):
      assert channels % 2 == 0, "channels should be divisible by 2"
      super().__init__()
      self.channels = channels
      self.hidden_channels = hidden_channels
      self.kernel_size = kernel_size
      self.dilation_rate = dilation_rate
      self.n_layers = n_layers
      self.half_channels = channels // 2
      self.mean_only = mean_only
      #vits2
      self.pre_transformer = attentions.Encoder(
          self.half_channels,
          self.half_channels,
          n_heads=2,
          n_layers=2,
          kernel_size=3,
          p_dropout=0.1,
          window_size=None
          )

      self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
      self.enc = modules.WN_original(
          hidden_channels,
          kernel_size,
          dilation_rate,
          n_layers,
          p_dropout=p_dropout,
          gin_channels=gin_channels,
      )
      #vits2
      self.post_transformer = attentions.Encoder( 
          self.hidden_channels,
          self.hidden_channels,
          n_heads=2,
          n_layers=2,
          kernel_size=3,
          p_dropout=0.1,
          window_size=None
          )
      
      self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
      self.post.weight.data.zero_()
      self.post.bias.data.zero_()

  def forward(self, x, x_mask, g=None, reverse=False):
      x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
      x0_ = self.pre_transformer(x0 * x_mask, x_mask) #vits2
      x0_ = x0_ + x0 #vits2 residual connection
      h = self.pre(x0_) * x_mask #changed from x0 to x0_ to retain x0 for the flow
      h = self.enc(h, x_mask, g=g)

      #vits2 - (experimental;uncomment the following 2 line to use)
      # h_ = self.post_transformer(h, x_mask) 
      # h = h + h_ #vits2 residual connection 

      stats = self.post(h) * x_mask
      if not self.mean_only:
          m, logs = torch.split(stats, [self.half_channels] * 2, 1)
      else:
          m = stats
          logs = torch.zeros_like(m)
      if not reverse:
          x1 = m + x1 * torch.exp(logs) * x_mask
          x = torch.cat([x0, x1], 1)
          logdet = torch.sum(logs, [1, 2])
          return x, logdet
      else:
          x1 = (x1 - m) * torch.exp(-logs) * x_mask
          x = torch.cat([x0, x1], 1)
          return x

class FFTransformerCouplingLayer(nn.Module): #vits2
  def __init__(self,
      channels,
      hidden_channels,
      kernel_size,
      n_layers,
      n_heads,
      p_dropout=0,
      filter_channels=768,
      mean_only=False,
      gin_channels = 0
      ):
    assert channels % 2 == 0, "channels should be divisible by 2"
    super().__init__()
    self.channels = channels
    self.hidden_channels = hidden_channels
    self.kernel_size = kernel_size
    self.n_layers = n_layers
    self.half_channels = channels // 2
    self.mean_only = mean_only

    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
    self.enc = attentions.FFT(
       hidden_channels, 
       filter_channels, 
       n_heads, 
       n_layers, 
       kernel_size, 
       p_dropout, 
       isflow = True, 
       gin_channels = gin_channels
       )
    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
    self.post.weight.data.zero_()
    self.post.bias.data.zero_()

  def forward(self, x, x_mask, g=None, reverse=False):
    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
    h = self.pre(x0) * x_mask
    h_ = self.enc(h, x_mask, g=g)
    h = h_ + h
    stats = self.post(h) * x_mask
    if not self.mean_only:
      m, logs = torch.split(stats, [self.half_channels]*2, 1)
    else:
      m = stats
      logs = torch.zeros_like(m)

    if not reverse:
      x1 = m + x1 * torch.exp(logs) * x_mask
      x = torch.cat([x0, x1], 1)
      logdet = torch.sum(logs, [1,2])
      return x, logdet
    else:
      x1 = (x1 - m) * torch.exp(-logs) * x_mask
      x = torch.cat([x0, x1], 1)
      return x

class MonoTransformerFlowLayer(nn.Module): #vits2
  def __init__(
      self,
      channels,
      hidden_channels,
      mean_only=False,
  ):
    assert channels % 2 == 0, "channels should be divisible by 2"
    super().__init__()
    self.channels = channels
    self.hidden_channels = hidden_channels
    self.half_channels = channels // 2
    self.mean_only = mean_only
    #vits2
    self.pre_transformer = attentions.Encoder(
        self.half_channels,
        self.half_channels,
        n_heads=2,
        n_layers=2,
        kernel_size=3,
        p_dropout=0.1,
        window_size=None
        )
    
    self.post = nn.Conv1d(self.half_channels, self.half_channels * (2 - mean_only), 1)
    self.post.weight.data.zero_()
    self.post.bias.data.zero_()

  def forward(self, x, x_mask, g=None, reverse=False):
    x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
    x0_ = self.pre_transformer(x0 * x_mask, x_mask) #vits2
    h = x0_ + x0 #vits2
    stats = self.post(h) * x_mask
    if not self.mean_only:
        m, logs = torch.split(stats, [self.half_channels] * 2, 1)
    else:
        m = stats
        logs = torch.zeros_like(m)
    if not reverse:
        x1 = m + x1 * torch.exp(logs) * x_mask
        x = torch.cat([x0, x1], 1)
        logdet = torch.sum(logs, [1, 2])
        return x, logdet
    else:
        x1 = (x1 - m) * torch.exp(-logs) * x_mask
        x = torch.cat([x0, x1], 1)
        return x
        
class ResidualCouplingTransformersBlock(nn.Module): #vits2
  def __init__(self,
      channels,
      hidden_channels,
      kernel_size,
      dilation_rate,
      n_layers,
      n_flows=4,
      gin_channels=0,
      use_transformer_flows=False,
      transformer_flow_type="pre_conv",
      ):
    super().__init__()
    self.channels = channels
    self.hidden_channels = hidden_channels
    self.kernel_size = kernel_size
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.n_flows = n_flows
    self.gin_channels = gin_channels

    self.flows = nn.ModuleList()
    if use_transformer_flows:
      if transformer_flow_type == "pre_conv":
        for i in range(n_flows):
          self.flows.append(
             ResidualCouplingTransformersLayer(
             channels, 
             hidden_channels, 
             kernel_size, 
             dilation_rate, 
             n_layers, 
             gin_channels=gin_channels, 
             mean_only=False
             )
            )
          self.flows.append(modules.Flip())
      elif transformer_flow_type == "fft":
         for i in range(n_flows):
          self.flows.append(
             FFTransformerCouplingLayer(
             channels, 
             hidden_channels, 
             kernel_size, 
             dilation_rate, 
             n_layers, 
             gin_channels=gin_channels, 
             mean_only=True
             )
            )
          self.flows.append(modules.Flip())
      elif transformer_flow_type == "mono_layer":
        for i in range(n_flows):
          self.flows.append(
             modules.ResidualCouplingLayer(
             channels, 
             hidden_channels, 
             kernel_size, 
             dilation_rate, 
             n_layers, 
             gin_channels=gin_channels, 
             mean_only=True
             )
            )
          self.flows.append(
            MonoTransformerFlowLayer(
            channels, hidden_channels, mean_only=True
            )
          )
          self.flows.append(modules.Flip())
          self.flows.append(
           MonoTransformerFlowLayer(
           channels, hidden_channels, mean_only=True
           )
           )
    else:
      for i in range(n_flows):
        self.flows.append(
           modules.ResidualCouplingLayer(
           channels, 
           hidden_channels,
           kernel_size, 
           dilation_rate, 
           n_layers, 
           gin_channels=gin_channels, 
           mean_only=True
           )
          )
        self.flows.append(modules.Flip())

  def forward(self, x, x_mask, g=None, reverse=False):
    if not reverse:
      for flow in self.flows:
        x, _ = flow(x, x_mask, g=g, reverse=reverse)
    else:
      for flow in reversed(self.flows):
        x = flow(x, x_mask, g=g, reverse=reverse)
    return x
  
class ResidualCouplingBlock(nn.Module):
  def __init__(self,
      channels,
      hidden_channels,
      kernel_size,
      dilation_rate,
      n_layers,
      n_flows=4,
      gin_channels=0):
    super().__init__()
    self.channels = channels
    self.hidden_channels = hidden_channels
    self.kernel_size = kernel_size
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.n_flows = n_flows
    self.gin_channels = gin_channels

    self.flows = nn.ModuleList()
    for i in range(n_flows):
      self.flows.append(
         modules.ResidualCouplingLayer(
         channels, 
         hidden_channels, 
         kernel_size, 
         dilation_rate, 
         n_layers, 
         gin_channels=gin_channels, 
         mean_only=True
         )
         )
      self.flows.append(modules.Flip())

  def forward(self, x, x_mask, g=None, reverse=False):
    if not reverse:
      for flow in self.flows:
        x, _ = flow(x, x_mask, g=g, reverse=reverse)
    else:
      for flow in reversed(self.flows):
        x = flow(x, x_mask, g=g, reverse=reverse)
    return x

class unofficialRMSSinger(nn.Module):
    def __init__(self,hps,
                 vocab_size = 50,
                 ):
        super(unofficialRMSSinger, self).__init__()
        self.segment_size = hps["segments_size"] // hps["hop_length"]
        self.f0_max = int(hps["f0_max"])
        self.z_slice_segment = None
        self.eval_slice_gain =  int(hps["slice_gain"])

        self.ph_enc = TextEncoder(n_vocab=vocab_size,
                                 out_channels   =hps["ph_encoder"]["out_channels"],
                                 hidden_channels=hps["ph_encoder"]["hidden_channels"],
                                 filter_channels=hps["ph_encoder"]["filter_channels"],
                                 n_heads        =hps["ph_encoder"]["n_heads"],
                                 n_layers       =hps["ph_encoder"]["n_layers"],
                                 kernel_size    =hps["ph_encoder"]["kernel_size"],
                                 p_dropout      =hps["ph_encoder"]["p_dropout"])

        self.note_enc = NoteEncoder(n_note      =hps["note_encoder"]["n_note"]+1,
                                    hidden_channels=hps["note_encoder"]["hidden_channels"],
                                    hps         =hps)
        
        self.pooling = Pooling_from_ph_to_Word()

        self.spec_encoder = PosteriorEncoder(in_channels=hps["spec_encoder"]["spec_channels"], 
                                      out_channels=hps["spec_encoder"]["out_channels"],  
                                      hidden_channels=hps["spec_encoder"]["hidden_channels"],  
                                      kernel_size=5, 
                                      dilation_rate=1, 
                                      n_layers=16,
                                      gin_channels=hps["spec_encoder"]["gin_channels"])
        
        self.ph_note_encoder = PosteriorEncoder(in_channels=hps["ph_note_encoder"]["spec_channels"], 
                                      out_channels=hps["ph_note_encoder"]["out_channels"],  
                                      hidden_channels=hps["ph_note_encoder"]["hidden_channels"],  
                                      kernel_size=5, 
                                      dilation_rate=1, 
                                      n_layers=16,
                                      gin_channels=hps["ph_note_encoder"]["gin_channels"])

        self.flow = ResidualCouplingBlock(channels=hps["flow"]["inter_channels"], 
                                          hidden_channels=hps["flow"]["hidden_channels"], 
                                          kernel_size=5, 
                                          dilation_rate=1, 
                                          n_layers=4, 
                                          gin_channels=hps["flow"]["gin_channels"])

        self.duration_predictor = VariancePredictor(hps=hps, 
                                                    input_size      =hps["dur_predictor"]["input_size"],
                                                    filter_size     =hps["dur_predictor"]["filter_size"],
                                                    kernel          =hps["dur_predictor"]["kernel_size"],
                                                    conv_output_size=hps["dur_predictor"]["filter_size"],
                                                    dropout         =hps["dur_predictor"]["dropout"],
                                                    n_speaker=hps["common"]["n_speaker"])

        self.word_attn = WordLevelPositionalAttention(hps=hps,
                                                      word_vocab=hps["wordlevel_posattn"]["word_vocab"], # 後でちゃんとした値を入れる
                                                      emb_dim=hps["wordlevel_posattn"]["emb_dim"])
        
        self.upsampling = GaussianUpsampling(delta=hps["gaussian_upsampler"]["delta"])
        
        self.diffusion = DiffusionModels(hps=hps, vocab_size=vocab_size)

        # NSF Decoder from so-vits-svc
        self.decoder = Generator(h=hps["hifi_gan"], 
                                 sampling_rate=hps["sampling_rate"],)

        self.spk_emb = nn.Embedding(num_embeddings=hps["n_speaker"], embedding_dim=int(hps["NoisePredictor"]["inner_channels"]*5 + hps["hifi_gan"]["gin_channels"]))
        self.divide_len = int(hps["NoisePredictor"]["inner_channels"])

        self.proj = nn.Conv1d(in_channels=hps["ph_encoder"]["out_channels"]*2, 
                              out_channels=hps["ph_encoder"]["out_channels"]*2,
                              kernel_size=1)

    def forward(self,
                spec, spec_lengths,
                f0, f0_lengths, 
                ph_IDs, ph_IDs_lengths,
                noteID, noteID_lengths, 
                word_dur, word_dur_lengths,
                ph_idx_in_a_word, ph_idx_in_a_word_lengths,
                frame_word_id, frame_word_id_lengths,
                n_ph_pooling, n_ph_pooling_lengths,
                speakerID):
        
        emb_g = self.spk_emb (speakerID)
        
        # spec to z
        z, m_q, logs_q, y_mask = self.spec_encoder(spec, spec_lengths, 
                                                   g=emb_g[:, self.divide_len*2:self.divide_len*3].unsqueeze(2))
        # z to z_text
        z_p = self.flow(z, y_mask, g=emb_g[:, self.divide_len*3:self.divide_len*4].unsqueeze(2))
        
        # generate f0 mask 
        f0_mask = torch.unsqueeze(commons.sequence_mask(f0_lengths, f0.size(2)), 1).to(f0.dtype)
        word_mask = torch.unsqueeze(commons.sequence_mask(word_dur_lengths, word_dur.size(1)), 1).to(f0.dtype)

        ####################### ph_dur estimat
        # ph encode
        H_ph, H_ph_mask = self.ph_enc(ph_IDs, ph_IDs_lengths)

        # average pooling from ph to word & permute  # ph→poolingじゃなくて、
        H_word = self.pooling(ph=H_ph, 
                              ph_mask=H_ph_mask, 
                              n_ph_pooling=n_ph_pooling, 
                              word_mask=word_mask)

        # note encode ここは本物のword_durで生成する
        H_note, H_note_mask = self.note_enc(noteID, noteID_lengths, word_dur) # note_embedding Only now

        # predict word duration 
        word_dur_pd = self.duration_predictor(x      = H_note + H_word + torch.unsqueeze(emb_g[:, 0:self.divide_len],dim=1), 
                                              x_mask = torch.squeeze(H_note_mask, dim=1))
        #######################

        # word level pos attention
        H_epd, H_epd_mask = self.word_attn( hid_ph=H_ph,
                                            hid_ph_mask=H_ph_mask, 
                                            n_ph_pos=ph_idx_in_a_word,   
                                            n_ph_pos_lengths=ph_idx_in_a_word_lengths, 
                                            frame_word_id=frame_word_id, 
                                            frame_word_id_lengths=frame_word_id_lengths)

        # word level gaussian upsampler 本物のword_durを使って予測
        H_note_up = self.upsampling(hs=H_note, ds=word_dur, 
                                    h_masks=torch.squeeze(f0_mask>0,dim=1), 
                                    d_masks=torch.squeeze(H_note_mask>0,dim=1)).permute(0,2,1)

        # F0 training (possible independent training)
        loss_f0 = self.diffusion(f0         =f0,  
                                 f0_len     =f0_lengths,
                                 hid_ph     =H_epd + H_note_up, 
                                 hid_ph_mask=H_epd_mask,
                                 g          =emb_g[:, self.divide_len:self.divide_len*2])

        hidden = torch.concat([H_epd, H_note_up],dim=1)
        lengths = torch.sum(H_epd_mask, dim=2).squeeze(1)
        z_hid, m_p, logs_p, hid_mask = self.ph_note_encoder(hidden, lengths , g=emb_g[:, self.divide_len*4:self.divide_len*5].unsqueeze(2))
        #m_p, logs_p = torch.split(hidden, 192, dim=1)

        # decoder
        z_slice, pitch_slice, ids_slice = commons.rand_slice_segments_with_pitch(x=z, 
                                                                                 pitch=torch.squeeze(f0, dim=1), 
                                                                                 x_lengths=f0.size(2), 
                                                                                 segment_size=self.segment_size)

        # nsf decoder
        y_gen = self.decoder(z_slice, g=emb_g[:,self.divide_len*5:], f0=pitch_slice*1100)
        
        return y_gen, ids_slice, word_dur_pd, loss_f0, (z, z_p, m_p, logs_p, m_q, logs_q, y_mask)
    
    def inference(self, ph_IDs, ph_IDs_lengths,
                        noteID, noteID_lengths, 
                        note_dur, note_dur_lengths,
                        ph_idx_in_a_word, ph_idx_in_a_word_lengths,
                        frame_word_id, frame_word_id_lengths,
                        n_ph_pooling, n_ph_pooling_lengths,
                        speakerID,
                        f0=None, f0_lengths=None, slice_idx=None):
        assert ph_IDs.size(0) == 1
        
        emb_g = self.spk_emb (speakerID)

        # generate f0 mask 
        note_dur_mask = torch.unsqueeze(commons.sequence_mask(note_dur_lengths, note_dur.size(1)), 1).to(f0.dtype)

        # ph encode
        H_ph, H_ph_mask = self.ph_enc(ph_IDs, ph_IDs_lengths)

        # average pooling from ph to word(note) & permute
        H_word = self.pooling(ph=H_ph, 
                              ph_mask=H_ph_mask, 
                              n_ph_pooling=n_ph_pooling, 
                              word_mask=note_dur_mask)

        # note encode
        H_note, H_note_mask = self.note_enc(noteID, noteID_lengths, note_dur)

        # predict word duration
        word_dur_pd = self.duration_predictor(x      = H_note + H_word + torch.unsqueeze(emb_g[:, 0:self.divide_len],dim=1), 
                                              x_mask = torch.squeeze(H_note_mask, dim=1))
        word_dur = convert_logdur_to_intdur(word_dur_pd)
        
        # generate f0 mask 
        f0_mask = torch.ones(size=(1,int(torch.sum(word_dur)))) # 全てTrue

        # word level pos attention
        H_epd, H_epd_mask = self.word_attn( hid_ph=H_ph,
                                            hid_ph_mask=H_ph_mask, 
                                            n_ph_pos=ph_idx_in_a_word,   
                                            n_ph_pos_lengths=ph_idx_in_a_word_lengths, 
                                            frame_word_id=frame_word_id, 
                                            frame_word_id_lengths=frame_word_id_lengths)
        
        # word level gaussian upsampler
        H_note_up = self.upsampling(hs=H_note, ds=word_dur, 
                                 h_masks=torch.squeeze(f0_mask>0,dim=1), 
                                 d_masks=torch.squeeze(H_note_mask>0,dim=1)).permute(0,2,1)
        
        # F0の指定が有れば使用する
        if f0 is None or f0_lengths is None:
            # NoteIDs, NoteID_len, word_dur, speakerID = condition
            condition = [H_epd+H_note_up,  0,  0,  emb_g[:, self.divide_len:self.divide_len*2]]
            f0 = self.diffusion.sampling(condition).view(1,1,-1) * self.f0_max
            f0_mask = torch.unsqueeze(commons.sequence_mask(f0_lengths, f0.size(2)), 1).to(f0.dtype)
        else:
            f0_mask = torch.unsqueeze(commons.sequence_mask(f0_lengths, f0.size(2)), 1).to(f0.dtype)

        # nsf decoder
        z = torch.concat([H_epd, H_note_up],dim=1)
        y_gen = self.decoder(z, g=emb_g[:,self.divide_len*4:], f0=torch.unsqueeze(f0, dim=1)) 
        
        return y_gen
    
    def evaluation(self, ph_IDs, ph_IDs_lengths,
                   noteID, noteID_lengths, 
                   word_dur, word_dur_lengths, # note_dur,との連携は後で
                   ph_idx_in_a_word, ph_idx_in_a_word_lengths,
                   frame_word_id, frame_word_id_lengths,
                   n_ph_pooling, n_ph_pooling_lengths,
                   speakerID,
                   f0=None, f0_lengths=None):
        
        emb_g = self.spk_emb (speakerID)
        
        # generate mask 
        #note_dur_sum = torch.sum(note_dur)
        #note_mask = commons.sequence_mask(note_dur_sum.view(1,-1), note_dur_sum).to(f0.dtype).cuda()
        #note_mask = commons.sequence_mask(word_dur_lengths, word_dur.size(1)).to(f0.dtype)

        # ph encode
        H_ph, H_ph_mask = self.ph_enc(ph_IDs, ph_IDs_lengths)

        # average pooling from ph to word & permute
        word_mask = torch.unsqueeze(commons.sequence_mask(word_dur_lengths, word_dur.size(1)), 1).to(f0.dtype)
        H_word = self.pooling(ph=H_ph, 
                              ph_mask=H_ph_mask, 
                              n_ph_pooling=n_ph_pooling, 
                              word_mask=word_mask)          # [B, word_len, hidden]

        # note encode
        H_note, H_note_mask = self.note_enc(noteID, noteID_lengths, word_dur)   # [B, word_len, hidden]

        # predict word duration
        word_dur_pd = self.duration_predictor(x      = H_note.detach() + H_word.detach() + torch.unsqueeze(emb_g[:, 0:self.divide_len],dim=1), 
                                              x_mask = torch.squeeze(H_note_mask, dim=1))
        
        # word level pos attention
        H_epd, H_epd_mask = self.word_attn( hid_ph=H_ph,
                                            hid_ph_mask=H_ph_mask, 
                                            n_ph_pos=ph_idx_in_a_word,   
                                            n_ph_pos_lengths=ph_idx_in_a_word_lengths, 
                                            frame_word_id=frame_word_id,                    # ここをword_dur_pdで埋める必要アリ
                                            frame_word_id_lengths=frame_word_id_lengths)
        
        # word level gaussian upsampler
        H_note_up = self.upsampling(hs=H_note, ds=word_dur,  ### 分岐01、ここは予測値を使うこともできる
                                 h_masks=torch.squeeze(H_epd_mask>0,dim=1), 
                                 d_masks=torch.squeeze(H_note_mask>0,dim=1)).permute(0,2,1)
        
        # F0 sampling NoteIDs, NoteID_len, NoteDur, speakerID = condition
        condition = [H_epd+H_note_up,  noteID_lengths,  word_dur,  emb_g[:, self.divide_len:self.divide_len*2]]
        #condition = [H_epd+H_note_up,  noteID_lengths,  convert_logdur_to_intdur(word_dur_pd) ,  emb_g[:, self.divide_len:self.divide_len*2]]
        f0_pd = self.diffusion.sampling(condition).view(1,1,-1) * self.f0_max
        #f0_pd_mask = torch.unsqueeze(commons.sequence_mask(f0_lengths, f0.size(2)), 1).to(f0.dtype)
        #f0_mask = torch.unsqueeze(commons.sequence_mask(f0_lengths, f0.size(2)), 1).to(f0.dtype)

        # decoder
        
        hidden = torch.concat([H_epd, H_note_up],dim=1)
        lengths = torch.sum(H_epd_mask, dim=2).squeeze(1)
        z_hid, m_p, logs_p, hid_mask = self.ph_note_encoder(hidden, lengths , g=emb_g[:, self.divide_len*4:self.divide_len*5].unsqueeze(2))
        #m_p, logs_p = torch.split(hidden, 192, dim=1)

        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * 1 # noisescale
        z = self.flow(z_p, H_epd_mask, g=emb_g[:, self.divide_len*3:self.divide_len*4].unsqueeze(2), reverse=True)
        
        ### 分岐02
        # for eval with f0gt
        z_slice, pitch_slice, ids_slice = commons.rand_slice_segments_with_pitch(x=z, 
                                                                                 pitch=torch.squeeze(f0, dim=1), 
                                                                                 x_lengths=f0.size(2), 
                                                                                 segment_size=self.segment_size*self.eval_slice_gain)
        y_gen_f0gt = self.decoder(z_slice, g=emb_g[:,self.divide_len*5:], f0=pitch_slice*1100) # nsf decoder

        # for eval with f0pd
        z_slice, pitch_slice, ids_slice = commons.slice_segments_with_pitch(x=z, 
                                                                            pitch=torch.squeeze(f0_pd, dim=1), 
                                                                            ids_str=ids_slice, 
                                                                            segment_size=self.segment_size*self.eval_slice_gain)
        y_gen_f0pd = self.decoder(z_slice, g=emb_g[:,self.divide_len*5:], f0=pitch_slice*1100) # nsf decoder

        return y_gen_f0gt, y_gen_f0pd, f0_pd, ids_slice
    
# Copyright 2022 Dan Lim
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
class GaussianUpsampling(torch.nn.Module):
    """Gaussian upsampling with fixed temperature as in:

    https://arxiv.org/abs/2010.04301

    """

    def __init__(self, delta=0.1):
        super().__init__()
        self.delta = delta # deltaは1/(2*sigma^2)の値だと思われる

    def forward(self, hs, ds, h_masks=None, d_masks=None):
        """Upsample hidden states according to durations.

        Args:
            hs (Tensor): Batched hidden state to be expanded (B, T_text, adim).
            ds (Tensor): Batched token duration (B, T_text).
            h_masks (Tensor): Mask tensor (B, T_feats).
            d_masks (Tensor): Mask tensor (B, T_text).

        Returns:
            Tensor: Expanded hidden state (B, T_feat, adim).

        """
        B = ds.size(0)
        device = ds.device

        if ds.sum() == 0:
            logging.warning(
                "predicted durations includes all 0 sequences. "
                "fill the first element with 1."
            )
            # NOTE(kan-bayashi): This case must not be happened in teacher forcing.
            #   It will be happened in inference with a bad duration predictor.
            #   So we do not need to care the padded sequence case here.
            ds[ds.sum(dim=1).eq(0)] = 1

        if h_masks is None:
            T_feats = ds.sum().int()
        else:
            T_feats = h_masks.size(-1)
        t = torch.arange(0, T_feats).unsqueeze(0).repeat(B, 1).to(device).float()   # [B, frame]
        if h_masks is not None:
            t = t * h_masks.float()

        c = ds.cumsum(dim=-1) - ds / 2 # ds.cumsumはeのこと
        energy = -1 * self.delta * (t.unsqueeze(-1) - c.unsqueeze(1)) ** 2 # energy = w_t^n
        if d_masks is not None:
            #energy = energy.masked_fill(
            #    ~(d_masks.unsqueeze(1).repeat(1, T_feats, 1)), -float("inf")
            #)
            d_masks = d_masks.unsqueeze(1).repeat(1, T_feats, 1)
            energy = energy.masked_fill(d_masks==0, -float("inf"))

        p_attn = torch.softmax(energy, dim=2)  # (B, T_feats, T_text)、これが重み
        hs = torch.matmul(p_attn, hs) # 重み適応 hs=[B, T_feats, hidden]
        return hs

class Pooling_from_ph_to_Word(nn.Module):
    def __init__(self,):
        super(Pooling_from_ph_to_Word, self).__init__()
    def forward(self,ph,            # [B, hidden,   ph_len]
                     ph_mask,       # [B,      1,   ph_len]
                     n_ph_pooling,  # [B, word_len]
                     word_mask):    # [B,      1,   word_len]
        attn_mask = torch.unsqueeze(word_mask, 2) * torch.unsqueeze(ph_mask, -1)            # attn_mask= [B, 1, ph_len, word_len]
        attn = generate_path(duration=torch.unsqueeze(n_ph_pooling,dim=1), mask=attn_mask ) # attn= [B, 1, ph_len, word_len]
        attn = torch.squeeze(attn, dim=1)
        attn = torch.div(attn, torch.unsqueeze(n_ph_pooling, dim=1)) # NaN occurs
        attn = torch.nan_to_num(attn) # replace NaN with 0.
        H_word = torch.matmul(ph, attn.detach()) # detach and pooling process 
        return H_word.permute(0,2,1)

class WordLevelPositionalAttention(nn.Module):
    def __init__(self,
                 hps,
                 word_vocab,
                 emb_dim
                 ):
        super(WordLevelPositionalAttention, self).__init__()
        self.n_ph_emb = nn.Embedding(num_embeddings=5, embedding_dim=int(emb_dim//4)) 
        self.word_id_emb = nn.Embedding(num_embeddings=word_vocab, embedding_dim=emb_dim) 

        self.linear_proj = nn.Conv1d(in_channels=emb_dim + int(emb_dim//4),
                                     out_channels=emb_dim, kernel_size=1)
        self.emb_dim = emb_dim
        
    def forward(self, 
                hid_ph,                 # [B, hidden, ph_len]
                hid_ph_mask,            # [B,      1, ph_len]
                n_ph_pos,               # [B, ph_len]
                n_ph_pos_lengths,       # [B,] 
                frame_word_id,          # [B, frame]
                frame_word_id_lengths): # [B, ]
        
        ph_len = hid_ph.size(2)
        frame_len = frame_word_id.size(1)

        n_ph_mask = torch.unsqueeze(commons.sequence_mask(n_ph_pos_lengths, n_ph_pos.size(1)), 1).to(n_ph_pos.dtype) # [B, 1, ph_len]
        emb_n_ph_pos = self.n_ph_emb(n_ph_pos).permute(0,2,1) # emb_n_ph_pos=[B, temp_hidden, ph_len]
        emb_n_ph_pos = emb_n_ph_pos * n_ph_mask

        frame_word_id_mask = torch.unsqueeze(commons.sequence_mask(frame_word_id_lengths, frame_word_id.size(1)), 1).to(frame_word_id.dtype) #[B, 1, frame]
        emb_frame_word_id = self.word_id_emb(frame_word_id).permute(0,2,1) # [B, hidden, frame]
        emb_frame_word_id = emb_frame_word_id * frame_word_id_mask
        
        H_k = torch.cat([hid_ph, emb_n_ph_pos], dim=1)
        H_k = self.linear_proj(H_k)
        H_k  = H_k.permute(0,2,1) # [B, ph_len, hidden]

        attn_score = torch.matmul( H_k, emb_frame_word_id) / math.sqrt(self.emb_dim) # attn_score = [B, ph_len, frame]
        attn_score = attn_score.masked_fill(n_ph_mask.squeeze(1).unsqueeze(2).repeat(1, 1, frame_len)==0, -float("inf"))
        attn_score = attn_score.masked_fill(frame_word_id_mask.repeat(1, ph_len, 1)==0, -float("1e20")) # NaN対策
        attn_score = torch.softmax(attn_score, dim=1) 

        H_epd = torch.matmul(hid_ph, attn_score)
        return H_epd * frame_word_id_mask, frame_word_id_mask        

class NoteEncoder(nn.Module):
    def __init__(self,
                 n_note,
                 hidden_channels,
                 hps,
                 ):
        super(NoteEncoder, self).__init__()
        self.hidden_ch = hidden_channels
        #self.note_emb = nn.Embedding(num_embeddings=n_note, 
        #                             embedding_dim=hidden_channels)
        self.dur_emb  = nn.Linear   (in_features=1,  # max_len
                                     out_features=hidden_channels)
        self.note_emb  = nn.Linear   (in_features=1,  # max_len
                                     out_features=hidden_channels)
        self.hidden_channels=hidden_channels
        
        ### wordIDをselfAttnして加算してもよいかも
    def forward(self, noteID, noteID_lengths, note_dur=None):
        noteID_mask = torch.unsqueeze(commons.sequence_mask(noteID_lengths, noteID.size(1)), 1).to(noteID.dtype)
        
        emb = self.note_emb(noteID.unsqueeze(2)) / math.sqrt(self.hidden_ch) 
        if note_dur is not None:
          emb += self.dur_emb(note_dur.unsqueeze(2)) / math.sqrt(self.hidden_ch) 
        
        return emb.permute(0,2,1), noteID_mask


####################
# from FastSpeech2 #
####################
class VariancePredictor(nn.Module):
    """Duration, Pitch and Energy Predictor"""

    def __init__(self, hps,
                 input_size,
                 filter_size,
                 kernel,
                 conv_output_size,
                 dropout,
                 n_speaker):
        super(VariancePredictor, self).__init__()

        self.input_size         = input_size
        self.filter_size        = filter_size
        self.kernel             = kernel
        self.conv_output_size   = conv_output_size
        self.dropout            = dropout
        

        self.emb_g = nn.Conv1d(input_size, self.filter_size, kernel_size=1)

        self.conv_layer = nn.Sequential(
            OrderedDict(
                [
                    (
                        "conv1d_1",
                        Conv(
                            self.input_size,
                            self.filter_size,
                            kernel_size=self.kernel,
                            padding=(self.kernel - 1) // 2,
                        ),
                    ),
                    #("relu_2", nn.ReLU()),
                    ("relu_2", nn.Mish()),
                    ("layer_norm_1", nn.LayerNorm(self.filter_size)),
                    ("dropout_1", nn.Dropout(self.dropout)),
                    (
                        "conv1d_2",
                        Conv(
                            self.filter_size,
                            self.filter_size,
                            kernel_size=self.kernel,
                            padding=1,
                        ),
                    ),
                    #("relu_2", nn.ReLU()),
                    ("relu_2", nn.Mish()),
                    ("layer_norm_2", nn.LayerNorm(self.filter_size)),
                    ("dropout_2", nn.Dropout(self.dropout)),
                ]
            )
        )

        self.linear_layer = nn.Linear(self.conv_output_size, 1)

    def forward(self, x, x_mask, g=None):
        if g is not None:
            x = x + self.emb_g(g)
        x = self.conv_layer(x.permute(0,2,1))
        x = self.linear_layer(x)
        x = x.permute(0,2,1)

        if x_mask is not None:
            x = x.masked_fill(x_mask==0, 0.0)

        return x


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,
        )

    def forward(self, x):
        x = x.contiguous().transpose(1, 2)
        x = self.conv(x)
        x = x.contiguous().transpose(1, 2)

        return x

class DiffusionModels(nn.Module):
    def __init__(self,
                 hps,
                 vocab_size,
                 training=True):
        super(DiffusionModels, self).__init__()

        self.num_timesteps = hps["Diffusion"]["T"]
        if training is True:
          self.infer_timesteps = hps["Diffusion"]["T"]
          self.ddim = False
        else:
          self.infer_timesteps = hps["Diffusion"]["N"]
          self.ddim = hps["Diffusion"]["ddim"]

        self.f0_max = float(hps["f0_max"])
        self.lf0_max = 2595. * torch.log10(1. + torch.tensor(hps["f0_max"]).float() / 700.) / 500

        # F0 diffusion modules
        self.f0_diff = F0_Diffusion(hps=hps)

        # Voice/UnVoice diffusion modules
        #self.vuv_diff = MultinomialDiffusion(num_classes=3,   # vとuvとmaskの3通り
        #                                     timesteps=self.num_timesteps,
        #                                     loss_type="vb_stochastic",
        #                                     parametrization="x0")
        self.noise_schedule = None
        self.N = None 
        self.step_infer = None

        # Noise Predictor
        self.noise_predictor = NoisePredictor(hps=hps, 
                                              out_channels            = hps["NoisePredictor"]['out_channels'],
                                              inner_channels          = hps["NoisePredictor"]['inner_channels'],
                                              WN_in_channels          = hps["NoisePredictor"]['WN_in_channels'],
                                              WN_kernel_size          = hps["NoisePredictor"]['WN_kernel_size'],
                                              WN_dilation_rate        = hps["NoisePredictor"]['WN_dilation_rate'],
                                              WN_n_layers             = hps["NoisePredictor"]['WN_n_layers'],
                                              WN_p_dropout            = hps["NoisePredictor"]['WN_p_dropout'],
                                              Attn_filter_channels    = hps["NoisePredictor"]['Attn_filter_channels'],
                                              Attn_kernel_size        = hps["NoisePredictor"]['Attn_kernel_size'],
                                              Attn_n_layers           = hps["NoisePredictor"]['Attn_n_layers'],
                                              Attn_n_heads            = hps["NoisePredictor"]['Attn_n_heads'],
                                              Attn_p_dropout          = hps["NoisePredictor"]['Attn_p_dropout'],
                                              Attn_vocab_size         = vocab_size+1, # mask用に0を取っておく。
                                              n_speakers              = hps["NoisePredictor"]['n_speakers'],
                                              Diff_step_embed_in      = hps["NoisePredictor"]['Diff_step_embed_in'],
                                              Diff_step_embed_mid     = hps["NoisePredictor"]['Diff_step_embed_mid'],
                                              Diff_step_embed_out     = hps["NoisePredictor"]['Diff_step_embed_out'],
                                              gin_channels            = hps["NoisePredictor"]['n_speakers'])

    # 入力を受け取り、ロスを返す。
    def forward(self, f0,  f0_len,
                      IDs, IDs_len, IDs_dur,
                      NoteIDs = None, NoteID_len = None,
                      g=None):
        
        f0 = self.norm_f0(f0)
        if NoteIDs is not None:
          NoteIDs = self.norm_f0(NoteIDs)

        # Generate timesteps 
        B, _, _ = f0.shape  
        timesteps= torch.randint(self.num_timesteps, size=(B, 1, 1)).cuda()
        #timesteps, pt = self.vuv_diff.sample_time(B, device="cuda", method='importance')

        # F0 preprocess (F0 → NoisyF0)
        noisy_f0, noise_gt = self.f0_diff.get_noisy_f0(f0=f0, ts=timesteps)

        # vuv preprocess (VUV → NoisyVUV)
        # noisy_vuv, vuv_diff_params = self.vuv_diff.preprocess(x=vuv, t_int=timesteps.view(-1))

        # Noise prediction 
        f0_noise = self.noise_predictor(f0=noisy_f0,  
                                        f0_len=f0_len,    # noisyF0
                                        ph_IDs=IDs, 
                                        ph_IDs_len=IDs_len,   # Condition (ph)
                                        ph_ID_dur=IDs_dur, 
                                        NoteIDs=NoteIDs, 
                                        NoteIDS_len=NoteID_len,
                                        timesteps=timesteps.view(B, 1),
                                        g=g)

        # Loss calculation for F0
        loss_f0 = self.f0_diff.get_loss( f0_noise_pd=f0_noise, 
                                         noisy_f0=noisy_f0, 
                                         noise_gt=noise_gt)

        # Loss calculation for VUV
        #loss_vuv = self.vuv_diff.postprocess(model_out=vuv_noise, 
        #                                     parameters=vuv_diff_params,
        #                                     t_int=timesteps.view(-1), 
        #                                     t_float=pt)
        #loss_vuv = loss_vuv.sum() / (math.log(2) * torch.sum(vuv_len))

        return loss_f0  #, loss_vuv 
    
    @torch.inference_mode()
    def sampling(self, condition):
        #ph_IDs, ph_IDs_len, ph_IDs_dur, NoteIDs, NoteID_len, NoteDur,speakerID = condition
        f0_len = int(torch.sum(condition[2][0]))

        if self.noise_schedule is None:
          self.noise_schedule = self.f0_diff.get_noise_schedule(timesteps=self.infer_timesteps, training=False) 
        # sampling using DDPM reverse process. 
        #f0_pd, vuv_pd = self.sampling_given_noise_schedule(size=(1, 1, f0_len),  
        #                                                   inference_noise_schedule=noise_schedule,
        #                                                   condition=condition, 
        #                                                   ddim=False, 
        #                                                   return_sequence=False)
        f0_pd = self.sampling_given_noise_schedule(size=(1, 1, f0_len),  
                                                           inference_noise_schedule=self.noise_schedule,
                                                           condition=condition, 
                                                           ddim=self.ddim, 
                                                           return_sequence=False)

        return self.reverse_norm_f0(f0_pd)
        #return f0_pd, vuv_pd
    
    @torch.inference_mode()
    def sampling_given_noise_schedule(
        self,
        size,
        inference_noise_schedule,
        condition=None,
        ddim=False,
        return_sequence=False):
        """
        Perform the complete sampling step according to p(x_0|x_T) = \prod_{t=1}^T p_{\theta}(x_{t-1}|x_t)

        Parameters:
        net (torch network):            the wavenet models
        size (tuple):                   size of tensor to be generated,
                                        usually is (number of audios to generate, channels=1, length of audio)
        diffusion_hyperparams (dict):   dictionary of diffusion hyperparameters returned by calc_diffusion_hyperparams
                                        note, the tensors need to be cuda tensors
        condition (torch.tensor):       ground truth mel spectrogram read from disk
                                        None if used for unconditional generation

        Returns:
        the generated audio(s) in torch.tensor, shape=size
        """

        #IDs_gt, IDs_gt_len, dur_gt, speaker_ID = condition
        ph_IDs, ph_IDs_len, ph_IDs_dur, NoteIDs, NoteID_len, _, speakerID = condition
        if NoteIDs is not None:
          NoteIDs = self.norm_f0(NoteIDs)
        N , steps_infer = self.f0_diff.get_step_infer(size=size,
                                                     inference_noise_schedule=inference_noise_schedule)
          
        x = self.f0_diff.generate_x(size).cuda()
        #log_z = self.vuv_diff.generate_log_z(size).cuda()

        if return_sequence:
            x_ = copy.deepcopy(x)
            xs = [x_]

        # need  N と steps_infer
        with torch.no_grad():
            for n in tqdm(range(N - 1, -1, -1),desc="Sampling..."):
                diffusion_steps = (steps_infer[n] * torch.ones((size[0], 1))).cuda()

                # Multinomial diffusion preprocess
                #model_in, log_z = self.vuv_diff.sample_preprocess(log_z=log_z, t=n)

                # Noise Prediction
                #f0_noise, vuv_noise = self.noise_predictor(f0=x,  
                #                                           f0_len=torch.tensor([dur_gt.size(1)], dtype=torch.int64, device="cuda:0"),
                #                                           IDs=IDs_gt, 
                #                                           IDs_len=IDs_gt_len, 
                #                                           vuv=model_in,                  
                #                                           vuv_len=torch.tensor([dur_gt.size(1)], dtype=torch.int64, device="cuda:0"),               
                #                                           attn=dur_gt, 
                #                                           timesteps=diffusion_steps,
                #                                           g=speaker_ID)
                
                f0_noise = self.noise_predictor(f0=x,  
                                                f0_len=torch.tensor([size[2]], dtype=torch.int64, device="cuda:0"),    # noisyF0
                                                ph_IDs=ph_IDs, 
                                                ph_IDs_len=ph_IDs_len,   # Condition (ph)
                                                ph_ID_dur=ph_IDs_dur, 
                                                timesteps=diffusion_steps,
                                                NoteIDs=NoteIDs, 
                                                NoteIDS_len=NoteID_len,
                                                g=speakerID)

                # Multinomial diffusion postprocess
                #log_z = self.vuv_diff.sample_postprocess(model_out=vuv_noise, log_z=log_z, t_int=n)

                # Denoising
                x = self.f0_diff.denoising(x=x, noise_pd=f0_noise, ddim=ddim, n=n)

                if return_sequence:
                    x_ = copy.deepcopy(x)
                    xs.append(x_)

            # VUV Decoding
            #vuv_pd = self.vuv_diff.decode_log_z(log_z)

        if return_sequence:
            return xs
        return x
        #return x, vuv_pd
    def norm_f0(self,f0):
        #lf0 = 2595. * torch.log10(1. + f0 / 700.) / 500 / self.lf0_max
        f0 = f0 / self.f0_max
        return f0 
    def reverse_norm_f0(self,f0):
        #lf0[lf0<0] = 0
        #f0 = (700 * (torch.pow(10, lf0 *self.lf0_max * 500 / 2595) - 1))
        f0[f0<0] = 0
        f0 = f0 * self.f0_max
        return f0  

class NoisePredictor(nn.Module):
    def __init__(self,hps,
        out_channels,
        inner_channels,
        WN_in_channels,
        WN_kernel_size,
        WN_dilation_rate,
        WN_n_layers,
        WN_p_dropout,
        Attn_filter_channels,
        Attn_kernel_size,
        Attn_n_layers,
        Attn_n_heads,
        Attn_p_dropout,
        Attn_vocab_size,
        n_speakers,
        Diff_step_embed_in,
        Diff_step_embed_mid,
        Diff_step_embed_out,
        gin_channels=0):
        super().__init__()

        self.out_channels=out_channels
        self.diffusion_step_embed_dim_in = Diff_step_embed_in

        # Timestep Embedding
        self.fc_t = nn.ModuleList()
        self.fc_t1 = nn.Linear(Diff_step_embed_in, Diff_step_embed_mid)
        self.fc_t2 = nn.Linear(Diff_step_embed_mid, Diff_step_embed_out)

        # Speaker / VUV Embedding
        #self.emb_uv = nn.Embedding(3, inner_channels) # vとuvとmaskの３つ
        self.emb_g  = nn.Embedding(n_speakers, gin_channels)

        # NoteEncoder
        self.note_enc = NoteEncoder(n_note=hps["note_encoder"]["n_note"]+1, 
                                    hidden_channels=hps["note_encoder"]["hidden_channels"],
                                    hps=hps)

        # Attention
        self.text_encoder = TextEncoder(n_vocab         =Attn_vocab_size,
                                        out_channels    =inner_channels,
                                        hidden_channels =inner_channels,
                                        filter_channels =Attn_filter_channels,
                                        n_heads         =Attn_n_heads,
                                        n_layers        =Attn_n_layers,
                                        kernel_size     =Attn_kernel_size,
                                        p_dropout       =Attn_p_dropout   )

        # WaveNet
        self.pre = nn.Conv1d(WN_in_channels, inner_channels, 1)
        self.WaveNet = modules.WN(hidden_channels=inner_channels,
                                  Diff_step_embed_out=Diff_step_embed_out,
                                  kernel_size=WN_kernel_size,
                                  dilation_rate=WN_dilation_rate, 
                                  n_layers=WN_n_layers, 
                                  gin_channels=gin_channels,
                                  p_dropout=WN_p_dropout)

        #self.rezero = Rezero() 

        # projection
        output_ch = 1 # 4 = [f0 vuv]
        self.proj_1 = nn.Conv1d(inner_channels, output_ch, 1) 
        self.relu = nn.Mish()
        self.proj_2 = nn.Conv1d(output_ch, output_ch, 1) 
        #self.vuv_proj = nn.Conv1d(1,4,1)

    def forward(self,   f0,  
                        f0_len,    # noisyF0
                        ph_IDs, 
                        ph_IDs_len,   # Condition (ph)
                        ph_ID_dur, 
                        #vuv, vuv_len, 
                        timesteps,
                        NoteIDs=None, 
                        NoteIDS_len=None, 
                        g=None):

        # Embedding timesteps
        emb_timesteps = calc_diffusion_step_embedding(timesteps, self.diffusion_step_embed_dim_in)
        emb_timesteps = swish(self.fc_t1(emb_timesteps))
        emb_timesteps = swish(self.fc_t2(emb_timesteps))

        # Embedding speakerID 
        g = torch.unsqueeze(self.emb_g(g),dim=2) # [Batch, n_speakers] to [Batch, 1, gin_channels]

        # Projection 
        f0_mask = torch.unsqueeze(commons.sequence_mask(f0_len, f0.size(2)), 1).to(f0.dtype)
        f0 = self.pre(f0) * f0_mask                 # [Batch, Channel=1, f0_len] to [Batch, inner_channels, f0_len]

        # Embedding vuv labels
        #vuv_mask = torch.unsqueeze(commons.sequence_mask(vuv_len, vuv.size(1)), 1).to(vuv.dtype)
        #vuv = self.emb_uv(vuv).permute(0,2,1) * vuv_mask    # [Batch, f0_len]     to [Batch, inner_channel, f0_len]


        # Encoding (output is masked) Attention Embedding
        ph_IDs, ph_IDs_mask = self.text_encoder(ph_IDs, ph_IDs_len)                 # [Batch, inner_channel, IDs_len]   

        if NoteIDs is not None:
          noteIDs, noteIDs_mask = self.note_enc(  noteID=NoteIDs, 
                                                noteID_lengths=NoteIDS_len)

        # expand ph_len to f0_len 
        attn_mask = torch.unsqueeze(ph_IDs_mask, 2) * torch.unsqueeze(f0_mask, -1)
        attn = generate_path(duration=torch.unsqueeze(ph_ID_dur,dim=1), mask=attn_mask )
        attn = torch.squeeze(attn, dim=1).permute(0,2,1)                          # [Batch, IDs_len, f0_len] 
        ph_IDs = torch.matmul(ph_IDs, torch.tensor(attn, dtype=torch.float32))    # to [Batch, inner_channel, f0_len] 
        
        if NoteIDs is not None:
          noteIDs = torch.matmul(noteIDs, torch.tensor(attn, dtype=torch.float32))  # to [Batch, inner_channel, f0_len] 
          hidden = ph_IDs + noteIDs
        else:
          hidden = ph_IDs
        # NoisePrediction Process
        f0 = self.WaveNet(  x         =f0 ,#+ vuv,      
                            x_mask    =f0_mask, 
                            IDs       =hidden, 
                            IDs_mask  =ph_IDs_mask, 
                            timesteps =emb_timesteps, 
                            g         =g)  

        # Projection
        f0 = self.proj_1(f0)
        f0 = self.relu  (f0)
        f0 = self.proj_2(f0) * f0_mask

        # ReZero Regularization
        # f0_vuv = self.rezero(f0_vuv)
        
        # f0, vuv_vector = torch.split(tensor=f0_vuv, 
        #                             split_size_or_sections=1,  # ここlabel３つ
        #                             dim=1 )
        # vuv_vector = self.vuv_proj(vuv_vector)
        

        return f0 # , torch.zeros(size=(f0.size(0), 3, f0.size(2))).cuda()


class U_Net_1D(torch.nn.Module):
    def __init__(self,in_ch=192,             # Hidden
                    inner_ch=256,          # 
                    filter_ch=768,
                    out_ch=4,              # f0, mask, v, uv  
                    time_embed_dim=512,
                    attn_filter_ch=768,
                    ):
        super().__init__()
        
        ########################
        ### For Down Process ###
        ########################
        self.ResBlock_D1 = ResBlock1D(inp_channels=in_ch, 
                                      out_channels=inner_ch,
                                      time_embed_dim=time_embed_dim)
        self.CrossAttn_D1 = CrossAttn(hidden_channels=inner_ch,
                                      filter_channels=filter_ch,
                                      n_heads=4,
                                      n_layers=1,
                                      kernel_size=3,
                                      p_dropout=0.1)
        self.DownSample_D1 =Downsample1D(channels=inner_ch,
                                         use_conv=True,
                                         out_channels=inner_ch)

        self.ResBlock_D2 = ResBlock1D(inp_channels=inner_ch, 
                                      out_channels=inner_ch,
                                      time_embed_dim=time_embed_dim)
        self.CrossAttn_D2 = CrossAttn(hidden_channels=inner_ch,
                                      filter_channels=filter_ch,
                                      n_heads=4,
                                      n_layers=1,
                                      kernel_size=3,
                                      p_dropout=0.1)
        self.DownSample_D2 =Downsample1D(channels=inner_ch,
                                         use_conv=True,
                                         out_channels=inner_ch)
        
        
        #######################
        ### For Mid Process ###
        #######################
        self.ResBlock_M1 = ResBlock1D(inp_channels=inner_ch, 
                                      out_channels=inner_ch,
                                      time_embed_dim=time_embed_dim)
        self.SelfAttn    = SelfAttn(hidden_channels=inner_ch,
                                    filter_channels=filter_ch,
                                    n_heads=4,
                                    n_layers=1,
                                    kernel_size=3,
                                    p_dropout=0.1)
        self.ResBlock_M2 = ResBlock1D(inp_channels=inner_ch, 
                                      out_channels=inner_ch,
                                      time_embed_dim=time_embed_dim)

        
        ######################
        ### For Up Process ###
        ######################
        self.UpSample_U1 = Upsample1D(channels=inner_ch,
                                      use_conv=True,
                                      out_channels=inner_ch)
        self.ResBlock_U1 = ResBlock1D(inp_channels=inner_ch*2, 
                                      out_channels=inner_ch,
                                      time_embed_dim=time_embed_dim)
        self.CrossAttn_U1 = CrossAttn(hidden_channels=inner_ch,
                                      filter_channels=filter_ch,
                                      n_heads=4,
                                      n_layers=1,
                                      kernel_size=3,
                                      p_dropout=0.1)
        
        self.UpSample_U2 = Upsample1D(channels=inner_ch,
                                      use_conv=True,
                                      out_channels=inner_ch)
        self.ResBlock_U2 = ResBlock1D(inp_channels=in_ch*2, 
                                      out_channels=inner_ch,
                                      time_embed_dim=time_embed_dim)
        self.CrossAttn_U2 = CrossAttn(hidden_channels=inner_ch,
                                      filter_channels=filter_ch,
                                      n_heads=4,
                                      n_layers=1,
                                      kernel_size=3,
                                      p_dropout=0.1)
        
        self.final = OutConv1DBlock(num_groups_out=8,
                                    out_channels=out_ch,
                                    embed_dim=inner_ch,
                                    act_fn="mish")
        
    
    def forward(self, x, x_mask, h, h_mask, t):

        # Down Process 
        x = self.ResBlock_D1(inputs=x, t=t)
        x_0 = self.CrossAttn_D1(x=x, x_mask=x_mask, h=h, h_mask=h_mask)
        _, _, x_0_len = x_0.shape
        x = self.DownSample_D1(inputs=x_0) # length/2 
        
        x = self.ResBlock_D2(inputs=x, t=t)
        x_mask_1 = self.mask_downsample(downsample_x=x, mask=x_mask)
        x_1 = self.CrossAttn_D2(x=x, x_mask=x_mask_1, h=h, h_mask=h_mask)
        _, _, x_1_len = x_1.shape
        x = self.DownSample_D2(inputs=x_1) # length/4

        # Mid Process
        x = self.ResBlock_M1(inputs=x, t=t)
        x_mask_2 = self.mask_downsample(downsample_x=x, mask=x_mask_1)
        x = self.SelfAttn(x=x,x_mask=x_mask_2)
        x = self.ResBlock_M2(inputs=x, t=t) # length/4

        # Up Process
        x = self.UpSample_U1(inputs=x) # length/2
        x = torch.cat(tensors=[x_1, x[:,:,:x_1_len]], dim=1) # concat in "channel direction"
        x = self.ResBlock_U1(inputs=x, t=t)
        x = self.CrossAttn_U1(x=x, x_mask=x_mask_1, h=h, h_mask=h_mask)

        x = self.UpSample_U2(inputs=x)
        x = torch.cat(tensors=[x_0, x[:,:,:x_0_len]], dim=1) # concat in "channel direction"
        x = self.ResBlock_U2(inputs=x, t=t)
        x = self.CrossAttn_U2(x=x, x_mask=x_mask, h=h, h_mask=h_mask)

        # Final Process
        x = self.final(x)
        return x
    
    def mask_downsample(self, downsample_x, mask):
        _,_,mask_length = mask.shape
        _,_,x_length = downsample_x.shape
        if mask_length % 2 == 0: # even 
            mask = mask[:,:,1::2]
        else:
            mask = mask[:,:,::2]

        if x_length == mask.size(2):
            return mask

def swish(x):
    return x * torch.sigmoid(x)

class PosteriorEncoder(nn.Module):
  def __init__(self,
      in_channels,
      out_channels,
      hidden_channels,
      kernel_size,
      dilation_rate,
      n_layers,
      gin_channels=0):
    super().__init__()
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.hidden_channels = hidden_channels
    self.kernel_size = kernel_size
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.gin_channels = gin_channels

    self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
    self.enc = modules.WN_original(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)

  def forward(self, x, x_lengths, g=None):
    x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
    x = self.pre(x) * x_mask
    x = self.enc(x, x_mask, g=g)
    stats = self.proj(x) * x_mask
    m, logs = torch.split(stats, self.out_channels, dim=1)
    z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
    return z, m, logs, x_mask

class TextEncoder_VITS2(nn.Module):
  def __init__(self,
      n_vocab,
      n_note, 
      out_channels,
      hidden_channels,
      filter_channels,
      n_heads,
      n_layers,
      kernel_size,
      p_dropout,
      gin_channels=0):
    super().__init__()
    self.n_vocab = n_vocab
    self.out_channels = out_channels
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout
    self.gin_channels = gin_channels
    self.x_emb = nn.Embedding(n_vocab, hidden_channels)
    self.ph_w_idx_emb = nn.Embedding(2,  hidden_channels)
    self.dur_emb = nn.Linear(1, hidden_channels)
    nn.init.normal_(self.x_emb.weight, 0.0, hidden_channels**-0.5)
    nn.init.normal_(self.ph_w_idx_emb.weight, 0.0, hidden_channels**-0.5)

    self.encoder = attentions.Encoder_VITS2(
      hidden_channels,
      filter_channels,
      n_heads,
      n_layers,
      kernel_size,
      p_dropout,
      gin_channels=self.gin_channels)
    self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1)

  def forward(self, x, x_lengths, w_dur_ms=None, ph_w_idx=None,  g=None):
    x = self.x_emb(x) / math.sqrt(self.hidden_channels) # [b, t, h]
    x = torch.transpose(x, 1, -1) # [b, h, t]

    x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)

    if w_dur_ms is not None:
      w_dur_ms = self.dur_emb(w_dur_ms.unsqueeze(2)).permute(0,2,1) / math.sqrt(self.hidden_channels)
      x = x + w_dur_ms
    if ph_w_idx is not None:
      ph_w_idx = torch.diff(ph_w_idx, dim=1, prepend=ph_w_idx.new_zeros(x.size(0), 1)) >= 0
      ph_w_idx = self.ph_w_idx_emb(ph_w_idx.long()).permute(0,2,1) / math.sqrt(self.hidden_channels)  # [B, T_ph, H]
      x = x + ph_w_idx
    
    enc_in = x * x_mask 
    x = self.encoder(enc_in, x_mask, g=g)
    stats = self.proj(x) * x_mask

    m, logs = torch.split(stats, self.out_channels, dim=1)
    return x, m, logs, x_mask

class TextEncoder(nn.Module):
  def __init__(self,
      n_vocab,
      out_channels,
      hidden_channels,
      filter_channels,
      n_heads,
      n_layers,
      kernel_size,
      p_dropout):
    super().__init__()
    self.n_vocab = n_vocab
    self.out_channels = out_channels
    self.hidden_channels = hidden_channels
    self.filter_channels = filter_channels
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.kernel_size = kernel_size
    self.p_dropout = p_dropout

    self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1)
    self.emb = nn.Embedding(n_vocab, hidden_channels)

    self.encoder = attentions.Encoder(
      hidden_channels,
      filter_channels,
      n_heads,
      n_layers,
      kernel_size,
      p_dropout)

  def forward(self, x, x_lengths, noise_scale=1):
    x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]
    x = torch.transpose(x, 1, -1) # [b, h, t]
    x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)

    # Process attention
    x = self.encoder(x * x_mask, x_mask)
    x = self.proj(x) * x_mask
    m, logs = torch.split(x, self.out_channels, dim=1)
    x = (m + torch.randn_like(m) * torch.exp(logs) * noise_scale) * x_mask
    return x, x_mask

def swish(x):
    return x * torch.sigmoid(x)

class DiscriminatorP(torch.nn.Module):
    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
        super(DiscriminatorP, self).__init__()
        self.period = period
        self.use_spectral_norm = use_spectral_norm
        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
        self.convs = nn.ModuleList([
            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))),
        ])
        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))

    def forward(self, x):
        fmap = []

        # 1d to 2d
        b, c, t = x.shape
        if t % self.period != 0: # pad first
            n_pad = self.period - (t % self.period)
            x = nn.functional.pad(x, (0, n_pad), "reflect")
            t = t + n_pad
        x = x.view(b, c, t // self.period, self.period)

        for l in self.convs:
            x = l(x)
            x = nn.functional.leaky_relu(x, modules.LRELU_SLOPE)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        x = torch.flatten(x, 1, -1)

        return x, fmap


class DiscriminatorS(torch.nn.Module):
    def __init__(self, use_spectral_norm=False):
        super(DiscriminatorS, self).__init__()
        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
        self.convs = nn.ModuleList([
            norm_f(Conv1d(1, 16, 15, 1, padding=7)),
            norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
            norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
            norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
            norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
        ])
        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))

    def forward(self, x):
        fmap = []

        for l in self.convs:
            x = l(x) 
            x = nn.functional.leaky_relu(x, modules.LRELU_SLOPE)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        x = torch.flatten(x, 1, -1)

        return x, fmap


class MultiPeriodDiscriminator(torch.nn.Module):
    def __init__(self, periods, use_spectral_norm=False):
        super(MultiPeriodDiscriminator, self).__init__()
        #periods = [2,3,5,7,11] 

        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
        discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
        self.discriminators = nn.ModuleList(discs)

    def forward(self, y, y_hat):
        y_d_rs = []
        y_d_gs = []
        fmap_rs = []
        fmap_gs = []
        for i, d in enumerate(self.discriminators):
            y_d_r, fmap_r = d(y)
            y_d_g, fmap_g = d(y_hat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)

        return y_d_rs, y_d_gs, fmap_rs, fmap_gs


if __name__ == "__main__":

    from utils import get_hparams
    hps_path = "./configs/config.yaml"
    hparams = get_hparams(hps_path)


models.py
tonnetonne814/SiFi-VITS2-44100-Ja
import copy
import math
import numpy as np
import scipy
import torch
from torch import nn
from torch.nn import functional as F

from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
from torch.nn.utils import weight_norm, remove_weight_norm

import commons
from commons import init_weights, get_padding
from transforms import piecewise_rational_quadratic_transform


LRELU_SLOPE = 0.1


class LayerNorm(nn.Module):
  def __init__(self, channels, eps=1e-5):
    super().__init__()
    self.channels = channels
    self.eps = eps

    self.gamma = nn.Parameter(torch.ones(channels))
    self.beta = nn.Parameter(torch.zeros(channels))

  def forward(self, x):
    x = x.transpose(1, -1)
    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
    return x.transpose(1, -1)

 
class ConvReluNorm(nn.Module):
  def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
    super().__init__()
    self.in_channels = in_channels
    self.hidden_channels = hidden_channels
    self.out_channels = out_channels
    self.kernel_size = kernel_size
    self.n_layers = n_layers
    self.p_dropout = p_dropout
    assert n_layers > 1, "Number of layers should be larger than 0."

    self.conv_layers = nn.ModuleList()
    self.norm_layers = nn.ModuleList()
    self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
    self.norm_layers.append(LayerNorm(hidden_channels))
    self.relu_drop = nn.Sequential(
        nn.ReLU(),
        nn.Dropout(p_dropout))
    for _ in range(n_layers-1):
      self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
      self.norm_layers.append(LayerNorm(hidden_channels))
    self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
    self.proj.weight.data.zero_()
    self.proj.bias.data.zero_()

  def forward(self, x, x_mask):
    x_org = x
    for i in range(self.n_layers):
      x = self.conv_layers[i](x * x_mask)
      x = self.norm_layers[i](x)
      x = self.relu_drop(x)
    x = x_org + self.proj(x)
    return x * x_mask


class DDSConv(nn.Module):
  """
  Dialted and Depth-Separable Convolution
  """
  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
    super().__init__()
    self.channels = channels
    self.kernel_size = kernel_size
    self.n_layers = n_layers
    self.p_dropout = p_dropout

    self.drop = nn.Dropout(p_dropout)
    self.convs_sep = nn.ModuleList()
    self.convs_1x1 = nn.ModuleList()
    self.norms_1 = nn.ModuleList()
    self.norms_2 = nn.ModuleList()
    for i in range(n_layers):
      dilation = kernel_size ** i
      padding = (kernel_size * dilation - dilation) // 2
      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, 
          groups=channels, dilation=dilation, padding=padding
      ))
      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
      self.norms_1.append(LayerNorm(channels))
      self.norms_2.append(LayerNorm(channels))

  def forward(self, x, x_mask, g=None):
    if g is not None:
      x = x + g
    for i in range(self.n_layers):
      y = self.convs_sep[i](x * x_mask)
      y = self.norms_1[i](y)
      y = F.gelu(y)
      y = self.convs_1x1[i](y)
      y = self.norms_2[i](y)
      y = F.gelu(y)
      y = self.drop(y)
      x = x + y
    return x * x_mask


class WN(torch.nn.Module):
  def __init__(self, hidden_channels, Diff_step_embed_out, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
    super(WN, self).__init__()
    assert(kernel_size % 2 == 1)
    self.hidden_channels =hidden_channels
    self.kernel_size = kernel_size,
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.gin_channels = gin_channels
    self.p_dropout = p_dropout

    self.in_layers = torch.nn.ModuleList()
    self.res_skip_layers = torch.nn.ModuleList()
    self.drop = nn.Dropout(p_dropout)

    # timestep embeddings
    #self.timestep_embed = nn.Linear(Diff_step_embed_out, hidden_channels)
    cond_layer_timesteps = torch.nn.Linear(Diff_step_embed_out, hidden_channels*n_layers)
    self.timestep_embed = torch.nn.utils.weight_norm(cond_layer_timesteps, name='weight')

    if gin_channels != 0:
      cond_layer_g = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
      self.cond_layer_g = torch.nn.utils.weight_norm(cond_layer_g, name='weight')
    
    cond_layer_IDs = torch.nn.Conv1d(hidden_channels, 2*hidden_channels*n_layers, 1)
    self.cond_layer_IDs = torch.nn.utils.weight_norm(cond_layer_IDs, name='weight')

    for i in range(n_layers):
      dilation = dilation_rate ** i
      padding = int((kernel_size * dilation - dilation) / 2)
      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
                                 dilation=dilation, padding=padding)
      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
      self.in_layers.append(in_layer)

      # last one is not necessary
      if i < n_layers - 1:
        res_skip_channels = 2 * hidden_channels
      else:
        res_skip_channels = hidden_channels

      res_skip_layer = torch.nn.Conv1d(hidden_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, x, x_mask, IDs, IDs_mask, timesteps, g=None, **kwargs):
    output = torch.zeros_like(x)
    n_channels_tensor = torch.IntTensor([self.hidden_channels])

    # Speaker Embedding
    if g is not None:
      g = self.cond_layer_g(g) # conv 1x1
    
    # Text Embedding
    IDs = self.cond_layer_IDs(IDs) # conv 1x1
    
    timesteps = self.timestep_embed(timesteps) # Linear
    timesteps = torch.unsqueeze(timesteps, dim=2)

    # Residual process
    for i in range(self.n_layers):
      timesteps_offset = i * self.hidden_channels # offset
      cond_offset = i * 2 * self.hidden_channels # offset
      
      timesteps_l = timesteps[:,timesteps_offset:timesteps_offset+self.hidden_channels, :]
      x_ts = x + timesteps_l # 
      x_in = self.in_layers[i](x_ts) # conv 3x3

      # condition process
      if g is not None:
        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
      else:
        g_l = torch.zeros_like(x_in)
      
      IDs_l = IDs[:,cond_offset:cond_offset+2*self.hidden_channels,:]

      acts = commons.fused_add_tanh_sigmoid_multiply(
          x_in,
          IDs_l + g_l,
          n_channels_tensor)
      acts = self.drop(acts)

      res_skip_acts = self.res_skip_layers[i](acts)
      if i < self.n_layers - 1: # other layers
        res_acts = res_skip_acts[:,:self.hidden_channels,:]
        x = (x + res_acts) * x_mask                                 # residual out "x". x is reInputed 
        output = output + res_skip_acts[:,self.hidden_channels:,:]  # skip out ==> sum(skip out) is final output
      else: # last layer
        output = output + res_skip_acts # final output
    return output * x_mask

  def remove_weight_norm(self):
    if self.gin_channels != 0:
      torch.nn.utils.remove_weight_norm(self.cond_layer)
    for l in self.in_layers:
      torch.nn.utils.remove_weight_norm(l)
    for l in self.res_skip_layers:
     torch.nn.utils.remove_weight_norm(l)

class WN_original(torch.nn.Module):
  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
    super(WN_original, self).__init__()
    assert(kernel_size % 2 == 1)
    self.hidden_channels =hidden_channels
    self.kernel_size = kernel_size,
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.gin_channels = gin_channels
    self.p_dropout = p_dropout

    self.in_layers = torch.nn.ModuleList()
    self.res_skip_layers = torch.nn.ModuleList()
    self.drop = nn.Dropout(p_dropout)

    if gin_channels != 0:
      cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
      self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

    for i in range(n_layers):
      dilation = dilation_rate ** i
      padding = int((kernel_size * dilation - dilation) / 2)
      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
                                 dilation=dilation, padding=padding)
      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
      self.in_layers.append(in_layer)

      # last one is not necessary
      if i < n_layers - 1:
        res_skip_channels = 2 * hidden_channels
      else:
        res_skip_channels = hidden_channels

      res_skip_layer = torch.nn.Conv1d(hidden_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, x, x_mask, g=None, **kwargs):
    output = torch.zeros_like(x)
    n_channels_tensor = torch.IntTensor([self.hidden_channels])

    if g is not None:
      g = self.cond_layer(g)

    for i in range(self.n_layers):
      x_in = self.in_layers[i](x)
      if g is not None:
        cond_offset = i * 2 * self.hidden_channels
        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
      else:
        g_l = torch.zeros_like(x_in)

      acts = commons.fused_add_tanh_sigmoid_multiply(
          x_in,
          g_l,
          n_channels_tensor)
      acts = self.drop(acts)

      res_skip_acts = self.res_skip_layers[i](acts)
      if i < self.n_layers - 1:
        res_acts = res_skip_acts[:,:self.hidden_channels,:]
        x = (x + res_acts) * x_mask
        output = output + res_skip_acts[:,self.hidden_channels:,:]
      else:
        output = output + res_skip_acts
    return output * x_mask

  def remove_weight_norm(self):
    if self.gin_channels != 0:
      torch.nn.utils.remove_weight_norm(self.cond_layer)
    for l in self.in_layers:
      torch.nn.utils.remove_weight_norm(l)
    for l in self.res_skip_layers:
     torch.nn.utils.remove_weight_norm(l)


class ResBlock1(torch.nn.Module):
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(ResBlock1, self).__init__()
        self.convs1 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)

    def forward(self, x, x_mask=None):
        for c1, c2 in zip(self.convs1, self.convs2):
            xt = F.leaky_relu(x, LRELU_SLOPE)
            if x_mask is not None:
                xt = xt * x_mask
            xt = c1(xt)
            xt = F.leaky_relu(xt, LRELU_SLOPE)
            if x_mask is not None:
                xt = xt * x_mask
            xt = c2(xt)
            x = xt + x
        if x_mask is not None:
            x = x * x_mask
        return x

    def remove_weight_norm(self):
        for l in self.convs1:
            remove_weight_norm(l)
        for l in self.convs2:
            remove_weight_norm(l)


class ResBlock2(torch.nn.Module):
    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
        super(ResBlock2, self).__init__()
        self.convs = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1])))
        ])
        self.convs.apply(init_weights)

    def forward(self, x, x_mask=None):
        for c in self.convs:
            xt = F.leaky_relu(x, LRELU_SLOPE)
            if x_mask is not None:
                xt = xt * x_mask
            xt = c(xt)
            x = xt + x
        if x_mask is not None:
            x = x * x_mask
        return x

    def remove_weight_norm(self):
        for l in self.convs:
            remove_weight_norm(l)


class Log(nn.Module):
  def forward(self, x, x_mask, reverse=False, **kwargs):
    if not reverse:
      y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
      logdet = torch.sum(-y, [1, 2])
      return y, logdet
    else:
      x = torch.exp(x) * x_mask
      return x
    

class Flip(nn.Module):
  def forward(self, x, *args, reverse=False, **kwargs):
    x = torch.flip(x, [1])
    if not reverse:
      logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
      return x, logdet
    else:
      return x


class ElementwiseAffine(nn.Module):
  def __init__(self, channels):
    super().__init__()
    self.channels = channels
    self.m = nn.Parameter(torch.zeros(channels,1))
    self.logs = nn.Parameter(torch.zeros(channels,1))

  def forward(self, x, x_mask, reverse=False, **kwargs):
    if not reverse:
      y = self.m + torch.exp(self.logs) * x
      y = y * x_mask
      logdet = torch.sum(self.logs * x_mask, [1,2])
      return y, logdet
    else:
      x = (x - self.m) * torch.exp(-self.logs) * x_mask
      return x


class ResidualCouplingLayer(nn.Module):
  def __init__(self,
      channels,
      hidden_channels,
      kernel_size,
      dilation_rate,
      n_layers,
      p_dropout=0,
      gin_channels=0,
      mean_only=False):
    assert channels % 2 == 0, "channels should be divisible by 2"
    super().__init__()
    self.channels = channels
    self.hidden_channels = hidden_channels
    self.kernel_size = kernel_size
    self.dilation_rate = dilation_rate
    self.n_layers = n_layers
    self.half_channels = channels // 2
    self.mean_only = mean_only

    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
    self.enc = WN_original(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
    self.post.weight.data.zero_()
    self.post.bias.data.zero_()

  def forward(self, x, x_mask, g=None, reverse=False):
    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
    h = self.pre(x0) * x_mask
    h = self.enc(h, x_mask, g=g)
    stats = self.post(h) * x_mask
    if not self.mean_only:
      m, logs = torch.split(stats, [self.half_channels]*2, 1)
    else:
      m = stats
      logs = torch.zeros_like(m)

    if not reverse:
      x1 = m + x1 * torch.exp(logs) * x_mask
      x = torch.cat([x0, x1], 1)
      logdet = torch.sum(logs, [1,2])
      return x, logdet
    else:
      x1 = (x1 - m) * torch.exp(-logs) * x_mask
      x = torch.cat([x0, x1], 1)
      return x


class ConvFlow(nn.Module):
  def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
    super().__init__()
    self.in_channels = in_channels
    self.filter_channels = filter_channels
    self.kernel_size = kernel_size
    self.n_layers = n_layers
    self.num_bins = num_bins
    self.tail_bound = tail_bound
    self.half_channels = in_channels // 2

    self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
    self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
    self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
    self.proj.weight.data.zero_()
    self.proj.bias.data.zero_()

  def forward(self, x, x_mask, g=None, reverse=False):
    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
    h = self.pre(x0)
    h = self.convs(h, x_mask, g=g)
    h = self.proj(h) * x_mask

    b, c, t = x0.shape
    h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]

    unnormalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
    unnormalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
    unnormalized_derivatives = h[..., 2 * self.num_bins:]

    x1, logabsdet = piecewise_rational_quadratic_transform(x1,
        unnormalized_widths,
        unnormalized_heights,
        unnormalized_derivatives,
        inverse=reverse,
        tails='linear',
        tail_bound=self.tail_bound
    )

    x = torch.cat([x0, x1], 1) * x_mask
    logdet = torch.sum(logabsdet * x_mask, [1,2])
    if not reverse:
        return x, logdet
    else:
        return x

modules.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import torchaudio
import torch

from tqdm import tqdm
import random
import argparse
from f0_extractor import F0_extractor
import subprocess
from singDB_loader import get_annotated_data, get_lab_info

target_sr = 44100
hop_size = 512
song_min_s = 5000 
split_ratio = 0.005

def preprocess(folder, f0_method, out_dir, audio_norm):
    os.makedirs(out_dir, exist_ok=True)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    song_dir = os.path.join(folder, "DATABASE")
    oto2lab_path = os.path.join(folder,"kana2phonemes_002_oto2lab.table" )
    pitch_extractor = F0_extractor(samplerate=target_sr,
                                   hop_size=hop_size,
                                   device=device)
    count = 0
    filelist = list()
    for song_name in tqdm(os.listdir(song_dir), desc="Preprocessing..."):
        lab_path        = os.path.join(song_dir, song_name, song_name + ".lab")
        ust_path        = os.path.join(song_dir, song_name, song_name + ".ust")

        if audio_norm is True:
            orig_wav_path = os.path.join(song_dir, song_name, song_name + ".wav")
            norm_wav_path = os.path.join(song_dir, song_name, song_name + "_norm.wav")
            audio_norm_process(orig_wav_path, norm_wav_path)
            wav_path = norm_wav_path
        else:
            wav_path = os.path.join(song_dir, song_name, song_name + ".wav")

        anno_datalist = get_annotated_data(ust_path = ust_path,
                                           lab_path=lab_path,
                                           ono2lab_table_path=oto2lab_path)
        
        f0, vuv = pitch_extractor.compute_f0(path=wav_path, f0_method=f0_method)
        wav, sr = torchaudio.load(wav_path)
        if sr != target_sr:
            wav = torchaudio.functional.resample(wav, orig_freq=sr, new_freq=target_sr)
        if wav.size(0) > 1:
            print("[ERROR] Audio Channel is not 1ch.")
            exit()
        wav = torch.squeeze(wav, dim=0)
        
        t_s = 0
        t_e = 0
        ph_list             = list()
        ph_s_ms_list        = list()
        ph_e_ms_list        = list()
        ph_dur_ms_list      = list()
        word_s_ms_list      = list()
        word_e_ms_list      = list()
        word_dur_ms_list    = list()
        word_list           = list()
        noteID_list         = list()
        notename_list       = list()
        notedur_list        = list()
        n_ph_in_a_word      = list()

        for idx , line in enumerate(anno_datalist):

            # 無音時　かつ　長さが足りている時（開始時終了時にpauは来ない方式
            if line["ph"][0][0] == "pau":
                t_e = line["start_ms"]

                if idx == 0:
                    t_s = line["end_ms"]
                    continue
                
                # 保存データが短すぎる、かつ今回のpau_durも短いなら、再ループへ
                if t_e - t_s < song_min_s and line["dur_ms"] < song_min_s:
                    for data in line["ph"]:
                        ph_list             .append(data[0])
                        ph_s_ms_list        .append(data[1] - t_s)
                        ph_e_ms_list        .append(data[2] - t_s)
                        ph_dur_ms_list      .append(data[3])
                    word_s_ms_list      .append(line["start_ms"] - t_s)
                    word_e_ms_list      .append(line["end_ms"]   - t_s)
                    word_dur_ms_list    .append(line["dur_ms"])
                    word_list           .append(line["word"])
                    noteID_list         .append(line["note_num"])
                    notename_list       .append(line["notename"])
                    notedur_list        .append(line["note_dur"])
                    n_ph_in_a_word      .append(len(line["ph"]))
                    continue
                
                # check length
                word_len = len(word_list)
                if  word_len != len(word_s_ms_list) or \
                    word_len != len(word_e_ms_list) or \
                    word_len != len(word_dur_ms_list) or \
                    word_len != len(noteID_list) or \
                    word_len != len(notename_list) or \
                    word_len != len(n_ph_in_a_word):
                    print("[ERROR] word_length error. check preprocess.py")
                    exit()
                ph_len = len(ph_list)
                if ph_len != len(ph_s_ms_list) or \
                    ph_len != len(ph_e_ms_list) or \
                    ph_len != len(ph_dur_ms_list) or \
                    ph_len != int(torch.sum(torch.tensor(n_ph_in_a_word, dtype=torch.int64))):
                    print("[ERROR] ph_length error. check preprocess.py")
                    exit()

                # save
                basename = os.path.join(out_dir, str(count).zfill(5))
                w_s_idx = int((t_s/1000) * target_sr)  
                w_e_idx = int((t_e/1000) * target_sr)  
                h_s_idx = int((t_s/1000) * target_sr / hop_size)
                h_e_idx = int((t_e/1000) * target_sr / hop_size)

                # last check (diff is below 3 frame)
                diff = abs(abs(t_e -t_s) - int(torch.sum(torch.tensor(word_dur_ms_list, dtype=torch.int64))))
                if diff > 3 *(hop_size/target_sr) * 1000:
                    print("[WARNING] Detects errors of 3 frames or more. {:.5f} [ms]".format(diff))

                torchaudio.save(filepath=basename + ".wav", 
                                src=torch.unsqueeze(wav[w_s_idx:w_e_idx], dim=0) ,
                                sample_rate=target_sr)
                torch.save(f0[h_s_idx:h_e_idx], basename + "_f0.pt")
                torch.save(vuv[h_s_idx:h_e_idx],basename + "_vuv.pt")
                torch.save(ph_list,             basename + "_ph.pt" )
                torch.save(ph_s_ms_list,        basename + "_ph_s_ms.pt" )
                torch.save(ph_e_ms_list,        basename + "_ph_e_ms.pt" )
                torch.save(ph_dur_ms_list,      basename + "_ph_dur_ms.pt" )
                torch.save(word_s_ms_list,      basename + "_word_s_ms.pt" )
                torch.save(word_e_ms_list,      basename + "_word_e_ms.pt" )
                torch.save(word_dur_ms_list,    basename + "_word_dur_ms.pt" )
                torch.save(word_list,           basename + "_word.pt" )
                torch.save(noteID_list,         basename + "_noteID.pt" )
                torch.save(notename_list,       basename + "_notename.pt" )
                torch.save(notedur_list,        basename + "_notedur.pt" )
                torch.save(n_ph_in_a_word,      basename + "_n_ph_in_a_word.pt" )
                filelist.append(basename + "\n")
                count+=1 

                # init
                ph_list             = list()
                word_s_ms_list      = list()
                word_e_ms_list      = list()
                word_dur_ms_list    = list()
                word_list           = list()
                noteID_list         = list()
                notename_list       = list()
                notedur_list        = list()
                ph_s_ms_list        = list()
                ph_e_ms_list        = list()
                ph_dur_ms_list      = list()
                n_ph_in_a_word      = list()
                
                # next start time
                t_s = line["end_ms"]

            else:
                for data in line["ph"]:
                    ph_list             .append(data[0])
                    ph_s_ms_list        .append(data[1] - t_s)
                    ph_e_ms_list        .append(data[2] - t_s)
                    ph_dur_ms_list      .append(data[3])
                word_s_ms_list      .append(line["start_ms"] - t_s)
                word_e_ms_list      .append(line["end_ms"]   - t_s)
                word_dur_ms_list    .append(line["dur_ms"])
                word_list           .append(line["word"])
                noteID_list         .append(line["note_num"])
                notename_list       .append(line["notename"])
                notedur_list        .append(line["note_dur"])
                n_ph_in_a_word      .append(len(line["ph"]))

        if audio_norm is True:
            os.remove(wav_path)

    filelist_split(filelist=filelist)

def audio_norm_process(in_path, out_path):
    if os.path.isfile(in_path) is True:
        pass
    else:
        print("[ERROR] File is not existed : ", in_path)
        exit()
    cmd = ["ffmpeg-normalize", in_path,   "-o",  out_path]
    subprocess.run(cmd, encoding='utf-8', stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def filelist_split(filelist):
    os.makedirs("./filelists/",exist_ok=True)   
    max_n = len(filelist)
    test_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(filelist)
        idx = random.randint(9, int(n-1))
        txt = filelist.pop(idx)
        test_list.append(txt)

    max_n = len(filelist)
    val_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(filelist)
        idx = random.randint(9, int(n-1))
        txt = filelist.pop(idx)
        val_list.append(txt)

    write_txt(f"./filelists/SVS_train_{target_sr}.txt", sorted(filelist))
    write_txt(f"./filelists/SVS_val_{target_sr}.txt",   sorted(val_list))
    write_txt(f"./filelists/SVS_test_{target_sr}.txt",  sorted(test_list))

    return 0

def read_txt(path):
    with open(path, mode="r", encoding="utf-8")as f:
        lines = f.readlines()
    return lines

def write_txt(path, lines):
    with open(path, mode="w", encoding="utf-8")as f:
        f.writelines(lines)

import glob 
def generate_word_list(dataset_dir, out_path):
    pathlist = glob.glob(f"{dataset_dir}/*.wav")
    word_list = list()
    for path in pathlist:
        word_path = path.replace(".wav", "_word.pt")
        word = torch.load(word_path)
        word_list.extend(word)
    word_list = sorted(set(word_list))
    torch.save(word_list, out_path)
    print(f"[INFO] Word list is saved at {out_path}")
    print(f"[INFO] Word list : {word_list}")
    return 0

def take_ph_statistics(dataset_dir, out_path, threshold_ms=50):
    pathlist = glob.glob(f"{dataset_dir}/**/*.lab", recursive=True)

    ph_list = list()
    for path in pathlist:
        for data in get_lab_info(path):
            ph_list.append(data["ph"])
    ph_to_id = {s.replace("\n", ""): i for i, s in enumerate(sorted(set(ph_list)))}
    id_to_ph = {i: s.replace("\n", "") for i, s in enumerate(sorted(set(ph_list)))}
    data_lists = [[] for _ in range(len(ph_to_id))]
    for path in pathlist:
        for data in get_lab_info(path):
            ph = data["ph"]
            ms = data["dur_ms"]
            data_lists[ph_to_id[ph]].append(ms)

    import numpy as np
    statistic_dict = {}
    for idx in range(len(ph_to_id)):
        data = np.array(data_lists[idx])
        ph = id_to_ph[idx]
        mean = np.mean(data)
        std = np.std(data)
        # 大体同じような値の奴は、そのまま使う
        if std<threshold_ms:
            statistic_dict[ph] = [mean, std]

    torch.save(statistic_dict, out_path)
    print(f"[INFO] ph duration statstics is saved at {out_path}")
    print(f"[INFO] ph statstics : {statistic_dict}")
    return 0


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--song_db_path',
                        type=str,
                        #required=True, 
                        default="./「波音リツ」歌声データベースVer2/")
    
    parser.add_argument('--f0_method',
                        type=str,
                        #required=True,
                        default="crepe")
    
    parser.add_argument('--dataset_out_dir',
                        type=str,
                        #required=True,
                        default="./dataset_SVS/")
    
    parser.add_argument('--audio_normalize',
                        type=str,
                        #required=True,
                        default=True)

    args = parser.parse_args()

    preprocess(folder= args.song_db_path,
               f0_method=args.f0_method,
               out_dir=args.dataset_out_dir,
               audio_norm=args.audio_normalize)
    
    take_ph_statistics(dataset_dir=args.song_db_path,
                       out_path=os.path.join(args.dataset_out_dir, "ph_statistics.pt"))

    #import glob
    #filelist = list()
    #for path in glob.glob("./dataset_SVS/*.wav"):
    #    filelist.append(str(path).replace(".wav", "") + "\n")
    #filelist_split(filelist)
    
preprocess_svs.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import argparse
import torch
import yaml
from models import VITS2_based_SiFiTTS
from singDB_loader import get_ust_info, get_g2p_dict_from_tabledata
import pyopenjtalk
import re
import jaconv
from mel_processing import spectrogram_torch
from utils import load_wav_to_torch
from tqdm import tqdm

target_sr = 44100
hop_size = 512
song_min_s = 5000 # ms
split_ratio = 0.005

device =torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
os.makedirs("./filelists/",exist_ok=True)

def get_duration_from_MAS(hps):

    oto2lab, \
    ph_symbol_to_id, id_to_ph_symbol, \
    word_symbol_to_id,id_to_word_symbol = get_g2p_dict_from_tabledata(hps["oto2lab_path"], include_converter=True)

    vocab_size = len(ph_symbol_to_id) + 1

    net_g = VITS2_based_SiFiTTS(
                hps = hps,
                n_vocab=vocab_size,
                spec_channels=hps["spec_encoder"]["spec_channels"],
                segment_size=int(hps["segments_size"] // hps["hop_length"]),
                n_speakers=hps["common"]["n_speaker"],
                gin_channels=hps["common"]["gin_channels"],
                **hps["hifi_gan"]).cuda()    
    
    lines = []
    lines += read_txt(hps["train_data_path"])
    lines += read_txt(hps["eval_data_path"])
    lines += read_txt(hps["test_data_path"])
    count = 0
    for line in tqdm(lines, desc="get duration from MAS..."):
        count += 1
        if count < 4300:
            continue
        path, text = line.split("|")
        spec_path = path + "_spec.pt"
        if os.path.exists(spec_path):
            spec = torch.load(spec_path)
        else:
            wav_path = path + ".wav"
            audio_norm, sr = load_wav_to_torch(wav_path)
            assert sr == hps["sampling_rate"]

            spec = spectrogram_torch(audio_norm.unsqueeze(0), hps["filter_length"],
                hps["sampling_rate"], hps["hop_length"], hps["win_length"],
                center=False)
            spec = torch.squeeze(spec, 0)
            if spec.size(1) == -1:
                print("ERROR SPEC")
            torch.save(spec, spec_path)
        spec = spec.unsqueeze(0).cuda()

        spec_lengths = torch.tensor(spec.size(2), dtype=torch.int64).view(-1).cuda()
        ph_IDs = get_ph_ID(text, oto2lab, ph_symbol_to_id).view(1,-1).cuda()
        ph_IDs_lengths = torch.tensor(ph_IDs.size(1), dtype=torch.int64).view(-1).cuda()
        speakerID = torch.zeros(size=(1,), dtype=torch.int64).view(-1).cuda() # speaker id is fixed

        with torch.inference_mode():
            
            ph_dur, attn = net_g.get_mas_output(spec, 
                                                spec_lengths,
                                                ph_IDs, 
                                                ph_IDs_lengths,
                                                speakerID)
        
        torch.save(ph_dur[0][0].to('cpu').detach().numpy().copy(), path + "_ph_frame_dur.pt")
        torch.save(attn[0][0].to('cpu').detach().numpy().copy(),   path + "_ph_spec_attn.pt")

    return 0

def read_txt(path):
    with open(path, mode="r", encoding="utf-8")as f:
        lines = f.readlines()
    return lines

def write_txt(path, lines):
    with open(path, mode="w", encoding="utf-8")as f:
        f.writelines(lines)

def get_ph_ID(text, oto2lab, ph_symbol_to_id):
    small_letters = ["ぁ", "ぃ", "ぅ", "ぇ", "ぉ", 
                     "ァ", "ィ", "ゥ", "ェ", "ォ",
                     "ゃ", "ゅ", "ょ", "っ", "ゎ",
                     "ャ", "ュ", "ョ", "ッ", "ヮ"]
    find_desc = "|".join(small_letters)
   
    kana = pyopenjtalk.g2p(text, kana=True)
    hira = " ".join(jaconv.kata2hira(kana))
    hira_back= hira

    skip = 0
    for idx, m in enumerate(re.finditer(find_desc, hira)):
        pos = m.span()[0]-(idx-skip)-1 
        if pos < 0:
            skip += 1
            pass
        else:
            hira = hira[:pos] + hira[pos+1:]
    hira = hira.split(" ")
    ph_list = []
    for word in hira:
        try:
            # jsut base covert rule
            if word == "・" or word == "-" or word == "−" or word == "？" :
                continue
            if word == '．':
                word = "。"
            if word == '乃':
                word = "の"
            elif word == '珂':
                ph_list += [oto2lab["か"]]
                continue
            elif word == '是':
                ph_list += [oto2lab["ぜ"]]
                continue

            elif word == '哉':
                ph_list += [oto2lab["や"]]
                continue
            elif word == '孺':
                ph_list += [oto2lab["じゅ"]]
                continue

            if word == "吐":
                ph_list += [oto2lab["ば"]]
                ph_list += [oto2lab["き"]]
                continue
            elif word == '熏':
                ph_list += [oto2lab["く"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '撃':
                ph_list += [oto2lab["げ"]]
                ph_list += [oto2lab["き"]]
                continue
            elif word == '哨':
                ph_list += [oto2lab["しょ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '逓':
                ph_list += [oto2lab["て"]]
                ph_list += [oto2lab["い"]]
                continue
            elif word == '騎':
                ph_list += [oto2lab["き"]]
                continue
            elif word == '閃':
                ph_list += [oto2lab["せ"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '抽':
                ph_list += [oto2lab["ちゅ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '分':
                ph_list += [oto2lab["ぶ"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '斬':
                ph_list += [oto2lab["ざ"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '醸':
                ph_list += [oto2lab["じょ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '勃':
                ph_list += [oto2lab["ぼ"]]
                ph_list += [oto2lab["つ"]]
                continue
            elif word == '煖':
                ph_list += [oto2lab["え"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '四':
                ph_list += [oto2lab["し"]]
                continue
            elif word == '川':
                ph_list += [oto2lab["せ"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '省':
                ph_list += [oto2lab["しょ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '雅':
                ph_list += [oto2lab["が"]]
                continue
            elif word == '安':
                ph_list += [oto2lab["あ"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '県':
                ph_list += [oto2lab["け"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '高':
                ph_list += [oto2lab["こ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '頤':
                ph_list += [oto2lab["い"]]
                continue
            elif word == '墓':
                ph_list += [oto2lab["ぼ"]]
                continue
            elif word == '闕':
                ph_list += [oto2lab["け"]]
                ph_list += [oto2lab["つ"]]
                continue
            elif word == '窒':
                ph_list += [oto2lab["ち"]]
                ph_list += [oto2lab["つ"]]
                continue
            elif word == '菌':
                ph_list += [oto2lab["き"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '聚':
                ph_list += [oto2lab["じゅ"]]
                continue
            elif word == '慎':
                ph_list += [oto2lab["し"]]
                ph_list += [oto2lab["ん"]]
                continue
            elif word == '果':
                ph_list += [oto2lab["か"]]
                continue
            elif word == '胞':
                ph_list += [oto2lab["ほ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '溶':
                ph_list += [oto2lab["よ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '結':
                ph_list += [oto2lab["け"]]
                ph_list += [oto2lab["つ"]]
                continue
            elif word == '腔':
                ph_list += [oto2lab["こ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '療':
                ph_list += [oto2lab["りょ"]]
                ph_list += [oto2lab["う"]]
                continue
            elif word == '々': 
                ph_list += [oto2lab["し"]]
                continue
            elif word == '禕': 
                ph_list += [oto2lab["い"]]
                continue

            if word == 'ー':
                c_ph = ph_list[-1]
                if len(c_ph) != 1:
                    ph_list += c_ph[-1]
                else:
                    ph_list += c_ph
                continue

            elif word == "ーっ":
                c_ph = ph_list[-1]
                if len(c_ph) != 1:
                    ph_list += [c_ph[-1]+" "+ oto2lab["っ"]]
                else:
                    ph_list += [c_ph     + " " + oto2lab["っ"]]
                continue

            try:
                word_ph = oto2lab[word]
            except:
                txt = ""   
                word_ph = ""   
                # check cl
                for w in word:
                    if w == "っ":
                        word_ph += oto2lab[txt] + " " + oto2lab[w] # pull from dict
                        txt = "" #reset
                    else:
                        txt += w # stack
                # last pull
                if txt != "":
                    word_ph += oto2lab[txt]
            ph_list += [word_ph]
        except:
            tmp = []
            for w in word:
                ph_list += [oto2lab[w]]
                tmp += [oto2lab[w]]
            print(f"[WARNING] {word} is not existed in oto2lab.")
            print(f"[WARNING] Full text is {hira_back}.")
            print(f"[WARNING] {word} is converted to {tmp}.")

    sequence = []
    for symbols in ph_list:
        for symbol in symbols.split(" "):
            symbol_id = ph_symbol_to_id[symbol]
            sequence += [symbol_id]
    return torch.tensor(sequence, dtype=torch.int64)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--config',
                        type=str,
                        default="./configs/config.yaml")
    
    args = parser.parse_args()

    config_path = args.config
    with open(config_path, mode="r", encoding="utf-8") as f:
      hparams = yaml.safe_load(f)

    get_duration_from_MAS(hps=hparams)

preprocess_tts_F0Diffusion.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
from tqdm import tqdm
import random
import argparse
import torchaudio 
from f0_extractor import F0_extractor
import torch

target_sr = 44100
hop_size = 512
song_min_s = 5000 # ms
split_ratio = 0.005

device =torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
os.makedirs("./filelists/",exist_ok=True)

def jsut_preprocess(dataset_dir:str = "./jsut_ver1.1/basic5000/", 
                    f0_method:str = "crepe",
                    results_folder="./dataset_TTS/jsut/",
                     normalize = False ):
    
    os.makedirs(results_folder, exist_ok=True)
    
    pitch_extractor = F0_extractor(samplerate=target_sr,
                                   hop_size=hop_size,
                                   device=device)

    wav_dir = os.path.join(dataset_dir, "wav")
    #"""
    for filename in tqdm(os.listdir(wav_dir)):
        wav_path = os.path.join(wav_dir, filename)
        if normalize is True:
            out_path = wav_path.replace(".wav", "_norm.wav")
            audio_norm_process(wav_path, out_path)
            wav_path = out_path

        y, sr = torchaudio.load(wav_path)
        y_converted = torchaudio.functional.resample(waveform=y, orig_freq=sr, new_freq=target_sr)
        save_path = os.path.join(results_folder, filename)
        torchaudio.save(save_path, y_converted, target_sr) 
        
        f0, vuv = pitch_extractor.compute_f0(path=save_path, f0_method=f0_method)
        torch.save(f0, save_path.replace(".wav", "_f0.pt") )
        torch.save(vuv, save_path.replace(".wav", "_vuv.pt") )
        if normalize is True:
            os.remove(wav_path)
    #"""
    txt_path = os.path.join(dataset_dir, "transcript_utf8.txt")
    text_list = list()
    for txt in read_txt(txt_path):
        if txt == "\n":
            continue
        name, sentence = txt.split(":")
        sentence = sentence.replace("\n", "")
        wav_filepath = os.path.join(results_folder, name )
        out_txt = wav_filepath + "|" + sentence + "\n"
        text_list.append(out_txt)

    max_n = len(text_list)
    test_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(text_list)
        idx = random.randint(9, int(n-1))
        txt = text_list.pop(idx)
        test_list.append(txt)

    max_n = len(text_list)
    val_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(text_list)
        idx = random.randint(9, int(n-1))
        txt = text_list.pop(idx)
        val_list.append(txt)

    write_txt(f"./filelists/TTS_jsut_train_{target_sr}.txt", text_list)
    write_txt(f"./filelists/TTS_jsut_val_{target_sr}.txt", val_list)
    write_txt(f"./filelists/TTS_jsut_test_{target_sr}.txt", test_list)
    
    return 0

def ita_preprocess(dataset_dir:str = "./path/to/ita_corpus", 
                    f0_method:str = "crepe",results_folder="./dataset_TTS/ita/", normalize=False):
    
    os.makedirs(results_folder, exist_ok=True)
    pitch_extractor = F0_extractor(samplerate=target_sr,
                                   hop_size=hop_size,
                                   device=device)
    
    folder_list = ["recitation", "emotion"]
    #"""
    for folder in folder_list:
        wav_dir = os.path.join(dataset_dir, folder)
        filelist = os.listdir(wav_dir)
        results_folder_dir = os.path.join(results_folder,folder)
        os.makedirs(results_folder_dir, exist_ok=True)
        
        for filename in tqdm(filelist):
            wav_path = os.path.join(wav_dir, filename)
            if normalize is True:
                out_path = wav_path.replace(".wav", "_norm.wav")
                audio_norm_process(wav_path, out_path)
                wav_path = out_path
            y, sr = torchaudio.load(wav_path)
            y_converted = torchaudio.functional.resample(waveform=y, orig_freq=sr, new_freq=target_sr)
            for idx in range(999):
                if str(idx).zfill(3) in filename:
                    break
                
            if folder == "recitation":  
                filename_out = "RECITATION324_" + str(idx).zfill(3) + ".wav"
            elif folder == "emotion":
                filename_out = "EMOTION100_"+ str(idx).zfill(3) + ".wav"
            else:
                print("ERROR. Check ITA corpus.")
                continue
            
            save_path = os.path.join(results_folder_dir, filename_out)
            torchaudio.save(save_path, y_converted, target_sr) 

            f0, vuv = pitch_extractor.compute_f0(path=save_path, f0_method=f0_method)
            torch.save(f0, save_path.replace(".wav", "_f0.pt") )
            torch.save(vuv, save_path.replace(".wav", "_vuv.pt") )
            
            if normalize is True:
                os.remove(wav_path)

    #"""
    txt_path = os.path.join(results_folder, "transcript_utf8.txt")
    text_list = list()
    for txt in read_txt(txt_path):
        if txt == "\n":
            continue
        name, sentence = txt.split(":")
        sentence, kana = sentence.split(",")
        sentence = sentence.replace("\n", "")
        if "RECITATION" in name:
            wav_filepath = os.path.join(results_folder,"recitation", name )
        elif "EMOTION" in name:
            wav_filepath = os.path.join(results_folder,"emotion", name )
            
        out_txt = wav_filepath + "|" + sentence + "\n"
        text_list.append(out_txt)

    max_n = len(text_list)
    test_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(text_list)
        idx = random.randint(9, int(n-1))
        txt = text_list.pop(idx)
        test_list.append(txt)

    max_n = len(text_list)
    val_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(text_list)
        idx = random.randint(9, int(n-1))
        txt = text_list.pop(idx)
        val_list.append(txt)

    write_txt(f"./filelists/TTS_ita_train_{target_sr}.txt", text_list)
    write_txt(f"./filelists/TTS_ita_val_{target_sr}.txt", val_list)
    write_txt(f"./filelists/TTS_ita_test_{target_sr}.txt", test_list)
    
    return 0

def homebrew_preprocess(dataset_dir:str = "./homebrew/", 
                    f0_method:str = "crepe",results_folder="./dataset_TTS/homebrew/", normalize=False ):
 
    os.makedirs(results_folder, exist_ok=True)
    pitch_extractor = F0_extractor(samplerate=target_sr,
                                   hop_size=hop_size,
                                   device=device)

    wav_dir = dataset_dir
    for filename in tqdm(os.listdir(wav_dir)):
        wav_path = os.path.join(wav_dir, filename)
        
        if normalize is True:
            out_path = wav_path.replace(".wav", "_norm.wav")
            audio_norm_process(wav_path, out_path)
            wav_path = out_path
        y, sr = torchaudio.load(wav_path)
        y_converted = torchaudio.functional.resample(waveform=y, orig_freq=sr, new_freq=target_sr)
        save_path = os.path.join(results_folder, filename)
        torchaudio.save(save_path, y_converted, target_sr) 

        f0, vuv = pitch_extractor.compute_f0(path=save_path, f0_method=f0_method)
        torch.save(f0, save_path.replace(".wav", "_f0.pt") )
        torch.save(vuv, save_path.replace(".wav", "_vuv.pt") )
        if normalize is True:
            os.remove(wav_path)
        
    txt_path = os.path.join(results_folder, "transcript_utf8.txt")
    text_list = list()
    for txt in read_txt(txt_path):
        if txt == "\n":
            continue
        name, sentence = txt.split(":")
        sentence = sentence.replace("\n", "")
        wav_filepath = os.path.join(results_folder, name)
        out_txt = wav_filepath + "|" + sentence + "\n"
        text_list.append(out_txt)

    max_n = len(text_list)
    test_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(text_list)
        idx = random.randint(9, int(n-1))
        txt = text_list.pop(idx)
        test_list.append(txt)

    max_n = len(text_list)
    val_list = list()
    for _ in range(int(max_n * split_ratio)):
        n = len(text_list)
        idx = random.randint(9, int(n-1))
        txt = text_list.pop(idx)
        val_list.append(txt)

    write_txt(f"./filelists/TTS_homebrew_train_{target_sr}.txt", text_list)
    write_txt(f"./filelists/TTS_homebrew_val_{target_sr}.txt", val_list)
    write_txt(f"./filelists/TTS_homebrew_test_{target_sr}.txt", test_list)

    return 0

def read_txt(path):
    with open(path, mode="r", encoding="utf-8")as f:
        lines = f.readlines()
    return lines

def write_txt(path, lines):
    with open(path, mode="w", encoding="utf-8")as f:
        f.writelines(lines)

import subprocess
def audio_norm_process(in_path, out_path):
    if os.path.isfile(in_path) is True:
        pass
    else:
        print("[ERROR] File is not existed : ", in_path)
        exit()
    cmd = ["ffmpeg-normalize", in_path,   "-o",  out_path]
    subprocess.run(cmd, encoding='utf-8', stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--dataset_name',
                        type=str,
                        #required=True, 
                        default="jsut",
                        help='jsut / ita / homebrew')
    parser.add_argument('--folder_path',
                        type=str,
                        #required=True, 
                        default="./jsut_ver1.1/basic5000/",
                        help='Path to corpus folder')
    parser.add_argument('--f0_method',
                        type=str,
                        #required=True,
                        default="crepe")
    parser.add_argument('--dataset_out_dir',
                        type=str,
                        #required=True,
                        default="./dataset_TTS/")
    parser.add_argument('--audio_normalize',
                        type=str,
                        #required=True,
                        default=True)

    
    args = parser.parse_args()
    dataset_out_dir = os.path.join(args.dataset_out_dir, args.dataset_name)

    if args.dataset_name == "jsut":
        jsut_preprocess(dataset_dir=args.folder_path, f0_method=args.f0_method, results_folder=dataset_out_dir, normalize=args.audio_normalize)
    elif args.dataset_name == "ita":
        ita_preprocess(dataset_dir=args.folder_path, f0_method=args.f0_method, results_folder=dataset_out_dir, normalize=args.audio_normalize)
    elif args.dataset_name == "homebrew":
        homebrew_preprocess(dataset_dir=args.folder_path, f0_method=args.f0_method, results_folder=dataset_out_dir, normalize=args.audio_normalize)
    else:
        print("ERROR. Check dataset_name.")

preprocess_tts_synthesizer.py
tonnetonne814/SiFi-VITS2-44100-Ja
import pdb, os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from librosa.util import normalize, pad_center, tiny
from scipy.signal import get_window

import logging

logger = logging.getLogger(__name__)


###stft codes from https://github.com/pseeth/torch-stft/blob/master/torch_stft/util.py
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 = normalize(win_sq, norm=norm) ** 2
    win_sq = 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


class STFT(torch.nn.Module):
    def __init__(
        self, filter_length=1024, hop_length=512, win_length=None, window="hann"
    ):
        """
        This module implements an STFT using 1D convolution and 1D transpose convolutions.
        This is a bit tricky so there are some cases that probably won't work as working
        out the same sizes before and after in all overlap add setups is tough. Right now,
        this code should work with hop lengths that are half the filter length (50% overlap
        between frames).

        Keyword Arguments:
            filter_length {int} -- Length of filters used (default: {1024})
            hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
            win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
                equals the filter length). (default: {None})
            window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
                (default: {'hann'})
        """
        super(STFT, self).__init__()
        self.filter_length = filter_length
        self.hop_length = hop_length
        self.win_length = win_length if win_length else filter_length
        self.window = window
        self.forward_transform = None
        self.pad_amount = int(self.filter_length / 2)
        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, :]
        )

        assert filter_length >= self.win_length
        # get window and zero center pad it to filter_length
        fft_window = get_window(window, self.win_length, fftbins=True)
        fft_window = pad_center(fft_window, size=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):
        """Take input data (audio) to STFT domain.

        Arguments:
            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)

        Returns:
            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
                num_frequencies, num_frames)
            phase {tensor} -- Phase of STFT with shape (num_batch,
                num_frequencies, num_frames)
        """
        num_batches = input_data.shape[0]
        num_samples = input_data.shape[-1]

        self.num_samples = num_samples

        # similar to librosa, reflect-pad the input
        input_data = input_data.view(num_batches, 1, num_samples)
        # print(1234,input_data.shape)
        input_data = F.pad(
            input_data.unsqueeze(1),
            (self.pad_amount, self.pad_amount, 0, 0, 0, 0),
            mode="reflect",
        ).squeeze(1)
        # print(2333,input_data.shape,self.forward_basis.shape,self.hop_length)
        # pdb.set_trace()
        forward_transform = F.conv1d(
            input_data, self.forward_basis, 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.atan2(imag_part.data, real_part.data)

        return magnitude  # , phase

    def inverse(self, magnitude, phase):
        """Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
        by the ```transform``` function.

        Arguments:
            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
                num_frequencies, num_frames)
            phase {tensor} -- Phase of STFT with shape (num_batch,
                num_frequencies, num_frames)

        Returns:
            inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
                shape (num_batch, num_samples)
        """
        recombine_magnitude_phase = torch.cat(
            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
        )

        inverse_transform = F.conv_transpose1d(
            recombine_magnitude_phase,
            self.inverse_basis,
            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.from_numpy(window_sum).to(inverse_transform.device)
            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[..., self.pad_amount :]
        inverse_transform = inverse_transform[..., : self.num_samples]
        inverse_transform = inverse_transform.squeeze(1)

        return inverse_transform

    def forward(self, input_data):
        """Take input data (audio) to STFT domain and then back to audio.

        Arguments:
            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)

        Returns:
            reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
                shape (num_batch, num_samples)
        """
        self.magnitude, self.phase = self.transform(input_data)
        reconstruction = self.inverse(self.magnitude, self.phase)
        return reconstruction


from time import time as ttime


class BiGRU(nn.Module):
    def __init__(self, input_features, hidden_features, num_layers):
        super(BiGRU, self).__init__()
        self.gru = nn.GRU(
            input_features,
            hidden_features,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
        )

    def forward(self, x):
        return self.gru(x)[0]


class ConvBlockRes(nn.Module):
    def __init__(self, in_channels, out_channels, momentum=0.01):
        super(ConvBlockRes, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=(3, 3),
                stride=(1, 1),
                padding=(1, 1),
                bias=False,
            ),
            nn.BatchNorm2d(out_channels, momentum=momentum),
            nn.ReLU(),
            nn.Conv2d(
                in_channels=out_channels,
                out_channels=out_channels,
                kernel_size=(3, 3),
                stride=(1, 1),
                padding=(1, 1),
                bias=False,
            ),
            nn.BatchNorm2d(out_channels, momentum=momentum),
            nn.ReLU(),
        )
        if in_channels != out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
            self.is_shortcut = True
        else:
            self.is_shortcut = False

    def forward(self, x):
        if self.is_shortcut:
            return self.conv(x) + self.shortcut(x)
        else:
            return self.conv(x) + x


class Encoder(nn.Module):
    def __init__(
        self,
        in_channels,
        in_size,
        n_encoders,
        kernel_size,
        n_blocks,
        out_channels=16,
        momentum=0.01,
    ):
        super(Encoder, self).__init__()
        self.n_encoders = n_encoders
        self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
        self.layers = nn.ModuleList()
        self.latent_channels = []
        for i in range(self.n_encoders):
            self.layers.append(
                ResEncoderBlock(
                    in_channels, out_channels, kernel_size, n_blocks, momentum=momentum
                )
            )
            self.latent_channels.append([out_channels, in_size])
            in_channels = out_channels
            out_channels *= 2
            in_size //= 2
        self.out_size = in_size
        self.out_channel = out_channels

    def forward(self, x):
        concat_tensors = []
        x = self.bn(x)
        for i in range(self.n_encoders):
            _, x = self.layers[i](x)
            concat_tensors.append(_)
        return x, concat_tensors


class ResEncoderBlock(nn.Module):
    def __init__(
        self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01
    ):
        super(ResEncoderBlock, self).__init__()
        self.n_blocks = n_blocks
        self.conv = nn.ModuleList()
        self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
        for i in range(n_blocks - 1):
            self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
        self.kernel_size = kernel_size
        if self.kernel_size is not None:
            self.pool = nn.AvgPool2d(kernel_size=kernel_size)

    def forward(self, x):
        for i in range(self.n_blocks):
            x = self.conv[i](x)
        if self.kernel_size is not None:
            return x, self.pool(x)
        else:
            return x


class Intermediate(nn.Module):  #
    def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
        super(Intermediate, self).__init__()
        self.n_inters = n_inters
        self.layers = nn.ModuleList()
        self.layers.append(
            ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)
        )
        for i in range(self.n_inters - 1):
            self.layers.append(
                ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)
            )

    def forward(self, x):
        for i in range(self.n_inters):
            x = self.layers[i](x)
        return x


class ResDecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
        super(ResDecoderBlock, self).__init__()
        out_padding = (0, 1) if stride == (1, 2) else (1, 1)
        self.n_blocks = n_blocks
        self.conv1 = nn.Sequential(
            nn.ConvTranspose2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=(3, 3),
                stride=stride,
                padding=(1, 1),
                output_padding=out_padding,
                bias=False,
            ),
            nn.BatchNorm2d(out_channels, momentum=momentum),
            nn.ReLU(),
        )
        self.conv2 = nn.ModuleList()
        self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
        for i in range(n_blocks - 1):
            self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))

    def forward(self, x, concat_tensor):
        x = self.conv1(x)
        x = torch.cat((x, concat_tensor), dim=1)
        for i in range(self.n_blocks):
            x = self.conv2[i](x)
        return x


class Decoder(nn.Module):
    def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList()
        self.n_decoders = n_decoders
        for i in range(self.n_decoders):
            out_channels = in_channels // 2
            self.layers.append(
                ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)
            )
            in_channels = out_channels

    def forward(self, x, concat_tensors):
        for i in range(self.n_decoders):
            x = self.layers[i](x, concat_tensors[-1 - i])
        return x


class DeepUnet(nn.Module):
    def __init__(
        self,
        kernel_size,
        n_blocks,
        en_de_layers=5,
        inter_layers=4,
        in_channels=1,
        en_out_channels=16,
    ):
        super(DeepUnet, self).__init__()
        self.encoder = Encoder(
            in_channels, 128, en_de_layers, kernel_size, n_blocks, en_out_channels
        )
        self.intermediate = Intermediate(
            self.encoder.out_channel // 2,
            self.encoder.out_channel,
            inter_layers,
            n_blocks,
        )
        self.decoder = Decoder(
            self.encoder.out_channel, en_de_layers, kernel_size, n_blocks
        )

    def forward(self, x):
        x, concat_tensors = self.encoder(x)
        x = self.intermediate(x)
        x = self.decoder(x, concat_tensors)
        return x


class E2E(nn.Module):
    def __init__(
        self,
        n_blocks,
        n_gru,
        kernel_size,
        en_de_layers=5,
        inter_layers=4,
        in_channels=1,
        en_out_channels=16,
    ):
        super(E2E, self).__init__()
        self.unet = DeepUnet(
            kernel_size,
            n_blocks,
            en_de_layers,
            inter_layers,
            in_channels,
            en_out_channels,
        )
        self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
        if n_gru:
            self.fc = nn.Sequential(
                BiGRU(3 * 128, 256, n_gru),
                nn.Linear(512, 360),
                nn.Dropout(0.25),
                nn.Sigmoid(),
            )
        else:
            self.fc = nn.Sequential(
                nn.Linear(3 * nn.N_MELS, nn.N_CLASS), nn.Dropout(0.25), nn.Sigmoid()
            )

    def forward(self, mel):
        # print(mel.shape)
        mel = mel.transpose(-1, -2).unsqueeze(1)
        x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
        x = self.fc(x)
        # print(x.shape)
        return x


from librosa.filters import mel


class MelSpectrogram(torch.nn.Module):
    def __init__(
        self,
        is_half,
        n_mel_channels,
        sampling_rate,
        win_length,
        hop_length,
        n_fft=None,
        mel_fmin=0,
        mel_fmax=None,
        clamp=1e-5,
    ):
        super().__init__()
        n_fft = win_length if n_fft is None else n_fft
        self.hann_window = {}
        mel_basis = mel(
            sr=sampling_rate,
            n_fft=n_fft,
            n_mels=n_mel_channels,
            fmin=mel_fmin,
            fmax=mel_fmax,
            htk=True,
        )
        mel_basis = torch.from_numpy(mel_basis).float()
        self.register_buffer("mel_basis", mel_basis)
        self.n_fft = win_length if n_fft is None else n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.sampling_rate = sampling_rate
        self.n_mel_channels = n_mel_channels
        self.clamp = clamp
        self.is_half = is_half

    def forward(self, audio, keyshift=0, speed=1, center=True):
        factor = 2 ** (keyshift / 12)
        n_fft_new = int(np.round(self.n_fft * factor))
        win_length_new = int(np.round(self.win_length * factor))
        hop_length_new = int(np.round(self.hop_length * speed))
        keyshift_key = str(keyshift) + "_" + str(audio.device)
        if keyshift_key not in self.hann_window:
            self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
                # "cpu"if(audio.device.type=="privateuseone") else audio.device
                audio.device
            )
        # fft = torch.stft(#doesn't support pytorch_dml
        #     # audio.cpu() if(audio.device.type=="privateuseone")else audio,
        #     audio,
        #     n_fft=n_fft_new,
        #     hop_length=hop_length_new,
        #     win_length=win_length_new,
        #     window=self.hann_window[keyshift_key],
        #     center=center,
        #     return_complex=True,
        # )
        # magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
        # print(1111111111)
        # print(222222222222222,audio.device,self.is_half)
        if hasattr(self, "stft") == False:
            # print(n_fft_new,hop_length_new,win_length_new,audio.shape)
            self.stft = STFT(
                filter_length=n_fft_new,
                hop_length=hop_length_new,
                win_length=win_length_new,
                window="hann",
            ).to(audio.device)
        magnitude = self.stft.transform(audio)  # phase
        # if (audio.device.type == "privateuseone"):
        #     magnitude=magnitude.to(audio.device)
        if keyshift != 0:
            size = self.n_fft // 2 + 1
            resize = magnitude.size(1)
            if resize < size:
                magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
        mel_output = torch.matmul(self.mel_basis, magnitude)
        if self.is_half == True:
            mel_output = mel_output.half()
        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
        # print(log_mel_spec.device.type)
        return log_mel_spec


class RMVPE:
    def __init__(self, model_path, 
                 is_half,
                 sampling_rate, 
                 hop_length,
                 win_size, 
                 f0_mel_min,
                 f0_mel_max,device=None):
        self.resample_kernel = {}
        self.resample_kernel = {}
        self.is_half = is_half
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        self.mel_extractor = MelSpectrogram(
            is_half, 128, sampling_rate, win_size, hop_length, win_size, f0_mel_min, f0_mel_max
        ).to(device)
        if "privateuseone" in str(device):
            import onnxruntime as ort

            ort_session = ort.InferenceSession(
                "%s/rmvpe.onnx" % os.environ["rmvpe_root"],
                providers=["DmlExecutionProvider"],
            )
            self.model = ort_session
        else:
            model = E2E(4, 1, (2, 2))
            ckpt = torch.load(model_path, map_location="cpu")
            model.load_state_dict(ckpt)
            model.eval()
            if is_half == True:
                model = model.half()
            self.model = model
            self.model = self.model.to(device)
        cents_mapping = 20 * np.arange(360) + 1997.3794084376191
        self.cents_mapping = np.pad(cents_mapping, (4, 4))  # 368

    def mel2hidden(self, mel):
        with torch.no_grad():
            n_frames = mel.shape[-1]
            mel = F.pad(
                mel, (0, 32 * ((n_frames - 1) // 32 + 1) - n_frames), mode="constant"
            )
            if "privateuseone" in str(self.device):
                onnx_input_name = self.model.get_inputs()[0].name
                onnx_outputs_names = self.model.get_outputs()[0].name
                hidden = self.model.run(
                    [onnx_outputs_names],
                    input_feed={onnx_input_name: mel.cpu().numpy()},
                )[0]
            else:
                hidden = self.model(mel)
            return hidden[:, :n_frames]

    def decode(self, hidden, thred=0.03):
        cents_pred = self.to_local_average_cents(hidden, thred=thred)
        f0 = 10 * (2 ** (cents_pred / 1200))
        f0[f0 == 10] = 0
        # f0 = np.array([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred])
        return f0

    def infer_from_audio(self, audio, thred=0.03):
        # torch.cuda.synchronize()
        t0 = ttime()
        mel = self.mel_extractor(
            torch.from_numpy(audio).float().to(self.device).unsqueeze(0), center=True
        )
        # print(123123123,mel.device.type)
        # torch.cuda.synchronize()
        t1 = ttime()
        hidden = self.mel2hidden(mel)
        # torch.cuda.synchronize()
        t2 = ttime()
        # print(234234,hidden.device.type)
        if "privateuseone" not in str(self.device):
            hidden = hidden.squeeze(0).cpu().numpy()
        else:
            hidden = hidden[0]
        if self.is_half == True:
            hidden = hidden.astype("float32")

        f0 = self.decode(hidden, thred=thred)
        # torch.cuda.synchronize()
        t3 = ttime()
        # print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
        return f0

    def to_local_average_cents(self, salience, thred=0.05):
        # t0 = ttime()
        center = np.argmax(salience, axis=1)  # 帧长#index
        salience = np.pad(salience, ((0, 0), (4, 4)))  # 帧长,368
        # t1 = ttime()
        center += 4
        todo_salience = []
        todo_cents_mapping = []
        starts = center - 4
        ends = center + 5
        for idx in range(salience.shape[0]):
            todo_salience.append(salience[:, starts[idx] : ends[idx]][idx])
            todo_cents_mapping.append(self.cents_mapping[starts[idx] : ends[idx]])
        # t2 = ttime()
        todo_salience = np.array(todo_salience)  # 帧长，9
        todo_cents_mapping = np.array(todo_cents_mapping)  # 帧长，9
        product_sum = np.sum(todo_salience * todo_cents_mapping, 1)
        weight_sum = np.sum(todo_salience, 1)  # 帧长
        devided = product_sum / weight_sum  # 帧长
        # t3 = ttime()
        maxx = np.max(salience, axis=1)  # 帧长
        devided[maxx <= thred] = 0
        # t4 = ttime()
        # print("decode:%s\t%s\t%s\t%s" % (t1 - t0, t2 - t1, t3 - t2, t4 - t3))
        return devided


if __name__ == "__main__":
    import librosa
    import soundfile as sf

    audio, sampling_rate = sf.read(r"C:\Users\liujing04\Desktop\Z\冬之花clip1.wav")
    if len(audio.shape) > 1:
        audio = librosa.to_mono(audio.transpose(1, 0))
    audio_bak = audio.copy()
    if sampling_rate != 16000:
        audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000)
    model_path = r"D:\BaiduNetdiskDownload\RVC-beta-v2-0727AMD_realtime\rmvpe.pt"
    thred = 0.03  # 0.01
    device = "cuda" if torch.cuda.is_available() else "cpu"
    rmvpe = RMVPE(model_path, is_half=False, device=device)
    t0 = ttime()
    f0 = rmvpe.infer_from_audio(audio, thred=thred)
    # f0 = rmvpe.infer_from_audio(audio, thred=thred)
    # f0 = rmvpe.infer_from_audio(audio, thred=thred)
    # f0 = rmvpe.infer_from_audio(audio, thred=thred)
    # f0 = rmvpe.infer_from_audio(audio, thred=thred)
    t1 = ttime()
    logger.info("%s %.2f", f0.shape, t1 - t0)

rmvpe.py
tonnetonne814/SiFi-VITS2-44100-Ja
import torch 
import yaml

def main3():
    from sifigan.models.generator import SiFiGANGenerator
    from sifigan.models.discriminator import UnivNetMultiResolutionMultiPeriodDiscriminator
    from sifigan.utils import dilated_factor        
    from sifigan.utils.features import SignalGenerator
    
    config_path = "./configs/sifigan.yaml"
    with open(config_path, mode="r", encoding="utf-8") as f:
      hps = yaml.safe_load(f)
    
    # model load
    model = SiFiGANGenerator(**hps)

    # data prep
    B = 2
    T = 97
    hop_size = 512 
    dense_factors= [0.5, 1, 4, 8, 16]   
    df_f0_type = "cf0"
    sr = 44100
    in_channels = 43 # features for recon

    out_size = (B, 1, int(hop_size*T)) # answer

    x = torch.rand(size=(B, 1, int(hop_size*T)), dtype=torch.float32) 
    c = torch.rand(size=(B, in_channels, T), dtype=torch.float32)
    f0 = torch.rand(size=(B, 1, T), dtype=torch.float32)
    g  = torch.zeros(size=(B,256, 1))
    cf0 = torch.rand(size=(B, 1, T), dtype=torch.float32)
    singlef0 = torch.rand(size=(T,), dtype=torch.float32)
    single_x = torch.rand(size=(int(hop_size*T),), dtype=torch.float32)
    single_c = torch.rand(size=(in_channels, T), dtype=torch.float32)

    single_vuv, single_cf0, flag = get_continuos_f0(singlef0)# from dataset config
    signal_types= ["sine"]              # from dataset config
    noise_amp= 0.003                    # from dataset config
    sine_amp= 0.1                    # Sine amplitude.# from dataset config

    signal_generator = SignalGenerator(
            sample_rate=sr,
            hop_size=hop_size,
            sine_amp=sine_amp,
            noise_amp=noise_amp,
            signal_types=signal_types,
        )
    single_x_check = signal_generator(cf0)

    upsample_scales= hps["upsample_scales"]
    prod_upsample_scales = np.cumprod(upsample_scales)
    df_sample_rates = [sr / hop_size * s for s in prod_upsample_scales]
    dfs = []
    for df, us in zip(dense_factors, prod_upsample_scales):
        dfs += [
            np.repeat(dilated_factor(single_cf0, sr, df), us)
            if df_f0_type == "cf0"
            else np.repeat(dilated_factor(single_cf0, sr, df), us)
        ]
    for i in range(len(dense_factors)):
        assert len(single_x) * df_sample_rates[i] == len(dfs[i]) * sr

    dfs_batch = [[] for _ in range(len(dense_factors))]
    for i in range(len(dense_factors)):
        dfs_batch[i] += [
            dfs[i].astype(np.float32).reshape(-1, 1)
        ]  # list = [(T', 1), ...]
    
    for i in range(len(dense_factors)):
        dfs_batch[i] = torch.FloatTensor(np.array(dfs_batch[i])).transpose(
                2, 1
            )  # (B, 1, T')
    Voice, Exicitation = model(x=x, c=c, d=dfs_batch, g=g)
    print(Voice.shape)
    
    config_path = "./configs/univnet.yaml"
    with open(config_path, mode="r", encoding="utf-8") as f:
      hps_d = yaml.safe_load(f)
    discriminator = UnivNetMultiResolutionMultiPeriodDiscriminator(**hps_d)
    out = discriminator(Voice)

    return 0


import copy
import numpy as np
from scipy.interpolate import interp1d
def get_continuos_f0(f0):
    """Convert F0 to continuous F0
    Args:
        f0 (ndarray): original f0 sequence with the shape (T)
    Return:
        (ndarray): continuous f0 with the shape (T)
    """
    # get uv information as binary
    uv = np.float32(f0 != 0)
    # get start and end of f0
    start_f0 = f0[f0 != 0][0]
    end_f0 = f0[f0 != 0][-1]
    # padding start and end of f0 sequence
    cf0 = copy.deepcopy(f0)
    start_idx = np.where(cf0 == start_f0)[0][0]
    end_idx = np.where(cf0 == end_f0)[0][-1]
    cf0[:start_idx] = start_f0
    cf0[end_idx:] = end_f0
    # get non-zero frame index
    nz_frames = np.where(cf0 != 0)[0]
    # perform linear interpolation
    f = interp1d(nz_frames, cf0[nz_frames])
    cf0 = f(np.arange(0, cf0.shape[0]))
    return uv, cf0, True


if __name__ == "__main__":
    main()
sifi_gan_test.py
tonnetonne814/SiFi-VITS2-44100-Ja
from sifigan.layers.cheaptrick import *  # NOQA
from sifigan.layers.snake import *  # NOQA
from sifigan.layers.residual_block import *  # NOQA

sifigan/layers/__init__.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Spectral envelopes estimation module based on CheapTrick.

References:
    - https://www.sciencedirect.com/science/article/pii/S0167639314000697
    - https://github.com/mmorise/World

"""

import math

import torch
import torch.fft
import torch.nn as nn


class AdaptiveWindowing(nn.Module):
    """CheapTrick F0 adptive windowing module."""

    def __init__(
        self,
        sample_rate,
        hop_size,
        fft_size,
        f0_floor,
        f0_ceil,
    ):
        """Initilize AdaptiveWindowing module.

        Args:
            sample_rate (int): Sampling rate.
            hop_size (int): Hop size.
            fft_size (int): FFT size.
            f0_floor (int): Minimum value of F0.
            f0_ceil (int): Maximum value of F0.

        """
        super(AdaptiveWindowing, self).__init__()

        self.sample_rate = sample_rate
        self.hop_size = hop_size
        self.fft_size = fft_size
        self.register_buffer("window", torch.zeros((f0_ceil + 1, fft_size)))
        self.zero_padding = nn.ConstantPad2d((fft_size // 2, fft_size // 2, 0, 0), 0)

        # Pre-calculation of the window functions
        for f0 in range(f0_floor, f0_ceil + 1):
            half_win_len = round(1.5 * self.sample_rate / f0)
            base_index = torch.arange(
                -half_win_len, half_win_len + 1, dtype=torch.int64
            )
            position = base_index / 1.5 / self.sample_rate
            left = fft_size // 2 - half_win_len
            right = fft_size // 2 + half_win_len + 1
            window = torch.zeros(fft_size)
            window[left:right] = 0.5 * torch.cos(math.pi * position * f0) + 0.5
            average = torch.sum(window * window).pow(0.5)
            self.window[f0] = window / average

    def forward(self, x, f, power=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Waveform (B, fft_size // 2 + 1, T).
            f (Tensor): F0 sequence (B, T').
            power (boot): Whether to use power or magnitude.

        Returns:
            Tensor: Power spectrogram (B, bin_size, T').

        """
        # Get the matrix of window functions corresponding to F0
        x = self.zero_padding(x).unfold(1, self.fft_size, self.hop_size)
        windows = self.window[f]

        # Adaptive windowing and calculate power spectrogram.
        # In test, change x[:, : -1, :] to x.
        x = torch.abs(torch.fft.rfft(x[:, :-1, :] * windows))
        x = x.pow(2) if power else x

        return x


class AdaptiveLiftering(nn.Module):
    """CheapTrick F0 adptive windowing module."""

    def __init__(
        self,
        sample_rate,
        fft_size,
        f0_floor,
        f0_ceil,
        q1=-0.15,
    ):
        """Initilize AdaptiveLiftering module.

        Args:
            sample_rate (int): Sampling rate.
            fft_size (int): FFT size.
            f0_floor (int): Minimum value of F0.
            f0_ceil (int): Maximum value of F0.
            q1 (float): Parameter to remove effect of adjacent harmonics.

        """
        super(AdaptiveLiftering, self).__init__()

        self.sample_rate = sample_rate
        self.bin_size = fft_size // 2 + 1
        self.q1 = q1
        self.q0 = 1.0 - 2.0 * q1
        self.register_buffer(
            "smoothing_lifter", torch.zeros((f0_ceil + 1, self.bin_size))
        )
        self.register_buffer(
            "compensation_lifter", torch.zeros((f0_ceil + 1, self.bin_size))
        )

        # Pre-calculation of the smoothing lifters and compensation lifters
        for f0 in range(f0_floor, f0_ceil + 1):
            smoothing_lifter = torch.zeros(self.bin_size)
            compensation_lifter = torch.zeros(self.bin_size)
            quefrency = torch.arange(1, self.bin_size) / sample_rate
            smoothing_lifter[0] = 1.0
            smoothing_lifter[1:] = torch.sin(math.pi * f0 * quefrency) / (
                math.pi * f0 * quefrency
            )
            compensation_lifter[0] = self.q0 + 2.0 * self.q1
            compensation_lifter[1:] = self.q0 + 2.0 * self.q1 * torch.cos(
                2.0 * math.pi * f0 * quefrency
            )
            self.smoothing_lifter[f0] = smoothing_lifter
            self.compensation_lifter[f0] = compensation_lifter

    def forward(self, x, f, elim_0th=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Power spectrogram (B, bin_size, T').
            f (Tensor): F0 sequence (B, T').
            elim_0th (bool): Whether to eliminate cepstram 0th component.

        Returns:
            Tensor: Estimated spectral envelope (B, bin_size, T').

        """
        # Setting the smoothing lifter and compensation lifter
        smoothing_lifter = self.smoothing_lifter[f]
        compensation_lifter = self.compensation_lifter[f]

        # Calculating cepstrum
        tmp = torch.cat((x, torch.flip(x[:, :, 1:-1], [2])), dim=2)
        cepstrum = torch.fft.rfft(torch.log(torch.clamp(tmp, min=1e-7))).real

        # Set the 0th cepstrum to 0
        if elim_0th:
            cepstrum[..., 0] = 0

        # Liftering cepstrum with the lifters
        liftered_cepstrum = cepstrum * smoothing_lifter * compensation_lifter

        # Return the result to the spectral domain
        x = torch.fft.irfft(liftered_cepstrum)[:, :, : self.bin_size]

        return x


class CheapTrick(nn.Module):
    """CheapTrick based spectral envelope estimation module."""

    def __init__(
        self,
        sample_rate,
        hop_size,
        fft_size,
        f0_floor=70,
        f0_ceil=340,
        uv_threshold=0,
        q1=-0.15,
    ):
        """Initilize AdaptiveLiftering module.

        Args:
            sample_rate (int): Sampling rate.
            hop_size (int): Hop size.
            fft_size (int): FFT size.
            f0_floor (int): Minimum value of F0.
            f0_ceil (int): Maximum value of F0.
            uv_threshold (float): V/UV determining threshold.
            q1 (float): Parameter to remove effect of adjacent harmonics.

        """
        super(CheapTrick, self).__init__()

        # fft_size must be larger than 3.0 * sample_rate / f0_floor
        assert fft_size > 3.0 * sample_rate / f0_floor
        self.f0_floor = f0_floor
        self.f0_ceil = f0_ceil
        self.uv_threshold = uv_threshold

        self.ada_wind = AdaptiveWindowing(
            sample_rate,
            hop_size,
            fft_size,
            f0_floor,
            f0_ceil,
        )
        self.ada_lift = AdaptiveLiftering(
            sample_rate,
            fft_size,
            f0_floor,
            f0_ceil,
            q1,
        )

    def forward(self, x, f, power=False, elim_0th=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Power spectrogram (B, T).
            f (Tensor): F0 sequence (B, T').
            power (boot): Whether to use power or magnitude spectrogram.
            elim_0th (bool): Whether to eliminate cepstram 0th component.

        Returns:
            Tensor: Estimated spectral envelope (B, bin_size, T').

        """
        # Step0: Round F0 values to integers.
        voiced = (f > self.uv_threshold) * torch.ones_like(f)
        f = voiced * f + (1.0 - voiced) * self.f0_ceil
        f = torch.round(torch.clamp(f, min=self.f0_floor, max=self.f0_ceil)).to(
            torch.int64
        )

        # Step1: Adaptive windowing and calculate power or amplitude spectrogram.
        x = self.ada_wind(x, f, power)

        # Step3: Smoothing (log axis) and spectral recovery on the cepstrum domain.
        x = self.ada_lift(x, f, elim_0th)

        return x

sifigan/layers/cheaptrick.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Residual block modules.

References:
    - https://github.com/kan-bayashi/ParallelWaveGAN
    - https://github.com/bigpon/QPPWG
    - https://github.com/r9y9/wavenet_vocoder

"""

from logging import getLogger

import torch
import torch.nn as nn
from sifigan.layers import Snake
from sifigan.utils import index_initial, pd_indexing

# A logger for this file
logger = getLogger(__name__)


class Conv1d(nn.Conv1d):
    """Conv1d module with customized initialization."""

    def __init__(self, *args, **kwargs):
        """Initialize Conv1d module."""
        super(Conv1d, self).__init__(*args, **kwargs)

    def reset_parameters(self):
        """Reset parameters."""
        nn.init.kaiming_normal_(self.weight, nonlinearity="relu")
        if self.bias is not None:
            nn.init.constant_(self.bias, 0.0)


class Conv1d1x1(Conv1d):
    """1x1 Conv1d with customized initialization."""

    def __init__(self, in_channels, out_channels, bias=True):
        """Initialize 1x1 Conv1d module."""
        super(Conv1d1x1, self).__init__(
            in_channels, out_channels, kernel_size=1, padding=0, dilation=1, bias=bias
        )


class Conv2d(nn.Conv2d):
    """Conv2d module with customized initialization."""

    def __init__(self, *args, **kwargs):
        """Initialize Conv2d module."""
        super(Conv2d, self).__init__(*args, **kwargs)

    def reset_parameters(self):
        """Reset parameters."""
        nn.init.kaiming_normal_(self.weight, mode="fan_out", nonlinearity="relu")
        if self.bias is not None:
            nn.init.constant_(self.bias, 0.0)


class Conv2d1x1(Conv2d):
    """1x1 Conv2d with customized initialization."""

    def __init__(self, in_channels, out_channels, bias=True):
        """Initialize 1x1 Conv2d module."""
        super(Conv2d1x1, self).__init__(
            in_channels, out_channels, kernel_size=1, padding=0, dilation=1, bias=bias
        )


class ResidualBlock(nn.Module):
    """Residual block module in HiFiGAN."""

    def __init__(
        self,
        kernel_size=3,
        channels=512,
        dilations=(1, 3, 5),
        bias=True,
        use_additional_convs=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
    ):
        """Initialize ResidualBlock module.

        Args:
            kernel_size (int): Kernel size of dilation convolution layer.
            channels (int): Number of channels for convolution layer.
            dilations (List[int]): List of dilation factors.
            use_additional_convs (bool): Whether to use additional convolution layers.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.

        """
        super().__init__()
        self.use_additional_convs = use_additional_convs
        self.convs1 = nn.ModuleList()
        if use_additional_convs:
            self.convs2 = nn.ModuleList()
        assert kernel_size % 2 == 1, "Kernel size must be odd number."
        for dilation in dilations:
            if nonlinear_activation == "Snake":
                nonlinear = Snake(channels, **nonlinear_activation_params)
            else:
                nonlinear = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
            self.convs1 += [
                nn.Sequential(
                    nonlinear,
                    nn.Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        dilation=dilation,
                        bias=bias,
                        padding=(kernel_size - 1) // 2 * dilation,
                    ),
                )
            ]
            if use_additional_convs:
                if nonlinear_activation == "Snake":
                    nonlinear = Snake(channels, **nonlinear_activation_params)
                else:
                    nonlinear = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
                self.convs2 += [
                    nn.Sequential(
                        nonlinear,
                        nn.Conv1d(
                            channels,
                            channels,
                            kernel_size,
                            dilation=1,
                            bias=bias,
                            padding=(kernel_size - 1) // 2,
                        ),
                    )
                ]

    def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, channels, T).

        Returns:
            Tensor: Output tensor (B, channels, T).

        """
        for idx in range(len(self.convs1)):
            xt = self.convs1[idx](x)
            if self.use_additional_convs:
                xt = self.convs2[idx](xt)
            x = xt + x
        return x


class AdaptiveResidualBlock(nn.Module):
    """Residual block module in HiFiGAN."""

    def __init__(
        self,
        kernel_size=3,
        channels=512,
        dilations=(1, 2, 4),
        bias=True,
        use_additional_convs=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
    ):
        """Initialize ResidualBlock module.

        Args:
            kernel_size (int): Kernel size of dilation convolution layer.
            channels (int): Number of channels for convolution layer.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.

        """
        super().__init__()
        self.use_additional_convs = use_additional_convs
        assert kernel_size == 3, "Currently only kernel_size = 3 is supported."
        self.channels = channels
        self.dilations = dilations
        self.nonlinears = nn.ModuleList()
        self.convsC = nn.ModuleList()
        self.convsP = nn.ModuleList()
        self.convsF = nn.ModuleList()
        if use_additional_convs:
            self.convsA = nn.ModuleList()
        for _ in dilations:
            if nonlinear_activation == "Snake":
                self.nonlinears += [Snake(channels, **nonlinear_activation_params)]
            else:
                self.nonlinears += [getattr(nn, nonlinear_activation)(**nonlinear_activation_params)]
            self.convsC += [
                Conv1d1x1(
                    channels,
                    channels,
                    bias=bias,
                ),
            ]
            self.convsP += [
                Conv1d1x1(
                    channels,
                    channels,
                    bias=bias,
                ),
            ]
            self.convsF += [
                Conv1d1x1(
                    channels,
                    channels,
                    bias=bias,
                ),
            ]
            if use_additional_convs:
                if nonlinear_activation == "Snake":
                    nonlinear = Snake(channels, **nonlinear_activation_params)
                else:
                    nonlinear = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
                self.convsA += [
                    nn.Sequential(
                        nonlinear,
                        nn.Conv1d(
                            channels,
                            channels,
                            kernel_size,
                            dilation=1,
                            bias=bias,
                            padding=(kernel_size - 1) // 2,
                        ),
                    )
                ]

    def forward(self, x, d):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, channels, T).
            d (Tensor): Input pitch-dependent dilated factors (B, 1, T).

        Returns:
            Tensor: Output tensor (B, channels, T).

        """
        batch_index, ch_index = index_initial(x.size(0), self.channels)
        for i, dilation in enumerate(self.dilations):
            xt = self.nonlinears[i](x)
            xP, xF = pd_indexing(xt, d, dilation, batch_index, ch_index)
            xt = self.convsC[i](xt) + self.convsP[i](xP) + self.convsF[i](xF)
            if self.use_additional_convs:
                xt = self.convsA[i](xt)
            x = xt + x
        return x

sifigan/layers/residual_block.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Snake Activation Function Module.

References:
    - Neural Networks Fail to Learn Periodic Functions and How to Fix It
        https://arxiv.org/pdf/2006.08195.pdf
    - BigVGAN: A Universal Neural Vocoder with Large-Scale Training
        https://arxiv.org/pdf/2206.04658.pdf

"""

import torch
import torch.nn as nn


class Snake(nn.Module):
    """Snake activation function module."""

    def __init__(self, channels, init=50):
        """Initialize Snake module.

        Args:
            channels (int): Number of feature channels.
            init (float): Initial value of the learnable parameter alpha.
                          According to the original paper, 5 ~ 50 would be
                          suitable for periodic data (i.e. voices).

        """
        super(Snake, self).__init__()
        alpha = init * torch.ones(1, channels, 1)
        self.alpha = nn.Parameter(alpha)

    def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, channels, T).

        Returns:
            Tensor: Output tensor (B, channels, T).

        """
        return x + torch.sin(self.alpha * x) ** 2 / self.alpha

sifigan/layers/snake.py
tonnetonne814/SiFi-VITS2-44100-Ja
from sifigan.losses.adv import *  # NOQA
from sifigan.losses.mel import *  # NOQA
from sifigan.losses.reg import *  # NOQA

sifigan/losses/__init__.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Adversarial loss modules.

References:
    - https://github.com/kan-bayashi/ParallelWaveGAN

"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class AdversarialLoss(nn.Module):
    """Adversarial loss module."""

    def __init__(
        self,
        average_by_discriminators=False,
        loss_type="mse",
    ):
        """Initialize AversarialLoss module."""
        super(AdversarialLoss, self).__init__()
        self.average_by_discriminators = average_by_discriminators

        assert loss_type in ["mse", "hinge"], f"{loss_type} is not supported."
        if loss_type == "mse":
            self.adv_criterion = self._mse_adv_loss
            self.fake_criterion = self._mse_fake_loss
            self.real_criterion = self._mse_real_loss
        else:
            self.adv_criterion = self._hinge_adv_loss
            self.fake_criterion = self._hinge_fake_loss
            self.real_criterion = self._hinge_real_loss

    def forward(self, p_fakes, p_reals=None):
        """Calcualate generator/discriminator adversarial loss.

        Args:
            p_fakes (list): List of
                discriminator outputs calculated from generator outputs.
            p_reals (list): List of
                discriminator outputs calculated from groundtruth.

        Returns:
            Tensor: Generator adversarial loss value.
            Tensor: Discriminator real loss value.
            Tensor: Discriminator fake loss value.

        """
        # generator adversarial loss
        if p_reals is None:
            adv_loss = 0.0
            for p_fake in p_fakes:
                adv_loss += self.adv_criterion(p_fake)

            if self.average_by_discriminators:
                adv_loss /= len(p_fakes)

            return adv_loss

        # discriminator adversarial loss
        else:
            fake_loss = 0.0
            real_loss = 0.0
            for p_fake, p_real in zip(p_fakes, p_reals):
                fake_loss += self.fake_criterion(p_fake)
                real_loss += self.real_criterion(p_real)

            if self.average_by_discriminators:
                fake_loss /= len(p_fakes)
                real_loss /= len(p_reals)

            return fake_loss, real_loss

    def _mse_adv_loss(self, x):
        return F.mse_loss(x, x.new_ones(x.size()))

    def _mse_real_loss(self, x):
        return F.mse_loss(x, x.new_ones(x.size()))

    def _mse_fake_loss(self, x):
        return F.mse_loss(x, x.new_zeros(x.size()))

    def _hinge_adv_loss(self, x):
        return -x.mean()

    def _hinge_real_loss(self, x):
        return -torch.mean(torch.min(x - 1, x.new_zeros(x.size())))

    def _hinge_fake_loss(self, x):
        return -torch.mean(torch.min(-x - 1, x.new_zeros(x.size())))


class FeatureMatchLoss(nn.Module):
    # Feature matching loss module.

    def __init__(
        self,
        average_by_layers=False,
    ):
        """Initialize FeatureMatchLoss module."""
        super(FeatureMatchLoss, self).__init__()
        self.average_by_layers = average_by_layers

    def forward(self, fmaps_fake, fmaps_real):
        """Calculate forward propagation.

        Args:
            fmaps_fake (list): List of discriminator outputs
                calcuated from generater outputs.
            fmaps_real (list): List of discriminator outputs
                calcuated from groundtruth.

        Returns:
            Tensor: Feature matching loss value.

        """
        fm_loss = 0.0
        for feat_fake, feat_real in zip(fmaps_fake, fmaps_real):
            fm_loss += F.l1_loss(feat_fake, feat_real.detach())

        if self.average_by_layers:
            fm_loss /= len(fmaps_fake)

        return fm_loss

sifigan/losses/adv.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""STFT-based loss modules.

References:
    - https://github.com/kan-bayashi/ParallelWaveGAN

"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from librosa.filters import mel as librosa_mel


def stft(
    x, fft_size, hop_size, win_length, window, center=True, onesided=True, power=False
):
    """Perform STFT and convert to magnitude spectrogram.

    Args:
        x (Tensor): Input signal tensor (B, T).
        fft_size (int): FFT size.
        hop_size (int): Hop size.
        win_length (int): Window length.
        window (str): Window function type.

    Returns:
        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).

    """
    x_stft = torch.stft(
        x,
        fft_size,
        hop_size,
        win_length,
        window,
        center=center,
        onesided=onesided,
        return_complex=False,
    )
    real = x_stft[..., 0]
    imag = x_stft[..., 1]

    if power:
        return torch.clamp(real**2 + imag**2, min=1e-7).transpose(2, 1)
    else:
        return torch.sqrt(torch.clamp(real**2 + imag**2, min=1e-7)).transpose(2, 1)


class MelSpectralLoss(nn.Module):
    """Mel-spectral L1 loss module."""

    def __init__(
        self,
        fft_size=1024,
        hop_size=120,
        win_length=1024,
        window="hann_window",
        sample_rate=24000,
        n_mels=80,
        fmin=0,
        fmax=None,
    ):
        """Initialize MelSpectralLoss loss.

        Args:
            fft_size (int): FFT points.
            hop_length (int): Hop length.
            win_length (Optional[int]): Window length.
            window (str): Window type.
            sample_rate (int): Sampling rate.
            n_mels (int): Number of Mel basis.
            fmin (Optional[int]): Minimum frequency of mel-filter-bank.
            fmax (Optional[int]): Maximum frequency of mel-filter-bank.

        """
        super().__init__()
        self.fft_size = fft_size
        self.hop_size = hop_size
        self.win_length = win_length if win_length is not None else fft_size
        self.register_buffer("window", getattr(torch, window)(self.win_length))
        self.sample_rate = sample_rate
        self.n_mels = n_mels
        self.fmin = fmin
        self.fmax = fmax if fmax is not None else sample_rate / 2
        melmat = librosa_mel(
            sr=sample_rate, n_fft=fft_size, n_mels=n_mels, fmin=fmin, fmax=fmax
        ).T
        self.register_buffer("melmat", torch.from_numpy(melmat).float())

    def forward(self, x, y):
        """Calculate Mel-spectral L1 loss.

        Args:
            x (Tensor): Generated waveform tensor (B, 1, T).
            y (Tensor): Groundtruth waveform tensor (B, 1, T).

        Returns:
            Tensor: Mel-spectral L1 loss value.

        """
        x = x.squeeze(1)
        y = y.squeeze(1)
        x_mag = stft(x, self.fft_size, self.hop_size, self.win_length, self.window)
        y_mag = stft(y, self.fft_size, self.hop_size, self.win_length, self.window)
        x_log_mel = torch.log(torch.clamp(torch.matmul(x_mag, self.melmat), min=1e-7))
        y_log_mel = torch.log(torch.clamp(torch.matmul(y_mag, self.melmat), min=1e-7))
        mel_loss = F.l1_loss(x_log_mel, y_log_mel)

        return mel_loss

sifigan/losses/mel.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Source regularization loss modules."""

import sifigan.losses
import torch
import torch.nn as nn
import torch.nn.functional as F
from librosa.filters import mel as librosa_mel
from sifigan.layers import CheapTrick


class FlattenLoss(nn.Module):
    """The regularization loss of uSFGAN."""

    def __init__(
        self,
        sample_rate=24000,
        hop_size=120,
        fft_size=2048,
        f0_floor=70,
        f0_ceil=340,
        power=False,
        elim_0th=False,
        l2_norm=False,
    ):
        """Initialize spectral envelope regularlization loss module.

        Args:
            sample_rate (int): Sampling rate.
            hop_size (int): Hop size.
            fft_size (int): FFT size.
            f0_floor (int): Minimum F0 value.
            f0_ceil (int): Maximum F0 value.
            power (bool): Whether to use power or magnitude spectrogram.
            elim_0th (bool): Whether to exclude 0th cepstrum in CheapTrick.
                If set to true, power is estimated by source-network.
            l2_norm (bool): Whether to regularize the spectral envelopes with L2 norm.

        """
        super(FlattenLoss, self).__init__()
        self.hop_size = hop_size
        self.power = power
        self.elim_0th = elim_0th
        self.cheaptrick = CheapTrick(
            sample_rate=sample_rate,
            hop_size=hop_size,
            fft_size=fft_size,
            f0_floor=f0_floor,
            f0_ceil=f0_ceil,
        )
        if l2_norm:
            self.loss = nn.MSELoss()
        else:
            self.loss = nn.L1Loss()

    def forward(self, s, f):
        """Calculate forward propagation.

        Args:
            s (Tensor): Predicted source excitation signal (B, 1, T).
            f (Tensor): F0 sequence (B, 1, T // hop_size).

        Returns:
            loss (Tensor): Loss value.

        """
        s, f = s.squeeze(1), f.squeeze(1)
        e = self.cheaptrick.forward(s, f, self.power, self.elim_0th)
        loss = self.loss(e, e.new_zeros(e.size()))
        return loss


class ResidualLoss(nn.Module):
    """The regularization loss of hn-uSFGAN."""

    def __init__(
        self,
        sample_rate=24000,
        fft_size=2048,
        hop_size=120,
        f0_floor=100,
        f0_ceil=840,
        n_mels=80,
        fmin=0,
        fmax=None,
        power=False,
        elim_0th=True,
    ):
        """Initialize ResidualLoss module.

        Args:
            sample_rate (int): Sampling rate.
            fft_size (int): FFT size.
            hop_size (int): Hop size.
            f0_floor (int): Minimum F0 value.
            f0_ceil (int): Maximum F0 value.
            n_mels (int): Number of Mel basis.
            fmin (int): Minimum frequency for Mel.
            fmax (int): Maximum frequency for Mel.
            power (bool): Whether to use power or magnitude spectrogram.
            elim_0th (bool): Whether to exclude 0th cepstrum in CheapTrick.
                If set to true, power is estimated by source-network.

        """
        super(ResidualLoss, self).__init__()
        self.sample_rate = sample_rate
        self.fft_size = fft_size
        self.hop_size = hop_size
        self.cheaptrick = CheapTrick(
            sample_rate=sample_rate,
            hop_size=hop_size,
            fft_size=fft_size,
            f0_floor=f0_floor,
            f0_ceil=f0_ceil,
        )
        self.win_length = fft_size
        self.register_buffer("window", torch.hann_window(self.win_length))

        # define mel-filter-bank
        self.n_mels = n_mels
        self.fmin = fmin
        self.fmax = fmax if fmax is not None else sample_rate / 2
        melmat = librosa_mel(sr=sample_rate, n_fft=fft_size, n_mels=n_mels, fmin=fmin, fmax=self.fmax).T
        self.register_buffer("melmat", torch.from_numpy(melmat).float())

        self.power = power
        self.elim_0th = elim_0th

    def forward(self, s, y, f):
        """Calculate forward propagation.

        Args:
            s (Tensor): Predicted source excitation signal (B, 1, T).
            y (Tensor): Ground truth signal (B, 1, T).
            f (Tensor): F0 sequence (B, 1, T // hop_size).

        Returns:
            Tensor: Loss value.

        """
        s, y, f = s.squeeze(1), y.squeeze(1), f.squeeze(1)

        with torch.no_grad():
            # calculate log power (or magnitude) spectrograms
            e = self.cheaptrick.forward(y, f, self.power, self.elim_0th)
            y = sifigan.losses.stft(
                y,
                self.fft_size,
                self.hop_size,
                self.win_length,
                self.window,
                power=self.power,
            )
            # adjust length, (B, T', C)
            minlen = min(e.size(1), y.size(1))
            e, y = e[:, :minlen, :], y[:, :minlen, :]

            # calculate mean power (or magnitude) of y
            if self.elim_0th:
                y_mean = y.mean(dim=-1, keepdim=True)

            # calculate target of output source signal
            y = torch.log(torch.clamp(y, min=1e-7))
            t = (y - e).exp()
            if self.elim_0th:
                t_mean = t.mean(dim=-1, keepdim=True)
                t = y_mean / t_mean * t

            # apply mel-filter-bank and log
            t = torch.matmul(t, self.melmat)
            t = torch.log(torch.clamp(t, min=1e-7))

        # calculate power (or magnitude) spectrogram
        s = sifigan.losses.stft(
            s,
            self.fft_size,
            self.hop_size,
            self.win_length,
            self.window,
            power=self.power,
        )
        # adjust length, (B, T', C)
        minlen = min(minlen, s.size(1))
        s, t = s[:, :minlen, :], t[:, :minlen, :]

        # apply mel-filter-bank and log
        s = torch.matmul(s, self.melmat)
        s = torch.log(torch.clamp(s, min=1e-7))

        loss = F.l1_loss(s, t.detach())

        return loss

sifigan/losses/reg.py
tonnetonne814/SiFi-VITS2-44100-Ja
from sifigan.models.discriminator import *  # NOQA
from sifigan.models.generator import *  # NOQA

sifigan/models/__init__.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Discriminator modules.

References:
    - https://github.com/kan-bayashi/ParallelWaveGAN
    - https://github.com/jik876/hifi-gan
    - UnivNet: A Neural Vocoder with Multi-Resolution Spectrogram Discriminators for High-Fidelity Waveform Generation
      https://www.isca-speech.org/archive/interspeech_2021/jang21_interspeech.html

"""

import copy
from logging import getLogger

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchaudio.functional import spectrogram

# A logger for this file
logger = getLogger(__name__)


class HiFiGANPeriodDiscriminator(nn.Module):
    """HiFiGAN period discriminator module."""

    def __init__(
        self,
        in_channels=1,
        out_channels=1,
        period=3,
        kernel_sizes=[5, 3],
        channels=32,
        downsample_scales=[3, 3, 3, 3, 1],
        max_downsample_channels=1024,
        bias=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
        use_weight_norm=True,
        use_spectral_norm=False,
    ):
        """Initialize HiFiGANPeriodDiscriminator module.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            period (int): Period.
            kernel_sizes (list): Kernel sizes of initial conv layers and the final conv layer.
            channels (int): Number of initial channels.
            downsample_scales (list): List of downsampling scales.
            max_downsample_channels (int): Number of maximum downsampling channels.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.
            use_weight_norm (bool): Whether to use weight norm.
                If set to true, it will be applied to all of the conv layers.
            use_spectral_norm (bool): Whether to use spectral norm.
                If set to true, it will be applied to all of the conv layers.

        """
        super().__init__()
        assert len(kernel_sizes) == 2
        assert kernel_sizes[0] % 2 == 1, "Kernel size must be odd number."
        assert kernel_sizes[1] % 2 == 1, "Kernel size must be odd number."

        self.period = period
        self.convs = nn.ModuleList()
        in_chs = in_channels
        out_chs = channels
        for downsample_scale in downsample_scales:
            self.convs += [
                nn.Sequential(
                    nn.Conv2d(
                        in_chs,
                        out_chs,
                        (kernel_sizes[0], 1),
                        (downsample_scale, 1),
                        padding=((kernel_sizes[0] - 1) // 2, 0),
                        bias=bias,
                    ),
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                )
            ]
            in_chs = out_chs
            # NOTE(kan-bayashi): Use downsample_scale + 1?
            out_chs = min(out_chs * 4, max_downsample_channels)
        self.output_conv = nn.Conv2d(
            out_chs,
            out_channels,
            (kernel_sizes[1] - 1, 1),
            1,
            padding=((kernel_sizes[1] - 1) // 2, 0),
            bias=bias,
        )

        if use_weight_norm and use_spectral_norm:
            raise ValueError("Either use use_weight_norm or use_spectral_norm.")

        # apply weight norm
        if use_weight_norm:
            self.apply_weight_norm()

        # apply spectral norm
        if use_spectral_norm:
            self.apply_spectral_norm()

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            c (Tensor): Input tensor (B, in_channels, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            list: List of each layer's tensors.

        """
        # transform 1d to 2d -> (B, C, T/P, P)
        b, c, t = x.shape
        if t % self.period != 0:
            n_pad = self.period - (t % self.period)
            x = F.pad(x, (0, n_pad), "reflect")
            t += n_pad
        x = x.view(b, c, t // self.period, self.period)

        # forward conv
        fmap = []
        for f in self.convs:
            x = f(x)
            if return_fmaps:
                fmap.append(x)
        x = self.output_conv(x)
        out = torch.flatten(x, 1, -1)

        if return_fmaps:
            return out, fmap
        else:
            return out

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv2d):
                nn.utils.weight_norm(m)
                logger.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)

    def apply_spectral_norm(self):
        """Apply spectral normalization module from all of the layers."""

        def _apply_spectral_norm(m):
            if isinstance(m, nn.Conv2d):
                nn.utils.spectral_norm(m)
                logger.debug(f"Spectral norm is applied to {m}.")

        self.apply(_apply_spectral_norm)


class HiFiGANMultiPeriodDiscriminator(nn.Module):
    """HiFiGAN multi-period discriminator module."""

    def __init__(
        self,
        periods=[2, 3, 5, 7, 11],
        discriminator_params={
            "in_channels": 1,
            "out_channels": 1,
            "kernel_sizes": [5, 3],
            "channels": 32,
            "downsample_scales": [3, 3, 3, 3, 1],
            "max_downsample_channels": 1024,
            "bias": True,
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.1},
            "use_weight_norm": True,
            "use_spectral_norm": False,
        },
    ):
        """Initialize HiFiGANMultiPeriodDiscriminator module.

        Args:
            periods (list): List of periods.
            discriminator_params (dict): Parameters for hifi-gan period discriminator module.
                The period parameter will be overwritten.

        """
        super().__init__()
        self.discriminators = nn.ModuleList()
        for period in periods:
            params = copy.deepcopy(discriminator_params)
            params["period"] = period
            self.discriminators += [HiFiGANPeriodDiscriminator(**params)]

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of list of each discriminator outputs, which consists of each layer output tensors.

        """
        outs, fmaps = [], []
        for f in self.discriminators:
            if return_fmaps:
                out, fmap = f(x, return_fmaps)
                fmaps.extend(fmap)
            else:
                out = f(x)
            outs.append(out)

        if return_fmaps:
            return outs, fmaps
        else:
            return outs


class HiFiGANScaleDiscriminator(nn.Module):
    """HiFi-GAN scale discriminator module."""

    def __init__(
        self,
        in_channels=1,
        out_channels=1,
        kernel_sizes=[15, 41, 5, 3],
        channels=128,
        max_downsample_channels=1024,
        max_groups=16,
        bias=True,
        downsample_scales=[2, 2, 4, 4, 1],
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
        use_weight_norm=True,
        use_spectral_norm=False,
    ):
        """Initilize HiFiGAN scale discriminator module.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            kernel_sizes (list): List of four kernel sizes. The first will be used for the first conv layer,
                and the second is for downsampling part, and the remaining two are for output layers.
            channels (int): Initial number of channels for conv layer.
            max_downsample_channels (int): Maximum number of channels for downsampling layers.
            bias (bool): Whether to add bias parameter in convolution layers.
            downsample_scales (list): List of downsampling scales.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.
            use_weight_norm (bool): Whether to use weight norm.
                If set to true, it will be applied to all of the conv layers.
            use_spectral_norm (bool): Whether to use spectral norm.
                If set to true, it will be applied to all of the conv layers.

        """
        super().__init__()
        self.layers = nn.ModuleList()

        # check kernel size is valid
        assert len(kernel_sizes) == 4
        for ks in kernel_sizes:
            assert ks % 2 == 1

        # add first layer
        self.layers += [
            nn.Sequential(
                nn.Conv1d(
                    in_channels,
                    channels,
                    # NOTE(kan-bayashi): Use always the same kernel size
                    kernel_sizes[0],
                    bias=bias,
                    padding=(kernel_sizes[0] - 1) // 2,
                ),
                getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
            )
        ]

        # add downsample layers
        in_chs = channels
        out_chs = channels
        # NOTE(kan-bayashi): Remove hard coding?
        groups = 4
        for downsample_scale in downsample_scales:
            self.layers += [
                nn.Sequential(
                    nn.Conv1d(
                        in_chs,
                        out_chs,
                        kernel_size=kernel_sizes[1],
                        stride=downsample_scale,
                        padding=(kernel_sizes[1] - 1) // 2,
                        groups=groups,
                        bias=bias,
                    ),
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                )
            ]
            in_chs = out_chs
            # NOTE(kan-bayashi): Remove hard coding?
            out_chs = min(in_chs * 2, max_downsample_channels)
            # NOTE(kan-bayashi): Remove hard coding?
            groups = min(groups * 4, max_groups)

        # add final layers
        out_chs = min(in_chs * 2, max_downsample_channels)
        self.layers += [
            nn.Sequential(
                nn.Conv1d(
                    in_chs,
                    out_chs,
                    kernel_size=kernel_sizes[2],
                    stride=1,
                    padding=(kernel_sizes[2] - 1) // 2,
                    bias=bias,
                ),
                getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
            )
        ]
        self.last_layer = nn.Conv1d(
            out_chs,
            out_channels,
            kernel_size=kernel_sizes[3],
            stride=1,
            padding=(kernel_sizes[3] - 1) // 2,
            bias=bias,
        )

        if use_weight_norm and use_spectral_norm:
            raise ValueError("Either use use_weight_norm or use_spectral_norm.")

        # apply weight norm
        if use_weight_norm:
            self.apply_weight_norm()

        # apply spectral norm
        if use_spectral_norm:
            self.apply_spectral_norm()

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of output tensors of each layer.

        """
        fmap = []
        for f in self.layers:
            x = f(x)
            if return_fmaps:
                fmap.append(x)
        out = self.last_layer(x)

        if return_fmaps:
            return out, fmap
        else:
            return out

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv2d):
                nn.utils.weight_norm(m)
                logger.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)

    def apply_spectral_norm(self):
        """Apply spectral normalization module from all of the layers."""

        def _apply_spectral_norm(m):
            if isinstance(m, nn.Conv2d):
                nn.utils.spectral_norm(m)
                logger.debug(f"Spectral norm is applied to {m}.")

        self.apply(_apply_spectral_norm)


class HiFiGANMultiScaleDiscriminator(nn.Module):
    """HiFi-GAN multi-scale discriminator module."""

    def __init__(
        self,
        scales=3,
        downsample_pooling="AvgPool1d",
        # follow the official implementation setting
        downsample_pooling_params={
            "kernel_size": 4,
            "stride": 2,
            "padding": 2,
        },
        discriminator_params={
            "in_channels": 1,
            "out_channels": 1,
            "kernel_sizes": [15, 41, 5, 3],
            "channels": 128,
            "max_downsample_channels": 1024,
            "max_groups": 16,
            "bias": True,
            "downsample_scales": [2, 2, 4, 4, 1],
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.1},
        },
        follow_official_norm=False,
    ):
        """Initilize HiFiGAN multi-scale discriminator module.

        Args:
            scales (int): Number of multi-scales.
            downsample_pooling (str): Pooling module name for downsampling of the inputs.
            downsample_pooling_params (dict): Parameters for the above pooling module.
            discriminator_params (dict): Parameters for hifi-gan scale discriminator module.
            follow_official_norm (bool): Whether to follow the norm setting of the official
                implementaion. The first discriminator uses spectral norm and the other
                discriminators use weight norm.

        """
        super().__init__()
        self.discriminators = nn.ModuleList()

        # add discriminators
        for i in range(scales):
            params = copy.deepcopy(discriminator_params)
            if follow_official_norm:
                if i == 0:
                    params["use_weight_norm"] = False
                    params["use_spectral_norm"] = True
                else:
                    params["use_weight_norm"] = True
                    params["use_spectral_norm"] = False
            self.discriminators += [HiFiGANScaleDiscriminator(**params)]
        self.pooling = getattr(nn, downsample_pooling)(**downsample_pooling_params)

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of list of each discriminator outputs, which consists of each layer output tensors.

        """
        outs, fmaps = [], []
        for f in self.discriminators:
            if return_fmaps:
                out, fmap = f(x, return_fmaps)
                fmaps.extend(fmap)
            else:
                out = f(x)
            outs.append(out)
            x = self.pooling(x)

        if return_fmaps:
            return outs, fmaps
        else:
            return outs


class HiFiGANMultiScaleMultiPeriodDiscriminator(nn.Module):
    """HiFi-GAN multi-scale + multi-period discriminator module."""

    def __init__(
        self,
        # Multi-scale discriminator related
        scales=3,
        scale_downsample_pooling="AvgPool1d",
        scale_downsample_pooling_params={
            "kernel_size": 4,
            "stride": 2,
            "padding": 2,
        },
        scale_discriminator_params={
            "in_channels": 1,
            "out_channels": 1,
            "kernel_sizes": [15, 41, 5, 3],
            "channels": 128,
            "max_downsample_channels": 1024,
            "max_groups": 16,
            "bias": True,
            "downsample_scales": [2, 2, 4, 4, 1],
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.1},
        },
        follow_official_norm=True,
        # Multi-period discriminator related
        periods=[2, 3, 5, 7, 11],
        period_discriminator_params={
            "in_channels": 1,
            "out_channels": 1,
            "kernel_sizes": [5, 3],
            "channels": 32,
            "downsample_scales": [3, 3, 3, 3, 1],
            "max_downsample_channels": 1024,
            "bias": True,
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.1},
            "use_weight_norm": True,
            "use_spectral_norm": False,
        },
    ):
        """Initilize HiFiGAN multi-scale + multi-period discriminator module.

        Args:
            scales (int): Number of multi-scales.
            scale_downsample_pooling (str): Pooling module name for downsampling of the inputs.
            scale_downsample_pooling_params (dict): Parameters for the above pooling module.
            scale_discriminator_params (dict): Parameters for hifi-gan scale discriminator module.
            follow_official_norm (bool): Whether to follow the norm setting of the official
                implementaion. The first discriminator uses spectral norm and the other
                discriminators use weight norm.
            periods (list): List of periods.
            period_discriminator_params (dict): Parameters for hifi-gan period discriminator module.
                The period parameter will be overwritten.

        """
        super().__init__()
        self.msd = HiFiGANMultiScaleDiscriminator(
            scales=scales,
            downsample_pooling=scale_downsample_pooling,
            downsample_pooling_params=scale_downsample_pooling_params,
            discriminator_params=scale_discriminator_params,
            follow_official_norm=follow_official_norm,
        )
        self.mpd = HiFiGANMultiPeriodDiscriminator(
            periods=periods,
            discriminator_params=period_discriminator_params,
        )

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of list of each discriminator outputs,
                which consists of each layer output tensors.
                Multi scale and multi period ones are concatenated.

        """
        if return_fmaps:
            msd_outs, msd_fmaps = self.msd(x, return_fmaps)
            mpd_outs, mpd_fmaps = self.mpd(x, return_fmaps)
            outs = msd_outs + mpd_outs
            fmaps = msd_fmaps + mpd_fmaps
            return outs, fmaps
        else:
            msd_outs = self.msd(x)
            mpd_outs = self.mpd(x)
            outs = msd_outs + mpd_outs
            return outs


class UnivNetSpectralDiscriminator(nn.Module):
    """UnivNet spectral discriminator module."""

    def __init__(
        self,
        fft_size,
        hop_size,
        win_length,
        window="hann_window",
        kernel_sizes=[(3, 9), (3, 9), (3, 9), (3, 9), (3, 3), (3, 3)],
        strides=[(1, 1), (1, 2), (1, 2), (1, 2), (1, 1), (1, 1)],
        channels=32,
        bias=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.2},
        use_weight_norm=True,
    ):
        """Initilize HiFiGAN scale discriminator module.

        Args:
            fft_size (list): FFT size.
            hop_size (int): Hop size.
            win_length (int): Window length.
            window (stt): Name of window function.
            kernel_sizes (list): List of kernel sizes in down-sampling CNNs.
            strides (list): List of stride sizes in down-sampling CNNs.
            channels (int): Number of channels for conv layer.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.
            use_weight_norm (bool): Whether to use weight norm.
                If set to true, it will be applied to all of the conv layers.

        """
        super().__init__()

        self.fft_size = fft_size
        self.hop_size = hop_size
        self.win_length = win_length
        self.register_buffer("window", getattr(torch, window)(win_length))

        self.layers = nn.ModuleList()

        # check kernel size is valid
        assert len(kernel_sizes) == len(strides)

        # add first layer
        self.layers += [
            nn.Sequential(
                nn.Conv2d(
                    1,
                    channels,
                    kernel_sizes[0],
                    stride=strides[0],
                    bias=bias,
                ),
                getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
            )
        ]

        for i in range(1, len(kernel_sizes) - 2):
            self.layers += [
                nn.Sequential(
                    nn.Conv2d(
                        channels,
                        channels,
                        kernel_size=kernel_sizes[i],
                        stride=strides[i],
                        bias=bias,
                    ),
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                )
            ]

        # add final layers
        self.layers += [
            nn.Sequential(
                nn.Conv2d(
                    channels,
                    channels,
                    kernel_size=kernel_sizes[-2],
                    stride=strides[-2],
                    bias=bias,
                ),
                getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
            )
        ]
        self.layers += [
            nn.Conv2d(
                channels,
                1,
                kernel_size=kernel_sizes[-1],
                stride=strides[-1],
                bias=bias,
            )
        ]

        # apply weight norm
        if use_weight_norm:
            self.apply_weight_norm()

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of output tensors of each layer.

        """
        x = spectrogram(
            x,
            pad=self.win_length // 2,
            window=self.window,
            n_fft=self.fft_size,
            hop_length=self.hop_size,
            win_length=self.win_length,
            power=1.0,
            normalized=False,
        ).transpose(-1, -2)

        fmap = []
        for f in self.layers:
            x = f(x)
            if return_fmaps:
                fmap.append(x)

        if return_fmaps:
            return x, fmap
        else:
            return x

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv2d):
                nn.utils.weight_norm(m)
                logger.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)


class UnivNetMultiResolutionSpectralDiscriminator(nn.Module):
    """UnivNet multi-resolution spectral discriminator module."""

    def __init__(
        self,
        fft_sizes=[1024, 2048, 512],
        hop_sizes=[120, 240, 50],
        win_lengths=[600, 1200, 240],
        window="hann_window",
        discriminator_params={
            "channels": 32,
            "kernel_sizes": [(3, 9), (3, 9), (3, 9), (3, 9), (3, 3), (3, 3)],
            "strides": [(1, 1), (1, 2), (1, 2), (1, 2), (1, 1), (1, 1)],
            "bias": True,
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.2},
        },
    ):
        """Initilize UnivNetMultiResolutionSpectralDiscriminator module.

        Args:
            fft_sizes (list): FFT sizes for each spectral discriminator.
            hop_sizes (list): Hop sizes for each spectral discriminator.
            win_lengths (list): Window lengths for each spectral discriminator.
            window (stt): Name of window function.
            discriminator_params (dict): Parameters for univ-net spectral discriminator module.

        """
        super().__init__()
        assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
        self.discriminators = nn.ModuleList()

        # add discriminators
        for i in range(len(fft_sizes)):
            params = copy.deepcopy(discriminator_params)
            self.discriminators += [
                UnivNetSpectralDiscriminator(
                    fft_size=fft_sizes[i],
                    hop_size=hop_sizes[i],
                    win_length=win_lengths[i],
                    window=window,
                    **params,
                )
            ]

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of list of each discriminator outputs, which consists of each layer output tensors.

        """
        outs, fmaps = [], []
        for f in self.discriminators:
            if return_fmaps:
                out, fmap = f(x, return_fmaps)
                fmaps.extend(fmap)
            else:
                out = f(x)
            outs.append(out)

        if return_fmaps:
            return outs, fmaps
        else:
            return outs


class UnivNetMultiResolutionMultiPeriodDiscriminator(nn.Module):
    """UnivNet multi-resolution + multi-period discriminator module."""

    def __init__(
        self,
        # Multi-resolution discriminator related
        fft_sizes=[1024, 2048, 512],
        hop_sizes=[120, 240, 50],
        win_lengths=[600, 1200, 240],
        window="hann_window",
        spectral_discriminator_params={
            "channels": 32,
            "kernel_sizes": [(3, 9), (3, 9), (3, 9), (3, 9), (3, 3), (3, 3)],
            "strides": [(1, 1), (1, 2), (1, 2), (1, 2), (1, 1), (1, 1)],
            "bias": True,
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.2},
        },
        # Multi-period discriminator related
        periods=[2, 3, 5, 7, 11],
        period_discriminator_params={
            "in_channels": 1,
            "out_channels": 1,
            "kernel_sizes": [5, 3],
            "channels": 32,
            "downsample_scales": [3, 3, 3, 3, 1],
            "max_downsample_channels": 1024,
            "bias": True,
            "nonlinear_activation": "LeakyReLU",
            "nonlinear_activation_params": {"negative_slope": 0.1},
            "use_weight_norm": True,
            "use_spectral_norm": False,
        },
    ):
        """Initilize UnivNetMultiResolutionMultiPeriodDiscriminator module.

        Args:
            fft_sizes (list): FFT sizes for each spectral discriminator.
            hop_sizes (list): Hop sizes for each spectral discriminator.
            win_lengths (list): Window lengths for each spectral discriminator.
            window (stt): Name of window function.
            sperctral_discriminator_params (dict): Parameters for hifi-gan scale discriminator module.
            periods (list): List of periods.
            period_discriminator_params (dict): Parameters for hifi-gan period discriminator module.
                The period parameter will be overwritten.

        """
        super().__init__()
        self.mrd = UnivNetMultiResolutionSpectralDiscriminator(
            fft_sizes=fft_sizes,
            hop_sizes=hop_sizes,
            win_lengths=win_lengths,
            window=window,
            discriminator_params=spectral_discriminator_params,
        )
        self.mpd = HiFiGANMultiPeriodDiscriminator(
            periods=periods,
            discriminator_params=period_discriminator_params,
        )

    def forward(self, x, return_fmaps=False):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input noise signal (B, 1, T).
            return_fmaps (bool): Whether to return feature maps.

        Returns:
            List: List of list of each discriminator outputs,
                which consists of each layer output tensors.
                Multi scale and multi period ones are concatenated.

        """
        if return_fmaps:
            mrd_outs, mrd_fmaps = self.mrd(x, return_fmaps)
            mpd_outs, mpd_fmaps = self.mpd(x, return_fmaps)
            outs = mrd_outs + mpd_outs
            fmaps = mrd_fmaps + mpd_fmaps

            return outs, fmaps
        else:
            mrd_outs = self.mrd(x)
            mpd_outs = self.mpd(x)
            outs = mrd_outs + mpd_outs

            return outs

sifigan/models/discriminator.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""HiFiGAN and SiFiGAN Generator modules.

References:
    - https://github.com/kan-bayashi/ParallelWaveGAN
    - https://github.com/bigpon/QPPWG
    - https://github.com/jik876/hifi-gan

"""

from logging import getLogger

import torch
import torch.nn as nn
from sifigan.layers import AdaptiveResidualBlock, Conv1d, ResidualBlock

# A logger for this file
logger = getLogger(__name__)


class HiFiGANGenerator(nn.Module):
    """HiFiGAN generator module with customizable pitch-dependent mechanisms."""

    def __init__(
        self,
        in_channels,
        out_channels=1,
        channels=512,
        kernel_size=7,
        upsample_scales=(5, 4, 3, 2),
        upsample_kernel_sizes=(10, 8, 6, 4),
        qp_resblock_kernel_size=3,
        qp_resblock_dilations=[(1), (1, 2), (1, 2, 4), (1, 2, 4, 8)],
        qp_use_additional_convs=True,
        resblock_kernel_sizes=(3, 7, 11),
        resblock_dilations=[(1, 3, 5), (1, 3, 5), (1, 3, 5)],
        use_additional_convs=False,
        use_sine_embs=False,
        use_qp_resblocks=False,
        bias=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
        use_weight_norm=True,
    ):
        """Initialize QpHiFiGANGenerator module.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            channels (int): Number of hidden representation channels.
            kernel_size (int): Kernel size of initial and final conv layer.
            upsample_scales (list): List of upsampling scales.
            upsample_kernel_sizes (list): List of kernel sizes for upsampling layers.
            qp_resblock_kernel_size (int): Kernel size for quasi-periodic residual blocks.
            qp_resblock_dilations (list): Dilations for quasi-periodic residual blocks.
            qp_use_additional_convs (bool): Whether to use additional conv layers.
            resblock_kernel_sizes (list): List of kernel sizes for residual blocks.
            resblock_dilations (list): List of dilation list for residual blocks.
            use_additional_convs (bool): Whether to use additional conv layers in residual blocks.
            use_sine_embs (bool): Whether to use sine embeddings.
            use_qp_resblocks (bool): Whether to use quasi-periodic residual blocks.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.
            use_weight_norm (bool): Whether to use weight norm.
                If set to true, it will be applied to all of the conv layers.

        """
        super().__init__()
        # check hyperparameters are valid
        assert kernel_size % 2 == 1, "Kernel size must be odd number."
        assert len(upsample_scales) == len(upsample_kernel_sizes)
        assert len(resblock_dilations) == len(resblock_kernel_sizes)

        # define modules
        self.num_upsamples = len(upsample_kernel_sizes)
        self.qp_resdual_dilations = qp_resblock_dilations
        self.num_blocks = len(resblock_kernel_sizes)
        self.input_conv = Conv1d(
            in_channels,
            channels,
            kernel_size,
            bias=bias,
            padding=(kernel_size - 1) // 2,
        )
        self.upsamples = nn.ModuleList()
        self.use_qp_resblocks = use_qp_resblocks
        if use_qp_resblocks:
            self.qp_blocks = nn.ModuleList()
        self.blocks = nn.ModuleList()
        for i in range(len(upsample_kernel_sizes)):
            assert upsample_kernel_sizes[i] == 2 * upsample_scales[i]
            self.upsamples += [
                nn.Sequential(
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                    nn.ConvTranspose1d(
                        channels // (2 ** i),
                        channels // (2 ** (i + 1)),
                        upsample_kernel_sizes[i],
                        upsample_scales[i],
                        padding=upsample_scales[i] // 2 + upsample_scales[i] % 2,
                        output_padding=upsample_scales[i] % 2,
                        bias=bias,
                    ),
                )
            ]
            if use_qp_resblocks:
                self.qp_blocks += [
                    AdaptiveResidualBlock(
                        kernel_size=qp_resblock_kernel_size,
                        channels=channels // (2 ** (i + 1)),
                        dilations=qp_resblock_dilations[i],
                        bias=bias,
                        use_additional_convs=qp_use_additional_convs,
                        nonlinear_activation=nonlinear_activation,
                        nonlinear_activation_params=nonlinear_activation_params,
                    )
                ]
            for j in range(len(resblock_kernel_sizes)):
                self.blocks += [
                    ResidualBlock(
                        kernel_size=resblock_kernel_sizes[j],
                        channels=channels // (2 ** (i + 1)),
                        dilations=resblock_dilations[j],
                        bias=bias,
                        use_additional_convs=use_additional_convs,
                        nonlinear_activation=nonlinear_activation,
                        nonlinear_activation_params=nonlinear_activation_params,
                    )
                ]
        self.output_conv = nn.Sequential(
            nn.LeakyReLU(),
            nn.Conv1d(
                channels // (2 ** (i + 1)),
                out_channels,
                kernel_size,
                bias=bias,
                padding=(kernel_size - 1) // 2,
            ),
            nn.Tanh(),
        )

        # sine embedding layers
        self.use_sine_embs = use_sine_embs
        if use_sine_embs:
            self.emb = Conv1d(
                1,
                channels // (2 ** len(upsample_kernel_sizes)),
                kernel_size,
                bias=bias,
                padding=(kernel_size - 1) // 2,
            )
            self.downsamples = nn.ModuleList()
            for i in reversed(range(len(upsample_kernel_sizes))):
                self.downsamples += [
                    nn.Sequential(
                        nn.Conv1d(
                            channels // (2 ** (i + 1)),
                            channels // (2 ** i),
                            upsample_kernel_sizes[i],
                            upsample_scales[i],
                            padding=upsample_scales[i] - (upsample_kernel_sizes[i] % 2 == 0),
                            bias=bias,
                        ),
                        getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                    )
                ]

        # apply weight norm
        if use_weight_norm:
            self.apply_weight_norm()

        # reset parameters
        self.reset_parameters()

    def forward(self, x, c, d=None):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input sine signal (B, 1, T).
            c (Tensor): Input tensor (B, in_channels, T).
            d (List): F0-dependent dilation factors [(B, 1, T) x num_upsamples].

        Returns:
            Tensor: Output tensor (B, out_channels, T).

        """
        c = self.input_conv(c)

        if self.use_sine_embs:
            x = self.emb(x)
            embs = [x]
            for i in range(self.num_upsamples - 1):
                x = self.downsamples[i](x)
                embs += [x]

        for i in range(self.num_upsamples):
            c = self.upsamples[i](c)
            if self.use_sine_embs:
                c = c + embs[-i - 1]
            if self.use_qp_resblocks:
                c = self.qp_blocks[i](c, d[i])
            cs = 0.0  # initialize
            for j in range(self.num_blocks):
                cs += self.blocks[i * self.num_blocks + j](c)
            c = cs / self.num_blocks
        c = self.output_conv(c)

        return (c,)

    def reset_parameters(self):
        """Reset parameters.

        This initialization follows the official implementation manner.
        https://github.com/jik876/hifi-gan/blob/master/models.py

        """

        def _reset_parameters(m):
            if isinstance(m, (nn.Conv1d, nn.ConvTranspose1d)):
                m.weight.data.normal_(0.0, 0.01)
                logger.debug(f"Reset parameters in {m}.")

        self.apply(_reset_parameters)

    def remove_weight_norm(self):
        """Remove weight normalization module from all of the layers."""

        def _remove_weight_norm(m):
            try:
                logger.debug(f"Weight norm is removed from {m}.")
                nn.utils.remove_weight_norm(m)
            except ValueError:  # this module didn't have weight norm
                return

        self.apply(_remove_weight_norm)

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
                nn.utils.weight_norm(m)
                logger.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)


class SiFiGANGenerator(nn.Module):
    """SiFiGAN generator module."""

    def __init__(
        self,
        in_channels,
        out_channels=1,
        channels=512,
        gin_channels=256,
        kernel_size=7,
        upsample_scales=(5, 4, 3, 2),
        upsample_kernel_sizes=(10, 8, 6, 4),
        source_network_params={
            "resblock_kernel_size": 3,  # currently only 3 is supported.
            "resblock_dilations": [(1), (1, 2), (1, 2, 4), (1, 2, 4, 8)],
            "use_additional_convs": True,
        },
        filter_network_params={
            "resblock_kernel_sizes": (3, 5, 7),
            "resblock_dilations": [(1, 3, 5), (1, 3, 5), (1, 3, 5)],
            "use_additional_convs": False,
        },
        share_upsamples=False,
        share_downsamples=False,
        bias=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
        use_weight_norm=True,
    ):
        """Initialize SiFiGANGenerator module.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            channels (int): Number of hidden representation channels.
            kernel_size (int): Kernel size of initial and final conv layer.
            upsample_scales (list): List of upsampling scales.
            upsample_kernel_sizes (list): List of kernel sizes for upsampling layers.
            source_network_params (dict): Parameters for source-network.
            filter_network_params (dict): Parameters for filter-network.
            share_upsamples (bool): Whether to share up-sampling transposed CNNs.
            share_downsamples (bool): Whether to share down-sampling CNNs.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.
            use_weight_norm (bool): Whether to use weight norm.
                If set to true, it will be applied to all of the conv layers.

        """
        super().__init__()
        # check hyperparameters are valid
        assert kernel_size % 2 == 1, "Kernel size must be odd number."
        assert len(upsample_scales) == len(upsample_kernel_sizes)

        # define modules
        self.num_upsamples = len(upsample_kernel_sizes)
        self.source_network_params = source_network_params
        self.filter_network_params = filter_network_params
        self.share_upsamples = share_upsamples
        self.share_downsamples = share_downsamples
        self.sn = nn.ModuleDict()
        self.fn = nn.ModuleDict()
        self.input_conv = Conv1d(
            in_channels,
            channels,
            kernel_size,
            bias=bias,
            padding=(kernel_size - 1) // 2,
        )
        self.sn["upsamples"] = nn.ModuleList()
        self.fn["upsamples"] = nn.ModuleList()
        self.sn["blocks"] = nn.ModuleList()
        self.fn["blocks"] = nn.ModuleList()
        for i in range(len(upsample_kernel_sizes)):
            assert upsample_kernel_sizes[i] == 2 * upsample_scales[i]
            self.sn["upsamples"] += [
                nn.Sequential(
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                    nn.ConvTranspose1d(
                        channels // (2 ** i),
                        channels // (2 ** (i + 1)),
                        upsample_kernel_sizes[i],
                        upsample_scales[i],
                        padding=upsample_scales[i] // 2 + upsample_scales[i] % 2,
                        output_padding=upsample_scales[i] % 2,
                        bias=bias,
                    ),
                )
            ]
            if not share_upsamples:
                self.fn["upsamples"] += [
                    nn.Sequential(
                        getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                        nn.ConvTranspose1d(
                            channels // (2 ** i),
                            channels // (2 ** (i + 1)),
                            upsample_kernel_sizes[i],
                            upsample_scales[i],
                            padding=upsample_scales[i] // 2 + upsample_scales[i] % 2,
                            output_padding=upsample_scales[i] % 2,
                            bias=bias,
                        ),
                    )
                ]
            self.sn["blocks"] += [
                AdaptiveResidualBlock(
                    kernel_size=source_network_params["resblock_kernel_size"],
                    channels=channels // (2 ** (i + 1)),
                    dilations=source_network_params["resblock_dilations"][i],
                    bias=bias,
                    use_additional_convs=source_network_params["use_additional_convs"],
                    nonlinear_activation=nonlinear_activation,
                    nonlinear_activation_params=nonlinear_activation_params,
                )
            ]
            for j in range(len(filter_network_params["resblock_kernel_sizes"])):
                self.fn["blocks"] += [
                    ResidualBlock(
                        kernel_size=filter_network_params["resblock_kernel_sizes"][j],
                        channels=channels // (2 ** (i + 1)),
                        dilations=filter_network_params["resblock_dilations"][j],
                        bias=bias,
                        use_additional_convs=filter_network_params["use_additional_convs"],
                        nonlinear_activation=nonlinear_activation,
                        nonlinear_activation_params=nonlinear_activation_params,
                    )
                ]
        self.sn["output_conv"] = nn.Sequential(
            nn.LeakyReLU(),
            nn.Conv1d(
                channels // (2 ** (i + 1)),
                out_channels,
                kernel_size,
                bias=bias,
                padding=(kernel_size - 1) // 2,
            ),
        )
        self.fn["output_conv"] = nn.Sequential(
            nn.LeakyReLU(),
            nn.Conv1d(
                channels // (2 ** (i + 1)),
                out_channels,
                kernel_size,
                bias=bias,
                padding=(kernel_size - 1) // 2,
            ),
            nn.Tanh(),
        )

        # sine embedding layers
        self.sn["emb"] = Conv1d(
            1,
            channels // (2 ** len(upsample_kernel_sizes)),
            kernel_size,
            bias=bias,
            padding=(kernel_size - 1) // 2,
        )
        # down-sampling CNNs
        self.sn["downsamples"] = nn.ModuleList()
        for i in reversed(range(len(upsample_kernel_sizes))):
            self.sn["downsamples"] += [
                nn.Sequential(
                    nn.Conv1d(
                        channels // (2 ** (i + 1)),
                        channels // (2 ** i),
                        upsample_kernel_sizes[i],
                        upsample_scales[i],
                        padding=upsample_scales[i] - (upsample_kernel_sizes[i] % 2 == 0),
                        bias=bias,
                    ),
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                )
            ]
        if not share_downsamples:
            self.fn["downsamples"] = nn.ModuleList()
            for i in reversed(range(len(upsample_kernel_sizes))):
                self.fn["downsamples"] += [
                    nn.Sequential(
                        nn.Conv1d(
                            channels // (2 ** (i + 1)),
                            channels // (2 ** i),
                            upsample_kernel_sizes[i],
                            upsample_scales[i],
                            padding=upsample_scales[i] - (upsample_kernel_sizes[i] % 2 == 0),
                            bias=bias,
                        ),
                        getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                    )
                ]
        # apply speaker embedding
        self.cond = nn.Conv1d(gin_channels, channels, 1)

        # apply weight norm
        if use_weight_norm:
            self.apply_weight_norm()

        # reset parameters
        self.reset_parameters()

    def forward(self, x, c, d,g=None):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input sine signal (B, 1, T).
            c (Tensor): Input tensor (B, in_channels, T').
            d (List): F0-dependent dilation factors [(B, 1, T') x num_upsamples].

        Returns:
            Tensor: Output tensor (B, out_channels, T).

        """
        # currently, same input feature is input to each network
        c = self.input_conv(c)
        if g is not None:
            g = self.cond(g)
            c = c + g
        e = c

        # source-network forward
        x = self.sn["emb"](x)
        embs = [x]
        for i in range(self.num_upsamples - 1):
            x = self.sn["downsamples"][i](x)
            embs += [x]
        for i in range(self.num_upsamples):
            # excitation generation network
            e = self.sn["upsamples"][i](e) 
            e = e + embs[-i - 1]
            e = self.sn["blocks"][i](e, d[i])
        e_ = self.sn["output_conv"](e)

        # filter-network forward
        embs = [e]
        for i in range(self.num_upsamples - 1):
            if self.share_downsamples:
                e = self.sn["downsamples"][i](e)
            else:
                e = self.fn["downsamples"][i](e)
            embs += [e]
        num_blocks = len(self.filter_network_params["resblock_kernel_sizes"])
        for i in range(self.num_upsamples):
            # resonance filtering network
            if self.share_upsamples:
                c = self.sn["upsamples"][i](c) + embs[-i - 1]
            else:
                c = self.fn["upsamples"][i](c) + embs[-i - 1]
            cs = 0.0  # initialize
            for j in range(num_blocks):
                cs += self.fn["blocks"][i * num_blocks + j](c)
            c = cs / num_blocks
        c = self.fn["output_conv"](c)

        return c, e_

    def reset_parameters(self):
        """Reset parameters.

        This initialization follows the official implementation manner.
        https://github.com/jik876/hifi-gan/blob/master/models.py

        """

        def _reset_parameters(m):
            if isinstance(m, (nn.Conv1d, nn.ConvTranspose1d)):
                m.weight.data.normal_(0.0, 0.01)
                logger.debug(f"Reset parameters in {m}.")

        self.apply(_reset_parameters)

    def remove_weight_norm(self):
        """Remove weight normalization module from all of the layers."""

        def _remove_weight_norm(m):
            try:
                logger.debug(f"Weight norm is removed from {m}.")
                nn.utils.remove_weight_norm(m)
            except ValueError:  # this module didn't have weight norm
                return

        self.apply(_remove_weight_norm)

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
                nn.utils.weight_norm(m)
                logger.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)


class SiFiGANDirectGenerator(nn.Module):
    """SiFiGAN Direct generator module."""

    def __init__(
        self,
        in_channels,
        out_channels=1,
        channels=512,
        kernel_size=7,
        upsample_scales=(5, 4, 3, 2),
        upsample_kernel_sizes=(10, 8, 6, 4),
        source_network_params={
            "resblock_kernel_size": 3,  # currently only 3 is supported.
            "resblock_dilations": [(1), (1, 2), (1, 2, 4), (1, 2, 4, 8)],
            "use_additional_convs": True,
        },
        filter_network_params={
            "resblock_kernel_sizes": (3, 5, 7),
            "resblock_dilations": [(1, 3, 5), (1, 3, 5), (1, 3, 5)],
            "use_additional_convs": False,
        },
        share_upsamples=False,
        bias=True,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.1},
        use_weight_norm=True,
    ):
        """Initialize SiFiGANDirectGenerator module.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            channels (int): Number of hidden representation channels.
            kernel_size (int): Kernel size of initial and final conv layer.
            upsample_scales (list): List of upsampling scales.
            upsample_kernel_sizes (list): List of kernel sizes for upsampling layers.
            source_network_params (dict): Parameters for source-network.
            filter_network_params (dict): Parameters for filter-network.
            share_upsamples (bool): Whether to share up-sampling transposed CNNs.
            share_downsamples (bool): Whether to share down-sampling CNNs.
            bias (bool): Whether to add bias parameter in convolution layers.
            nonlinear_activation (str): Activation function module name.
            nonlinear_activation_params (dict): Hyperparameters for activation function.
            use_weight_norm (bool): Whether to use weight norm.
                If set to true, it will be applied to all of the conv layers.

        """
        super().__init__()
        # check hyperparameters are valid
        assert kernel_size % 2 == 1, "Kernel size must be odd number."
        assert len(upsample_scales) == len(upsample_kernel_sizes)

        # define modules
        self.num_upsamples = len(upsample_kernel_sizes)
        self.source_network_params = source_network_params
        self.filter_network_params = filter_network_params
        self.share_upsamples = share_upsamples
        self.sn = nn.ModuleDict()
        self.fn = nn.ModuleDict()
        self.input_conv = Conv1d(
            in_channels,
            channels,
            kernel_size,
            bias=bias,
            padding=(kernel_size - 1) // 2,
        )
        self.sn["upsamples"] = nn.ModuleList()
        self.fn["upsamples"] = nn.ModuleList()
        self.sn["blocks"] = nn.ModuleList()
        self.fn["blocks"] = nn.ModuleList()
        for i in range(len(upsample_kernel_sizes)):
            assert upsample_kernel_sizes[i] == 2 * upsample_scales[i]
            self.sn["upsamples"] += [
                nn.Sequential(
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                    nn.ConvTranspose1d(
                        channels // (2 ** i),
                        channels // (2 ** (i + 1)),
                        upsample_kernel_sizes[i],
                        upsample_scales[i],
                        padding=upsample_scales[i] // 2 + upsample_scales[i] % 2,
                        output_padding=upsample_scales[i] % 2,
                        bias=bias,
                    ),
                )
            ]
            if not share_upsamples:
                self.fn["upsamples"] += [
                    nn.Sequential(
                        getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                        nn.ConvTranspose1d(
                            channels // (2 ** i),
                            channels // (2 ** (i + 1)),
                            upsample_kernel_sizes[i],
                            upsample_scales[i],
                            padding=upsample_scales[i] // 2 + upsample_scales[i] % 2,
                            output_padding=upsample_scales[i] % 2,
                            bias=bias,
                        ),
                    )
                ]
            self.sn["blocks"] += [
                AdaptiveResidualBlock(
                    kernel_size=source_network_params["resblock_kernel_size"],
                    channels=channels // (2 ** (i + 1)),
                    dilations=source_network_params["resblock_dilations"][i],
                    bias=bias,
                    use_additional_convs=source_network_params["use_additional_convs"],
                    nonlinear_activation=nonlinear_activation,
                    nonlinear_activation_params=nonlinear_activation_params,
                )
            ]
            for j in range(len(filter_network_params["resblock_kernel_sizes"])):
                self.fn["blocks"] += [
                    ResidualBlock(
                        kernel_size=filter_network_params["resblock_kernel_sizes"][j],
                        channels=channels // (2 ** (i + 1)),
                        dilations=filter_network_params["resblock_dilations"][j],
                        bias=bias,
                        use_additional_convs=filter_network_params["use_additional_convs"],
                        nonlinear_activation=nonlinear_activation,
                        nonlinear_activation_params=nonlinear_activation_params,
                    )
                ]
        self.sn["output_conv"] = nn.Sequential(
            nn.LeakyReLU(),
            nn.Conv1d(
                channels // (2 ** (i + 1)),
                out_channels,
                kernel_size,
                bias=bias,
                padding=(kernel_size - 1) // 2,
            ),
        )
        self.fn["output_conv"] = nn.Sequential(
            nn.LeakyReLU(),
            nn.Conv1d(
                channels // (2 ** (i + 1)),
                out_channels,
                kernel_size,
                bias=bias,
                padding=(kernel_size - 1) // 2,
            ),
            nn.Tanh(),
        )

        # sine embedding layers
        self.sn["emb"] = Conv1d(
            1,
            channels // (2 ** len(upsample_kernel_sizes)),
            kernel_size,
            bias=bias,
            padding=(kernel_size - 1) // 2,
        )
        # down-sampling CNNs
        self.sn["downsamples"] = nn.ModuleList()
        for i in reversed(range(len(upsample_kernel_sizes))):
            self.sn["downsamples"] += [
                nn.Sequential(
                    nn.Conv1d(
                        channels // (2 ** (i + 1)),
                        channels // (2 ** i),
                        upsample_kernel_sizes[i],
                        upsample_scales[i],
                        padding=upsample_scales[i] - (upsample_kernel_sizes[i] % 2 == 0),
                        bias=bias,
                    ),
                    getattr(nn, nonlinear_activation)(**nonlinear_activation_params),
                )
            ]

        # apply weight norm
        if use_weight_norm:
            self.apply_weight_norm()

        # reset parameters
        self.reset_parameters()

    def forward(self, x, c, d):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input sine signal (B, 1, T).
            c (Tensor): Input tensor (B, in_channels, T).
            d (List): F0-dependent dilation factors [(B, 1, T) x num_upsamples].

        Returns:
            Tensor: Output tensor (B, out_channels, T).

        """
        # currently, same input feature is input to each network
        c = self.input_conv(c)
        e = c

        # source-network forward
        x = self.sn["emb"](x)
        embs = [x]
        for i in range(self.num_upsamples - 1):
            x = self.sn["downsamples"][i](x)
            embs += [x]
        embs2 = []
        for i in range(self.num_upsamples):
            # excitation generation network
            e = self.sn["upsamples"][i](e) + embs[-i - 1]
            e = self.sn["blocks"][i](e, d[i])
            embs2 += [e]
        e = self.sn["output_conv"](e)

        # filter-network forward
        num_blocks = len(self.filter_network_params["resblock_kernel_sizes"])
        for i in range(self.num_upsamples):
            # resonance filtering network
            if self.share_upsamples:
                c = self.sn["upsamples"][i](c) + embs2[i]
            else:
                c = self.fn["upsamples"][i](c) + embs2[i]
            cs = 0.0  # initialize
            for j in range(num_blocks):
                cs += self.fn["blocks"][i * num_blocks + j](c)
            c = cs / num_blocks
        c = self.fn["output_conv"](c)

        return c, e

    def reset_parameters(self):
        """Reset parameters.

        This initialization follows the official implementation manner.
        https://github.com/jik876/hifi-gan/blob/master/models.py

        """

        def _reset_parameters(m):
            if isinstance(m, (nn.Conv1d, nn.ConvTranspose1d)):
                m.weight.data.normal_(0.0, 0.01)
                logger.debug(f"Reset parameters in {m}.")

        self.apply(_reset_parameters)

    def remove_weight_norm(self):
        """Remove weight normalization module from all of the layers."""

        def _remove_weight_norm(m):
            try:
                logger.debug(f"Weight norm is removed from {m}.")
                nn.utils.remove_weight_norm(m)
            except ValueError:  # this module didn't have weight norm
                return

        self.apply(_remove_weight_norm)

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
                nn.utils.weight_norm(m)
                logger.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)

sifigan/models/generator.py
tonnetonne814/SiFi-VITS2-44100-Ja
from sifigan.utils.features import *  # NOQA
from sifigan.utils.index import *  # NOQA
from sifigan.utils.utils import *  # NOQA

sifigan/utils/__init__.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2022 Reo Yoneyama (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Feature-related functions.

References:
    - https://github.com/bigpon/QPPWG
    - https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts

"""

import sys
from logging import getLogger

import numpy as np
import torch
from torch.nn.functional import interpolate

# A logger for this file
logger = getLogger(__name__)


def validate_length(xs, ys=None, hop_size=None):
    """Validate length

    Args:
        xs (ndarray): numpy array of features
        ys (ndarray): numpy array of audios
        hop_size (int): upsampling factor

    Returns:
        (ndarray): length adjusted features

    """
    min_len_x = min([x.shape[0] for x in xs])
    if ys is not None:
        min_len_y = min([y.shape[0] for y in ys])
        if min_len_y < min_len_x * hop_size:
            min_len_x = min_len_y // hop_size
        if min_len_y > min_len_x * hop_size:
            min_len_y = min_len_x * hop_size
        ys = [y[:min_len_y] for y in ys]
    xs = [x[:min_len_x] for x in xs]

    return xs + ys if ys is not None else xs


def dilated_factor(batch_f0, fs, dense_factor):
    """Pitch-dependent dilated factor

    Args:
        batch_f0 (ndarray): the f0 sequence (T)
        fs (int): sampling rate
        dense_factor (int): the number of taps in one cycle

    Return:
        dilated_factors(np array):
            float array of the pitch-dependent dilated factors (T)

    """
    batch_f0[batch_f0 == 0] = fs / dense_factor
    dilated_factors = np.ones(batch_f0.shape) * fs / dense_factor / batch_f0
    assert np.all(dilated_factors > 0)

    return dilated_factors


class SignalGenerator:
    """Input signal generator module."""

    def __init__(
        self,
        sample_rate=24000,
        hop_size=120,
        sine_amp=0.1,
        noise_amp=0.003,
        signal_types=["sine", "noise"],
    ):
        """Initialize WaveNetResidualBlock module.

        Args:
            sample_rate (int): Sampling rate.
            hop_size (int): Hop size of input F0.
            sine_amp (float): Sine amplitude for NSF-based sine generation.
            noise_amp (float): Noise amplitude for NSF-based sine generation.
            signal_types (list): List of input signal types for generator.

        """
        self.sample_rate = sample_rate
        self.hop_size = hop_size
        self.signal_types = signal_types
        self.sine_amp = sine_amp
        self.noise_amp = noise_amp

        for signal_type in signal_types:
            if not signal_type in ["noise", "sine", "sines", "uv"]:
                logger.info(f"{signal_type} is not supported type for generator input.")
                sys.exit(0)
        #logger.info(f"Use {signal_types} for generator input signals.")

    @torch.no_grad()
    def __call__(self, f0):
        signals = []
        for typ in self.signal_types:
            if "noise" == typ:
                signals.append(self.random_noise(f0))
            if "sine" == typ:
                signals.append(self.sinusoid(f0))
            if "sines" == typ:
                signals.append(self.sinusoids(f0))
            if "uv" == typ:
                signals.append(self.vuv_binary(f0))

        input_batch = signals[0]
        for signal in signals[1:]:
            input_batch = torch.cat([input_batch, signal], axis=1)

        return input_batch

    @torch.no_grad()
    def random_noise(self, f0):
        """Calculate noise signals.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: Gaussian noise signals (B, 1, T).

        """
        B, _, T = f0.size()
        noise = torch.randn((B, 1, T * self.hop_size), device=f0.device)

        return noise

    @torch.no_grad()
    def sinusoid(self, f0):
        """Calculate sine signals.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: Sines generated following NSF (B, 1, T).

        """
        B, _, T = f0.size()
        vuv = interpolate((f0 > 0) * torch.ones_like(f0), T * self.hop_size)
        radious = (interpolate(f0.to(torch.float64), T * self.hop_size) / self.sample_rate) % 1
        sine = vuv * torch.sin(torch.cumsum(radious, dim=2) * 2 * np.pi) * self.sine_amp
        if self.noise_amp > 0:
            noise_amp = vuv * self.noise_amp + (1.0 - vuv) * self.noise_amp / 3.0
            noise = torch.randn((B, 1, T * self.hop_size), device=f0.device) * noise_amp
            sine = sine + noise

        return sine

    @torch.no_grad()
    def sinusoids(self, f0):
        """Calculate sines.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: Sines generated following NSF (B, 1, T).

        """
        B, _, T = f0.size()
        vuv = interpolate((f0 > 0) * torch.ones_like(f0), T * self.hop_size)
        f0 = interpolate(f0.to(torch.float64), T * self.hop_size)
        sines = torch.zeros_like(f0, device=f0.device)
        harmonics = 5  # currently only fixed number of harmonics is supported
        for i in range(harmonics):
            radious = (f0 * (i + 1) / self.sample_rate) % 1
            sines += torch.sin(torch.cumsum(radious, dim=2) * 2 * np.pi)
        sines = self.sine_amp * sines * vuv / harmonics
        if self.noise_amp > 0:
            noise_amp = vuv * self.noise_amp + (1.0 - vuv) * self.noise_amp / 3.0
            noise = torch.randn((B, 1, T * self.hop_size), device=f0.device) * noise_amp
            sines = sines + noise

        return sines

    @torch.no_grad()
    def vuv_binary(self, f0):
        """Calculate V/UV binary sequences.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: V/UV binary sequences (B, 1, T).

        """
        _, _, T = f0.size()
        uv = interpolate((f0 > 0) * torch.ones_like(f0), T * self.hop_size)

        return uv

sifigan/utils/features.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2020 Yi-Chiao Wu (Nagoya University)
#  MIT License (https://opensource.org/licenses/MIT)

"""Indexing-related functions."""

import torch


def pd_indexing(x, d, dilation, batch_index, ch_index):
    """Pitch-dependent indexing of past and future samples.

    Args:
        x (Tensor): Input feature map (B, C, T).
        d (Tensor): Input pitch-dependent dilated factors (B, 1, T).
        dilation (Int): Dilation size.
        batch_index (Tensor): Batch index
        ch_index (Tensor): Channel index

    Returns:
        Tensor: Past output tensor (B, out_channels, T)
        Tensor: Future output tensor (B, out_channels, T)

    """
    B, C, T = x.size()
    batch_index = torch.arange(0, B, dtype=torch.long, device=x.device).reshape(B, 1, 1)
    ch_index = torch.arange(0, C, dtype=torch.long, device=x.device).reshape(1, C, 1)
    dilations = torch.clamp((d * dilation).long(), min=1)

    # get past index (assume reflect padding)
    idx_base = torch.arange(0, T, dtype=torch.long, device=x.device).reshape(1, 1, T)
    idxP = (idx_base - dilations).abs() % T
    idxP = (batch_index, ch_index, idxP)

    # get future index (assume reflect padding)
    idxF = idx_base + dilations
    overflowed = idxF >= T
    idxF[overflowed] = -(idxF[overflowed] % T)
    idxF = (batch_index, ch_index, idxF)

    return x[idxP], x[idxF]


def index_initial(n_batch, n_ch, tensor=True):
    """Tensor batch and channel index initialization.

    Args:
        n_batch (Int): Number of batch.
        n_ch (Int): Number of channel.
        tensor (bool): Return tensor or numpy array

    Returns:
        Tensor: Batch index
        Tensor: Channel index

    """
    batch_index = []
    for i in range(n_batch):
        batch_index.append([[i]] * n_ch)
    ch_index = []
    for i in range(n_ch):
        ch_index += [[i]]
    ch_index = [ch_index] * n_batch

    if tensor:
        batch_index = torch.tensor(batch_index)
        ch_index = torch.tensor(ch_index)
        if torch.cuda.is_available():
            batch_index = batch_index.cuda()
            ch_index = ch_index.cuda()
    return batch_index, ch_index

sifigan/utils/index.py
tonnetonne814/SiFi-VITS2-44100-Ja
# -*- coding: utf-8 -*-

# Copyright 2020 Yi-Chiao Wu (Nagoya University)
# based on a Parallel WaveGAN script by Tomoki Hayashi (Nagoya University)
# (https://github.com/kan-bayashi/ParallelWaveGAN)
#  MIT License (https://opensource.org/licenses/MIT)

"""Utility functions."""

import os
import sys
from logging import getLogger

import h5py
import numpy as np

# A logger for this file
logger = getLogger(__name__)


def read_hdf5(hdf5_name, hdf5_path):
    """Read hdf5 dataset.

    Args:
        hdf5_name (str): Filename of hdf5 file.
        hdf5_path (str): Dataset name in hdf5 file.

    Return:
        any: Dataset values.

    """
    if not os.path.exists(hdf5_name):
        logger.error(f"There is no such a hdf5 file ({hdf5_name}).")
        sys.exit(1)

    hdf5_file = h5py.File(hdf5_name, "r")

    if hdf5_path not in hdf5_file:
        logger.error(f"There is no such a data in hdf5 file. ({hdf5_path})")
        sys.exit(1)

    hdf5_data = hdf5_file[hdf5_path][()]
    hdf5_file.close()

    return hdf5_data


def write_hdf5(hdf5_name, hdf5_path, write_data, is_overwrite=True):
    """Write dataset to hdf5.

    Args:
        hdf5_name (str): Hdf5 dataset filename.
        hdf5_path (str): Dataset path in hdf5.
        write_data (ndarray): Data to write.
        is_overwrite (bool): Whether to overwrite dataset.

    """
    # convert to numpy array
    write_data = np.array(write_data)

    # check folder existence
    folder_name, _ = os.path.split(hdf5_name)
    if not os.path.exists(folder_name) and len(folder_name) != 0:
        os.makedirs(folder_name)

    # check hdf5 existence
    if os.path.exists(hdf5_name):
        # if already exists, open with r+ mode
        hdf5_file = h5py.File(hdf5_name, "r+")
        # check dataset existence
        if hdf5_path in hdf5_file:
            if is_overwrite:
                hdf5_file.__delitem__(hdf5_path)
            else:
                logger.error(
                    "Dataset in hdf5 file already exists. "
                    "if you want to overwrite, please set is_overwrite = True."
                )
                hdf5_file.close()
                sys.exit(1)
    else:
        # if not exists, open with w mode
        hdf5_file = h5py.File(hdf5_name, "w")

    # write data to hdf5
    hdf5_file.create_dataset(hdf5_path, data=write_data)
    hdf5_file.flush()
    hdf5_file.close()


def check_hdf5(hdf5_name, hdf5_path):
    """Check hdf5 file existence

    Args:
        hdf5_name (str): filename of hdf5 file
        hdf5_path (str): dataset name in hdf5 file

    Return:
        (bool): dataset exists then return true

    """
    if not os.path.exists(hdf5_name):
        return False
    else:
        with h5py.File(hdf5_name, "r") as f:
            if hdf5_path in f:
                return True
            else:
                return False


def read_txt(file_list):
    """Read .txt file list

    Arg:
        file_list (str): txt file filename

    Return:
        (list): list of read lines

    """
    with open(file_list, "r") as f:
        filenames = f.readlines()
    return [filename.replace("\n", "") for filename in filenames]


def check_filename(list1, list2):
    """Check the filenames of two list are matched

    Arg:
        list1 (list): file list 1
        list2 (list): file list 2

    Return:
        (bool): matched (True) or not (False)

    """

    def _filename(x):
        return os.path.basename(x).split(".")[0]

    list1 = list(map(_filename, list1))
    list2 = list(map(_filename, list2))

    return list1 == list2

sifigan/utils/utils.py
tonnetonne814/SiFi-VITS2-44100-Ja
import utaupy
import music21 as m21
import mido
import pyopenjtalk
import shutil
from time import sleep
import os
import torch

def get_midi_info(path):
    midi = mido.MidiFile(path)
    for tracks in midi.tracks:
        track_name = tracks.name

        # GET BPM
        lines = list()
        if track_name == "Control":
            for track in tracks:
                try:
                    if "@set tempo" in track.text:
                        _, BPM =  track.text .split("=")
                except:
                    pass 
        # GET notes and dur
        else:
            for track in tracks:
                try:
                    if float(track.time) > 0:
                        tick_dur = float(track.time)
                        try:
                            note = str(track.note)
                        except:
                            note = "rest"
                        line = {"Note" : note, "tick_dur" : tick_dur, "ms_dur":tick2ms(float(BPM), tick_dur)}
                        lines.append(line)
                except:
                    continue 
    #print(midi)
    return BPM, lines

def get_ust_info(path): 
# Note Number : https://www.asahi-net.or.jp/~hb9t-ktd/music/Japan/Research/DTM/freq_map.html
    ustobj = utaupy.ust.load(path)
    lines = list()
    for note in ustobj.notes:
        length = note.length
        length_ms = note.length_ms
        notename = note.notename
        note_MIDI = note.notenum
        lyric = note.lyric

        line = {"length"    : length,
                "length_ms" : length_ms,
                "notename"  : notename,
                "note_MIDI" : note_MIDI,
                "lyric"     : lyric}
        lines.append(line)
    
    return lines

def get_lab_info(path):
    with open(path, mode="r", encoding="utf-8") as f:
        lines_r = f.readlines()
    lines = list()
    for line_r in lines_r:
        line_r = line_r.replace("\n", "")
        start_time_100us, end_time_100us, ph = line_r.split(" ")
        t_start = float(start_time_100us)/10000
        t_end = float(end_time_100us)/10000
        line = {"start_ms":t_start,
                "end_ms": t_end,
                "dur_ms": float(t_end - t_start),
                "ph": ph}
        lines.append(line)

    return lines


def get_lab_info_no_scaling(path):
    with open(path, mode="r", encoding="utf-8") as f:
        lines_r = f.readlines()
    lines = list()
    for line_r in lines_r:
        line_r = line_r.replace("\n", "")
        start_time_100us, end_time_100us, ph = line_r.split(" ")
        t_start = int(start_time_100us)
        t_end = int(end_time_100us)
        line = {"start_ms":t_start,
                "end_ms": t_end,
                "dur_ms": int(t_end - t_start),
                "ph": ph}
        lines.append(line)

    return lines

def get_musicxml_info(filepath):
    try:
        piece = m21.converter.parse(filepath)
        lines = list()

        #音またはrestをlist_noteに、長さをlist_durationに格納する
        for n in piece.flat.notesAndRests:
            if type(n) == m21.note.Note:
                pitch = str(n.pitch)
                beats_dur = str(n.duration.quarterLength)
                word_lyric = str(n.lyric)
                sec_dur = float(n.seconds)
            elif type(n) == m21.note.Rest:
                pitch = str(n.name)
                beats_dur = str(n.duration.quarterLength)
                word_lyric = "rest"
                sec_dur = float(n.seconds)
            else:
                print(f"ERROR : {filepath}  END.")
                exit()
            line = {"pitch" : pitch, 
                    "beats_dur" : beats_dur,
                    "word_lyric" : word_lyric,
                    "ms_dur" : str(sec_dur*1000)}
            lines.append(line)

    except ZeroDivisionError:
        print('Zero Division Error')

    return lines 

def get_g2p_dict_from_training_data(word_list_path):
    word_list = torch.load(word_list_path)
    word_symbol_to_id = {s.replace("\n", ""): i for i, s in enumerate(sorted(set(word_list)))} # oto table用
    id_to_word_symbol = {i: s.replace("\n", "") for i, s in enumerate(sorted(set(word_list)))} # oto table用
    return word_symbol_to_id, id_to_word_symbol

def get_g2p_dict_from_tabledata(table_path, include_converter=False):
    with open(table_path, mode="r", encoding="utf-8") as f:
        lines = f.readlines()
    g2p_dict = {}
    ph_list = list()
    oto_list = list()
    for line in lines:
        line = line.split(" ")
        oto = line.pop(0).replace("\n","")
        oto_list.append(oto)
        ph = " ".join(line).replace("\n","")
        ph_list.extend(line)
        g2p_dict[oto] = ph
    if include_converter is True:
        ph_symbol_to_id = {s.replace("\n", ""): i for i, s in enumerate(sorted(set(ph_list)))}
        id_to_ph_symbol = {i: s.replace("\n", "") for i, s in enumerate(sorted(set(ph_list)))}
        word_symbol_to_id = {s.replace("\n", ""): i for i, s in enumerate(sorted(set(oto_list)))} # oto table用
        id_to_word_symbol = {i: s.replace("\n", "") for i, s in enumerate(sorted(set(oto_list)))} # oto table用
        return g2p_dict, ph_symbol_to_id, id_to_ph_symbol, word_symbol_to_id, id_to_word_symbol
    else:
        return g2p_dict

def g2p_pyopenjtalk(kana):
    return pyopenjtalk.g2p(kana)

def tick2ms(BPM, tick):
    tick_units = (60/BPM) *1000 /480
    ms = tick_units * tick
    return ms

def get_annotated_data(ust_path, lab_path, ono2lab_table_path):
    lab_collate(lab_path) # silを完全消去→pauへ
    data_UST = get_ust_info(ust_path)
    data_LAB = get_lab_info(lab_path)
    g2p_dict = get_g2p_dict_from_tabledata(ono2lab_table_path)

    lab_lines = list()
    for line in data_LAB:
        lab_lines.append({"ph":line["ph"] , "ph_start_point_ms" :line["start_ms"], "ph_end_point_ms" :line["end_ms"]})

    word_lines = list()
    for line in data_UST:
        text = line["lyric"]
        # "とっ"のような単語が入ってくるが、"と"、"っ"の二つで対応するための処理
        if "っ" in text and len(text)>1:
            try:
                word_ph = g2p_dict[text] # もし辞書に有ればそのまま使う。
            except:
                txt = ""   
                word_ph = ""   
                # check cl
                for w in text:
                    if w == "っ":
                        word_ph += g2p_dict[txt] + " " + g2p_dict[w] # pull from dict
                        txt = "" #reset
                    else:
                        txt += w # stack
                # last pull
                if txt != "":
                    word_ph += g2p_dict[txt]
        else:      
            word_ph = g2p_dict[text]

        if text == "R":
            note_num = -1
            notename = "rest"
        else:
            note_num = line["note_MIDI"]
            notename = line["notename"]
        note_dur = line["length_ms"]

        # 大文字で判別
        start_point = 0 
        end_point = 0
        ph_list = list()
        for idx , WORD_PH_UNIT in enumerate(word_ph.upper().split(" ")):
            line_lab = lab_lines.pop(0) # 一個ずつ出していく
            LAB_PH = line_lab["ph"].upper()

            # 初回なら開始位置保存
            if idx == 0:
                start_point = line_lab["ph_start_point_ms"]
            
            # ワードが一致するまで終了位置を保存
            if WORD_PH_UNIT == LAB_PH:
                end_point = line_lab["ph_end_point_ms"]

            elif LAB_PH == "SIL":
                end_point = line_lab["ph_end_point_ms"]
                text = "R" 

            # 一致しなかったら強制終了
            else:
                print("ERROR ph mismatch")
                exit()
            ph_list.append([   line_lab["ph"],                  \
                               line_lab["ph_start_point_ms"],   \
                               line_lab["ph_end_point_ms"],     \
                               line_lab["ph_end_point_ms"] - line_lab["ph_start_point_ms"] \
                               ])
        word_line = {"word": text, 
                     "ph": ph_list ,
                     "start_ms":start_point, 
                     "end_ms": end_point, 
                     "dur_ms": end_point-start_point,
                     "note_num": note_num,
                     "notename": notename,
                     "note_dur": note_dur}
        word_lines.append(word_line)
    return word_lines


def lab_collate(path):

    lab_list = get_lab_info_no_scaling(path)
    temp_txt = "temp.lab"
    z_line   = lab_list.pop(0)
    z_s_time = z_line["start_ms"]
    z_e_time = z_line["end_ms"]
    z_ph     = z_line["ph"]
    with open(temp_txt, mode="a", encoding="utf-8") as f:
        f.write(str(z_s_time) + " ")
        for  line in lab_list:
            #s_time, e_time, ph = line
            s_time = line["start_ms"]
            e_time = line["end_ms"]
            ph = line["ph"]

            # z_e_timeとs_timeの書き込みは、pau or silが連続していない時のみ。
            if z_ph == "sil" or z_ph =="pau":
                if ph == "sil" or ph == "pau":
                    pass
                else:
                    f.write(str(z_e_time) + " " + z_ph + "\n" + str(s_time) + " ")
            else:
                f.write(str(z_e_time) + " " + z_ph + "\n" + str(s_time) + " ")
            z_s_time = s_time
            z_e_time = e_time
            z_ph = ph
        f.write(str(e_time) + " " + ph + "\n")
    sleep(1)
    shutil.copy(temp_txt, path)
    os.remove(temp_txt)
    


if __name__ == "__main__":
    path = "./kana2phonemes_002_oto2lab.table"
    get_g2p_dict_from_tabledata(path)
    #lab_collate("./DATABASE/1st_color/1st_color.lab")
    


singDB_loader.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler

from models import DiffusionModels

######### export TORCH_CUDNN_V8_API_DISABLED=1
torch.cuda.empty_cache()
import torch._dynamo.config
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.verbose=True

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

import utils
from data_utils import TextAudioLoader_SVS_f0diff, TextAudioCollate_SVS_f0diff, DistributedBucketSampler

#torch.autograd.set_detect_anomaly(True) # 逆誤差伝搬のエラー部分を自動で見つける機能
torch.backends.cudnn.benchmark = True # 学習早くなるらしい（知らんけど
global_step = 0

import warnings
warnings.filterwarnings(action='ignore')

def main():

  """ GPUのみの学習を許可 """
  assert torch.cuda.is_available(), "CPU training is not allowed."

  n_gpus = torch.cuda.device_count()
  os.environ['MASTER_ADDR'] = 'localhost'
  os.environ['MASTER_PORT'] = '65520'

  hps = utils.get_hparams()   # config.yamlやargparseは、ここで読み込まれている。

  if hps["DDP_use"] is True:
    mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))  # DDP ON
  else:
    run(rank=0, n_gpus=1, hps=hps)                     # DDP OFF

def run(rank, n_gpus, hps):
  global global_step
  if rank == 0: # 1番目のGPUであれば入る（GPUが複数ある場合を想定
    logger = utils.get_logger(hps["model_dir"])
    logger.info(hps)
    writer = SummaryWriter(log_dir=hps["model_dir"])
    writer_eval = SummaryWriter(log_dir=os.path.join(hps["model_dir"], "eval"))

  # Linux Only
  if hps["DDP_use"] is True: 
    dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)   

  # 乱数の再現性確保
  torch.manual_seed(hps["seed"])
  torch.cuda.set_device(rank)

  # データセットの読み込み君 
  train_dataset = TextAudioLoader_SVS_f0diff(audiopaths_and_text= hps["train_data_path"] ,
                                  hparams            = hps)
  vocab_size = train_dataset.get_ph_vocab_size()
  
  # 似た長さのデータをモデルに流す君
  train_sampler = DistributedBucketSampler(dataset     = train_dataset,
                                           batch_size  = hps["batch_size"],
                                           boundaries  = [32,300,400,500,600,700,800,900,1000],
                                           num_replicas= n_gpus,
                                           rank        = rank,
                                           shuffle     = True) # シャッフル指定

  # 入力するデータに0を加えて長さを統一する君
  collate_fn = TextAudioCollate_SVS_f0diff()

  # 上三つを統べる
  train_loader = DataLoader(dataset=train_dataset, 
                            num_workers=hps["num_workers"],     # データ読み込みに使用するCPUのスレッド数
                            shuffle=False, 
                            pin_memory=True,    # Trueだとなんか動作早くなるらしい（知らんけど
                            collate_fn=collate_fn, 
                            batch_sampler=train_sampler)

  # 1番目のGPUであれば
  if rank == 0:
    eval_dataset = TextAudioLoader_SVS_f0diff(audiopaths_and_text=hps["eval_data_path"], 
                                              hparams            =hps)
    
    eval_loader = DataLoader(dataset      = eval_dataset, 
                             num_workers  = 1, 
                             shuffle      = False,
                             batch_size   = hps["batch_size"], 
                             pin_memory   = True,
                             drop_last    = False, 
                             collate_fn   = collate_fn)

  # モデルを作成
  net_g = DiffusionModels(hps=hps, vocab_size=vocab_size).cuda(rank)
  #net_g = torch.compile(net_g)

  # 最適化手法
  optim_g = torch.optim.AdamW(params      = net_g.parameters(),
                              lr          = float(hps['learning_rate']), 
                              betas       =       hps["betas"], 
                              eps         = float(hps["eps"]))
 
  # checkpoint 確認
  in_progress_checkpoint_path = utils.latest_checkpoint_path(hps["model_dir"], "G_*.pth")
  if os.path.isfile(in_progress_checkpoint_path): # 学習再開
    try:
      _, _, _, epoch_str = utils.load_checkpoint(in_progress_checkpoint_path, net_g, optim_g)
      global_step = (epoch_str - 1) * len(train_loader)
      print("### Resume learning mode ###")
    except:
      net_g = utils.load_model_diffsize(in_progress_checkpoint_path, net_g, hps, optim_g)
      epoch_str = 1
      global_step = 0.
      print("### Resume learning mode ###")
  elif hps["is_finetune"] is True: # ファインチューニング
    net_g = utils.load_model_diffsize(hps["finetune_G"], net_g, hps, optim_g)
    epoch_str = 1
    global_step = 0
    print("### FineTuning mode ###")
  else: # どちらでもない→初期状態から
    epoch_str = 1
    global_step = 0
    print("### Initial state mode ###")
  global_step = (epoch_str - 1) * len(train_loader)

  # エポックに応じて自動で学習率を変える
  scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_g, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)

  # float16で計算して，高速化するか
  scaler = GradScaler(enabled=hps["fp16_run"])

  # Linux Only
  if hps["DDP_use"] is True: 
    net_g = DDP(net_g, device_ids=[rank]) 

  # 学習/評価処理
  for epoch in range(epoch_str, int(hps["max_epochs"]) + 1):
    if rank==0:
      train_and_evaluate(rank, epoch, hps, net_g, optim_g, scheduler_g, scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
    else:
      train_and_evaluate(rank, epoch, hps, net_g, optim_g, scheduler_g, scaler, [train_loader, None], None, None)
    scheduler_g.step()

  # 最終モデルの保存
  utils.save_checkpoint(net_g, optim_g, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "G_LastModel.pth".format(global_step)))

def train_and_evaluate(rank, epoch, hps, net_g, optim_g, scheduler_g, scaler, loaders, logger, writers):

  train_loader, eval_loader = loaders
  if writers is not None:
    writer, writer_eval = writers

  train_loader.batch_sampler.set_epoch(epoch)
  global global_step

  net_g.train() # 学習モードへ
  for batch_idx, (f0,f0_lengths,\
                  ph_IDs,ph_IDs_lengths,\
                  ph_frame_dur,\
                  noteID,noteID_lengths,
                  speakerID) in enumerate(train_loader):

    # Forward Process
    with autocast(enabled=hps["fp16_run"]):
      loss_f0 = net_g(f0=f0.cuda(rank),  
                      f0_len=f0_lengths.cuda(rank),
                      IDs=ph_IDs.cuda(rank), 
                      IDs_len=ph_IDs_lengths.cuda(rank),
                      IDs_dur=ph_frame_dur.cuda(rank),
                      NoteIDs=noteID.cuda(rank),
                      NoteID_len=noteID_lengths.cuda(rank),
                      g=speakerID.cuda(rank))
    #losses = loss_f0 + loss_vuv
    losses = loss_f0
    optim_g.zero_grad()
    scaler.scale(losses).backward()   
    scaler.unscale_(optim_g)
    scaler.step(optim_g)
    scaler.update()

    # 記録 / 評価
    if rank==0:
      if global_step % hps["log_interval"] == 0:
        lr = optim_g.param_groups[0]['lr']
        #losses = 
        logger.info('Train Epoch: {} [{:.0f}%]'.format(
          epoch,
          100. * batch_idx / len(train_loader)))
        #logger.info([x.item() for x in [losses, loss_f0, loss_vuv]] + [global_step, lr])
        logger.info([x.item() for x in [losses]] + [global_step, lr])
        #scalar_dict = {"loss": losses, "loss_f0" : loss_f0, "loss_vuv":loss_vuv}
        scalar_dict = {"loss_f0": losses}
        
        utils.summarize(
          writer=writer,
          global_step=global_step, 
          scalars=scalar_dict)
        
      if global_step % hps["eval_interval"] == 0:
        evaluate(hps=hps,
                 net_g=net_g,
                 eval_loader=eval_loader, 
                 writer_eval=writer_eval)
        utils.save_checkpoint(model=net_g, 
                              optimizer=optim_g, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "G_{}.pth".format(global_step)))
    global_step += 1

  if rank == 0:
    logger.info('====> Epoch: {}'.format(epoch))
  
def evaluate(hps, net_g, eval_loader, writer_eval):
    net_g.eval()
    rank = 0
    count = 0
    loss_f0_sum = 0
    with torch.no_grad():
      for batch_idx, (f0,f0_lengths,\
                      ph_IDs,ph_IDs_lengths,\
                      ph_frame_dur,\
                      noteID,noteID_lengths,
                      speakerID) in enumerate(eval_loader):
        # Forward Process
        with autocast(enabled=hps["fp16_run"]):
          loss_f0 = net_g(f0=f0.cuda(rank),  
                          f0_len=f0_lengths.cuda(rank),
                          IDs=ph_IDs.cuda(rank), 
                          IDs_len=ph_IDs_lengths.cuda(rank),
                          IDs_dur=ph_frame_dur.cuda(rank),
                          NoteIDs=noteID.cuda(rank),
                          NoteID_len=noteID_lengths.cuda(rank),
                          g=speakerID.cuda(rank))
          loss_f0_sum += loss_f0
          count += 1

    f0_gt=f0                  [0][0].view(1, 1,-1).cuda(rank)      
    f0_gt_len=f0_lengths      [0]   .view(-1)     .cuda(rank)            
    ph_IDs=ph_IDs             [0]   .view(1,-1)   .cuda(rank)                
    ph_IDs_len=ph_IDs_lengths [0]   .view(-1)     .cuda(rank)       
    ph_IDs_dur=ph_frame_dur   [0]   .view(1,-1)   .cuda(rank)  
    NoteIDs=noteID            [0]   .view(1,-1)   .cuda(rank)   
    NoteID_len=noteID_lengths [0]   .view(-1)     .cuda(rank)     
    speakerID=speakerID       [0]   .view(-1)     .cuda(rank)     

    f0_pd = net_g.sampling(condition=[ph_IDs, ph_IDs_len, ph_IDs_dur, 
                                      NoteIDs, NoteID_len, None, speakerID])

    scalar_dict = {"loss_f0": loss_f0_sum/count}
    #scalar_dict = {"loss": loss_f0_sum + loss_vuv_sum,
    #               "loss_f0": loss_f0_sum,
    #               "loss_vuv" : loss_vuv_sum}

    f0_gt = f0_gt[0][0].to('cpu').detach().numpy().copy()
    f0_pd = f0_pd[0][0].to('cpu').detach().numpy().copy() 
    #vuv_gt= vuv_gt.view(-1).to('cpu').detach().numpy().copy() 
    #vuv_pd= vuv_pd.view(-1).to('cpu').detach().numpy().copy() 

    f0_gt_image= generate_graph(vector=f0_gt, 
                                label="F0 gt",
                                color="red", 
                                x_label = 'Frames',
                                y_label = "f0")

    f0_pd_image= generate_graph(vector=f0_pd , 
                                label="F0 pd",
                                color="blue", 
                                x_label = 'Frames',
                                y_label = "f0")
    #vuv_gt_image=generate_graph(vector=vuv_gt, 
    #                            label="Voice Unvoice Value",
    #                            color="red", 
    #                            x_label = 'Frames',
    #                            y_label = "0=mask, 1=uv, 2=v")
    #vuv_pd_image= generate_graph(vector=vuv_pd, 
    #                            label="Voice Unvoice Value",
    #                            color="blue", 
    #                            x_label = 'Frames',
    #                            y_label = "0=mask, 1=uv, 2=v",)
    
    image_dict = { 
    "f0/gt":f0_gt_image ,
    "f0/pd":f0_pd_image ,
    #"vuv/gt":vuv_gt_image ,
    #"vuv/pd":vuv_pd_image ,
        }
    
    utils.summarize(
      writer=writer_eval,
      global_step=global_step, 
      scalars=scalar_dict,
      images=image_dict
    )
    net_g.train()

def generate_graph(vector, 
                   label="NoName",
                   color="blue", 
                   title="Title",
                   x_label = 'Frames',
                   y_label = "y_labels",
                   figsize=(20,5)):
   
    fig = plt.figure(figsize=figsize) 
    x = np.arange(0, len(vector))
    plt.plot(x, vector, label=label, color=color)
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    #plt.ylim(y_val_min, y_val_max)
    fig.canvas.draw() 
    plot_image = fig.canvas.renderer._renderer  
    image_numpy = np.array(plot_image)
    plt.clf()
    plt.close()
    return image_numpy

# このファイルを実行したら、ココから始まる                  
if __name__ == "__main__":
  os.environ[
        "TORCH_DISTRIBUTED_DEBUG"
    ] = "DETAIL"
  main()

train_SVS_F0Diffusion.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler

from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
import commons

######### export TORCH_CUDNN_V8_API_DISABLED=1
torch.cuda.empty_cache()
import torch._dynamo.config
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.verbose=True

from models import VITS2_based_SiFiSinger
from sifigan.models.discriminator import UnivNetMultiResolutionMultiPeriodDiscriminator
from models import DurationDiscriminator

from sifigan.losses.reg import ResidualLoss

from losses import generator_loss, discriminator_loss, feature_loss,kl_loss
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import utils
from data_utils import TextAudioLoader_SVS_synth, TextAudioCollate_SVS_synth, DistributedBucketSampler

# torch.autograd.set_detect_anomaly(True) # 逆誤差伝搬のエラー部分を自動で見つける機能
torch.backends.cudnn.benchmark = True # 学習早くなるらしい（知らんけど
global_step = 0

import warnings
warnings.filterwarnings(action='ignore')

def main():

  """ GPUのみの学習を許可 """
  assert torch.cuda.is_available(), "CPU training is not allowed."

  n_gpus = torch.cuda.device_count()
  os.environ['MASTER_ADDR'] = 'localhost'
  os.environ['MASTER_PORT'] = '65520'

  hps = utils.get_hparams()   # config.yamlやargparseは、ここで読み込まれている。

  if hps["DDP_use"] is True:
    mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))  # DDP ON(未確認)
  else:
    run(rank=0, n_gpus=1, hps=hps)                     # DDP OFF

def run(rank, n_gpus, hps):
  global global_step
  if rank == 0: # 1番目のGPUであれば入る（GPUが複数ある場合を想定
    logger = utils.get_logger(hps["model_dir"])
    logger.info(hps)
    # utils.check_git_hash(hps["model_dir"])
    writer = SummaryWriter(log_dir=hps["model_dir"])
    writer_eval = SummaryWriter(log_dir=os.path.join(hps["model_dir"], "eval"))

  # Linux Only
  if hps["DDP_use"] is True: 
    dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)   

  # 乱数の再現性確保
  torch.manual_seed(hps["seed"])
  torch.cuda.set_device(rank)

  # データセットの読み込み君 
  train_dataset = TextAudioLoader_SVS_synth(audiopaths_and_text= hps["train_data_path"] ,
                                  hparams            = hps)
  vocab_size = train_dataset.get_ph_vocab_size()
  
  # 似た長さのデータをモデルに流す君
  train_sampler = DistributedBucketSampler(dataset     = train_dataset,
                                           batch_size  = hps["batch_size"],
                                           boundaries  = [32,300,400,500,600,700,800,900,1000],
                                           num_replicas= n_gpus,
                                           rank        = rank,
                                           shuffle     = True) # シャッフル指定

  # 入力するデータに0を加えて長さを統一する君
  collate_fn = TextAudioCollate_SVS_synth(hparams=hps)

  # 上三つを統べる
  train_loader = DataLoader(dataset=train_dataset, 
                            num_workers=hps["num_workers"],     # データ読み込みに使用するCPUのスレッド数
                            shuffle=False, 
                            pin_memory=True,    # Trueだとなんか動作早くなるらしい（知らんけど
                            collate_fn=collate_fn, 
                            batch_sampler=train_sampler)

  # 1番目のGPUであれば
  if rank == 0:
    eval_dataset = TextAudioLoader_SVS_synth(audiopaths_and_text=hps["eval_data_path"], 
                                   hparams            =hps)

    eval_loader = DataLoader(dataset      = eval_dataset, 
                             num_workers  = 1, 
                             shuffle      = False,
                             batch_size   = hps["batch_size"], 
                             pin_memory   = True,
                             drop_last    = False, 
                             collate_fn   = collate_fn)

  net_dur_disc = DurationDiscriminator(
    hps["dur_discriminator"]["hidden_channels"], 
    hps["dur_discriminator"]["hidden_channels"], 
    3, 
    0.1, 
    gin_channels=hps["common"]["gin_channels"] if hps["common"]["n_speaker"] != 0 else 0,
    ).cuda(rank)

  net_g = VITS2_based_SiFiSinger(
      hps = hps,
      n_vocab=vocab_size,
      spec_channels=hps["spec_encoder"]["spec_channels"],
      segment_size=int(hps["segments_size"] // hps["hop_length"]),
      n_speakers=hps["common"]["n_speaker"],
      gin_channels=hps["common"]["gin_channels"],
      **hps["hifi_gan"]).cuda(rank)


  net_d = UnivNetMultiResolutionMultiPeriodDiscriminator(**hps["UnivNet_MRMPDiscriminator"]).cuda(rank)
  
  # torch.compile 学習高速化
  #net_g = torch.compile(net_g)
  #net_d = torch.compile(net_d)
  #if net_dur_disc is not None:
  #  net_dur_disc = torch.compile(net_dur_disc)
  
  # 最適化手法
  optim_g = torch.optim.AdamW(params      = net_g.parameters(),
                              lr          = float(hps['learning_rate']), 
                              betas       =       hps["betas"], 
                              eps         = float(hps["eps"]))
  optim_d = torch.optim.AdamW(params      = net_d.parameters(),
                              lr          = float(hps['learning_rate']), 
                              betas       =       hps["betas"], 
                              eps         = float(hps["eps"]))
  if net_dur_disc is not None:
    optim_dur_disc = torch.optim.AdamW(params      = net_dur_disc.parameters(),
                                       lr          = float(hps['learning_rate']), 
                                       betas       =       hps["betas"], 
                                       eps         = float(hps["eps"]))
  else:
    optim_dur_disc = None
 
  # checkpoint 確認
  in_progress_checkpoint_path = utils.latest_checkpoint_path(hps["model_dir"], "G_*.pth")
  if os.path.isfile(in_progress_checkpoint_path): # 学習再開
    try:
      _, _, _, epoch_str = utils.load_checkpoint(in_progress_checkpoint_path, net_g, optim_g)
      _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps["model_dir"], "D_*.pth"), net_d, optim_d)
      if net_dur_disc is not None:
        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps["model_dir"], "DUR_*.pth"), net_dur_disc, optim_dur_disc)
      global_step = (epoch_str - 1) * len(train_loader)
      print("### Resume learning mode ###")
    except:
      net_g = utils.load_model_diffsize(in_progress_checkpoint_path, net_g, hps, optim_g)
      net_d = utils.load_model_diffsize(utils.latest_checkpoint_path(hps["model_dir"], "D_*.pth"), net_d, hps, optim_d)
      if net_dur_disc is not None:
        net_dur_disc = utils.load_model_diffsize(utils.latest_checkpoint_path(hps["model_dir"], "DUR_*.pth"), net_dur_disc, hps, optim_dur_disc)
      epoch_str = 1
      global_step = 0.
      print("### Resume learning mode ###")
  elif hps["is_finetune"] is True: # ファインチューニング
    net_g = utils.load_model_diffsize(hps["finetune_G"], net_g, hps, optim_g)
    net_d = utils.load_model_diffsize(hps["finetune_D"], net_d, hps, optim_d)
    if net_dur_disc is not None:
      net_dur_disc = utils.load_model_diffsize(hps["finetune_DUR"], net_dur_disc, hps, optim_dur_disc)
    epoch_str = 1
    global_step = 0
    print("### FineTuning mode ###")
  else: # どちらでもない→初期状態から
    epoch_str = 1
    global_step = 0
    print("### Initial state mode ###")
  global_step = (epoch_str - 1) * len(train_loader)

  # エポックに応じて自動で学習率を変える
  scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_g, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)
  scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_d, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)
  if net_dur_disc is not None:
    scheduler_dur_disc = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_dur_disc, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)
  else:
    scheduler_dur_disc = None

  # float16で計算して，高速化するか
  scaler = GradScaler(enabled=hps["fp16_run"])

  # Linux Only
  if hps["DDP_use"] is True: 
    net_g = DDP(net_g, device_ids=[rank]) 
    net_d = DDP(net_d, device_ids=[rank]) 
    if net_dur_disc is not None:
      net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True)

  # 学習/評価処理
  for epoch in range(epoch_str, int(hps["max_epochs"]) + 1):
    if rank==0:
      train_and_evaluate(rank, epoch, hps, [net_g, net_d, net_dur_disc], [optim_g, optim_d, optim_dur_disc], [scheduler_g, scheduler_d, scheduler_dur_disc], scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
    else:
      train_and_evaluate(rank, epoch, hps, [net_g, net_d, net_dur_disc], [optim_g, optim_d, optim_dur_disc], [scheduler_g, scheduler_d, scheduler_dur_disc], scaler, [train_loader, None], None, None)
    scheduler_g.step()
    scheduler_d.step()
    if net_dur_disc is not None:
      scheduler_dur_disc.step()

  # 最終モデルの保存
  utils.save_checkpoint(net_g, optim_g, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "G_LastModel.pth".format(global_step)))
  utils.save_checkpoint(net_d, optim_d, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "D_LastModel.pth".format(global_step)))
  if net_dur_disc is not None:
    utils.save_checkpoint(net_dur_disc, optim_dur_disc, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "DUR_LastModel.pth".format(global_step)))

def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers):
  net_g, net_d, net_dur_disc = nets
  optim_g, optim_d, optim_dur_disc = optims
  scheduler_g, scheduler_d, scheduler_dur_disc = schedulers
  
  train_loader, eval_loader = loaders
  if writers is not None:
    writer, writer_eval = writers

  train_loader.batch_sampler.set_epoch(epoch)
  global global_step
  
  reg_loss = ResidualLoss(sample_rate=hps["sampling_rate"],
                          fft_size=hps["filter_length"],
                          hop_size=hps["hop_length"],
                          f0_floor=hps["f0_min"],
                          f0_ceil=hps["f0_max"],
                          n_mels=hps["n_mel_channels"]).cuda()

  net_g.train() # 学習モードへ
  net_d.train() # 学習モードへ
  if net_dur_disc is not None:
    net_dur_disc.train()
  for batch_idx, (wav,\
                  spec,            spec_lengths,\
                  f0,\
                  ph_IDs,          ph_IDs_lengths,\
                  ph_frame_dur,\
                  word_dur_ms,     \
                  word_frame_dur,  word_frame_dur_lengths,\
                  ph_idx_in_a_word,  \
                  n_ph_pooling,\
                  dfs, \
                  Sinewaves, speakerID) in enumerate(train_loader):

    # Forward Process
    with autocast(enabled=hps["fp16_run"]):
      
      # y_hat is sliced 
      y_hat, excitation, l2_dur_loss, attn_gt, ids_slice, ph_IDs_mask, spec_mask, \
      (z, z_p, m_p, logs_p, m_q, logs_q), \
      (hidden_x, logw, logw_) = net_g(spec=spec.cuda(rank),
                                              spec_lengths=spec_lengths.cuda(rank),
                                              ph_IDs=ph_IDs.cuda(rank), 
                                              ph_IDs_lengths=ph_IDs_lengths.cuda(rank),
                                              ph_dur=ph_frame_dur.cuda(rank), 
                                              word_frame_dur = word_frame_dur.cuda(rank),
                                              word_frame_dur_lenngths=word_frame_dur_lengths.cuda(rank),
                                              word_dur_ms = word_dur_ms.cuda(rank),
                                              ph_word_flag=ph_idx_in_a_word.cuda(rank),
                                              n_ph_pool=n_ph_pooling.cuda(rank),
                                              dfs=dfs,  # slice後にcudaに乗せる
                                              sinewave=Sinewaves,
                                              speakerID=speakerID.cuda(rank))
      mel = spec_to_mel_torch(
          spec.cuda(rank), 
          hps["filter_length"], 
          hps["n_mel_channels"],  
          hps["sampling_rate"], 
          hps["mel_fmin"],  
          hps["mel_fmax"])

      # slice process
      y = commons.slice_segments(wav.cuda(rank), ids_slice * hps["hop_length"], hps["segments_size"])
      y_mel = commons.slice_segments(mel, ids_slice, hps["segments_size"] // hps["hop_length"] )
      f0_slice = commons.slice_segments(f0.cuda(rank), ids_slice , hps["segments_size"]//hps["hop_length"])

      y_hat_mel = mel_spectrogram_torch(y_hat.squeeze(1),
                                        hps["filter_length"],
                                        hps["n_mel_channels"],
                                        hps["sampling_rate"],
                                        hps["hop_length"],
                                        hps["win_length"],
                                        hps["mel_fmin"],
                                        hps["mel_fmax"])
      
      ### Duration Discriminator ###
      if net_dur_disc is not None:  
        y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x.detach(), ph_IDs_mask.detach(), logw.detach(), logw_.detach())
        with autocast(enabled=False):
          # TODO: I think need to mean using the mask, but for now, just mean all
          loss_dur_disc, losses_dur_disc_r, losses_dur_disc_g = discriminator_loss(y_dur_hat_r, y_dur_hat_g)
          loss_dur_disc_all = loss_dur_disc
        optim_dur_disc.zero_grad()
        scaler.scale(loss_dur_disc_all).backward()
        scaler.unscale_(optim_dur_disc)
        grad_norm_dur_disc = commons.clip_grad_value_(net_dur_disc.parameters(), None)
        scaler.step(optim_dur_disc)

      ### Discriminator ### 
      y_d_hat_g = net_d(y_hat.detach()) 
      y_d_hat_r = net_d(y)
      with autocast(enabled=False):
        loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
        loss_disc_all = loss_disc    

    optim_d.zero_grad()
    scaler.scale(loss_disc_all).backward()
    scaler.unscale_(optim_d)
    grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
    scaler.step(optim_d) 

    ### Generator ###
    with autocast(enabled=hps["fp16_run"]):
      y_d_hat_g, fmap_g = net_d(y_hat, return_fmaps=True)  # fake feature
      with torch.no_grad():
        _, fmap_r = net_d(y, return_fmaps=True) # real feature

      if net_dur_disc is not None:
        y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x, ph_IDs_mask, logw, logw_)
      with autocast(enabled=False):
        loss_ph_dur    = torch.sum(l2_dur_loss[0]).float() * hps["gain_dur_ph"]
        loss_word_dur  = torch.sum(l2_dur_loss[1]).float() * hps["gain_dur_word"]
        loss_length    = torch.sum(l2_dur_loss[2]).float() * hps["gain_dur_length"]
        loss_length_l2 = loss_ph_dur + loss_word_dur + loss_length
        loss_mel = torch.nn.functional.l1_loss(y_mel, y_hat_mel) * hps["gain_mel"]
        loss_kl_spec = kl_loss(z_p, logs_q, m_p, logs_p, spec_mask) * hps["gain_kl_spec"]
        loss_fm = feature_loss(fmap_r, fmap_g) * hps["gain_fm"]
        loss_gen, losses_gen = generator_loss(y_d_hat_g) 
        loss_reg = reg_loss(s=excitation, y=y, f=f0_slice) * hps["gain_reg"] # sifi gan
        ############################################################
        loss_gen_all = loss_gen + loss_fm + loss_mel + loss_length_l2 + loss_kl_spec + loss_reg#+ loss_f0
        ############################################################
        if net_dur_disc is not None:
          loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
          loss_gen_all += loss_dur_gen

    losses = loss_gen_all
    optim_g.zero_grad()
    scaler.scale(losses).backward()   
    grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
    scaler.unscale_(optim_g)
    scaler.step(optim_g)
    scaler.update()

    # 記録 / 評価
    if rank==0:
      if global_step % hps["log_interval"] == 0:
        lr = optim_g.param_groups[0]['lr']
        #losses = 
        logger.info('Train Epoch: {} [{:.0f}%]'.format(
          epoch,
          100. * batch_idx / len(train_loader)))
        logger.info("steps:{}   [Generator]   ALL:{:.5f} | mel:{:.5f} | fm:{:.5f} | reg:{:.5f}| kl:{:.5f} | dur:{:.5f}"
                    .format(global_step, loss_gen_all.item(), loss_mel.item(), loss_fm.item(),loss_reg.item(), loss_kl_spec.item(),loss_length_l2.item()
                    ))
        logger.info("steps:{} [Discriminator] d_ALL:{:.5f} | d_dur_ALL:{:.5f} | "
                    .format(global_step, loss_disc_all.item(), loss_dur_disc_all.item()
                    ))
        
        scalar_dict = {"loss/g/ALL": loss_gen_all, 
                       "loss/d/ALL": loss_disc_all,
                       "loss/g/AdvDiscDur" :loss_dur_gen,
                       "loss/g/AdvDisc" :loss_gen,
                       "loss/g/Regularization" : loss_reg,
                       "loss/g/KL_div_spec" : loss_kl_spec,
                       "loss/g/FeatureMatching": loss_fm, 
                       "loss/g/MelReconstruction": loss_mel, 
                       "loss/g/Duration_SUM": loss_length_l2, 
                       "loss/g/Duration_ph": loss_ph_dur, 
                       "loss/g/Duration_word": loss_word_dur, 
                       "loss/g/Duration_length": loss_length, 
                       "learning_rate": lr, 
                       "grad_norm/discriminator": grad_norm_d, 
                       "grad_norm/generator": grad_norm_g}
        
        scalar_dict.update({"loss/g/Adversarial_Deceive{}".format(i): v for i, v in enumerate(losses_gen)})
        scalar_dict.update({"loss/d/Adversarial_Real_{}".format(i): v for i, v in enumerate(losses_disc_r)})
        scalar_dict.update({"loss/d/Adversarial_Fake_{}".format(i): v for i, v in enumerate(losses_disc_g)})
        
        if net_dur_disc is not None:
          scalar_dict.update({"loss/d_dur/ALL"    : loss_dur_disc_all,})
          scalar_dict.update({"grad_norm/duration_discriminator": grad_norm_dur_disc})
          scalar_dict.update({"loss/d_dur/Adversarial_Real_{}".format(i): v for i, v in enumerate(losses_dur_disc_r)})
          scalar_dict.update({"loss/d_dur/Adversarial_Fake_{}".format(i): v for i, v in enumerate(losses_dur_disc_g)})

        # attnマップの作成
        length_scale=1
        w = torch.exp(logw) * ph_IDs_mask * length_scale
        w_ceil = torch.ceil(w)
        spec_mask = torch.unsqueeze(commons.sequence_mask(spec_lengths, spec.size(2)), 1).to(spec.dtype) # [B, 1, ph_len]
        ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
        attn_mask     = torch.unsqueeze(spec_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
        attn_pd       = commons.generate_path(duration=torch.unsqueeze(w_ceil,dim=1), mask=attn_mask.permute(0,1,3,2).cuda(rank))

        image_dict = { 
            "slice/mel_orgiginal": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
            "slice/mel_generated": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()), 
            "all/mel_full_length": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
            "all/attn_pd": utils.plot_alignment_to_numpy(attn_pd[0][0].data.cpu().numpy()),
            "all/attn_gt": utils.plot_alignment_to_numpy(attn_gt.permute(0,2,1)[0].data.cpu().numpy()),
            "all/attn_overlap": utils.plot_alignment_to_numpy(0.5*(attn_pd[0][0].data.cpu().numpy()) + 0.5*attn_gt.permute(0,2,1)[0].data.cpu().numpy())
        }

        utils.summarize(
          writer=writer,
          global_step=global_step, 
          images=image_dict,
          scalars=scalar_dict)

      if global_step % hps["eval_interval"] == 0:
        evaluate(hps=hps,
                 epoch=epoch,
                 nets=[net_g,net_d,net_dur_disc],
                 eval_loader=eval_loader, 
                 writer_eval=writer_eval)

        utils.save_checkpoint(model=net_g, 
                              optimizer=optim_g, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "G_{}.pth".format(global_step)))
        utils.save_checkpoint(model=net_d, 
                              optimizer=optim_d, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "D_{}.pth".format(global_step)))
        if net_dur_disc is not None:
          utils.save_checkpoint(model=net_dur_disc, 
                              optimizer=optim_dur_disc, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "DUR_{}.pth".format(global_step)))
    global_step += 1
    #torch.cuda.empty_cache()

  if rank == 0:
    logger.info('====> Epoch: {}'.format(epoch))
  
def evaluate(hps, epoch, nets, eval_loader, writer_eval):
    net_g,net_d,net_dur_disc = nets

    net_g.eval()
    net_d.eval()
    if net_dur_disc is not None:
      net_d.eval()

    reg_loss = ResidualLoss(sample_rate=hps["sampling_rate"],
                          fft_size=hps["filter_length"],
                          hop_size=hps["hop_length"],
                          f0_floor=hps["f0_min"],
                          f0_ceil=hps["f0_max"],
                          n_mels=hps["n_mel_channels"]).cuda()
    
    rank = 0
    
    count = 0
    loss_gen_all = 0
    loss_gen_adv_all = 0
    loss_disc_all= 0
    loss_dur_gen_all= 0
    #loss_f0_all= 0
    loss_reg_all = 0
    loss_kl_spec_all= 0
    loss_fm_all= 0
    loss_mel_all= 0
    loss_length_l2_all= 0
    loss_dur_ph_all = 0
    loss_dur_word_all = 0
    loss_dur_length_all = 0
    loss_dur_disc_all = 0
    with torch.no_grad():
      for batch_idx, (wav,\
                      spec,            spec_lengths,\
                      f0,\
                      ph_IDs,          ph_IDs_lengths,\
                      ph_frame_dur,\
                      word_dur_ms,     \
                      word_frame_dur,  word_frame_dur_lengths,\
                      ph_idx_in_a_word,  \
                      n_ph_pooling,\
                      dfs, \
                      Sinewaves, speakerID) in enumerate(eval_loader):
        
        count += 1

        # Forward Process
        with autocast(enabled=hps["fp16_run"]):

          # y_hat is sliced 
          y_hat, excitation, l2_dur_loss, attn_gt, ids_slice, ph_IDs_mask, spec_mask, \
          (z, z_p, m_p, logs_p, m_q, logs_q), \
          (hidden_x, logw, logw_) = net_g(spec=spec.cuda(rank),
                                              spec_lengths=spec_lengths.cuda(rank),
                                              ph_IDs=ph_IDs.cuda(rank), 
                                              ph_IDs_lengths=ph_IDs_lengths.cuda(rank),
                                              ph_dur=ph_frame_dur.cuda(rank), 
                                              word_frame_dur = word_frame_dur.cuda(rank),
                                              word_frame_dur_lenngths=word_frame_dur_lengths.cuda(rank),
                                              word_dur_ms = word_dur_ms.cuda(rank),
                                              ph_word_flag=ph_idx_in_a_word.cuda(rank),
                                              n_ph_pool=n_ph_pooling.cuda(rank),
                                              dfs=dfs,  # slice後にcudaに乗せる
                                              sinewave=Sinewaves,
                                              speakerID=speakerID.cuda(rank))

          mel = spec_to_mel_torch(
              spec.cuda(rank), 
              hps["filter_length"], 
              hps["n_mel_channels"],  
              hps["sampling_rate"], 
              hps["mel_fmin"],  
              hps["mel_fmax"])

          # sliced 
          y_mel = commons.slice_segments(mel, ids_slice, hps["segments_size"] // hps["hop_length"] )
          f0_slice = commons.slice_segments(f0.cuda(rank), ids_slice , hps["segments_size"]//hps["hop_length"])

          y_hat_mel = mel_spectrogram_torch(y_hat.squeeze(1),
                                            hps["filter_length"],
                                            hps["n_mel_channels"],
                                            hps["sampling_rate"],
                                            hps["hop_length"],
                                            hps["win_length"],
                                            hps["mel_fmin"],
                                            hps["mel_fmax"])

          y = commons.slice_segments(wav.cuda(rank), ids_slice * hps["hop_length"], hps["segments_size"])

          ### Duration Discriminator ###
          if net_dur_disc is not None:  
            y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x.detach(), ph_IDs_mask.detach(), logw.detach(), logw_.detach())
            with autocast(enabled=False):
              # TODO: I think need to mean using the mask, but for now, just mean all
              loss_dur_disc, losses_dur_disc_r, losses_dur_disc_g = discriminator_loss(y_dur_hat_r, y_dur_hat_g)
              loss_dur_disc_all += loss_dur_disc
              
          ### Discriminator ### 
          y_d_hat_g = net_d(y_hat.detach()) 
          y_d_hat_r = net_d(y)
          with autocast(enabled=False):
            loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
            loss_disc_all += loss_disc    

        ### Generator ###
        with autocast(enabled=hps["fp16_run"]):

          #y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)

          y_d_hat_g, fmap_g = net_d(y_hat, return_fmaps=True)  # fake feature
          with torch.no_grad():
            _, fmap_r = net_d(y, return_fmaps=True) # real feature

          if net_dur_disc is not None:
            y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x, ph_IDs_mask, logw, logw_)
          with autocast(enabled=False):
          
            loss_ph_dur    = torch.sum(l2_dur_loss[0]).float()  * hps["gain_dur_ph"]
            loss_word_dur  = torch.sum(l2_dur_loss[1]).float()  * hps["gain_dur_word"]
            loss_length    = torch.sum(l2_dur_loss[2]).float()  * hps["gain_dur_length"]
            loss_length_l2 = loss_ph_dur + loss_word_dur + loss_length
            loss_dur_ph_all += loss_ph_dur
            loss_dur_word_all += loss_word_dur
            loss_dur_length_all += loss_length
            loss_length_l2_all += loss_length_l2

            loss_mel = torch.nn.functional.l1_loss(y_mel, y_hat_mel) * hps["gain_mel"]
            loss_mel_all += loss_mel

            loss_kl_spec = kl_loss(z_p, logs_q, m_p, logs_p, spec_mask) * hps["gain_kl_spec"]
            loss_kl_spec_all += loss_kl_spec
            
            loss_fm = feature_loss(fmap_r, fmap_g) * hps["gain_fm"]
            loss_fm_all += loss_fm

            loss_gen, losses_gen = generator_loss(y_d_hat_g)
            loss_gen_adv_all += loss_gen 

            loss_reg = reg_loss(s=excitation, y=y, f = f0_slice) * hps["gain_reg"]
            loss_reg_all += loss_reg

            if net_dur_disc is not None:
              loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
              loss_dur_gen_all += loss_dur_gen

      loss_gen_all = loss_gen + loss_fm + loss_mel + loss_length_l2 + loss_kl_spec + loss_reg#+ loss_f0
            
      scalar_dict = {"loss/g/ALL"               : loss_gen_all / count, 
                     "loss/d/ALL"               : loss_disc_all/ count,
                     "loss/g/AdvDiscDur"        : loss_dur_gen_all/ count,
                     "loss/g/AdvDisc"           : loss_gen_adv_all/ count,
                     "loss/g/Regularization"    : loss_reg_all/ count,
                     "loss/g/KL_div_spec"       : loss_kl_spec_all/ count,
                     "loss/g/FeatureMatching"   : loss_fm_all/ count, 
                     "loss/g/MelReconstruction" : loss_mel_all/ count, 
                     "loss/g/Duration_SUM": loss_length_l2_all/ count, 
                     "loss/g/Duration_ph": loss_dur_ph_all/ count, 
                     "loss/g/Duration_word": loss_dur_word_all/ count, 
                     "loss/g/Duration_length": loss_dur_length_all/ count
                     }
      if net_dur_disc is not None:
        scalar_dict.update({"loss/d_dur/ALL"    : loss_dur_disc_all})

    wav_gt                 = wav[0].view(1, 1,-1)
    ph_IDs                 = ph_IDs[0].view(1,-1).cuda(rank)                
    ph_IDs_len             = ph_IDs_lengths[0].view(-1).cuda(rank)       
    ph_IDs_dur             = ph_frame_dur[0].view(1,-1).cuda(rank)  
    f0_gt                  = f0  [0][0].view(1,1,-1).cuda(rank)
    f0_gt_len              = torch.tensor(torch.sum(ph_IDs_dur), dtype=torch.int64).view(1).cuda(rank)
    word_frame_dur         = word_frame_dur[0].view(1,-1).cuda(rank)  
    speakerID              = speakerID[0].view(-1).cuda(rank)      
    word_frame_dur_lengths = word_frame_dur_lengths[0].view(-1)
    word_dur_ms            = word_dur_ms[0].view(1,-1)
    ph_idx_in_a_word       = ph_idx_in_a_word[0].view(1,-1)
    n_ph_pooling           = n_ph_pooling[0].view(1,-1)
    Sinewaves              = Sinewaves[0].view(1,1,-1)
    for idx, df in enumerate(dfs):
      dfs[idx] = df[0].view(1,1,-1)

    with torch.no_grad():
      # with f0_gt
      y_hat_with_f0gt, attn, y_mask = net_g.eval_infer(ph_IDs=ph_IDs.cuda(rank), 
                                                  ph_IDs_lengths=ph_IDs_len.cuda(rank),
                                                  word_frame_dur = word_frame_dur.cuda(rank),
                                                  word_frame_dur_lenngths=word_frame_dur_lengths.cuda(rank),
                                                  word_dur_ms = word_dur_ms.cuda(rank),
                                                  ph_word_flag=ph_idx_in_a_word.cuda(rank),
                                                  n_ph_pool=n_ph_pooling.cuda(rank),
                                                  speakerID=speakerID.cuda(rank),
                                                  dfs=dfs, 
                                                  sinewave=Sinewaves
                                                  ) 

    try:
      y_hat_with_f0gt = y_hat_with_f0gt[0].to('cpu').detach().numpy().copy()
    except: 
      y_hat_with_f0gt = None

    f0_mask = torch.unsqueeze(commons.sequence_mask(f0_gt_len, f0_gt.size(2)), 1).to(f0_gt.dtype) # [B, 1, ph_len]
    ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_len, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
    attn_mask     = torch.unsqueeze(f0_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
    attn_gt       = commons.generate_path(duration=torch.unsqueeze(ph_IDs_dur,dim=1), mask=attn_mask.permute(0,1,3,2))

    image_dict = { 
    "all/attn": utils.plot_alignment_to_numpy(attn.permute(0,2,1)[0].data.cpu().numpy()),
    "all/attn_GT": utils.plot_alignment_to_numpy(attn_gt[0][0].data.cpu().numpy()),
    "all/attn_overlap": utils.plot_alignment_to_numpy(0.5*(attn_gt[0][0].data.cpu().numpy()) + 0.5*attn.permute(0,2,1)[0].data.cpu().numpy())
    }
    
    audio_dict = {f"GroundTruth/audio_all" : wav_gt}
    if y_hat_with_f0gt is not None:
      audio_dict.update({f"gen/audio_with_f0gt_{batch_idx}"        : y_hat_with_f0gt})

    utils.summarize(
      writer=writer_eval,
      global_step=global_step, 
      scalars=scalar_dict,
      images=image_dict,
      audios = audio_dict,
      audio_sampling_rate=hps["sampling_rate"]
    )

    net_g.train()
    net_d.train()
    if net_dur_disc is not None:
      net_dur_disc.train()

def generate_graph(vector, 
                   label="NoName",
                   color="blue", 
                   title="Title",
                   x_label = 'Frames',
                   y_label = "y_labels",
                   figsize=(36,8)):
   
    fig = plt.figure(figsize=figsize) 
    x = np.arange(0, len(vector))
    plt.plot(x, vector, label=label, color=color)
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    #plt.ylim(y_val_min, y_val_max)
    fig.canvas.draw() 
    plot_image = fig.canvas.renderer._renderer  
    image_numpy = np.array(plot_image)
    plt.clf()
    plt.close()
    return image_numpy

def moving_average(loss,n=5):

    # Build a tensor list
    loss_list = []

    if len(loss_list) >= n:
        # drop the first loss
        loss_list =  loss_list[1:,:]
  
    loss_list.append(loss.data)

    baseline = loss-sum(loss_list)/len(loss_list)
    
    if baseline.all() == 0:
        return loss.detach()
    else:
        return baseline


# このファイルを実行したら、ココから始まる                  
if __name__ == "__main__":
  os.environ[
        "TORCH_DISTRIBUTED_DEBUG"
    ] = "DETAIL"
  main()

train_SVS_synthesizer.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler

from models import DiffusionModels

######### export TORCH_CUDNN_V8_API_DISABLED=1
torch.cuda.empty_cache()
import torch._dynamo.config
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.verbose=True

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

import utils
from data_utils import TextAudioLoader_TTS_f0diff, TextAudioCollate_TTS_f0diff, DistributedBucketSampler

# torch.autograd.set_detect_anomaly(True) # 逆誤差伝搬のエラー部分を自動で見つける機能
torch.backends.cudnn.benchmark = True # 学習早くなるらしい（知らんけど
global_step = 0

import warnings
warnings.filterwarnings(action='ignore')

def main():

  """ GPUのみの学習を許可 """
  assert torch.cuda.is_available(), "CPU training is not allowed."

  n_gpus = torch.cuda.device_count()
  os.environ['MASTER_ADDR'] = 'localhost'
  os.environ['MASTER_PORT'] = '65520'

  hps = utils.get_hparams()   # config.yamlやargparseは、ここで読み込まれている。

  if hps["DDP_use"] is True:
    mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))  # DDP ON
  else:
    run(rank=0, n_gpus=1, hps=hps)                     # DDP OFF


def run(rank, n_gpus, hps):
  global global_step
  if rank == 0: # 1番目のGPUであれば入る（GPUが複数ある場合を想定
    logger = utils.get_logger(hps["model_dir"])
    logger.info(hps)
    writer = SummaryWriter(log_dir=hps["model_dir"])
    writer_eval = SummaryWriter(log_dir=os.path.join(hps["model_dir"], "eval"))

  # Linux Only
  if hps["DDP_use"] is True: 
    dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)   

  # 乱数の再現性確保
  torch.manual_seed(hps["seed"])
  torch.cuda.set_device(rank)

  # データセットの読み込み君 
  train_dataset = TextAudioLoader_TTS_f0diff(audiopaths_and_text= hps["train_data_path"] ,
                                  hparams            = hps)
  vocab_size = train_dataset.get_ph_vocab_size()
  
  # 似た長さのデータをモデルに流す君
  train_sampler = DistributedBucketSampler(dataset     = train_dataset,
                                           batch_size  = hps["batch_size"],
                                           boundaries  = [32,300,400,500,600,700,800,900,1000],
                                           num_replicas= n_gpus,
                                           rank        = rank,
                                           shuffle     = True) # シャッフル指定

  # 入力するデータに0を加えて長さを統一する君
  collate_fn = TextAudioCollate_TTS_f0diff()

  # 上三つを統べる
  train_loader = DataLoader(dataset=train_dataset, 
                            num_workers=hps["num_workers"],     # データ読み込みに使用するCPUのスレッド数
                            shuffle=False, 
                            pin_memory=True,    # Trueだとなんか動作早くなるらしい（知らんけど
                            collate_fn=collate_fn, 
                            batch_sampler=train_sampler)

  # 1番目のGPUであれば
  if rank == 0:
    eval_dataset = TextAudioLoader_TTS_f0diff(audiopaths_and_text=hps["eval_data_path"], 
                                              hparams            =hps)
    
    eval_loader = DataLoader(dataset      = eval_dataset, 
                             num_workers  = 1, 
                             shuffle      = False,
                             batch_size   = hps["batch_size"], 
                             pin_memory   = True,
                             drop_last    = False, 
                             collate_fn   = collate_fn)

  # モデルを作成
  net_g = DiffusionModels(hps=hps, vocab_size=vocab_size).cuda(rank)
  #net_g = torch.compile(net_g)

  # 最適化手法
  optim_g = torch.optim.AdamW(params      = net_g.parameters(),
                              lr          = float(hps['learning_rate']), 
                              betas       =       hps["betas"], 
                              eps         = float(hps["eps"]))

  # checkpoint 確認
  in_progress_checkpoint_path = utils.latest_checkpoint_path(hps["model_dir"], "G_*.pth")
  if os.path.isfile(in_progress_checkpoint_path): # 学習再開
    try:
      _, _, _, epoch_str = utils.load_checkpoint(in_progress_checkpoint_path, net_g, optim_g)
      global_step = (epoch_str - 1) * len(train_loader)
      print("### Resume learning mode ###")
    except:
      net_g = utils.load_model_diffsize(in_progress_checkpoint_path, net_g, hps, optim_g)
      epoch_str = 1
      global_step = 0.
      print("### Resume learning mode ###")
  elif hps["is_finetune"] is True: # ファインチューニング
    net_g = utils.load_model_diffsize(hps["finetune_G"], net_g, hps, optim_g)
    epoch_str = 1
    global_step = 0
    print("### FineTuning mode ###")
  else: # どちらでもない→初期状態から
    epoch_str = 1
    global_step = 0
    print("### Initial state mode ###")
  global_step = (epoch_str - 1) * len(train_loader)

  # エポックに応じて自動で学習率を変える
  scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_g, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)

  # float16で計算して，高速化するか
  scaler = GradScaler(enabled=hps["fp16_run"])

  # Linux Only
  if hps["DDP_use"] is True: 
    net_g = DDP(net_g, device_ids=[rank]) 

  # 学習/評価処理
  for epoch in range(epoch_str, int(hps["max_epochs"]) + 1):
    if rank==0:
      train_and_evaluate(rank, epoch, hps, net_g, optim_g, scheduler_g, scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
    else:
      train_and_evaluate(rank, epoch, hps, net_g, optim_g, scheduler_g, scaler, [train_loader, None], None, None)
    scheduler_g.step()

  # 最終モデルの保存
  utils.save_checkpoint(net_g, optim_g, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "G_LastModel.pth".format(global_step)))
    
def train_and_evaluate(rank, epoch, hps, net_g, optim_g, scheduler_g, scaler, loaders, logger, writers):
  
  train_loader, eval_loader = loaders
  if writers is not None:
    writer, writer_eval = writers

  train_loader.batch_sampler.set_epoch(epoch)
  global global_step

  net_g.train() # 学習モードへ
  for batch_idx, (f0,f0_lengths,\
                  ph_IDs,ph_IDs_lengths,\
                  ph_frame_dur,\
                  speakerID) in enumerate(train_loader):

    # Forward Process
    with autocast(enabled=hps["fp16_run"]):
      loss_f0 = net_g(f0=f0.cuda(rank),  
                      f0_len=f0_lengths.cuda(rank),
                      IDs=ph_IDs.cuda(rank), 
                      IDs_len=ph_IDs_lengths.cuda(rank),
                      IDs_dur=ph_frame_dur.cuda(rank),
                      g=speakerID.cuda(rank))
    #losses = loss_f0 + loss_vuv
    losses = loss_f0
    optim_g.zero_grad()
    scaler.scale(losses).backward()   
    scaler.unscale_(optim_g)
    scaler.step(optim_g)
    scaler.update()

    # 記録 / 評価
    if rank==0:
      if global_step % hps["log_interval"] == 0:
        lr = optim_g.param_groups[0]['lr']
        #losses = 
        logger.info('Train Epoch: {} [{:.0f}%]'.format(
          epoch,
          100. * batch_idx / len(train_loader)))
        #logger.info([x.item() for x in [losses, loss_f0, loss_vuv]] + [global_step, lr])
        logger.info([x.item() for x in [losses]] + [global_step, lr])
        #scalar_dict = {"loss": losses, "loss_f0" : loss_f0, "loss_vuv":loss_vuv}
        scalar_dict = {"loss_f0": losses}
        
        utils.summarize(
          writer=writer,
          global_step=global_step, 
          scalars=scalar_dict)
        
      if global_step % hps["eval_interval"] == 0:
        evaluate(hps=hps,
                 net_g=net_g,
                 eval_loader=eval_loader, 
                 writer_eval=writer_eval)
        utils.save_checkpoint(model=net_g, 
                              optimizer=optim_g, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "G_{}.pth".format(global_step)))
    global_step += 1

  if rank == 0:
    logger.info('====> Epoch: {}'.format(epoch))
  
def evaluate(hps, net_g, eval_loader, writer_eval):
    net_g.eval()
    rank = 0
    count = 0
    loss_f0_sum = 0
    with torch.no_grad():
      for batch_idx, (f0,f0_lengths,\
                      ph_IDs,ph_IDs_lengths,\
                      ph_frame_dur,\
                      speakerID) in enumerate(eval_loader):
        # Forward Process
        with autocast(enabled=hps["fp16_run"]):
          loss_f0 = net_g(f0=f0.cuda(rank),  
                          f0_len=f0_lengths.cuda(rank),
                          IDs=ph_IDs.cuda(rank), 
                          IDs_len=ph_IDs_lengths.cuda(rank),
                          IDs_dur=ph_frame_dur.cuda(rank),
                          g=speakerID.cuda(rank))
          loss_f0_sum += loss_f0
          count += 1

    f0_gt=f0                  [0][0].view(1, 1,-1).cuda(rank)
    f0_gt_len=f0_lengths      [0]   .view(-1)     .cuda(rank)
    ph_IDs=ph_IDs             [0]   .view(1,-1)   .cuda(rank)   
    ph_IDs_len=ph_IDs_lengths [0]   .view(-1)     .cuda(rank)
    ph_IDs_dur=ph_frame_dur   [0]   .view(1,-1)   .cuda(rank)
    speakerID=speakerID       [0]   .view(-1)     .cuda(rank)

    #f0_pd, vuv_pd = net_g.sampling(condition=[IDs_gt, IDs_gt_len, dur_gt, speaker_ID])
    f0_pd = net_g.sampling(condition=[ph_IDs, ph_IDs_len, ph_IDs_dur, 
                                      None, None, None, speakerID])

    scalar_dict = {"loss_f0": loss_f0_sum/count}
    #scalar_dict = {"loss": loss_f0_sum + loss_vuv_sum,
    #               "loss_f0": loss_f0_sum,
    #               "loss_vuv" : loss_vuv_sum}

    f0_gt = f0_gt[0][0].to('cpu').detach().numpy().copy()
    f0_pd = f0_pd[0][0].to('cpu').detach().numpy().copy() 
    #vuv_gt= vuv_gt.view(-1).to('cpu').detach().numpy().copy() 
    #vuv_pd= vuv_pd.view(-1).to('cpu').detach().numpy().copy() 

    f0_gt_image= generate_graph(vector=f0_gt, 
                                label="F0 gt",
                                color="red", 
                                x_label = 'Frames',
                                y_label = "f0")

    f0_pd_image= generate_graph(vector=f0_pd, 
                                label="F0 pd",
                                color="blue", 
                                x_label = 'Frames',
                                y_label = "f0")
    #vuv_gt_image=generate_graph(vector=vuv_gt, 
    #                            label="Voice Unvoice Value",
    #                            color="red", 
    #                            x_label = 'Frames',
    #                            y_label = "0=mask, 1=uv, 2=v")
    #vuv_pd_image= generate_graph(vector=vuv_pd, 
    #                            label="Voice Unvoice Value",
    #                            color="blue", 
    #                            x_label = 'Frames',
    #                            y_label = "0=mask, 1=uv, 2=v",)
    
    image_dict = { 
    "f0/gt":f0_gt_image ,
    "f0/pd":f0_pd_image ,
    #"vuv/gt":vuv_gt_image ,
    #"vuv/pd":vuv_pd_image ,
        }
    
    utils.summarize(
      writer=writer_eval,
      global_step=global_step, 
      scalars=scalar_dict,
      images=image_dict
    )
    net_g.train()

def generate_graph(vector, 
                   label="NoName",
                   color="blue", 
                   title="Title",
                   x_label = 'Frames',
                   y_label = "y_labels",
                   figsize=(20,5)):
   
    fig = plt.figure(figsize=figsize) 
    x = np.arange(0, len(vector))
    plt.plot(x, vector, label=label, color=color)
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    #plt.ylim(y_val_min, y_val_max)
    fig.canvas.draw() 
    plot_image = fig.canvas.renderer._renderer  
    image_numpy = np.array(plot_image)
    plt.clf()
    plt.close()
    return image_numpy

# このファイルを実行したら、ココから始まる                  
if __name__ == "__main__":
  os.environ[
        "TORCH_DISTRIBUTED_DEBUG"
    ] = "DETAIL"
  main()

train_TTS_F0Diffusion.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler

from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
import commons

######### export TORCH_CUDNN_V8_API_DISABLED=1
torch.cuda.empty_cache()
import torch._dynamo.config
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.verbose=True

from models import VITS2_based_SiFiTTS
from sifigan.models.discriminator import UnivNetMultiResolutionMultiPeriodDiscriminator
from models import DurationDiscriminator

from sifigan.losses.reg import ResidualLoss

from losses import generator_loss, discriminator_loss, feature_loss,kl_loss
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import utils
from data_utils import TextAudioLoader_TTS_synth, TextAudioCollate_TTS_synth, DistributedBucketSampler

# torch.autograd.set_detect_anomaly(True) # 逆誤差伝搬のエラー部分を自動で見つける機能
torch.backends.cudnn.benchmark = True # 学習早くなるらしい（知らんけど
global_step = 0

import warnings
warnings.filterwarnings(action='ignore')

def main():

  """ GPUのみの学習を許可 """
  assert torch.cuda.is_available(), "CPU training is not allowed."

  n_gpus = torch.cuda.device_count()
  os.environ['MASTER_ADDR'] = 'localhost'
  os.environ['MASTER_PORT'] = '65520'

  hps = utils.get_hparams()   # config.yamlやargparseは、ここで読み込まれている。

  if hps["DDP_use"] is True:
    mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))  # DDP ON(未確認)
  else:
    run(rank=0, n_gpus=1, hps=hps)                     # DDP OFF

def run(rank, n_gpus, hps):
  global global_step
  if rank == 0: # 1番目のGPUであれば入る（GPUが複数ある場合を想定
    logger = utils.get_logger(hps["model_dir"])
    logger.info(hps)
    # utils.check_git_hash(hps["model_dir"])
    writer = SummaryWriter(log_dir=hps["model_dir"])
    writer_eval = SummaryWriter(log_dir=os.path.join(hps["model_dir"], "eval"))

  # Linux Only
  if hps["DDP_use"] is True: 
    dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)   

  # 乱数の再現性確保
  torch.manual_seed(hps["seed"])
  torch.cuda.set_device(rank)

  # データセットの読み込み君 
  train_dataset = TextAudioLoader_TTS_synth(audiopaths_and_text= hps["train_data_path"] ,
                                  hparams            = hps)
  vocab_size = train_dataset.get_ph_vocab_size()
  
  # 似た長さのデータをモデルに流す君
  train_sampler = DistributedBucketSampler(dataset     = train_dataset,
                                           batch_size  = hps["batch_size"],
                                           boundaries  = [32,300,400,500,600,700,800,900,1000],
                                           num_replicas= n_gpus,
                                           rank        = rank,
                                           shuffle     = True) # シャッフル指定

  # 入力するデータに0を加えて長さを統一する君
  collate_fn = TextAudioCollate_TTS_synth(hparams=hps)

  # 上三つを統べる
  train_loader = DataLoader(dataset=train_dataset, 
                            num_workers=hps["num_workers"],     # データ読み込みに使用するCPUのスレッド数
                            shuffle=False, 
                            pin_memory=True,    # Trueだとなんか動作早くなるらしい（知らんけど
                            collate_fn=collate_fn, 
                            batch_sampler=train_sampler)

  # 1番目のGPUであれば
  if rank == 0:
    eval_dataset = TextAudioLoader_TTS_synth(audiopaths_and_text=hps["eval_data_path"], 
                                   hparams            =hps)

    eval_loader = DataLoader(dataset      = eval_dataset, 
                             num_workers  = 1, 
                             shuffle      = False,
                             batch_size   = hps["batch_size"], 
                             pin_memory   = True,
                             drop_last    = False, 
                             collate_fn   = collate_fn)

  net_dur_disc = DurationDiscriminator(
    hps["dur_discriminator"]["hidden_channels"], 
    hps["dur_discriminator"]["hidden_channels"], 
    3, 
    0.1, 
    gin_channels=hps["common"]["gin_channels"] if hps["common"]["n_speaker"] != 0 else 0,
    ).cuda(rank)

  net_g = VITS2_based_SiFiTTS(
      hps = hps,
      n_vocab=vocab_size,
      spec_channels=hps["spec_encoder"]["spec_channels"],
      segment_size=int(hps["segments_size"] // hps["hop_length"]),
      n_speakers=hps["common"]["n_speaker"],
      gin_channels=hps["common"]["gin_channels"],
      **hps["hifi_gan"]).cuda(rank)

  net_d = UnivNetMultiResolutionMultiPeriodDiscriminator(**hps["UnivNet_MRMPDiscriminator"]).cuda(rank)
  
  # torch.compile 学習高速化
  #net_g = torch.compile(net_g)
  #net_d = torch.compile(net_d)
  #if net_dur_disc is not None:
  #  net_dur_disc = torch.compile(net_dur_disc)
  
  # 最適化手法
  optim_g = torch.optim.AdamW(params      = net_g.parameters(),
                              lr          = float(hps['learning_rate']), 
                              betas       =       hps["betas"], 
                              eps         = float(hps["eps"]))
  optim_d = torch.optim.AdamW(params      = net_d.parameters(),
                              lr          = float(hps['learning_rate']), 
                              betas       =       hps["betas"], 
                              eps         = float(hps["eps"]))
  if net_dur_disc is not None:
    optim_dur_disc = torch.optim.AdamW(params      = net_dur_disc.parameters(),
                                       lr          = float(hps['learning_rate']), 
                                       betas       =       hps["betas"], 
                                       eps         = float(hps["eps"]))
  else:
    optim_dur_disc = None
 
  # checkpoint 確認
  in_progress_checkpoint_path = utils.latest_checkpoint_path(hps["model_dir"], "G_*.pth")
  if os.path.isfile(in_progress_checkpoint_path): # 学習再開
    try:
      _, _, _, epoch_str = utils.load_checkpoint(in_progress_checkpoint_path, net_g, optim_g)
      _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps["model_dir"], "D_*.pth"), net_d, optim_d)
      if net_dur_disc is not None:
        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps["model_dir"], "DUR_*.pth"), net_dur_disc, optim_dur_disc)
      global_step = (epoch_str - 1) * len(train_loader)
      print("### Resume learning mode ###")
    except:
      net_g = utils.load_model_diffsize(in_progress_checkpoint_path, net_g, hps, optim_g)
      net_d = utils.load_model_diffsize(utils.latest_checkpoint_path(hps["model_dir"], "D_*.pth"), net_d, hps, optim_d)
      if net_dur_disc is not None:
        net_dur_disc = utils.load_model_diffsize(utils.latest_checkpoint_path(hps["model_dir"], "DUR_*.pth"), net_dur_disc, hps, optim_dur_disc)
      epoch_str = 1
      global_step = 0.
      print("### Resume learning mode ###")
  elif hps["is_finetune"] is True: # ファインチューニング
    net_g = utils.load_model_diffsize(hps["finetune_G"], net_g, hps, optim_g)
    net_d = utils.load_model_diffsize(hps["finetune_D"], net_d, hps, optim_d)
    if net_dur_disc is not None:
      net_dur_disc = utils.load_model_diffsize(hps["finetune_DUR"], net_dur_disc, hps, optim_dur_disc)
    epoch_str = 1
    global_step = 0
    print("### FineTuning mode ###")
  else: # どちらでもない→初期状態から
    epoch_str = 1
    global_step = 0
    print("### Initial state mode ###")
  global_step = (epoch_str - 1) * len(train_loader)

  # エポックに応じて自動で学習率を変える
  scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_g, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)
  scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_d, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)
  if net_dur_disc is not None:
    scheduler_dur_disc = torch.optim.lr_scheduler.ExponentialLR(optimizer  = optim_dur_disc, 
                                                       gamma      = hps["lr_decay"], 
                                                       last_epoch = epoch_str-2)
  else:
    scheduler_dur_disc = None

  # float16で計算して，高速化するか
  scaler = GradScaler(enabled=hps["fp16_run"])

  # Linux Only
  if hps["DDP_use"] is True: 
    net_g = DDP(net_g, device_ids=[rank]) 
    net_d = DDP(net_d, device_ids=[rank]) 
    if net_dur_disc is not None:
      net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True)

  # 学習/評価処理
  for epoch in range(epoch_str, int(hps["max_epochs"]) + 1):
    if rank==0:
      train_and_evaluate(rank, epoch, hps, [net_g, net_d, net_dur_disc], [optim_g, optim_d, optim_dur_disc], [scheduler_g, scheduler_d, scheduler_dur_disc], scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
    else:
      train_and_evaluate(rank, epoch, hps, [net_g, net_d, net_dur_disc], [optim_g, optim_d, optim_dur_disc], [scheduler_g, scheduler_d, scheduler_dur_disc], scaler, [train_loader, None], None, None)
    scheduler_g.step()
    scheduler_d.step()
    if net_dur_disc is not None:
      scheduler_dur_disc.step()

  # 最終モデルの保存
  utils.save_checkpoint(net_g, optim_g, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "G_LastModel.pth".format(global_step)))
  utils.save_checkpoint(net_d, optim_d, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "D_LastModel.pth".format(global_step)))
  if net_dur_disc is not None:
    utils.save_checkpoint(net_dur_disc, optim_dur_disc, hps["learning_rate"], epoch, os.path.join(hps["model_dir"], "DUR_LastModel.pth".format(global_step)))

def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers):
  net_g, net_d, net_dur_disc = nets
  optim_g, optim_d, optim_dur_disc = optims
  scheduler_g, scheduler_d, scheduler_dur_disc = schedulers

  train_loader, eval_loader = loaders
  if writers is not None:
    writer, writer_eval = writers

  train_loader.batch_sampler.set_epoch(epoch)
  global global_step

  reg_loss = ResidualLoss(sample_rate=hps["sampling_rate"],
                          fft_size=hps["filter_length"],
                          hop_size=hps["hop_length"],
                          f0_floor=hps["f0_min"],
                          f0_ceil=hps["f0_max"],
                          n_mels=hps["n_mel_channels"]).cuda()

  net_g.train() # 学習モードへ
  net_d.train() # 学習モードへ
  if net_dur_disc is not None:
    net_dur_disc.train()
  for batch_idx, (wav,\
                  spec,            spec_lengths,\
                  f0,\
                  ph_IDs,          ph_IDs_lengths,\
                  dfs, \
                  Sinewaves, speakerID) in enumerate(train_loader):

    # Forward Process
    with autocast(enabled=hps["fp16_run"]):
      # y_hat is sliced 
      y_hat, excitation, l2_dur_loss, attn_gt, ids_slice, ph_IDs_mask, spec_mask, \
      (z, z_p, m_p, logs_p, m_q, logs_q), \
      (hidden_x, logw, logw_) = net_g(spec=spec.cuda(rank),
                                      spec_lengths=spec_lengths.cuda(rank),
                                      ph_IDs=ph_IDs.cuda(rank), 
                                      ph_IDs_lengths=ph_IDs_lengths.cuda(rank),
                                      dfs=dfs,  # slice後にcudaに乗せる
                                      sinewave=Sinewaves,
                                      speakerID=speakerID.cuda(rank))
      mel = spec_to_mel_torch(
          spec.cuda(rank), 
          hps["filter_length"], 
          hps["n_mel_channels"],  
          hps["sampling_rate"], 
          hps["mel_fmin"],  
          hps["mel_fmax"])

      # slice process
      y = commons.slice_segments(wav.cuda(rank), ids_slice * hps["hop_length"], hps["segments_size"])
      y_mel = commons.slice_segments(mel, ids_slice, hps["segments_size"] // hps["hop_length"] )
      f0_slice = commons.slice_segments(f0.cuda(rank), ids_slice , hps["segments_size"]//hps["hop_length"])

      y_hat_mel = mel_spectrogram_torch(y_hat.squeeze(1),
                                        hps["filter_length"],
                                        hps["n_mel_channels"],
                                        hps["sampling_rate"],
                                        hps["hop_length"],
                                        hps["win_length"],
                                        hps["mel_fmin"],
                                        hps["mel_fmax"])

      ### Duration Discriminator ###
      if net_dur_disc is not None:  
        y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x.detach(), ph_IDs_mask.detach(), logw.detach(), logw_.detach())
        with autocast(enabled=False):
          # TODO: I think need to mean using the mask, but for now, just mean all
          loss_dur_disc, losses_dur_disc_r, losses_dur_disc_g = discriminator_loss(y_dur_hat_r, y_dur_hat_g)
          loss_dur_disc_all = loss_dur_disc
        optim_dur_disc.zero_grad()
        scaler.scale(loss_dur_disc_all).backward()
        scaler.unscale_(optim_dur_disc)
        grad_norm_dur_disc = commons.clip_grad_value_(net_dur_disc.parameters(), None)
        scaler.step(optim_dur_disc)

      ### Discriminator ### 
      y_d_hat_g = net_d(y_hat.detach()) 
      y_d_hat_r = net_d(y)
      with autocast(enabled=False):
        loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
        loss_disc_all = loss_disc    

    optim_d.zero_grad()
    scaler.scale(loss_disc_all).backward()
    scaler.unscale_(optim_d)
    grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
    scaler.step(optim_d) 

    ### Generator ###
    with autocast(enabled=hps["fp16_run"]):
      y_d_hat_g, fmap_g = net_d(y_hat, return_fmaps=True)  # fake feature
      with torch.no_grad():
        _, fmap_r = net_d(y, return_fmaps=True) # real feature

      if net_dur_disc is not None:
        y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x, ph_IDs_mask, logw, logw_)
      with autocast(enabled=False):
        loss_length_l2 = torch.sum(l2_dur_loss).float()
        loss_mel = torch.nn.functional.l1_loss(y_mel, y_hat_mel) * hps["gain_mel"]
        loss_kl_spec = kl_loss(z_p, logs_q, m_p, logs_p, spec_mask) * hps["gain_kl_spec"]
        loss_fm = feature_loss(fmap_r, fmap_g) * hps["gain_fm"]
        loss_gen, losses_gen = generator_loss(y_d_hat_g) 
        loss_reg = reg_loss(s=excitation, y=y, f=f0_slice) * hps["gain_reg"] # sifi gan
        ############################################################
        loss_gen_all = loss_gen + loss_fm + loss_mel + loss_length_l2 + loss_kl_spec + loss_reg#+ loss_f0
        ############################################################
        if net_dur_disc is not None:
          loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
          loss_gen_all += loss_dur_gen

    losses = loss_gen_all
    optim_g.zero_grad()
    scaler.scale(losses).backward()   
    grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
    scaler.unscale_(optim_g)
    scaler.step(optim_g)
    scaler.update()

    # 記録 / 評価
    if rank==0:
      if global_step % hps["log_interval"] == 0:
        lr = optim_g.param_groups[0]['lr']
        #losses = 
        logger.info('Train Epoch: {} [{:.0f}%]'.format(
          epoch,
          100. * batch_idx / len(train_loader)))
        logger.info("steps:{}   [Generator]   ALL:{:.5f} | mel:{:.5f} | fm:{:.5f} | reg:{:.5f}| kl:{:.5f} | dur:{:.5f}"
                    .format(global_step, loss_gen_all.item(), loss_mel.item(), loss_fm.item(),loss_reg.item(), loss_kl_spec.item(),loss_length_l2.item()
                    ))
        logger.info("steps:{} [Discriminator] d_ALL:{:.5f} | d_dur_ALL:{:.5f} | "
                    .format(global_step, loss_disc_all.item(), loss_dur_disc_all.item()
                    ))
        
        scalar_dict = {"loss/g/ALL": loss_gen_all, 
                       "loss/d/ALL": loss_disc_all,
                       "loss/g/AdvDiscDur" :loss_dur_gen,
                       "loss/g/AdvDisc" :loss_gen,
                       "loss/g/Regularization" : loss_reg,
                       "loss/g/KL_div_spec" : loss_kl_spec,
                       "loss/g/FeatureMatching": loss_fm, 
                       "loss/g/MelReconstruction": loss_mel, 
                       "loss/g/Duration_SUM": loss_length_l2,  
                       "learning_rate": lr, 
                       "grad_norm/discriminator": grad_norm_d, 
                       "grad_norm/generator": grad_norm_g}
        
        scalar_dict.update({"loss/g/Adversarial_Deceive{}".format(i): v for i, v in enumerate(losses_gen)})
        scalar_dict.update({"loss/d/Adversarial_Real_{}".format(i): v for i, v in enumerate(losses_disc_r)})
        scalar_dict.update({"loss/d/Adversarial_Fake_{}".format(i): v for i, v in enumerate(losses_disc_g)})
        
        if net_dur_disc is not None:
          scalar_dict.update({"loss/d_dur/ALL"    : loss_dur_disc_all,})
          scalar_dict.update({"grad_norm/duration_discriminator": grad_norm_dur_disc})
          scalar_dict.update({"loss/d_dur/Adversarial_Real_{}".format(i): v for i, v in enumerate(losses_dur_disc_r)})
          scalar_dict.update({"loss/d_dur/Adversarial_Fake_{}".format(i): v for i, v in enumerate(losses_dur_disc_g)})

        # attnマップの作成
        length_scale=1
        w = torch.exp(logw) * ph_IDs_mask * length_scale
        w_ceil = torch.ceil(w)
        spec_mask = torch.unsqueeze(commons.sequence_mask(spec_lengths, spec.size(2)), 1).to(spec.dtype) # [B, 1, ph_len]
        ph_IDs_mask = torch.unsqueeze(commons.sequence_mask(ph_IDs_lengths, ph_IDs.size(1)), 1).to(ph_IDs.dtype) # [B, 1, ph_len]
        attn_mask     = torch.unsqueeze(spec_mask, 2) * torch.unsqueeze(ph_IDs_mask, -1)    # attn_mask = [B, 1, ph_len, note(word)_len]
        attn_pd       = commons.generate_path(duration=torch.unsqueeze(w_ceil,dim=1), mask=attn_mask.permute(0,1,3,2).cuda(rank))

        image_dict = { 
            "slice/mel_orgiginal": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
            "slice/mel_generated": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()), 
            "all/mel_full_length": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
            "all/attn_pd": utils.plot_alignment_to_numpy(attn_pd[0][0].data.cpu().numpy()),
            "all/attn_gt": utils.plot_alignment_to_numpy(attn_gt.permute(0,2,1)[0].data.cpu().numpy()),
            "all/attn_overlap": utils.plot_alignment_to_numpy(0.5*(attn_pd[0][0].data.cpu().numpy()) + 0.5*attn_gt.permute(0,2,1)[0].data.cpu().numpy())
        }

        utils.summarize(
          writer=writer,
          global_step=global_step, 
          images=image_dict,
          scalars=scalar_dict)

      if global_step % hps["eval_interval"] == 0:
        evaluate(hps=hps,
                 epoch=epoch,
                 nets=[net_g,net_d,net_dur_disc],
                 eval_loader=eval_loader, 
                 writer_eval=writer_eval)

        utils.save_checkpoint(model=net_g, 
                              optimizer=optim_g, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "G_{}.pth".format(global_step)))
        utils.save_checkpoint(model=net_d, 
                              optimizer=optim_d, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "D_{}.pth".format(global_step)))
        if net_dur_disc is not None:
          utils.save_checkpoint(model=net_dur_disc, 
                              optimizer=optim_dur_disc, 
                              learning_rate=hps["learning_rate"], 
                              iteration=epoch, 
                              checkpoint_path=os.path.join(hps["model_dir"], "DUR_{}.pth".format(global_step)))
    global_step += 1
    #torch.cuda.empty_cache()

  if rank == 0:
    logger.info('====> Epoch: {}'.format(epoch))
  
def evaluate(hps, epoch, nets, eval_loader, writer_eval):
    net_g,net_d,net_dur_disc = nets

    net_g.eval()
    net_d.eval()
    if net_dur_disc is not None:
      net_d.eval()

    reg_loss = ResidualLoss(sample_rate=hps["sampling_rate"],
                          fft_size=hps["filter_length"],
                          hop_size=hps["hop_length"],
                          f0_floor=hps["f0_min"],
                          f0_ceil=hps["f0_max"],
                          n_mels=hps["n_mel_channels"]).cuda()
    
    rank = 0
    
    count = 0
    loss_gen_all = 0
    loss_gen_adv_all = 0
    loss_disc_all= 0
    loss_dur_gen_all= 0
    #loss_f0_all= 0
    loss_reg_all = 0
    loss_kl_spec_all= 0
    loss_fm_all= 0
    loss_mel_all= 0
    loss_length_l2_all= 0
    loss_dur_disc_all = 0
    with torch.no_grad():
      for batch_idx, (wav,\
                      spec,            spec_lengths,\
                      f0,\
                      ph_IDs,          ph_IDs_lengths,\
                      dfs, \
                      Sinewaves, speakerID) in enumerate(eval_loader):
        
        count += 1

        # Forward Process
        with autocast(enabled=hps["fp16_run"]):

          # y_hat is sliced 
          y_hat, excitation, l2_dur_loss, attn_gt, ids_slice, ph_IDs_mask, spec_mask, \
          (z, z_p, m_p, logs_p, m_q, logs_q), \
          (hidden_x, logw, logw_) = net_g(spec=spec.cuda(rank),
                                              spec_lengths=spec_lengths.cuda(rank),
                                              ph_IDs=ph_IDs.cuda(rank), 
                                              ph_IDs_lengths=ph_IDs_lengths.cuda(rank),
                                              dfs=dfs,  # slice後にcudaに乗せる
                                              sinewave=Sinewaves,
                                              speakerID=speakerID.cuda(rank))

          mel = spec_to_mel_torch(
              spec.cuda(rank), 
              hps["filter_length"], 
              hps["n_mel_channels"],  
              hps["sampling_rate"], 
              hps["mel_fmin"],  
              hps["mel_fmax"])

          # sliced 
          y_mel = commons.slice_segments(mel, ids_slice, hps["segments_size"] // hps["hop_length"] )
          f0_slice = commons.slice_segments(f0.cuda(rank), ids_slice , hps["segments_size"]//hps["hop_length"])

          y_hat_mel = mel_spectrogram_torch(y_hat.squeeze(1),
                                            hps["filter_length"],
                                            hps["n_mel_channels"],
                                            hps["sampling_rate"],
                                            hps["hop_length"],
                                            hps["win_length"],
                                            hps["mel_fmin"],
                                            hps["mel_fmax"])

          y = commons.slice_segments(wav.cuda(rank), ids_slice * hps["hop_length"], hps["segments_size"])

          ### Duration Discriminator ###
          if net_dur_disc is not None:  
            y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x.detach(), ph_IDs_mask.detach(), logw.detach(), logw_.detach())
            with autocast(enabled=False):
              # TODO: I think need to mean using the mask, but for now, just mean all
              loss_dur_disc, losses_dur_disc_r, losses_dur_disc_g = discriminator_loss(y_dur_hat_r, y_dur_hat_g)
              loss_dur_disc_all += loss_dur_disc
              
          ### Discriminator ### 
          y_d_hat_g = net_d(y_hat.detach()) 
          y_d_hat_r = net_d(y)
          with autocast(enabled=False):
            loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
            loss_disc_all += loss_disc    

        ### Generator ###
        with autocast(enabled=hps["fp16_run"]):

          #y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)

          y_d_hat_g, fmap_g = net_d(y_hat, return_fmaps=True)  # fake feature
          with torch.no_grad():
            _, fmap_r = net_d(y, return_fmaps=True) # real feature

          if net_dur_disc is not None:
            y_dur_hat_r, y_dur_hat_g = net_dur_disc(hidden_x, ph_IDs_mask, logw, logw_)
          with autocast(enabled=False):
          
            loss_length_l2    = torch.sum(l2_dur_loss).float()
            loss_length_l2_all += loss_length_l2

            loss_mel = torch.nn.functional.l1_loss(y_mel, y_hat_mel) * hps["gain_mel"]
            loss_mel_all += loss_mel

            loss_kl_spec = kl_loss(z_p, logs_q, m_p, logs_p, spec_mask) * hps["gain_kl_spec"]
            loss_kl_spec_all += loss_kl_spec
            
            loss_fm = feature_loss(fmap_r, fmap_g) * hps["gain_fm"]
            loss_fm_all += loss_fm

            loss_gen, losses_gen = generator_loss(y_d_hat_g)
            loss_gen_adv_all += loss_gen 

            loss_reg = reg_loss(s=excitation, y=y, f = f0_slice) * hps["gain_reg"]
            loss_reg_all += loss_reg

            if net_dur_disc is not None:
              loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
              loss_dur_gen_all += loss_dur_gen

      loss_gen_all = loss_gen + loss_fm + loss_mel + loss_length_l2 + loss_kl_spec + loss_reg#+ loss_f0
            
      scalar_dict = {"loss/g/ALL"               : loss_gen_all / count, 
                     "loss/d/ALL"               : loss_disc_all/ count,
                     "loss/g/AdvDiscDur"        : loss_dur_gen_all/ count,
                     "loss/g/AdvDisc"           : loss_gen_adv_all/ count,
                     "loss/g/Regularization"    : loss_reg_all/ count,
                     "loss/g/KL_div_spec"       : loss_kl_spec_all/ count,
                     "loss/g/FeatureMatching"   : loss_fm_all/ count, 
                     "loss/g/MelReconstruction" : loss_mel_all/ count, 
                     "loss/g/Duration_SUM": loss_length_l2_all/ count, 
                     }
      if net_dur_disc is not None:
        scalar_dict.update({"loss/d_dur/ALL"    : loss_dur_disc_all})

    wav_gt                 = wav[0].view(1, 1,-1)
    f0_lengths             = torch.tensor(f0.size(2), dtype=torch.int64).view(-1).cuda(rank)
    ph_IDs                 = ph_IDs[0].view(1,-1).cuda(rank)                
    ph_IDs_len             = ph_IDs_lengths[0].view(-1).cuda(rank)   
    speakerID              = speakerID[0].view(-1).cuda(rank)      
    Sinewaves              = Sinewaves[0].view(1,1,-1)
    for idx, df in enumerate(dfs):
      dfs[idx] = df[0].view(1,1,-1)

    with torch.no_grad():
      # with f0_gt
      y_hat_with_f0gt, attn, y_mask = net_g.eval_infer(ph_IDs=ph_IDs.cuda(rank), 
                                                  ph_IDs_lengths=ph_IDs_len.cuda(rank),
                                                  speakerID=speakerID.cuda(rank),
                                                  dfs=dfs, 
                                                  sinewave=Sinewaves,
                                                  f0_lengths=f0_lengths
                                                  ) 

    try:
      y_hat_with_f0gt = y_hat_with_f0gt[0].to('cpu').detach().numpy().copy()
    except: 
      y_hat_with_f0gt = None

    image_dict = { 
    "all/attn": utils.plot_alignment_to_numpy(attn.permute(0,2,1)[0].data.cpu().numpy()),
    }

    audio_dict = {f"GroundTruth/audio_all" : wav_gt}
    if y_hat_with_f0gt is not None:
      audio_dict.update({f"gen/audio_with_f0gt_{batch_idx}"        : y_hat_with_f0gt})

    utils.summarize(
      writer=writer_eval,
      global_step=global_step, 
      scalars=scalar_dict,
      images=image_dict,
      audios = audio_dict,
      audio_sampling_rate=hps["sampling_rate"]
    )

    net_g.train()
    net_d.train()
    if net_dur_disc is not None:
      net_dur_disc.train()

def generate_graph(vector, 
                   label="NoName",
                   color="blue", 
                   title="Title",
                   x_label = 'Frames',
                   y_label = "y_labels",
                   figsize=(36,8)):
   
    fig = plt.figure(figsize=figsize) 
    x = np.arange(0, len(vector))
    plt.plot(x, vector, label=label, color=color)
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    #plt.ylim(y_val_min, y_val_max)
    fig.canvas.draw() 
    plot_image = fig.canvas.renderer._renderer  
    image_numpy = np.array(plot_image)
    plt.clf()
    plt.close()
    return image_numpy

def moving_average(loss,n=5):

    # Build a tensor list
    loss_list = []

    if len(loss_list) >= n:
        # drop the first loss
        loss_list =  loss_list[1:,:]
  
    loss_list.append(loss.data)

    baseline = loss-sum(loss_list)/len(loss_list)
    
    if baseline.all() == 0:
        return loss.detach()
    else:
        return baseline


# このファイルを実行したら、ココから始まる                  
if __name__ == "__main__":
  os.environ[
        "TORCH_DISTRIBUTED_DEBUG"
    ] = "DETAIL"
  main()

train_TTS_synthesizer.py
tonnetonne814/SiFi-VITS2-44100-Ja
import os
import glob
import sys
import argparse
import logging
import json
import subprocess
import numpy as np
from scipy.io.wavfile import read
import torch
import yaml

MATPLOTLIB_FLAG = False

logging.basicConfig(stream=sys.stdout, level=logging.WARNING)
logger = logging


def normalize_f0(f0, x_mask, uv, random_scale=True):
    # calculate means based on x_mask
    uv_sum = torch.sum(uv, dim=1, keepdim=True)
    uv_sum[uv_sum == 0] = 9999
    means = torch.sum(f0[:, 0, :] * uv, dim=1, keepdim=True) / uv_sum

    if random_scale:
        factor = torch.Tensor(f0.shape[0], 1).uniform_(0.8, 1.2).to(f0.device)
    else:
        factor = torch.ones(f0.shape[0], 1).to(f0.device)
    # normalize f0 based on means and factor
    f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)
    if torch.isnan(f0_norm).any():
        exit(0)
    return f0_norm * x_mask

def load_checkpoint(checkpoint_path, model, optimizer=None):
  assert os.path.isfile(checkpoint_path)
  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
  iteration = checkpoint_dict['iteration']
  learning_rate = checkpoint_dict['learning_rate']
  if optimizer is not None:
    optimizer.load_state_dict(checkpoint_dict['optimizer'])
  saved_state_dict = checkpoint_dict['model']
  if hasattr(model, 'module'):
    state_dict = model.module.state_dict()
  else:
    state_dict = model.state_dict()
  new_state_dict= {}
  for k, v in state_dict.items():
    try:
      new_state_dict[k] = saved_state_dict[k]
    except:
      try:
        new_state_dict[k] = saved_state_dict[f"_orig_mod.{k}"] # compile用にorig_modを付け加える
      except:
        logger.info("%s is not in the checkpoint" % k)
        new_state_dict[k] = v
  if hasattr(model, 'module'):
    model.module.load_state_dict(new_state_dict)
  else:
    model.load_state_dict(new_state_dict)
  logger.info("Loaded checkpoint '{}' (iteration {})" .format(
    checkpoint_path, iteration))
  return model, optimizer, learning_rate, iteration


def load_checkpoint_for_finetuning(checkpoint_path, model,hps, optimizer=None):
  assert os.path.isfile(checkpoint_path)
  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
  iteration = 1
  learning_rate = hps.train.learning_rate
  if optimizer is not None:
    optimizer.load_state_dict(checkpoint_dict['optimizer'])
  saved_state_dict = checkpoint_dict['model']
  if hasattr(model, 'module'):
    state_dict = model.module.state_dict()
  else:
    state_dict = model.state_dict()
  new_state_dict= {}
  for k, v in state_dict.items():
    try:
      new_state_dict[k] = saved_state_dict[k]
    except:
      logger.info("%s is not in the checkpoint" % k)
      new_state_dict[k] = v
  if hasattr(model, 'module'):
    model.module.load_state_dict(new_state_dict)
  else:
    model.load_state_dict(new_state_dict)
  logger.info("Loaded checkpoint for finetuning'{}' (iteration {})" .format(
    checkpoint_path, iteration))
  return model, optimizer, learning_rate, iteration

def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
  logger.info("Saving model and optimizer state at iteration {} to {}".format(
    iteration, checkpoint_path))
  if hasattr(model, 'module'):
    state_dict = model.module.state_dict()
  else:
    state_dict = model.state_dict()
  torch.save({'model': state_dict,
              'iteration': iteration,
              'optimizer': optimizer.state_dict(),
              'learning_rate': learning_rate}, checkpoint_path)

def load_model_diffsize(checkpoint_path, model,hps, optimizer=None):
    assert os.path.isfile(checkpoint_path)
    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')["model"]

    if hasattr(model, 'module'):
        state_dict = model.module.state_dict()
    else:
        state_dict = model.state_dict()

    for k, v in checkpoint_dict.items():
        if k in state_dict and state_dict[k].size() == v.size():
            state_dict[k] = v
        else:
          k=k.replace("_orig_mod.", "")
          if k in state_dict and state_dict[k].size() == v.size():
              state_dict[k] = v
          else:
            print("[WARNING] Parameter mismatch :",k)
    
    if hasattr(model, 'module'):
        model.module.load_state_dict(state_dict, strict=False)
    else:
        model.load_state_dict(state_dict, strict=False)
        
    return model


def summarize(writer, global_step, scalars={}, vectors={}, histograms={}, images={}, audios={}, audio_sampling_rate=44100):
  for k, v in scalars.items():
    writer.add_scalar(k, v, global_step)
  for k, v in vectors.items():
    for i, d in enumerate(v):
      writer.add_scalar(k, d, i)
  for k, v in histograms.items():
    writer.add_histogram(k, v, global_step)
  for k, v in images.items():
    writer.add_image(k, v, global_step, dataformats='HWC')
  for k, v in audios.items():
    writer.add_audio(k, v, global_step, audio_sampling_rate)


def latest_checkpoint_path(dir_path, regex="G_*.pth"):
  f_list = glob.glob(os.path.join(dir_path, regex))
  f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
  if len(f_list) == 0:
    return "checkpoint_path_is_None"
  x = f_list[-1]
  print(x)
  return x


def plot_spectrogram_to_numpy(spectrogram):
  global MATPLOTLIB_FLAG
  if not MATPLOTLIB_FLAG:
    import matplotlib
    matplotlib.use("Agg")
    MATPLOTLIB_FLAG = True
    mpl_logger = logging.getLogger('matplotlib')
    mpl_logger.setLevel(logging.WARNING)
  import matplotlib.pylab as plt
  import numpy as np
  
  fig, ax = plt.subplots(figsize=(14,6))
  im = ax.imshow(spectrogram, aspect="auto", origin="lower",
                  interpolation='none')
  plt.colorbar(im, ax=ax)
  plt.xlabel("Frames")
  plt.ylabel("Channels")
  plt.tight_layout()

  fig.canvas.draw()
  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
  plt.close()
  return data


def plot_alignment_to_numpy(alignment, info=None):
  global MATPLOTLIB_FLAG
  if not MATPLOTLIB_FLAG:
    import matplotlib
    matplotlib.use("Agg")
    MATPLOTLIB_FLAG = True
    mpl_logger = logging.getLogger('matplotlib')
    mpl_logger.setLevel(logging.WARNING)
  import matplotlib.pylab as plt
  import numpy as np

  fig, ax = plt.subplots(figsize=(6, 4))
  im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',
                  interpolation='none')
  fig.colorbar(im, ax=ax)
  xlabel = 'Decoder timestep'
  if info is not None:
      xlabel += '\n\n' + info
  plt.xlabel(xlabel)
  plt.ylabel('Encoder timestep')
  plt.tight_layout()

  fig.canvas.draw()
  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
  plt.close()
  return data

def load_wav_to_torch(full_path):
  sampling_rate, wav = read(full_path.replace("\\", "/")) ### modify .replace("\\", "/") ###
  
  if len(wav.shape) == 2:
      wav = wav[:, 0]
  if wav.dtype == np.int16:
      wav = wav / 32768.0
  elif wav.dtype == np.int32:
      wav = wav / 2147483648.0
  elif wav.dtype == np.uint8:
      wav = (wav - 128) / 128.0
  wav = wav.astype(np.float32)
  
  if sampling_rate != 44100:
    print("ERROR SAMPLINGRATE")
    
  return torch.FloatTensor(wav), sampling_rate

def load_filepaths_and_text(filename, split="|"):
  with open(filename, encoding='utf-8') as f:
    filepaths_and_text = [line.strip().split(split) for line in f]
  return filepaths_and_text


def get_hparams(init=True):

  ### .pyを実行する時の引数はここで処理する ###
  parser = argparse.ArgumentParser()
  parser.add_argument('-c', '--config_yaml', type=str  , default="./configs/config.yaml")
  parser.add_argument('-m', '--model'      , type=str  , default='NoNameModel')
  args = parser.parse_args()
  ##########################################

  model_dir = os.path.join("./logs", args.model)
  if not os.path.exists(model_dir):
    os.makedirs(model_dir)

  config_path = args.config_yaml
  config_save_path = os.path.join(model_dir, "config_backup.yaml")

  if init:
    with open(config_path, mode="r", encoding="utf-8") as f:
      hparams = yaml.safe_load(f)
    with open(config_save_path, "w", encoding="utf-8") as yf:
        yaml.dump(hparams, yf, default_flow_style=False)
  else:
    with open(config_save_path, "r") as f:
      hparams = f.read()

  hparams["model_dir"] = model_dir

  return hparams


def get_hparams_from_dir(model_dir):
  config_save_path = os.path.join(model_dir, "config.json")
  with open(config_save_path, "r") as f:
    data = f.read()
  config = json.loads(data)

  hparams =HParams(**config)
  hparams.model_dir = model_dir
  return hparams


def get_hparams_from_file(config_path):
  with open(config_path, "r") as f:
    data = f.read()
  config = json.loads(data)

  hparams =HParams(**config)
  return hparams


def check_git_hash(model_dir):
  source_dir = os.path.dirname(os.path.realpath(__file__))
  if not os.path.exists(os.path.join(source_dir, ".git")):
    logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(
      source_dir
    ))
    return

  cur_hash = subprocess.getoutput("git rev-parse HEAD")

  path = os.path.join(model_dir, "githash")
  if os.path.exists(path):
    saved_hash = open(path).read()
    if saved_hash != cur_hash:
      logger.warn("git hash values are different. {}(saved) != {}(current)".format(
        saved_hash[:8], cur_hash[:8]))
  else:
    open(path, "w").write(cur_hash)


def get_logger(model_dir, filename="train.log"):
  global logger
  logger = logging.getLogger(os.path.basename(model_dir))
  logger.setLevel(logging.DEBUG)
  
  formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
  if not os.path.exists(model_dir):
    os.makedirs(model_dir)
  h = logging.FileHandler(os.path.join(model_dir, filename))
  h.setLevel(logging.DEBUG)
  h.setFormatter(formatter)
  logger.addHandler(h)
  return logger


class HParams():
  def __init__(self, **kwargs):
    for k, v in kwargs.items():
      if type(v) == dict:
        v = HParams(**v)
      self[k] = v
    
  def keys(self):
    return self.__dict__.keys()

  def items(self):
    return self.__dict__.items()

  def values(self):
    return self.__dict__.values()

  def __len__(self):
    return len(self.__dict__)

  def __getitem__(self, key):
    return getattr(self, key)

  def __setitem__(self, key, value):
    return setattr(self, key, value)

  def __contains__(self, key):
    return key in self.__dict__

  def __repr__(self):
    return self.__dict__.__repr__()

utils.py
tonnetonne814/SiFi-VITS2-44100-Ja