codingutf8 2023 the huggingface 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 clvp model import copy import math from dataclasses import dataclass from typing import dict optional tuple union import torch import torch utils checkpoint from torch import nn from torch nn import crossentropyloss from activations import act2fn from generation import generationconfig from modelingattnmaskutils import prepare4dattentionmask prepare4dcausalattentionmask from modelingoutputs import basemodeloutput basemodeloutputwithpastandcrossattentions basemodeloutputwithpooling causallmoutputwithcrossattentions from modelingutils import pretrainedmodel sequencesummary from pytorchutils import conv1d from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationclvp import clvpconfig clvpdecoderconfig clvpencoderconfig logger logging getloggername checkpointfordoc susnatoclvpdev clvppretrainedmodelarchivelist susnatoclvpdev see all clvp models at https huggingface comodels filterclvp copied from transformers models clip modelingclip contrastiveloss def contrastivelosslogits torch tensor torch tensor return nn functional crossentropylogits torch arangelenlogits devicelogits device copied from transformers models clip modelingclip cliploss with clipclvp imagelossspeechloss def clvplosssimilarity torch tensor torch tensor captionloss contrastivelosssimilarity speechloss contrastivelosssimilarity t return captionloss speechloss 2 0 copied from transformers models llama modelingllama rotatehalf def rotatehalfx applies rotary position embedding to the query and key tensors args q torch tensor the query tensor k torch tensor the key tensor cos torch tensor the cosine part of the rotary embedding sin torch tensor the sine part of the rotary embedding positionids torch tensor the position indices of the tokens corresponding to the query and key tensors for example this can be used to pass offsetted position ids when working with a kvcache unsqueezedim int optional defaults to 1 the unsqueezedim argument specifies the dimension along which to unsqueeze cospositionids and sinpositionids so that they can be properly broadcasted to the dimensions of q and k for example note that cospositionids and sinpositionids have the shape batchsize seqlen headdim then if q and k have the shape batchsize heads seqlen headdim then setting unsqueezedim1 makes cospositionids and sinpositionids broadcastable to the shapes of q and k similarly if q and k have the shape batchsize seqlen heads headdim then set unsqueezedim2 returns tupletorch tensor comprising of the query and key tensors rotated using the rotary position embedding this method adds extra bos and eos tokens to inputids and accordingly modifies the attentionmask which is used in clvpconditioningencoder and the generation loop of the clvpmodelforconditionalgeneration add the bos token at the beginning locate where the valid tokens end and then add the eos token if there are no pad tokens present then add eos to the end base class for clvp encoder s outputs that contains a pooling of the last hidden states as well as a projection output a linear layer on top of the pooled output args embeds torch floattensor of shape batchsize outputdim optional returned when model is initialized with withprojectiontrue the embeddings obtained by applying the projection layer to the pooleroutput lasthiddenstate torch floattensor of shape batchsize sequencelength hiddensize the hidden state of the last layer of the model pooleroutput torch floattensor of shape batchsize hiddensize pooled output of the lasthiddenstate hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue 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 batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the optional initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads args loss torch floattensor of shape 1 optional returned when returnloss is true contrastive loss for speechtext similarity speechids torch longtensor optional speechids or speech candidates generated by the clvpforcausallm model logitsperspeech torch floattensor of shape speechbatchsize textbatchsize the scaled dot product scores between speechembeds and textembeds this represents the speechtext similarity scores logitspertext torch floattensor of shape textbatchsize speechbatchsize the scaled dot product scores between textembeds and speechembeds this represents the textspeech similarity scores textembeds torch floattensor of shape batchsize outputdim the text embeddings obtained by applying the projection layer to the pooled output of the text encoder model speechembeds torch floattensor of shape batchsize outputdim the speech embeddings obtained by applying the projection layer to the pooled output of the speech encoder model textmodeloutput basemodeloutputwithpooling the pooled output of the lasthiddenstate of the text encoder model speechmodeloutput basemodeloutputwithpooling the pooled output of the lasthiddenstate of the speech encoder model decoderhiddenstates torch floattensor optional the hidden states of the decoder model textencoderhiddenstates torch floattensor optional the hidden states of the text encoder model speechencoderhiddenstates torch floattensor optional the hidden states of the speech encoder model copied from transformers models llama modelingllama llamarmsnorm with llamaclvp clvprmsnorm is equivalent to t5layernorm rotary position embedding class for clvp it was proposed in the paper roformer enhanced transformer with rotary position embedding please see https arxiv orgpdf2104 09864v1 pdf multiheaded attention to combine absolute and rotary positional embeddings into a single attention module copied from transformers models clip modelingclip clipattention shape raise error when positionids is none but rotaryposemb is provided because we need that when applying rotaryposemb to query and key states get query proj partial rotary embedding batchsize numheads seqlength headdim mask heads if we want to clvpgatedlinearunit uses the second half of the hiddenstates to act as a gate for the first half of the hiddenstates which controls the flow of data from the first of the tensor this mlp is used in clvp speech or text encoder models args hiddenstates torch floattensor of shape batch seqlen embeddim input to the layer rotaryposemb torch floattensor rotary position embeddings generated by clvprotarypositionalembedding module attentionmask torch floattensor of shape batch 1 tgtlen srclen attention mask where padding elements are indicated by very large negative values positionids torch longtensor denotes position ids of the input tokens outputattentions bool optional defaults to false whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail copied from transformers models gpt2 modelinggpt2 gpt2mlp with gpt2clvpdecodermlp residual connection residual connection this class processes the logmel spectrogramsextracted by the feature extractor and text tokensproduced by the tokenizer as inputs for the decoder model first each logmel spectrogram is processed into a single vector which captures valuable characteristics from each of them then the text tokens are converted into token embeddings and position embeddings are added afterwards both of these vectors are concatenated and then passed to the decoder model the text tokens helps to incorporate the text information and the logmel spectrogram is used to specify the voice characteristics into the generated mel tokens define group norms to be used before each attention layer define the attention layers calculates the value of numgroups for nn groupnorm this logic is taken from the official tortoise repository link https github comneonbjbtortoisettsblob4003544b6ff4b68c09856e04d3eff9da26d023c2tortoisemodelsarchutil pyl26 process text construct attention mask if not given we add bos and eos inputids in the modeling file instead of the tokenizer file to keep the logic simple this logic is specific to clvpconditioningencoder and not used by other modules process each logmel spectrogram into a single vector process each logmel spectrogram into a single vector repeat if there is either 1 text vs n audios or n texts vs 1 audio if there is n texts and m audios we will raise error since the number of text and audio must be same an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights factor self config initializerfactor if isinstancemodule nn embedding module weight data normalmean0 0 stdfactor 0 02 elif isinstancemodule nn linear conv1d nn conv1d module weight data normalmean0 0 stdfactor 0 02 if module bias is not none module bias data zero elif isinstancemodule clvpencodermlp factor self config initializerfactor inprojstd module config hiddensize0 5 2 module config numhiddenlayers 0 5 factor fcstd 2 module config hiddensize 0 5 factor nn init normalmodule fc1 proj weight if getattrmodule fc1 proj else module fc1 weight stdfcstd nn init normalmodule fc2 weight stdinprojstd elif isinstancemodule clvpencoder config self config textconfig if hasattrself config textconfig else self config factor config initializerfactor module projection weight data normalmean0 0 stdfactor config hiddensize0 5 elif isinstancemodule clvpconditioningencoder module melconv weight data normalmean0 0 stdfactor module melconv bias data zero elif isinstancemodule clvpforcausallm for name p in module namedparameters if name cproj weight p data normal mean0 0 stdself config initializerrange math sqrt2 self config numhiddenlayers if isinstancemodule nn layernorm module bias data zero module weight data fill1 0 clvpstartdocstring 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 clvpconfig 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 clvpdecoderinputsdocstring r args inputids torch longtensor of shape batchsize inputidslength indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids pastkeyvalues tupletupletorch tensor of length config nlayers contains precomputed hiddenstates key and values in the attention blocks as computed by the model see pastkeyvalues output below can be used to speed up sequential decoding the inputids which have their past given to this model should not be passed as inputids as they have already been computed attentionmask torch floattensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked if pastkeyvalues is used attentionmask needs to contain the masking strategy that was used for pastkeyvalues in other words the attentionmask always has to have the length lenpastkeyvalues leninputids what are attention masks glossaryattentionmask tokentypeids torch longtensor of shape batchsize inputidslength optional segment token indices to indicate first and second portions of the inputs indices are selected in 0 1 0 corresponds to a sentence a token 1 corresponds to a sentence b token what are token type ids glossarytokentypeids positionids torch longtensor of shape batchsize sequencelength optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config maxpositionembeddings 1 what are position ids glossarypositionids headmask torch floattensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the selfattention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optional optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix if pastkeyvalues is used optionally only the last inputsembeds have to be input see pastkeyvalues usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple transformer encoder consisting of config numhiddenlayers self attention layers each layer is a clvpencoderlayer args config clvpconfig outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both inputids and inputsembeds at the same time elif inputids is not none self warnifpaddingandnoattentionmaskinputids attentionmask inputshape inputids size inputids inputids view1 inputshape1 inputsembeds self tokenembeddinginputids elif inputsembeds is not none inputshape inputsembeds size 1 else raise valueerroryou have to specify either inputids or inputsembeds expand attentionmask and create positionids if needed if attentionmask is not none bsz seqlen bsz 1 tgtseqlen srcseqlen attentionmask prepare4dattentionmaskattentionmask inputsembeds dtype if positionids is none device inputids device if inputids is not none else inputsembeds device positionids torch arangeinputshape1 dtypetorch long devicedevice positionids positionids unsqueeze0 encoderstates if outputhiddenstates else none allattentions if outputattentions else none rotaryposemb self rotaryposembinputsembeds if self rotaryposemb is not none else none hiddenstates inputsembeds for idx encoderlayer in enumerateself layers if outputhiddenstates encoderstates encoderstates hiddenstates if self gradientcheckpointing and self training layeroutputs torch utils checkpoint checkpoint encoderlayer call hiddenstates rotaryposemb attentionmask positionids else layeroutputs encoderlayer hiddenstates rotaryposemb attentionmask positionids outputattentionsoutputattentions hiddenstates layeroutputs0 if outputattentions allattentions allattentions layeroutputs1 if outputhiddenstates encoderstates encoderstates hiddenstates lasthiddenstate hiddenstates lasthiddenstate self finallayernormlasthiddenstate take the mean over axis 1 and get pooled output pooledoutput self sequencesummarylasthiddenstate apply the projection layer embeds self projectionpooledoutput if not returndict return tuple v for v in embeds lasthiddenstate pooledoutput encoderstates allattentions if v is not none return clvpencoderoutput embedsembeds lasthiddenstatelasthiddenstate pooleroutputpooledoutput hiddenstatesencoderstates attentionsallattentions class clvpdecoderclvppretrainedmodel def initself config super initconfig self config config self inputembedslayer nn embeddingself config vocabsize self config hiddensize self positionembedslayer nn embeddingself config maxpositionembeddings self config hiddensize self drop nn dropoutself config embdpdrop self layers nn modulelistclvpdecoderlayerself config for in rangeself config numhiddenlayers self layernorm nn layernormself config hiddensize epsself config layernormepsilon self gradientcheckpointing false initialize weights and apply final processing self postinit def getinputembeddingsself return self inputembedslayer def setinputembeddingsself newembeddings self inputembedslayer newembeddings def pruneheadsself headstoprune for layer heads in headstoprune items self layerslayer attn pruneheadsheads addstartdocstringstomodelforwardclvpdecoderinputsdocstring def forward self inputids optionaltorch longtensor none attentionmask optionaltorch floattensor none tokentypeids optionaltorch longtensor none positionids optionaltorch longtensor none headmask optionaltorch floattensor none pastkeyvalues optionaltupletupletorch tensor none inputsembeds optionaltorch floattensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpastandcrossattentions outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates usecache usecache if usecache is not none else self config usecache returndict returndict if returndict is not none else self config usereturndict if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both inputids and inputsembeds at the same time elif inputids is not none self warnifpaddingandnoattentionmaskinputids attentionmask inputshape inputids size inputids inputids view1 inputshape1 inputids shape0 elif inputsembeds is not none inputshape inputsembeds size 1 inputsembeds shape0 else raise valueerroryou have to specify either inputids or inputsembeds device inputids device if inputids is not none else inputsembeds device if tokentypeids is not none tokentypeids tokentypeids view1 inputshape1 if pastkeyvalues is none pastkeyvalueslength 0 pastkeyvalues tuplenone lenself layers else pastkeyvalueslength pastkeyvalues00 size2 if positionids is none positionids torch arange pastkeyvalueslength inputshape1 pastkeyvalueslength dtypetorch long devicedevice positionids positionids unsqueeze0 view1 inputshape1 if inputsembeds is none inputsembeds self inputembedslayerinputids positionembeds self positionembedslayerpositionids inputsembeds inputsembeds positionembeds attentionmask prepare4dcausalattentionmask attentionmask inputshape inputsembeds pastkeyvalueslength prepare head mask if needed 1 0 in headmask indicate we keep the head attentionprobs has shape bsz x numattentionheads x n x n headmask has shape numhiddenlayers x batch x numattentionheads x n x n headmask self getheadmaskheadmask self config numhiddenlayers hiddenstates inputsembeds if tokentypeids is not none tokentypeembeds self inputembedslayertokentypeids hiddenstates hiddenstates tokentypeembeds hiddenstates self drophiddenstates outputshape 1 inputshape1 hiddenstates size1 if self gradientcheckpointing and self training if usecache logger warningonce usecachetrue is incompatible with gradient checkpointing setting usecachefalse usecache false presents if usecache else none allselfattentions if outputattentions else none allcrossattentions if outputattentions and self config addcrossattention else none allhiddenstates if outputhiddenstates else none for i block pastkeyvalue in enumeratezipself layers pastkeyvalues if outputhiddenstates allhiddenstates allhiddenstates hiddenstates if self gradientcheckpointing and self training outputs torch utils checkpoint checkpoint block call hiddenstates none attentionmask positionids headmaski else outputs block hiddenstates pastkeyvaluepastkeyvalue attentionmaskattentionmask positionidspositionids headmaskheadmaski usecacheusecache outputattentionsoutputattentions hiddenstates outputs0 if usecache is true presents presents outputs1 if outputattentions allselfattentions allselfattentions outputs2 if usecache else 1 if self config addcrossattention allcrossattentions allcrossattentions outputs3 if usecache else 2 hiddenstates self layernormhiddenstates hiddenstates hiddenstates viewoutputshape add last hidden state if outputhiddenstates allhiddenstates allhiddenstates hiddenstates if not returndict return tuple v for v in hiddenstates presents allhiddenstates allselfattentions allcrossattentions if v is not none return basemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluespresents hiddenstatesallhiddenstates attentionsallselfattentions crossattentionsallcrossattentions addstartdocstrings the bare clvp decoder model outputting raw hiddenstates without any specific head on top clvpstartdocstring class clvpmodelclvppretrainedmodel def initself config clvpdecoderconfig super initconfig self config config self decoder clvpdecoderself config initialize weights and apply final processing self postinit def getinputembeddingsself return self decoder inputembedslayer def setinputembeddingsself value self decoder inputembedslayer value def getdecoderself return self decoder addstartdocstringstomodelforwardclvpdecoderinputsdocstring def forward self inputids optionaltorch longtensor none attentionmask optionaltorch floattensor none tokentypeids optionaltorch longtensor none positionids optionaltorch longtensor none headmask optionaltorch floattensor none pastkeyvalues optionaltupletupletorch tensor none inputsembeds optionaltorch floattensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpastandcrossattentions outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates usecache usecache if usecache is not none else self config usecache returndict returndict if returndict is not none else self config usereturndict decoder outputs consists of decfeatures pastkeyvalue dechidden decattn decoderoutputs self decoder inputidsinputids attentionmaskattentionmask tokentypeidstokentypeids positionidspositionids headmaskheadmask pastkeyvaluespastkeyvalues inputsembedsinputsembeds usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict if not returndict return decoderoutputs return basemodeloutputwithpastandcrossattentions lasthiddenstatedecoderoutputs lasthiddenstate pastkeyvaluesdecoderoutputs pastkeyvalues hiddenstatesdecoderoutputs hiddenstates attentionsdecoderoutputs attentions crossattentionsdecoderoutputs crossattentions addstartdocstrings the clvp decoder model with a language modelling head on top clvpstartdocstring class clvpforcausallmclvppretrainedmodel def initself config super initconfig self config config self model clvpmodelself config self finalnorm nn layernormself config hiddensize self lmhead nn linearself config hiddensize self config vocabsize biastrue initialize weights and apply final processing self postinit def getinputembeddingsself return self model decoder inputembedslayer def setinputembeddingsself newembeddings self model decoder inputembedslayer newembeddings def preparemodelinputs self inputs optionaltorch tensor none bostokenid optionalint none modelkwargs optionaldictstr torch tensor none tupletorch tensor optionalstr dictstr torch tensor inputname self maininputname modelkwargs k v for k v in modelkwargs items if v is not none inputskwarg modelkwargs popinputname none if inputskwarg is not none and inputs is not none raise valueerror finputs inputs were passed alongside inputname which is not allowed fmake sure to either pass inputs or inputname elif inputskwarg is not none inputs inputskwarg if inputname inputids and inputsembeds in modelkwargs modelkwargsinputids self maybeinitializeinputidsforgeneration inputs bostokenid modelkwargsmodelkwargs inputs inputname modelkwargsinputsembeds inputsembeds check if conditioningembeds are provided or not if yes then concatenate the bostokenid at the end of the conditioningembeds then we must subtract the positionalids because during the forward pass it will be added anyways so we must cancel them out here conditioningembeds modelkwargs getconditioningembeds none if conditioningembeds is not none melstarttokenembedding self model decoder inputembedslayer torch full conditioningembeds shape0 1 fillvalueself config bostokenid deviceconditioningembeds device melstarttokenembedding self model decoder positionembedslayer torch fullconditioningembeds shape0 1 fillvalue0 deviceconditioningembeds device conditioningembeds torch concatconditioningembeds melstarttokenembedding dim1 subtract the positionalids here if hasattrmodelkwargs attentionmask positionids modelkwargsattentionmask long cumsum1 1 else positionids torch range 0 conditioningembeds shape1 1 dtypetorch long deviceconditioningembeds device positionids positionids unsqueeze0 repeatconditioningembeds shape0 1 modelkwargsinputsembeds conditioningembeds self model decoder positionembedslayer positionids modelkwargsinputids torch onesmodelkwargsinputsembeds shape0 1 dtypetorch long deviceself device self config bostokenid return modelkwargsinputsembeds inputsembeds modelkwargs inputs self maybeinitializeinputidsforgenerationinputs bostokenid modelkwargs return inputs inputname modelkwargs def prepareinputsforgeneration self inputids pastkeyvaluesnone inputsembedsnone conditioningembedsnone kwargs inputidslength inputids shape1 tokentypeids kwargs gettokentypeids none only last token for inputsids if past is defined in kwargs if pastkeyvalues pastlength pastkeyvalues00 shape2 some generation methods already pass only the last input id if inputids shape1 pastlength removeprefixlength pastlength else default to old behavior keep only final id removeprefixlength inputids shape1 1 inputids inputids removeprefixlength if tokentypeids is not none tokentypeids tokentypeids inputids shape1 attentionmask kwargs getattentionmask none positionids kwargs getpositionids none if attentionmask is not none and positionids is none create positionids on the fly for batch generation positionids attentionmask long cumsum1 1 positionids maskedfillattentionmask 0 1 if pastkeyvalues positionids positionids 1 unsqueeze1 else positionids none if conditioningembeds is not none and pastkeyvalues is not none positionids torch tensorinputidslength dtypetorch long deviceinputids device if inputsembeds are passed we only want to use them in the 1st generation step if inputsembeds is not none and pastkeyvalues is none modelinputs inputsembeds inputsembeds else modelinputs inputids inputids modelinputs update pastkeyvalues pastkeyvalues usecache kwargs getusecache positionids positionids tokentypeids tokentypeids return modelinputs addstartdocstringstomodelforwardclvpdecoderinputsdocstring def forward self inputids optionaltorch longtensor none pastkeyvalues optionaltupletupletorch tensor none attentionmask optionaltorch floattensor none tokentypeids optionaltorch longtensor none positionids optionaltorch longtensor none headmask optionaltorch floattensor none inputsembeds optionaltorch floattensor none labels optionaltorch longtensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple causallmoutputwithcrossattentions r labels torch longtensor of shape batchsize sequencelength optional labels for language modeling note that the labels are shifted inside the model i e you can set labels inputids indices are selected in 100 0 config vocabsize all labels set to 100 are ignored masked the loss is only computed for labels in 0 config vocabsize shift so that tokens n predict n flatten the tokens copied from transformers models gpt2 modelinggpt2 gpt2lmheadmodel reordercache this function is used to reorder the pastkeyvalues cache if pretrainedmodel beamsearch or pretrainedmodel beamsample is called this is required to match pastkeyvalues with the correct beamidx at every generation step initialize weights and apply final processing taken from the original repo link https github comneonbjbtortoisettsblob4003544b6ff4b68c09856e04d3eff9da26d023c2tortoiseapi pyl117 this method modifies the output of the decoder model such as replacing the eostokenid and changing the last few tokens of each sequence args speechids torch longtensor this refers to the output of the decoder model this means that no stop tokens were found so the sentence was still being generated in that case we don t need to apply any padding so just skip to the next sequence of tokens outputs self textencodermodel inputidsinputids inputsembedstextencoderinputsembeds attentionmaskattentionmask return outputs0 def getspeechfeatures self speechids optionaltorch longtensor none inputids optionaltorch longtensor none inputfeatures optionaltorch floattensor none conditioningencoderinputsembeds optionaltorch floattensor none attentionmask optionaltorch tensor none generationconfig optionalgenerationconfig none kwargs torch floattensor r this method can be used to extract speechembeds the speech embeddings are obtained by applying the speech model on speechids if speechids is not present but both inputids and inputfeatures are given then the decoder model will be used to first generate the speechids and then applying the speech model args speechids torch longtensor of shape batchsize numspeechids optional speech tokens padding will be ignored by default should you provide it if speechids are provided then inputids and inputfeatures will be automatically ignored inputids torch longtensor of shape batchsize sequencelength optional input text tokens processed from the clvptokenizer if speechids is not provided then inputids and inputfeatures will be used inputfeatures torch floattensor of shape batchsize featuresize timedim optional indicates logmelspectrogram representations for audio returned by clvpfeatureextractor if speechids is not provided then inputids and inputfeatures will be used conditioningencoderinputsembeds torch floattensor optional inputsembeds for clvpconditioningencoder can be used in place of inputids attentionmask torch longtensor of shape batchsize sequencelength optional mask to avoid performing attention on padding speech 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 glossaryattentionmask generationconfig generationconfig optional generation config to control the generation of speechids if they are not provided returns torch floattensor of shape batchsize outputdim the speech embeddings obtained by applying the projection layer to the pooled output of the clvp speech model examples python import datasets from transformers import clvpprocessor clvpmodelforconditionalgeneration define the text and load the audio we are taking an audio example from huggingface hub using datasets library text this is an example text ds datasets loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation ds ds castcolumnaudio datasets audiosamplingrate22050 audio sr ds sortid selectrange1 1audio0 values define processor and model processor clvpprocessor frompretrainedsusnatoclvpdev model clvpmodelforconditionalgeneration frompretrainedsusnatoclvpdev generate processor output and model output processoroutput processorrawspeechaudio samplingratesr texttext returntensorspt speechembeds model getspeechfeatures inputidsprocessoroutputinputids inputfeaturesprocessoroutputinputfeatures use clvp model s config for some fields if specified instead of those of speech text components outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict conditioningembeds self conditioningencoder inputfeaturesinputfeatures inputidsinputids inputsembedsconditioningencoderinputsembeds attentionmaskattentionmask decoderoutputs self speechdecodermodel inputsembedsconditioningembeds outputhiddenstatesoutputhiddenstates returndictreturndict speechids decoderoutputs0 since we will get the embeds of shape batchsize seqlen embeddingdim during the forward pass we must convert it to tokens to make it compaitable with speechtransformer if speechids ndim 3 speechids speechids argmax2 speechids self fixspeechdecoderoutputspeechids speechoutputs self speechencodermodel inputidsspeechids outputhiddenstatesoutputhiddenstates returndictreturndict textoutputs self textencodermodel inputidsinputids inputsembedstextencoderinputsembeds attentionmaskattentionmask outputhiddenstatesoutputhiddenstates returndictreturndict speechembeds speechoutputs0 textembeds textoutputs0 normalized features speechembeds speechembeds speechembeds normp2 dim1 keepdimtrue textembeds textembeds textembeds normp2 dim1 keepdimtrue cosine similarity as logits logitscale self logitscale exp logitspertext torch matmultextembeds speechembeds t logitscale logitsperspeech logitspertext t loss none if returnloss loss clvplosslogitspertext if not returndict output logitsperspeech logitspertext textembeds speechembeds textoutputs2 speechoutputs2 if outputhiddenstates output decoderoutputs1 textoutputs1 speechoutputs1 return loss output if loss is not none else output return clvpoutput lossloss logitsperspeechlogitsperspeech logitspertextlogitspertext textembedstextembeds speechembedsspeechembeds textmodeloutputtextoutputs2 speechmodeloutputspeechoutputs2 decoderhiddenstatesdecoderoutputs hiddenstates textencoderhiddenstatestextoutputs hiddenstates speechencoderhiddenstatesspeechoutputs hiddenstates torch nograd def generate self inputids torch longtensor none inputfeatures torch floattensor none attentionmask optionaltorch longtensor none generationconfig optionalgenerationconfig none padtomaxmeltokens optionalint none outputhiddenstates optionalbool none kwargs if the input sequences are larger than self config decoderconfig maxtexttokens 3 then raise error because we need to add 3 tokens 1 bos tokens and 2 eos tokens to the inputids in clvpconditioningencoder to properly sample sequencelength inputids shape1 if sequencelength self config decoderconfig maxtexttokens 3 raise valueerror fmaximum sequence length reached found inputids of length sequencelength fplease make sure that the maximum length of inputids is self config decoderconfig maxtexttokens 3 if generationconfig is none generationconfig self generationconfig generationconfig copy deepcopygenerationconfig modelkwargs generationconfig updatekwargs all unused kwargs must be model kwargs generationconfig validate self validatemodelkwargsmodelkwargs copy pad inputids as specified in the original repo link https github comneonbjbtortoisettsblob80f89987a5abda5e2b082618cd74f9c7411141dctortoiseapi pyl380 inputids attentionmask padextraboseostokens inputids attentionmask addbostokenfalse bostokenidself config textconfig bostokenid eostokenidself config textconfig eostokenid conditioningembeds self conditioningencoder inputfeaturesinputfeatures inputidsinputids attentionmaskattentionmask decoderoutputs self speechdecodermodel generate conditioningembedsconditioningembeds generationconfiggenerationconfig outputhiddenstatesoutputhiddenstates returndictgenerationconfig returndictingenerate if isinstancedecoderoutputs modeloutput speechids decoderoutputs sequences pad to padtomaxmeltokens if given to replicate the original repo logic link https github comneonbjbtortoisettsblob80f89987a5abda5e2b082618cd74f9c7411141dctortoiseapi pyl430 if padtomaxmeltokens is not none paddingneeded padtomaxmeltokens speechids shape1 speechids torch nn functional pad speechids 0 paddingneeded valueself generationconfig eostokenid speechids self fixspeechdecoderoutputspeechids speechoutputs self speechencodermodel inputidsspeechids outputhiddenstatesoutputhiddenstates returndictgenerationconfig returndictingenerate textoutputs self textencodermodel inputidsinputids attentionmaskattentionmask outputhiddenstatesoutputhiddenstates returndictgenerationconfig returndictingenerate speechembeds speechoutputs0 textembeds textoutputs0 normalized features speechembeds speechembeds speechembeds normp2 dim1 keepdimtrue textembeds textembeds textembeds normp2 dim1 keepdimtrue cosine similarity as logits logitscale self logitscale exp logitspertext torch matmultextembeds speechembeds t logitscale logitsperspeech logitspertext t if not generationconfig returndictingenerate output speechids logitsperspeech logitspertext textembeds speechembeds textoutputs2 speechoutputs2 if outputhiddenstates output decoderoutputs1 textoutputs1 speechoutputs1 return output return clvpoutput speechidsspeechids logitsperspeechlogitsperspeech logitspertextlogitspertext textembedstextembeds speechembedsspeechembeds textmodeloutputtextoutputs2 speechmodeloutputspeechoutputs2 decoderhiddenstatesdecoderoutputs hiddenstates textencoderhiddenstatestextoutputs hiddenstates speechencoderhiddenstatesspeechoutputs hiddenstates coding utf 8 2023 the huggingface 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 clvp model see all clvp models at https huggingface co models filter clvp copied from transformers models clip modeling_clip contrastive_loss copied from transformers models clip modeling_clip clip_loss with clip clvp image_loss speech_loss copied from transformers models llama modeling_llama rotate_half rotates half the hidden dims of the input applies rotary position embedding to the query and key tensors args q torch tensor the query tensor k torch tensor the key tensor cos torch tensor the cosine part of the rotary embedding sin torch tensor the sine part of the rotary embedding position_ids torch tensor the position indices of the tokens corresponding to the query and key tensors for example this can be used to pass offsetted position ids when working with a kv cache unsqueeze_dim int optional defaults to 1 the unsqueeze_dim argument specifies the dimension along which to unsqueeze cos position_ids and sin position_ids so that they can be properly broadcasted to the dimensions of q and k for example note that cos position_ids and sin position_ids have the shape batch_size seq_len head_dim then if q and k have the shape batch_size heads seq_len head_dim then setting unsqueeze_dim 1 makes cos position_ids and sin position_ids broadcastable to the shapes of q and k similarly if q and k have the shape batch_size seq_len heads head_dim then set unsqueeze_dim 2 returns tuple torch tensor comprising of the query and key tensors rotated using the rotary position embedding this method adds extra bos and eos tokens to input_ids and accordingly modifies the attention_mask which is used in clvpconditioningencoder and the generation loop of the clvpmodelforconditionalgeneration add the bos token at the beginning locate where the valid tokens end and then add the eos token if there are no pad tokens present then add eos to the end base class for clvp encoder s outputs that contains a pooling of the last hidden states as well as a projection output a linear layer on top of the pooled output args embeds torch floattensor of shape batch_size output_dim optional returned when model is initialized with with_projection true the embeddings obtained by applying the projection layer to the pooler_output last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size the hidden state of the last layer of the model pooler_output torch floattensor of shape batch_size hidden_size pooled output of the last_hidden_state 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 attentions weights after the attention softmax used to compute the weighted average in the self attention heads args loss torch floattensor of shape 1 optional returned when return_loss is true contrastive loss for speech text similarity speech_ids torch longtensor optional speech_ids or speech candidates generated by the clvpforcausallm model logits_per_speech torch floattensor of shape speech_batch_size text_batch_size the scaled dot product scores between speech_embeds and text_embeds this represents the speech text similarity scores logits_per_text torch floattensor of shape text_batch_size speech_batch_size the scaled dot product scores between text_embeds and speech_embeds this represents the text speech similarity scores text_embeds torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of the text encoder model speech_embeds torch floattensor of shape batch_size output_dim the speech embeddings obtained by applying the projection layer to the pooled output of the speech encoder model text_model_output basemodeloutputwithpooling the pooled output of the last_hidden_state of the text encoder model speech_model_output basemodeloutputwithpooling the pooled output of the last_hidden_state of the speech encoder model decoder_hidden_states torch floattensor optional the hidden states of the decoder model text_encoder_hidden_states torch floattensor optional the hidden states of the text encoder model speech_encoder_hidden_states torch floattensor optional the hidden states of the speech encoder model copied from transformers models llama modeling_llama llamarmsnorm with llama clvp clvprmsnorm is equivalent to t5layernorm rotary position embedding class for clvp it was proposed in the paper roformer enhanced transformer with rotary position embedding please see https arxiv org pdf 2104 09864v1 pdf multi headed attention to combine absolute and rotary positional embeddings into a single attention module copied from transformers models clip modeling_clip clipattention _shape raise error when position_ids is none but rotary_pos_emb is provided because we need that when applying rotary_pos_emb to query and key states get query proj partial rotary embedding batch_size num_heads seq_length head_dim mask heads if we want to clvpgatedlinearunit uses the second half of the hidden_states to act as a gate for the first half of the hidden_states which controls the flow of data from the first of the tensor this mlp is used in clvp speech or text encoder models args hidden_states torch floattensor of shape batch seq_len embed_dim input to the layer rotary_pos_emb torch floattensor rotary position embeddings generated by clvprotarypositionalembedding module attention_mask torch floattensor of shape batch 1 tgt_len src_len attention mask where padding elements are indicated by very large negative values position_ids torch longtensor denotes position ids of the input tokens output_attentions bool optional defaults to false whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail copied from transformers models gpt2 modeling_gpt2 gpt2mlp with gpt2 clvpdecodermlp residual connection residual connection this class processes the log mel spectrograms extracted by the feature extractor and text tokens produced by the tokenizer as inputs for the decoder model first each log mel spectrogram is processed into a single vector which captures valuable characteristics from each of them then the text tokens are converted into token embeddings and position embeddings are added afterwards both of these vectors are concatenated and then passed to the decoder model the text tokens helps to incorporate the text information and the log mel spectrogram is used to specify the voice characteristics into the generated mel tokens define group norms to be used before each attention layer define the attention layers calculates the value of num_groups for nn groupnorm this logic is taken from the official tortoise repository link https github com neonbjb tortoise tts blob 4003544b6ff4b68c09856e04d3eff9da26d023c2 tortoise models arch_util py l26 process text construct attention mask if not given we add bos and eos input_ids in the modeling file instead of the tokenizer file to keep the logic simple this logic is specific to clvpconditioningencoder and not used by other modules process each log mel spectrogram into a single vector process each log mel spectrogram into a single vector repeat if there is either 1 text vs n audios or n texts vs 1 audio if there is n texts and m audios we will raise error since the number of text and audio must be same 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 clvpconfig 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 optional 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 input_features torch floattensor of shape batch_size feature_size time_dim indicates log mel spectrogram representations for audio returned by clvpfeatureextractor conditioning_encoder_inputs_embeds torch floattensor optional inputs_embeds for clvpconditioningencoder can be used in place of input_ids text_encoder_inputs_embeds torch floattensor optional inputs_embeds for the text encoder model passed in place of input_ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding text 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 return_loss bool optional whether or not to return the contrastive loss 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 args input_ids torch longtensor of shape batch_size input_ids_length indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids past_key_values tuple tuple torch tensor of length config n_layers contains precomputed hidden states key and values in the attention blocks as computed by the model see past_key_values output below can be used to speed up sequential decoding the input_ids which have their past given to this model should not be passed as input_ids as they have already been computed attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked if past_key_values is used attention_mask needs to contain the masking strategy that was used for past_key_values in other words the attention_mask always has to have the length len past_key_values len input_ids what are attention masks glossary attention mask token_type_ids torch longtensor of shape batch_size input_ids_length optional segment token indices to indicate first and second portions of the inputs indices are selected in 0 1 0 corresponds to a sentence a token 1 corresponds to a sentence b token what are token type ids glossary token type ids position_ids torch longtensor of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary position ids head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix if past_key_values is used optionally only the last inputs_embeds have to be input see past_key_values use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values 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 transformer encoder consisting of config num_hidden_layers self attention layers each layer is a clvpencoderlayer args config clvpconfig args input_ids torch longtensor of shape batch_size input_ids_length optional indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional input embeddings for the model this bypasses the model s internal embedding lookup matrix attention_mask torch longtensor of shape batch_size sequence_length optional mask to avoid performing 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 position_ids torch longtensor optional denotes the position ids of input_ids 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 expand attention_mask and create position_ids if needed bsz seq_len bsz 1 tgt_seq_len src_seq_len take the mean over axis 1 and get pooled output apply the projection layer transformer decoder consisting of config num_hidden_layers layers each layer is a clvpdecoderlayer initialize weights and apply final processing prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer prepare head mask if needed 1 0 in head_mask indicate we keep the head attention_probs has shape bsz x num_attention_heads x n x n head_mask has shape num_hidden_layers x batch x num_attention_heads x n x n add last hidden state initialize weights and apply final processing decoder outputs consists of dec_features past_key_value dec_hidden dec_attn initialize weights and apply final processing this function extracts the model specific inputs for generation check if conditioning_embeds are provided or not if yes then concatenate the bos_token_id at the end of the conditioning_embeds then we must subtract the positional_ids because during the forward pass it will be added anyways so we must cancel them out here subtract the positional_ids here only last token for inputs_ids if past is defined in kwargs some generation methods already pass only the last input id default to old behavior keep only final id create position_ids on the fly for batch generation if inputs_embeds are passed we only want to use them in the 1st generation step labels torch longtensor of shape batch_size sequence_length optional labels for language modeling note that the labels are shifted inside the model i e you can set labels input_ids indices are selected in 100 0 config vocab_size all labels set to 100 are ignored masked the loss is only computed for labels in 0 config vocab_size shift so that tokens n predict n flatten the tokens copied from transformers models gpt2 modeling_gpt2 gpt2lmheadmodel _reorder_cache this function is used to re order the past_key_values cache if pretrainedmodel beam_search or pretrainedmodel beam_sample is called this is required to match past_key_values with the correct beam_idx at every generation step initialize weights and apply final processing taken from the original repo link https github com neonbjb tortoise tts blob 4003544b6ff4b68c09856e04d3eff9da26d023c2 tortoise api py l117 this method modifies the output of the decoder model such as replacing the eos_token_id and changing the last few tokens of each sequence args speech_ids torch longtensor this refers to the output of the decoder model this means that no stop tokens were found so the sentence was still being generated in that case we don t need to apply any padding so just skip to the next sequence of tokens this method can be used to extract text_embeds from a text the text embeddings obtained by applying the projection layer to the pooled output of the clvp text encoder model 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 what are input ids glossary input ids text_encoder_inputs_embeds torch floattensor optional inputs_embeds for the text encoder model passed in place of input_ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing 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 returns torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of the clvp text model examples python from transformers import clvpprocessor clvpmodelforconditionalgeneration define the text text this is an example text define processor and model processor clvpprocessor from_pretrained susnato clvp_dev model clvpmodelforconditionalgeneration from_pretrained susnato clvp_dev generate processor output and text embeds processor_output processor text text return_tensors pt text_embeds model get_text_features input_ids processor_output input_ids this method can be used to extract speech_embeds the speech embeddings are obtained by applying the speech model on speech_ids if speech_ids is not present but both input_ids and input_features are given then the decoder model will be used to first generate the speech_ids and then applying the speech model args speech_ids torch longtensor of shape batch_size num_speech_ids optional speech tokens padding will be ignored by default should you provide it if speech_ids are provided then input_ids and input_features will be automatically ignored input_ids torch longtensor of shape batch_size sequence_length optional input text tokens processed from the clvptokenizer if speech_ids is not provided then input_ids and input_features will be used input_features torch floattensor of shape batch_size feature_size time_dim optional indicates log melspectrogram representations for audio returned by clvpfeatureextractor if speech_ids is not provided then input_ids and input_features will be used conditioning_encoder_inputs_embeds torch floattensor optional inputs_embeds for clvpconditioningencoder can be used in place of input_ids attention_mask torch longtensor of shape batch_size sequence_length optional mask to avoid performing attention on padding speech 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 generation_config generationconfig optional generation config to control the generation of speech_ids if they are not provided returns torch floattensor of shape batch_size output_dim the speech embeddings obtained by applying the projection layer to the pooled output of the clvp speech model examples python import datasets from transformers import clvpprocessor clvpmodelforconditionalgeneration define the text and load the audio we are taking an audio example from huggingface hub using datasets library text this is an example text ds datasets load_dataset hf internal testing librispeech_asr_dummy clean split validation ds ds cast_column audio datasets audio sampling_rate 22050 _ audio sr ds sort id select range 1 1 audio 0 values define processor and model processor clvpprocessor from_pretrained susnato clvp_dev model clvpmodelforconditionalgeneration from_pretrained susnato clvp_dev generate processor output and model output processor_output processor raw_speech audio sampling_rate sr text text return_tensors pt speech_embeds model get_speech_features input_ids processor_output input_ids input_features processor_output input_features returns examples python import datasets from transformers import clvpprocessor clvpmodelforconditionalgeneration define the text and load the audio we are taking an audio example from huggingface hub using datasets library text this is an example text ds datasets load_dataset hf internal testing librispeech_asr_dummy clean split validation ds ds cast_column audio datasets audio sampling_rate 22050 _ audio sr ds sort id select range 1 1 audio 0 values define processor and model processor clvpprocessor from_pretrained susnato clvp_dev model clvpmodelforconditionalgeneration from_pretrained susnato clvp_dev processor outputs and model outputs processor_output processor raw_speech audio sampling_rate sr text text return_tensors pt outputs model input_ids processor_output input_ids input_features processor_output input_features return_dict true use clvp model s config for some fields if specified instead of those of speech text components since we will get the embeds of shape batch_size seq_len embedding_dim during the forward pass we must convert it to tokens to make it compaitable with speech_transformer normalized features cosine similarity as logits generate method for clvpmodelforconditionalgeneration this method calls the generate method of clvpforcausallm and then uses those generated speech_ids to process text_embeds and speech_embeds using clvpencoder args input_ids torch floattensor of shape batch_size sequence_length optional input text tokens processed from the clvptokenizer input_features torch floattensor of shape batch_size feature_size time_dim optional indicates log melspectrogram representations for audio returned by clvpfeatureextractor attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding text 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 generation_config generation generationconfig optional the generation configuration to be used as base parametrization for the generation call kwargs passed to generate matching the attributes of generation_config will override them if generation_config is not provided the default will be used which had the following loading priority 1 from the generation_config json model file if it exists 2 from the model configuration please note that unspecified parameters will inherit generation generationconfig s default values whose documentation should be checked to parameterize generation pad_to_max_mel_tokens int optional pads generated speech_ids to the specified value this is to implement the same logic from the official repo link https github com neonbjb tortoise tts blob 80f89987a5abda5e2b082618cd74f9c7411141dc tortoise api py l430 and to make sure the logits are same this does not affect generation quality so please don t consider using it since it is less efficient output_hidden_states bool optional whether or not to return the hidden states of decoder model text encoder and speech encoder models returns clvpoutput or tuple a clvpoutput if return_dict_in_generate true or when config return_dict_in_generate true or a tuple if the input sequences are larger than self config decoder_config max_text_tokens 3 then raise error because we need to add 3 tokens 1 bos tokens and 2 eos tokens to the input_ids in clvpconditioningencoder to properly sample all unused kwargs must be model kwargs pad input_ids as specified in the original repo link https github com neonbjb tortoise tts blob 80f89987a5abda5e2b082618cd74f9c7411141dc tortoise api py l380 pad to pad_to_max_mel_tokens if given to replicate the original repo logic link https github com neonbjb tortoise tts blob 80f89987a5abda5e2b082618cd74f9c7411141dc tortoise api py l430 normalized features cosine similarity as logits
import copy
import math
from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss

