codingutf8 2023 meta ai and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch musicgen model import copy import inspect import math import random from dataclasses import dataclass from typing import typechecking any dict optional tuple union import torch import torch nn as nn from torch nn import crossentropyloss from activations import act2fn from generation configurationutils import generationconfig from generation logitsprocess import classifierfreeguidancelogitsprocessor logitsprocessorlist from generation stoppingcriteria import stoppingcriterialist from modelingattnmaskutils import prepare4dattentionmask prepare4dcausalattentionmask from modelingoutputs import basemodeloutput basemodeloutputwithpastandcrossattentions causallmoutputwithcrossattentions modeloutput seq2seqlmoutput from modelingutils import pretrainedmodel from utils import addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from auto configurationauto import autoconfig from auto modelingauto import automodel from configurationmusicgen import musicgenconfig musicgendecoderconfig if typechecking from generation streamers import basestreamer logger logging getloggername configfordoc musicgenconfig checkpointfordoc facebookmusicgensmall musicgenpretrainedmodelarchivelist facebookmusicgensmall see all musicgen models at https huggingface comodels filtermusicgen dataclass class musicgenunconditionalinputmodeloutput encoderoutputs tupletorch floattensor none attentionmask torch longtensor none guidancescale float none copied from transformers models encoderdecoder modelingencoderdecoder shifttokensright def shifttokensrightinputids torch tensor padtokenid int decoderstarttokenid int shiftedinputids inputids newzerosinputids shape shiftedinputids 1 inputids 1 clone if decoderstarttokenid is none raise valueerrormake sure to set the decoderstarttokenid attribute of the model s configuration shiftedinputids 0 decoderstarttokenid if padtokenid is none raise valueerrormake sure to set the padtokenid attribute of the model s configuration replace possible 100 values in labels by padtokenid shiftedinputids maskedfillshiftedinputids 100 padtokenid return shiftedinputids class musicgensinusoidalpositionalembeddingnn module in forward put the weights on the correct dtype and device of the param build sinusoidal embeddings this matches the implementation in tensor2tensor but differs slightly from the description in section 3 5 of attention is all you need zero pad create the position ids from the input token ids expand embeddings if needed copied from transformers models bart modelingbart bartattention with bartmusicgen multiheaded attention from attention is all you need paper def init self embeddim int numheads int dropout float 0 0 isdecoder bool false bias bool true iscausal bool false config optionalmusicgenconfig none super init self embeddim embeddim self numheads numheads self dropout dropout self headdim embeddim numheads self config config if self headdim numheads self embeddim raise valueerror fembeddim must be divisible by numheads got embeddim self embeddim f and numheads numheads self scaling self headdim0 5 self isdecoder isdecoder self iscausal iscausal self kproj nn linearembeddim embeddim biasbias self vproj nn linearembeddim embeddim biasbias self qproj nn linearembeddim embeddim biasbias self outproj nn linearembeddim embeddim biasbias def shapeself tensor torch tensor seqlen int bsz int return tensor viewbsz seqlen self numheads self headdim transpose1 2 contiguous def forward self hiddenstates torch tensor keyvaluestates optionaltorch tensor none pastkeyvalue optionaltupletorch tensor none attentionmask optionaltorch tensor none layerheadmask optionaltorch tensor none outputattentions bool false tupletorch tensor optionaltorch tensor optionaltupletorch tensor if keyvaluestates are provided this layer is used as a crossattention layer for the decoder get query proj get key value proj pastkeyvalue0 shape2 keyvaluestates shape1 is checking that the sequencelength of the pastkeyvalue is the same as the provided keyvaluestates to support prefix tuning reuse k v crossattentions crossattentions reuse k v selfattention selfattention if crossattention save tupletorch tensor torch tensor of all cross attention keyvaluestates further calls to crossattention layer can then reuse all crossattention keyvaluestates first if case if unidirectional selfattention decoder save tupletorch tensor torch tensor of all previous decoder keyvaluestates further calls to unidirectional selfattention can concat previous decoder keyvaluestates to current projected keyvaluestates third elif case if encoder bidirectional selfattention pastkeyvalue is always none this operation is a bit awkward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to be reshaped twice and have to be reused in the following use the embeddim from the config stored in the class rather than hiddenstate because attnoutput can be partitioned across gpus when using tensorparallelism copied from transformers models mbart modelingmbart mbartdecoderlayer forward args hiddenstates torch floattensor input to the layer of shape batch seqlen embeddim attentionmask torch floattensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values encoderhiddenstates torch floattensor cross attention input to the layer of shape batch seqlen embeddim encoderattentionmask torch floattensor encoder attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values layerheadmask torch floattensor mask for attention heads in a given layer of size encoderattentionheads crossattnlayerheadmask torch floattensor mask for crossattention heads in a given layer of size decoderattentionheads pastkeyvalue tupletorch floattensor cached past key and value projection states outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail self attention decoder unidirectional selfattention cached keyvalues tuple is at positions 1 2 add present selfattn cache to positions 1 2 of presentkeyvalue tuple crossattention block crossattn cached keyvalues tuple is at positions 3 4 of presentkeyvalue tuple add crossattn to positions 3 4 of presentkeyvalue tuple fully connected an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models musicgeninputsdocstring r args inputids torch longtensor of shape batchsize sequencelength 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 glossaryinputids attentionmask torch tensor 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 what are attention masks glossaryattentionmask decoderinputids torch longtensor of shape batchsize numcodebooks targetsequencelength optional indices of decoder input sequence tokens in the vocabulary corresponding to the sequence of audio codes indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes such as with the encodecmodel see encodecmodel encode for details what are decoder input ids glossarydecoderinputids tip warningtrue the decoderinputids will automatically be converted from shape batchsize numcodebooks targetsequencelength to batchsize numcodebooks targetsequencelength in the forward pass if you obtain audio codes from an audio encoding model such as encodecmodel ensure that the number of frames is equal to 1 and that you reshape the audio codes from frames batchsize numcodebooks targetsequencelength to batchsize numcodebooks targetsequencelength prior to passing them as decoderinputids tip decoderattentionmask torch longtensor of shape batchsize targetsequencelength optional default behavior generate a tensor that ignores pad tokens in decoderinputids causal mask will also be used by default headmask torch tensor of shape encoderlayers encoderattentionheads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoderheadmask torch tensor of shape decoderlayers decoderattentionheads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked crossattnheadmask torch tensor of shape decoderlayers decoderattentionheads optional mask to nullify selected heads of the crossattention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoderoutputs tupletupletorch floattensor optional tuple consists of lasthiddenstate optional hiddenstates optional attentions lasthiddenstate of shape batchsize sequencelength hiddensize optional is a sequence of hiddenstates at the output of the last layer of the encoder used in the crossattention of the decoder pastkeyvalues tupletupletorch floattensor optional returned when usecachetrue is passed or when config usecachetrue tuple of tupletorch floattensor of length config nlayers with each tuple having 2 tensors of shape batchsize numheads sequencelength embedsizeperhead and 2 additional tensors of shape batchsize numheads encodersequencelength embedsizeperhead contains precomputed hiddenstates key and values in the selfattention blocks and in the crossattention blocks that can be used see pastkeyvalues input to speed up sequential decoding if pastkeyvalues are used the user can optionally input only the last decoderinputids those that don t have their past key value states given to this model of shape batchsize 1 instead of all decoderinputids of shape batchsize sequencelength 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 decoderinputsembeds torch floattensor of shape batchsize targetsequencelength hiddensize optional optionally instead of passing decoderinputids you can choose to directly pass an embedded representation if pastkeyvalues is used optionally only the last decoderinputsembeds have to be input see pastkeyvalues this is useful if you want more control over how to convert decoderinputids indices into associated vectors than the model s internal embedding lookup matrix if decoderinputids and decoderinputsembeds are both unset decoderinputsembeds takes the value of inputsembeds 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 class musicgendecodermusicgenpretrainedmodel def initself config musicgendecoderconfig super initconfig self dropout config dropout self layerdrop config layerdrop self maxtargetpositions config maxpositionembeddings self dmodel config hiddensize self numcodebooks config numcodebooks self embedscale math sqrtconfig hiddensize if config scaleembedding else 1 0 embeddim config vocabsize 1 self embedtokens nn modulelist nn embeddingembeddim config hiddensize for in rangeconfig numcodebooks self embedpositions musicgensinusoidalpositionalembedding config maxpositionembeddings config hiddensize self layers nn modulelistmusicgendecoderlayerconfig for in rangeconfig numhiddenlayers self layernorm nn layernormconfig hiddensize self gradientcheckpointing false initialize weights and apply final processing self postinit def getinputembeddingsself return self embedtokens def setinputembeddingsself value self embedtokens value addstartdocstringstomodelforwardmusicgendecoderinputsdocstring def forward self inputids torch longtensor none attentionmask optionaltorch tensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch longtensor none headmask optionaltorch tensor none crossattnheadmask optionaltorch tensor none pastkeyvalues optionaltupletupletorch floattensor 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 retrieve inputids and inputsembeds if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both decoderinputids and decoderinputsembeds at the same time elif inputids is not none bsz codebooks seqlen bsz codebooks seqlen input inputids reshape1 self numcodebooks inputids shape1 bsz numcodebooks seqlen input shape inputshape bsz seqlen elif inputsembeds is not none inputshape inputsembeds size 1 input inputsembeds 1 else raise valueerroryou have to specify either decoderinputids or decoderinputsembeds pastkeyvalueslength pastkeyvalueslength pastkeyvalues00 shape2 if pastkeyvalues is not none else 0 if inputsembeds is none inputsembeds sumself embedtokenscodebookinput codebook for codebook in rangenumcodebooks attentionmask prepare4dcausalattentionmask attentionmask inputshape inputsembeds pastkeyvalueslength expand encoder attention mask if encoderhiddenstates is not none and encoderattentionmask is not none bsz seqlen bsz 1 tgtseqlen srcseqlen encoderattentionmask prepare4dattentionmask encoderattentionmask inputsembeds dtype tgtleninputshape1 embed positions positions self embedpositionsinput pastkeyvalueslength hiddenstates inputsembeds positions toinputsembeds device hiddenstates nn functional dropouthiddenstates pself dropout trainingself training if self gradientcheckpointing and self training if usecache logger warningonce usecachetrue is incompatible with gradient checkpointing setting usecachefalse usecache false decoder layers allhiddenstates if outputhiddenstates else none allselfattns if outputattentions else none allcrossattentions if outputattentions and encoderhiddenstates is not none else none nextdecodercache if usecache else none check if headmaskcrossattnheadmask has a correct number of layers specified if desired for attnmask maskname in zipheadmask crossattnheadmask headmask crossattnheadmask if attnmask is not none if attnmask size0 lenself layers raise valueerror fthe maskname should be specified for lenself layers layers but it is for f attnmask size0 for idx decoderlayer in enumerateself layers add layerdrop see https arxiv orgabs1909 11556 for description if outputhiddenstates allhiddenstates hiddenstates dropoutprobability random uniform0 1 if self training and dropoutprobability self layerdrop continue pastkeyvalue pastkeyvaluesidx if pastkeyvalues is not none else none if self gradientcheckpointing and self training layeroutputs self gradientcheckpointingfunc decoderlayer forward hiddenstates attentionmask encoderhiddenstates encoderattentionmask headmaskidx if headmask is not none else none crossattnheadmaskidx if crossattnheadmask is not none else none none outputattentions usecache else layeroutputs decoderlayer hiddenstates attentionmaskattentionmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask layerheadmaskheadmaskidx if headmask is not none else none crossattnlayerheadmask crossattnheadmaskidx if crossattnheadmask is not none else none pastkeyvaluepastkeyvalue outputattentionsoutputattentions usecacheusecache hiddenstates layeroutputs0 if usecache nextdecodercache layeroutputs3 if outputattentions else 1 if outputattentions allselfattns layeroutputs1 if encoderhiddenstates is not none allcrossattentions layeroutputs2 hiddenstates self layernormhiddenstates add hidden states from the last decoder layer if outputhiddenstates allhiddenstates hiddenstates nextcache nextdecodercache if usecache else none if not returndict return tuple v for v in hiddenstates nextcache allhiddenstates allselfattns allcrossattentions if v is not none return basemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluesnextcache hiddenstatesallhiddenstates attentionsallselfattns crossattentionsallcrossattentions addstartdocstrings the bare musicgen decoder model outputting raw hiddenstates without any specific head on top musicgenstartdocstring class musicgenmodelmusicgenpretrainedmodel def initself config musicgendecoderconfig super initconfig self decoder musicgendecoderconfig initialize weights and apply final processing self postinit def getinputembeddingsself return self decoder embedtokens def setinputembeddingsself value self decoder embedtokens value def getdecoderself return self decoder addstartdocstringstomodelforwardmusicgendecoderinputsdocstring def forward self inputids torch longtensor none attentionmask optionaltorch tensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch longtensor none headmask optionaltorch tensor none crossattnheadmask optionaltorch tensor none pastkeyvalues optionaltupletupletorch floattensor 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 encoderattentionmaskencoderattentionmask encoderhiddenstatesencoderhiddenstates headmaskheadmask crossattnheadmaskcrossattnheadmask 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 musicgen decoder model with a language modelling head on top musicgenstartdocstring class musicgenforcausallmmusicgenpretrainedmodel def initself config musicgendecoderconfig super initconfig self model musicgenmodelconfig self numcodebooks config numcodebooks self lmheads nn modulelist nn linearconfig hiddensize config vocabsize biasfalse for in rangeconfig numcodebooks initialize weights and apply final processing self postinit def getinputembeddingsself return self model decoder embedtokens def setinputembeddingsself value self model decoder embedtokens value def getoutputembeddingsself return self lmheads def setoutputembeddingsself newembeddings self lmheads newembeddings def setdecoderself decoder self model decoder decoder def getdecoderself return self model decoder addstartdocstringstomodelforwardmusicgendecoderinputsdocstring replacereturndocstringsoutputtypeseq2seqlmoutput configclassconfigfordoc def forward self inputids torch longtensor none attentionmask optionaltorch tensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch longtensor none headmask optionaltorch tensor none crossattnheadmask optionaltorch tensor none pastkeyvalues optionaltupletupletorch 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 returns bsz numcodebooks seqlen vocabsize bsz numcodebooks seqlen vocabsize apply the delay pattern mask for classifier free guidance we need to replicate the decoder args across the batch dim we ll split these before sampling build a delayed pattern mask to the inputids each codebook is offset by the previous codebook by one giving a delayed pattern mask at the start of sequence and end of sequence take the example where there are 4 codebooks and a max sequence length of 8 we have the delayed pattern mask of shape codebooks seqlen p 1 1 1 1 p p p p p 1 1 1 1 p p p p p 1 1 1 1 p p p p p 1 1 1 1 where p is the special padding token id and 1 indicates that the token is valid for prediction if we include a prompt decoder input ids the 1 positions indicate where new tokens should be predicted otherwise the mask is set to the value in the prompt p a b 1 1 p p p p p c d 1 1 p p p p p e f 1 1 p p p p p g h 1 1 where ah indicate the input prompt decoder input ids that are offset by 1 now we only override the 1 tokens in our prediction bsz numcodebooks seqlen bsz numcodebooks seqlen we only apply the mask if we have a large enough seq len otherwise we return as is fill the shifted ids with the prompt entries offset by the codebook idx mono channel loop over the codebooks onebyone leftright channels are interleaved in the generated codebooks so handle one then the other construct a pattern mask that indicates the positions of padding tokens for each codebook first fill the upper triangular part the eos padding then fill the lower triangular part the bos padding for leftright channel we need to duplicate every row of the pattern mask in an interleaved fashion find the first position to start generating this is the first place we have the 1 token and will always be in the first codebook since it has no codebook offset we have no tokens that need to be filled return entire matrix of input ids bsz numcodebooks seqlen bsz numcodebooks seqlen apply a delay pattern mask to the decoder input ids only preserving predictions where the mask is set to 1 and otherwise setting to the value detailed in the mask seqlen inputids shape1 decoderpadtokenmask decoderpadtokenmask seqlen inputids torch wheredecoderpadtokenmask 1 inputids decoderpadtokenmask return inputids torch nograd def generate self inputs optionaltorch tensor none generationconfig optionalgenerationconfig none logitsprocessor optionallogitsprocessorlist none stoppingcriteria optionalstoppingcriterialist none syncedgpus optionalbool none streamer optionalbasestreamer none kwargs 1 handle generationconfig and kwargs that might update it and validate the resulting objects 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 2 set generation parameters if not already defined logitsprocessor logitsprocessor if logitsprocessor is not none else logitsprocessorlist stoppingcriteria stoppingcriteria if stoppingcriteria is not none else stoppingcriterialist if generationconfig padtokenid is none and generationconfig eostokenid is not none if modelkwargs getattentionmask none is none logger warning the attention mask and the pad token id were not set as a consequence you may observe unexpected behavior please pass your input s attentionmask to obtain reliable results eostokenid generationconfig eostokenid if isinstanceeostokenid list eostokenid eostokenid0 logger warningfsetting padtokenid to eostokenid eostokenid for openend generation generationconfig padtokenid eostokenid 3 define model inputs inputstensor has to be defined modelinputname is defined if modelspecific keyword input is passed otherwise modelinputname is none all modelspecific keyword inputs are removed from modelkwargs inputids modelinputname modelkwargs self preparemodelinputs inputs generationconfig bostokenid modelkwargs batchsize inputids shape0 self numcodebooks 4 define other model kwargs modelkwargsoutputattentions generationconfig outputattentions modelkwargsoutputhiddenstates generationconfig outputhiddenstates modelkwargsusecache generationconfig usecache modelkwargsguidancescale generationconfig guidancescale requiresattentionmask encoderoutputs not in modelkwargs if modelkwargs getattentionmask none is none and requiresattentionmask modelkwargsattentionmask self prepareattentionmaskforgeneration inputids generationconfig padtokenid generationconfig eostokenid 5 prepare maxlength depending on other stopping criteria inputidsseqlength inputids shape1 hasdefaultmaxlength kwargs getmaxlength is none and generationconfig maxlength is not none if hasdefaultmaxlength and generationconfig maxnewtokens is none and generationconfig maxlength 20 logger warning fusing the modelagnostic default maxlength generationconfig maxlength to control the generation length recommend setting maxnewtokens to control the maximum length of the generation elif generationconfig maxnewtokens is not none if not hasdefaultmaxlength logger warning fboth maxnewtokens generationconfig maxnewtokens and maxlength fgenerationconfig maxlength seem to have been set maxnewtokens will take precedence please refer to the documentation for more information https huggingface codocstransformersmainenmainclassestextgeneration generationconfig maxlength generationconfig maxnewtokens inputidsseqlength if generationconfig minlength is not none and generationconfig minlength generationconfig maxlength raise valueerror funfeasible length constraints the minimum length generationconfig minlength is larger than f the maximum length generationconfig maxlength if inputidsseqlength generationconfig maxlength logger warning finput length of decoderinputids is inputidsseqlength but maxlength is set to f generationconfig maxlength this can lead to unexpected behavior you should consider increasing maxnewtokens 6 prepare inputids which will be used for autoregressive generation build the delay pattern mask for offsetting each codebook prediction by 1 this behaviour is specific to musicgen inputids delaypatternmask self builddelaypatternmask inputids padtokenidgenerationconfig decoderstarttokenid maxlengthgenerationconfig maxlength if streamer is not none streamer putinputids cpu stash the delay mask so that we don t have to recompute it in each forward pass modelkwargsdelaypatternmask delaypatternmask 7 determine generation mode isgreedygenmode generationconfig numbeams 1 and generationconfig numbeamgroups 1 and generationconfig dosample is false issamplegenmode generationconfig numbeams 1 and generationconfig numbeamgroups 1 and generationconfig dosample is true 8 prepare batched cfg externally to enable coexistance with the unbatched cfg if generationconfig guidancescale is not none and generationconfig guidancescale 1 logitsprocessor appendclassifierfreeguidancelogitsprocessorgenerationconfig guidancescale generationconfig guidancescale none 9 prepare distribution preprocessing samplers logitsprocessor self getlogitsprocessor generationconfiggenerationconfig inputidsseqlengthinputidsseqlength encoderinputidsinputids prefixallowedtokensfnnone logitsprocessorlogitsprocessor 10 prepare stopping criteria stoppingcriteria self getstoppingcriteria generationconfiggenerationconfig stoppingcriteriastoppingcriteria if isgreedygenmode if generationconfig numreturnsequences 1 raise valueerror numreturnsequences has to be 1 when doing greedy search fbut is generationconfig numreturnsequences 11 run greedy search outputs self greedysearch inputids logitsprocessorlogitsprocessor stoppingcriteriastoppingcriteria padtokenidgenerationconfig padtokenid eostokenidgenerationconfig eostokenid outputscoresgenerationconfig outputscores returndictingenerategenerationconfig returndictingenerate syncedgpussyncedgpus streamerstreamer modelkwargs elif issamplegenmode 11 prepare logits warper logitswarper self getlogitswarpergenerationconfig expand inputids with numreturnsequences additional sequences per batch inputids modelkwargs self expandinputsforgeneration inputidsinputids expandsizegenerationconfig numreturnsequences modelkwargs 12 run sample outputs self sample inputids logitsprocessorlogitsprocessor logitswarperlogitswarper stoppingcriteriastoppingcriteria padtokenidgenerationconfig padtokenid eostokenidgenerationconfig eostokenid outputscoresgenerationconfig outputscores returndictingenerategenerationconfig returndictingenerate syncedgpussyncedgpus streamerstreamer modelkwargs else raise valueerror got incompatible mode for generation should be one of greedy or sampling ensure that beam search is deactivated by setting numbeams1 and numbeamgroups1 if generationconfig returndictingenerate outputids outputs sequences else outputids outputs apply the pattern mask to the final ids outputids self applydelaypatternmaskoutputids modelkwargsdelaypatternmask revert the pattern delay mask by filtering the pad token id outputids outputidsoutputids generationconfig padtokenid reshape batchsize self numcodebooks 1 if generationconfig returndictingenerate outputs sequences outputids return outputs else return outputids addstartdocstrings the composite musicgen model with a text encoder audio encoder and musicgen decoder for music generation tasks with one or both of text and audio prompts musicgenstartdocstring class musicgenforconditionalgenerationpretrainedmodel configclass musicgenconfig basemodelprefix encoderdecoder maininputname inputids supportsgradientcheckpointing true def init self config optionalmusicgenconfig none textencoder optionalpretrainedmodel none audioencoder optionalpretrainedmodel none decoder optionalmusicgenforcausallm none if config is none and textencoder is none or audioencoder is none or decoder is none raise valueerror either a configuration has to be provided or all three of text encoder audio encoder and musicgen decoder if config is none config musicgenconfig fromsubmodelsconfigtextencoder config audioencoder config decoder config else if not isinstanceconfig self configclass raise valueerrorfconfig config has to be of type self configclass if config decoder crossattentionhiddensize is not none if config decoder crossattentionhiddensize config textencoder hiddensize raise valueerror if crossattentionhiddensize is specified in the musicgen decoder s configuration it has to be equal f to the text encoder s hiddensize got config decoder crossattentionhiddensize for f config decoder crossattentionhiddensize and config textencoder hiddensize for config textencoder hiddensize initialize with config super initconfig if textencoder is none from auto modelingauto import automodelfortextencoding textencoder automodelfortextencoding fromconfigconfig textencoder if audioencoder is none from auto modelingauto import automodel audioencoder automodel fromconfigconfig audioencoder if decoder is none decoder musicgenforcausallmconfig decoder self textencoder textencoder self audioencoder audioencoder self decoder decoder if self textencoder config todict self config textencoder todict logger warning fconfig of the textencoder self textencoder class is overwritten by shared textencoder config f self config textencoder if self audioencoder config todict self config audioencoder todict logger warning fconfig of the audioencoder self audioencoder class is overwritten by shared audioencoder config f self config audioencoder if self decoder config todict self config decoder todict logger warning fconfig of the decoder self decoder class is overwritten by shared decoder config f self config decoder make sure that the individual model s config refers to the shared config so that the updates to the config will be synced self textencoder config self config textencoder self audioencoder config self config audioencoder self decoder config self config decoder text encoder outputs might need to be projected to different dimension for decoder if self textencoder config hiddensize self decoder config hiddensize and self decoder config crossattentionhiddensize is none self enctodecproj nn linearself textencoder config hiddensize self decoder config hiddensize if self textencoder getoutputembeddings is not none raise valueerror fthe encoder self textencoder should not have a lm head please use a model without and lm head decodersignature setinspect signatureself decoder forward parameters keys if encoderhiddenstates not in decodersignature raise valueerror the selected decoder is not prepared for the encoder hidden states to be passed please see the following discussion on github https github comhuggingfacetransformersissues23350 tie text encoder decoder weights if config set accordingly self tieweights def tieweightsself tie text encoder decoder if needed if self config tieencoderdecoder tie text encoder and decoder base model decoderbasemodelprefix self decoder basemodelprefix self tieencoderdecoderweights self textencoder self decoder modulesdecoderbasemodelprefix self decoder basemodelprefix def getaudioencoderself return self audioencoder def gettextencoderself return self textencoder def getencoderself get the text encoder to compute the encoder hiddenstates for generation return self gettextencoder def getdecoderself return self decoder def getinputembeddingsself return self textencoder getinputembeddings def getoutputembeddingsself return self decoder getoutputembeddings def setoutputembeddingsself newembeddings return self decoder setoutputembeddingsnewembeddings classmethod def frompretrainedcls pretrainedmodelnameorpath modelargs kwargs r example python from transformers import musicgenforconditionalgeneration model musicgenforconditionalgeneration frompretrainedfacebookmusicgensmall at the moment fast initialization is not supported for composite models if kwargs getfastinit false logger warning fast initialization is currently not supported for musicgenforconditionalgeneration falling back to slow initialization kwargsfastinit false return super frompretrainedpretrainedmodelnameorpath modelargs kwargs classmethod def fromsubmodelspretrained cls textencoderpretrainedmodelnameorpath str none audioencoderpretrainedmodelnameorpath str none decoderpretrainedmodelnameorpath str none modelargs kwargs pretrainedmodel r instantiate a text encoder an audio encoder and a musicgen decoder from one two or three base classes of the library from pretrained model checkpoints the model is set in evaluation mode by default using model eval dropout modules are deactivated to train the model you need to first set it back in training mode with model train params textencoderpretrainedmodelnameorpath str optional information necessary to initiate the text encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the rootlevel like t5base or namespaced under a user or organization name like googleflant5base a path to a directory containing model weights saved using pretrainedmodel savepretrained e g mymodeldirectory audioencoderpretrainedmodelnameorpath str optional information necessary to initiate the audio encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the rootlevel like bertbaseuncased or namespaced under a user or organization name like facebookencodec24khz a path to a directory containing model weights saved using pretrainedmodel savepretrained e g mymodeldirectory decoderpretrainedmodelnameorpath str optional defaults to none information necessary to initiate the decoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the rootlevel like gpt2 or namespaced under a user or organization name like facebookmusicgensmall a path to a directory containing model weights saved using pretrainedmodel savepretrained e g mymodeldirectory modelargs remaining positional arguments optional all remaining positional arguments will be passed to the underlying model s init method kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g outputattentionstrue to update the text encoder configuration use the prefix textencoder for each configuration parameter to update the audio encoder configuration use the prefix audioencoder for each configuration parameter to update the decoder configuration use the prefix decoder for each configuration parameter to update the parent model configuration do not use a prefix for each configuration parameter behaves differently depending on whether a config is provided or automatically loaded example python from transformers import musicgenforconditionalgeneration initialize a musicgen model from a t5 text encoder encodec audio encoder and musicgen decoder model musicgenforconditionalgeneration fromsubmodelspretrained textencoderpretrainedmodelnameorpatht5base audioencoderpretrainedmodelnameorpathfacebookencodec24khz decoderpretrainedmodelnameorpathfacebookmusicgensmall saving model after finetuning model savepretrained musicgenft load finetuned model model musicgenforconditionalgeneration frompretrained musicgenft kwargstextencoder argumentlentextencoder value for argument value in kwargs items if argument startswithtextencoder kwargsaudioencoder argumentlenaudioencoder value for argument value in kwargs items if argument startswithaudioencoder kwargsdecoder argumentlendecoder value for argument value in kwargs items if argument startswithdecoder remove text encoder audio encoder and decoder kwargs from kwargs for key in kwargstextencoder keys del kwargstextencoder key for key in kwargsaudioencoder keys del kwargsaudioencoder key for key in kwargsdecoder keys del kwargsdecoder key load and initialize the encoder and decoder the distinction between encoder and decoder at the model level is made by the value of the flag isdecoder that we need to set correctly textencoder kwargstextencoder popmodel none if textencoder is none if textencoderpretrainedmodelnameorpath is none raise valueerror if textencodermodel is not defined as an argument a textencoderpretrainedmodelnameorpath has to be defined if config not in kwargstextencoder encoderconfig kwargstextencoder autoconfig frompretrained textencoderpretrainedmodelnameorpath kwargstextencoder returnunusedkwargstrue if encoderconfig isdecoder is true or encoderconfig addcrossattention is true logger info finitializing textencoderpretrainedmodelnameorpath as a textencoder model from a decoder model crossattention and casual mask are disabled encoderconfig isdecoder false encoderconfig addcrossattention false kwargstextencoderconfig encoderconfig textencoder automodel frompretrained textencoderpretrainedmodelnameorpath modelargs kwargstextencoder audioencoder kwargsaudioencoder popmodel none if audioencoder is none if audioencoderpretrainedmodelnameorpath is none raise valueerror if audioencodermodel is not defined as an argument an audioencoderpretrainedmodelnameorpath has to be defined if config not in kwargsaudioencoder encoderconfig kwargsaudioencoder autoconfig frompretrained audioencoderpretrainedmodelnameorpath kwargsaudioencoder returnunusedkwargstrue if encoderconfig isdecoder is true or encoderconfig addcrossattention is true logger info finitializing audioencoderpretrainedmodelnameorpath as an audioencoder model from a decoder model crossattention and casual mask are disabled encoderconfig isdecoder false encoderconfig addcrossattention false kwargsaudioencoderconfig encoderconfig audioencoder automodel frompretrained audioencoderpretrainedmodelnameorpath modelargs kwargsaudioencoder decoder kwargsdecoder popmodel none if decoder is none if decoderpretrainedmodelnameorpath is none raise valueerror if decodermodel is not defined as an argument a decoderpretrainedmodelnameorpath has to be defined if config not in kwargsdecoder decoderconfig kwargsdecoder autoconfig frompretrained decoderpretrainedmodelnameorpath kwargsdecoder returnunusedkwargstrue if isinstancedecoderconfig musicgenconfig decoderconfig decoderconfig decoder if decoderconfig isdecoder is false or decoderconfig addcrossattention is false logger info finitializing decoderpretrainedmodelnameorpath as a decoder model cross attention f layers are added to decoderpretrainedmodelnameorpath and randomly initialized if f decoderpretrainedmodelnameorpath s architecture allows for cross attention layers decoderconfig isdecoder true decoderconfig addcrossattention true kwargsdecoderconfig decoderconfig if kwargsdecoderconfig isdecoder is false or kwargsdecoderconfig addcrossattention is false logger warning fdecoder model decoderpretrainedmodelnameorpath is not initialized as a decoder fin order to initialize decoderpretrainedmodelnameorpath as a decoder make sure that the attributes isdecoder and addcrossattention of decoderconfig passed to fromsubmodelspretrained are set to true or do not pass a decoderconfig to fromsubmodelspretrained decoder musicgenforcausallm frompretraineddecoderpretrainedmodelnameorpath kwargsdecoder instantiate config with corresponding kwargs config musicgenconfig fromsubmodelsconfig textencoder config audioencoder config decoder config kwargs return clstextencodertextencoder audioencoderaudioencoder decoderdecoder configconfig addstartdocstringstomodelforwardmusicgeninputsdocstring replacereturndocstringsoutputtypeseq2seqlmoutput configclassconfigfordoc def forward self inputids optionaltorch longtensor none attentionmask optionaltorch booltensor none inputvalues optionaltorch floattensor none paddingmask optionaltorch booltensor none decoderinputids optionaltorch longtensor none decoderattentionmask optionaltorch booltensor none encoderoutputs optionaltupletorch floattensor none pastkeyvalues tupletupletorch floattensor none inputsembeds optionaltorch floattensor none decoderinputsembeds optionaltorch floattensor none labels optionaltorch longtensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none kwargs uniontuple seq2seqlmoutput r returns examples python from transformers import autoprocessor musicgenforconditionalgeneration import torch processor autoprocessor frompretrainedfacebookmusicgensmall model musicgenforconditionalgeneration frompretrainedfacebookmusicgensmall inputs processor text80s pop track with bassy drums and synth 90s rock song with loud guitars and heavy drums paddingtrue returntensorspt padtokenid model generationconfig padtokenid decoderinputids torch onesinputs inputids shape0 model decoder numcodebooks 1 dtypetorch long padtokenid logits modelinputs decoderinputidsdecoderinputids logits logits shape bsz numcodebooks tgtlen vocabsize torch size8 1 2048 returndict returndict if returndict is not none else self config usereturndict kwargstextencoder argumentlentextencoder value for argument value in kwargs items if argument startswithtextencoder kwargsaudioencoder argumentlenaudioencoder value for argument value in kwargs items if argument startswithaudioencoder kwargsdecoder argumentlendecoder value for argument value in kwargs items if argument startswithdecoder if encoderoutputs is none encoderoutputs self textencoder inputidsinputids attentionmaskattentionmask inputsembedsinputsembeds outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict kwargstextencoder elif isinstanceencoderoutputs tuple encoderoutputs basemodeloutputencoderoutputs encoderhiddenstates encoderoutputs0 optionally project encoderhiddenstates if self textencoder config hiddensize self decoder config hiddensize and self decoder config crossattentionhiddensize is none encoderhiddenstates self enctodecprojencoderhiddenstates if attentionmask is not none encoderhiddenstates encoderhiddenstates attentionmask none if labels is not none and decoderinputids is none and decoderinputsembeds is none decoderinputids shifttokensright labels self config padtokenid self config decoderstarttokenid elif decoderinputids is none and decoderinputsembeds is none audioencoderoutputs self audioencoder inputvaluesinputvalues paddingmaskpaddingmask kwargsaudioencoder audiocodes audioencoderoutputs audiocodes frames bsz codebooks seqlen audiocodes shape if frames 1 raise valueerror fexpected 1 frame in the audio code outputs got frames frames ensure chunking is disabled by setting chunklengthnone in the audio encoder if self config decoder audiochannels 2 and audiocodes shape2 self decoder numcodebooks 2 mono input through encodec that we convert to stereo audiocodes audiocodes repeatinterleave2 dim2 decoderinputids audiocodes0 reshapebsz self decoder numcodebooks seqlen decode decoderoutputs self decoder inputidsdecoderinputids attentionmaskdecoderattentionmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskattentionmask inputsembedsdecoderinputsembeds outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates usecacheusecache pastkeyvaluespastkeyvalues returndictreturndict kwargsdecoder loss none if labels is not none logits decoderoutputs logits if returndict else decoderoutputs0 lossfct crossentropyloss loss lossfctlogits view1 self config vocabsize labels view1 if not returndict if loss is not none return loss decoderoutputs encoderoutputs else return decoderoutputs encoderoutputs return seq2seqlmoutput lossloss logitsdecoderoutputs logits pastkeyvaluesdecoderoutputs pastkeyvalues decoderhiddenstatesdecoderoutputs hiddenstates decoderattentionsdecoderoutputs attentions crossattentionsdecoderoutputs crossattentions encoderlasthiddenstateencoderoutputs lasthiddenstate encoderhiddenstatesencoderoutputs hiddenstates encoderattentionsencoderoutputs attentions def prepareinputsforgeneration self decoderinputids pastkeyvaluesnone attentionmasknone headmasknone decoderattentionmasknone decoderheadmasknone crossattnheadmasknone usecachenone encoderoutputsnone decoderdelaypatternmasknone guidancescalenone kwargs if decoderdelaypatternmask is none decoderinputids decoderdelaypatternmask self decoder builddelaypatternmask decoderinputids self generationconfig padtokenid maxlengthself generationconfig maxlength apply the delay pattern mask decoderinputids self decoder applydelaypatternmaskdecoderinputids decoderdelaypatternmask if guidancescale is not none and guidancescale 1 for classifier free guidance we need to replicate the decoder args across the batch dim we ll split these before sampling decoderinputids decoderinputids repeat2 1 if decoderattentionmask is not none decoderattentionmask decoderattentionmask repeat2 1 if pastkeyvalues is not none pastlength pastkeyvalues00 shape2 some generation methods already pass only the last input id if decoderinputids shape1 pastlength removeprefixlength pastlength else default to old behavior keep only final id removeprefixlength decoderinputids shape1 1 decoderinputids decoderinputids removeprefixlength return inputids none encoderoutputs is defined inputids not needed encoderoutputs encoderoutputs pastkeyvalues pastkeyvalues decoderinputids decoderinputids attentionmask attentionmask decoderattentionmask decoderattentionmask headmask headmask decoderheadmask decoderheadmask crossattnheadmask crossattnheadmask usecache usecache def preparedecoderinputidsforgeneration self batchsize int modelinputname str modelkwargs dictstr torch tensor decoderstarttokenid int none bostokenid int none device torch device none tupletorch longtensor dictstr torch tensor 1 check whether the user has defined decoderinputids manually to facilitate in terms of input naming we also allow the user to pass it under inputids if the encoder does not use it as the main input 2 encoderdecoder models expect the decoderinputids to start with a special token let s ensure that no user input use decoderstarttokenid as decoderinputids user input but doesn t start with decoderstarttokenid prepend decoderstarttokenid and adjust decoderattentionmask if provided 1 get text encoder compatibility with accelerate big model inference we need the encoder to outputs stuff on the same device as the inputs 2 prepare encoder args and encoder kwargs from model kwargs 3 make sure that encoder returns modeloutput for classifier free guidance we need to add a null input to our encoder hidden states 1 get audio encoder compatibility with accelerate big model inference we need the encoder to outputs stuff on the same device as the inputs 2 prepare encoder args and encoder kwargs from model kwargs 3 make sure that encoder returns modeloutput copy alternating leftright channel codes into stereo codebook initializes input ids for generation if necessary if inputs is not none return inputs encoderoutputs modelkwargs getencoderoutputs if encoderoutputs is not none make dummy inputids with value 100 as a sanity check ensuring that they won t be used for encoding shape encoderoutputs0 size 1 return torch onesshape dtypetorch long deviceself device 100 if bostokenid is none raise valueerrorbostokenid has to be defined when no inputids are provided if there is some tensor in modelkwargs we can infer the batch size from it this is helpful with softprompting or in multimodal implementations built on top of decoderonly language models batchsize 1 for value in modelkwargs values if isinstancevalue torch tensor batchsize value shape0 break return torch onesbatchsize 1 dtypetorch long deviceself device bostokenid torch nograd def generate self inputs optionaltorch tensor none generationconfig optionalgenerationconfig none logitsprocessor optionallogitsprocessorlist none stoppingcriteria optionalstoppingcriterialist none syncedgpus optionalbool none streamer optionalbasestreamer none kwargs 1 handle generationconfig and kwargs that might update it and validate the resulting objects 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 if modelkwargs getencoderoutputs is not none and typemodelkwargsencoderoutputs tuple wrap the unconditional outputs as a basemodeloutput for compatibility with the rest of generate modelkwargsencoderoutputs basemodeloutputlasthiddenstatemodelkwargsencoderoutputs0 2 set generation parameters if not already defined logitsprocessor logitsprocessor if logitsprocessor is not none else logitsprocessorlist stoppingcriteria stoppingcriteria if stoppingcriteria is not none else stoppingcriterialist if generationconfig padtokenid is none and generationconfig eostokenid is not none if modelkwargs getattentionmask none is none logger warning the attention mask and the pad token id were not set as a consequence you may observe unexpected behavior please pass your input s attentionmask to obtain reliable results eostokenid generationconfig eostokenid if isinstanceeostokenid list eostokenid eostokenid0 logger warningfsetting padtokenid to eostokenid eostokenid for openend generation generationconfig padtokenid eostokenid 3 define model inputs inputstensor has to be defined modelinputname is defined if modelspecific keyword input is passed otherwise modelinputname is none all modelspecific keyword inputs are removed from modelkwargs inputstensor modelinputname modelkwargs self preparemodelinputs inputs generationconfig bostokenid modelkwargs batchsize inputstensor shape0 4 define other model kwargs modelkwargsoutputattentions generationconfig outputattentions modelkwargsoutputhiddenstates generationconfig outputhiddenstates modelkwargsusecache generationconfig usecache modelkwargsguidancescale generationconfig guidancescale requiresattentionmask encoderoutputs not in modelkwargs if modelkwargs getattentionmask none is none and requiresattentionmask modelkwargsattentionmask self prepareattentionmaskforgeneration inputstensor generationconfig padtokenid generationconfig eostokenid if encoderoutputs not in modelkwargs encoderoutputs are created and added to modelkwargs modelkwargs self preparetextencoderkwargsforgeneration inputstensor modelkwargs modelinputname guidancescalegenerationconfig guidancescale if decoderinputids not in modelkwargs and inputvalues in modelkwargs modelkwargs self prepareaudioencoderkwargsforgeneration modelkwargsinputvalues modelkwargs 5 prepare inputids which will be used for autoregressive generation inputids modelkwargs self preparedecoderinputidsforgeneration batchsizebatchsize modelinputnamemodelinputname modelkwargsmodelkwargs decoderstarttokenidgenerationconfig decoderstarttokenid bostokenidgenerationconfig bostokenid deviceinputstensor device 6 prepare maxlength depending on other stopping criteria inputidsseqlength inputids shape1 hasdefaultmaxlength kwargs getmaxlength is none and generationconfig maxlength is not none if hasdefaultmaxlength and generationconfig maxnewtokens is none logger warning fusing the modelagnostic default maxlength generationconfig maxlength to control the generation length we recommend setting maxnewtokens to control the maximum length of the generation elif generationconfig maxnewtokens is not none if not hasdefaultmaxlength logger warning fboth maxnewtokens generationconfig maxnewtokens and maxlength fgenerationconfig maxlength seem to have been set maxnewtokens will take precedence please refer to the documentation for more information https huggingface codocstransformersmainenmainclassestextgeneration generationconfig maxlength generationconfig maxnewtokens inputidsseqlength if generationconfig minlength is not none and generationconfig minlength generationconfig maxlength raise valueerror funfeasible length constraints the minimum length generationconfig minlength is larger than f the maximum length generationconfig maxlength if inputidsseqlength generationconfig maxlength logger warning finput length of decoderinputids is inputidsseqlength but maxlength is set to f generationconfig maxlength this can lead to unexpected behavior you should consider increasing maxnewtokens build the delay pattern mask for offsetting each codebook prediction by 1 this behaviour is specific to musicgen inputids decoderdelaypatternmask self decoder builddelaypatternmask inputids padtokenidgenerationconfig decoderstarttokenid maxlengthgenerationconfig maxlength stash the delay mask so that we don t have to recompute in each forward pass modelkwargsdecoderdelaypatternmask decoderdelaypatternmask inputids are ready to be placed on the streamer if used if streamer is not none streamer putinputids cpu 7 determine generation mode isgreedygenmode generationconfig numbeams 1 and generationconfig numbeamgroups 1 and generationconfig dosample is false issamplegenmode generationconfig numbeams 1 and generationconfig numbeamgroups 1 and generationconfig dosample is true 8 prepare batched cfg externally to enable coexistance with the unbatched cfg if generationconfig guidancescale is not none and generationconfig guidancescale 1 logitsprocessor appendclassifierfreeguidancelogitsprocessorgenerationconfig guidancescale generationconfig guidancescale none 9 prepare distribution preprocessing samplers logitsprocessor self getlogitsprocessor generationconfiggenerationconfig inputidsseqlengthinputidsseqlength encoderinputidsinputstensor prefixallowedtokensfnnone logitsprocessorlogitsprocessor 10 prepare stopping criteria stoppingcriteria self getstoppingcriteria generationconfiggenerationconfig stoppingcriteriastoppingcriteria if isgreedygenmode if generationconfig numreturnsequences 1 raise valueerror numreturnsequences has to be 1 when doing greedy search fbut is generationconfig numreturnsequences 11 run greedy search outputs self greedysearch inputids logitsprocessorlogitsprocessor stoppingcriteriastoppingcriteria padtokenidgenerationconfig padtokenid eostokenidgenerationconfig eostokenid outputscoresgenerationconfig outputscores returndictingenerategenerationconfig returndictingenerate syncedgpussyncedgpus streamerstreamer modelkwargs elif issamplegenmode 11 prepare logits warper logitswarper self getlogitswarpergenerationconfig expand inputids with numreturnsequences additional sequences per batch inputids modelkwargs self expandinputsforgeneration inputidsinputids expandsizegenerationconfig numreturnsequences isencoderdecoderself config isencoderdecoder modelkwargs 12 run sample outputs self sample inputids logitsprocessorlogitsprocessor logitswarperlogitswarper stoppingcriteriastoppingcriteria padtokenidgenerationconfig padtokenid eostokenidgenerationconfig eostokenid outputscoresgenerationconfig outputscores returndictingenerategenerationconfig returndictingenerate syncedgpussyncedgpus streamerstreamer modelkwargs else raise valueerror got incompatible mode for generation should be one of greedy or sampling ensure that beam search is deactivated by setting numbeams1 and numbeamgroups1 if generationconfig returndictingenerate outputids outputs sequences else outputids outputs apply the pattern mask to the final ids outputids self decoder applydelaypatternmaskoutputids modelkwargsdecoderdelaypatternmask revert the pattern delay mask by filtering the pad token id outputids outputidsoutputids generationconfig padtokenid reshape batchsize self decoder numcodebooks 1 append the frame dimension back to the audio codes outputids outputidsnone audioscales modelkwargs getaudioscales if audioscales is none audioscales none batchsize if self decoder config audiochannels 1 outputvalues self audioencoder decode outputids audioscalesaudioscales audiovalues else codecoutputsleft self audioencoder decodeoutputids 2 audioscalesaudioscales outputvaluesleft codecoutputsleft audiovalues codecoutputsright self audioencoder decodeoutputids 1 2 audioscalesaudioscales outputvaluesright codecoutputsright audiovalues outputvalues torch catoutputvaluesleft outputvaluesright dim1 if generationconfig returndictingenerate outputs sequences outputvalues return outputs else return outputvalues def getunconditionalinputsself numsamples1 coding utf 8 2023 meta ai and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch musicgen model see all musicgen models at https huggingface co models filter musicgen args encoder_outputs tuple torch floattensor of length 1 with tensor shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the text encoder model attention_mask torch longtensor of shape batch_size sequence_length optional encoder attention 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 guidance_scale float optional guidance scale for classifier free guidance setting the balance between the conditional logits predicted from the prompts and the unconditional logits predicted without prompts copied from transformers models encoder_decoder modeling_encoder_decoder shift_tokens_right shift input ids one token to the right replace possible 100 values in labels by pad_token_id this module produces sinusoidal positional embeddings of any length in forward put the weights on the correct dtype and device of the param build sinusoidal embeddings this matches the implementation in tensor2tensor but differs slightly from the description in section 3 5 of attention is all you need zero pad create the position ids from the input token ids expand embeddings if needed copied from transformers models bart modeling_bart bartattention with bart musicgen multi headed attention from attention is all you need paper input shape batch x time x channel if key_value_states are provided this layer is used as a cross attention layer for the decoder get query proj get key value proj past_key_value 0 shape 2 key_value_states shape 1 is checking that the sequence_length of the past_key_value is the same as the provided key_value_states to support prefix tuning reuse k v cross_attentions cross_attentions reuse k v self_attention self_attention if cross_attention save tuple torch tensor torch tensor of all cross attention key value_states further calls to cross_attention layer can then reuse all cross attention key value_states first if case if uni directional self attention decoder save tuple torch tensor torch tensor of all previous decoder key value_states further calls to uni directional self attention can concat previous decoder key value_states to current projected key value_states third elif case if encoder bi directional self attention past_key_value is always none this operation is a bit awkward but it s required to make sure that attn_weights keeps its gradient in order to do so attn_weights have to be reshaped twice and have to be reused in the following use the embed_dim from the config stored in the class rather than hidden_state because attn_output can be partitioned across gpus when using tensor parallelism copied from transformers models mbart modeling_mbart mbartdecoderlayer forward args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values encoder_hidden_states torch floattensor cross attention input to the layer of shape batch seq_len embed_dim encoder_attention_mask torch floattensor encoder attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values layer_head_mask torch floattensor mask for attention heads in a given layer of size encoder_attention_heads cross_attn_layer_head_mask torch floattensor mask for cross attention heads in a given layer of size decoder_attention_heads past_key_value tuple torch floattensor cached past key and value projection states output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail self attention decoder uni directional self attention cached key values tuple is at positions 1 2 add present self attn cache to positions 1 2 of present_key_value tuple cross attention block cross_attn cached key values tuple is at positions 3 4 of present_key_value tuple add cross attn to positions 3 4 of present_key_value tuple fully connected an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models the musicgen model was proposed in simple and controllable music generation https arxiv org abs 2306 05284 by jade copet felix kreuk itai gat tal remez david kant gabriel synnaeve yossi adi alexandre défossez it is an encoder decoder transformer trained on the task of conditional music generation 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 musicgenconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing 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 decoder_input_ids torch longtensor of shape batch_size num_codebooks target_sequence_length optional indices of decoder input sequence tokens in the vocabulary corresponding to the sequence of audio codes indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes such as with the encodecmodel see encodecmodel encode for details what are decoder input ids glossary decoder input ids tip warning true the decoder_input_ids will automatically be converted from shape batch_size num_codebooks target_sequence_length to batch_size num_codebooks target_sequence_length in the forward pass if you obtain audio codes from an audio encoding model such as encodecmodel ensure that the number of frames is equal to 1 and that you reshape the audio codes from frames batch_size num_codebooks target_sequence_length to batch_size num_codebooks target_sequence_length prior to passing them as decoder_input_ids tip decoder_attention_mask torch longtensor of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default head_mask torch tensor of shape encoder_layers encoder_attention_heads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoder_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoder_outputs tuple tuple torch floattensor optional tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder past_key_values tuple tuple torch floattensor optional returned when use_cache true is passed or when config use_cache true tuple of tuple torch floattensor of length config n_layers with each tuple having 2 tensors of shape batch_size num_heads sequence_length embed_size_per_head and 2 additional tensors of shape batch_size num_heads encoder_sequence_length embed_size_per_head contains pre computed hidden states key and values in the self attention blocks and in the cross attention blocks that can be used see past_key_values input to speed up sequential decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length 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 decoder_inputs_embeds torch floattensor of shape batch_size target_sequence_length hidden_size optional optionally instead of passing decoder_input_ids you can choose to directly pass an embedded representation if past_key_values is used optionally only the last decoder_inputs_embeds have to be input see past_key_values this is useful if you want more control over how to convert decoder_input_ids indices into associated vectors than the model s internal embedding lookup matrix if decoder_input_ids and decoder_inputs_embeds are both unset decoder_inputs_embeds takes the value of inputs_embeds 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 args input_ids torch longtensor of shape batch_size num_codebooks sequence_length indices of input sequence tokens in the vocabulary corresponding to the sequence of audio codes indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes such as with the encodecmodel see encodecmodel encode for details what are input ids glossary input ids tip warning true the input_ids will automatically be converted from shape batch_size num_codebooks target_sequence_length to batch_size num_codebooks target_sequence_length in the forward pass if you obtain audio codes from an audio encoding model such as encodecmodel ensure that the number of frames is equal to 1 and that you reshape the audio codes from frames batch_size num_codebooks target_sequence_length to batch_size num_codebooks target_sequence_length prior to passing them as input_ids tip 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 encoder_hidden_states torch floattensor of shape batch_size encoder_sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder encoder_attention_mask torch longtensor of shape batch_size encoder_sequence_length optional mask to avoid performing cross attention on padding tokens indices of encoder input_ids 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 head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules in the decoder to avoid performing cross attention on hidden heads mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked past_key_values tuple tuple torch floattensor optional returned when use_cache true is passed or when config use_cache true tuple of tuple torch floattensor of length config n_layers with each tuple having 2 tensors of shape batch_size num_heads sequence_length embed_size_per_head and 2 additional tensors of shape batch_size num_heads encoder_sequence_length embed_size_per_head contains pre computed hidden states key and values in the self attention blocks and in the cross attention blocks that can be used see past_key_values input to speed up sequential decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length 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 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 decoder consisting of config num_hidden_layers layers each layer is a musicgendecoderlayer initialize weights and apply final processing retrieve input_ids and inputs_embeds bsz codebooks seq_len bsz codebooks seq_len past_key_values_length expand encoder attention mask bsz seq_len bsz 1 tgt_seq_len src_seq_len embed positions decoder layers check if head_mask cross_attn_head_mask has a correct number of layers specified if desired add layerdrop see https arxiv org abs 1909 11556 for description add hidden states from the last decoder layer 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 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 returns bsz num_codebooks seq_len vocab_size bsz num_codebooks seq_len vocab_size apply the delay pattern mask for classifier free guidance we need to replicate the decoder args across the batch dim we ll split these before sampling build a delayed pattern mask to the input_ids each codebook is offset by the previous codebook by one giving a delayed pattern mask at the start of sequence and end of sequence take the example where there are 4 codebooks and a max sequence length of 8 we have the delayed pattern mask of shape codebooks seq_len p 1 1 1 1 p p p p p 1 1 1 1 p p p p p 1 1 1 1 p p p p p 1 1 1 1 where p is the special padding token id and 1 indicates that the token is valid for prediction if we include a prompt decoder input ids the 1 positions indicate where new tokens should be predicted otherwise the mask is set to the value in the prompt p a b 1 1 p p p p p c d 1 1 p p p p p e f 1 1 p p p p p g h 1 1 where a h indicate the input prompt decoder input ids that are offset by 1 now we only override the 1 tokens in our prediction bsz num_codebooks seq_len bsz num_codebooks seq_len we only apply the mask if we have a large enough seq len otherwise we return as is fill the shifted ids with the prompt entries offset by the codebook idx mono channel loop over the codebooks one by one left right channels are interleaved in the generated codebooks so handle one then the other construct a pattern mask that indicates the positions of padding tokens for each codebook first fill the upper triangular part the eos padding then fill the lower triangular part the bos padding for left right channel we need to duplicate every row of the pattern mask in an interleaved fashion find the first position to start generating this is the first place we have the 1 token and will always be in the first codebook since it has no codebook offset we have no tokens that need to be filled return entire matrix of input ids bsz num_codebooks seq_len bsz num_codebooks seq_len apply a delay pattern mask to the decoder input ids only preserving predictions where the mask is set to 1 and otherwise setting to the value detailed in the mask generates sequences of token ids for models with a language modeling head tip warning true most generation controlling parameters are set in generation_config which if not passed will be set to the model s default generation configuration you can override any generation_config by passing the corresponding parameters to generate e g generate inputs num_beams 4 do_sample true for an overview of generation strategies and code examples check out the following guide generation_strategies tip parameters inputs torch tensor of varying shape depending on the modality optional the sequence used as a prompt for the generation or as model inputs to the encoder if none the method initializes it with bos_token_id and a batch size of 1 for decoder only models inputs should be in the format input_ids for encoder decoder models inputs can represent any of input_ids input_values input_features or pixel_values 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 logits_processor logitsprocessorlist optional custom logits processors that complement the default logits processors built from arguments and generation config if a logit processor is passed that is already created with the arguments or a generation config an error is thrown this feature is intended for advanced users stopping_criteria stoppingcriterialist optional custom stopping criteria that complement the default stopping criteria built from arguments and a generation config if a stopping criteria is passed that is already created with the arguments or a generation config an error is thrown this feature is intended for advanced users synced_gpus bool optional defaults to false whether to continue running the while loop until max_length needed for zero stage 3 streamer basestreamer optional streamer object that will be used to stream the generated sequences generated tokens are passed through streamer put token_ids and the streamer is responsible for any further processing kwargs dict str any optional ad hoc parametrization of generate_config and or additional model specific kwargs that will be forwarded to the forward function of the model if the model is an encoder decoder model encoder specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with decoder_ return utils modeloutput or torch longtensor a utils modeloutput if return_dict_in_generate true or when config return_dict_in_generate true or a torch floattensor if the model is not an encoder decoder model model config is_encoder_decoder false the possible utils modeloutput types are generation greedysearchdecoderonlyoutput generation sampledecoderonlyoutput generation beamsearchdecoderonlyoutput generation beamsampledecoderonlyoutput if the model is an encoder decoder model model config is_encoder_decoder true the possible utils modeloutput types are generation greedysearchencoderdecoderoutput generation sampleencoderdecoderoutput generation beamsearchencoderdecoderoutput generation beamsampleencoderdecoderoutput 1 handle generation_config and kwargs that might update it and validate the resulting objects all unused kwargs must be model kwargs 2 set generation parameters if not already defined 3 define model inputs inputs_tensor has to be defined model_input_name is defined if model specific keyword input is passed otherwise model_input_name is none all model specific keyword inputs are removed from model_kwargs 4 define other model kwargs 5 prepare max_length depending on other stopping criteria 6 prepare input_ids which will be used for auto regressive generation build the delay pattern mask for offsetting each codebook prediction by 1 this behaviour is specific to musicgen stash the delay mask so that we don t have to recompute it in each forward pass 7 determine generation mode 8 prepare batched cfg externally to enable coexistance with the unbatched cfg 9 prepare distribution pre_processing samplers 10 prepare stopping criteria 11 run greedy search 11 prepare logits warper expand input_ids with num_return_sequences additional sequences per batch 12 run sample apply the pattern mask to the final ids revert the pattern delay mask by filtering the pad token id initialize with config make sure that the individual model s config refers to the shared config so that the updates to the config will be synced text encoder outputs might need to be projected to different dimension for decoder tie text encoder decoder weights if config set accordingly tie text encoder decoder if needed tie text encoder and decoder base model get the text encoder to compute the encoder hidden states for generation example python from transformers import musicgenforconditionalgeneration model musicgenforconditionalgeneration from_pretrained facebook musicgen small at the moment fast initialization is not supported for composite models instantiate a text encoder an audio encoder and a musicgen decoder from one two or three base classes of the library from pretrained model checkpoints the model is set in evaluation mode by default using model eval dropout modules are deactivated to train the model you need to first set it back in training mode with model train params text_encoder_pretrained_model_name_or_path str optional information necessary to initiate the text encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like t5 base or namespaced under a user or organization name like google flan t5 base a path to a directory containing model weights saved using pretrainedmodel save_pretrained e g my_model_directory audio_encoder_pretrained_model_name_or_path str optional information necessary to initiate the audio encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like facebook encodec_24khz a path to a directory containing model weights saved using pretrainedmodel save_pretrained e g my_model_directory decoder_pretrained_model_name_or_path str optional defaults to none information necessary to initiate the decoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like gpt2 or namespaced under a user or organization name like facebook musicgen small a path to a directory containing model weights saved using pretrainedmodel save_pretrained e g my_model_directory model_args remaining positional arguments optional all remaining positional arguments will be passed to the underlying model s __init__ method kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g output_attentions true to update the text encoder configuration use the prefix text_encoder_ for each configuration parameter to update the audio encoder configuration use the prefix audio_encoder_ for each configuration parameter to update the decoder configuration use the prefix decoder_ for each configuration parameter to update the parent model configuration do not use a prefix for each configuration parameter behaves differently depending on whether a config is provided or automatically loaded example python from transformers import musicgenforconditionalgeneration initialize a musicgen model from a t5 text encoder encodec audio encoder and musicgen decoder model musicgenforconditionalgeneration from_sub_models_pretrained text_encoder_pretrained_model_name_or_path t5 base audio_encoder_pretrained_model_name_or_path facebook encodec_24khz decoder_pretrained_model_name_or_path facebook musicgen small saving model after fine tuning model save_pretrained musicgen ft load fine tuned model model musicgenforconditionalgeneration from_pretrained musicgen ft remove text encoder audio encoder and decoder kwargs from kwargs load and initialize the encoder and decoder the distinction between encoder and decoder at the model level is made by the value of the flag is_decoder that we need to set correctly instantiate config with corresponding kwargs returns examples python from transformers import autoprocessor musicgenforconditionalgeneration import torch processor autoprocessor from_pretrained facebook musicgen small model musicgenforconditionalgeneration from_pretrained facebook musicgen small inputs processor text 80s pop track with bassy drums and synth 90s rock song with loud guitars and heavy drums padding true return_tensors pt pad_token_id model generation_config pad_token_id decoder_input_ids torch ones inputs input_ids shape 0 model decoder num_codebooks 1 dtype torch long pad_token_id logits model inputs decoder_input_ids decoder_input_ids logits logits shape bsz num_codebooks tgt_len vocab_size torch size 8 1 2048 optionally project encoder_hidden_states mono input through encodec that we convert to stereo decode apply the delay pattern mask for classifier free guidance we need to replicate the decoder args across the batch dim we ll split these before sampling some generation methods already pass only the last input id default to old behavior keep only final id encoder_outputs is defined input_ids not needed prepares decoder_input_ids for generation with encoder decoder models 1 check whether the user has defined decoder_input_ids manually to facilitate in terms of input naming we also allow the user to pass it under input_ids if the encoder does not use it as the main input 2 encoder decoder models expect the decoder_input_ids to start with a special token let s ensure that no user input use decoder_start_token_id as decoder_input_ids user input but doesn t start with decoder_start_token_id prepend decoder_start_token_id and adjust decoder_attention_mask if provided 1 get text encoder compatibility with accelerate big model inference we need the encoder to outputs stuff on the same device as the inputs 2 prepare encoder args and encoder kwargs from model kwargs 3 make sure that encoder returns modeloutput for classifier free guidance we need to add a null input to our encoder hidden states 1 get audio encoder compatibility with accelerate big model inference we need the encoder to outputs stuff on the same device as the inputs 2 prepare encoder args and encoder kwargs from model kwargs 3 make sure that encoder returns modeloutput copy alternating left right channel codes into stereo codebook initializes input ids for generation if necessary make dummy input_ids with value 100 as a sanity check ensuring that they won t be used for encoding if there is some tensor in model_kwargs we can infer the batch size from it this is helpful with soft prompting or in multimodal implementations built on top of decoder only language models generates sequences of token ids for models with a language modeling head tip warning true most generation controlling parameters are set in generation_config which if not passed will be set to the model s default generation configuration you can override any generation_config by passing the corresponding parameters to generate e g generate inputs num_beams 4 do_sample true for an overview of generation strategies and code examples check out the following guide generation_strategies tip parameters inputs torch tensor of varying shape depending on the modality optional the sequence used as a prompt for the generation or as model inputs to the encoder if none the method initializes it with bos_token_id and a batch size of 1 for decoder only models inputs should be in the format input_ids for encoder decoder models inputs can represent any of input_ids input_values input_features or pixel_values 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 logits_processor logitsprocessorlist optional custom logits processors that complement the default logits processors built from arguments and generation config if a logit processor is passed that is already created with the arguments or a generation config an error is thrown this feature is intended for advanced users stopping_criteria stoppingcriterialist optional custom stopping criteria that complement the default stopping criteria built from arguments and a generation config if a stopping criteria is passed that is already created with the arguments or a generation config an error is thrown this feature is intended for advanced users synced_gpus bool optional defaults to false whether to continue running the while loop until max_length needed for zero stage 3 streamer basestreamer optional streamer object that will be used to stream the generated sequences generated tokens are passed through streamer put token_ids and the streamer is responsible for any further processing kwargs dict str any optional ad hoc parametrization of generate_config and or additional model specific kwargs that will be forwarded to the forward function of the model if the model is an encoder decoder model encoder specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with decoder_ return utils modeloutput or torch longtensor a utils modeloutput if return_dict_in_generate true or when config return_dict_in_generate true or a torch floattensor if the model is not an encoder decoder model model config is_encoder_decoder false the possible utils modeloutput types are generation greedysearchdecoderonlyoutput generation sampledecoderonlyoutput generation beamsearchdecoderonlyoutput generation beamsampledecoderonlyoutput if the model is an encoder decoder model model config is_encoder_decoder true the possible utils modeloutput types are generation greedysearchencoderdecoderoutput generation sampleencoderdecoderoutput generation beamsearchencoderdecoderoutput generation beamsampleencoderdecoderoutput 1 handle generation_config and kwargs that might update it and validate the resulting objects all unused kwargs must be model kwargs wrap the unconditional outputs as a basemodeloutput for compatibility with the rest of generate 2 set generation parameters if not already defined 3 define model inputs inputs_tensor has to be defined model_input_name is defined if model specific keyword input is passed otherwise model_input_name is none all model specific keyword inputs are removed from model_kwargs 4 define other model kwargs encoder_outputs are created and added to model_kwargs 5 prepare input_ids which will be used for auto regressive generation 6 prepare max_length depending on other stopping criteria build the delay pattern mask for offsetting each codebook prediction by 1 this behaviour is specific to musicgen stash the delay mask so that we don t have to recompute in each forward pass input_ids are ready to be placed on the streamer if used 7 determine generation mode 8 prepare batched cfg externally to enable coexistance with the unbatched cfg 9 prepare distribution pre_processing samplers 10 prepare stopping criteria 11 run greedy search 11 prepare logits warper expand input_ids with num_return_sequences additional sequences per batch 12 run sample apply the pattern mask to the final ids revert the pattern delay mask by filtering the pad token id append the frame dimension back to the audio codes helper function to get null inputs for unconditional generation enabling the model to be used without the feature extractor or tokenizer args num_samples int optional number of audio samples to unconditionally generate max_new_tokens int optional number of tokens to generate for each sample more tokens means longer audio samples at the expense of longer inference since more audio tokens need to be generated per sample example python from transformers import musicgenforconditionalgeneration model musicgenforconditionalgeneration from_pretrained facebook musicgen small get the unconditional or null inputs for the model unconditional_inputs model get_unconditional_inputs num_samples 1 audio_samples model generate unconditional_inputs max_new_tokens 256
import copy
import inspect
import math
import random
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss

