codingutf8 2021 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 testing suite for the pytorch musicgen model import copy import inspect import math import unittest import numpy as np from transformers import encodecconfig musicgenconfig musicgendecoderconfig musicgenprocessor pretrainedconfig t5config from transformers testingutils import istorchavailable requiretorch requiretorchfp16 slow torchdevice from transformers utils import cachedproperty from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin floatstensor idstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import musicgenforcausallm musicgenforconditionalgeneration musicgenmodel setseed from transformers generation import greedysearchdecoderonlyoutput greedysearchencoderdecoderoutput infnanremovelogitsprocessor logitsprocessorlist sampledecoderonlyoutput sampleencoderdecoderoutput def configzeroinitconfig configsnoinit copy deepcopyconfig for key in configsnoinit dict keys if range in key or std in key or initializerfactor in key or layerscale in key setattrconfigsnoinit key 1e10 if isinstancegetattrconfigsnoinit key none pretrainedconfig noinitsubconfig configzeroinitgetattrconfigsnoinit key setattrconfigsnoinit key noinitsubconfig return configsnoinit def preparemusicgendecoderinputsdict config inputids attentionmasknone headmasknone encoderhiddenstatesnone encoderattentionmasknone crossattnheadmasknone if attentionmask is none attentionmask inputids reshape1 config numcodebooks inputids shape1 0 attentionmask attentionmask neconfig padtokenid if headmask is none headmask torch onesconfig numhiddenlayers config numattentionheads devicetorchdevice if encoderattentionmask is none and encoderhiddenstates is not none encoderattentionmask torch onesencoderhiddenstates shape 2 devicetorchdevice if crossattnheadmask is none crossattnheadmask torch onesconfig numhiddenlayers config numattentionheads devicetorchdevice return inputids inputids attentionmask attentionmask encoderhiddenstates encoderhiddenstates encoderattentionmask encoderattentionmask headmask headmask crossattnheadmask crossattnheadmask class musicgendecodertester def init self parent batchsize2 seqlength7 istrainingfalse uselabelsfalse vocabsize99 hiddensize16 numhiddenlayers2 numattentionheads4 intermediatesize4 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings100 padtokenid99 bostokenid99 numcodebooks4 self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self maxpositionembeddings maxpositionembeddings self padtokenid padtokenid self bostokenid bostokenid self numcodebooks numcodebooks def prepareconfigandinputsself inputids idstensorself batchsize self numcodebooks self seqlength self vocabsize encoderhiddenstates floatstensorself batchsize self seqlength self hiddensize config self getconfig inputsdict preparemusicgendecoderinputsdict config inputids encoderhiddenstatesencoderhiddenstates return config inputsdict def getconfigself config musicgendecoderconfig vocabsizeself vocabsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads dffself intermediatesize padtokenidself padtokenid decoderstarttokenidself bostokenid bostokenidself bostokenid numcodebooksself numcodebooks tiewordembeddingsfalse return config def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict requiretorch class musicgendecodertestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses musicgenmodel musicgenforcausallm if istorchavailable else greedysamplemodelclasses musicgenforcausallm if istorchavailable else we don t want to run all the generation tests only a specific subset pipelinemodelmapping testpruning false testresizeembeddings false def setupself self modeltester musicgendecodertesterself self configtester configtesterself configclassmusicgendecoderconfig hiddensize16 def testconfigself self configtester runcommontests override since we have to compute the input embeddings over codebooks def testinputsembedsself config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig model totorchdevice model eval inputs copy deepcopyself prepareforclassinputsdict modelclass inputids inputsinputids del inputsinputids embedtokens model getinputembeddings inputids inputids reshape1 config numcodebooks inputids shape1 inputsinputsembeds sum embedtokenscodebookinputids codebook for codebook in rangeconfig numcodebooks with torch nograd modelinputs0 override since we have embeddings lm heads over multiple codebooks def testmodelcommonattributesself config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig firstembed model getinputembeddings0 self assertisinstancefirstembed torch nn embedding lmheads model getoutputembeddings self asserttruelmheads is none or isinstancelmheads0 torch nn linear skip as this model doesn t support all arguments tested def testmodeloutputsequivalenceself pass skip as this model has multiple inputs embeds and lm heads that should not be tied def testtiemodelweightsself pass skip as this model has multiple inputs embeds and lm heads that should not be tied def testtiedweightskeysself pass def getinputidsandconfigself batchsize2 config inputsdict self modeltester prepareconfigandinputsforcommon inputids inputsdictinputids take max batchsize sequencelength inputids shape1 inputids inputids batchsize config numcodebooks generate max 3 tokens maxlength inputids shape1 3 attentionmask torch onesbatchsize sequencelength dtypetorch long return config inputids attentionmask maxlength staticmethod def getlogitsprocessorandkwargs inputlength eostokenid forcedbostokenidnone forcedeostokenidnone maxlengthnone diversitypenaltynone processkwargs minlength inputlength 1 if maxlength is none else maxlength 1 logitsprocessor logitsprocessorlist return processkwargs logitsprocessor override since we don t expect the outputs of generate and greedysearch to be the same since we perform additional postprocessing in the former def testgreedygeneratedictoutputsself for modelclass in self greedysamplemodelclasses disable cache config inputids attentionmask maxlength self getinputidsandconfig config usecache false model modelclassconfig totorchdevice eval outputgreedy outputgenerate self greedygenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice maxlengthmaxlength outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputgreedy greedysearchdecoderonlyoutput self assertisinstanceoutputgenerate greedysearchdecoderonlyoutput self assertnotinconfig padtokenid outputgenerate override since we don t expect the outputs of generate and greedysearch to be the same since we perform additional postprocessing in the former def testgreedygeneratedictoutputsusecacheself for modelclass in self greedysamplemodelclasses enable cache config inputids attentionmask maxlength self getinputidsandconfig config usecache true config isdecoder true model modelclassconfig totorchdevice eval outputgreedy outputgenerate self greedygenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice maxlengthmaxlength outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputgreedy greedysearchdecoderonlyoutput self assertisinstanceoutputgenerate greedysearchdecoderonlyoutput override since we don t expect the outputs of generate and sample to be the same since we perform additional postprocessing in the former def testsamplegenerateself for modelclass in self greedysamplemodelclasses config inputids attentionmask maxlength self getinputidsandconfig model modelclassconfig totorchdevice eval processkwargs logitsprocessor self getlogitsprocessorandkwargs inputids shape1 model config eostokenid forcedbostokenidmodel config forcedbostokenid forcedeostokenidmodel config forcedeostokenid maxlengthmaxlength logitswarperkwargs logitswarper self getwarperandkwargsnumbeams2 check generate and sample are equal outputsample outputgenerate self samplegenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice maxlengthmaxlength numreturnsequences3 logitsprocessorlogitsprocessor logitswarperlogitswarper logitswarperkwargslogitswarperkwargs processkwargsprocesskwargs self assertisinstanceoutputsample torch tensor self assertisinstanceoutputgenerate torch tensor override since we don t expect the outputs of generate and sample to be the same since we perform additional postprocessing in the former def testsamplegeneratedictoutputself for modelclass in self greedysamplemodelclasses disable cache config inputids attentionmask maxlength self getinputidsandconfig config usecache false model modelclassconfig totorchdevice eval processkwargs logitsprocessor self getlogitsprocessorandkwargs inputids shape1 model config eostokenid forcedbostokenidmodel config forcedbostokenid forcedeostokenidmodel config forcedeostokenid maxlengthmaxlength logitswarperkwargs logitswarper self getwarperandkwargsnumbeams1 outputsample outputgenerate self samplegenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice maxlengthmaxlength numreturnsequences1 logitsprocessorlogitsprocessor logitswarperlogitswarper logitswarperkwargslogitswarperkwargs processkwargsprocesskwargs outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputsample sampledecoderonlyoutput self assertisinstanceoutputgenerate sampledecoderonlyoutput def testgreedygeneratestereooutputsself for modelclass in self greedysamplemodelclasses config inputids attentionmask maxlength self getinputidsandconfig config audiochannels 2 model modelclassconfig totorchdevice eval outputgreedy outputgenerate self greedygenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice maxlengthmaxlength outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputgreedy greedysearchdecoderonlyoutput self assertisinstanceoutputgenerate greedysearchdecoderonlyoutput self assertnotinconfig padtokenid outputgenerate def preparemusicgeninputsdict config inputids decoderinputids attentionmasknone decoderattentionmasknone headmasknone decoderheadmasknone crossattnheadmasknone if decoderattentionmask is none decoderattentionmask decoderinputids reshape 1 config decoder numcodebooks decoderinputids shape1 0 decoderattentionmask decoderattentionmask neconfig decoder padtokenid if headmask is none headmask torch ones config textencoder numhiddenlayers config textencoder numattentionheads devicetorchdevice if decoderheadmask is none decoderheadmask torch ones config decoder numhiddenlayers config decoder numattentionheads devicetorchdevice if crossattnheadmask is none crossattnheadmask torch ones config decoder numhiddenlayers config decoder numattentionheads devicetorchdevice return inputids inputids attentionmask attentionmask decoderinputids decoderinputids decoderattentionmask decoderattentionmask headmask headmask decoderheadmask decoderheadmask crossattnheadmask crossattnheadmask class musicgentester def init self parent batchsize2 seqlength7 istrainingfalse uselabelsfalse vocabsize99 hiddensize16 numhiddenlayers2 numattentionheads4 intermediatesize4 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings100 padtokenid99 bostokenid99 numcodebooks4 numfilters4 codebooksize128 self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self maxpositionembeddings maxpositionembeddings self padtokenid padtokenid self bostokenid bostokenid self numcodebooks numcodebooks self numfilters numfilters self codebooksize codebooksize def prepareconfigandinputsself inputids idstensorself batchsize self seqlength self vocabsize decoderinputids idstensorself batchsize self numcodebooks self seqlength self vocabsize config self getconfig inputsdict preparemusicgeninputsdictconfig inputids decoderinputidsdecoderinputids return config inputsdict def getconfigself textencoderconfig t5config vocabsizeself vocabsize dmodelself hiddensize dffself intermediatesize numlayersself numhiddenlayers numheadsself numattentionheads audioencoderconfig encodecconfig hiddensizeself vocabsize compress1 numfiltersself numfilters codebooksizeself codebooksize codebookdimself vocabsize decoderconfig musicgendecoderconfig vocabsizeself vocabsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads ffndimself intermediatesize padtokenidself padtokenid decoderstarttokenidself bostokenid bostokenidself bostokenid numcodebooksself numcodebooks tiewordembeddingsfalse config musicgenconfig fromsubmodelsconfigtextencoderconfig audioencoderconfig decoderconfig return config def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict requiretorch class musicgentestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses musicgenforconditionalgeneration if istorchavailable else greedysamplemodelclasses musicgenforconditionalgeneration if istorchavailable else pipelinemodelmapping texttoaudio musicgenforconditionalgeneration if istorchavailable else testpruning false training is not supported yet for musicgen testheadmasking false testresizeembeddings false not to test torchscript as the model tester doesn t prepare inputvalues and paddingmask and torchscript hates none values testtorchscript false def setupself self modeltester musicgentesterself def checkoutputwithattentionsself outputs config inputids decoderinputids textencoderconfig config textencoder decoderconfig config decoder encoderattentions outputsencoderattentions self assertequallenencoderattentions textencoderconfig numhiddenlayers self assertequal encoderattentions0 shape3 textencoderconfig numattentionheads inputids shape1 inputids shape1 decoderattentions outputsdecoderattentions numdecoderlayers decoderconfig numhiddenlayers self assertequallendecoderattentions numdecoderlayers self assertequal decoderattentions0 shape3 decoderconfig numattentionheads decoderinputids shape1 decoderinputids shape1 crossattentions outputscrossattentions self assertequallencrossattentions numdecoderlayers crossattentioninputseqlen decoderinputids shape1 self assertequal crossattentions0 shape3 decoderconfig numattentionheads crossattentioninputseqlen inputids shape1 def checkmusicgenmodeloutputattentions self modelclass config inputids attentionmask decoderinputids decoderattentionmask kwargs model modelclassconfig model totorchdevice model eval with torch nograd outputs model inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask outputattentionstrue kwargs self checkoutputwithattentionsoutputs config inputids decoderinputids def checkmusicgenmodeloutputattentionsfromconfig self modelclass config inputids attentionmask decoderinputids decoderattentionmask kwargs similar to checkmusicgenmodeloutputattentions but with outputattentions triggered from the config file contrarily to most models changing the model s config won t work the defaults are loaded from the inner models configurations config outputattentions true model config won t work model modelclassconfig model totorchdevice model eval with torch nograd outputs model inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask kwargs self asserttrue allkey not in outputs for key in encoderattentions decoderattentions crossattentions config textencoder outputattentions true inner model config will work config audioencoder outputattentions true config decoder outputattentions true model modelclassconfig model totorchdevice model eval with torch nograd outputs model inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask kwargs self checkoutputwithattentionsoutputs config inputids decoderinputids override since changing outputattentions from the toplevel model config won t work def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses self checkmusicgenmodeloutputattentionsmodelclass config inputsdict self checkmusicgenmodeloutputattentionsfromconfigmodelclass config inputsdict override since we have a specific forward signature for musicgen def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel forward signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames inputids attentionmask inputvalues paddingmask decoderinputids decoderattentionmask expectedargnames extend headmask decoderheadmask crossattnheadmask encoderoutputs if headmask and decoderheadmask and crossattnheadmask in argnames else encoderoutputs self assertlistequalargnames lenexpectedargnames expectedargnames override since changing gradientcheckpointing from the toplevel model config won t work def testgradientcheckpointingbackwardcompatibilityself config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses if not modelclass supportsgradientcheckpointing continue config textencoder gradientcheckpointing true config audioencoder gradientcheckpointing true config decoder gradientcheckpointing true model modelclassconfig self asserttruemodel isgradientcheckpointing skip as this model has multiple inputs embeds and lm heads that should not be tied def testtiemodelweightsself pass skip as this model has multiple inputs embeds and lm heads that should not be tied def testtiedmodelweightskeyignoreself pass skip as this model has multiple inputs embeds and lm heads that should not be tied def testtiedweightskeysself pass override since changing outputhiddenstates outputattentions from the toplevel model config won t work def testretaingradhiddenstatesattentionsself config inputsdict self modeltester prepareconfigandinputsforcommon config textencoder outputhiddenstates true config audioencoder outputhiddenstates true config decoder outputhiddenstates true config textencoder outputattentions true config decoder outputattentions true no need to test all models as different heads yield the same functionality modelclass self allmodelclasses0 model modelclassconfig model totorchdevice inputs self prepareforclassinputsdict modelclass outputs modelinputs output outputs0 encoderhiddenstates outputs encoderhiddenstates0 encoderhiddenstates retaingrad decoderhiddenstates outputs decoderhiddenstates0 decoderhiddenstates retaingrad if self hasattentions encoderattentions outputs encoderattentions0 encoderattentions retaingrad decoderattentions outputs decoderattentions0 decoderattentions retaingrad crossattentions outputs crossattentions0 crossattentions retaingrad output flatten0 backwardretaingraphtrue self assertisnotnoneencoderhiddenstates grad self assertisnotnonedecoderhiddenstates grad if self hasattentions self assertisnotnoneencoderattentions grad self assertisnotnonedecoderattentions grad self assertisnotnonecrossattentions grad override since changing outputhiddenstates from the toplevel model config won t work def testhiddenstatesoutputself def checkhiddenstatesoutputinputsdict config modelclass model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass hiddenstates outputs encoderhiddenstates expectednumlayers self modeltester numhiddenlayers 1 self assertequallenhiddenstates expectednumlayers seqlength self modeltester seqlength self assertlistequal listhiddenstates0 shape2 seqlength self modeltester hiddensize hiddenstates outputs decoderhiddenstates self assertisinstancehiddenstates list tuple self assertequallenhiddenstates expectednumlayers self assertlistequal listhiddenstates0 shape2 seqlength self modeltester hiddensize config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses inputsdictoutputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config textencoder outputhiddenstates true config audioencoder outputhiddenstates true config decoder outputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass override since the conv layers and lstm s in encodec are exceptions def testinitializationself config inputsdict self modeltester prepareconfigandinputsforcommon configsnoinit configzeroinitconfig for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit for name param in model namedparameters uniforminitparms conv ignoreinit lstm if param requiresgrad if anyx in name for x in uniforminitparms self asserttrue 1 0 param data mean 1e9 round 1e9 item 1 0 msgfparameter name of model modelclass seems not properly initialized elif not anyx in name for x in ignoreinit self assertin param data mean 1e9 round 1e9 item 0 0 1 0 msgfparameter name of model modelclass seems not properly initialized override since we have embeddings lm heads over multiple codebooks def testmodelcommonattributesself config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig self assertisinstancemodel getinputembeddings torch nn embedding lmheads model getoutputembeddings self asserttruelmheads is none or isinstancelmheads0 torch nn linear def getinputidsandconfigself batchsize2 config inputsdict self modeltester prepareconfigandinputsforcommon inputids inputsdictinputids take max batchsize sequencelength inputids shape1 inputids inputids batchsize attentionmask torch onesbatchsize sequencelength dtypetorch long generate max 3 tokens decoderinputids inputsdictdecoderinputids maxlength decoderinputids shape1 3 decoderinputids decoderinputids batchsize config decoder numcodebooks return config inputids attentionmask decoderinputids maxlength override since the inputids cannot be used as the decoderinputids for musicgen input outputs are different modalities different shapes def greedygenerate self model inputids attentionmask decoderinputids maxlength outputscoresfalse outputattentionsfalse outputhiddenstatesfalse returndictingeneratefalse logitsprocesskwargs logitsprocessor self getlogitsprocessorandkwargs inputids shape1 eostokenidmodel config eostokenid forcedbostokenidmodel config forcedbostokenid forcedeostokenidmodel config forcedeostokenid maxlengthmaxlength modelkwargs attentionmask attentionmask if attentionmask is not none else outputgenerate model generate inputids dosamplefalse numbeams1 maxlengthmaxlength outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates outputscoresoutputscores returndictingeneratereturndictingenerate removeinvalidvaluestrue logitsprocesskwargs modelkwargs encoderoutputs inputids attentionmask self getencoderoutputs model inputids attentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates with torch nograd modelkwargs attentionmask attentionmask if attentionmask is not none else outputgreedy model greedysearch decoderinputids maxlengthmaxlength logitsprocessorlogitsprocessor outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates outputscoresoutputscores returndictingeneratereturndictingenerate encoderoutputsencoderoutputs modelkwargs return outputgreedy outputgenerate override since the inputids cannot be used as the decoderinputids for musicgen input outputs are different modalities different shapes def samplegenerate self model inputids attentionmask decoderinputids maxlength numreturnsequences logitsprocessor logitswarper logitswarperkwargs processkwargs outputscoresfalse outputattentionsfalse outputhiddenstatesfalse returndictingeneratefalse torch manualseed0 modelkwargs attentionmask attentionmask if attentionmask is not none else outputgenerate model generate inputids dosampletrue numbeams1 maxlengthmaxlength numreturnsequencesnumreturnsequences outputscoresoutputscores outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictingeneratereturndictingenerate removeinvalidvaluestrue logitswarperkwargs processkwargs modelkwargs torch manualseed0 encoderoutputs inputids attentionmask self getencoderoutputs model inputids attentionmask numinterleavenumreturnsequences outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates prevent flaky generation test failures logitsprocessor appendinfnanremovelogitsprocessor with torch nograd modelkwargs attentionmask attentionmask if attentionmask is not none else outputsample model sample decoderinputids repeatinterleavenumreturnsequences dim0 maxlengthmaxlength logitsprocessorlogitsprocessor logitswarperlogitswarper outputscoresoutputscores outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictingeneratereturndictingenerate encoderoutputsencoderoutputs modelkwargs return outputsample outputgenerate staticmethod def getlogitsprocessorandkwargs inputlength eostokenid forcedbostokenidnone forcedeostokenidnone maxlengthnone diversitypenaltynone processkwargs minlength inputlength 1 if maxlength is none else maxlength 1 logitsprocessor logitsprocessorlist return processkwargs logitsprocessor def testgreedygeneratedictoutputsself for modelclass in self greedysamplemodelclasses disable cache config inputids attentionmask decoderinputids maxlength self getinputidsandconfig config usecache false model modelclassconfig totorchdevice eval outputgreedy outputgenerate self greedygenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice decoderinputidsdecoderinputids maxlengthmaxlength outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputgreedy greedysearchencoderdecoderoutput self assertisinstanceoutputgenerate greedysearchencoderdecoderoutput self assertnotinconfig padtokenid outputgenerate def testgreedygeneratedictoutputsusecacheself for modelclass in self greedysamplemodelclasses enable cache config inputids attentionmask decoderinputids maxlength self getinputidsandconfig config usecache true config isdecoder true model modelclassconfig totorchdevice eval outputgreedy outputgenerate self greedygenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice decoderinputidsdecoderinputids maxlengthmaxlength outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputgreedy greedysearchencoderdecoderoutput self assertisinstanceoutputgenerate greedysearchencoderdecoderoutput def testsamplegenerateself for modelclass in self greedysamplemodelclasses config inputids attentionmask decoderinputids maxlength self getinputidsandconfig model modelclassconfig totorchdevice eval processkwargs logitsprocessor self getlogitsprocessorandkwargs inputids shape1 model config eostokenid forcedbostokenidmodel config forcedbostokenid forcedeostokenidmodel config forcedeostokenid maxlengthmaxlength logitswarperkwargs logitswarper self getwarperandkwargsnumbeams2 check generate and sample are equal outputsample outputgenerate self samplegenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice decoderinputidsdecoderinputids maxlengthmaxlength numreturnsequences1 logitsprocessorlogitsprocessor logitswarperlogitswarper logitswarperkwargslogitswarperkwargs processkwargsprocesskwargs self assertisinstanceoutputsample torch tensor self assertisinstanceoutputgenerate torch tensor def testsamplegeneratedictoutputself for modelclass in self greedysamplemodelclasses disable cache config inputids attentionmask decoderinputids maxlength self getinputidsandconfig config usecache false model modelclassconfig totorchdevice eval processkwargs logitsprocessor self getlogitsprocessorandkwargs inputids shape1 model config eostokenid forcedbostokenidmodel config forcedbostokenid forcedeostokenidmodel config forcedeostokenid maxlengthmaxlength logitswarperkwargs logitswarper self getwarperandkwargsnumbeams1 outputsample outputgenerate self samplegenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice decoderinputidsdecoderinputids maxlengthmaxlength numreturnsequences3 logitsprocessorlogitsprocessor logitswarperlogitswarper logitswarperkwargslogitswarperkwargs processkwargsprocesskwargs outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputsample sampleencoderdecoderoutput self assertisinstanceoutputgenerate sampleencoderdecoderoutput def testgeneratewithoutinputidsself config maxlength self getinputidsandconfig if no bos token id cannot generate from none if config bostokenid is none return for modelclass in self greedysamplemodelclasses model modelclassconfig totorchdevice model eval outputidsgenerate model generatedosamplefalse maxlengthmaxlength removeinvalidvaluestrue self assertisnotnoneoutputidsgenerate requiretorchfp16 def testgeneratefp16self config inputdict self modeltester prepareconfigandinputs for modelclass in self greedysamplemodelclasses model modelclassconfig eval totorchdevice model half greedy model generateinputdictinputids attentionmaskinputdictattentionmask maxnewtokens10 sampling model generate inputdictinputids attentionmaskinputdictattentionmask dosampletrue maxnewtokens10 def testgreedygeneratestereooutputsself for modelclass in self greedysamplemodelclasses config inputids attentionmask decoderinputids maxlength self getinputidsandconfig config audiochannels 2 model modelclassconfig totorchdevice eval outputgreedy outputgenerate self greedygenerate modelmodel inputidsinputids totorchdevice attentionmaskattentionmask totorchdevice decoderinputidsdecoderinputids maxlengthmaxlength outputscorestrue outputhiddenstatestrue outputattentionstrue returndictingeneratetrue self assertisinstanceoutputgreedy greedysearchencoderdecoderoutput self assertisinstanceoutputgenerate greedysearchencoderdecoderoutput self assertnotinconfig padtokenid outputgenerate def getbipbipbipduration0 125 duration0 5 samplerate32000 prepare the encoder inputs prepare the decoder inputs fmt off fmt on prepare the text encoder inputs prepare the audio encoder inputs fmt off fmt on only generate 1 sample with greedy since it s deterministic all elements of the batch will be the same fmt off fmt on for stochastic sampling we can generate multiple outputs fmt off fmt on prepare the encoder inputs fmt off fmt on prepare the encoder inputs fmt off fmt on prepare the encoder inputs fmt off fmt on fmt off fmt on only generate 1 sample with greedy since it s deterministic all elements of the batch will be the same fmt off fmt on bsz channels seqlen create stereo inputs fmt off fmt on bsz channels seqlen input values take shape 32000 and we generate from there we check the last generated values coding utf 8 2021 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 testing suite for the pytorch musicgen model we don t want to run all the generation tests only a specific subset override since we have to compute the input embeddings over codebooks override since we have embeddings lm heads over multiple codebooks skip as this model doesn t support all arguments tested skip as this model has multiple inputs embeds and lm heads that should not be tied skip as this model has multiple inputs embeds and lm heads that should not be tied take max batch_size generate max 3 tokens override since we don t expect the outputs of generate and greedy_search to be the same since we perform additional post processing in the former disable cache override since we don t expect the outputs of generate and greedy_search to be the same since we perform additional post processing in the former enable cache override since we don t expect the outputs of generate and sample to be the same since we perform additional post processing in the former check generate and sample are equal override since we don t expect the outputs of generate and sample to be the same since we perform additional post processing in the former disable cache training is not supported yet for musicgen not to test torchscript as the model tester doesn t prepare input_values and padding_mask and torchscript hates none values similar to check_musicgen_model_output_attentions but with output_attentions triggered from the config file contrarily to most models changing the model s config won t work the defaults are loaded from the inner models configurations model config won t work inner model config will work override since changing output_attentions from the top level model config won t work override since we have a specific forward signature for musicgen signature parameters is an ordereddict so arg_names order is deterministic override since changing gradient_checkpointing from the top level model config won t work skip as this model has multiple inputs embeds and lm heads that should not be tied skip as this model has multiple inputs embeds and lm heads that should not be tied skip as this model has multiple inputs embeds and lm heads that should not be tied override since changing output_hidden_states output_attentions from the top level model config won t work no need to test all models as different heads yield the same functionality override since changing output_hidden_states from the top level model config won t work check that output_hidden_states also work using config override since the conv layers and lstm s in encodec are exceptions override since we have embeddings lm heads over multiple codebooks take max batch_size generate max 3 tokens override since the input_ids cannot be used as the decoder_input_ids for musicgen input outputs are different modalities different shapes override since the input_ids cannot be used as the decoder_input_ids for musicgen input outputs are different modalities different shapes prevent flaky generation test failures disable cache enable cache check generate and sample are equal disable cache if no bos token id cannot generate from none greedy sampling produces a series of bip bip sounds at a given frequency prepare the encoder inputs prepare the decoder inputs fmt off fmt on prepare the text encoder inputs prepare the audio encoder inputs fmt off fmt on only generate 1 sample with greedy since it s deterministic all elements of the batch will be the same fmt off fmt on for stochastic sampling we can generate multiple outputs fmt off fmt on prepare the encoder inputs fmt off fmt on prepare the encoder inputs fmt off fmt on prepare the encoder inputs fmt off fmt on fmt off fmt on input values take shape 32000 and we generate from there only generate 1 sample with greedy since it s deterministic all elements of the batch will be the same fmt off fmt on bsz channels seq_len create stereo inputs fmt off fmt on bsz channels seq_len input values take shape 32000 and we generate from there we check the last generated values
import copy
import inspect
import math
import unittest