from ...activations import ACT2FN
from ...generation import GenerationConfig
from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
from ...modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPooling,
    CausalLMOutputWithCrossAttentions,
)
from ...modeling_utils import PreTrainedModel, SequenceSummary
from ...pytorch_utils import Conv1D
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_clvp import (
    ClvpConfig,
    ClvpDecoderConfig,
    ClvpEncoderConfig,
)


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "susnato/clvp_dev"

CLVP_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "susnato/clvp_dev",
    
]



def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
    return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))



def clvp_loss(similarity: torch.Tensor) -> torch.Tensor:
    caption_loss = contrastive_loss(similarity)
    speech_loss = contrastive_loss(similarity.t())
    return (caption_loss + speech_loss) / 2.0



def rotate_half(x):
    
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, v, cos, sin, position_ids, unsqueeze_dim=1):
    
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    v_embed = (v * cos) + (rotate_half(v) * sin)
    return q_embed, k_embed, v_embed


def _pad_extra_bos_eos_tokens(
    input_ids,
    attention_mask=None,
    pad_token_id=0,
    bos_token_id=255,
    eos_token_id=0,
    add_bos_token=True,
    add_eos_token=True,
):
    

    
    if add_bos_token:
        input_ids = torch.nn.functional.pad(input_ids, (1, 0), value=bos_token_id)
        attention_mask = (
            torch.nn.functional.pad(attention_mask, (1, 0), value=1) if attention_mask is not None else attention_mask
        )

    modified_input_ids = input_ids
    if add_eos_token:
        modified_input_ids = torch.zeros(
            (input_ids.shape[0], input_ids.shape[1] + 1), dtype=input_ids.dtype, device=input_ids.device
        )
        for i, each_input_id in enumerate(input_ids):
            
            if torch.isin(each_input_id, pad_token_id).sum():
                pos = torch.where(each_input_id == pad_token_id)[0].min()
                modified_input_ids[i] = torch.concatenate(
                    [each_input_id[:pos], torch.tensor([eos_token_id], device=input_ids.device), each_input_id[pos:]]
                )
            else:
                
                modified_input_ids[i] = torch.nn.functional.pad(each_input_id, (0, 1), value=eos_token_id)
        attention_mask = (
            torch.nn.functional.pad(attention_mask, (1, 0), value=1) if attention_mask is not None else attention_mask
        )

    return modified_input_ids, attention_mask


