From 81be46d594709e35bbd730c01dc5890e1bd7e2f0 Mon Sep 17 00:00:00 2001
From: keonlee9420 <keonlee9420@gmail.com>
Date: Fri, 8 Oct 2021 23:18:45 +0900
Subject: [PATCH] add Relative Multi-Head Attention and unify masking

---
 README.md                      |   3 +-
 config/LJSpeech/train.yaml     |   4 +-
 model/PortaSpeech.py           |   6 +-
 model/blocks.py                | 280 ++++++++++++---------------------
 model/linguistic_encoder.py    | 101 ++++--------
 model/loss.py                  |   3 -
 model/variational_generator.py |  41 ++---
 utils/tools.py                 |   8 +-
 8 files changed, 156 insertions(+), 290 deletions(-)

diff --git a/README.md b/README.md
index 2088da0..e218c96 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@ PyTorch Implementation of [PortaSpeech: Portable and High-Quality Generative Tex
 | Module | Normal | Small | Normal (paper) | Small (paper) |
 | :----- | :-----: | :-----: | :-----: | :-----: |
 | *Total* | 34.3M | 9.6M | 21.8M | 6.7M
-| *LinguisticEncoder* | 14M | 3.4M | - | -
+| *LinguisticEncoder* | 14M | 3.5M | - | -
 | *VariationalGenerator* | 11M | 2.8M | - | -
 | *FlowPostNet* | 9.3M | 3.4M | - | -
 
@@ -122,7 +122,6 @@ to serve TensorBoard on your localhost.
 - For vocoder, **HiFi-GAN** and **MelGAN** are supported.
 - Add convolution layer and residual layer in **VariationalGenerator** to match the shape of conditioner and output.
 - No ReLU activation and LayerNorm in **VariationalGenerator** for convergence of word-to-phoneme alignment of **LinguisticEncoder**.
-- Use absolute positional encoding in **LinguisticEncoder** instead of relative positional encoding.
 - Will be extended to a **multi-speaker TTS**.
 <!-- - Two options for embedding for the **multi-speaker TTS** setting: training speaker embedder from scratch or using a pre-trained [philipperemy's DeepSpeaker](https://github.com/philipperemy/deep-speaker) model (as [STYLER](https://github.com/keonlee9420/STYLER) did). You can toggle it by setting the config (between `'none'` and `'DeepSpeaker'`).
 - DeepSpeaker on VCTK dataset shows clear identification among speakers. The following figure shows the T-SNE plot of extracted speaker embedding.
diff --git a/config/LJSpeech/train.yaml b/config/LJSpeech/train.yaml
index 8e234ff..fc354a4 100644
--- a/config/LJSpeech/train.yaml
+++ b/config/LJSpeech/train.yaml
@@ -15,10 +15,10 @@ optimizer:
   grad_clip_thresh: 1.0
   grad_acc_step: 1
   warm_up_step: 4000
-  anneal_steps: [300000, 400000, 500000]
+  anneal_steps: [100000, 200000, 300000]
   anneal_rate: 0.3
 step:
-  total_step: 900000
+  total_step: 500000
   log_step: 100
   synth_step: 1000
   val_step: 1000
diff --git a/model/PortaSpeech.py b/model/PortaSpeech.py
index 78c38ad..4e47817 100644
--- a/model/PortaSpeech.py
+++ b/model/PortaSpeech.py
@@ -17,7 +17,7 @@ def __init__(self, preprocess_config, model_config):
         super(PortaSpeech, self).__init__()
         self.model_config = model_config
 
-        self.linguistic_encoder = LinguisticEncoder(model_config, abs_mha=True)
+        self.linguistic_encoder = LinguisticEncoder(model_config)
         self.variational_generator = VariationalGenerator(
             preprocess_config, model_config)
         self.postnet = FlowPostNet(preprocess_config, model_config)
@@ -104,14 +104,14 @@ def forward(
                 mels, mel_lens, mel_masks, output)
             postnet_output = self.postnet(
                 mels.transpose(1, 2),
-                ~mel_masks.unsqueeze(1),
+                mel_masks.unsqueeze(1),
                 g=(out_residual + residual).transpose(1, 2),
             )
         else:
             _, out_residual, dist_info = self.variational_generator.inference(
                 mel_lens, mel_masks, output)
             output = self.postnet.inference(
-                ~mel_masks.unsqueeze(1),
+                mel_masks.unsqueeze(1),
                 g=(out_residual + residual).transpose(1, 2),
             )
             postnet_output = None
diff --git a/model/blocks.py b/model/blocks.py
index 72e15b5..3ef488e 100644
--- a/model/blocks.py
+++ b/model/blocks.py
@@ -22,6 +22,8 @@ def convert_pad_shape(pad_shape):
 
 
 class Flip(nn.Module):
+    """ Flip Layer """
+
     def forward(self, x, *args, reverse=False, **kwargs):
         x = torch.flip(x, [1])
         if not reverse:
@@ -212,159 +214,9 @@ def forward(self, x):
         return x
 
 
-class FFTBlock(nn.Module):
-    """ FFT Block with AbsMHA """
-
-    def __init__(self, d_model, n_head, d_k, d_v, d_inner, kernel_size, dropout=0.0):
-        super(FFTBlock, self).__init__()
-        self.slf_attn = MultiHeadAttention(
-            n_head, d_model, d_k, d_v, dropout=dropout)
-        self.pos_ffn = PositionwiseFeedForward(
-            d_model, d_inner, kernel_size, dropout=dropout
-        )
-
-    def forward(self, enc_input, mask=None, slf_attn_mask=None):
-        enc_output, enc_slf_attn = self.slf_attn(
-            enc_input, enc_input, enc_input, mask_1=slf_attn_mask
-        )
-        if mask is not None:
-            enc_output = enc_output.masked_fill(mask.unsqueeze(-1), 0)
-
-        enc_output = self.pos_ffn(enc_output)
-        if mask is not None:
-            enc_output = enc_output.masked_fill(mask.unsqueeze(-1), 0)
-
-        return enc_output, enc_slf_attn
-
-
-class MultiHeadAttention(nn.Module):
-    """ Multi-Head Attention """
-
-    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.0):
-        super(MultiHeadAttention, self).__init__()
-
-        self.n_head = n_head
-        self.d_k = d_k
-        self.d_v = d_v
-
-        self.w_qs = LinearNorm(d_model, n_head * d_k)
-        self.w_ks = LinearNorm(d_model, n_head * d_k)
-        self.w_vs = LinearNorm(d_model, n_head * d_v)
-
-        self.attention = ScaledDotProductAttention(
-            temperature=np.power(d_k, 0.5))
-        self.layer_norm = nn.LayerNorm(d_model)
-
-        self.fc = LinearNorm(n_head * d_v, d_model)
-
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, q, k, v, mask_1=None, mask_2=None, mapping_mask=None, indivisual_attn=False):
-
-        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
-
-        sz_b, len_q, _ = q.size()
-        sz_b, len_k, _ = k.size()
-        sz_b, len_v, _ = v.size()
-
-        residual = q
-
-        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
-        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
-        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
-        q = q.permute(2, 0, 1, 3).contiguous().view(-1,
-                                                    len_q, d_k)  # (n*b) x lq x dk
-        k = k.permute(2, 0, 1, 3).contiguous().view(-1,
-                                                    len_k, d_k)  # (n*b) x lk x dk
-        v = v.permute(2, 0, 1, 3).contiguous().view(-1,
-                                                    len_v, d_v)  # (n*b) x lv x dv
-
-        if mask_1 is not None:
-            mask_1 = mask_1.repeat(n_head, 1, 1)  # (n*b) x .. x ..
-        if mask_2 is not None:
-            mask_2 = mask_2.repeat(n_head, 1, 1)  # (n*b) x .. x ..
-        if mapping_mask is not None:
-            mapping_mask = mapping_mask.repeat(n_head, 1, 1)  # (n*b) x .. x ..
-        output, attn = self.attention(
-            q, k, v, mask_1=mask_1, mask_2=mask_2, mapping_mask=mapping_mask)
-
-        output = output.view(n_head, sz_b, len_q, d_v)
-        output = (
-            output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1)
-        )  # b x lq x (n*dv)
-
-        output = self.dropout(self.fc(output))
-        output = self.layer_norm(output + residual)
-
-        if indivisual_attn:
-            attn = attn.view(n_head, sz_b, len_q, len_k)
-
-        return output, attn
-
-
-class ScaledDotProductAttention(nn.Module):
-    """ Scaled Dot-Product Attention """
-
-    def __init__(self, temperature):
-        super(ScaledDotProductAttention, self).__init__()
-        self.temperature = temperature
-        self.softmax = nn.Softmax(dim=2)
-
-    def forward(self, q, k, v, mask_1=None, mask_2=None, mapping_mask=None):
-
-        attn = torch.bmm(q, k.transpose(1, 2))
-        attn = attn / self.temperature
-
-        if mask_1 is not None:
-            attn = attn.masked_fill(mask_1, -np.inf)
-        attn = self.softmax(attn)
-
-        if mask_2 is not None:
-            attn = attn.masked_fill(mask_2, 0.)
-        if mapping_mask is not None:
-            attn = attn.masked_fill(mapping_mask, 0.)
-        output = torch.bmm(attn, v)
-
-        return output, attn
-
-
-class PositionwiseFeedForward(nn.Module):
-    """ A two-feed-forward-layer """
-
-    def __init__(self, d_in, d_hid, kernel_size, dropout=0.0):
-        super(PositionwiseFeedForward, self).__init__()
-
-        # Use Conv1D
-        # position-wise
-        self.w_1 = nn.Conv1d(
-            d_in,
-            d_hid,
-            kernel_size=kernel_size,
-            padding=(kernel_size - 1) // 2,
-        )
-        # position-wise
-        self.w_2 = nn.Conv1d(
-            d_hid,
-            d_in,
-            kernel_size=kernel_size,
-            padding=(kernel_size - 1) // 2,
-        )
-
-        self.layer_norm = nn.LayerNorm(d_in)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        residual = x
-        output = x.transpose(1, 2)
-        output = self.w_2(F.relu(self.w_1(output)))
-        output = output.transpose(1, 2)
-        output = self.dropout(output)
-        output = self.layer_norm(output + residual)
-
-        return output
-
-
 class NonCausalWaveNet(torch.nn.Module):
