U
    ,-ezi                  O   @   sf  d Z ddlZddlmZmZ ddlmZ ddlmZ e	e
Zddd	Zd
ddd
ddd
ddd
ddd
ddd
dddd
ddd
ddd
dddd
ddd
ddd
dddd
ddd
ddd
dddd
ddd
ddd
dddd
ddd
ddd
dddd
ddd
ddd
dddgOZd
ddgZdgZdddgZdd Zdd Zdd Zdd ZeeeedZG dd deZG dd deZG dd deZdS ) z Jukebox configuration    N)ListUnion   )PretrainedConfig)loggingzEhttps://huggingface.co/openai/jukebox-5b-lyrics/blob/main/config.jsonzEhttps://huggingface.co/openai/jukebox-1b-lyrics/blob/main/config.json)zopenai/jukebox-5b-lyricszopenai/jukebox-1b-lyricsZ
block_attnZtranspose_block_attnZprev_block_attnZcross_attentionZdense_attentionZ
prime_attnZ
dense_attnc                 C   s   t d S )Nr   )_FullDenseAttentionlayer r
   r/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/transformers/models/jukebox/configuration_jukebox.pyfull_dense_attentionu   s    r   c                 C   s   t | d  S )Nr   )_RawColumnPreviousRowAttentionr   r
   r
   r   !raw_column_previous_row_attentiony   s    r   c                 C   s   t | d  S )NO   )_LARGE_ATTENTIONr   r
   r
   r    large_separated_enc_dec_w_lyrics}   s    r   c                 C   s$   | d dkrt | d  S t| d  S )N      r   )_PrimePrimeDenseAttentionr   r   r
   r
   r   enc_dec_with_lyrics   s    r   )r   r   r   r   c                *   @   s   e Zd ZdZdZdddZdddd	d
