codingutf8 2023 the kakao enterprise s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license 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 activations import act2fn from integrations deepspeed import isdeepspeedzero3enabled from modelingattnmaskutils import prepare4dattentionmask from modelingoutputs import basemodeloutput modeloutput from modelingutils import pretrainedmodel from utils import addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationvits import vitsconfig logger logging getloggername general docstring configfordoc vitsconfig vitspretrainedmodelarchivelist facebookmmsttseng see all vits models at https huggingface comodels filtervits and all mms models at https huggingface comodels sorttrendingsearchfacebook2fmmstts dataclass class vitsmodeloutputmodeloutput waveform torch floattensor none sequencelengths torch floattensor none spectrogram optionaltupletorch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none dataclass class vitstextencoderoutputmodeloutput lasthiddenstate torch floattensor none priormeans torch floattensor none priorlogvariances torch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none torch jit script def fusedaddtanhsigmoidmultiplyinputa inputb numchannels inact inputa inputb tact torch tanhinact numchannels sact torch sigmoidinact numchannels acts tact sact return acts def unconstrainedrationalquadraticspline inputs unnormalizedwidths unnormalizedheights unnormalizedderivatives reversefalse tailbound5 0 minbinwidth1e3 minbinheight1e3 minderivative1e3 insideintervalmask inputs tailbound inputs tailbound outsideintervalmask insideintervalmask outputs torch zeroslikeinputs logabsdet torch zeroslikeinputs constant np lognp exp1 minderivative 1 unnormalizedderivatives nn functional padunnormalizedderivatives pad1 1 unnormalizedderivatives 0 constant unnormalizedderivatives 1 constant outputsoutsideintervalmask inputsoutsideintervalmask logabsdetoutsideintervalmask 0 0 outputsinsideintervalmask logabsdetinsideintervalmask rationalquadraticspline inputsinputsinsideintervalmask unnormalizedwidthsunnormalizedwidthsinsideintervalmask unnormalizedheightsunnormalizedheightsinsideintervalmask unnormalizedderivativesunnormalizedderivativesinsideintervalmask reversereverse tailboundtailbound minbinwidthminbinwidth minbinheightminbinheight minderivativeminderivative return outputs logabsdet def rationalquadraticspline inputs unnormalizedwidths unnormalizedheights unnormalizedderivatives reverse tailbound minbinwidth minbinheight minderivative upperbound tailbound lowerbound tailbound if torch mininputs lowerbound or torch maxinputs upperbound raise valueerrorinput to a transform is not within its domain numbins unnormalizedwidths shape1 if minbinwidth numbins 1 0 raise valueerrorfminimal bin width minbinwidth too large for the number of bins numbins if minbinheight numbins 1 0 raise valueerrorfminimal bin height minbinheight too large for the number of bins numbins widths nn functional softmaxunnormalizedwidths dim1 widths minbinwidth 1 minbinwidth numbins widths cumwidths torch cumsumwidths dim1 cumwidths nn functional padcumwidths pad1 0 modeconstant value0 0 cumwidths upperbound lowerbound cumwidths lowerbound cumwidths 0 lowerbound cumwidths 1 upperbound widths cumwidths 1 cumwidths 1 derivatives minderivative nn functional softplusunnormalizedderivatives heights nn functional softmaxunnormalizedheights dim1 heights minbinheight 1 minbinheight numbins heights cumheights torch cumsumheights dim1 cumheights nn functional padcumheights pad1 0 modeconstant value0 0 cumheights upperbound lowerbound cumheights lowerbound cumheights 0 lowerbound cumheights 1 upperbound heights cumheights 1 cumheights 1 binlocations cumheights if reverse else cumwidths binlocations 1 1e6 binidx torch suminputs none binlocations dim1 1 binidx binidx none inputcumwidths cumwidths gather1 binidx 0 inputbinwidths widths gather1 binidx 0 inputcumheights cumheights gather1 binidx 0 delta heights widths inputdelta delta gather1 binidx 0 inputderivatives derivatives gather1 binidx 0 inputderivativesplusone derivatives 1 gather1 binidx 0 inputheights heights gather1 binidx 0 intermediate1 inputderivatives inputderivativesplusone 2 inputdelta if not reverse theta inputs inputcumwidths inputbinwidths thetaoneminustheta theta 1 theta numerator inputheights inputdelta theta pow2 inputderivatives thetaoneminustheta denominator inputdelta intermediate1 thetaoneminustheta outputs inputcumheights numerator denominator derivativenumerator inputdelta pow2 inputderivativesplusone theta pow2 2 inputdelta thetaoneminustheta inputderivatives 1 theta pow2 logabsdet torch logderivativenumerator 2 torch logdenominator return outputs logabsdet else find the roots of a quadratic equation intermediate2 inputs inputcumheights intermediate3 intermediate2 intermediate1 a inputheights inputdelta inputderivatives intermediate3 b inputheights inputderivatives intermediate3 c inputdelta intermediate2 discriminant b pow2 4 a c if not discriminant 0 all raise runtimeerrorfinvalid discriminant discriminant root 2 c b torch sqrtdiscriminant outputs root inputbinwidths inputcumwidths thetaoneminustheta root 1 root denominator inputdelta intermediate1 thetaoneminustheta derivativenumerator inputdelta pow2 inputderivativesplusone root pow2 2 inputdelta thetaoneminustheta inputderivatives 1 root pow2 logabsdet torch logderivativenumerator 2 torch logdenominator return outputs logabsdet class vitswavenettorch nn module def initself config vitsconfig numlayers int super init self hiddensize config hiddensize self numlayers numlayers self inlayers torch nn modulelist self resskiplayers torch nn modulelist self dropout nn dropoutconfig wavenetdropout if hasattrnn utils parametrizations weightnorm weightnorm nn utils parametrizations weightnorm else weightnorm nn utils weightnorm if config speakerembeddingsize 0 condlayer torch nn conv1dconfig speakerembeddingsize 2 config hiddensize numlayers 1 self condlayer weightnormcondlayer nameweight for i in rangenumlayers dilation config wavenetdilationratei padding config wavenetkernelsize dilation dilation 2 inlayer torch nn conv1d inchannelsconfig hiddensize outchannels2 config hiddensize kernelsizeconfig wavenetkernelsize dilationdilation paddingpadding inlayer weightnorminlayer nameweight self inlayers appendinlayer last one is not necessary if i numlayers 1 resskipchannels 2 config hiddensize else resskipchannels config hiddensize resskiplayer torch nn conv1dconfig hiddensize resskipchannels 1 resskiplayer weightnormresskiplayer nameweight self resskiplayers appendresskiplayer def forwardself inputs paddingmask globalconditioningnone outputs torch zeroslikeinputs numchannelstensor torch inttensorself hiddensize if globalconditioning is not none globalconditioning self condlayerglobalconditioning for i in rangeself numlayers hiddenstates self inlayersiinputs if globalconditioning is not none condoffset i 2 self hiddensize globalstates globalconditioning condoffset condoffset 2 self hiddensize else globalstates torch zeroslikehiddenstates acts fusedaddtanhsigmoidmultiplyhiddenstates globalstates numchannelstensor0 acts self dropoutacts resskipacts self resskiplayersiacts if i self numlayers 1 resacts resskipacts self hiddensize inputs inputs resacts paddingmask outputs outputs resskipacts self hiddensize else outputs outputs resskipacts return outputs paddingmask def removeweightnormself if self speakerembeddingsize 0 torch nn utils removeweightnormself condlayer for layer in self inlayers torch nn utils removeweightnormlayer for layer in self resskiplayers torch nn utils removeweightnormlayer class vitsposteriorencodernn module def initself config vitsconfig super init self outchannels config flowsize self convpre nn conv1dconfig spectrogrambins config hiddensize 1 self wavenet vitswavenetconfig numlayersconfig posteriorencodernumwavenetlayers self convproj nn conv1dconfig hiddensize self outchannels 2 1 def forwardself inputs paddingmask globalconditioningnone inputs self convpreinputs paddingmask inputs self wavenetinputs paddingmask globalconditioning stats self convprojinputs paddingmask mean logstddev torch splitstats self outchannels dim1 sampled mean torch randnlikemean torch explogstddev paddingmask return sampled mean logstddev copied from transformers models speecht5 modelingspeecht5 hifiganresidualblock class hifiganresidualblocknn module def initself channels kernelsize3 dilation1 3 5 leakyreluslope0 1 super init self leakyreluslope leakyreluslope self convs1 nn modulelist nn conv1d channels channels kernelsize stride1 dilationdilationi paddingself getpaddingkernelsize dilationi for i in rangelendilation self convs2 nn modulelist nn conv1d channels channels kernelsize stride1 dilation1 paddingself getpaddingkernelsize 1 for in rangelendilation def getpaddingself kernelsize dilation1 return kernelsize dilation dilation 2 def applyweightnormself for layer in self convs1 nn utils weightnormlayer for layer in self convs2 nn utils weightnormlayer def removeweightnormself for layer in self convs1 nn utils removeweightnormlayer for layer in self convs2 nn utils removeweightnormlayer def forwardself hiddenstates for conv1 conv2 in zipself convs1 self convs2 residual hiddenstates hiddenstates nn functional leakyreluhiddenstates self leakyreluslope hiddenstates conv1hiddenstates hiddenstates nn functional leakyreluhiddenstates self leakyreluslope hiddenstates conv2hiddenstates hiddenstates hiddenstates residual return hiddenstates class vitshifigannn module def initself config vitsconfig super init self config config self numkernels lenconfig resblockkernelsizes self numupsamples lenconfig upsamplerates self convpre nn conv1d config flowsize config upsampleinitialchannel kernelsize7 stride1 padding3 self upsampler nn modulelist for i upsamplerate kernelsize in enumeratezipconfig upsamplerates config upsamplekernelsizes self upsampler append nn convtranspose1d config upsampleinitialchannel 2i config upsampleinitialchannel 2 i 1 kernelsizekernelsize strideupsamplerate paddingkernelsize upsamplerate 2 self resblocks nn modulelist for i in rangelenself upsampler channels config upsampleinitialchannel 2 i 1 for kernelsize dilation in zipconfig resblockkernelsizes config resblockdilationsizes self resblocks appendhifiganresidualblockchannels kernelsize dilation config leakyreluslope self convpost nn conv1dchannels 1 kernelsize7 stride1 padding3 biasfalse if config speakerembeddingsize 0 self cond nn conv1dconfig speakerembeddingsize config upsampleinitialchannel 1 def applyweightnormself for layer in self upsampler nn utils weightnormlayer for layer in self resblocks layer applyweightnorm def removeweightnormself for layer in self upsampler nn utils removeweightnormlayer for layer in self resblocks layer removeweightnorm def forward self spectrogram torch floattensor globalconditioning optionaltorch floattensor none torch floattensor r converts a spectrogram into a speech waveform args spectrogram torch floattensor of shape batchsize config spectrogrambins sequencelength tensor containing the spectrograms globalconditioning torch floattensor of shape batchsize config speakerembeddingsize 1 optional tensor containing speaker embeddings for multispeaker models returns torch floattensor tensor of shape shape batchsize 1 numframes containing the speech waveform multiheaded attention with relative positional representation def initself config vitsconfig super init self embeddim config hiddensize self numheads config numattentionheads self dropout config attentiondropout self windowsize config windowsize self headdim self embeddim self numheads self scaling self headdim0 5 if self headdim self numheads self embeddim raise valueerror fhiddensize must be divisible by numattentionheads got hiddensize self embeddim f and numattentionheads self numheads self kproj nn linearself embeddim self embeddim biasconfig usebias self vproj nn linearself embeddim self embeddim biasconfig usebias self qproj nn linearself embeddim self embeddim biasconfig usebias self outproj nn linearself embeddim self embeddim biasconfig usebias if self windowsize self embrelk nn parametertorch randn1 self windowsize 2 1 self headdim self scaling self embrelv nn parametertorch randn1 self windowsize 2 1 self headdim self scaling def shapeself tensor torch tensor seqlen int bsz int return tensor viewbsz seqlen self numheads self headdim transpose1 2 contiguous def forward self hiddenstates torch tensor keyvaluestates optionaltorch tensor none attentionmask optionaltorch tensor none layerheadmask optionaltorch tensor none outputattentions bool false tupletorch tensor optionaltorch tensor if keyvaluestates are provided this layer is used as a crossattention layer for the decoder get query proj selfattention this operation is a bit awkward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to be reshaped twice and have to be reused in the following use the embeddim from the config stored in the class rather than hiddenstate because attnoutput can be partitioned aross gpus when using tensorparallelism concat columns of pad to shift from relative to absolute indexing concat extra elements so to add up to shape len1 2len1 reshape and slice out the padded elements pad along column add 0 s in the beginning that will skew the elements after reshape expand attentionmask bsz seqlen bsz 1 tgtseqlen srcseqlen add layerdrop see https arxiv orgabs1909 11556 for description under deepspeed zero3 all gpus must run in sync transformer encoder that uses relative positional representation instead of absolute positional encoding an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights if isinstancemodule nn linear module weight data normalmean0 0 stdself config initializerrange if module bias is not none module bias data zero elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 elif isinstancemodule nn conv1d nn init kaimingnormalmodule weight if module bias is not none k math sqrtmodule groups module inchannels module kernelsize0 nn init uniformmodule bias ak bk elif isinstancemodule nn embedding module weight data normalmean0 0 stdself config initializerrange if module paddingidx is not none module weight datamodule paddingidx zero vitsstartdocstring 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 modulehttps pytorch orgdocsstablenn htmltorch 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 vitsconfig 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 frompretrained method to load the model weights addstartdocstrings the complete vits model for texttospeech synthesis vitsstartdocstring class vitsmodelvitspretrainedmodel def initself config vitsconfig super initconfig self config config self textencoder vitstextencoderconfig self flow vitsresidualcouplingblockconfig self decoder vitshifiganconfig if config usestochasticdurationprediction self durationpredictor vitsstochasticdurationpredictorconfig else self durationpredictor vitsdurationpredictorconfig if config numspeakers 1 self embedspeaker nn embeddingconfig numspeakers config speakerembeddingsize this is used only for training self posteriorencoder vitsposteriorencoderconfig these parameters control the synthesised speech properties self speakingrate config speakingrate self noisescale config noisescale self noisescaleduration config noisescaleduration initialize weights and apply final processing self postinit def getencoderself return self textencoder addstartdocstringstomodelforwardvitsinputsdocstring replacereturndocstringsoutputtypevitsmodeloutput configclassconfigfordoc def forward self inputids optionaltorch tensor none attentionmask optionaltorch tensor none speakerid optionalint none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none labels optionaltorch floattensor none uniontupleany vitsmodeloutput r labels torch floattensor of shape batchsize config spectrogrambins sequencelength 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 setseed import torch tokenizer vitstokenizer frompretrainedfacebookmmsttseng model vitsmodel frompretrainedfacebookmmsttseng inputs tokenizertexthello my dog is cute returntensorspt setseed555 make deterministic with torch nograd outputs modelinputsinputids outputs waveform shape torch size1 45824 create a padding mask for the output lengths of shape batch 1 maxoutputlength reconstruct an attention tensor of shape batch 1 outlength inlength expand prior distribution coding utf 8 2023 the kakao enterprise s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch vits model general docstring see all vits models at https huggingface co models filter vits and all mms models at https huggingface co models sort trending search facebook 2fmms tts 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 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 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 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 find the roots of a quadratic equation last one is not necessary copied from transformers models speecht5 modeling_speecht5 hifiganresidualblock 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 remove a useless vflow multi headed attention with relative positional representation 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 get query proj self_attention 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 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 concat columns of pad to shift from relative to absolute indexing concat extra elements so to add up to shape len 1 2 len 1 reshape and slice out the padded elements pad along column add 0 s in the beginning that will skew the elements after reshape expand attention_mask bsz seq_len bsz 1 tgt_seq_len src_seq_len add layerdrop see https arxiv org abs 1909 11556 for description under deepspeed zero3 all gpus must run in sync transformer encoder that uses relative positional representation instead of absolute positional encoding an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights 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 vitsconfig 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 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 this is used only for training these parameters control the synthesised speech properties initialize weights and apply final processing 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 create a padding mask for the output lengths of shape batch 1 max_output_length reconstruct an attention tensor of shape batch 1 out_length in_length expand prior distribution
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 ...activations import ACT2FN
from ...integrations.deepspeed import is_deepspeed_zero3_enabled
from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
from ...modeling_outputs import (
    BaseModelOutput,
    ModelOutput,
)
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_vits import VitsConfig