+    """ Non-Causal WaveNet """
+
     def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, n_sqz=1):
         super(NonCausalWaveNet, self).__init__()
         assert(kernel_size % 2 == 1)
@@ -440,7 +292,7 @@ def remove_weight_norm(self):
 
 
 class RelativeFFTBlock(nn.Module):
-    """ FFT Block with RelMHA """
+    """ FFT Block with Relative Multi-Head Attention """
 
     def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=None, block_length=None):
         super(RelativeFFTBlock, self).__init__()
@@ -459,7 +311,7 @@ def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_s
         self.ffn_layers = nn.ModuleList()
         self.norm_layers_2 = nn.ModuleList()
         for i in range(self.n_layers):
-            self.attn_layers.append(RelativeMultiHeadAttention(hidden_channels, hidden_channels, n_heads,
+            self.attn_layers.append(RelativeSelfAttention(hidden_channels, hidden_channels, n_heads,
                                     window_size=window_size, p_dropout=p_dropout, block_length=block_length))
             self.norm_layers_1.append(LayerNorm(hidden_channels))
             self.ffn_layers.append(FFN(
@@ -481,9 +333,11 @@ def forward(self, x, x_mask):
         return x
 
 
-class RelativeMultiHeadAttention(nn.Module):
+class RelativeSelfAttention(nn.Module):
+    """ Relative Multi-Head Attention """
+
     def __init__(self, channels, out_channels, n_heads, window_size=None, heads_share=True, p_dropout=0., block_length=None, proximal_bias=False, proximal_init=False):
-        super().__init__()
+        super(RelativeSelfAttention, self).__init__()
         assert channels % n_heads == 0
 
         self.channels = channels
@@ -575,20 +429,10 @@ def attention(self, query, key, value, mask=None):
         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
 
@@ -609,10 +453,6 @@ def _get_relative_embeddings(self, relative_embeddings, length):
         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, convert_pad_shape(
@@ -629,10 +469,6 @@ def _relative_position_to_absolute_position(self, x):
         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, convert_pad_shape(
@@ -645,11 +481,8 @@ def _absolute_position_to_relative_position(self, x):
         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]
+        """
+        Bias for self-attention to encourage attention to close positions.
         """
         r = torch.arange(length, dtype=torch.float32)
         diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
@@ -702,3 +535,92 @@ def forward(self, x, x_mask):
         x = self.drop(x)
         x = self.conv_2(x * x_mask)
         return x * x_mask
+
+
+class WordToPhonemeAttention(nn.Module):
+    """ Word-to-Phoneme Multi-Head Attention """
+
+    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.0):
+        super(WordToPhonemeAttention, self).__init__()
+
+        self.n_head = n_head
+        self.d_k = d_k
+        self.d_v = d_v
+
+        self.w_qs = LinearNorm(d_model, n_head * d_k)
+        self.w_ks = LinearNorm(d_model, n_head * d_k)
+        self.w_vs = LinearNorm(d_model, n_head * d_v)
+
+        self.attention = ScaledDotProductAttention(
+            temperature=np.power(d_k, 0.5))
+        self.layer_norm = nn.LayerNorm(d_model)
+
+        self.fc = LinearNorm(n_head * d_v, d_model)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, q, k, v, mask_1=None, mask_2=None, mapping_mask=None, indivisual_attn=False):
+
+        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+
+        sz_b, len_q, _ = q.size()
+        sz_b, len_k, _ = k.size()
+        sz_b, len_v, _ = v.size()
+
+        residual = q
+
+        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
+        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
+        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
+        q = q.permute(2, 0, 1, 3).contiguous().view(-1,
+                                                    len_q, d_k)  # (n*b) x lq x dk
+        k = k.permute(2, 0, 1, 3).contiguous().view(-1,
+                                                    len_k, d_k)  # (n*b) x lk x dk
+        v = v.permute(2, 0, 1, 3).contiguous().view(-1,
+                                                    len_v, d_v)  # (n*b) x lv x dv
+
+        if mask_1 is not None:
+            mask_1 = mask_1.repeat(n_head, 1, 1)  # (n*b) x .. x ..
+        if mask_2 is not None:
+            mask_2 = mask_2.repeat(n_head, 1, 1)  # (n*b) x .. x ..
+        if mapping_mask is not None:
+            mapping_mask = mapping_mask.repeat(n_head, 1, 1)  # (n*b) x .. x ..
+        output, attn = self.attention(
+            q, k, v, mask_1=mask_1, mask_2=mask_2, mapping_mask=mapping_mask)
+
+        output = output.view(n_head, sz_b, len_q, d_v)
+        output = (
+            output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1)
+        )  # b x lq x (n*dv)
+
+        output = self.dropout(self.fc(output))
+        output = self.layer_norm(output + residual)
+
+        if indivisual_attn:
+            attn = attn.view(n_head, sz_b, len_q, len_k)
+
+        return output, attn
+
+
+class ScaledDotProductAttention(nn.Module):
+    def __init__(self, temperature):
+        super(ScaledDotProductAttention, self).__init__()
+        self.temperature = temperature
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, q, k, v, mask_1=None, mask_2=None, mapping_mask=None):
+
+        attn = torch.bmm(q, k.transpose(1, 2))
+        attn = attn / self.temperature
+
+        if mask_1 is not None:
+            attn = attn.masked_fill(mask_1==0., -np.inf)
+        attn = self.softmax(attn)
+
+        if mask_2 is not None:
+            attn = attn * mask_2
+        if mapping_mask is not None:
+            attn = attn * mapping_mask
+        output = torch.bmm(attn, v)
+
+        return output, attn
diff --git a/model/linguistic_encoder.py b/model/linguistic_encoder.py
index 7947999..b6c4f70 100644
--- a/model/linguistic_encoder.py
+++ b/model/linguistic_encoder.py
@@ -8,9 +8,8 @@
 
 from .blocks import (
     ConvNorm,
-    FFTBlock,
-    MultiHeadAttention,
     RelativeFFTBlock,
+    WordToPhonemeAttention,
 )
 from text.symbols import symbols
 
@@ -41,7 +40,7 @@ def get_posi_angle_vec(position):
 class LinguisticEncoder(nn.Module):
     """ Linguistic Encoder """
 
-    def __init__(self, config, abs_mha=True):
+    def __init__(self, config):
         super(LinguisticEncoder, self).__init__()
 
         n_position = config["max_seq_len"] + 1
@@ -79,47 +78,28 @@ def __init__(self, config, abs_mha=True):
             requires_grad=True,
         )
 