ddddddddddddddddddddgddddddddddddddgdddgdddddf*dd Zed%ee	e
jf d!d"d#d$ZdS )&JukeboxPriorConfiga"  
        This is the configuration class to store the configuration of a [`JukeboxPrior`]. It is used to instantiate a
        `JukeboxPrior` according to the specified arguments, defining the model architecture. Instantiating a
        configuration with the defaults will yield a similar configuration to that of the top level prior from the
        [openai/jukebox-1b-lyrics](https://huggingface.co/openai/jukebox
    -1b-lyrics) architecture.

        Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
        documentation from [`PretrainedConfig`] for more information.



    Args:
        act_fn (`str`, *optional*, defaults to `"quick_gelu"`):
            Activation function.
        alignment_head (`int`, *optional*, defaults to 2):
            Head that is responsible of the alignment between lyrics and music. Only used to compute the lyric to audio
            alignment
        alignment_layer (`int`, *optional*, defaults to 68):
            Index of the layer that is responsible of the alignment between lyrics and music. Only used to compute the
            lyric to audio alignment
        attention_multiplier (`float`, *optional*, defaults to 0.25):
            Multiplier coefficient used to define the hidden dimension of the attention layers. 0.25 means that
            0.25*width of the model will be used.
        attention_pattern (`str`, *optional*, defaults to `"enc_dec_with_lyrics"`):
            Which attention pattern to use for the decoder/
        attn_dropout (`int`, *optional*, defaults to 0):
            Dropout probability for the post-attention layer dropout in the decoder.
        attn_res_scale (`bool`, *optional*, defaults to `False`):
            Whether or not to scale the residuals in the attention conditioner block.
        blocks (`int`, *optional*, defaults to 64):
            Number of blocks used in the `block_attn`. A sequence of length seq_len is factored as `[blocks, seq_len //
            blocks]` in the `JukeboxAttention` layer.
        conv_res_scale (`int`, *optional*):
            Whether or not to scale the residuals in the conditioner block. Since the top level prior does not have a
            conditioner, the default value is to None and should not be modified.
        num_layers (`int`, *optional*, defaults to 72):
            Number of layers of the transformer architecture.
        emb_dropout (`int`, *optional*, defaults to 0):
            Embedding dropout used in the lyric decoder.
        encoder_config (`JukeboxPriorConfig`, *optional*) :
            Configuration of the encoder which models the prior on the lyrics.
        encoder_loss_fraction (`float`, *optional*, defaults to 0.4):
            Multiplication factor used in front of the lyric encoder loss.
        hidden_size (`int`, *optional*, defaults to 2048):
            Hidden dimension of the attention layers.
        init_scale (`float`, *optional*, defaults to 0.2):
            Initialization scales for the prior modules.
        is_encoder_decoder (`bool`, *optional*, defaults to `True`):
            Whether or not the prior is an encoder-decoder model. In case it is not, and `nb_relevant_lyric_tokens` is
            greater than 0, the `encoder` args should be specified for the lyric encoding.
        mask (`bool`, *optional*, defaults to `False`):
            Whether or not to mask the previous positions in the attention.
        max_duration (`int`, *optional*, defaults to 600):
            Maximum supported duration of the generated song in seconds.
        max_nb_genres (`int`, *optional*, defaults to 1):
            Maximum number of genres that can be used to condition the model.
        merged_decoder (`bool`, *optional*, defaults to `True`):
            Whether or not the decoder and the encoder inputs are merged. This is used for the separated
            encoder-decoder architecture
        metadata_conditioning (`bool`, *optional*, defaults to `True)`:
            Whether or not to condition on the artist and genre metadata.
        metadata_dims (`List[int]`, *optional*, defaults to `[604, 7898]`):
            Number of genres and the number of artists that were used to train the embedding layers of the prior
            models.
        min_duration (`int`, *optional*, defaults to 0):
            Minimum duration of the generated audio on which the model was trained.
        mlp_multiplier (`float`, *optional*, defaults to 1.0):
            Multiplier coefficient used to define the hidden dimension of the MLP layers. 0.25 means that 0.25*width of
            the model will be used.
        music_vocab_size (`int`, *optional*, defaults to 2048):
            Number of different music tokens. Should be similar to the `JukeboxVQVAEConfig.nb_discrete_codes`.
        n_ctx (`int`, *optional*, defaults to 6144):
            Number of context tokens for each prior. The context tokens are the music tokens that are attended to when
            generating music tokens.
        n_heads (`int`, *optional*, defaults to 2):
                Number of attention heads.
        nb_relevant_lyric_tokens (`int`, *optional*, defaults to 384):
            Number of lyric tokens that are used when sampling a single window of length `n_ctx`
        res_conv_depth (`int`, *optional*, defaults to 3):
            Depth of the `JukeboxDecoderConvBock` used to upsample the previously sampled audio in the
            `JukeboxMusicTokenConditioner`.
        res_conv_width (`int`, *optional*, defaults to 128):
            Width of the `JukeboxDecoderConvBock` used to upsample the previously sampled audio in the
            `JukeboxMusicTokenConditioner`.
        res_convolution_multiplier (`int`, *optional*, defaults to 1):
            Multiplier used to scale the `hidden_dim` of the `JukeboxResConv1DBlock`.
        res_dilation_cycle (`int`, *optional*):
            Dilation cycle used to define the `JukeboxMusicTokenConditioner`. Usually similar to the ones used in the
            corresponding level of the VQVAE. The first prior does not use it as it is not conditioned on upper level
            tokens.
        res_dilation_growth_rate (`int`, *optional*, defaults to 1):
            Dilation grow rate used between each convolutionnal block of the `JukeboxMusicTokenConditioner`
        res_downs_t (`List[int]`, *optional*, defaults to `[3, 2, 2]`):
            Downsampling rates used in the audio conditioning network
        res_strides_t (`List[int]`, *optional*, defaults to `[2, 2, 2]`):
            Striding used in the audio conditioning network
        resid_dropout (`int`, *optional*, defaults to 0):
            Residual dropout used in the attention pattern.
        sampling_rate (`int`, *optional*, defaults to 44100):
            Sampling rate used for training.
        spread (`int`, *optional*):
            Spread used in the `summary_spread_attention` pattern
        timing_dims (`int`, *optional*, defaults to 64):
            Dimension of the timing embedding.
        zero_out (`bool`, *optional*, defaults to `False`):
            Whether or not to zero out convolution weights when initializing.
    Zjukebox_priorZn_positionsZn_head)Zmax_position_embeddingsZnum_attention_headsZ
