U 9%e@sVdZddlZddlZddlZddlZddlmZddlmZm Z m Z m Z m Z m Z ddlZddlmZddlmZddlmZddlmZdd lmZdd lmZmZdd lmZdd lmZmZm Z m!Z!m"Z"dd l#m$Z$ddl%m&Z&m'Z'm(Z(m)Z)ddl*m+Z+ddl,m-Z-ddl.m/Z/m0Z0er*ddl1m2Z2e(3e4Z5dZ6dZ7dgZ8eGddde!Z9d8ej:ej;ejd9ej?ej;e e=dddZ@ej?e=e=dd d!ZAGd"d#d#ejBZCGd$d%d%ejBZDGd&d'd'ejBZEGd(d)d)e$ZFd*ZGd+ZHd,ZIGd-d.d.eFZJe&d/eGGd0d1d1eFZKe&d2eGGd3d4d4eFZLe&d5eGGd6d7d7e$ZMdS):z PyTorch Musicgen model.N) dataclass) TYPE_CHECKINGAnyDictOptionalTupleUnion)CrossEntropyLoss) checkpoint)ACT2FN)GenerationConfig)%ClassifierFreeGuidanceLogitsProcessorLogitsProcessorList)StoppingCriteriaList)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentions ModelOutputSeq2SeqLMOutput)PreTrainedModel)add_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) AutoConfig AutoModel)MusicgenConfigMusicgenDecoderConfig) BaseStreamerr zfacebook/musicgen-smallc@s>eZdZUdZdZeejed<dZ ej ed<dZ e ed<dS)MusicgenUnconditionalInputa% 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). Nencoder_outputsattention_maskguidance_scale) __name__ __module__ __qualname____doc__r$rtorch FloatTensor__annotations__r% LongTensorr&floatr0r0m/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/models/musicgen/modeling_musicgen.pyr#As  r#)input_ids_shapedtypedevicepast_key_values_lengthcCs|\}}tj||ft|j|d}tj|d|d}|||d|ddkd||}|dkrtj tj ||||d|gdd}|ddddddf |d|||S)zB Make causal mask used for bi-directional self-attention. )r4rrr3r4dimN) r+fullfinfominarangesize masked_fill_viewtocatzerosexpand)r2r3r4r5bsztgt_lenmaskZ mask_condr0r0r1_make_causal_maskUs"  rH)rGr3rFcCsj|\}}|dk r|n|}|ddddddf|d|||}d|}||tjt|jS)z_ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. Nr?)r>rDrAZ masked_fillr+boolr;r<)rGr3rFrEsrc_lenZ expanded_maskZ inverted_maskr0r0r1 _expand_maskgs  *rL) input_ids pad_token_iddecoder_start_token_idcCsx||j}|ddddf|ddddf<|dkrDtd||dddf<|dkrdtd||dk||S)z1 Shift input ids one token to the right. Nr6rzSMake sure to set the decoder_start_token_id attribute of the model's configuration.rzIMake sure to set the pad_token_id attribute of the model's configuration.)Z new_zerosshapeclone ValueErrorr?)rMrNrOZshifted_input_idsr0r0r1shift_tokens_rightvs (rTcsheZdZdZeedfdd ZeedddZeeeddd Ze de j ed d d Z Z S)%MusicgenSinusoidalPositionalEmbeddingzDThis module produces sinusoidal positional embeddings of any length.) num_positions embedding_dimcs t||_|||dSN)super__init__rW make_weights)selfrVrW __class__r0r1rZs z.MusicgenSinusoidalPositionalEmbedding.__init__)num_embeddingsrWcCsN|||}t|dr,|j|jj|jjd}t||_d|j_|j dS)Nweightsr7F) get_embeddinghasattrrAr`r3r4nn ParameterZ requires_gradZdetach_)r\r_rWZ emb_weightsr0r0r1r[s    z2MusicgenSinusoidalPositionalEmbedding.make_weightscCs|d}td|d}ttj|tjd| }tj|tjdd|d}tjt|t |gdd |d}|ddkrtj|t |dgdd}| t S)z 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". ri'rr3rr8r6)mathlogr+expr=r/Z unsqueezerBcossinr@rCrAZget_default_dtype)r_rWZhalf_dimZembr0r0r1ras & z3MusicgenSinusoidalPositionalEmbedding.get_embeddingr)rMr5cCs`|\}}}t|||j}||jdkrH|||j|j|j d| d S)Nrr6) r>r+r=rAr4r`r[offsetrWZ index_selectr@detach)r\rMr5rE codebooksseq_lenZ position_idsr0r0r1forwards z-MusicgenSinusoidalPositionalEmbedding.forward)r)r'r(r)r*intrZr[ staticmethodrar+no_gradTensorro __classcell__r0r0r]r1rUs rUc seZdZdZdeeeeedfdd Zej eedd d Z dej e ej e e ej e ej e ej ee ej e ej e e ej fd d dZ ZS)MusicgenAttentionz=Multi-headed attention from 'Attention Is All You Need' paperFT embed_dim num_headsdropout is_decoderbiascst||_||_||_|||_|j||jkrNtd|jd|d|jd|_||_t j |||d|_ t j |||d|_ t j |||d|_ t j |||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).g࿩r|)rYrZrxryrzhead_dimrSscalingr{rcLineark_projv_projq_projout_proj)r\rxryrzr{r|r]r0r1rZs   zMusicgenAttention.__init__)tensorrnrEcCs ||||j|jddS)Nrr)r@ryr~ transpose contiguous)r\rrnrEr0r0r1_shapeszMusicgenAttention._shapeN) hidden_stateskey_value_statespast_key_valuer%layer_head_maskoutput_attentionsreturncCsx|dk }|\}} } |||j} |r\|dk r\|djd|jdkr\|d} |d} n|r|||d|} |||d|} n|dk r|||d|} |||d|} tj|d| gdd} tj|d| gdd} n(|||d|} |||d|} |j r | | f}||j d|j f}|| | |j |} | j |} | j |} | d}t| | dd}|||j | |fkrtd||j | |fd||dk r ||d| |fkrtd |d| |fd|| ||j | ||}| ||j | |}tjj|dd}|dk r||j fkrhtd |j fd|| dddd| ||j | |}| ||j | |}|r| ||j | |}| ||j | |}nd}tjj||j|jd }t|| }|||j | |j fkr4td ||j | |j fd|| ||j | |j }|dd}| || |j}||}|||fS) z#Input shape: Batch x Time x ChannelNrrrr6r8z$Attention weights should be of size z , but is z!Attention mask should be of size z/Head mask for a single layer should be of size ptrainingz `attn_output` should be of size )r>rrrQrrrr+rBr{ryr~r@reshapeZbmmrrSrc functionalZsoftmaxrzrrxr)r\rrrr%rrZis_cross_attentionrErF_Z query_statesZ key_statesZ value_statesZ proj_shaperKZ attn_weightsZattn_weights_reshapedZ attn_probsZ attn_outputr0r0r1ros~       "    zMusicgenAttention.forward)rvFT)NNNNF)r'r(r)r*rpr/rJrZr+rsrrrrortr0r0r]r1rus4 ruc sxeZdZedfdd Zd ejeejeejeejeejeejeeejee ee ejd dd Z Z S) MusicgenDecoderLayerconfigcst|j|_t|j|j|jddd|_|j|_t |j |_ |j |_ t |j|_t|j|j|jddd|_t |j|_t j|j|jdd|_t j|j|jdd|_t |j|_dS)NTFrw)rzr{r|r})rYrZ hidden_sizerxruZnum_attention_headsZattention_dropout self_attnrzr Zactivation_function activation_fnactivation_dropoutrc LayerNormself_attn_layer_norm encoder_attnencoder_attn_layer_normrZffn_dimfc1fc2final_layer_normr\rr]r0r1rZQs0  zMusicgenDecoderLayer.__init__NFT) rr%encoder_hidden_statesencoder_attention_maskrcross_attn_layer_head_maskrr use_cacherc Cs^|} ||}|dk r"|ddnd} |j|| |||d\}} } tjj||j|jd}| |}d}d}|dk r|} ||}|dk r|ddnd}|j||||||d\}}}tjj||j|jd}| |}| |} |} ||}| | |}tjj||j |jd}| |}tjj||j|jd}| |}|f}|rJ|| |f7}| rZ|| f7}|S)a 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. Nr)rrr%rrr)rrr%rrr) rrrcrrzrrrrrrrr)r\rr%rrrrrrrZresidualZself_attn_past_key_valueZself_attn_weightsZpresent_key_valueZcross_attn_present_key_valueZcross_attn_weightsZcross_attn_past_key_valueoutputsr0r0r1ronsT        zMusicgenDecoderLayer.forward)NNNNNNFT) r'r(r)r!rZr+rsrrrJrortr0r0r]r1rPs*  rc@s6eZdZdZeZdZdZddgZddZ d d d Z d S) MusicgenPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. modelTrrucCs|jj}t|tjtjfrD|jjjd|d|j dk r~|j j n:t|tj r~|jjjd|d|j dk r~|jj|j  dS)Nrv)Zmeanstd) rZinitializer_factor isinstancercrZConv1dweightdataZnormal_r|Zzero_ EmbeddingZ padding_idx)r\modulerr0r0r1 _init_weightss   z%MusicgenPreTrainedModel._init_weightsFcCst|tr||_dSrX)rMusicgenDecodergradient_checkpointingr\rvaluer0r0r1_set_gradient_checkpointings z3MusicgenPreTrainedModel._set_gradient_checkpointingN)F) r'r(r)r*r! config_classbase_model_prefixsupports_gradient_checkpointingZ_no_split_modulesrrr0r0r0r1rs ru 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. as 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) 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`. 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. a 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) 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`. 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. cseZdZdZedfdd ZddZddZd d Ze e de j e e je e je e j e e je e je eee je e je ee ee ee eeeefd d dZZS)rzw Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MusicgenDecoderLayer`] rcstj|_j|_j|_j|_j|_j rFt jnd|_ j dtfddtjD|_tjj|_tfddtjD|_tj|_d|_|dS)NrIrcsg|]}tjqSr0)rcrr.0rrrxr0r1 sz,MusicgenDecoder.__init__..csg|] }tqSr0)rrrr0r1rsF)rYrZrz layerdropZmax_position_embeddingsZmax_target_positionsrZd_model num_codebooksZscale_embeddingrfsqrtZ embed_scale vocab_sizerc ModuleListrange embed_tokensrUembed_positionsZnum_hidden_layerslayersr layer_normr post_initrr]rr1rZs&   zMusicgenDecoder.__init__cCs|jSrXrr\r0r0r1get_input_embeddingssz$MusicgenDecoder.get_input_embeddingscCs ||_dSrXrr\rr0r0r1set_input_embeddingssz$MusicgenDecoder.set_input_embeddingscCs`d}|ddkr$t||j|j|d}|dk r\t||j|dd|j}|dkrT|n||}|S)Nr6r)r4r5rF)rHr3r4rLrA)r\r% input_shape inputs_embedsr5Zcombined_attention_maskZexpanded_attn_maskr0r0r1_prepare_decoder_attention_masks z/MusicgenDecoder._prepare_decoder_attention_maskN rMr%rr head_maskcross_attn_head_maskpast_key_valuesrrroutput_hidden_states return_dictrc ! sdk r njj| dk r | njj} dk r4njj| dk rH| njj} |dk rj|dk rjtdnp|dk r|dj|jdj\} }}| |f}n<|dk r| dd}|ddddddfntd|dk r|ddjdnd}|dkrt fddt |D} ||||}|dk rV|dk rVt ||j|dd}|}|||j}tjj|jjd }jrjrrtd d | rd nd}rd nd}r|dk rd nd}rd nd}t||gd dgD]V\}}|dk r| dtjkrtd|dtjd| ddqtjD]B\}}| r|||f7}tdd}jr|jkrqb|dk r||nd}jrjrfdd}t |||||||dk r||nd|dk r ||ndd}n>||||||dk r0||nd|dk rD||nd|d }|d}rx||rndndf7}rb||df7}|dk rb||df7}qb!|}| r||f7}r|nd} | st"dd|| |||fDSt#|| |||dS)NzTYou cannot specify both decoder_input_ids and decoder_inputs_embeds at the same timer6zEYou have to specify either decoder_input_ids or decoder_inputs_embedsrrcs&g|]}j|dd|fqSrXr)rcodebook)inputr\r0r1r sz+MusicgenDecoder.forward..rrz[`use_cache=True` is incompatible with gradient checkpointing`. Setting `use_cache=False`...Fr0rrzThe `z` should be specified for z layers, but it is for .rcsfdd}|S)Ncs|fSrXr0)inputs)rrrr0r1custom_forward?szNMusicgenDecoder.forward..create_custom_forward..custom_forwardr0)rr)rr)rr1create_custom_forward>sz6MusicgenDecoder.forward..create_custom_forward)r%rrrrrrrr css|]}|dk r|VqdSrXr0)rvr0r0r1 psz*MusicgenDecoder.forward..last_hidden_staterr attentionscross_attentions)$rrrruse_return_dictrSrrrQr>sumrrrLr3rrAr4rcrrzrrloggerZ warning_onceziplenr enumeraterandomuniformrr rtupler)!r\rMr%rrrrrrrrrrrErrnrr5Z positionsrZall_hidden_statesZall_self_attnsZall_cross_attentionsZnext_decoder_cacheZ attn_maskZ mask_nameidxZ decoder_layerZdropout_probabilityrrZ layer_outputsZ next_cacher0)rrr\rr1ros      $        zMusicgenDecoder.forward) NNNNNNNNNNNN)r'r(r)r*r!rZrrrr!MUSICGEN_DECODER_INPUTS_DOCSTRINGr+r.rrsr,rrJrrrortr0r0r]r1rsB rz^The bare Musicgen decoder model outputting raw hidden-states without any specific head on top.cseZdZedfdd ZddZddZdd Zee de j e e j e e je e j e e j e e j e eee je e je ee ee ee eeeefd d d ZZS) MusicgenModelrcs"t|t||_|dSrX)rYrZrdecoderrrr]r0r1rZs  zMusicgenModel.__init__cCs|jjSrXrrrr0r0r1rsz"MusicgenModel.get_input_embeddingscCs ||j_dSrXrrr0r0r1rsz"MusicgenModel.set_input_embeddingscCs|jSrXrrr0r0r1 get_decoderszMusicgenModel.get_decoderNrc Cs| dk r | n|jj} | dk r | n|jj} | dk r4| n|jj} | dk rH| n|jj} |j||||||||| | | | d } | sz| St| j| j| j | j | j dS)N) rMr%rrrrrrrrrrr) rrrrrrrrrrrr)r\rMr%rrrrrrrrrrdecoder_outputsr0r0r1ros8zMusicgenModel.forward) NNNNNNNNNNNN)r'r(r)r!rZrrrrrr+r.rrsr,rrJrrrortr0r0r]r1r~s@ rzAThe MusicGen decoder model with a language modelling head on top.csJeZdZedfdd ZddZddZdd Zd d Zd d Z ddZ e e e eedd!ejeejeejeejeejeejeeeejeejeejeeeeeeeeeeefdddZd"ddZd#ejeedddZeddZed$eejee ee!ee"eeedddd Z#Z$S)%MusicgenForCausalLMrcsJtt|_j|_tfddtjD|_| dS)Ncs g|]}tjjjddqS)Fr})rcrrrrrr0r1rsz0MusicgenForCausalLM.__init__..) rYrZrrrrcrrlm_headsrrr]rr1rZs  zMusicgenForCausalLM.__init__cCs |jjjSrXrrrrr0r0r1rsz(MusicgenForCausalLM.get_input_embeddingscCs||jj_dSrXrrr0r0r1rsz(MusicgenForCausalLM.set_input_embeddingscCs|jSrXrrr0r0r1get_output_embeddingssz)MusicgenForCausalLM.get_output_embeddingscCs ||_dSrXrr\Znew_embeddingsr0r0r1set_output_embeddingssz)MusicgenForCausalLM.set_output_embeddingscCs ||j_dSrXrr)r\rr0r0r1 set_decoderszMusicgenForCausalLM.set_decodercCs|jjSrXrrr0r0r1rszMusicgenForCausalLM.get_decoder output_typerN)rMr%rrrrrrlabelsrrrrrcs| dk r | n|jj} |j||||||||| | | | d }|dtjfdd|jDdd}d}| dk rptd|jd |jd d}| s|f|dd}|dk r|f|S|St |||j |j |j |j d S) a 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: N) r%rrrrrrrrrrrcsg|] }|qSr0r0)rheadrr0r1rsz/MusicgenForCausalLM.forward..rr8z)Training is not implemented for Musicgen.r6r)losslogitsrrrr)r6)rrrr+stackrNotImplementedErrorrrQrrrrr)r\rMr%rrrrrrrrrrrrZ lm_logitsroutputr0rr1ros@zMusicgenForCausalLM.forwardTc Ks| dkr$|j||jj|jjd\}} ||| }| dk r\| dkr\|d}|dk r\|d}|dk rx|ddddf}||||||||dS)NrN max_lengthrrrr6)rMr%rrrrrr)build_delay_pattern_maskgeneration_configrNr apply_delay_pattern_maskrepeat) r\rMr%rrrrrrdelay_pattern_maskr&kwargsr0r0r1prepare_inputs_for_generation,s,    z1MusicgenForCausalLM.prepare_inputs_for_generation)rMrNr cCs|d|j|jd}|j\}}}|dk r.|n|jj}tj|||ftj|jdd}|d|dkr|||d|||dfSt |D]*}|dd|f|dd||||f<qtj tj||ftj d||dd} | t tj||ftj d} | |j} | || |}|dddddf} | dkdddf} t| dkrbt| } n|} |||d}|d d| f||d}||fS) aDBuild 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. r6Nr7rrre)Zdiagonalr.)rrrQr r r+oneslongr4rZtriurJZtrilrAZnonzerorr<)r\rMrNr rErrnZinput_ids_shiftedrZ delay_patternrGZfirst_codebook_idsZ start_idsZfirst_start_idZ pattern_maskr0r0r1r Ys0   (  z,MusicgenForCausalLM.build_delay_pattern_maskcCs0|jd}|dd|f}t|dk||}|S)zApply 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.r6.N)rQr+where)rMZdecoder_pad_token_maskrnr0r0r1rs z,MusicgenForCausalLM.apply_delay_pattern_maskr"rr logits_processorstopping_criteria synced_gpusstreamerc Ks|dkr|j}t|}|jf|}||||dk rF|nt}|dk rX|nt}|jdkr|j dk r| dddkrt d|j } t | tr| d} t d| d| |_|||j|\} } }| jd|j} |j|d<|j|d<|j|d <|j|d <d |k} | dddkr@| r@|| |j|j |d<| jd }| d dkob|jdk }|r|jdkr|jdkrt d|jdn:|jdk r|st d|jd|jd|j||_|jdk r|j|jkrtd|jd|jd||jkr,t d|d|jd|j| |j|jd\} }|dk r\|| ||d<|j dko|j!dko|j"dk}|j dko|j!dko|j"dk}|jdk r|jdkr|#t$|jd|_|j%||| d|d}|j&||d }|rJ|j'dkrtd!|j'd"|j(| f|||j|j |j)|j*||d#|}nd|r|+|}|j,f| |j'd$|\} }|j-| f||||j|j |j)|j*||d% |}ntd&|j*r|j.}n|}|/||d}|||jk0| |jd }|j*r||_.|S|SdS)' Generates sequences of token ids for models with a language modeling head. 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). 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`] Nr%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.r)Setting `pad_token_id` to `eos_token_id`: for open-end generation.rrrr&r$r6r 0Using the model-agnostic default `max_length` (=zx) to control the generation length. recommend setting `max_new_tokens` to control the maximum length of the generation.Both `max_new_tokens` (=) and `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)3Unfeasible length constraints: the minimum length (%) is larger than the maximum length ()%Input length of decoder_input_ids is , but `max_length` is set to X. This can lead to unexpected behavior. You should consider increasing `max_new_tokens`.r rrFTr input_ids_seq_lengthZencoder_input_idsZprefix_allowed_tokens_fnrr rBnum_return_sequences has to be 1 when doing greedy search, but is rrrrN eos_token_id output_scoresreturn_dict_in_generaterr)rM expand_size r logits_warperrrNr/r0r1rrGot 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`.)1r copydeepcopyupdatevalidate_validate_model_kwargsrrrNr/getrwarningrlist_prepare_model_inputs bos_token_idrQrrrrr&&_prepare_attention_mask_for_generationr max_new_tokens min_lengthrSr rOputcpu num_beamsnum_beam_groups do_sampleappendr_get_logits_processor_get_stopping_criterianum_return_sequences greedy_searchr0r1_get_logits_warper_expand_inputs_for_generationsample sequencesrr)r\rr rrrrr model_kwargsr/rMmodel_input_name batch_sizerequires_attention_maskr+has_default_max_lengthris_greedy_gen_modeis_sample_gen_moderr4 output_idsr0r0r1generates"N                         zMusicgenForCausalLM.generate) NNNNNNNNNNNNN) NNNNNNTNN)N)NNNNNN)%r'r(r)r!rZrrrrrrrrrr_CONFIG_FOR_DOCr+r.rrsr,rrJrrrorrpr rqrrrr rrrYrtr0r0r]r1rs   D -= rzThe 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.csreZdZeZdZdZdZd>eeee ee ee dfdd Z dd Z d?d d Z d dZddZddZddZddZddZddZefddZed@eeee dddZeeeeed dAeejeej eej!eej eejeej ee"ej!e"e"ej!eej!eej!eejee#ee#ee#ee#e$e"efd!d"d#Z%dBd$d%Z&dCe'ee(eej)fe'e'ej*e"eje(eej)ffd&d'd(Z+dDej)eeee,e(ee-fd)d*d+Z.dEeed,d-d.Z/ej)d/d0d1Z0d2d3Z1dFeej)ee'ee(eej)fejd4d5d6Z2e3dGeej)ee4ee5ee6ee#ed7d8d9d:Z7dHd