-        self.abs_mha = abs_mha
-        if abs_mha:
-            self.phoneme_encoder = nn.ModuleList(
-                [
-                    FFTBlock(
-                        d_model, n_head, d_k, d_v, d_inner, kernel_size  # , dropout=dropout
-                    )
-                    for _ in range(n_layers)
-                ]
-            )
-            self.word_encoder = nn.ModuleList(
-                [
-                    FFTBlock(
-                        d_model, n_head, d_k, d_v, d_inner, kernel_size  # , dropout=dropout
-                    )
-                    for _ in range(n_layers)
-                ]
-            )
-        else:
-            self.phoneme_encoder = RelativeFFTBlock(
-                d_model,
-                d_inner,
-                n_head,
-                n_layers,
-                kernel_size,
-                # dropout,
-                window_size,
-            )
-            self.word_encoder = RelativeFFTBlock(
-                d_model,
-                d_inner,
-                n_head,
-                n_layers,
-                kernel_size,
-                # dropout,
-                window_size,
-            )
+        self.phoneme_encoder = RelativeFFTBlock(
+            hidden_channels=d_model,
+            filter_channels=d_inner,
+            n_heads=n_head,
+            n_layers=n_layers,
+            kernel_size=kernel_size,
+            # p_dropout=dropout,
+            window_size=window_size,
+        )
+        self.word_encoder = RelativeFFTBlock(
+            hidden_channels=d_model,
+            filter_channels=d_inner,
+            n_heads=n_head,
+            n_layers=n_layers,
+            kernel_size=kernel_size,
+            # p_dropout=dropout,
+            window_size=window_size,
+        )
         self.length_regulator = LengthRegulator()
         self.duration_predictor = VariancePredictor(config)
 
