From 076d0cb2583393646abe6de384a14ebf05d84d31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eren=20G=C3=B6lge?= <egolge@coqui.ai>
Date: Fri, 3 Sep 2021 13:28:26 +0000
Subject: [PATCH] Add tests for certain FastPitch functions

---
 tests/tts_tests/test_fast_pitch.py | 51 ++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 tests/tts_tests/test_fast_pitch.py

diff --git a/tests/tts_tests/test_fast_pitch.py b/tests/tts_tests/test_fast_pitch.py
new file mode 100644
index 00000000..ba6b0ce6
--- /dev/null
+++ b/tests/tts_tests/test_fast_pitch.py
@@ -0,0 +1,51 @@
+import unittest
+
+import torch as T
+
+from TTS.tts.layers.losses import L1LossMasked, SSIMLoss
+from TTS.tts.layers.tacotron.tacotron import CBHG, Decoder, Encoder, Prenet
+from TTS.tts.models.fast_pitch import FastPitch, FastPitchArgs, average_pitch
+from TTS.tts.utils.data import sequence_mask
+
+# pylint: disable=unused-variable
+
+
+class AveragePitchTests(unittest.TestCase):
+    def test_in_out(self):  # pylint: disable=no-self-use
+        pitch = T.rand(1, 1, 128)
+
+        durations = T.randint(1, 5, (1, 21))
+        coeff = 128.0 / durations.sum()
+        durations = T.round(durations * coeff)
+        diff = 128.0 - durations.sum()
+        durations[0, -1] += diff
+        durations = durations.long()
+
+        pitch_avg = average_pitch(pitch, durations)
+
+        index = 0
+        for idx, dur in enumerate(durations[0]):
+            assert abs(pitch_avg[0, 0, idx] - pitch[0, 0, index : index + dur.item()].mean()) < 1e-5
+            index += dur
+
+
+def expand_encoder_outputs_test():
+    model = FastPitch(FastPitchArgs(num_chars=10))
+
+    inputs = T.rand(2, 5, 57)
+    durations = T.randint(1, 4, (2, 57))
+
+    x_mask = T.ones(2, 1, 57)
+    y_mask = T.ones(2, 1, durations.sum(1).max())
+
+    expanded, attn = model.expand_encoder_outputs(inputs, durations, x_mask, y_mask)
+
+    for b in range(durations.shape[0]):
+        index = 0
+        for idx, dur in enumerate(durations[b]):
+            diff = (
+                expanded[b, :, index : index + dur.item()]
+                - inputs[b, :, idx].repeat(dur.item()).view(expanded[b, :, index : index + dur.item()].shape)
+            ).sum()
+            assert abs(diff) < 1e-6, diff
+            index += dur