import numpy as np

from transformers import (
    EncodecConfig,
    MusicgenConfig,
    MusicgenDecoderConfig,
    MusicgenProcessor,
    PretrainedConfig,
    T5Config,
)
from transformers.testing_utils import (
    is_torch_available,
    require_torch,
    require_torch_fp16,
    slow,
    torch_device,
)
from transformers.utils import cached_property

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MusicgenForCausalLM,
        MusicgenForConditionalGeneration,
        MusicgenModel,
        set_seed,
    )
    from transformers.generation import (
        GreedySearchDecoderOnlyOutput,
        GreedySearchEncoderDecoderOutput,
        InfNanRemoveLogitsProcessor,
        LogitsProcessorList,
        SampleDecoderOnlyOutput,
        SampleEncoderDecoderOutput,
    )


def _config_zero_init(config):
    configs_no_init = copy.deepcopy(config)
    for key in configs_no_init.__dict__.keys():
        if "_range" in key or "_std" in key or "initializer_factor" in key or "layer_scale" in key:
            setattr(configs_no_init, key, 1e-10)
        if isinstance(getattr(configs_no_init, key, None), PretrainedConfig):
            no_init_subconfig = _config_zero_init(getattr(configs_no_init, key))
            setattr(configs_no_init, key, no_init_subconfig)
    return configs_no_init


