Create models.py

models.py

@@ -0,0 +1,278 @@
+from loss import *
+
+class AlignmentEncoder(torch.nn.Module):
+    """
+    Module for alignment text and mel spectrogram.
+
+    Args:
+        n_mel_channels: Dimension of mel spectrogram.
+        n_text_channels: Dimension of text embeddings.
+        n_att_channels: Dimension of model
+        temperature: Temperature to scale distance by.
+            Suggested to be 0.0005 when using dist_type "l2" and 15.0 when using "cosine".
+        condition_types: List of types for nemo.collections.tts.modules.submodules.ConditionalInput.
+        dist_type: Distance type to use for similarity measurement. Supports "l2" and "cosine" distance.
+    """
+
+    def __init__(
+        self,
+        n_mel_channels=128,
+        n_text_channels=512,
+        n_att_channels=128,
+        temperature=0.0005,
+        condition_types=[],
+        dist_type="l2",
+    ):
+        super().__init__()
+        self.temperature = temperature
+        # self.cond_input = ConditionalInput(n_text_channels, n_text_channels, condition_types)
+        self.softmax = torch.nn.Softmax(dim=3)
+        self.log_softmax = torch.nn.LogSoftmax(dim=3)
+
+        self.key_proj = nn.Sequential(
+            ConvNorm(n_text_channels, n_text_channels * 2, kernel_size=3, bias=True, w_init_gain='relu'),
+            torch.nn.ReLU(),
+            ConvNorm(n_text_channels * 2, n_att_channels, kernel_size=1, bias=True),
+        )
+
+        self.query_proj = nn.Sequential(
+            ConvNorm(n_mel_channels, n_mel_channels * 2, kernel_size=3, bias=True, w_init_gain='relu'),
+            torch.nn.ReLU(),
+            ConvNorm(n_mel_channels * 2, n_mel_channels, kernel_size=1, bias=True),
+            torch.nn.ReLU(),
+            ConvNorm(n_mel_channels, n_att_channels, kernel_size=1, bias=True),
+        )
+        if dist_type == "l2":
+            self.dist_fn = self.get_euclidean_dist
+        elif dist_type == "cosine":
+            self.dist_fn = self.get_cosine_dist
+        else:
+            raise ValueError(f"Unknown distance type '{dist_type}'")
+
+    @staticmethod
+    def _apply_mask(inputs, mask, mask_value):
+        if mask is None:
+            return
+
+        mask = rearrange(mask, "B T2 1 -> B 1 1 T2")
+        inputs.data.masked_fill_(mask, mask_value)
+
+    def get_dist(self, keys, queries, mask=None):
+        """Calculation of distance matrix.
+
+        Args:
+            queries (torch.tensor): B x C1 x T1 tensor (probably going to be mel data).
+            keys (torch.tensor): B x C2 x T2 tensor (text data).
+            mask (torch.tensor): B x T2 x 1 tensor, binary mask for variable length entries and also can be used
+                for ignoring unnecessary elements from keys in the resulting distance matrix (True = mask element, False = leave unchanged).
+        Output:
+            dist (torch.tensor): B x T1 x T2 tensor.
+        """
+        # B x C x T1
+        queries_enc = self.query_proj(queries)
+        # B x C x T2
+        keys_enc = self.key_proj(keys)
+        # B x 1 x T1 x T2
+        dist = self.dist_fn(queries_enc=queries_enc, keys_enc=keys_enc)
+
+        self._apply_mask(dist, mask, float("inf"))
+
+        return dist.squeeze(1)
+
+    @staticmethod
+    def get_euclidean_dist(queries_enc, keys_enc):
+        queries_enc = rearrange(queries_enc, "B C T1 -> B C T1 1")
+        keys_enc = rearrange(keys_enc, "B C T2 -> B C 1 T2")
+        # B x C x T1 x T2
+        distance = (queries_enc - keys_enc) ** 2
+        # B x 1 x T1 x T2
+        l2_dist = distance.sum(axis=1, keepdim=True)
+        return l2_dist
+
+    @staticmethod
+    def get_cosine_dist(queries_enc, keys_enc):
+        queries_enc = rearrange(queries_enc, "B C T1 -> B C T1 1")
+        keys_enc = rearrange(keys_enc, "B C T2 -> B C 1 T2")
+        cosine_dist = -torch.nn.functional.cosine_similarity(queries_enc, keys_enc, dim=1)
+        cosine_dist = rearrange(cosine_dist, "B T1 T2 -> B 1 T1 T2")
+        return cosine_dist
+
+    @staticmethod
+    def get_durations(attn_soft, text_len, spect_len):
+        """Calculation of durations.
+
+        Args:
+            attn_soft (torch.tensor): B x 1 x T1 x T2 tensor.
+            text_len (torch.tensor): B tensor, lengths of text.
+            spect_len (torch.tensor): B tensor, lengths of mel spectrogram.
+        """
+        attn_hard = binarize_attention_parallel(attn_soft, text_len, spect_len)
+        durations = attn_hard.sum(2)[:, 0, :]
+        assert torch.all(torch.eq(durations.sum(dim=1), spect_len))
+        return durations
+
+    @staticmethod
+    def get_mean_dist_by_durations(dist, durations, mask=None):
+        """Select elements from the distance matrix for the given durations and mask and return mean distance.
+
+        Args:
+            dist (torch.tensor): B x T1 x T2 tensor.
+            durations (torch.tensor): B x T2 tensor. Dim T2 should sum to T1.
+            mask (torch.tensor): B x T2 x 1 binary mask for variable length entries and also can be used
+                for ignoring unnecessary elements in dist by T2 dim (True = mask element, False = leave unchanged).
+        Output:
+            mean_dist (torch.tensor): B x 1 tensor.
+        """
+        batch_size, t1_size, t2_size = dist.size()
+        assert torch.all(torch.eq(durations.sum(dim=1), t1_size))
+
+        AlignmentEncoder._apply_mask(dist, mask, 0)
+
+        # TODO(oktai15): make it more efficient
+        mean_dist_by_durations = []
+        for dist_idx in range(batch_size):
+            mean_dist_by_durations.append(
+                torch.mean(
+                    dist[
+                        dist_idx,
+                        torch.arange(t1_size),
+                        torch.repeat_interleave(torch.arange(t2_size), repeats=durations[dist_idx]),
+                    ]
+                )
+            )
+
+        return torch.tensor(mean_dist_by_durations, dtype=dist.dtype, device=dist.device)
+
+    @staticmethod
+    def get_mean_distance_for_word(l2_dists, durs, start_token, num_tokens):
+        """Calculates the mean distance between text and audio embeddings given a range of text tokens.
+
+        Args:
+            l2_dists (torch.tensor): L2 distance matrix from Aligner inference. T1 x T2 tensor.
+            durs (torch.tensor): List of durations corresponding to each text token. T2 tensor. Should sum to T1.
+            start_token (int): Index of the starting token for the word of interest.
+            num_tokens (int): Length (in tokens) of the word of interest.
+        Output:
+            mean_dist_for_word (float): Mean embedding distance between the word indicated and its predicted audio frames.
+        """
+        # Need to calculate which audio frame we start on by summing all durations up to the start token's duration
+        start_frame = torch.sum(durs[:start_token]).data
+
+        total_frames = 0
+        dist_sum = 0
+
+        # Loop through each text token
+        for token_ind in range(start_token, start_token + num_tokens):
+            # Loop through each frame for the given text token
+            for frame_ind in range(start_frame, start_frame + durs[token_ind]):
+                # Recall that the L2 distance matrix is shape [spec_len, text_len]
+                dist_sum += l2_dists[frame_ind, token_ind]
+
+            # Update total frames so far & the starting frame for the next token
+            total_frames += durs[token_ind]
+            start_frame += durs[token_ind]
+
+        return dist_sum / total_frames
+
+    def forward(self, queries, keys, mask=None, attn_prior=None, conditioning=None):
+        """Forward pass of the aligner encoder.
+
+        Args:
+            queries (torch.tensor): B x C1 x T1 tensor (probably going to be mel data).
+            keys (torch.tensor): B x C2 x T2 tensor (text data).
+            mask (torch.tensor): B x T2 x 1 tensor, binary mask for variable length entries (True = mask element, False = leave unchanged).
+            attn_prior (torch.tensor): prior for attention matrix.
+            conditioning (torch.tensor): B x 1 x C2 conditioning embedding
+        Output:
+            attn (torch.tensor): B x 1 x T1 x T2 attention mask. Final dim T2 should sum to 1.
+            attn_logprob (torch.tensor): B x 1 x T1 x T2 log-prob attention mask.
+        """
+        # keys = self.cond_input(keys.transpose(1, 2), conditioning).transpose(1, 2)
+        # B x C x T1
+        queries_enc = self.query_proj(queries)
+        # B x C x T2
+        keys_enc = self.key_proj(keys)
+        # B x 1 x T1 x T2
+        distance = self.dist_fn(queries_enc=queries_enc, keys_enc=keys_enc)
+        attn = -self.temperature * distance
+
+        if attn_prior is not None:
+            attn = self.log_softmax(attn) + torch.log(attn_prior[:, None] + 1e-8)
+
+        attn_logprob = attn.clone()
+
+        self._apply_mask(attn, mask, -float("inf"))
+
+        attn = self.softmax(attn)  # softmax along T2
+        return attn, attn_logprob
+
+
+
+def get_mask_from_lengths(
+    lengths: Optional[torch.Tensor] = None,
+    x: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Constructs binary mask from a 1D torch tensor of input lengths
+
+    Args:
+        lengths: Optional[torch.tensor] (torch.tensor): 1D tensor with lengths
+        x: Optional[torch.tensor] = tensor to be used on, last dimension is for mask
+    Returns:
+        mask (torch.tensor): num_sequences x max_length binary tensor
+    """
+    if lengths is None:
+        assert x is not None
+        return torch.ones(x.shape[-1], dtype=torch.bool, device=x.device)
+    else:
+        if x is None:
+            max_len = torch.max(lengths)
+        else:
+            max_len = x.shape[-1]
+    ids = torch.arange(0, max_len, device=lengths.device, dtype=lengths.dtype)
+    mask = ids < lengths.unsqueeze(1)
+    return mask
+
+class AlignerModel(torch.nn.Module):
+    """Speech-to-text alignment model (https://arxiv.org/pdf/2108.10447.pdf) that is used to learn alignments between mel spectrogram and text."""
+
+    def __init__(self):
+
+        # num_tokens = len(self.tokenizer.tokens)
+        # self.tokenizer_pad = self.tokenizer.pad
+        # self.tokenizer_unk = self.tokenizer.oov
+
+        super().__init__()
+
+        self.embed = nn.Embedding(214, 512)
+        self.alignment_encoder = AlignmentEncoder()
+
+
+        # self.bin_loss = BinLoss()
+        # self.add_bin_loss = False
+        # self.bin_loss_scale = 0.0
+        # self.bin_loss_start_ratio = cfg.bin_loss_start_ratio
+        # self.bin_loss_warmup_epochs = cfg.bin_loss_warmup_epochs
+
+    def forward(self, *, spec, spec_len, text, text_len, attn_prior=None):
+        # with torch.amp.autocast(self.device.type, enabled=False):
+          attn_soft, attn_logprob = self.alignment_encoder(
+              queries=spec,
+              keys=self.embed(text).transpose(1, 2),
+              mask=get_mask_from_lengths(text_len).unsqueeze(-1) == 0,
+              attn_prior=attn_prior,
+          )
+
+          return attn_soft, attn_logprob
+
+
+# mod = AlignerModel()
+
+# attn_soft, attn_logprob = mod(spec=mel_input, 
+#                               spec_len=mel_input_length, 
+#                               text=text_input, 
+#                               text_len=text_input_length,
+#                               attn_prior = attn_prior)
+
+
+# attn_soft.shape
+# text_input, text_input_length, mel_input, mel_input_length, attn_prior