from ...activations import ACT2FN
from ...generation.configuration_utils import GenerationConfig
from ...generation.logits_process import ClassifierFreeGuidanceLogitsProcessor, LogitsProcessorList
from ...generation.stopping_criteria import StoppingCriteriaList
from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
from ...modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPastAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
    ModelOutput,
    Seq2SeqLMOutput,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ..auto.configuration_auto import AutoConfig
from ..auto.modeling_auto import AutoModel
from .configuration_musicgen import MusicgenConfig, MusicgenDecoderConfig


if TYPE_CHECKING:
    from ...generation.streamers import BaseStreamer

logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "MusicgenConfig"
_CHECKPOINT_FOR_DOC = "facebook/musicgen-small"

MUSICGEN_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/musicgen-small",
    
]


@dataclass
class MusicgenUnconditionalInput(ModelOutput):
    

    encoder_outputs: Tuple[torch.FloatTensor] = None
    attention_mask: torch.LongTensor = None
    guidance_scale: float = None



def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
    
    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
    if decoder_start_token_id is None:
        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
    shifted_input_ids[:, 0] = decoder_start_token_id

    if pad_token_id is None:
        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
    
    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)

    return shifted_input_ids


class MusicgenSinusoidalPositionalEmbedding(nn.Module):
    

    def __init__(self, num_positions: int, embedding_dim: int):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.make_weights(num_positions, embedding_dim)

    def make_weights(self, num_embeddings: int, embedding_dim: int):
        emb_weights = self.get_embedding(num_embeddings, embedding_dim)
        if hasattr(self, "weights"):
            
            emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)

        self.weights = nn.Parameter(emb_weights)
        self.weights.requires_grad = False
        self.weights.detach_()

    @staticmethod
    def get_embedding(num_embeddings: int, embedding_dim: int):
        
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=1).view(num_embeddings, -1)
        if embedding_dim % 2 == 1:
            
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        return emb.to(torch.get_default_dtype())

    @torch.no_grad()
    def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
        bsz, codebooks, seq_len = input_ids.size()
        
        position_ids = (torch.arange(seq_len) + past_key_values_length).to(input_ids.device)
        
        if seq_len > self.weights.size(0):
            self.make_weights(seq_len + self.offset, self.embedding_dim)
        return self.weights.index_select(0, position_ids.view(-1)).detach()



