Update modeling_vits.py

modeling_vits.py

@@ -26,5 +26,9 @@ class ModVitsModel(VitsModel):
     
     @torch.inference_mode()
     def forward(self, input_ids = None, attention_mask = None, speaker_id = None, output_attentions = None, output_hidden_states = None, return_dict = None, labels = None, **kwargs):
-        speaker_id = self.config.speaker_names.index(speaker_id) * len(self.config.emotion_names) + self.config.emotion_names.index(kwargs['emotion_id'])
-        return super().forward(input_ids, attention_mask, speaker_id, output_attentions, output_hidden_states, return_dict, labels, **kwargs)
+        # speaker_id = self.config.speaker_names.index(speaker_id) * len(self.config.emotion_names) + self.config.emotion_names.index(kwargs['emotion_id'])
+        audio = super().forward(input_ids, attention_mask, speaker_id, output_attentions, output_hidden_states, return_dict, labels, **kwargs)
+        B, T = audio.waveform.shape
+        mask = torch.arange(T, device=audio.waveform.device).expand(B, T) < audio.sequence_lengths.unsqueeze(1)
+        audio.waveform.masked_fill_(~mask, 0)
+        return audio