quick_gelur      D   g      ?r   F@   NH   g?   皙?TP   iX     i\  i  g      ?i   i  r      D  c+           ,      K   s  || _ || _|| _|| _|| _|| _|| _|	| _|
| _|| _	|| _
|| _|d k	r^tf || _nd | _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _|| _| | _|!| _ |"| _!|#| _"|$| _#|%| _$|&| _%|'| _&|(| _'|)| _(|| _)|*| _*d S N)+act_fnalignment_headalignment_layerattention_multiplierattention_patternattn_dropoutattn_res_scaleblocksconv_res_scale
num_layersemb_dropoutmusic_vocab_sizer   encoder_configencoder_loss_fraction
init_scaleis_encoder_decoderlyric_vocab_sizelevelmaskmax_durationmax_nb_genresmerged_decodermetadata_conditioningmetadata_dimsmin_durationmlp_multipliern_ctxn_headsnb_relevant_lyric_tokensres_conv_depthres_conv_widthres_convolution_multiplierres_dilation_cycleres_dilation_growth_rateres_downs_tres_strides_tresid_dropoutsampling_ratespreadtiming_dimshidden_sizezero_out),selfr"   r3   r#   r$   r%   r&   r'   r(   r)   r*   r+   r,   r.   r/   rJ   r0   r1   r2   r4   r5   r6   r7   r8   r9   r:   r;   r-   r<   r=   r>   r?   r@   rA   rB   rC   rD   rE   rF   rG   rH   rI   rK   kwargsr
   r
   r   __init__  sX    .zJukeboxPriorConfig.__init__r   pretrained_model_name_or_pathreturnc                 K   s   |  | | j|f|\}}|ddkr8|d|  }d|krvt| drv|d | jkrvtd|d  d| j d | j|f|S )N
model_typejukeboxprior_You are using a model of type   to instantiate a model of type N. This is not supported for all configurations of models and can yield errors.Z_set_token_in_kwargsZget_config_dictgethasattrrR   loggerwarning	from_dict)clsrP   r3   rM   config_dictr
   r
   r   from_pretrained_  s    
 z"JukeboxPriorConfig.from_pretrained)r   )__name__