class MusicgenAttention(nn.Module):
    

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        is_decoder: bool = False,
        bias: bool = True,
        is_causal: bool = False,
        config: Optional[MusicgenConfig] = None,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        self.config = config

        if (self.head_dim * num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {num_heads})."
            )
        self.scaling = self.head_dim**-0.5
        self.is_decoder = is_decoder
        self.is_causal = is_causal

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

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_value_states: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        

        
        
        is_cross_attention = key_value_states is not None

        bsz, tgt_len, _ = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scaling
        
        
        
        
        if (
            is_cross_attention
            and past_key_value is not None
            and past_key_value[0].shape[2] == key_value_states.shape[1]
        ):
            
            key_states = past_key_value[0]
            value_states = past_key_value[1]
        elif is_cross_attention:
            
            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
        elif past_key_value is not None:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)
        else:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        if self.is_decoder:
            
            
            
            
            
            
            
            past_key_value = (key_states, value_states)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
        key_states = key_states.reshape(*proj_shape)
        value_states = value_states.reshape(*proj_shape)

        src_len = key_states.size(1)
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                f" {attn_weights.size()}"
            )

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

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

        if layer_head_mask is not None:
            if layer_head_mask.size() != (self.num_heads,):
                raise ValueError(
                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
                    f" {layer_head_mask.size()}"
                )
            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if output_attentions:
            
            
            
            
            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
        else:
            attn_weights_reshaped = None

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.bmm(attn_probs, value_states)

        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)

        
        
        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped, past_key_value