def prepare_musicgen_decoder_inputs_dict(
    config,
    input_ids,
    attention_mask=None,
    head_mask=None,
    encoder_hidden_states=None,
    encoder_attention_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = input_ids.reshape(-1, config.num_codebooks, input_ids.shape[-1])[:, 0, :]
        attention_mask = attention_mask.ne(config.pad_token_id)
    if head_mask is None:
        head_mask = torch.ones(config.num_hidden_layers, config.num_attention_heads, device=torch_device)
    if encoder_attention_mask is None and encoder_hidden_states is not None:
        encoder_attention_mask = torch.ones(encoder_hidden_states.shape[:2], device=torch_device)
    if cross_attn_head_mask is None:
        cross_attn_head_mask = torch.ones(config.num_hidden_layers, config.num_attention_heads, device=torch_device)
    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,
    }


class MusicgenDecoderTester:
    def __init__(
        self,
        parent,
        batch_size=2,
        seq_length=7,
        is_training=False,
        use_labels=False,
        vocab_size=99,
        hidden_size=16,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=4,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=100,
        pad_token_id=99,
        bos_token_id=99,
        num_codebooks=4,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.num_codebooks = num_codebooks

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size * self.num_codebooks, self.seq_length], self.vocab_size)
        encoder_hidden_states = floats_tensor([self.batch_size, self.seq_length, self.hidden_size])

        config = self.get_config()
        inputs_dict = prepare_musicgen_decoder_inputs_dict(
            config, input_ids, encoder_hidden_states=encoder_hidden_states
        )
        return config, inputs_dict

    def get_config(self):
        config = MusicgenDecoderConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            d_ff=self.intermediate_size,
            pad_token_id=self.pad_token_id,
            decoder_start_token_id=self.bos_token_id,
            bos_token_id=self.bos_token_id,
            num_codebooks=self.num_codebooks,
            tie_word_embeddings=False,
        )
        return config

    def prepare_config_and_inputs_for_common(self):
        config, inputs_dict = self.prepare_config_and_inputs()
        return config, inputs_dict