__module____qualname____doc__rR   Zattribute_maprN   classmethodr   strosPathLiker`   r
   r
   r
   r   r      sl   m
\ r   c                   @   s|   e Zd ZdZdZddddddd	d
d
gdddddgddddddddgdddgdddfddZeeee	j
f ddddZdS )JukeboxVQVAEConfiga  
    This is the configuration class to store the configuration of a [`JukeboxVQVAE`]. It is used to instantiate a
    `JukeboxVQVAE` according to the specified arguments, defining the model architecture. Instantiating a configuration
    with the defaults will yield a similar configuration to that of the VQVAE from
    [openai/jukebox-1b-lyrics](https://huggingface.co/openai/jukebox-1b-lyrics) architecture.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        act_fn (`str`, *optional*, defaults to `"relu"`):
            Activation function of the model.
        nb_discrete_codes (`int`, *optional*, defaults to 2048):
            Number of codes of the VQVAE.
        commit (`float`, *optional*, defaults to 0.02):
            Commit loss multiplier.
        conv_input_shape (`int`, *optional*, defaults to 1):
            Number of audio channels.
        conv_res_scale (`bool`, *optional*, defaults to `False`):
            Whether or not to scale the residuals of the `JukeboxResConv1DBlock`.
        embed_dim (`int`, *optional*, defaults to 64):
            Embedding dimension of the codebook vectors.
        hop_fraction (`List[int]`, *optional*, defaults to `[0.125, 0.5, 0.5]`):
            Fraction of non-intersecting window used when continuing the sampling process.
        levels (`int`, *optional*, defaults to 3):
            Number of hierarchical levels that used in the VQVAE.
        lmu (`float`, *optional*, defaults to 0.99):
            Used in the codebook update, exponential moving average coefficient. For more detail refer to Appendix A.1
            of the original [VQVAE paper](https://arxiv.org/pdf/1711.00937v2.pdf)
        multipliers (`List[int]`, *optional*, defaults to `[2, 1, 1]`):
            Depth and width multipliers used for each level. Used on the `res_conv_width` and `res_conv_depth`
        res_conv_depth (`int`, *optional*, defaults to 4):
            Depth of the encoder and decoder block. If no `multipliers` are used, this is the same for each level.
        res_conv_width (`int`, *optional*, defaults to 32):
            Width of the encoder and decoder block. If no `multipliers` are used, this is the same for each level.
        res_convolution_multiplier (`int`, *optional*, defaults to 1):
            Scaling factor of the hidden dimension used in the `JukeboxResConv1DBlock`.
        res_dilation_cycle (`int`, *optional*):
            Dilation cycle value used in the `JukeboxResnet`. If an int is used, each new Conv1 block will have a depth
            reduced by a power of `res_dilation_cycle`.
        res_dilation_growth_rate (`int`, *optional*, defaults to 3):
            Resnet dilation growth rate used in the VQVAE (dilation_growth_rate ** depth)
        res_downs_t (`List[int]`, *optional*, defaults to `[3, 2, 2]`):
            Downsampling rate for each level of the hierarchical VQ-VAE.
        res_strides_t (`List[int]`, *optional*, defaults to `[2, 2, 2]`):
            Stride used for each level of the hierarchical VQ-VAE.
        sample_length (`int`, *optional*, defaults to 1058304):
            Provides the max input shape of the VQVAE. Is used to compute the input shape of each level.
        init_scale (`float`, *optional*, defaults to 0.2):
            Initialization scale.
        zero_out (`bool`, *optional*, defaults to `False`):
            Whether or not to zero out convolution weights when initializing.
    Zjukebox_vqvaeZrelur   g{Gz?r   Fr   g      ?g      ?r   gGz?r          Ni & r   c                 K   s|   || _ || _|| _|| _|| _|| _|| _|| _|| _|| _	|| _
|
| _|| _|| _|	| _|| _|| _|| _|| _|| _d S r!   )hop_fractionconv_input_shapesample_lengthlevels	embed_dimnb_discrete_codesr@   r?   rA   rC   rB   multipliersrD   rE   lmucommitr*   r"   r0   rK   )rL   r"   rq   rt   rm   r*   rp   rl   ro   rs   rr   r?   r@   rA   rB   rC   rD   rE   rn   r0   rK   rM   r
   r
   r   rN     s(    zJukeboxVQVAEConfig.__init__r   rO   c                 K   s~   |  | | j|f|\}}|ddkr2|d }d|krpt| drp|d | jkrptd|d  d| j d | j|f|S )NrR   rS   vqvae_configrU   rV   rW   rX   )r^   rP   rM   r_   r
   r
   r   r`     s    
 z"JukeboxVQVAEConfig.from_pretrained)ra   rb   rc   rd   rR   rN   re   r   rf   rg   rh   r`   r
   r
   r
   r   ri   t  s2   6