@dataclass
class ClvpEncoderOutput(ModelOutput):
    

    embeds: Optional[torch.FloatTensor] = None
    last_hidden_state: torch.FloatTensor = None
    pooler_output: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class ClvpOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    speech_ids: Optional[torch.LongTensor] = None
    logits_per_speech: torch.FloatTensor = None
    logits_per_text: torch.FloatTensor = None
    text_embeds: torch.FloatTensor = None
    speech_embeds: torch.FloatTensor = None
    text_model_output: BaseModelOutputWithPooling = None
    speech_model_output: BaseModelOutputWithPooling = None
    decoder_hidden_states: torch.FloatTensor = None
    text_encoder_hidden_states: torch.FloatTensor = None
    speech_encoder_hidden_states: torch.FloatTensor = None



class ClvpRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


class ClvpRotaryPositionalEmbedding(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        dim = max(config.projection_dim // (config.num_attention_heads * 2), 32)
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))

        self.register_buffer("inv_freq", inv_freq)
        self.cached_sequence_length = None
        self.cached_rotary_positional_embedding = None

    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        sequence_length = hidden_states.shape[1]

        if sequence_length == self.cached_sequence_length and self.cached_rotary_positional_embedding is not None:
            return self.cached_rotary_positional_embedding

        self.cached_sequence_length = sequence_length
        time_stamps = torch.arange(sequence_length, device=hidden_states.device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", time_stamps, self.inv_freq)
        embeddings = torch.cat((freqs, freqs), dim=-1)

        self.cached_rotary_positional_embedding = embeddings.unsqueeze(0)
        return self.cached_rotary_positional_embedding


class ClvpSelfAttention(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {self.num_heads})."
            )
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_dropout

        if hasattr(config, "max_position_embeddings"):
            max_positions = config.max_position_embeddings
            bias = torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool))
            bias = bias.view(1, 1, max_positions, max_positions)
            self.register_buffer("bias", bias, persistent=False)

        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_attention_bias)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_attention_bias)
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_attention_bias)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

    
    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.FloatTensor,
        rotary_pos_emb: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        use_cache: Optional[bool] = False,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:
        
        
        if rotary_pos_emb is not None and position_ids is None:
            raise ValueError("`position_ids` must be provided when `rotary_pos_emb` is not None.")

        bsz, _, embed_dim = hidden_states.size()

        
        query_states = self._shape(self.q_proj(hidden_states), -1, bsz) * self.scale
        key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
        value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        if past_key_value is not None:
            past_key, past_value = past_key_value
            key_states = torch.cat((past_key, key_states), dim=-2)
            value_states = torch.cat((past_value, value_states), dim=-2)

        if use_cache is True:
            present = (key_states, value_states)
        else:
            present = None

        if rotary_pos_emb is not None:
            rotary_emb_dim = rotary_pos_emb.shape[-1]

            
            query_rot, query_pass = (
                query_states[..., :rotary_emb_dim],
                query_states[..., rotary_emb_dim:],
            )
            key_rot, key_pass = (
                key_states[..., :rotary_emb_dim],
                key_states[..., rotary_emb_dim:],
            )
            value_rot, value_pass = (
                value_states[..., :rotary_emb_dim],
                value_states[..., rotary_emb_dim:],
            )

            cos, sin = rotary_pos_emb.cos().squeeze(0), rotary_pos_emb.sin().squeeze(0)
            query_rot, key_rot, value_rot = apply_rotary_pos_emb(query_rot, key_rot, value_rot, cos, sin, position_ids)

            
            query_states = torch.cat((query_rot, query_pass), dim=-1)
            key_states = torch.cat((key_rot, key_pass), dim=-1)
            value_states = torch.cat((value_rot, value_pass), dim=-1)

        tgt_len = query_states.shape[2]
        src_len = key_states.shape[2]
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3))

        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 + attention_mask

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)

        
        if head_mask is not None:
            attn_weights = attn_weights * head_mask

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
        attn_output = torch.matmul(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()}"
            )

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)

        attn_output = self.out_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, present, attn_weights


