returns y_mask in VITS inference (#1540)

* returns y_mask * make style
2022-04-27 10:42:26 +02:00 · 2022-04-27 10:42:26 +02:00 · 379ccd7ba6
parent a41e860a66
commit 379ccd7ba6
1 changed files with 9 additions and 1 deletions
--- a/TTS/tts/models/vits.py
+++ b/TTS/tts/models/vits.py
@ -1050,7 +1050,15 @@ class Vits(BaseTTS):

        o = self.waveform_decoder((z * y_mask)[:, :, : self.max_inference_len], g=g)

-        outputs = {"model_outputs": o, "alignments": attn.squeeze(1), "z": z, "z_p": z_p, "m_p": m_p, "logs_p": logs_p}
+        outputs = {
+            "model_outputs": o,
+            "alignments": attn.squeeze(1),
+            "z": z,
+            "z_p": z_p,
+            "m_p": m_p,
+            "logs_p": logs_p,
+            "y_mask": y_mask,
+        }
        return outputs

    @torch.no_grad()