/ri   c                	       sJ   e Zd ZdZdZd fdd	Zeee e	dddZ
 fddZ  ZS )JukeboxConfigaW  
    This is the configuration class to store the configuration of a [`JukeboxModel`].

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information. Instantiating a configuration with the defaults will
    yield a similar configuration to that of
    [openai/jukebox-1b-lyrics](https://huggingface.co/openai/jukebox-1b-lyrics) architecture.


    The downsampling and stride are used to determine downsampling of the input sequence. For example, downsampling =
    (5,3), and strides = (2, 2) will downsample the audio by 2^5 = 32 to get the first level of codes, and 2**8 = 256
    to get the second level codes. This is mostly true for training the top level prior and the upsamplers.

    Args:
        vqvae_config (`JukeboxVQVAEConfig`, *optional*):
            Configuration for the `JukeboxVQVAE` model.
        prior_config_list (`List[JukeboxPriorConfig]`, *optional*):
            List of the configs for each of the `JukeboxPrior` of the model. The original architecture uses 3 priors.
        nb_priors (`int`, *optional*, defaults to 3):
            Number of prior models that will sequentially sample tokens. Each prior is conditional auto regressive
            (decoder) model, apart from the top prior, which can include a lyric encoder. The available models were
            trained using a top prior and 2 upsampler priors.
        sampling_rate (`int`, *optional*, defaults to 44100):
            Sampling rate of the raw audio.
        timing_dims (`int`, *optional*, defaults to 64):
            Dimensions of the JukeboxRangeEmbedding layer which is equivalent to traditional positional embedding
            layer. The timing embedding layer converts the absolute and relative position in the currently sampled
            audio to a tensor of length `timing_dims` that will be added to the music tokens.
        min_duration (`int`, *optional*, defaults to 0):
            Minimum duration of the audios to generate
        max_duration (`float`, *optional*, defaults to 600.0):
            Maximum duration of the audios to generate
        max_nb_genres (`int`, *optional*, defaults to 5):
            Maximum number of genres that can be used to condition a single sample.
        metadata_conditioning (`bool`, *optional*, defaults to `True`):
            Whether or not to use metadata conditioning, corresponding to the artist, the genre and the min/maximum
            duration.

    Example:

    ```python
    >>> from transformers import JukeboxModel, JukeboxConfig

    >>> # Initializing a Jukebox configuration
    >>> configuration = JukeboxConfig()

    >>> # Initializing a model from the configuration
    >>> model = JukeboxModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    rS   Nr   r    r   r        @   Tc
                    s   |d kri }t d tf || _|d k	r<dd |D | _nVg | _t|D ]F}|
d| d }|d kr~i }t d| d | jtf | qJ| jj	| _	|| _
|| _|| _|| _|| _|| _|	| _t jf |
 d S )NzHvqvae_config is None. initializing the JukeboxVQVAE with default values.c                 S   s   g | ]}t f |qS r
   )r   ).0prior_configr
   r
   r   
<listcomp>;  s     z*JukeboxConfig.__init__.<locals>.<listcomp>rT   zQ's  config is None. Initializing the JukeboxPriorConfig list with default values.)r[   infori   ru   prior_configsrangepopappendr   rl   	nb_priorsr6   rG   rI   r:   r5   r8   superrN   )rL   ru   prior_config_listr   rG   rI   r:   r5   r6   r8   rM   Z	prior_idxrz   	__class__r
   r   rN   (  s0    


zJukeboxConfig.__init__)r}   ru   c                 K   s&   dd |D }| f ||  d|S )z
        Instantiate a [`JukeboxConfig`] (or a derived class) from clip text model configuration and clip vision model
        configuration.

        Returns:
            [`JukeboxConfig`]: An instance of a configuration object
        c                 S   s   g | ]}|  qS r
   to_dictry   configr
   r
   r   r{   _  s     z.JukeboxConfig.from_configs.<locals>.<listcomp>)r   Zvqvae_config_dictr   )r^   r}   ru   rM   r   r
   r
   r   from_configsV  s    	zJukeboxConfig.from_configsc                    s&   t   }dd |dD |d< |S )Nc                 S   s   g | ]}|  qS r
   r   r   r
   r
   r   r{   e  s     z)JukeboxConfig.to_dict.<locals>.<listcomp>r}   r   )r   r   r   )rL   resultr   r
   r   r   b  s    
zJukeboxConfig.to_dict)	NNr   r    r   r   rw   rx   T)ra   rb   rc   rd   rR   rN   re   r   r   ri   r   r   __classcell__r
   r
   r   r   rv     s   6         .rv   )rd   rg   typingr   r   Zconfiguration_utilsr   utilsr   Z
get_loggerra   r[   Z%JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAPr   r   r   r   r   r   r   r   ZATTENTION_PATTERNSr   ri   rv   r
   r
   r
   r   <module>   s   
Q

 f{