chore: import upstream snapshot with attribution

modules/losses/__init__.py

@@ -0,0 +1,3 @@
+from .diff_loss import DiffusionLoss
+from .reflow_loss import RectifiedFlowLoss
+from .dur_loss import DurationLoss

modules/losses/diff_loss.py

@@ -0,0 +1,34 @@
+import torch.nn as nn
+from torch import Tensor
+
+
+class DiffusionLoss(nn.Module):
+    def __init__(self, loss_type):
+        super().__init__()
+        self.loss_type = loss_type
+        if self.loss_type == 'l1':
+            self.loss = nn.L1Loss(reduction='none')
+        elif self.loss_type == 'l2':
+            self.loss = nn.MSELoss(reduction='none')
+        else:
+            raise NotImplementedError()
+
+    @staticmethod
+    def _mask_non_padding(x_recon, noise, non_padding=None):
+        if non_padding is not None:
+            non_padding = non_padding.transpose(1, 2).unsqueeze(1)
+            return x_recon * non_padding, noise * non_padding
+        else:
+            return x_recon, noise
+
+    def _forward(self, x_recon, noise):
+        return self.loss(x_recon, noise)
+
+    def forward(self, x_recon: Tensor, noise: Tensor, non_padding: Tensor = None) -> Tensor:
+        """
+        :param x_recon: [B, 1, M, T]
+        :param noise: [B, 1, M, T]
+        :param non_padding: [B, T, M]
+        """
+        x_recon, noise = self._mask_non_padding(x_recon, noise, non_padding)
+        return self._forward(x_recon, noise).mean()

modules/losses/dur_loss.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class DurationLoss(nn.Module):
+    """
+    Loss module as combination of phone duration loss, word duration loss and sentence duration loss.
+    """
+
+    def __init__(self, offset, loss_type,
+                 lambda_pdur=0.6, lambda_wdur=0.3, lambda_sdur=0.1):
+        super().__init__()
+        self.loss_type = loss_type
+        if self.loss_type == 'mse':
+            self.loss = nn.MSELoss()
+        elif self.loss_type == 'huber':
+            self.loss = nn.HuberLoss()
+        else:
+            raise NotImplementedError()
+        self.offset = offset
+
+        self.lambda_pdur = lambda_pdur
+        self.lambda_wdur = lambda_wdur
+        self.lambda_sdur = lambda_sdur
+
+    def linear2log(self, any_dur):
+        return torch.log(any_dur + self.offset)
+
+    def forward(self, dur_pred: Tensor, dur_gt: Tensor, ph2word: Tensor) -> Tensor:
+        dur_gt = dur_gt.to(dtype=dur_pred.dtype)
+
+        # pdur_loss
+        pdur_loss = self.lambda_pdur * self.loss(self.linear2log(dur_pred), self.linear2log(dur_gt))
+
+        dur_pred = dur_pred.clamp(min=0.)  # clip to avoid NaN loss
+
+        # wdur loss
+        shape = dur_pred.shape[0], ph2word.max() + 1
+        wdur_pred = dur_pred.new_zeros(*shape).scatter_add(
+            1, ph2word, dur_pred
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        wdur_gt = dur_gt.new_zeros(*shape).scatter_add(
+            1, ph2word, dur_gt
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        wdur_loss = self.lambda_wdur * self.loss(self.linear2log(wdur_pred), self.linear2log(wdur_gt))
+
+        # sdur loss
+        sdur_pred = dur_pred.sum(dim=1)
+        sdur_gt = dur_gt.sum(dim=1)
+        sdur_loss = self.lambda_sdur * self.loss(self.linear2log(sdur_pred), self.linear2log(sdur_gt))
+
+        # combine
+        dur_loss = pdur_loss + wdur_loss + sdur_loss
+
+        return dur_loss

modules/losses/reflow_loss.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class RectifiedFlowLoss(nn.Module):
+    def __init__(self, loss_type, log_norm=True):
+        super().__init__()
+        self.loss_type = loss_type
+        self.log_norm = log_norm
+        if self.loss_type == 'l1':
+            self.loss = nn.L1Loss(reduction='none')
+        elif self.loss_type == 'l2':
+            self.loss = nn.MSELoss(reduction='none')
+        else:
+            raise NotImplementedError()
+
+    @staticmethod
+    def _mask_non_padding(v_pred, v_gt, non_padding=None):
+        if non_padding is not None:
+            non_padding = non_padding.transpose(1, 2).unsqueeze(1)
+            return v_pred * non_padding, v_gt * non_padding
+        else:
+            return v_pred, v_gt
+
+    @staticmethod
+    def get_weights(t):
+        eps = 1e-7
+        t = t.float()
+        t = torch.clip(t, 0 + eps, 1 - eps)
+        weights = 0.398942 / t / (1 - t) * torch.exp(
+            -0.5 * torch.log(t / (1 - t)) ** 2
+        ) + eps
+        return weights[:, None, None, None]
+
+    def _forward(self, v_pred, v_gt, t=None):
+        if self.log_norm:
+            return self.get_weights(t) * self.loss(v_pred, v_gt)
+        else:
+            return self.loss(v_pred, v_gt)
+
+    def forward(self, v_pred: Tensor, v_gt: Tensor, t: Tensor, non_padding: Tensor = None) -> Tensor:
+        """
+        :param v_pred: [B, 1, M, T]
+        :param v_gt: [B, 1, M, T]
+        :param t: [B,]
+        :param non_padding: [B, T, M]
+        """
+        v_pred, v_gt = self._mask_non_padding(v_pred, v_gt, non_padding)
+        return self._forward(v_pred, v_gt, t=t).mean()