@require_torch
class MusicgenDecoderTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (MusicgenModel, MusicgenForCausalLM) if is_torch_available() else ()
    greedy_sample_model_classes = (
        (MusicgenForCausalLM,) if is_torch_available() else ()
    )  
    pipeline_model_mapping = {}
    test_pruning = False
    test_resize_embeddings = False

    def setUp(self):
        self.model_tester = MusicgenDecoderTester(self)
        self.config_tester = ConfigTester(self, config_class=MusicgenDecoderConfig, hidden_size=16)

    def test_config(self):
        self.config_tester.run_common_tests()

    
    def test_inputs_embeds(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            inputs = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))

            input_ids = inputs["input_ids"]
            del inputs["input_ids"]

            embed_tokens = model.get_input_embeddings()

            input_ids = input_ids.reshape(-1, config.num_codebooks, input_ids.shape[-1])

            inputs["inputs_embeds"] = sum(
                [embed_tokens[codebook](input_ids[:, codebook]) for codebook in range(config.num_codebooks)]
            )

            with torch.no_grad():
                model(**inputs)[0]

    
    def test_model_common_attributes(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            first_embed = model.get_input_embeddings()[0]
            self.assertIsInstance(first_embed, torch.nn.Embedding)
            lm_heads = model.get_output_embeddings()
            self.assertTrue(lm_heads is None or isinstance(lm_heads[0], torch.nn.Linear))

    
    def test_model_outputs_equivalence(self):
        pass

    
    def test_tie_model_weights(self):
        pass

    
    def test_tied_weights_keys(self):
        pass

    def _get_input_ids_and_config(self, batch_size=2):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        input_ids = inputs_dict["input_ids"]

        
        sequence_length = input_ids.shape[-1]
        input_ids = input_ids[: batch_size * config.num_codebooks, :]

        
        max_length = input_ids.shape[-1] + 3
        attention_mask = torch.ones((batch_size, sequence_length), dtype=torch.long)
        return config, input_ids, attention_mask, max_length

    @staticmethod
    def _get_logits_processor_and_kwargs(
        input_length,
        eos_token_id,
        forced_bos_token_id=None,
        forced_eos_token_id=None,
        max_length=None,
        diversity_penalty=None,
    ):
        process_kwargs = {
            "min_length": input_length + 1 if max_length is None else max_length - 1,
        }
        logits_processor = LogitsProcessorList()
        return process_kwargs, logits_processor

    
    
    def test_greedy_generate_dict_outputs(self):
        for model_class in self.greedy_sample_model_classes:
            
            config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
            output_greedy, output_generate = self._greedy_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                max_length=max_length,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_greedy, GreedySearchDecoderOnlyOutput)
            self.assertIsInstance(output_generate, GreedySearchDecoderOnlyOutput)

            self.assertNotIn(config.pad_token_id, output_generate)

    
    
    def test_greedy_generate_dict_outputs_use_cache(self):
        for model_class in self.greedy_sample_model_classes:
            
            config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()

            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
            output_greedy, output_generate = self._greedy_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                max_length=max_length,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_greedy, GreedySearchDecoderOnlyOutput)
            self.assertIsInstance(output_generate, GreedySearchDecoderOnlyOutput)

    
    
    def test_sample_generate(self):
        for model_class in self.greedy_sample_model_classes:
            config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
            model = model_class(config).to(torch_device).eval()

            process_kwargs, logits_processor = self._get_logits_processor_and_kwargs(
                input_ids.shape[-1],
                model.config.eos_token_id,
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
                max_length=max_length,
            )
            logits_warper_kwargs, logits_warper = self._get_warper_and_kwargs(num_beams=2)

            
            output_sample, output_generate = self._sample_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                max_length=max_length,
                num_return_sequences=3,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
            )
            self.assertIsInstance(output_sample, torch.Tensor)
            self.assertIsInstance(output_generate, torch.Tensor)

    
    
    def test_sample_generate_dict_output(self):
        for model_class in self.greedy_sample_model_classes:
            
            config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()

            process_kwargs, logits_processor = self._get_logits_processor_and_kwargs(
                input_ids.shape[-1],
                model.config.eos_token_id,
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
                max_length=max_length,
            )
            logits_warper_kwargs, logits_warper = self._get_warper_and_kwargs(num_beams=1)

            output_sample, output_generate = self._sample_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                max_length=max_length,
                num_return_sequences=1,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_sample, SampleDecoderOnlyOutput)
            self.assertIsInstance(output_generate, SampleDecoderOnlyOutput)

    def test_greedy_generate_stereo_outputs(self):
        for model_class in self.greedy_sample_model_classes:
            config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
            config.audio_channels = 2
            model = model_class(config).to(torch_device).eval()
            output_greedy, output_generate = self._greedy_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                max_length=max_length,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_greedy, GreedySearchDecoderOnlyOutput)
            self.assertIsInstance(output_generate, GreedySearchDecoderOnlyOutput)

            self.assertNotIn(config.pad_token_id, output_generate)