class MusicgenDecoderLayer(nn.Module):
    def __init__(self, config: MusicgenDecoderConfig):
        super().__init__()
        self.embed_dim = config.hidden_size

        self.self_attn = MusicgenAttention(
            embed_dim=self.embed_dim,
            num_heads=config.num_attention_heads,
            dropout=config.attention_dropout,
            is_decoder=True,
            bias=False,
        )
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.encoder_attn = MusicgenAttention(
            self.embed_dim,
            config.num_attention_heads,
            dropout=config.attention_dropout,
            is_decoder=True,
            bias=False,
        )
        self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, config.ffn_dim, bias=False)
        self.fc2 = nn.Linear(config.ffn_dim, self.embed_dim, bias=False)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        cross_attn_layer_head_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = True,
    ) -> torch.Tensor:
        
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        
        
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=self_attn_past_key_value,
            attention_mask=attention_mask,
            layer_head_mask=layer_head_mask,
            output_attentions=output_attentions,
        )
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        
        cross_attn_present_key_value = None
        cross_attn_weights = None
        if encoder_hidden_states is not None:
            residual = hidden_states
            hidden_states = self.encoder_attn_layer_norm(hidden_states)

            
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
                hidden_states=hidden_states,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
                layer_head_mask=cross_attn_layer_head_mask,
                past_key_value=cross_attn_past_key_value,
                output_attentions=output_attentions,
            )
            hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
            hidden_states = residual + hidden_states

            
            present_key_value = present_key_value + cross_attn_present_key_value

        
        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights, cross_attn_weights)

        if use_cache:
            outputs += (present_key_value,)

        return outputs


