Update models.py

models.py

@@ -14,6 +14,74 @@ from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
 
 from commons import init_weights, get_padding
 from text import symbols, num_tones, num_languages
+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
+
+  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 TransformerCouplingBlock(nn.Module):
     def __init__(self,
@@ -615,7 +683,7 @@ class SynthesizerTrn(nn.Module):
 
         z_slice, ids_slice = commons.rand_slice_segments(z, y_lengths, self.segment_size)
         o = self.dec(z_slice, g=g)
-        return o, l_length, attn, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+        return o, l_length, attn, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q), (logw, logw_)
 
     def infer(self, x, x_lengths, sid, tone, language, bert, noise_scale=.667, length_scale=1, noise_scale_w=0.8, max_len=None, sdp_ratio=0,y=None):
         x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)