def prepare_musicgen_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if decoder_attention_mask is None:
        decoder_attention_mask = decoder_input_ids.reshape(
            -1, config.decoder.num_codebooks, decoder_input_ids.shape[-1]
        )[:, 0, :]
        decoder_attention_mask = decoder_attention_mask.ne(config.decoder.pad_token_id)
    if head_mask is None:
        head_mask = torch.ones(
            config.text_encoder.num_hidden_layers, config.text_encoder.num_attention_heads, device=torch_device
        )
    if decoder_head_mask is None:
        decoder_head_mask = torch.ones(
            config.decoder.num_hidden_layers, config.decoder.num_attention_heads, device=torch_device
        )
    if cross_attn_head_mask is None:
        cross_attn_head_mask = torch.ones(
            config.decoder.num_hidden_layers, config.decoder.num_attention_heads, device=torch_device
        )
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "decoder_input_ids": decoder_input_ids,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }


class MusicgenTester:
    def __init__(
        self,
        parent,
        batch_size=2,
        seq_length=7,
        is_training=False,
        use_labels=False,
        vocab_size=99,
        hidden_size=16,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=4,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=100,
        pad_token_id=99,
        bos_token_id=99,
        num_codebooks=4,
        num_filters=4,
        codebook_size=128,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.num_codebooks = num_codebooks
        self.num_filters = num_filters
        self.codebook_size = codebook_size

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
        decoder_input_ids = ids_tensor([self.batch_size * self.num_codebooks, self.seq_length], self.vocab_size)

        config = self.get_config()
        inputs_dict = prepare_musicgen_inputs_dict(config, input_ids, decoder_input_ids=decoder_input_ids)
        return config, inputs_dict

    def get_config(self):
        text_encoder_config = T5Config(
            vocab_size=self.vocab_size,
            d_model=self.hidden_size,
            d_ff=self.intermediate_size,
            num_layers=self.num_hidden_layers,
            num_heads=self.num_attention_heads,
        )
        audio_encoder_config = EncodecConfig(
            hidden_size=self.vocab_size,
            compress=1,
            num_filters=self.num_filters,
            codebook_size=self.codebook_size,
            codebook_dim=self.vocab_size,
        )
        decoder_config = MusicgenDecoderConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            ffn_dim=self.intermediate_size,
            pad_token_id=self.pad_token_id,
            decoder_start_token_id=self.bos_token_id,
            bos_token_id=self.bos_token_id,
            num_codebooks=self.num_codebooks,
            tie_word_embeddings=False,
        )
        config = MusicgenConfig.from_sub_models_config(text_encoder_config, audio_encoder_config, decoder_config)
        return config

    def prepare_config_and_inputs_for_common(self):
        config, inputs_dict = self.prepare_config_and_inputs()
        return config, inputs_dict