logger = logging.get_logger(__name__)



_CONFIG_FOR_DOC = "VitsConfig"


VITS_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/mms-tts-eng",
    
    
]


@dataclass
class VitsModelOutput(ModelOutput):
    

    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):
    

    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,
):
    
    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,
):
    
    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:
        
        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: VitsConfig, 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)

            
            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.remove_weight_norm(self.cond_layer)
        for layer in self.in_layers:
            torch.nn.utils.remove_weight_norm(layer)
        for layer in self.res_skip_layers:
            torch.nn.utils.remove_weight_norm(layer)


class VitsPosteriorEncoder(nn.Module):
    def __init__(self, config: VitsConfig):
        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



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

        self.convs1 = nn.ModuleList(
            [
                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(
            [
                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):
        for layer in self.convs1:
            nn.utils.weight_norm(layer)
        for layer in self.convs2:
            nn.utils.weight_norm(layer)

    def remove_weight_norm(self):
        for layer in self.convs1:
            nn.utils.remove_weight_norm(layer)
        for layer in self.convs2:
            nn.utils.remove_weight_norm(layer)

    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: VitsConfig):
        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,
        )

        self.upsampler = nn.ModuleList()
        for i, (upsample_rate, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)):
            self.upsampler.append(
                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)

    def apply_weight_norm(self):
        for layer in self.upsampler:
            nn.utils.weight_norm(layer)
        for layer in self.resblocks:
            layer.apply_weight_norm()

    def remove_weight_norm(self):
        for layer in self.upsampler:
            nn.utils.remove_weight_norm(layer)
        for layer in self.resblocks:
            layer.remove_weight_norm()

    def forward(
        self, spectrogram: torch.FloatTensor, global_conditioning: Optional[torch.FloatTensor] = None
    ) -> torch.FloatTensor:
        r
        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: VitsConfig):
        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: VitsConfig):
        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: VitsConfig, 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: VitsConfig):
        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: VitsConfig):
        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]]  

            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):
    

    def __init__(self, config: VitsConfig):
        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]]:
        

        
        

        bsz, tgt_len, _ = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scaling

        
        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:
            
            
            
            
            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)

        
        
        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()

        
        x = nn.functional.pad(x, [0, 1, 0, 0, 0, 0])

        
        x_flat = x.view([batch_heads, length * 2 * length])
        x_flat = nn.functional.pad(x_flat, [0, length - 1, 0, 0])

        
        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()

        
        x = nn.functional.pad(x, [0, length - 1, 0, 0, 0, 0])
        x_flat = x.view([batch_heads, length**2 + length * (length - 1)])

        
        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: VitsConfig):
        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: VitsConfig):
        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

        
        if attention_mask is not None:
            
            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)

        hidden_states = hidden_states * padding_mask

        deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()

        for encoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            
            dropout_probability = np.random.uniform(0, 1)

            skip_the_layer = self.training and (dropout_probability < self.layerdrop)
            if not skip_the_layer or deepspeed_zero3_is_enabled:
                
                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):
    

    def __init__(self, config: VitsConfig):
        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):
    

    config_class = VitsConfig
    base_model_prefix = "vits"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        
        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


VITS_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The complete VITS model, for text-to-speech synthesis.",
    VITS_START_DOCSTRING,
)
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)

        
        self.posterior_encoder = VitsPosteriorEncoder(config)

        
        self.speaking_rate = config.speaking_rate
        self.noise_scale = config.noise_scale
        self.noise_scale_duration = config.noise_scale_duration

        
        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,
        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
        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 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()

        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_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
        else:
            speaker_embeddings = None

        if labels is not None:
            raise NotImplementedError("Training of VITS is not supported yet.")

        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()

        
        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)

        
        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

        
        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,
        )