-        self.w2p_attn = MultiHeadAttention(
+        self.w2p_attn = WordToPhonemeAttention(
             n_head, d_model, d_k, d_v  # , dropout=dropout
         )
 
@@ -137,7 +117,7 @@ def get_mapping_mask(self, q, kv, dur_w, wb, src_w_len):
             for i in range(1, len(w)):
                 mask[b, w[i-1]:w[i], p[i-1]:p[i]
                      ] = torch.zeros(w[i]-w[i-1], p[i]-p[i-1], device=device)
-        return mask == 1.
+        return mask == 0.
 
     def add_position_enc(self, src_seq, position_enc=None, coef=None):
         batch_size, max_len = src_seq.shape[0], src_seq.shape[1]
@@ -174,22 +154,7 @@ def get_rel_coef(self, dur, dur_len, mask):
                 idx_b += list(range(d_i))
             idx.append(torch.tensor(idx_b).to(device))
             # assert L[-1].shape == idx[-1].shape
-        return torch.div(pad(idx).to(device), pad(L).masked_fill(mask, 1.).to(device))
-
-    def encode(self, encoder, src_seq, mask, src_emb=True, return_attns=False):
-        slf_attn_list = []
-        slf_attn_mask = mask.unsqueeze(1).expand(-1, src_seq.shape[1], -1)
-        if src_emb:
-            src_seq = self.src_emb(src_seq)
-        # Will be replaced with Relative PE
-        enc_out = self.add_position_enc(src_seq)
-        for enc_layer in encoder:
-            enc_out, enc_slf_attn = enc_layer(
-                enc_out, mask=mask, slf_attn_mask=slf_attn_mask
-            )
-            if return_attns:
-                slf_attn_list += [enc_slf_attn]
-        return enc_out, slf_attn_list
+        return torch.div(pad(idx).to(device), pad(L).masked_fill(mask==0., 1.).to(device))
 
     def forward(
         self,
@@ -206,25 +171,17 @@ def forward(
         return_attns=False,
     ):
         # Phoneme Encoding
-        if self.abs_mha:
-            enc_out_p, slf_attn_list_p = self.encode(
-                self.phoneme_encoder, src_seq, p_mask, return_attns=return_attns)
-        else:
-            src_seq = self.src_emb(src_seq)
-            enc_out_p = self.phoneme_encoder(src_seq.transpose(
-                1, 2), p_mask.unsqueeze(1)).transpose(1, 2)
+        src_seq = self.src_emb(src_seq)
+        enc_out_p = self.phoneme_encoder(src_seq.transpose(
+            1, 2), p_mask.unsqueeze(1)).transpose(1, 2)
 
         # Word-level Pooing
         src_seq_w = word_level_pooling(
             enc_out_p, src_len, wb, src_w_len, reduce_mean=True)
 
         # Word Encoding
-        if self.abs_mha:
-            enc_out_w, slf_attn_list_w = self.encode(
-                self.word_encoder, src_seq_w, w_mask, src_emb=False, return_attns=return_attns)
-        else:
-            enc_out_w = self.word_encoder(src_seq_w.transpose(
-                1, 2), w_mask.unsqueeze(1)).transpose(1, 2)
+        enc_out_w = self.word_encoder(src_seq_w.transpose(
+            1, 2), w_mask.unsqueeze(1)).transpose(1, 2)
 
         # Phoneme-level Duration Prediction
         log_duration_p_prediction = self.duration_predictor(enc_out_p, p_mask)
@@ -374,6 +331,6 @@ def forward(self, encoder_output, mask):
         out = out.squeeze(-1)
 
         if mask is not None:
-            out = out.masked_fill(mask, 0.0)
+            out = out * mask
 
         return out
diff --git a/model/loss.py b/model/loss.py
index 6ed4b36..7bc0d2f 100644
--- a/model/loss.py
+++ b/model/loss.py
@@ -62,9 +62,6 @@ def forward(self, inputs, predictions):
             src_w_masks,
             _,
         ) = predictions
-        src_masks = ~src_masks
-        src_w_masks = ~src_w_masks
-        mel_masks = ~mel_masks
         log_duration_targets = torch.log(duration_roundeds.float() + 1)
         mel_targets = mel_targets[:, : mel_masks.shape[1], :]
         mel_masks = mel_masks[:, :mel_masks.shape[1]]
diff --git a/model/variational_generator.py b/model/variational_generator.py
index 4a22ec2..ef1302a 100644
--- a/model/variational_generator.py
+++ b/model/variational_generator.py
@@ -173,40 +173,34 @@ def forward(self, mel, mel_len, mel_mask, h_text):
         h_text_e = self.cond_layer_e(h_text)  # [B, H, L']
         h_text_d = self.cond_layer_d(h_text)  # [B, H, L']
         mel_mask_conv = self.get_conv_mask(
-            mel_len, h_text_f.shape[2], mel_mask)
+            mel_len, h_text_f.shape[2], mel_mask).unsqueeze(-1)
 
         # Encoding
         x = self.enc_conv(mel)
         x = x.contiguous().transpose(1, 2)