@require_torch
class MusicgenTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (MusicgenForConditionalGeneration,) if is_torch_available() else ()
    greedy_sample_model_classes = (MusicgenForConditionalGeneration,) if is_torch_available() else ()
    pipeline_model_mapping = {"text-to-audio": MusicgenForConditionalGeneration} if is_torch_available() else {}
    test_pruning = False  
    test_headmasking = False
    test_resize_embeddings = False
    
    
    test_torchscript = False

    def setUp(self):
        self.model_tester = MusicgenTester(self)

    def _check_output_with_attentions(self, outputs, config, input_ids, decoder_input_ids):
        text_encoder_config = config.text_encoder
        decoder_config = config.decoder

        encoder_attentions = outputs["encoder_attentions"]
        self.assertEqual(len(encoder_attentions), text_encoder_config.num_hidden_layers)

        self.assertEqual(
            encoder_attentions[0].shape[-3:],
            (text_encoder_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1]),
        )

        decoder_attentions = outputs["decoder_attentions"]
        num_decoder_layers = decoder_config.num_hidden_layers
        self.assertEqual(len(decoder_attentions), num_decoder_layers)

        self.assertEqual(
            decoder_attentions[0].shape[-3:],
            (decoder_config.num_attention_heads, decoder_input_ids.shape[-1], decoder_input_ids.shape[-1]),
        )

        cross_attentions = outputs["cross_attentions"]
        self.assertEqual(len(cross_attentions), num_decoder_layers)

        cross_attention_input_seq_len = decoder_input_ids.shape[-1]
        self.assertEqual(
            cross_attentions[0].shape[-3:],
            (decoder_config.num_attention_heads, cross_attention_input_seq_len, input_ids.shape[-1]),
        )

    def check_musicgen_model_output_attentions(
        self,
        model_class,
        config,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        **kwargs,
    ):
        model = model_class(config)
        model.to(torch_device)
        model.eval()

        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                decoder_input_ids=decoder_input_ids,
                attention_mask=attention_mask,
                decoder_attention_mask=decoder_attention_mask,
                output_attentions=True,
                **kwargs,
            )
        self._check_output_with_attentions(outputs, config, input_ids, decoder_input_ids)

    def check_musicgen_model_output_attentions_from_config(
        self,
        model_class,
        config,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        **kwargs,
    ):
        
        
        
        config.output_attentions = True  

        model = model_class(config)
        model.to(torch_device)
        model.eval()

        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                decoder_input_ids=decoder_input_ids,
                attention_mask=attention_mask,
                decoder_attention_mask=decoder_attention_mask,
                **kwargs,
            )
        self.assertTrue(
            all(key not in outputs for key in ["encoder_attentions", "decoder_attentions", "cross_attentions"])
        )
        config.text_encoder.output_attentions = True  
        config.audio_encoder.output_attentions = True
        config.decoder.output_attentions = True

        model = model_class(config)
        model.to(torch_device)
        model.eval()

        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                decoder_input_ids=decoder_input_ids,
                attention_mask=attention_mask,
                decoder_attention_mask=decoder_attention_mask,
                **kwargs,
            )
        self._check_output_with_attentions(outputs, config, input_ids, decoder_input_ids)

    
    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            self.check_musicgen_model_output_attentions(model_class, config, **inputs_dict)
            self.check_musicgen_model_output_attentions_from_config(model_class, config, **inputs_dict)

    
    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = [
                "input_ids",
                "attention_mask",
                "input_values",
                "padding_mask",
                "decoder_input_ids",
                "decoder_attention_mask",
            ]
            expected_arg_names.extend(
                ["head_mask", "decoder_head_mask", "cross_attn_head_mask", "encoder_outputs"]
                if "head_mask" and "decoder_head_mask" and "cross_attn_head_mask" in arg_names
                else ["encoder_outputs"]
            )
            self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names)

    
    def test_gradient_checkpointing_backward_compatibility(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            if not model_class.supports_gradient_checkpointing:
                continue

            config.text_encoder.gradient_checkpointing = True
            config.audio_encoder.gradient_checkpointing = True
            config.decoder.gradient_checkpointing = True
            model = model_class(config)
            self.assertTrue(model.is_gradient_checkpointing)

    
    def test_tie_model_weights(self):
        pass

    
    def test_tied_model_weights_key_ignore(self):
        pass

    
    def test_tied_weights_keys(self):
        pass

    
    def test_retain_grad_hidden_states_attentions(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.text_encoder.output_hidden_states = True
        config.audio_encoder.output_hidden_states = True
        config.decoder.output_hidden_states = True

        config.text_encoder.output_attentions = True
        config.decoder.output_attentions = True

        
        model_class = self.all_model_classes[0]
        model = model_class(config)
        model.to(torch_device)

        inputs = self._prepare_for_class(inputs_dict, model_class)

        outputs = model(**inputs)

        output = outputs[0]

        encoder_hidden_states = outputs.encoder_hidden_states[0]
        encoder_hidden_states.retain_grad()

        decoder_hidden_states = outputs.decoder_hidden_states[0]
        decoder_hidden_states.retain_grad()

        if self.has_attentions:
            encoder_attentions = outputs.encoder_attentions[0]
            encoder_attentions.retain_grad()

            decoder_attentions = outputs.decoder_attentions[0]
            decoder_attentions.retain_grad()

            cross_attentions = outputs.cross_attentions[0]
            cross_attentions.retain_grad()

        output.flatten()[0].backward(retain_graph=True)

        self.assertIsNotNone(encoder_hidden_states.grad)
        self.assertIsNotNone(decoder_hidden_states.grad)

        if self.has_attentions:
            self.assertIsNotNone(encoder_attentions.grad)
            self.assertIsNotNone(decoder_attentions.grad)
            self.assertIsNotNone(cross_attentions.grad)

    
    def test_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            hidden_states = outputs.encoder_hidden_states

            expected_num_layers = self.model_tester.num_hidden_layers + 1
            self.assertEqual(len(hidden_states), expected_num_layers)

            seq_length = self.model_tester.seq_length
            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [seq_length, self.model_tester.hidden_size],
            )

            hidden_states = outputs.decoder_hidden_states
            self.assertIsInstance(hidden_states, (list, tuple))
            self.assertEqual(len(hidden_states), expected_num_layers)

            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [seq_length, self.model_tester.hidden_size],
            )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            
            del inputs_dict["output_hidden_states"]
            config.text_encoder.output_hidden_states = True
            config.audio_encoder.output_hidden_states = True
            config.decoder.output_hidden_states = True

            check_hidden_states_output(inputs_dict, config, model_class)

    
    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                uniform_init_parms = ["conv"]
                ignore_init = ["lstm"]
                if param.requires_grad:
                    if any(x in name for x in uniform_init_parms):
                        self.assertTrue(
                            -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )
                    elif not any(x in name for x in ignore_init):
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item(),
                            [0.0, 1.0],
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )

    
    def test_model_common_attributes(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertIsInstance(model.get_input_embeddings(), torch.nn.Embedding)
            lm_heads = model.get_output_embeddings()
            self.assertTrue(lm_heads is None or isinstance(lm_heads[0], torch.nn.Linear))

    def _get_input_ids_and_config(self, batch_size=2):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        input_ids = inputs_dict["input_ids"]

        
        sequence_length = input_ids.shape[-1]
        input_ids = input_ids[:batch_size, :]
        attention_mask = torch.ones((batch_size, sequence_length), dtype=torch.long)

        
        decoder_input_ids = inputs_dict["decoder_input_ids"]
        max_length = decoder_input_ids.shape[-1] + 3
        decoder_input_ids = decoder_input_ids[: batch_size * config.decoder.num_codebooks, :]
        return config, input_ids, attention_mask, decoder_input_ids, max_length

    
    
    def _greedy_generate(
        self,
        model,
        input_ids,
        attention_mask,
        decoder_input_ids,
        max_length,
        output_scores=False,
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs(
            input_ids.shape[-1],
            eos_token_id=model.config.eos_token_id,
            forced_bos_token_id=model.config.forced_bos_token_id,
            forced_eos_token_id=model.config.forced_eos_token_id,
            max_length=max_length,
        )

        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
        output_generate = model.generate(
            input_ids,
            do_sample=False,
            num_beams=1,
            max_length=max_length,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_scores=output_scores,
            return_dict_in_generate=return_dict_in_generate,
            remove_invalid_values=True,
            **logits_process_kwargs,
            **model_kwargs,
        )

        encoder_outputs, input_ids, attention_mask = self._get_encoder_outputs(
            model,
            input_ids,
            attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        with torch.no_grad():
            model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
            output_greedy = model.greedy_search(
                decoder_input_ids,
                max_length=max_length,
                logits_processor=logits_processor,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                output_scores=output_scores,
                return_dict_in_generate=return_dict_in_generate,
                encoder_outputs=encoder_outputs,
                **model_kwargs,
            )
        return output_greedy, output_generate

    
    
    def _sample_generate(
        self,
        model,
        input_ids,
        attention_mask,
        decoder_input_ids,
        max_length,
        num_return_sequences,
        logits_processor,
        logits_warper,
        logits_warper_kwargs,
        process_kwargs,
        output_scores=False,
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        torch.manual_seed(0)
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
        output_generate = model.generate(
            input_ids,
            do_sample=True,
            num_beams=1,
            max_length=max_length,
            num_return_sequences=num_return_sequences,
            output_scores=output_scores,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            remove_invalid_values=True,
            **logits_warper_kwargs,
            **process_kwargs,
            **model_kwargs,
        )

        torch.manual_seed(0)
        encoder_outputs, input_ids, attention_mask = self._get_encoder_outputs(
            model,
            input_ids,
            attention_mask,
            num_interleave=num_return_sequences,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        
        logits_processor.append(InfNanRemoveLogitsProcessor())

        with torch.no_grad():
            model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
            output_sample = model.sample(
                decoder_input_ids.repeat_interleave(num_return_sequences, dim=0),
                max_length=max_length,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                output_scores=output_scores,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict_in_generate=return_dict_in_generate,
                encoder_outputs=encoder_outputs,
                **model_kwargs,
            )

        return output_sample, output_generate

    @staticmethod
    def _get_logits_processor_and_kwargs(
        input_length,
        eos_token_id,
        forced_bos_token_id=None,
        forced_eos_token_id=None,
        max_length=None,
        diversity_penalty=None,
    ):
        process_kwargs = {
            "min_length": input_length + 1 if max_length is None else max_length - 1,
        }
        logits_processor = LogitsProcessorList()
        return process_kwargs, logits_processor

    def test_greedy_generate_dict_outputs(self):
        for model_class in self.greedy_sample_model_classes:
            
            config, input_ids, attention_mask, decoder_input_ids, max_length = self._get_input_ids_and_config()
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
            output_greedy, output_generate = self._greedy_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                decoder_input_ids=decoder_input_ids,
                max_length=max_length,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_greedy, GreedySearchEncoderDecoderOutput)
            self.assertIsInstance(output_generate, GreedySearchEncoderDecoderOutput)

            self.assertNotIn(config.pad_token_id, output_generate)

    def test_greedy_generate_dict_outputs_use_cache(self):
        for model_class in self.greedy_sample_model_classes:
            
            config, input_ids, attention_mask, decoder_input_ids, max_length = self._get_input_ids_and_config()

            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
            output_greedy, output_generate = self._greedy_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                decoder_input_ids=decoder_input_ids,
                max_length=max_length,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_greedy, GreedySearchEncoderDecoderOutput)
            self.assertIsInstance(output_generate, GreedySearchEncoderDecoderOutput)

    def test_sample_generate(self):
        for model_class in self.greedy_sample_model_classes:
            config, input_ids, attention_mask, decoder_input_ids, max_length = self._get_input_ids_and_config()
            model = model_class(config).to(torch_device).eval()

            process_kwargs, logits_processor = self._get_logits_processor_and_kwargs(
                input_ids.shape[-1],
                model.config.eos_token_id,
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
                max_length=max_length,
            )
            logits_warper_kwargs, logits_warper = self._get_warper_and_kwargs(num_beams=2)

            
            output_sample, output_generate = self._sample_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                decoder_input_ids=decoder_input_ids,
                max_length=max_length,
                num_return_sequences=1,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
            )
            self.assertIsInstance(output_sample, torch.Tensor)
            self.assertIsInstance(output_generate, torch.Tensor)

    def test_sample_generate_dict_output(self):
        for model_class in self.greedy_sample_model_classes:
            
            config, input_ids, attention_mask, decoder_input_ids, max_length = self._get_input_ids_and_config()
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()

            process_kwargs, logits_processor = self._get_logits_processor_and_kwargs(
                input_ids.shape[-1],
                model.config.eos_token_id,
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
                max_length=max_length,
            )
            logits_warper_kwargs, logits_warper = self._get_warper_and_kwargs(num_beams=1)

            output_sample, output_generate = self._sample_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                decoder_input_ids=decoder_input_ids,
                max_length=max_length,
                num_return_sequences=3,
                logits_processor=logits_processor,
                logits_warper=logits_warper,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_sample, SampleEncoderDecoderOutput)
            self.assertIsInstance(output_generate, SampleEncoderDecoderOutput)

    def test_generate_without_input_ids(self):
        config, _, _, _, max_length = self._get_input_ids_and_config()

        
        if config.bos_token_id is None:
            return

        for model_class in self.greedy_sample_model_classes:
            model = model_class(config).to(torch_device)
            model.eval()

            output_ids_generate = model.generate(do_sample=False, max_length=max_length, remove_invalid_values=True)
            self.assertIsNotNone(output_ids_generate)

    @require_torch_fp16
    def test_generate_fp16(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs()

        for model_class in self.greedy_sample_model_classes:
            model = model_class(config).eval().to(torch_device)
            model.half()
            
            model.generate(input_dict["input_ids"], attention_mask=input_dict["attention_mask"], max_new_tokens=10)
            
            model.generate(
                input_dict["input_ids"], attention_mask=input_dict["attention_mask"], do_sample=True, max_new_tokens=10
            )

    def test_greedy_generate_stereo_outputs(self):
        for model_class in self.greedy_sample_model_classes:
            config, input_ids, attention_mask, decoder_input_ids, max_length = self._get_input_ids_and_config()
            config.audio_channels = 2

            model = model_class(config).to(torch_device).eval()
            output_greedy, output_generate = self._greedy_generate(
                model=model,
                input_ids=input_ids.to(torch_device),
                attention_mask=attention_mask.to(torch_device),
                decoder_input_ids=decoder_input_ids,
                max_length=max_length,
                output_scores=True,
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            self.assertIsInstance(output_greedy, GreedySearchEncoderDecoderOutput)
            self.assertIsInstance(output_generate, GreedySearchEncoderDecoderOutput)

            self.assertNotIn(config.pad_token_id, output_generate)


def get_bip_bip(bip_duration=0.125, duration=0.5, sample_rate=32000):
    
    timesteps = np.arange(int(duration * sample_rate)) / sample_rate
    wav = np.cos(2 * math.pi * 440 * timesteps)
    time_period = (timesteps % (2 * bip_duration)) / (2 * bip_duration)
    envelope = time_period >= 0.5
    return wav * envelope


def place_dict_on_device(dict_to_place, device):
    for key in dict_to_place:
        if dict_to_place[key] is not None and isinstance(dict_to_place[key], torch.Tensor):
            dict_to_place[key] = dict_to_place[key].to(device)
    return dict_to_place


@require_torch
class MusicgenIntegrationTests(unittest.TestCase):
    @cached_property
    def model(self):
        return MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small").to(torch_device)

    @cached_property
    def processor(self):
        return MusicgenProcessor.from_pretrained("facebook/musicgen-small")

    @slow
    def test_logits_text_prompt(self):
        model = self.model
        processor = self.processor

        inputs = processor(text=["80s music", "Club techno"], padding=True, return_tensors="pt")

        
        input_ids = inputs.input_ids.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        
        pad_token_id = model.generation_config.pad_token_id
        decoder_input_ids = (
            torch.ones((input_ids.shape[0] * model.decoder.num_codebooks, 1), dtype=torch.long).to(torch_device)
            * pad_token_id
        )

        with torch.no_grad():
            logits = model(
                input_ids,
                attention_mask=attention_mask,
                decoder_input_ids=decoder_input_ids,
            ).logits

        
        EXPECTED_LOGITS = torch.tensor(
            [
                -0.9708, -3.0149, -4.6415, -1.4754, -0.2786, -2.3523, -2.6049, -6.7467,
                -1.0206, -3.2984, -3.3968, -1.5108, -1.5786, -3.1493, -1.1503, -0.0545,
            ]
        )
        

        self.assertTrue(logits.shape == (*decoder_input_ids.shape, model.decoder.config.vocab_size))
        self.assertTrue(torch.allclose(logits[0, 0, :16].cpu(), EXPECTED_LOGITS, atol=1e-4))

    @slow
    def test_logits_text_audio_prompt(self):
        model = self.model
        processor = self.processor

        audio = [get_bip_bip(duration=0.5), get_bip_bip(duration=1.0)]
        text = ["80s music", "Club techno"]

        inputs = processor(audio=audio, text=text, padding=True, return_tensors="pt")

        
        input_ids = inputs.input_ids.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        
        input_values = inputs.input_values.to(torch_device)
        padding_mask = inputs.padding_mask.to(torch_device)

        with torch.no_grad():
            logits = model(
                input_ids,
                attention_mask=attention_mask,
                input_values=input_values,
                padding_mask=padding_mask,
            ).logits

        
        EXPECTED_LOGITS = torch.tensor(
            [
                0.1841, -2.9324, -0.7898, 0.1857, 0.4971, -2.8685, -1.6525, -1.6541,
                2.7757, -2.5942, -3.0959, -1.0120, -1.0147, -0.4605, -0.8885, 0.6820,
            ]
        )
        

        self.assertTrue(logits.shape == (8, 50, 2048))
        self.assertTrue(torch.allclose(logits[0, -1, :16].cpu(), EXPECTED_LOGITS, atol=1e-4))

    @slow
    def test_generate_unconditional_greedy(self):
        model = self.model

        
        unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
        unconditional_inputs = place_dict_on_device(unconditional_inputs, device=torch_device)

        output_values = model.generate(**unconditional_inputs, do_sample=False, max_new_tokens=5)

        
        EXPECTED_VALUES = torch.tensor(
            [
                0.0056, 0.0064, 0.0063, 0.0054, 0.0042, 0.0033, 0.0024, 0.0015,
                0.0015, 0.0010, 0.0004, -0.0012, -0.0036, -0.0055, -0.0067, -0.0071,
            ]
        )
        

        self.assertTrue(output_values.shape == (1, 1, 3200))
        self.assertTrue(torch.allclose(output_values[0, 0, :16].cpu(), EXPECTED_VALUES, atol=1e-4))

    @slow
    def test_generate_unconditional_sampling(self):
        model = self.model

        
        unconditional_inputs = model.get_unconditional_inputs(num_samples=2)
        unconditional_inputs = place_dict_on_device(unconditional_inputs, device=torch_device)

        set_seed(0)
        output_values = model.generate(**unconditional_inputs, do_sample=True, max_new_tokens=10)

        
        EXPECTED_VALUES = torch.tensor(
            [
                -0.0099, -0.0140, 0.0079, 0.0080, -0.0046,  0.0065, -0.0068, -0.0185,
                 0.0105,  0.0059, 0.0329, 0.0249, -0.0204, -0.0341, -0.0465,  0.0053,
            ]
        )
        

        self.assertTrue(output_values.shape == (2, 1, 4480))
        self.assertTrue(torch.allclose(output_values[0, 0, :16].cpu(), EXPECTED_VALUES, atol=1e-4))

    @slow
    def test_generate_text_prompt_greedy(self):
        model = self.model
        processor = self.processor

        inputs = processor(text=["80s music", "Club techno"], padding=True, return_tensors="pt")

        
        input_ids = inputs.input_ids.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        output_values = model.generate(
            input_ids, attention_mask=attention_mask, do_sample=False, guidance_scale=None, max_new_tokens=10
        )

        
        EXPECTED_VALUES = torch.tensor(
            [
                -1.1998e-04, -2.2302e-04, 4.6296e-04, 1.0524e-03, 2.4827e-04,
                -4.0288e-05, -1.2468e-04, 4.9846e-05, 7.1485e-04, 4.4197e-04,
            ]
        )
        

        self.assertTrue(output_values.shape == (2, 1, 4480))
        self.assertTrue(torch.allclose(output_values[0, 0, :10].cpu(), EXPECTED_VALUES, atol=1e-4))

    @slow
    def test_generate_text_prompt_greedy_with_classifier_free_guidance(self):
        model = self.model
        processor = self.processor

        inputs = processor(text=["80s music", "Club techno"], padding=True, return_tensors="pt")

        
        input_ids = inputs.input_ids.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        output_values = model.generate(
            input_ids, attention_mask=attention_mask, do_sample=False, guidance_scale=3, max_new_tokens=10
        )

        
        EXPECTED_VALUES = torch.tensor(
            [
                0.0283, 0.0246, 0.0650, 0.0640, 0.0599, 0.0711, 0.0420, 0.0112,
                0.0511, 0.0746, 0.1363, 0.1213, 0.0185, -0.0578, -0.0908, 0.0443,
            ]
        )
        

        self.assertTrue(output_values.shape == (2, 1, 4480))
        self.assertTrue(torch.allclose(output_values[0, 0, :16].cpu(), EXPECTED_VALUES, atol=1e-4))

    @slow
    def test_generate_text_prompt_sampling(self):
        model = self.model
        processor = self.processor

        inputs = processor(text=["80s music", "Club techno"], padding=True, return_tensors="pt")

        
        input_ids = inputs.input_ids.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        set_seed(0)
        output_values = model.generate(
            input_ids, attention_mask=attention_mask, do_sample=True, guidance_scale=None, max_new_tokens=10
        )

        
        EXPECTED_VALUES = torch.tensor(
            [
                -0.0111, -0.0154, 0.0047, 0.0058, -0.0068,  0.0012, -0.0109, -0.0229,
                 0.0010, -0.0038, 0.0167, 0.0042, -0.0421, -0.0610, -0.0764, -0.0326,
            ]
        )
        

        self.assertTrue(output_values.shape == (2, 1, 4480))
        self.assertTrue(torch.allclose(output_values[0, 0, :16].cpu(), EXPECTED_VALUES, atol=1e-4))

    @slow
    def test_generate_text_audio_prompt(self):
        model = self.model
        processor = self.processor

        audio = [get_bip_bip(duration=0.5), get_bip_bip(duration=1.0)]
        text = ["80s music", "Club techno"]

        inputs = processor(audio=audio, text=text, padding=True, return_tensors="pt")
        inputs = place_dict_on_device(inputs, device=torch_device)

        output_values = model.generate(**inputs, do_sample=False, guidance_scale=None, max_new_tokens=10)

        
        EXPECTED_VALUES = torch.tensor(
            [
                -0.0036, -0.0130, -0.0261, -0.0384, -0.0557, -0.0718, -0.0680, -0.0632,
                -0.0529, -0.0403, -0.0289, -0.0198, -0.0136, -0.0101, -0.0095, -0.0040,
            ]
        )
        

        self.assertTrue(
            output_values.shape == (2, 1, 36480)
        )  
        self.assertTrue(torch.allclose(output_values[0, 0, -16:].cpu(), EXPECTED_VALUES, atol=1e-4))


@require_torch
class MusicgenStereoIntegrationTests(unittest.TestCase):
    @cached_property
    def model(self):
        return MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-stereo-small").to(torch_device)

    @cached_property
    def processor(self):
        return MusicgenProcessor.from_pretrained("facebook/musicgen-stereo-small")

    @slow
    def test_generate_unconditional_greedy(self):
        model = self.model

        
        unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
        unconditional_inputs = place_dict_on_device(unconditional_inputs, device=torch_device)

        output_values = model.generate(**unconditional_inputs, do_sample=False, max_new_tokens=12)

        
        EXPECTED_VALUES_LEFT = torch.tensor(
            [
                 0.0017,  0.0004,  0.0004,  0.0005,  0.0002,  0.0002, -0.0002, -0.0013,
                -0.0010, -0.0015, -0.0018, -0.0032, -0.0060, -0.0082, -0.0096, -0.0099,
            ]
        )
        EXPECTED_VALUES_RIGHT = torch.tensor(
            [
                0.0038, 0.0028, 0.0031,  0.0032,  0.0031,  0.0032,  0.0030,  0.0019,
                0.0021, 0.0015, 0.0009, -0.0008, -0.0040, -0.0067, -0.0087, -0.0096,
            ]
        )
        

        
        self.assertTrue(output_values.shape == (1, 2, 5760))
        self.assertTrue(torch.allclose(output_values[0, 0, :16].cpu(), EXPECTED_VALUES_LEFT, atol=1e-4))
        self.assertTrue(torch.allclose(output_values[0, 1, :16].cpu(), EXPECTED_VALUES_RIGHT, atol=1e-4))

    @slow
    def test_generate_text_audio_prompt(self):
        model = self.model
        processor = self.processor

        
        audio = [get_bip_bip(duration=0.5)[None, :].repeat(2, 0), get_bip_bip(duration=1.0)[None, :].repeat(2, 0)]
        text = ["80s music", "Club techno"]

        inputs = processor(audio=audio, text=text, padding=True, return_tensors="pt")
        inputs = place_dict_on_device(inputs, device=torch_device)

        output_values = model.generate(**inputs, do_sample=False, guidance_scale=3.0, max_new_tokens=12)

        
        EXPECTED_VALUES_LEFT = torch.tensor(
            [
                 0.2535,  0.2008,  0.1471,  0.0896,  0.0306, -0.0200, -0.0501, -0.0728,
                -0.0832, -0.0856, -0.0867, -0.0884, -0.0864, -0.0866, -0.0744, -0.0430,
            ]
        )
        EXPECTED_VALUES_RIGHT = torch.tensor(
            [
                 0.1695,  0.1213,  0.0732,  0.0239, -0.0264, -0.0705, -0.0935, -0.1103,
                -0.1163, -0.1139, -0.1104, -0.1082, -0.1027, -0.1004, -0.0900, -0.0614,
            ]
        )
        

        
        self.assertTrue(output_values.shape == (2, 2, 37760))
        
        self.assertTrue(torch.allclose(output_values[0, 0, -16:].cpu(), EXPECTED_VALUES_LEFT, atol=1e-4))
        self.assertTrue(torch.allclose(output_values[0, 1, -16:].cpu(), EXPECTED_VALUES_RIGHT, atol=1e-4))