class MusicgenPreTrainedModel(PreTrainedModel):
    

    config_class = MusicgenDecoderConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MusicgenDecoderLayer", "MusicgenAttention"]

    def _init_weights(self, module):
        std = self.config.initializer_factor
        if isinstance(module, (nn.Linear, nn.Conv1d)):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


MUSICGEN_START_DOCSTRING = r

MUSICGEN_INPUTS_DOCSTRING = r

MUSICGEN_DECODER_INPUTS_DOCSTRING = r


class MusicgenDecoder(MusicgenPreTrainedModel):
    

    def __init__(self, config: MusicgenDecoderConfig):
        super().__init__(config)
        self.dropout = config.dropout
        self.layerdrop = config.layerdrop
        self.max_target_positions = config.max_position_embeddings
        self.d_model = config.hidden_size
        self.num_codebooks = config.num_codebooks
        self.embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0

        embed_dim = config.vocab_size + 1
        self.embed_tokens = nn.ModuleList(
            [nn.Embedding(embed_dim, config.hidden_size) for _ in range(config.num_codebooks)]
        )

        self.embed_positions = MusicgenSinusoidalPositionalEmbedding(
            config.max_position_embeddings,
            config.hidden_size,
        )

        self.layers = nn.ModuleList([MusicgenDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.layer_norm = nn.LayerNorm(config.hidden_size)

        self.gradient_checkpointing = False
        
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    @add_start_docstrings_to_model_forward(MUSICGEN_DECODER_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = 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 decoder_input_ids and decoder_inputs_embeds at the same time")
        elif input_ids is not None:
            
            input = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
            bsz, num_codebooks, seq_len = input.shape
            input_shape = (bsz, seq_len)
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            input = inputs_embeds[:, :, -1:]
        else:
            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")

        
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if inputs_embeds is None:
            inputs_embeds = sum([self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks)])

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

        
        if encoder_hidden_states is not None and encoder_attention_mask is not None:
            
            encoder_attention_mask = _prepare_4d_attention_mask(
                encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
            )

        
        positions = self.embed_positions(input, past_key_values_length)

        hidden_states = inputs_embeds + positions.to(inputs_embeds.device)

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

        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

        
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
        next_decoder_cache = () if use_cache else None

        
        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
            if attn_mask is not None:
                if attn_mask.size()[0] != len(self.layers):
                    raise ValueError(
                        f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
                        f" {attn_mask.size()[0]}."
                    )
        for idx, decoder_layer in enumerate(self.layers):
            
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            dropout_probability = random.uniform(0, 1)
            if self.training and (dropout_probability < self.layerdrop):
                continue

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.forward,
                    hidden_states,
                    attention_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    head_mask[idx] if head_mask is not None else None,
                    cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
                    None,
                    output_attentions,
                    use_cache,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                    cross_attn_layer_head_mask=(
                        cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
                    ),
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                )
            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)

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

                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)

        hidden_states = self.layer_norm(hidden_states)

        
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
        )