-        x = self.enc_wn(x, g=h_text_e).masked_fill(
-            mel_mask_conv.unsqueeze(1), 0.)
+        x = self.enc_wn(x, g=h_text_e) * mel_mask_conv.transpose(1, 2)
         x = x.contiguous().transpose(1, 2)
         x = self.latent_enc_prj(x)
 
         # # Reparameterization
         m_q, logs_q = torch.split(x, self.latent_hidden, dim=-1)
-        m_q, logs_q = m_q.masked_fill(mel_mask_conv.unsqueeze(
-            -1), 0.), logs_q.masked_fill(mel_mask_conv.unsqueeze(-1), 0.)
-        z_q = reparameterize(m_q, logs_q).masked_fill(
-            mel_mask_conv.unsqueeze(-1), 0.)
+        m_q, logs_q = m_q * mel_mask_conv, logs_q * mel_mask_conv
+        z_q = reparameterize(m_q, logs_q) * mel_mask_conv
 
         # Prior VP FLow
-        z_p = self.flow(z_q.transpose(1, 2), x_mask=~
-                        mel_mask_conv.unsqueeze(1), g=h_text_f, reverse=False)
+        z_p = self.flow(z_q.transpose(1, 2), x_mask=mel_mask_conv.transpose(
+            1, 2), g=h_text_f, reverse=False)
 
         # Decoding
         x = self.latent_dec_prj(z_q)
         x = x.contiguous().transpose(1, 2)
