Update modeling_vits.py

modeling_vits.py

@@ -1,183 +1,186 @@
-from typing import Any, Optional, Union
-import numpy as np
-import torch
-from torch import nn
-from transformers import VitsPreTrainedModel , VitsConfig
-from transformers.models.vits.modeling_vits import VitsTextEncoder , VitsResidualCouplingBlock , VitsHifiGan , VitsStochasticDurationPredictor , VitsDurationPredictor , VitsPosteriorEncoder , VitsModelOutput
-from transformers.utils import auto_docstring
-from torch.nn.utils.parametrizations import weight_norm
-@auto_docstring(
-    custom_intro="""
-    The complete VITS model, for text-to-speech synthesis.
-    """
-)
-class VitsModel(VitsPreTrainedModel):
-    def __init__(self, config: VitsConfig):
-        super().__init__(config)
-        self.config = config
-        self.text_encoder = VitsTextEncoder(config)
-        self.flow = VitsResidualCouplingBlock(config)
-        self.decoder = VitsHifiGan(config)
-
-        if config.use_stochastic_duration_prediction:
-            self.duration_predictor = VitsStochasticDurationPredictor(config)
-        else:
-            self.duration_predictor = VitsDurationPredictor(config)
-
-        if config.num_speakers > 1:
-            self.embed_speaker = nn.Embedding(config.num_speakers, config.speaker_embedding_size)
-        
-        if config.num_emotions > 1:
-            self.embed_emotion = nn.Embedding(config.num_emotions, config.emotion_embedding_size)
-
-        # This is used only for training.
-        self.posterior_encoder = VitsPosteriorEncoder(config)
-
-        # These parameters control the synthesised speech properties
-        self.speaking_rate = config.speaking_rate
-        self.noise_scale = config.noise_scale
-        self.noise_scale_duration = config.noise_scale_duration
-
-        # Weight Norm Apply
-        for block in self.decoder.resblocks:
-            block.convs1 = nn.ModuleList([weight_norm(layer) for layer in block.convs1])
-            block.convs2 = nn.ModuleList([weight_norm(layer) for layer in block.convs2])
-        self.decoder.upsampler = nn.ModuleList([weight_norm(layer) for layer in self.decoder.upsampler])
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    @auto_docstring
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        speaker_id: Optional[int] = None,
-        emotion_id: Optional[int] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        labels: Optional[torch.FloatTensor] = None,
-        **kwargs,
-    ) -> Union[tuple[Any], VitsModelOutput]:
-        r"""
-        speaker_id (`int`, *optional*):
-            Which speaker embedding to use. Only used for multispeaker models.
-        labels (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`, *optional*):
-            Float values of target spectrogram. Timesteps set to `-100.0` are ignored (masked) for the loss
-            computation.
-
-        Example:
-
-        ```python
-        >>> from transformers import VitsTokenizer, VitsModel, set_seed
-        >>> import torch
-
-        >>> tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-eng")
-        >>> model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-
-        >>> inputs = tokenizer(text="Hello - my dog is cute", return_tensors="pt")
-
-        >>> set_seed(555)  # make deterministic
-
-        >>> with torch.no_grad():
-        ...     outputs = model(inputs["input_ids"])
-        >>> outputs.waveform.shape
-        torch.Size([1, 45824])
-        ```
-        """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if labels is not None:
-            raise NotImplementedError("Training of VITS is not supported yet.")
-
-        mask_dtype = self.text_encoder.embed_tokens.weight.dtype
-        if attention_mask is not None:
-            input_padding_mask = attention_mask.unsqueeze(-1).to(mask_dtype)
-        else:
-            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).to(mask_dtype)
-
-        if self.config.num_speakers > 1 and speaker_id is not None:
-            if not 0 <= speaker_id < self.config.num_speakers:
-                raise ValueError(f"Set `speaker_id` in the range 0-{self.config.num_speakers - 1}. or Set `emotion_id` in the range 0-{self.config.num_emotions - 1}.")
-            if isinstance(speaker_id, int) and isinstance(emotion_id, int):
-                speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device)
-                emotion_id = torch.full(size=(1,), fill_value=emotion_id, device=self.device)
-            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1) + self.embed_emotion(emotion_id).unsqueeze(-1)
-        else:
-            speaker_embeddings = None
-
-        text_encoder_output = self.text_encoder(
-            input_ids=input_ids,
-            padding_mask=input_padding_mask,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
-        hidden_states = hidden_states.transpose(1, 2)
-        input_padding_mask = input_padding_mask.transpose(1, 2)
-        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
-        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
-
-        if self.config.use_stochastic_duration_prediction:
-            log_duration = self.duration_predictor(
-                hidden_states,
-                input_padding_mask,
-                speaker_embeddings,
-                reverse=True,
-                noise_scale=self.noise_scale_duration,
-            )
-        else:
-            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
-
-        length_scale = 1.0 / self.speaking_rate
-        duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
-        predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()
-
-        # Create a padding mask for the output lengths of shape (batch, 1, max_output_length)
-        indices = torch.arange(predicted_lengths.max(), dtype=predicted_lengths.dtype, device=predicted_lengths.device)
-        output_padding_mask = indices.unsqueeze(0) < predicted_lengths.unsqueeze(1)
-        output_padding_mask = output_padding_mask.unsqueeze(1).to(input_padding_mask.dtype)
-
-        # Reconstruct an attention tensor of shape (batch, 1, out_length, in_length)
-        attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(output_padding_mask, -1)
-        batch_size, _, output_length, input_length = attn_mask.shape
-        cum_duration = torch.cumsum(duration, -1).view(batch_size * input_length, 1)
-        indices = torch.arange(output_length, dtype=duration.dtype, device=duration.device)
-        valid_indices = indices.unsqueeze(0) < cum_duration
-        valid_indices = valid_indices.to(attn_mask.dtype).view(batch_size, input_length, output_length)
-        padded_indices = valid_indices - nn.functional.pad(valid_indices, [0, 0, 1, 0, 0, 0])[:, :-1]
-        attn = padded_indices.unsqueeze(1).transpose(2, 3) * attn_mask
-
-        # Expand prior distribution
-        prior_means = torch.matmul(attn.squeeze(1), prior_means).transpose(1, 2)
-        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances).transpose(1, 2)
-
-        prior_latents = prior_means + torch.randn_like(prior_means) * torch.exp(prior_log_variances) * self.noise_scale
-        latents = self.flow(prior_latents, output_padding_mask, speaker_embeddings, reverse=True)
-
-        spectrogram = latents * output_padding_mask
-        waveform = self.decoder(spectrogram, speaker_embeddings)
-        waveform = waveform.squeeze(1)
-        sequence_lengths = predicted_lengths * np.prod(self.config.upsample_rates)
-
-        if not return_dict:
-            outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
-            return outputs
-
-        return VitsModelOutput(
-            waveform=waveform,
-            sequence_lengths=sequence_lengths,
-            spectrogram=spectrogram,
-            hidden_states=text_encoder_output.hidden_states,
-            attentions=text_encoder_output.attentions,
-        )
-
-
+from typing import Any, Optional, Union
+import numpy as np
+import torch
+from torch import nn
+from transformers import VitsPreTrainedModel , VitsConfig
+from transformers.models.vits.modeling_vits import VitsTextEncoder , VitsResidualCouplingBlock , VitsHifiGan , VitsStochasticDurationPredictor , VitsDurationPredictor , VitsPosteriorEncoder , VitsModelOutput
+from transformers.utils import auto_docstring
+from torch.nn.utils.parametrizations import weight_norm
+@auto_docstring(
+    custom_intro="""
+    The complete VITS model, for text-to-speech synthesis.
+    """
+)
+class VitsModel(VitsPreTrainedModel):
+    def __init__(self, config: VitsConfig):
+        super().__init__(config)
+        self.config = config
+        self.text_encoder = VitsTextEncoder(config)
+        self.flow = VitsResidualCouplingBlock(config)
+        self.decoder = VitsHifiGan(config)
+
+        if config.use_stochastic_duration_prediction:
+            self.duration_predictor = VitsStochasticDurationPredictor(config)
+        else:
+            self.duration_predictor = VitsDurationPredictor(config)
+
+        if config.num_speakers > 1:
+            self.embed_speaker = nn.Embedding(config.num_speakers, config.speaker_embedding_size)
+        
+        if config.num_emotions > 1:
+            self.embed_emotion = nn.Embedding(config.num_emotions, config.emotion_embedding_size)
+
+        # This is used only for training.
+        self.posterior_encoder = VitsPosteriorEncoder(config)
+
+        # These parameters control the synthesised speech properties
+        self.speaking_rate = config.speaking_rate
+        self.noise_scale = config.noise_scale
+        self.noise_scale_duration = config.noise_scale_duration
+
+        # Weight Norm Apply
+        for block in self.decoder.resblocks:
+            block.convs1 = nn.ModuleList([weight_norm(layer) for layer in block.convs1])
+            block.convs2 = nn.ModuleList([weight_norm(layer) for layer in block.convs2])
+        self.decoder.upsampler = nn.ModuleList([weight_norm(layer) for layer in self.decoder.upsampler])
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        speaker_id: Optional[int] = None,
+        emotion_id: Optional[int] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[tuple[Any], VitsModelOutput]:
+        r"""
+        speaker_id (`int`, *optional*):
+            Which speaker embedding to use. Only used for multispeaker models.
+        emotion_id (`int`, *optional*):
+            Which emotion embedding to use. Only used for multiemotion models.
+
+        labels (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`, *optional*):
+            Float values of target spectrogram. Timesteps set to `-100.0` are ignored (masked) for the loss
+            computation.
+
+        Example:
+
+        ```python
+        >>> from transformers import VitsTokenizer, VitsModel, set_seed
+        >>> import torch
+
+        >>> tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-eng")
+        >>> model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+
+        >>> inputs = tokenizer(text="Hello - my dog is cute", return_tensors="pt")
+
+        >>> set_seed(555)  # make deterministic
+
+        >>> with torch.no_grad():
+        ...     outputs = model(inputs["input_ids"])
+        >>> outputs.waveform.shape
+        torch.Size([1, 45824])
+        ```
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if labels is not None:
+            raise NotImplementedError("Training of VITS is not supported yet.")
+
+        mask_dtype = self.text_encoder.embed_tokens.weight.dtype
+        if attention_mask is not None:
+            input_padding_mask = attention_mask.unsqueeze(-1).to(mask_dtype)
+        else:
+            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).to(mask_dtype)
+
+        if self.config.num_speakers > 1 and speaker_id is not None:
+            if not 0 <= speaker_id < self.config.num_speakers:
+                raise ValueError(f"Set `speaker_id` in the range 0-{self.config.num_speakers - 1}. or Set `emotion_id` in the range 0-{self.config.num_emotions - 1}.")
+            if isinstance(speaker_id, int) and isinstance(emotion_id, int):
+                speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device)
+                emotion_id = torch.full(size=(1,), fill_value=emotion_id, device=self.device)
+            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1) + self.embed_emotion(emotion_id).unsqueeze(-1)
+        else:
+            speaker_embeddings = None
+
+        text_encoder_output = self.text_encoder(
+            input_ids=input_ids,
+            padding_mask=input_padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
+        hidden_states = hidden_states.transpose(1, 2)
+        input_padding_mask = input_padding_mask.transpose(1, 2)
+        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
+        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
+
+        if self.config.use_stochastic_duration_prediction:
+            log_duration = self.duration_predictor(
+                hidden_states,
+                input_padding_mask,
+                speaker_embeddings,
+                reverse=True,
+                noise_scale=self.noise_scale_duration,
+            )
+        else:
+            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
+
+        length_scale = 1.0 / self.speaking_rate
+        duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
+        predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()
+
+        # Create a padding mask for the output lengths of shape (batch, 1, max_output_length)
+        indices = torch.arange(predicted_lengths.max(), dtype=predicted_lengths.dtype, device=predicted_lengths.device)
+        output_padding_mask = indices.unsqueeze(0) < predicted_lengths.unsqueeze(1)
+        output_padding_mask = output_padding_mask.unsqueeze(1).to(input_padding_mask.dtype)
+
+        # Reconstruct an attention tensor of shape (batch, 1, out_length, in_length)
+        attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(output_padding_mask, -1)
+        batch_size, _, output_length, input_length = attn_mask.shape
+        cum_duration = torch.cumsum(duration, -1).view(batch_size * input_length, 1)
+        indices = torch.arange(output_length, dtype=duration.dtype, device=duration.device)
+        valid_indices = indices.unsqueeze(0) < cum_duration
+        valid_indices = valid_indices.to(attn_mask.dtype).view(batch_size, input_length, output_length)
+        padded_indices = valid_indices - nn.functional.pad(valid_indices, [0, 0, 1, 0, 0, 0])[:, :-1]
+        attn = padded_indices.unsqueeze(1).transpose(2, 3) * attn_mask
+
+        # Expand prior distribution
+        prior_means = torch.matmul(attn.squeeze(1), prior_means).transpose(1, 2)
+        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances).transpose(1, 2)
+
+        prior_latents = prior_means + torch.randn_like(prior_means) * torch.exp(prior_log_variances) * self.noise_scale
+        latents = self.flow(prior_latents, output_padding_mask, speaker_embeddings, reverse=True)
+
+        spectrogram = latents * output_padding_mask
+        waveform = self.decoder(spectrogram, speaker_embeddings)
+        waveform = waveform.squeeze(1)
+        sequence_lengths = predicted_lengths * np.prod(self.config.upsample_rates)
+
+        if not return_dict:
+            outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
+            return outputs
+
+        return VitsModelOutput(
+            waveform=waveform,
+            sequence_lengths=sequence_lengths,
+            spectrogram=spectrogram,
+            hidden_states=text_encoder_output.hidden_states,
+            attentions=text_encoder_output.attentions,
+        )
+
+
 __all__ = ["VitsModel"]