class ClvpGatedLinearUnit(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        self.activation_fn = ACT2FN[config.hidden_act]
        self.proj = nn.Linear(config.hidden_size, config.intermediate_size * 2)

    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
        return hidden_states * self.activation_fn(gate)


class ClvpEncoderMLP(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        self.config = config

        self.fc1 = ClvpGatedLinearUnit(config)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
        self.dropout_layer = nn.Dropout(config.dropout)

    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.dropout_layer(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states


class ClvpEncoderLayer(nn.Module):
    def __init__(self, config: ClvpConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.self_attn = ClvpSelfAttention(config)
        self.mlp = ClvpEncoderMLP(config)

        self.input_rmsnorm = ClvpRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.post_attention_rmsnorm = ClvpRMSNorm(self.embed_dim, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        rotary_pos_emb: torch.FloatTensor,
        attention_mask: torch.LongTensor,
        position_ids: torch.LongTensor,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor]:
        
        residual = hidden_states

        hidden_states = self.input_rmsnorm(hidden_states)

        attention_outputs = self.self_attn(
            hidden_states=hidden_states,
            rotary_pos_emb=rotary_pos_emb,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
        )

        hidden_states = attention_outputs[0]

        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_rmsnorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attention_outputs[-1],)

        return outputs



class ClvpDecoderMLP(nn.Module):
    def __init__(self, intermediate_size, config):
        super().__init__()
        embed_dim = config.hidden_size
        self.c_fc = Conv1D(intermediate_size, embed_dim)
        self.c_proj = Conv1D(embed_dim, intermediate_size)
        self.act = ACT2FN[config.activation_function]
        self.dropout = nn.Dropout(config.resid_pdrop)

    def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor:
        hidden_states = self.c_fc(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.c_proj(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


class ClvpDecoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_size = config.hidden_size
        inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size

        self.input_layernorm = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
        self.attn = ClvpSelfAttention(config)
        self.post_attention_layernorm = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)

        self.mlp = ClvpDecoderMLP(inner_dim, config)

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_outputs = self.attn(
            hidden_states,
            past_key_value=past_key_value,
            attention_mask=attention_mask,
            position_ids=position_ids,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        attn_output = attn_outputs[0]
        outputs = attn_outputs[1:]
        
        hidden_states = attn_output + residual

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        feed_forward_hidden_states = self.mlp(hidden_states)
        
        hidden_states = residual + feed_forward_hidden_states

        if use_cache:
            outputs = (hidden_states,) + outputs
        else:
            outputs = (hidden_states,) + outputs[1:]

        return outputs


class ClvpConditioningEncoder(nn.Module):
    

    def __init__(self, config: ClvpConfig):
        super().__init__()

        self.text_config = config.text_config
        self.decoder_config = config.decoder_config

        self.text_token_embedding = nn.Embedding(self.text_config.vocab_size, self.decoder_config.hidden_size)
        self.text_position_embedding = nn.Embedding(
            self.decoder_config.max_text_tokens, self.decoder_config.hidden_size
        )

        self.mel_conv = nn.Conv1d(self.decoder_config.feature_size, self.decoder_config.hidden_size, kernel_size=1)

        
        num_groups = self.compute_groupnorm_groups(self.decoder_config.hidden_size)
        self.group_norms = nn.ModuleList(
            [
                nn.GroupNorm(num_groups, self.decoder_config.hidden_size, eps=1e-5, affine=True)
                for _ in range(self.decoder_config.num_mel_attn_blocks)
            ]
        )

        
        self.mel_attn_blocks = nn.ModuleList(
            [ClvpSelfAttention(self.decoder_config) for _ in range(self.decoder_config.num_mel_attn_blocks)]
        )

        self.gradient_checkpointing = False

    def compute_groupnorm_groups(self, channels: int, groups: int = 32):
        
        if channels <= 16:
            groups = 8
        elif channels <= 64:
            groups = 16
        while channels % groups != 0:
            groups = int(groups / 2)

        if groups <= 2:
            raise ValueError(
                f"Number of groups for the GroupNorm must be greater than 2, but it is {groups}."
                f"Please consider using a different `hidden_size`"
            )

        return groups

    def forward(
        self,
        input_features: torch.FloatTensor,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
    ):
        
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            batch_size, seq_length = input_ids.size()
        elif inputs_embeds is not None:
            batch_size, seq_length = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        
        if attention_mask is None:
            attention_mask = torch.ones([batch_size, seq_length], dtype=torch.long, device=input_ids.device)

        
        
        input_ids, attention_mask = _pad_extra_bos_eos_tokens(
            input_ids,
            attention_mask,
            bos_token_id=self.text_config.bos_token_id,
            eos_token_id=self.text_config.eos_token_id,
        )

        inputs_embeds = self.text_token_embedding(input_ids)
        position_ids = attention_mask.cumsum(-1) - 1
        position_embeds = self.text_position_embedding(position_ids)
        text_embeds = inputs_embeds + position_embeds

        if self.gradient_checkpointing and self.training:
            
            mel_spec = torch.utils.checkpoint.checkpoint(self.mel_conv, input_features)

            for i, mel_attn_block in enumerate(self.mel_attn_blocks):
                residual_mel_spec = mel_spec.transpose(1, 2)

                mel_spec = torch.utils.checkpoint.checkpoint(self.group_norms[i], mel_spec).transpose(1, 2)
                mel_spec = torch.utils.checkpoint.checkpoint(mel_attn_block, mel_spec)[0] + residual_mel_spec
                mel_spec = mel_spec.transpose(1, 2)

        else:
            
            mel_spec = self.mel_conv(input_features)

            for i, mel_attn_block in enumerate(self.mel_attn_blocks):
                residual_mel_spec = mel_spec.transpose(1, 2)

                mel_spec = self.group_norms[i](mel_spec).transpose(1, 2)
                mel_spec = mel_attn_block(mel_spec)[0] + residual_mel_spec
                mel_spec = mel_spec.transpose(1, 2)

        mel_spec = mel_spec[:, :, 0]
        mel_spec = mel_spec.unsqueeze(1)

        
        if text_embeds.shape[0] == 1 and mel_spec.shape[0] != 1:
            text_embeds = text_embeds.repeat(mel_spec.shape[0], 1, 1)
        elif text_embeds.shape[0] != 1 and mel_spec.shape[0] == 1:
            mel_spec = mel_spec.repeat(text_embeds.shape[0], 1, 1)
        
        elif text_embeds.shape[0] != mel_spec.shape[0]:
            raise ValueError(
                f"The number of texts and number of audios must be same. "
                f"Found {text_embeds.shape[0]} texts vs {mel_spec.shape[0]} audios"
            )

        return torch.concat([mel_spec, text_embeds], dim=1)


class ClvpPreTrainedModel(PreTrainedModel):
    

    config_class = ClvpConfig
    base_model_prefix = "clvp"
    supports_gradient_checkpointing = True
    _skip_keys_device_placement = "past_key_values"

    def _init_weights(self, module):
        
        factor = self.config.initializer_factor
        if isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=factor * 0.02)
        elif isinstance(module, (nn.Linear, Conv1D, nn.Conv1d)):
            module.weight.data.normal_(mean=0.0, std=factor * 0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, ClvpEncoderMLP):
            factor = self.config.initializer_factor
            in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
            nn.init.normal_(module.fc1.proj.weight if getattr(module.fc1, "proj") else module.fc1.weight, std=fc_std)
            nn.init.normal_(module.fc2.weight, std=in_proj_std)
        elif isinstance(module, ClvpEncoder):
            config = self.config.text_config if hasattr(self.config, "text_config") else self.config
            factor = config.initializer_factor
            module.projection.weight.data.normal_(mean=0.0, std=factor * (config.hidden_size**-0.5))
        elif isinstance(module, ClvpConditioningEncoder):
            module.mel_conv.weight.data.normal_(mean=0.0, std=factor)
            module.mel_conv.bias.data.zero_()
        elif isinstance(module, ClvpForCausalLM):
            for name, p in module.named_parameters():
                if name == "c_proj.weight":
                    p.data.normal_(
                        mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers))
                    )
        if isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


CLVP_START_DOCSTRING = r


CLVP_INPUTS_DOCSTRING = r


CLVP_DECODER_INPUTS_DOCSTRING = r


class ClvpEncoder(ClvpPreTrainedModel):
    

    def __init__(self, config: ClvpConfig):
        super().__init__(config)

        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.rotary_pos_emb = ClvpRotaryPositionalEmbedding(config) if config.use_rotary_embedding else None
        self.layers = nn.ModuleList([ClvpEncoderLayer(config) for _ in range(config.num_hidden_layers)])

        self.sequence_summary = SequenceSummary(config)
        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

        self.projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False)

        self.gradient_checkpointing = False

        self.post_init()

    def get_input_embeddings(self):
        return self.token_embedding

    def set_input_embeddings(self, value):
        self.token_embedding = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        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 input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
            inputs_embeds = self.token_embedding(input_ids)
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

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

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(input_shape[1], dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0)

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        rotary_pos_emb = self.rotary_pos_emb(inputs_embeds) if self.rotary_pos_emb is not None else None

        hidden_states = inputs_embeds
        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            if self.gradient_checkpointing and self.training:
                layer_outputs = torch.utils.checkpoint.checkpoint(
                    encoder_layer.__call__,
                    hidden_states,
                    rotary_pos_emb,
                    attention_mask,
                    position_ids,
                )
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    rotary_pos_emb,
                    attention_mask,
                    position_ids,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        last_hidden_state = hidden_states
        last_hidden_state = self.final_layer_norm(last_hidden_state)

        
        pooled_output = self.sequence_summary(last_hidden_state)

        
        embeds = self.projection(pooled_output)

        if not return_dict:
            return tuple(
                v for v in [embeds, last_hidden_state, pooled_output, encoder_states, all_attentions] if v is not None
            )

        return ClvpEncoderOutput(
            embeds=embeds,
            last_hidden_state=last_hidden_state,
            pooler_output=pooled_output,
            hidden_states=encoder_states,
            attentions=all_attentions,
        )


class ClvpDecoder(ClvpPreTrainedModel):
    

    def __init__(self, config):
        super().__init__(config)

        self.config = config

        self.input_embeds_layer = nn.Embedding(self.config.vocab_size, self.config.hidden_size)
        self.position_embeds_layer = nn.Embedding(self.config.max_position_embeddings, self.config.hidden_size)

        self.drop = nn.Dropout(self.config.embd_pdrop)
        self.layers = nn.ModuleList([ClvpDecoderLayer(self.config) for _ in range(self.config.num_hidden_layers)])
        self.layer_norm = nn.LayerNorm(self.config.hidden_size, eps=self.config.layer_norm_epsilon)

        self.gradient_checkpointing = False

        
        self.post_init()

    def get_input_embeddings(self):
        return self.input_embeds_layer

    def set_input_embeddings(self, new_embeddings):
        self.input_embeds_layer = new_embeddings

    def _prune_heads(self, heads_to_prune):
        
        for layer, heads in heads_to_prune.items():
            self.layers[layer].attn.prune_heads(heads)

    @add_start_docstrings_to_model_forward(CLVP_DECODER_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
        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
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
            input_ids.shape[0]
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            inputs_embeds.shape[0]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])

        if past_key_values is None:
            past_key_values_length = 0
            past_key_values = tuple([None] * len(self.layers))
        else:
            past_key_values_length = past_key_values[0][0].size(-2)
        if position_ids is None:
            position_ids = torch.arange(
                past_key_values_length, input_shape[-1] + past_key_values_length, dtype=torch.long, device=device
            )
            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])

        if inputs_embeds is None:
            inputs_embeds = self.input_embeds_layer(input_ids)
        position_embeds = self.position_embeds_layer(position_ids)
        inputs_embeds = inputs_embeds + position_embeds

        attention_mask = _prepare_4d_causal_attention_mask(
            attention_mask, input_shape, inputs_embeds, past_key_values_length
        )

        
        
        
        
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        hidden_states = inputs_embeds

        if token_type_ids is not None:
            token_type_embeds = self.input_embeds_layer(token_type_ids)
            hidden_states = hidden_states + token_type_embeds

        hidden_states = self.drop(hidden_states)

        output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
        all_hidden_states = () if output_hidden_states else None
        for i, (block, past_key_value) in enumerate(zip(self.layers, past_key_values)):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                outputs = torch.utils.checkpoint.checkpoint(
                    block.__call__,
                    hidden_states,
                    None,
                    attention_mask,
                    position_ids,
                    head_mask[i],
                )
            else:
                outputs = block(
                    hidden_states,
                    past_key_value=past_key_value,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    head_mask=head_mask[i],
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                )

            hidden_states = outputs[0]
            if use_cache is True:
                presents = presents + (outputs[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)

        hidden_states = self.layer_norm(hidden_states)

        hidden_states = hidden_states.view(output_shape)

        
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions]
                if v is not None
            )

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )


@add_start_docstrings(
    "The bare Clvp decoder model outputting raw hidden-states without any specific head on top.",
    CLVP_START_DOCSTRING,
)
class ClvpModel(ClvpPreTrainedModel):
    def __init__(self, config: ClvpDecoderConfig):
        super().__init__(config)
        self.config = config
        self.decoder = ClvpDecoder(self.config)

        
        self.post_init()

    def get_input_embeddings(self):
        return self.decoder.input_embeds_layer

    def set_input_embeddings(self, value):
        self.decoder.input_embeds_layer = value

    def get_decoder(self):
        return self.decoder

    @add_start_docstrings_to_model_forward(CLVP_DECODER_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
        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
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        
        decoder_outputs = self.decoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            hidden_states=decoder_outputs.hidden_states,
            attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
        )


@add_start_docstrings(
    "The CLVP decoder model with a language modelling head on top.",
    CLVP_START_DOCSTRING,
)
class ClvpForCausalLM(ClvpPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.config = config
        self.model = ClvpModel(self.config)

        self.final_norm = nn.LayerNorm(self.config.hidden_size)
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=True)

        
        self.post_init()

    def get_input_embeddings(self):
        return self.model.decoder.input_embeds_layer

    def set_input_embeddings(self, new_embeddings):
        self.model.decoder.input_embeds_layer = new_embeddings

    def _prepare_model_inputs(
        self,
        inputs: Optional[torch.Tensor] = None,
        bos_token_id: Optional[int] = None,
        model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
    ) -> Tuple[torch.Tensor, Optional[str], Dict[str, torch.Tensor]]:
        
        input_name = self.main_input_name

        model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None}

        inputs_kwarg = model_kwargs.pop(input_name, None)
        if inputs_kwarg is not None and inputs is not None:
            raise ValueError(
                f"`inputs`: {inputs}` were passed alongside {input_name} which is not allowed."
                f"Make sure to either pass {inputs} or {input_name}=..."
            )
        elif inputs_kwarg is not None:
            inputs = inputs_kwarg

        if input_name == "input_ids" and "inputs_embeds" in model_kwargs:
            model_kwargs["input_ids"] = self._maybe_initialize_input_ids_for_generation(
                inputs, bos_token_id, model_kwargs=model_kwargs
            )
            inputs, input_name = model_kwargs["inputs_embeds"], "inputs_embeds"

        
        
        conditioning_embeds = model_kwargs.get("conditioning_embeds", None)

        if conditioning_embeds is not None:
            mel_start_token_embedding = self.model.decoder.input_embeds_layer(
                torch.full(
                    (conditioning_embeds.shape[0], 1),
                    fill_value=self.config.bos_token_id,
                    device=conditioning_embeds.device,
                )
            )
            mel_start_token_embedding += self.model.decoder.position_embeds_layer(
                torch.full((conditioning_embeds.shape[0], 1), fill_value=0, device=conditioning_embeds.device)
            )
            conditioning_embeds = torch.concat([conditioning_embeds, mel_start_token_embedding], dim=1)

            
            if hasattr(model_kwargs, "attention_mask"):
                position_ids = model_kwargs["attention_mask"].long().cumsum(-1) - 1
            else:
                position_ids = torch.range(
                    0, conditioning_embeds.shape[1] - 1, dtype=torch.long, device=conditioning_embeds.device
                )
            position_ids = position_ids.unsqueeze(0).repeat(conditioning_embeds.shape[0], 1)

            model_kwargs["inputs_embeds"] = conditioning_embeds - self.model.decoder.position_embeds_layer(
                position_ids
            )
            model_kwargs["input_ids"] = (
                torch.ones((model_kwargs["inputs_embeds"].shape[0], 1), dtype=torch.long, device=self.device)
                * self.config.bos_token_id
            )

            return model_kwargs["inputs_embeds"], "inputs_embeds", model_kwargs

        inputs = self._maybe_initialize_input_ids_for_generation(inputs, bos_token_id, model_kwargs)
        return inputs, input_name, model_kwargs

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, inputs_embeds=None, conditioning_embeds=None, **kwargs
    ):
        input_ids_length = input_ids.shape[-1]
        token_type_ids = kwargs.get("token_type_ids", None)
        
        if past_key_values:
            past_length = past_key_values[0][0].shape[2]

            
            if input_ids.shape[1] > past_length:
                remove_prefix_length = past_length
            else:
                
                remove_prefix_length = input_ids.shape[1] - 1

            input_ids = input_ids[:, remove_prefix_length:]
            if token_type_ids is not None:
                token_type_ids = token_type_ids[:, -input_ids.shape[1] :]

        attention_mask = kwargs.get("attention_mask", None)
        position_ids = kwargs.get("position_ids", None)

        if attention_mask is not None and position_ids is None:
            
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)
        else:
            position_ids = None

        if conditioning_embeds is not None and past_key_values is not None:
            position_ids = torch.tensor([input_ids_length], dtype=torch.long, device=input_ids.device)

        
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "position_ids": position_ids,
                "token_type_ids": token_type_ids,
            }
        )
        return model_inputs

    @add_start_docstrings_to_model_forward(CLVP_DECODER_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
        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
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model(
            input_ids=input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]

        lm_logits = self.final_norm(hidden_states)
        lm_logits = self.lm_head(lm_logits)

        loss = None
        if labels is not None:
            labels = labels.to(lm_logits.device)
            
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )

    @staticmethod
    
    def _reorder_cache(
        past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
    ) -> Tuple[Tuple[torch.Tensor]]:
        
        return tuple(
            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
            for layer_past in past_key_values
        )