-        x = self.dec_wn(x, g=h_text_d).masked_fill(
-            mel_mask_conv.unsqueeze(1), 0.)
+        x = self.dec_wn(x, g=h_text_d) * mel_mask_conv.transpose(1, 2)
         x = x.contiguous().transpose(1, 2)
         mel_res = self.dec_conv(x)
-        mel_res = self.trim_output(mel_res, mel_mask.shape[1])
-        mel_res = mel_res.masked_fill(mel_mask.unsqueeze(-1), 0.)
-        residual = self.residual_layer(
-            mel_res).masked_fill(mel_mask.unsqueeze(-1), 0.)
+        mel_res = self.trim_output(mel_res, mel_mask.shape[1]) * mel_mask.unsqueeze(-1)
+        residual = self.residual_layer(mel_res) * mel_mask.unsqueeze(-1)
 
-        return mel_res, residual, (z_p, logs_q.transpose(1, 2), ~mel_mask_conv.unsqueeze(1))
+        return mel_res, residual, (z_p, logs_q.transpose(1, 2), mel_mask_conv.transpose(1, 2))
 
     def inference(self, mel_len, mel_mask, h_text):
         """
@@ -222,25 +216,22 @@ def inference(self, mel_len, mel_mask, h_text):
         h_text_f = self.cond_layer_f(h_text)  # [B, H, L']
         h_text_d = self.cond_layer_d(h_text)  # [B, H, L']
         mel_mask_conv = self.get_conv_mask(
-            mel_len, h_text_f.shape[2], mel_mask)
+            mel_len, h_text_f.shape[2], mel_mask).unsqueeze(-1)
 
         # Sample from Prior
         z_n = torch.randn(h_text_f.shape[0], self.latent_hidden,
                           h_text_f.shape[2]).to(device=h_text_f.device, dtype=h_text.dtype)
-        z_q = self.flow(z_n, x_mask=~
-                        mel_mask_conv.unsqueeze(1), g=h_text_f, reverse=True)
+        z_q = self.flow(z_n, x_mask=mel_mask_conv.transpose(
+            1, 2), g=h_text_f, reverse=True)
 
         # Decoding
         x = self.latent_dec_prj(z_q.transpose(1, 2))
         x = x.contiguous().transpose(1, 2)
-        x = self.dec_wn(x, g=h_text_d).masked_fill(
-            mel_mask_conv.unsqueeze(1), 0.)
+        x = self.dec_wn(x, g=h_text_d) * mel_mask_conv.transpose(1, 2)
         x = x.contiguous().transpose(1, 2)
         mel_res = self.dec_conv(x)
-        mel_res = self.trim_output(mel_res, mel_mask.shape[1])
-        mel_res = mel_res.masked_fill(mel_mask.unsqueeze(-1), 0.)
-        residual = self.residual_layer(
-            mel_res).masked_fill(mel_mask.unsqueeze(-1), 0.)
+        mel_res = self.trim_output(mel_res, mel_mask.shape[1]) * mel_mask.unsqueeze(-1)
+        residual = self.residual_layer(mel_res) * mel_mask.unsqueeze(-1)
 
         return mel_res, residual, None
 
diff --git a/utils/tools.py b/utils/tools.py
index 7587b03..60f04a9 100644
--- a/utils/tools.py
+++ b/utils/tools.py
@@ -137,7 +137,7 @@ def get_mask_from_lengths(lengths, max_len=None):
         0).expand(batch_size, -1).to(lengths.device)
     mask = ids >= lengths.unsqueeze(1).expand(-1, max_len)
 
-    return mask
+    return ~mask
 
 
 def expand(values, durations):
@@ -168,7 +168,7 @@ def synth_one_sample(model, targets, predictions, vocoder, model_config, preproc
 
     # PostNet Inference on the reconstruction
     mel_reconst_pn = model.postnet.inference(
-        ~mel_mask.unsqueeze(1),
+        mel_mask.unsqueeze(1),
         g=(out_residual + residual).transpose(1, 2),
     )[0].float().detach()
 
@@ -179,7 +179,7 @@ def synth_one_sample(model, targets, predictions, vocoder, model_config, preproc
 
     # PostNet Inference on the inference
     mel_prediction_pn = model.postnet.inference(
-        ~mel_mask.unsqueeze(1),
+        mel_mask.unsqueeze(1),
         g=(out_residual + residual).transpose(1, 2),
     )[0].float().detach()
 
@@ -230,7 +230,7 @@ def synth_samples(targets, predictions, vocoder, model_config, preprocess_config
         basename = basenames[i]
         src_len = predictions[6][i].item()
         mel_len = predictions[7][i].item()
-        mel_prediction = predictions[0][i, :mel_len].detach()
+        mel_prediction = predictions[0][i, :, :mel_len].detach()
 
         fig = plot_mel(
             [