@add_start_docstrings(
    "The bare Musicgen decoder model outputting raw hidden-states without any specific head on top.",
    MUSICGEN_START_DOCSTRING,
)
class MusicgenModel(MusicgenPreTrainedModel):
    def __init__(self, config: MusicgenDecoderConfig):
        super().__init__(config)
        self.decoder = MusicgenDecoder(config)
        
        self.post_init()

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

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

    def get_decoder(self):
        return self.decoder

    @add_start_docstrings_to_model_forward(MUSICGEN_DECODER_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = 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,
            encoder_attention_mask=encoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            head_mask=head_mask,
            cross_attn_head_mask=cross_attn_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 MusicGen decoder model with a language modelling head on top.",
    MUSICGEN_START_DOCSTRING,
)
class MusicgenForCausalLM(MusicgenPreTrainedModel):
    def __init__(self, config: MusicgenDecoderConfig):
        super().__init__(config)

        self.model = MusicgenModel(config)

        self.num_codebooks = config.num_codebooks
        self.lm_heads = nn.ModuleList(
            [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_codebooks)]
        )

        
        self.post_init()

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

    def set_input_embeddings(self, value):
        self.model.decoder.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_heads

    def set_output_embeddings(self, new_embeddings):
        self.lm_heads = new_embeddings

    def set_decoder(self, decoder):
        self.model.decoder = decoder

    def get_decoder(self):
        return self.model.decoder

    @add_start_docstrings_to_model_forward(MUSICGEN_DECODER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[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

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            head_mask=head_mask,
            cross_attn_head_mask=cross_attn_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,
        )

        hidden_states = outputs[0]

        lm_logits = torch.stack([head(hidden_states) for head in self.lm_heads], dim=1)

        loss = None
        if labels is not None:
            raise NotImplementedError("Training is not implemented for Musicgen.")

        
        lm_logits = lm_logits.reshape(-1, *lm_logits.shape[2:])

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

    def prepare_inputs_for_generation(
        self,
        input_ids,
        attention_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        head_mask=None,
        cross_attn_head_mask=None,
        past_key_values=None,
        use_cache=True,
        delay_pattern_mask=None,
        guidance_scale=None,
        **kwargs,
    ):
        if delay_pattern_mask is None:
            input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
                input_ids,
                pad_token_id=self.generation_config.pad_token_id,
                max_length=self.generation_config.max_length,
            )

        
        input_ids = self.apply_delay_pattern_mask(input_ids, delay_pattern_mask)

        if guidance_scale is not None and guidance_scale > 1:
            
            
            input_ids = input_ids.repeat((2, 1))
            if attention_mask is not None:
                attention_mask = attention_mask.repeat((2, 1))

        if past_key_values is not None:
            input_ids = input_ids[:, -1:]

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "encoder_hidden_states": encoder_hidden_states,
            "encoder_attention_mask": encoder_attention_mask,
            "head_mask": head_mask,
            "cross_attn_head_mask": cross_attn_head_mask,
            "past_key_values": past_key_values,
            "use_cache": use_cache,
        }

    def build_delay_pattern_mask(self, input_ids: torch.LongTensor, pad_token_id: int, max_length: int = None):
        
        
        input_ids = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
        bsz, num_codebooks, seq_len = input_ids.shape

        max_length = max_length if max_length is not None else self.generation_config.max_length
        input_ids_shifted = (
            torch.ones((bsz, num_codebooks, max_length), dtype=torch.long, device=input_ids.device) * -1
        )

        channel_codebooks = num_codebooks // 2 if self.config.audio_channels == 2 else num_codebooks
        
        if max_length < 2 * channel_codebooks - 1:
            return input_ids.reshape(bsz * num_codebooks, -1), input_ids_shifted.reshape(bsz * num_codebooks, -1)

        
        for codebook in range(channel_codebooks):
            if self.config.audio_channels == 1:
                
                input_ids_shifted[:, codebook, codebook : seq_len + codebook] = input_ids[:, codebook]
            else:
                
                input_ids_shifted[:, 2 * codebook, codebook : seq_len + codebook] = input_ids[:, 2 * codebook]
                input_ids_shifted[:, 2 * codebook + 1, codebook : seq_len + codebook] = input_ids[:, 2 * codebook + 1]

        
        
        delay_pattern = torch.triu(
            torch.ones((channel_codebooks, max_length), dtype=torch.bool), diagonal=max_length - channel_codebooks + 1
        )
        
        delay_pattern = delay_pattern + torch.tril(torch.ones((channel_codebooks, max_length), dtype=torch.bool))

        if self.config.audio_channels == 2:
            
            delay_pattern = delay_pattern.repeat_interleave(2, dim=0)

        mask = ~delay_pattern.to(input_ids.device)
        input_ids = mask * input_ids_shifted + ~mask * pad_token_id

        
        
        first_codebook_ids = input_ids[:, 0, :]
        start_ids = (first_codebook_ids == -1).nonzero()[:, 1]
        if len(start_ids) > 0:
            first_start_id = min(start_ids)
        else:
            
            first_start_id = seq_len

        
        pattern_mask = input_ids.reshape(bsz * num_codebooks, -1)
        input_ids = input_ids[..., :first_start_id].reshape(bsz * num_codebooks, -1)
        return input_ids, pattern_mask

    @staticmethod
    def apply_delay_pattern_mask(input_ids, decoder_pad_token_mask):
        
        seq_len = input_ids.shape[-1]
        decoder_pad_token_mask = decoder_pad_token_mask[..., :seq_len]
        input_ids = torch.where(decoder_pad_token_mask == -1, input_ids, decoder_pad_token_mask)
        return input_ids

    @torch.no_grad()
    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        synced_gpus: Optional[bool] = None,
        streamer: Optional["BaseStreamer"] = None,
        **kwargs,
    ):
        
        
        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())

        
        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()

        if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:
            if model_kwargs.get("attention_mask", None) is None:
                logger.warning(
                    "The attention mask and the pad token id were not set. As a consequence, you may observe "
                    "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
                )
            eos_token_id = generation_config.eos_token_id
            if isinstance(eos_token_id, list):
                eos_token_id = eos_token_id[0]
            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.")
            generation_config.pad_token_id = eos_token_id

        
        
        
        
        
        input_ids, model_input_name, model_kwargs = self._prepare_model_inputs(
            inputs, generation_config.bos_token_id, model_kwargs
        )
        batch_size = input_ids.shape[0] // self.num_codebooks

        
        model_kwargs["output_attentions"] = generation_config.output_attentions
        model_kwargs["output_hidden_states"] = generation_config.output_hidden_states
        model_kwargs["use_cache"] = generation_config.use_cache
        model_kwargs["guidance_scale"] = generation_config.guidance_scale

        requires_attention_mask = "encoder_outputs" not in model_kwargs
        if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
            model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
                input_ids, generation_config.pad_token_id, generation_config.eos_token_id
            )

        
        input_ids_seq_length = input_ids.shape[-1]
        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:
            logger.warning(
                f"Using the model-agnostic default `max_length` (={generation_config.max_length}) "
                "to control the generation length.  recommend setting `max_new_tokens` to control the maximum length of the generation."
            )
        elif generation_config.max_new_tokens is not None:
            if not has_default_max_length:
                logger.warning(
                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
                    "Please refer to the documentation for more information. "
                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
                )
            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length

        if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
            raise ValueError(
                f"Unfeasible length constraints: the minimum length ({generation_config.min_length}) is larger than"
                f" the maximum length ({generation_config.max_length})"
            )
        if input_ids_seq_length >= generation_config.max_length:
            logger.warning(
                f"Input length of decoder_input_ids is {input_ids_seq_length}, but `max_length` is set to"
                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
                " increasing `max_new_tokens`."
            )

        
        
        input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
            input_ids,
            pad_token_id=generation_config.decoder_start_token_id,
            max_length=generation_config.max_length,
        )

        if streamer is not None:
            streamer.put(input_ids.cpu())

        
        model_kwargs["delay_pattern_mask"] = delay_pattern_mask

        
        is_greedy_gen_mode = (
            (generation_config.num_beams == 1)
            and (generation_config.num_beam_groups == 1)
            and generation_config.do_sample is False
        )
        is_sample_gen_mode = (
            (generation_config.num_beams == 1)
            and (generation_config.num_beam_groups == 1)
            and generation_config.do_sample is True
        )

        
        if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
            logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
            generation_config.guidance_scale = None

        
        logits_processor = self._get_logits_processor(
            generation_config=generation_config,
            input_ids_seq_length=input_ids_seq_length,
            encoder_input_ids=input_ids,
            prefix_allowed_tokens_fn=None,
            logits_processor=logits_processor,
        )

        
        stopping_criteria = self._get_stopping_criteria(
            generation_config=generation_config, stopping_criteria=stopping_criteria
        )

        if is_greedy_gen_mode:
            if generation_config.num_return_sequences > 1:
                raise ValueError(
                    "num_return_sequences has to be 1 when doing greedy search, "
                    f"but is {generation_config.num_return_sequences}."
                )

            
            outputs = self.greedy_search(
                input_ids,
                logits_processor=logits_processor,
                stopping_criteria=stopping_criteria,
                pad_token_id=generation_config.pad_token_id,
                eos_token_id=generation_config.eos_token_id,
                output_scores=generation_config.output_scores,
                return_dict_in_generate=generation_config.return_dict_in_generate,
                synced_gpus=synced_gpus,
                streamer=streamer,
                **model_kwargs,
            )

        elif is_sample_gen_mode:
            
            logits_warper = self._get_logits_warper(generation_config)

            
            input_ids, model_kwargs = self._expand_inputs_for_generation(
                input_ids=input_ids,
                expand_size=generation_config.num_return_sequences,
                **model_kwargs,
            )

            
            outputs = self.sample(
                input_ids,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                stopping_criteria=stopping_criteria,
                pad_token_id=generation_config.pad_token_id,
                eos_token_id=generation_config.eos_token_id,
                output_scores=generation_config.output_scores,
                return_dict_in_generate=generation_config.return_dict_in_generate,
                synced_gpus=synced_gpus,
                streamer=streamer,
                **model_kwargs,
            )

        else:
            raise ValueError(
                "Got incompatible mode for generation, should be one of greedy or sampling. "
                "Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`."
            )

        if generation_config.return_dict_in_generate:
            output_ids = outputs.sequences
        else:
            output_ids = outputs

        
        output_ids = self.apply_delay_pattern_mask(output_ids, model_kwargs["delay_pattern_mask"])

        
        output_ids = output_ids[output_ids != generation_config.pad_token_id].reshape(
            batch_size, self.num_codebooks, -1
        )

        if generation_config.return_dict_in_generate:
            outputs.sequences = output_ids
            return outputs
        else:
            return output_ids