@add_start_docstrings(
    "The composite CLVP model with a text encoder, speech encoder and speech decoder model."
    "The speech decoder model generates the speech_ids from the text and the text encoder and speech encoder works"
    "together to filter out the best speech_ids.",
    CLVP_START_DOCSTRING,
)
class ClvpModelForConditionalGeneration(ClvpPreTrainedModel):
    config_class = ClvpConfig

    def __init__(self, config: ClvpConfig):
        super().__init__(config)

        if not isinstance(config.text_config, ClvpEncoderConfig):
            raise ValueError(
                "config.text_config is expected to be of type `ClvpEncoderConfig` but is of type"
                f" {type(config.text_config)}."
            )

        if not isinstance(config.speech_config, ClvpEncoderConfig):
            raise ValueError(
                "config.speech_config is expected to be of type `ClvpEncoderConfig` but is of type"
                f" {type(config.speech_config)}."
            )

        if not isinstance(config.decoder_config, ClvpDecoderConfig):
            raise ValueError(
                "config.decoder_config is expected to be of type `ClvpDecoderConfig` but is of type"
                f" {type(config.decoder_config)}."
            )

        self.conditioning_encoder = ClvpConditioningEncoder(config)

        self.speech_decoder_model = ClvpForCausalLM(config.decoder_config)

        self.text_encoder_model = ClvpEncoder(config.text_config)
        self.speech_encoder_model = ClvpEncoder(config.speech_config)

        self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))

        
        self.post_init()

    
    
    def fix_speech_decoder_output(self, speech_ids: torch.LongTensor) -> torch.LongTensor:
        
        decoder_fixing_codes = self.config.decoder_config.decoder_fixing_codes
        speech_ids = speech_ids[:, 1:]

        stop_token_indices = torch.where(speech_ids == self.speech_decoder_model.config.eos_token_id, 1, 0)
        speech_ids = torch.masked_fill(speech_ids, mask=stop_token_indices.bool(), value=decoder_fixing_codes[0])

        for i, each_seq_stop_token_index in enumerate(stop_token_indices):
            
            
            if each_seq_stop_token_index.sum() == 0:
                continue

            stm = each_seq_stop_token_index.argmax()
            speech_ids[i, stm:] = decoder_fixing_codes[0]
            if stm - 3 < speech_ids.shape[1]:
                speech_ids[i, -3:] = torch.tensor(
                    [decoder_fixing_codes[1:]], device=speech_ids.device, dtype=torch.long
                )

        return speech_ids

    def get_text_features(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        text_encoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
    ) -> torch.FloatTensor:
        r

        outputs = self.text_encoder_model(
            input_ids=input_ids,
            inputs_embeds=text_encoder_inputs_embeds,
            attention_mask=attention_mask,
        )

        return outputs[0]

    def get_speech_features(
        self,
        speech_ids: Optional[torch.LongTensor] = None,
        input_ids: Optional[torch.LongTensor] = None,
        input_features: Optional[torch.FloatTensor] = None,
        conditioning_encoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        **kwargs,
    ) -> torch.FloatTensor:
        r

        if speech_ids is None:
            if (input_ids is None and conditioning_encoder_inputs_embeds is None) or input_features is None:
                raise ValueError(
                    "Either speech_ids or input_ids/conditioning_encoder_inputs_embeds and input_features must be provided."
                )

            if generation_config is None:
                generation_config = self.generation_config
            generation_config.update(**kwargs)

            conditioning_embeds = self.conditioning_encoder(
                input_features=input_features,
                input_ids=input_ids,
                inputs_embeds=conditioning_encoder_inputs_embeds,
                attention_mask=attention_mask,
            )

            speech_ids = self.speech_decoder_model.generate(
                conditioning_embeds=conditioning_embeds,
                generation_config=generation_config,
            )

            speech_ids = self.fix_speech_decoder_output(speech_ids[0])

        outputs = self.speech_encoder_model(
            input_ids=speech_ids,
            attention_mask=attention_mask,
        )

        return outputs[0]

    @add_start_docstrings_to_model_forward(CLVP_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=ClvpOutput, config_class=ClvpConfig)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        input_features: torch.FloatTensor = None,
        conditioning_encoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        text_encoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        return_loss: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, ClvpOutput]:
        r

        
        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

        conditioning_embeds = self.conditioning_encoder(
            input_features=input_features,
            input_ids=input_ids,
            inputs_embeds=conditioning_encoder_inputs_embeds,
            attention_mask=attention_mask,
        )

        decoder_outputs = self.speech_decoder_model(
            inputs_embeds=conditioning_embeds,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        speech_ids = decoder_outputs[0]

        
        
        if speech_ids.ndim == 3:
            speech_ids = speech_ids.argmax(2)
        speech_ids = self.fix_speech_decoder_output(speech_ids)

        speech_outputs = self.speech_encoder_model(
            input_ids=speech_ids,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        text_outputs = self.text_encoder_model(
            input_ids=input_ids,
            inputs_embeds=text_encoder_inputs_embeds,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        speech_embeds = speech_outputs[0]
        text_embeds = text_outputs[0]

        
        speech_embeds = speech_embeds / speech_embeds.norm(p=2, dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

        
        logit_scale = self.logit_scale.exp()
        logits_per_text = torch.matmul(text_embeds, speech_embeds.t()) * logit_scale
        logits_per_speech = logits_per_text.t()

        loss = None
        if return_loss:
            loss = clvp_loss(logits_per_text)

        if not return_dict:
            output = (
                logits_per_speech,
                logits_per_text,
                text_embeds,
                speech_embeds,
                text_outputs[2],
                speech_outputs[2],
            )
            if output_hidden_states:
                output += (
                    decoder_outputs[-1],
                    text_outputs[-1],
                    speech_outputs[-1],
                )

            return ((loss,) + output) if loss is not None else output

        return ClvpOutput(
            loss=loss,
            logits_per_speech=logits_per_speech,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            speech_embeds=speech_embeds,
            text_model_output=text_outputs[2],
            speech_model_output=speech_outputs[2],
            decoder_hidden_states=decoder_outputs.hidden_states,
            text_encoder_hidden_states=text_outputs.hidden_states,
            speech_encoder_hidden_states=speech_outputs.hidden_states,
        )

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.LongTensor = None,
        input_features: torch.FloatTensor = None,
        attention_mask: Optional[torch.LongTensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        pad_to_max_mel_tokens: Optional[int] = None,
        output_hidden_states: Optional[bool] = None,
        **kwargs,
    ):
        

        
        
        
        sequence_length = input_ids.shape[-1]
        if sequence_length > (self.config.decoder_config.max_text_tokens - 3):
            raise ValueError(
                f"Maximum sequence length reached! Found input_ids of length {sequence_length}."
                f"Please make sure that the maximum length of input_ids is {self.config.decoder_config.max_text_tokens - 3}"
            )

        if generation_config is None:
            generation_config = self.generation_config

        generation_config = copy.deepcopy(generation_config)
        model_kwargs = generation_config.update(**kwargs)  
        generation_config.validate()
        self._validate_model_kwargs(model_kwargs.copy())

        
        
        input_ids, attention_mask = _pad_extra_bos_eos_tokens(
            input_ids,
            attention_mask,
            add_bos_token=False,
            bos_token_id=self.config.text_config.bos_token_id,
            eos_token_id=self.config.text_config.eos_token_id,
        )

        conditioning_embeds = self.conditioning_encoder(
            input_features=input_features,
            input_ids=input_ids,
            attention_mask=attention_mask,
        )

        decoder_outputs = self.speech_decoder_model.generate(
            conditioning_embeds=conditioning_embeds,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            return_dict=generation_config.return_dict_in_generate,
        )
        if isinstance(decoder_outputs, ModelOutput):
            speech_ids = decoder_outputs.sequences

        
        
        if pad_to_max_mel_tokens is not None:
            padding_needed = pad_to_max_mel_tokens - speech_ids.shape[-1]
            speech_ids = torch.nn.functional.pad(
                speech_ids, (0, padding_needed), value=self.generation_config.eos_token_id
            )

        speech_ids = self.fix_speech_decoder_output(speech_ids)

        speech_outputs = self.speech_encoder_model(
            input_ids=speech_ids,
            output_hidden_states=output_hidden_states,
            return_dict=generation_config.return_dict_in_generate,
        )
        text_outputs = self.text_encoder_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            return_dict=generation_config.return_dict_in_generate,
        )

        speech_embeds = speech_outputs[0]
        text_embeds = text_outputs[0]

        
        speech_embeds = speech_embeds / speech_embeds.norm(p=2, dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

        
        logit_scale = self.logit_scale.exp()
        logits_per_text = torch.matmul(text_embeds, speech_embeds.t()) * logit_scale
        logits_per_speech = logits_per_text.t()

        if not generation_config.return_dict_in_generate:
            output = (
                speech_ids,
                logits_per_speech,
                logits_per_text,
                text_embeds,
                speech_embeds,
                text_outputs[2],
                speech_outputs[2],
            )
            if output_hidden_states:
                output += (
                    decoder_outputs[-1],
                    text_outputs[-1],
                    speech_outputs[-1],
                )

            return output

        return ClvpOutput(
            speech_ids=speech_ids,
            logits_per_speech=logits_per_speech,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            speech_embeds=speech_embeds,
            text_model_output=text_outputs[2],
            speech_model_output=speech_outputs[2],
            decoder_hidden_states=decoder_outputs.hidden_states,
            text_encoder_hidden_states=text_outputs.hidden_states,
            speech_encoder_hidden_states=speech_outputs.hidden_states,
        )