Create modeling_vits.py

modeling_vits.py

@@ -0,0 +1,1508 @@
+"""PyTorch VITS model."""
+
+import math
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.integrations.fsdp import is_fsdp_managed_module
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    ModelOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from .configuration_vits import IndicVitsConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# General docstring
+_CONFIG_FOR_DOC = "IndicVitsConfig"
+
+
+@dataclass
+class VitsModelOutput(ModelOutput):
+    """
+    Describes the outputs for the VITS model, with potential hidden states and attentions.
+
+    Args:
+        waveform (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+            The final audio waveform predicted by the model.
+        sequence_lengths  (`torch.FloatTensor` of shape `(batch_size,)`):
+            The length in samples of each element in the `waveform` batch.
+        spectrogram (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_bins)`):
+            The log-mel spectrogram predicted at the output of the flow model. This spectrogram is passed to the Hi-Fi
+            GAN decoder model to obtain the final audio waveform.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attention weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    waveform: torch.FloatTensor = None
+    sequence_lengths: torch.FloatTensor = None
+    spectrogram: Optional[Tuple[torch.FloatTensor]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+class VitsTextEncoderOutput(ModelOutput):
+    """
+    Describes the outputs for the VITS text encoder model, with potential hidden states and attentions.
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        prior_means (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            The predicted mean values of the prior distribution for the latent text variables.
+        prior_log_variances (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            The predicted log-variance values of the prior distribution for the latent text variables.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attention weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    last_hidden_state: torch.FloatTensor = None
+    prior_means: torch.FloatTensor = None
+    prior_log_variances: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, num_channels):
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :num_channels, :])
+    s_act = torch.sigmoid(in_act[:, num_channels:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def _unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    reverse=False,
+    tail_bound=5.0,
+    min_bin_width=1e-3,
+    min_bin_height=1e-3,
+    min_derivative=1e-3,
+):
+    """
+    This transformation represents a monotonically increasing piecewise rational quadratic function. Outside of the
+    `tail_bound`, the transform behaves as an identity function.
+
+    Args:
+        inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Second half of the hidden-states input to the Vits convolutional flow module.
+        unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        reverse (`bool`, *optional*, defaults to `False`):
+            Whether the model is being run in reverse mode.
+        tail_bound (`float`, *optional* defaults to 5):
+            Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the
+            transform behaves as an identity function.
+        min_bin_width (`float`, *optional*, defaults to 1e-3):
+            Minimum bin value across the width dimension for the piecewise rational quadratic function.
+        min_bin_height (`float`, *optional*, defaults to 1e-3):
+            Minimum bin value across the height dimension for the piecewise rational quadratic function.
+        min_derivative (`float`, *optional*, defaults to 1e-3):
+            Minimum bin value across the derivatives for the piecewise rational quadratic function.
+    Returns:
+        outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Hidden-states as transformed by the piecewise rational quadratic function with the `tail_bound` limits
+            applied.
+        log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Logarithm of the absolute value of the determinants corresponding to the `outputs` with the `tail_bound`
+            limits applied.
+    """
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    log_abs_det = torch.zeros_like(inputs)
+    constant = np.log(np.exp(1 - min_derivative) - 1)
+
+    unnormalized_derivatives = nn.functional.pad(unnormalized_derivatives, pad=(1, 1))
+    unnormalized_derivatives[..., 0] = constant
+    unnormalized_derivatives[..., -1] = constant
+
+    outputs[outside_interval_mask] = inputs[outside_interval_mask]
+    log_abs_det[outside_interval_mask] = 0.0
+
+    outputs[inside_interval_mask], log_abs_det[inside_interval_mask] = _rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        reverse=reverse,
+        tail_bound=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+    return outputs, log_abs_det
+
+
+def _rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    reverse,
+    tail_bound,
+    min_bin_width,
+    min_bin_height,
+    min_derivative,
+):
+    """
+    This transformation represents a monotonically increasing piecewise rational quadratic function. Unlike the
+    function `_unconstrained_rational_quadratic_spline`, the function behaves the same across the `tail_bound`.
+
+    Args:
+        inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Second half of the hidden-states input to the Vits convolutional flow module.
+        unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        reverse (`bool`):
+            Whether the model is being run in reverse mode.
+        tail_bound (`float`):
+            Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the
+            transform behaves as an identity function.
+        min_bin_width (`float`):
+            Minimum bin value across the width dimension for the piecewise rational quadratic function.
+        min_bin_height (`float`):
+            Minimum bin value across the height dimension for the piecewise rational quadratic function.
+        min_derivative (`float`):
+            Minimum bin value across the derivatives for the piecewise rational quadratic function.
+    Returns:
+        outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Hidden-states as transformed by the piecewise rational quadratic function.
+        log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Logarithm of the absolute value of the determinants corresponding to the `outputs`.
+    """
+    upper_bound = tail_bound
+    lower_bound = -tail_bound
+
+    if torch.min(inputs) < lower_bound or torch.max(inputs) > upper_bound:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError(f"Minimal bin width {min_bin_width} too large for the number of bins {num_bins}")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError(f"Minimal bin height {min_bin_height} too large for the number of bins {num_bins}")
+
+    widths = nn.functional.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = nn.functional.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (upper_bound - lower_bound) * cumwidths + lower_bound
+    cumwidths[..., 0] = lower_bound
+    cumwidths[..., -1] = upper_bound
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + nn.functional.softplus(unnormalized_derivatives)
+
+    heights = nn.functional.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = nn.functional.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (upper_bound - lower_bound) * cumheights + lower_bound
+    cumheights[..., 0] = lower_bound
+    cumheights[..., -1] = upper_bound
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    bin_locations = cumheights if reverse else cumwidths
+    bin_locations[..., -1] += 1e-6
+    bin_idx = torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+    bin_idx = bin_idx[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    intermediate1 = input_derivatives + input_derivatives_plus_one - 2 * input_delta
+    if not reverse:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
+        denominator = input_delta + intermediate1 * theta_one_minus_theta
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+        return outputs, log_abs_det
+    else:
+        # find the roots of a quadratic equation
+        intermediate2 = inputs - input_cumheights
+        intermediate3 = intermediate2 * intermediate1
+        a = input_heights * (input_delta - input_derivatives) + intermediate3
+        b = input_heights * input_derivatives - intermediate3
+        c = -input_delta * intermediate2
+
+        discriminant = b.pow(2) - 4 * a * c
+        if not (discriminant >= 0).all():
+            raise RuntimeError(f"invalid discriminant {discriminant}")
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + intermediate1 * theta_one_minus_theta
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+        return outputs, -log_abs_det
+
+
+class VitsWaveNet(torch.nn.Module):
+    def __init__(self, config: IndicVitsConfig, num_layers: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_layers = num_layers
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.dropout = nn.Dropout(config.wavenet_dropout)
+
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+        else:
+            weight_norm = nn.utils.weight_norm
+
+        if config.speaker_embedding_size != 0:
+            cond_layer = torch.nn.Conv1d(config.speaker_embedding_size, 2 * config.hidden_size * num_layers, 1)
+            self.cond_layer = weight_norm(cond_layer, name="weight")
+
+        for i in range(num_layers):
+            dilation = config.wavenet_dilation_rate**i
+            padding = (config.wavenet_kernel_size * dilation - dilation) // 2
+            in_layer = torch.nn.Conv1d(
+                in_channels=config.hidden_size,
+                out_channels=2 * config.hidden_size,
+                kernel_size=config.wavenet_kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < num_layers - 1:
+                res_skip_channels = 2 * config.hidden_size
+            else:
+                res_skip_channels = config.hidden_size
+
+            res_skip_layer = torch.nn.Conv1d(config.hidden_size, res_skip_channels, 1)
+            res_skip_layer = weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        outputs = torch.zeros_like(inputs)
+        num_channels_tensor = torch.IntTensor([self.hidden_size])
+
+        if global_conditioning is not None:
+            global_conditioning = self.cond_layer(global_conditioning)
+
+        for i in range(self.num_layers):
+            hidden_states = self.in_layers[i](inputs)
+
+            if global_conditioning is not None:
+                cond_offset = i * 2 * self.hidden_size
+                global_states = global_conditioning[:, cond_offset : cond_offset + 2 * self.hidden_size, :]
+            else:
+                global_states = torch.zeros_like(hidden_states)
+
+            acts = fused_add_tanh_sigmoid_multiply(hidden_states, global_states, num_channels_tensor[0])
+            acts = self.dropout(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.num_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_size, :]
+                inputs = (inputs + res_acts) * padding_mask
+                outputs = outputs + res_skip_acts[:, self.hidden_size :, :]
+            else:
+                outputs = outputs + res_skip_acts
+
+        return outputs * padding_mask
+
+    def remove_weight_norm(self):
+        if self.speaker_embedding_size != 0:
+            torch.nn.utils.parametrize.remove_parametrizations(self.cond_layer, "weight")
+        for layer in self.in_layers:
+            torch.nn.utils.parametrize.remove_parametrizations(layer, "weight")
+        for layer in self.res_skip_layers:
+            torch.nn.utils.parametrize.remove_parametrizations(layer, "weight")
+
+
+class VitsPosteriorEncoder(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.out_channels = config.flow_size
+
+        self.conv_pre = nn.Conv1d(config.spectrogram_bins, config.hidden_size, 1)
+        self.wavenet = VitsWaveNet(config, num_layers=config.posterior_encoder_num_wavenet_layers)
+        self.conv_proj = nn.Conv1d(config.hidden_size, self.out_channels * 2, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        inputs = self.conv_pre(inputs) * padding_mask
+        inputs = self.wavenet(inputs, padding_mask, global_conditioning)
+        stats = self.conv_proj(inputs) * padding_mask
+        mean, log_stddev = torch.split(stats, self.out_channels, dim=1)
+        sampled = (mean + torch.randn_like(mean) * torch.exp(log_stddev)) * padding_mask
+        return sampled, mean, log_stddev
+
+
+# Copied from transformers.models.speecht5.modeling_speecht5.HifiGanResidualBlock
+class HifiGanResidualBlock(nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), leaky_relu_slope=0.1):
+        super().__init__()
+        self.leaky_relu_slope = leaky_relu_slope
+
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+        else:
+            weight_norm = nn.utils.weight_norm
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=dilation[i],
+                        padding=self.get_padding(kernel_size, dilation[i]),
+                    )
+                )
+                for i in range(len(dilation))
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in range(len(dilation))
+            ]
+        )
+
+    def get_padding(self, kernel_size, dilation=1):
+        return (kernel_size * dilation - dilation) // 2
+
+    # def apply_weight_norm(self):
+    #     # Determine the correct weight_norm function to use
+    #     weight_norm = nn.utils.weight_norm
+    #     if hasattr(nn.utils.parametrizations, "weight_norm"):
+    #         weight_norm = nn.utils.parametrizations.weight_norm
+
+    #     # Apply weight_norm to each layer and replace the original layer with the wrapped layer
+    #     self.convs1 = nn.ModuleList([weight_norm(layer) for layer in self.convs1])
+    #     self.convs2 = nn.ModuleList([weight_norm(layer) for layer in self.convs2])
+
+
+    def remove_weight_norm(self):
+        for layer in self.convs1:
+            nn.utils.parametrize.remove_parametrizations(layer, "weight")
+        for layer in self.convs2:
+            nn.utils.parametrize.remove_parametrizations(layer, "weight")
+
+    def forward(self, hidden_states):
+        for conv1, conv2 in zip(self.convs1, self.convs2):
+            residual = hidden_states
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
+            hidden_states = conv1(hidden_states)
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
+            hidden_states = conv2(hidden_states)
+            hidden_states = hidden_states + residual
+        return hidden_states
+
+
+class VitsHifiGan(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.config = config
+        self.num_kernels = len(config.resblock_kernel_sizes)
+        self.num_upsamples = len(config.upsample_rates)
+        self.conv_pre = nn.Conv1d(
+            config.flow_size,
+            config.upsample_initial_channel,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+        )
+
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+        else:
+            weight_norm = nn.utils.weight_norm
+
+        self.upsampler = nn.ModuleList()
+        for i, (upsample_rate, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)):
+            self.upsampler.append(
+                weight_norm(
+                    nn.ConvTranspose1d(
+                        config.upsample_initial_channel // (2**i),
+                        config.upsample_initial_channel // (2 ** (i + 1)),
+                        kernel_size=kernel_size,
+                        stride=upsample_rate,
+                        padding=(kernel_size - upsample_rate) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.upsampler)):
+            channels = config.upsample_initial_channel // (2 ** (i + 1))
+            for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes):
+                self.resblocks.append(HifiGanResidualBlock(channels, kernel_size, dilation, config.leaky_relu_slope))
+
+        self.conv_post = nn.Conv1d(channels, 1, kernel_size=7, stride=1, padding=3, bias=False)
+
+        if config.speaker_embedding_size != 0:
+            self.cond = nn.Conv1d(config.speaker_embedding_size, config.upsample_initial_channel, 1)
+
+        # self.apply_weight_norm()
+
+    # def apply_weight_norm(self):
+    #     weight_norm = nn.utils.weight_norm
+    #     if hasattr(nn.utils.parametrizations, "weight_norm"):
+    #         weight_norm = nn.utils.parametrizations.weight_norm
+
+    #     self.upsampler = nn.ModuleList([weight_norm(layer) for layer in self.upsampler])
+    #     for layer in self.resblocks:
+    #         layer.apply_weight_norm()
+
+    def remove_weight_norm(self):
+        for layer in self.upsampler:
+            nn.utils.parametrize.remove_parametrizations(layer, "weight")
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+    def forward(
+        self, spectrogram: torch.FloatTensor, global_conditioning: Optional[torch.FloatTensor] = None
+    ) -> torch.FloatTensor:
+        r"""
+        Converts a spectrogram into a speech waveform.
+
+        Args:
+            spectrogram (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`):
+                Tensor containing the spectrograms.
+            global_conditioning (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_size, 1)`, *optional*):
+                Tensor containing speaker embeddings, for multispeaker models.
+
+        Returns:
+            `torch.FloatTensor`: Tensor of shape shape `(batch_size, 1, num_frames)` containing the speech waveform.
+        """
+        hidden_states = self.conv_pre(spectrogram)
+
+        if global_conditioning is not None:
+            hidden_states = hidden_states + self.cond(global_conditioning)
+
+        for i in range(self.num_upsamples):
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.config.leaky_relu_slope)
+            hidden_states = self.upsampler[i](hidden_states)
+
+            res_state = self.resblocks[i * self.num_kernels](hidden_states)
+            for j in range(1, self.num_kernels):
+                res_state += self.resblocks[i * self.num_kernels + j](hidden_states)
+            hidden_states = res_state / self.num_kernels
+
+        hidden_states = nn.functional.leaky_relu(hidden_states)
+        hidden_states = self.conv_post(hidden_states)
+        waveform = torch.tanh(hidden_states)
+        return waveform
+
+
+class VitsResidualCouplingLayer(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.half_channels = config.flow_size // 2
+
+        self.conv_pre = nn.Conv1d(self.half_channels, config.hidden_size, 1)
+        self.wavenet = VitsWaveNet(config, num_layers=config.prior_encoder_num_wavenet_layers)
+        self.conv_post = nn.Conv1d(config.hidden_size, self.half_channels, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        first_half, second_half = torch.split(inputs, [self.half_channels] * 2, dim=1)
+        hidden_states = self.conv_pre(first_half) * padding_mask
+        hidden_states = self.wavenet(hidden_states, padding_mask, global_conditioning)
+        mean = self.conv_post(hidden_states) * padding_mask
+        log_stddev = torch.zeros_like(mean)
+
+        if not reverse:
+            second_half = mean + second_half * torch.exp(log_stddev) * padding_mask
+            outputs = torch.cat([first_half, second_half], dim=1)
+            log_determinant = torch.sum(log_stddev, [1, 2])
+            return outputs, log_determinant
+        else:
+            second_half = (second_half - mean) * torch.exp(-log_stddev) * padding_mask
+            outputs = torch.cat([first_half, second_half], dim=1)
+            return outputs, None
+
+
+class VitsResidualCouplingBlock(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.flows = nn.ModuleList()
+        for _ in range(config.prior_encoder_num_flows):
+            self.flows.append(VitsResidualCouplingLayer(config))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                inputs, _ = flow(inputs, padding_mask, global_conditioning)
+                inputs = torch.flip(inputs, [1])
+        else:
+            for flow in reversed(self.flows):
+                inputs = torch.flip(inputs, [1])
+                inputs, _ = flow(inputs, padding_mask, global_conditioning, reverse=True)
+        return inputs
+
+
+class VitsDilatedDepthSeparableConv(nn.Module):
+    def __init__(self, config: IndicVitsConfig, dropout_rate=0.0):
+        super().__init__()
+        kernel_size = config.duration_predictor_kernel_size
+        channels = config.hidden_size
+        self.num_layers = config.depth_separable_num_layers
+
+        self.dropout = nn.Dropout(dropout_rate)
+        self.convs_dilated = nn.ModuleList()
+        self.convs_pointwise = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(self.num_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_dilated.append(
+                nn.Conv1d(
+                    in_channels=channels,
+                    out_channels=channels,
+                    kernel_size=kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_pointwise.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(nn.LayerNorm(channels))
+            self.norms_2.append(nn.LayerNorm(channels))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        if global_conditioning is not None:
+            inputs = inputs + global_conditioning
+
+        for i in range(self.num_layers):
+            hidden_states = self.convs_dilated[i](inputs * padding_mask)
+            hidden_states = self.norms_1[i](hidden_states.transpose(1, -1)).transpose(1, -1)
+            hidden_states = nn.functional.gelu(hidden_states)
+            hidden_states = self.convs_pointwise[i](hidden_states)
+            hidden_states = self.norms_2[i](hidden_states.transpose(1, -1)).transpose(1, -1)
+            hidden_states = nn.functional.gelu(hidden_states)
+            hidden_states = self.dropout(hidden_states)
+            inputs = inputs + hidden_states
+
+        return inputs * padding_mask
+
+
+class VitsConvFlow(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.filter_channels = config.hidden_size
+        self.half_channels = config.depth_separable_channels // 2
+        self.num_bins = config.duration_predictor_flow_bins
+        self.tail_bound = config.duration_predictor_tail_bound
+
+        self.conv_pre = nn.Conv1d(self.half_channels, self.filter_channels, 1)
+        self.conv_dds = VitsDilatedDepthSeparableConv(config)
+        self.conv_proj = nn.Conv1d(self.filter_channels, self.half_channels * (self.num_bins * 3 - 1), 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        first_half, second_half = torch.split(inputs, [self.half_channels] * 2, dim=1)
+
+        hidden_states = self.conv_pre(first_half)
+        hidden_states = self.conv_dds(hidden_states, padding_mask, global_conditioning)
+        hidden_states = self.conv_proj(hidden_states) * padding_mask
+
+        batch_size, channels, length = first_half.shape
+        hidden_states = hidden_states.reshape(batch_size, channels, -1, length).permute(0, 1, 3, 2)
+
+        unnormalized_widths = hidden_states[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = hidden_states[..., self.num_bins : 2 * self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_derivatives = hidden_states[..., 2 * self.num_bins :]
+
+        second_half, log_abs_det = _unconstrained_rational_quadratic_spline(
+            second_half,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            reverse=reverse,
+            tail_bound=self.tail_bound,
+        )
+
+        outputs = torch.cat([first_half, second_half], dim=1) * padding_mask
+        if not reverse:
+            log_determinant = torch.sum(log_abs_det * padding_mask, [1, 2])
+            return outputs, log_determinant
+        else:
+            return outputs, None
+
+
+class VitsElementwiseAffine(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.channels = config.depth_separable_channels
+        self.translate = nn.Parameter(torch.zeros(self.channels, 1))
+        self.log_scale = nn.Parameter(torch.zeros(self.channels, 1))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        if not reverse:
+            outputs = self.translate + torch.exp(self.log_scale) * inputs
+            outputs = outputs * padding_mask
+            log_determinant = torch.sum(self.log_scale * padding_mask, [1, 2])
+            return outputs, log_determinant
+        else:
+            outputs = (inputs - self.translate) * torch.exp(-self.log_scale) * padding_mask
+            return outputs, None
+
+
+class VitsStochasticDurationPredictor(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        embed_dim = config.speaker_embedding_size
+        filter_channels = config.hidden_size
+
+        self.conv_pre = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.conv_dds = VitsDilatedDepthSeparableConv(
+            config,
+            dropout_rate=config.duration_predictor_dropout,
+        )
+
+        if embed_dim != 0:
+            self.cond = nn.Conv1d(embed_dim, filter_channels, 1)
+
+        self.flows = nn.ModuleList()
+        self.flows.append(VitsElementwiseAffine(config))
+        for _ in range(config.duration_predictor_num_flows):
+            self.flows.append(VitsConvFlow(config))
+
+        self.post_conv_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_conv_dds = VitsDilatedDepthSeparableConv(
+            config,
+            dropout_rate=config.duration_predictor_dropout,
+        )
+
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(VitsElementwiseAffine(config))
+        for _ in range(config.duration_predictor_num_flows):
+            self.post_flows.append(VitsConvFlow(config))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, durations=None, reverse=False, noise_scale=1.0):
+        inputs = torch.detach(inputs)
+        inputs = self.conv_pre(inputs)
+
+        if global_conditioning is not None:
+            global_conditioning = torch.detach(global_conditioning)
+            inputs = inputs + self.cond(global_conditioning)
+
+        inputs = self.conv_dds(inputs, padding_mask)
+        inputs = self.conv_proj(inputs) * padding_mask
+
+        if not reverse:
+            hidden_states = self.post_conv_pre(durations)
+            hidden_states = self.post_conv_dds(hidden_states, padding_mask)
+            hidden_states = self.post_conv_proj(hidden_states) * padding_mask
+
+            random_posterior = (
+                torch.randn(durations.size(0), 2, durations.size(2)).to(device=inputs.device, dtype=inputs.dtype)
+                * padding_mask
+            )
+            log_determinant_posterior_sum = 0
+            latents_posterior = random_posterior
+            for flow in self.post_flows:
+                latents_posterior, log_determinant = flow(
+                    latents_posterior, padding_mask, global_conditioning=inputs + hidden_states
+                )
+                latents_posterior = torch.flip(latents_posterior, [1])
+                log_determinant_posterior_sum += log_determinant
+
+            first_half, second_half = torch.split(latents_posterior, [1, 1], dim=1)
+
+            log_determinant_posterior_sum += torch.sum(
+                (nn.functional.logsigmoid(first_half) + nn.functional.logsigmoid(-first_half)) * padding_mask, [1, 2]
+            )
+            logq = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (random_posterior**2)) * padding_mask, [1, 2])
+                - log_determinant_posterior_sum
+            )
+
+            first_half = (durations - torch.sigmoid(first_half)) * padding_mask
+            first_half = torch.log(torch.clamp_min(first_half, 1e-5)) * padding_mask
+            log_determinant_sum = torch.sum(-first_half, [1, 2])
+
+            latents = torch.cat([first_half, second_half], dim=1)
+            for flow in self.flows:
+                latents, log_determinant = flow(latents, padding_mask, global_conditioning=inputs)
+                latents = torch.flip(latents, [1])
+                log_determinant_sum += log_determinant
+
+            nll = torch.sum(0.5 * (math.log(2 * math.pi) + (latents**2)) * padding_mask, [1, 2]) - log_determinant_sum
+            return nll + logq
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
+
+            latents = (
+                torch.randn(inputs.size(0), 2, inputs.size(2)).to(device=inputs.device, dtype=inputs.dtype)
+                * noise_scale
+            )
+            for flow in flows:
+                latents = torch.flip(latents, [1])
+                latents, _ = flow(latents, padding_mask, global_conditioning=inputs, reverse=True)
+
+            log_duration, _ = torch.split(latents, [1, 1], dim=1)
+            return log_duration
+
+
+class VitsDurationPredictor(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        kernel_size = config.duration_predictor_kernel_size
+        filter_channels = config.duration_predictor_filter_channels
+
+        self.dropout = nn.Dropout(config.duration_predictor_dropout)
+        self.conv_1 = nn.Conv1d(config.hidden_size, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_1 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)
+        self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_2 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if config.speaker_embedding_size != 0:
+            self.cond = nn.Conv1d(config.speaker_embedding_size, config.hidden_size, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        inputs = torch.detach(inputs)
+
+        if global_conditioning is not None:
+            global_conditioning = torch.detach(global_conditioning)
+            inputs = inputs + self.cond(global_conditioning)
+
+        inputs = self.conv_1(inputs * padding_mask)
+        inputs = torch.relu(inputs)
+        inputs = self.norm_1(inputs.transpose(1, -1)).transpose(1, -1)
+        inputs = self.dropout(inputs)
+
+        inputs = self.conv_2(inputs * padding_mask)
+        inputs = torch.relu(inputs)
+        inputs = self.norm_2(inputs.transpose(1, -1)).transpose(1, -1)
+        inputs = self.dropout(inputs)
+
+        inputs = self.proj(inputs * padding_mask)
+        return inputs * padding_mask
+
+
+class VitsAttention(nn.Module):
+    """Multi-headed attention with relative positional representation."""
+
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.dropout = config.attention_dropout
+        self.window_size = config.window_size
+
+        self.head_dim = self.embed_dim // self.num_heads
+        self.scaling = self.head_dim**-0.5
+
+        if (self.head_dim * self.num_heads) != self.embed_dim:
+            raise ValueError(
+                f"hidden_size must be divisible by num_attention_heads (got `hidden_size`: {self.embed_dim}"
+                f" and `num_attention_heads`: {self.num_heads})."
+            )
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+
+        if self.window_size:
+            self.emb_rel_k = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)
+            self.emb_rel_v = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Input shape: Batch x Time x Channel"""
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+
+        bsz, tgt_len, _ = hidden_states.size()
+
+        # get query proj
+        query_states = self.q_proj(hidden_states) * self.scaling
+
+        # self_attention
+        key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+        value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
+        key_states = key_states.view(*proj_shape)
+        value_states = value_states.view(*proj_shape)
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if self.window_size is not None:
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, src_len)
+            relative_logits = torch.matmul(query_states, key_relative_embeddings.transpose(-2, -1))
+            rel_pos_bias = self._relative_position_to_absolute_position(relative_logits)
+            attn_weights += rel_pos_bias
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        if layer_head_mask is not None:
+            if layer_head_mask.size() != (self.num_heads,):
+                raise ValueError(
+                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
+                    f" {layer_head_mask.size()}"
+                )
+            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if output_attentions:
+            # this operation is a bit awkward, but it's required to
+            # make sure that attn_weights keeps its gradient.
+            # In order to do so, attn_weights have to be reshaped
+            # twice and have to be reused in the following
+            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.bmm(attn_probs, value_states)
+
+        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        if self.window_size is not None:
+            value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, src_len)
+            relative_weights = self._absolute_position_to_relative_position(attn_probs)
+            rel_pos_bias = torch.matmul(relative_weights, value_relative_embeddings)
+            attn_output += rel_pos_bias
+
+        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(1, 2)
+
+        # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
+        # partitioned aross GPUs when using tensor-parallelism.
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        pad_length = max(length - (self.window_size + 1), 0)
+        if pad_length > 0:
+            relative_embeddings = nn.functional.pad(relative_embeddings, [0, 0, pad_length, pad_length, 0, 0])
+
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        return relative_embeddings[:, slice_start_position:slice_end_position]
+
+    def _relative_position_to_absolute_position(self, x):
+        batch_heads, length, _ = x.size()
+
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = nn.functional.pad(x, [0, 1, 0, 0, 0, 0])
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch_heads, length * 2 * length])
+        x_flat = nn.functional.pad(x_flat, [0, length - 1, 0, 0])
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch_heads, length + 1, 2 * length - 1])
+        x_final = x_final[:, :length, length - 1 :]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        batch_heads, length, _ = x.size()
+
+        # Pad along column
+        x = nn.functional.pad(x, [0, length - 1, 0, 0, 0, 0])
+        x_flat = x.view([batch_heads, length * (2 * length - 1)])
+
+        # Add 0's in the beginning that will skew the elements after reshape
+        x_flat = nn.functional.pad(x_flat, [length, 0, 0, 0])
+        x_final = x_flat.view([batch_heads, length, 2 * length])[:, :, 1:]
+        return x_final
+
+
+class VitsFeedForward(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.conv_1 = nn.Conv1d(config.hidden_size, config.ffn_dim, config.ffn_kernel_size)
+        self.conv_2 = nn.Conv1d(config.ffn_dim, config.hidden_size, config.ffn_kernel_size)
+        self.dropout = nn.Dropout(config.activation_dropout)
+
+        if isinstance(config.hidden_act, str):
+            self.act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.act_fn = config.hidden_act
+
+        if config.ffn_kernel_size > 1:
+            pad_left = (config.ffn_kernel_size - 1) // 2
+            pad_right = config.ffn_kernel_size // 2
+            self.padding = [pad_left, pad_right, 0, 0, 0, 0]
+        else:
+            self.padding = None
+
+    def forward(self, hidden_states, padding_mask):
+        hidden_states = hidden_states.permute(0, 2, 1)
+        padding_mask = padding_mask.permute(0, 2, 1)
+
+        hidden_states = hidden_states * padding_mask
+        if self.padding is not None:
+            hidden_states = nn.functional.pad(hidden_states, self.padding)
+
+        hidden_states = self.conv_1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        hidden_states = hidden_states * padding_mask
+        if self.padding is not None:
+            hidden_states = nn.functional.pad(hidden_states, self.padding)
+
+        hidden_states = self.conv_2(hidden_states)
+        hidden_states = hidden_states * padding_mask
+
+        hidden_states = hidden_states.permute(0, 2, 1)
+        return hidden_states
+
+
+class VitsEncoderLayer(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.attention = VitsAttention(config)
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = VitsFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        padding_mask: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ):
+        residual = hidden_states
+        hidden_states, attn_weights = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.layer_norm(residual + hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.feed_forward(hidden_states, padding_mask)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.final_layer_norm(residual + hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class VitsEncoder(nn.Module):
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([VitsEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+        self.layerdrop = config.layerdrop
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        padding_mask: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        # expand attention_mask
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        hidden_states = hidden_states * padding_mask
+
+        synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)
+
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            dropout_probability = np.random.uniform(0, 1)
+
+            skip_the_layer = self.training and (dropout_probability < self.layerdrop)
+            if not skip_the_layer or synced_gpus:
+                # under fsdp or deepspeed zero3 all gpus must run in sync
+                if self.gradient_checkpointing and self.training:
+                    layer_outputs = self._gradient_checkpointing_func(
+                        encoder_layer.__call__,
+                        hidden_states,
+                        padding_mask,
+                        attention_mask,
+                        output_attentions,
+                    )
+                else:
+                    layer_outputs = encoder_layer(
+                        hidden_states,
+                        attention_mask=attention_mask,
+                        padding_mask=padding_mask,
+                        output_attentions=output_attentions,
+                    )
+                hidden_states = layer_outputs[0]
+
+            if skip_the_layer:
+                layer_outputs = (None, None)
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        hidden_states = hidden_states * padding_mask
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+class VitsTextEncoder(nn.Module):
+    """
+    Transformer encoder that uses relative positional representation instead of absolute positional encoding.
+    """
+
+    def __init__(self, config: IndicVitsConfig):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.encoder = VitsEncoder(config)
+        self.project = nn.Conv1d(config.hidden_size, config.flow_size * 2, kernel_size=1)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        padding_mask: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = True,
+    ) -> Union[Tuple[torch.Tensor], VitsTextEncoderOutput]:
+        hidden_states = self.embed_tokens(input_ids) * math.sqrt(self.config.hidden_size)
+
+        encoder_outputs = self.encoder(
+            hidden_states=hidden_states,
+            padding_mask=padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0] if not return_dict else encoder_outputs.last_hidden_state
+
+        stats = self.project(last_hidden_state.transpose(1, 2)).transpose(1, 2) * padding_mask
+        prior_means, prior_log_variances = torch.split(stats, self.config.flow_size, dim=2)
+
+        if not return_dict:
+            outputs = (last_hidden_state, prior_means, prior_log_variances) + encoder_outputs[1:]
+            return outputs
+
+        return VitsTextEncoderOutput(
+            last_hidden_state=last_hidden_state,
+            prior_means=prior_means,
+            prior_log_variances=prior_log_variances,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+class VitsPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = IndicVitsConfig
+    base_model_prefix = "vits"
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.kaiming_normal_(module.weight)
+            if module.bias is not None:
+                k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
+                nn.init.uniform_(module.bias, a=-k, b=k)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+VITS_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`IndicVitsConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+VITS_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0,
+            1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+        speaker_id (`int`, *optional*):
+            Which speaker embedding to use. Only used for multispeaker models.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The complete VITS model, for text-to-speech synthesis.",
+    VITS_START_DOCSTRING,
+)
+class IndicVitsModel(VitsPreTrainedModel):
+    def __init__(self, config: IndicVitsConfig):
+        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
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_encoder(self):
+        return self.text_encoder
+
+    @add_start_docstrings_to_model_forward(VITS_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=VitsModelOutput, config_class=_CONFIG_FOR_DOC)
+    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,
+    ) -> Union[Tuple[Any], VitsModelOutput]:
+        r"""
+        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.
+
+        Returns:
+
+        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.")
+
+        if attention_mask is not None:
+            input_padding_mask = attention_mask.unsqueeze(-1).float()
+        else:
+            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()
+        
+        speaker_and_style_embeddings = None
+
+        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}.")
+            if isinstance(speaker_id, int):
+                speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device)
+            speaker_and_style_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
+        
+        if self.config.num_emotions > 1 and emotion_id is not None:
+            if not 0 <= emotion_id < self.config.num_emotions:
+                raise ValueError(f"Set `emotion_id` in the range 0-{self.config.num_emotions - 1}.")
+            if isinstance(emotion_id, int):
+                emotion_id = torch.full(size=(1,), fill_value=emotion_id, device=self.device)
+            emotion_embeddings = self.embed_emotion(emotion_id).unsqueeze(-1)
+            if speaker_and_style_embeddings is not None:
+                speaker_and_style_embeddings += emotion_embeddings
+            else:
+                speaker_and_style_embeddings = emotion_embeddings
+
+        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_and_style_embeddings,
+                reverse=True,
+                noise_scale=self.noise_scale_duration,
+            )
+        else:
+            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_and_style_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_and_style_embeddings, reverse=True)
+
+        spectrogram = latents * output_padding_mask
+        waveform = self.decoder(spectrogram, speaker_and_style_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,
+        )