@add_start_docstrings(
    "The composite MusicGen model with a text encoder, audio encoder and Musicgen decoder, "
    "for music generation tasks with one or both of text and audio prompts.",
    MUSICGEN_START_DOCSTRING,
)
class MusicgenForConditionalGeneration(PreTrainedModel):
    config_class = MusicgenConfig
    base_model_prefix = "encoder_decoder"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = True

    def __init__(
        self,
        config: Optional[MusicgenConfig] = None,
        text_encoder: Optional[PreTrainedModel] = None,
        audio_encoder: Optional[PreTrainedModel] = None,
        decoder: Optional[MusicgenForCausalLM] = None,
    ):
        if config is None and (text_encoder is None or audio_encoder is None or decoder is None):
            raise ValueError(
                "Either a configuration has to be provided, or all three of text encoder, audio encoder and MusicGen decoder."
            )
        if config is None:
            config = MusicgenConfig.from_sub_models_config(text_encoder.config, audio_encoder.config, decoder.config)
        else:
            if not isinstance(config, self.config_class):
                raise ValueError(f"Config: {config} has to be of type {self.config_class}")

        if config.decoder.cross_attention_hidden_size is not None:
            if config.decoder.cross_attention_hidden_size != config.text_encoder.hidden_size:
                raise ValueError(
                    "If `cross_attention_hidden_size` is specified in the MusicGen decoder's configuration, it has to be equal"
                    f" to the text encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
                    f" `config.decoder.cross_attention_hidden_size` and {config.text_encoder.hidden_size} for"
                    " `config.text_encoder.hidden_size`."
                )

        
        super().__init__(config)

        if text_encoder is None:
            from ..auto.modeling_auto import AutoModelForTextEncoding

            text_encoder = AutoModelForTextEncoding.from_config(config.text_encoder)

        if audio_encoder is None:
            from ..auto.modeling_auto import AutoModel

            audio_encoder = AutoModel.from_config(config.audio_encoder)

        if decoder is None:
            decoder = MusicgenForCausalLM(config.decoder)

        self.text_encoder = text_encoder
        self.audio_encoder = audio_encoder
        self.decoder = decoder

        if self.text_encoder.config.to_dict() != self.config.text_encoder.to_dict():
            logger.warning(
                f"Config of the text_encoder: {self.text_encoder.__class__} is overwritten by shared text_encoder config:"
                f" {self.config.text_encoder}"
            )
        if self.audio_encoder.config.to_dict() != self.config.audio_encoder.to_dict():
            logger.warning(
                f"Config of the audio_encoder: {self.audio_encoder.__class__} is overwritten by shared audio_encoder config:"
                f" {self.config.audio_encoder}"
            )
        if self.decoder.config.to_dict() != self.config.decoder.to_dict():
            logger.warning(
                f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
                f" {self.config.decoder}"
            )

        
        
        self.text_encoder.config = self.config.text_encoder
        self.audio_encoder.config = self.config.audio_encoder
        self.decoder.config = self.config.decoder

        
        if (
            self.text_encoder.config.hidden_size != self.decoder.config.hidden_size
            and self.decoder.config.cross_attention_hidden_size is None
        ):
            self.enc_to_dec_proj = nn.Linear(self.text_encoder.config.hidden_size, self.decoder.config.hidden_size)

        if self.text_encoder.get_output_embeddings() is not None:
            raise ValueError(
                f"The encoder {self.text_encoder} should not have a LM Head. Please use a model without and LM Head"
            )

        decoder_signature = set(inspect.signature(self.decoder.forward).parameters.keys())
        if "encoder_hidden_states" not in decoder_signature:
            raise ValueError(
                "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "
                "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"
            )

        
        self.tie_weights()

    def tie_weights(self):
        
        if self.config.tie_encoder_decoder:
            
            decoder_base_model_prefix = self.decoder.base_model_prefix
            self._tie_encoder_decoder_weights(
                self.text_encoder, self.decoder._modules[decoder_base_model_prefix], self.decoder.base_model_prefix
            )

    def get_audio_encoder(self):
        return self.audio_encoder

    def get_text_encoder(self):
        return self.text_encoder

    def get_encoder(self):
        
        return self.get_text_encoder()

    def get_decoder(self):
        return self.decoder

    def get_input_embeddings(self):
        return self.text_encoder.get_input_embeddings()

    def get_output_embeddings(self):
        return self.decoder.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        return self.decoder.set_output_embeddings(new_embeddings)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        r

        
        if kwargs.get("_fast_init", False):
            logger.warning(
                "Fast initialization is currently not supported for MusicgenForConditionalGeneration. "
                "Falling back to slow initialization..."
            )
        kwargs["_fast_init"] = False

        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)

    @classmethod
    def from_sub_models_pretrained(
        cls,
        text_encoder_pretrained_model_name_or_path: str = None,
        audio_encoder_pretrained_model_name_or_path: str = None,
        decoder_pretrained_model_name_or_path: str = None,
        *model_args,
        **kwargs,
    ) -> PreTrainedModel:
        r

        kwargs_text_encoder = {
            argument[len("text_encoder_") :]: value
            for argument, value in kwargs.items()
            if argument.startswith("text_encoder_")
        }

        kwargs_audio_encoder = {
            argument[len("audio_encoder_") :]: value
            for argument, value in kwargs.items()
            if argument.startswith("audio_encoder_")
        }

        kwargs_decoder = {
            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
        }

        
        for key in kwargs_text_encoder.keys():
            del kwargs["text_encoder_" + key]
        for key in kwargs_audio_encoder.keys():
            del kwargs["audio_encoder_" + key]
        for key in kwargs_decoder.keys():
            del kwargs["decoder_" + key]

        
        
        
        text_encoder = kwargs_text_encoder.pop("model", None)
        if text_encoder is None:
            if text_encoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `text_encoder_model` is not defined as an argument, a `text_encoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_text_encoder:
                encoder_config, kwargs_text_encoder = AutoConfig.from_pretrained(
                    text_encoder_pretrained_model_name_or_path, **kwargs_text_encoder, return_unused_kwargs=True
                )

                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
                    logger.info(
                        f"Initializing {text_encoder_pretrained_model_name_or_path} as a text_encoder model "
                        "from a decoder model. Cross-attention and casual mask are disabled."
                    )
                    encoder_config.is_decoder = False
                    encoder_config.add_cross_attention = False

                kwargs_text_encoder["config"] = encoder_config

            text_encoder = AutoModel.from_pretrained(
                text_encoder_pretrained_model_name_or_path, *model_args, **kwargs_text_encoder
            )

        audio_encoder = kwargs_audio_encoder.pop("model", None)
        if audio_encoder is None:
            if audio_encoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `audio_encoder_model` is not defined as an argument, an `audio_encoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_audio_encoder:
                encoder_config, kwargs_audio_encoder = AutoConfig.from_pretrained(
                    audio_encoder_pretrained_model_name_or_path, **kwargs_audio_encoder, return_unused_kwargs=True
                )

                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
                    logger.info(
                        f"Initializing {audio_encoder_pretrained_model_name_or_path} as an audio_encoder model "
                        "from a decoder model. Cross-attention and casual mask are disabled."
                    )
                    encoder_config.is_decoder = False
                    encoder_config.add_cross_attention = False

                kwargs_audio_encoder["config"] = encoder_config

            audio_encoder = AutoModel.from_pretrained(
                audio_encoder_pretrained_model_name_or_path, *model_args, **kwargs_audio_encoder
            )

        decoder = kwargs_decoder.pop("model", None)
        if decoder is None:
            if decoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_decoder:
                decoder_config, kwargs_decoder = AutoConfig.from_pretrained(
                    decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
                )

                if isinstance(decoder_config, MusicgenConfig):
                    decoder_config = decoder_config.decoder

                if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
                    logger.info(
                        f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
                        f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
                        f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
                    )
                    decoder_config.is_decoder = True
                    decoder_config.add_cross_attention = True

                kwargs_decoder["config"] = decoder_config

            if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
                logger.warning(
                    f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
                    f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
                    "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
                    "passed to `.from_sub_models_pretrained(...)` are set to `True` or do not pass a "
                    "`decoder_config` to `.from_sub_models_pretrained(...)`"
                )

            decoder = MusicgenForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)

        
        config = MusicgenConfig.from_sub_models_config(
            text_encoder.config, audio_encoder.config, decoder.config, **kwargs
        )
        return cls(text_encoder=text_encoder, audio_encoder=audio_encoder, decoder=decoder, config=config)

    @add_start_docstrings_to_model_forward(MUSICGEN_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.BoolTensor] = None,
        input_values: Optional[torch.FloatTensor] = None,
        padding_mask: Optional[torch.BoolTensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.BoolTensor] = None,
        encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
        past_key_values: Tuple[Tuple[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_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,
        **kwargs,
    ) -> Union[Tuple, Seq2SeqLMOutput]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        kwargs_text_encoder = {
            argument[len("text_encoder_")]: value
            for argument, value in kwargs.items()
            if argument.startswith("text_encoder_")
        }

        kwargs_audio_encoder = {
            argument[len("audio_encoder_")]: value
            for argument, value in kwargs.items()
            if argument.startswith("audio_encoder_")
        }

        kwargs_decoder = {
            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
        }

        if encoder_outputs is None:
            encoder_outputs = self.text_encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
                **kwargs_text_encoder,
            )
        elif isinstance(encoder_outputs, tuple):
            encoder_outputs = BaseModelOutput(*encoder_outputs)

        encoder_hidden_states = encoder_outputs[0]

        
        if (
            self.text_encoder.config.hidden_size != self.decoder.config.hidden_size
            and self.decoder.config.cross_attention_hidden_size is None
        ):
            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)

        if attention_mask is not None:
            encoder_hidden_states = encoder_hidden_states * attention_mask[..., None]

        if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
            decoder_input_ids = shift_tokens_right(
                labels, self.config.pad_token_id, self.config.decoder_start_token_id
            )

        elif decoder_input_ids is None and decoder_inputs_embeds is None:
            audio_encoder_outputs = self.audio_encoder(
                input_values=input_values,
                padding_mask=padding_mask,
                **kwargs_audio_encoder,
            )
            audio_codes = audio_encoder_outputs.audio_codes
            frames, bsz, codebooks, seq_len = audio_codes.shape
            if frames != 1:
                raise ValueError(
                    f"Expected 1 frame in the audio code outputs, got {frames} frames. Ensure chunking is "
                    "disabled by setting `chunk_length=None` in the audio encoder."
                )

            if self.config.decoder.audio_channels == 2 and audio_codes.shape[2] == self.decoder.num_codebooks // 2:
                
                audio_codes = audio_codes.repeat_interleave(2, dim=2)

            decoder_input_ids = audio_codes[0, ...].reshape(bsz * self.decoder.num_codebooks, seq_len)

        
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=attention_mask,
            inputs_embeds=decoder_inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            use_cache=use_cache,
            past_key_values=past_key_values,
            return_dict=return_dict,
            **kwargs_decoder,
        )

        loss = None
        if labels is not None:
            logits = decoder_outputs.logits if return_dict else decoder_outputs[0]
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))

        if not return_dict:
            if loss is not None:
                return (loss,) + decoder_outputs + encoder_outputs
            else:
                return decoder_outputs + encoder_outputs

        return Seq2SeqLMOutput(
            loss=loss,
            logits=decoder_outputs.logits,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        past_key_values=None,
        attention_mask=None,
        head_mask=None,
        decoder_attention_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        use_cache=None,
        encoder_outputs=None,
        decoder_delay_pattern_mask=None,
        guidance_scale=None,
        **kwargs,
    ):
        if decoder_delay_pattern_mask is None:
            decoder_input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
                decoder_input_ids,
                self.generation_config.pad_token_id,
                max_length=self.generation_config.max_length,
            )

        
        decoder_input_ids = self.decoder.apply_delay_pattern_mask(decoder_input_ids, decoder_delay_pattern_mask)

        if guidance_scale is not None and guidance_scale > 1:
            
            
            decoder_input_ids = decoder_input_ids.repeat((2, 1))
            if decoder_attention_mask is not None:
                decoder_attention_mask = decoder_attention_mask.repeat((2, 1))

        if past_key_values is not None:
            past_length = past_key_values[0][0].shape[2]

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

            decoder_input_ids = decoder_input_ids[:, remove_prefix_length:]

        return {
            "input_ids": None,  
            "encoder_outputs": encoder_outputs,
            "past_key_values": past_key_values,
            "decoder_input_ids": decoder_input_ids,
            "attention_mask": attention_mask,
            "decoder_attention_mask": decoder_attention_mask,
            "head_mask": head_mask,
            "decoder_head_mask": decoder_head_mask,
            "cross_attn_head_mask": cross_attn_head_mask,
            "use_cache": use_cache,
        }

    def _prepare_decoder_input_ids_for_generation(
        self,
        batch_size: int,
        model_input_name: str,
        model_kwargs: Dict[str, torch.Tensor],
        decoder_start_token_id: int = None,
        bos_token_id: int = None,
        device: torch.device = None,
    ) -> Tuple[torch.LongTensor, Dict[str, torch.Tensor]]:
        

        
        
        if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
            decoder_input_ids = model_kwargs.pop("decoder_input_ids")
        elif "input_ids" in model_kwargs and model_input_name != "input_ids":
            decoder_input_ids = model_kwargs.pop("input_ids")
        else:
            decoder_input_ids = None

        
        decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)
        if device is None:
            device = self.device
        decoder_input_ids_start = (
            torch.ones((batch_size * self.decoder.num_codebooks, 1), dtype=torch.long, device=device)
            * decoder_start_token_id
        )

        
        if decoder_input_ids is None:
            decoder_input_ids = decoder_input_ids_start

        
        
        elif (decoder_input_ids[..., 0] != decoder_start_token_id).all().item():
            decoder_input_ids = torch.cat([decoder_input_ids_start, decoder_input_ids], dim=-1)
            if "decoder_attention_mask" in model_kwargs:
                decoder_attention_mask = model_kwargs["decoder_attention_mask"]
                decoder_attention_mask = torch.cat(
                    (torch.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask),
                    dim=-1,
                )
                model_kwargs["decoder_attention_mask"] = decoder_attention_mask

        return decoder_input_ids, model_kwargs

    def _prepare_text_encoder_kwargs_for_generation(
        self,
        inputs_tensor: torch.Tensor,
        model_kwargs,
        model_input_name: Optional[str] = None,
        guidance_scale: Optional[float] = None,
    ) -> Dict[str, Any]:
        
        encoder = self.get_text_encoder()
        
        
        if hasattr(encoder, "_hf_hook"):
            encoder._hf_hook.io_same_device = True

        
        irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
        encoder_kwargs = {
            argument: value
            for argument, value in model_kwargs.items()
            if not any(argument.startswith(p) for p in irrelevant_prefix)
        }
        encoder_signature = set(inspect.signature(encoder.forward).parameters)
        encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
        if not encoder_accepts_wildcard:
            encoder_kwargs = {
                argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
            }

        
        model_input_name = model_input_name if model_input_name is not None else self.text_encoder.main_input_name
        encoder_kwargs["return_dict"] = True
        encoder_kwargs[model_input_name] = inputs_tensor
        last_hidden_state = encoder(**encoder_kwargs).last_hidden_state

        
        if guidance_scale is not None and guidance_scale > 1:
            last_hidden_state = torch.concatenate([last_hidden_state, torch.zeros_like(last_hidden_state)], dim=0)
            if "attention_mask" in model_kwargs:
                model_kwargs["attention_mask"] = torch.concatenate(
                    [model_kwargs["attention_mask"], torch.zeros_like(model_kwargs["attention_mask"])], dim=0
                )

        model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=last_hidden_state)

        return model_kwargs

    def _prepare_audio_encoder_kwargs_for_generation(
        self, input_values, model_kwargs, model_input_name: Optional[str] = None
    ):
        
        encoder = self.get_audio_encoder()
        
        
        if hasattr(encoder, "_hf_hook"):
            encoder._hf_hook.io_same_device = True

        
        irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
        encoder_kwargs = {
            argument: value
            for argument, value in model_kwargs.items()
            if not any(argument.startswith(p) for p in irrelevant_prefix)
        }
        encoder_signature = set(inspect.signature(encoder.forward).parameters)
        encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
        if not encoder_accepts_wildcard:
            encoder_kwargs = {
                argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
            }

        
        model_input_name = model_input_name if model_input_name is not None else self.audio_encoder.main_input_name
        encoder_kwargs["return_dict"] = True

        if self.decoder.config.audio_channels == 1:
            encoder_kwargs[model_input_name] = input_values
            audio_encoder_outputs = encoder.encode(**encoder_kwargs)
            audio_codes = audio_encoder_outputs.audio_codes
            audio_scales = audio_encoder_outputs.audio_scales

            frames, bsz, codebooks, seq_len = audio_codes.shape

        else:
            if input_values.shape[1] != 2:
                raise ValueError(
                    f"Expected stereo audio (2-channels) but example has {input_values.shape[1]} channel."
                )

            encoder_kwargs[model_input_name] = input_values[:, :1, :]
            audio_encoder_outputs_left = encoder.encode(**encoder_kwargs)
            audio_codes_left = audio_encoder_outputs_left.audio_codes
            audio_scales_left = audio_encoder_outputs_left.audio_scales

            encoder_kwargs[model_input_name] = input_values[:, 1:, :]
            audio_encoder_outputs_right = encoder.encode(**encoder_kwargs)
            audio_codes_right = audio_encoder_outputs_right.audio_codes
            audio_scales_right = audio_encoder_outputs_right.audio_scales

            frames, bsz, codebooks, seq_len = audio_codes_left.shape
            
            audio_codes = audio_codes_left.new_ones((frames, bsz, 2 * codebooks, seq_len))

            audio_codes[:, :, ::2, :] = audio_codes_left
            audio_codes[:, :, 1::2, :] = audio_codes_right

            if audio_scales_left != [None] or audio_scales_right != [None]:
                audio_scales = torch.stack([audio_scales_left, audio_scales_right], dim=1)
            else:
                audio_scales = [None] * bsz

        if frames != 1:
            raise ValueError(
                f"Expected 1 frame in the audio code outputs, got {frames} frames. Ensure chunking is "
                "disabled by setting `chunk_length=None` in the audio encoder."
            )

        decoder_input_ids = audio_codes[0, ...].reshape(bsz * self.decoder.num_codebooks, seq_len)

        model_kwargs["decoder_input_ids"] = decoder_input_ids
        model_kwargs["audio_scales"] = audio_scales
        return model_kwargs

    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
        return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)

    def resize_token_embeddings(self, *args, **kwargs):
        raise NotImplementedError(
            "Resizing the embedding layers via the EncoderDecoderModel directly is not supported. Please use the"
            " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
            " model.decoder.resize_token_embeddings(...))"
        )

    def _maybe_initialize_input_ids_for_generation(
        self,
        inputs: Optional[torch.Tensor] = None,
        bos_token_id: Optional[int] = None,
        model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
    ) -> torch.LongTensor:
        
        if inputs is not None:
            return inputs

        encoder_outputs = model_kwargs.get("encoder_outputs")
        if encoder_outputs is not None:
            
            shape = encoder_outputs[0].size()[:-1]
            return torch.ones(shape, dtype=torch.long, device=self.device) * -100

        if bos_token_id is None:
            raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.")

        
        
        batch_size = 1
        for value in model_kwargs.values():
            if isinstance(value, torch.Tensor):
                batch_size = value.shape[0]
                break
        return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id

    @torch.no_grad()
    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        synced_gpus: Optional[bool] = None,
        streamer: Optional["BaseStreamer"] = None,
        **kwargs,
    ):
        
        
        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())

        if model_kwargs.get("encoder_outputs") is not None and type(model_kwargs["encoder_outputs"]) == tuple:
            
            model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=model_kwargs["encoder_outputs"][0])

        
        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()

        if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:
            if model_kwargs.get("attention_mask", None) is None:
                logger.warning(
                    "The attention mask and the pad token id were not set. As a consequence, you may observe "
                    "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
                )
            eos_token_id = generation_config.eos_token_id
            if isinstance(eos_token_id, list):
                eos_token_id = eos_token_id[0]
            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.")
            generation_config.pad_token_id = eos_token_id

        
        
        
        
        
        inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
            inputs, generation_config.bos_token_id, model_kwargs
        )
        batch_size = inputs_tensor.shape[0]

        
        model_kwargs["output_attentions"] = generation_config.output_attentions
        model_kwargs["output_hidden_states"] = generation_config.output_hidden_states
        model_kwargs["use_cache"] = generation_config.use_cache
        model_kwargs["guidance_scale"] = generation_config.guidance_scale

        requires_attention_mask = "encoder_outputs" not in model_kwargs

        if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
            model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
                inputs_tensor, generation_config.pad_token_id, generation_config.eos_token_id
            )

        if "encoder_outputs" not in model_kwargs:
            
            model_kwargs = self._prepare_text_encoder_kwargs_for_generation(
                inputs_tensor,
                model_kwargs,
                model_input_name,
                guidance_scale=generation_config.guidance_scale,
            )

        if "decoder_input_ids" not in model_kwargs and "input_values" in model_kwargs:
            model_kwargs = self._prepare_audio_encoder_kwargs_for_generation(
                model_kwargs["input_values"],
                model_kwargs,
            )

        
        input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
            batch_size=batch_size,
            model_input_name=model_input_name,
            model_kwargs=model_kwargs,
            decoder_start_token_id=generation_config.decoder_start_token_id,
            bos_token_id=generation_config.bos_token_id,
            device=inputs_tensor.device,
        )

        
        input_ids_seq_length = input_ids.shape[-1]
        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        if has_default_max_length and generation_config.max_new_tokens is None:
            logger.warning(
                f"Using the model-agnostic default `max_length` (={generation_config.max_length}) "
                "to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation."
            )
        elif generation_config.max_new_tokens is not None:
            if not has_default_max_length:
                logger.warning(
                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
                    "Please refer to the documentation for more information. "
                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
                )
            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length

        if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
            raise ValueError(
                f"Unfeasible length constraints: the minimum length ({generation_config.min_length}) is larger than"
                f" the maximum length ({generation_config.max_length})"
            )
        if input_ids_seq_length >= generation_config.max_length:
            logger.warning(
                f"Input length of decoder_input_ids is {input_ids_seq_length}, but `max_length` is set to"
                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
                " increasing `max_new_tokens`."
            )

        
        input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
            input_ids,
            pad_token_id=generation_config.decoder_start_token_id,
            max_length=generation_config.max_length,
        )
        
        model_kwargs["decoder_delay_pattern_mask"] = decoder_delay_pattern_mask

        
        if streamer is not None:
            streamer.put(input_ids.cpu())

        
        is_greedy_gen_mode = (
            (generation_config.num_beams == 1)
            and (generation_config.num_beam_groups == 1)
            and generation_config.do_sample is False
        )
        is_sample_gen_mode = (
            (generation_config.num_beams == 1)
            and (generation_config.num_beam_groups == 1)
            and generation_config.do_sample is True
        )

        
        if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
            logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
            generation_config.guidance_scale = None

        
        logits_processor = self._get_logits_processor(
            generation_config=generation_config,
            input_ids_seq_length=input_ids_seq_length,
            encoder_input_ids=inputs_tensor,
            prefix_allowed_tokens_fn=None,
            logits_processor=logits_processor,
        )

        
        stopping_criteria = self._get_stopping_criteria(
            generation_config=generation_config, stopping_criteria=stopping_criteria
        )

        if is_greedy_gen_mode:
            if generation_config.num_return_sequences > 1:
                raise ValueError(
                    "num_return_sequences has to be 1 when doing greedy search, "
                    f"but is {generation_config.num_return_sequences}."
                )

            
            outputs = self.greedy_search(
                input_ids,
                logits_processor=logits_processor,
                stopping_criteria=stopping_criteria,
                pad_token_id=generation_config.pad_token_id,
                eos_token_id=generation_config.eos_token_id,
                output_scores=generation_config.output_scores,
                return_dict_in_generate=generation_config.return_dict_in_generate,
                synced_gpus=synced_gpus,
                streamer=streamer,
                **model_kwargs,
            )

        elif is_sample_gen_mode:
            
            logits_warper = self._get_logits_warper(generation_config)

            
            input_ids, model_kwargs = self._expand_inputs_for_generation(
                input_ids=input_ids,
                expand_size=generation_config.num_return_sequences,
                is_encoder_decoder=self.config.is_encoder_decoder,
                **model_kwargs,
            )

            
            outputs = self.sample(
                input_ids,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                stopping_criteria=stopping_criteria,
                pad_token_id=generation_config.pad_token_id,
                eos_token_id=generation_config.eos_token_id,
                output_scores=generation_config.output_scores,
                return_dict_in_generate=generation_config.return_dict_in_generate,
                synced_gpus=synced_gpus,
                streamer=streamer,
                **model_kwargs,
            )

        else:
            raise ValueError(
                "Got incompatible mode for generation, should be one of greedy or sampling. "
                "Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`."
            )

        if generation_config.return_dict_in_generate:
            output_ids = outputs.sequences
        else:
            output_ids = outputs

        
        output_ids = self.decoder.apply_delay_pattern_mask(output_ids, model_kwargs["decoder_delay_pattern_mask"])

        
        output_ids = output_ids[output_ids != generation_config.pad_token_id].reshape(
            batch_size, self.decoder.num_codebooks, -1
        )

        
        output_ids = output_ids[None, ...]

        audio_scales = model_kwargs.get("audio_scales")
        if audio_scales is None:
            audio_scales = [None] * batch_size

        if self.decoder.config.audio_channels == 1:
            output_values = self.audio_encoder.decode(
                output_ids,
                audio_scales=audio_scales,
            ).audio_values
        else:
            codec_outputs_left = self.audio_encoder.decode(output_ids[:, :, ::2, :], audio_scales=audio_scales)
            output_values_left = codec_outputs_left.audio_values

            codec_outputs_right = self.audio_encoder.decode(output_ids[:, :, 1::2, :], audio_scales=audio_scales)
            output_values_right = codec_outputs_right.audio_values

            output_values = torch.cat([output_values_left, output_values_right], dim=1)

        if generation_config.return_dict_in_generate:
            outputs.sequences = output_values
            return outputs
        else:
            return output_values

    def get_unconditional_inputs(self, num_samples=1):
        
        last_hidden_state = torch.zeros(
            (num_samples, 1, self.config.text_encoder.hidden_size), device=self.device, dtype=self.dtype
        )

        attention_mask = torch.zeros((num_samples, 1), device=self.device, dtype=torch.long)

        return MusicgenUnconditionalInput(
            encoder_outputs=(last_hidden_state,),
            attention_mask=attention_mask,
            guidance_scale=1.0,
        )