codingutf8 2022 the openai team s and 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 jukebox configuration import os from typing import list union from configurationutils import pretrainedconfig from utils import logging logger logging getloggername jukeboxpretrainedconfigarchivemap openaijukebox5blyrics https huggingface coopenaijukebox5blyricsblobmainconfig json openaijukebox1blyrics https huggingface coopenaijukebox1blyricsblobmainconfig json largeattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn blockattn transposeblockattn prevblockattn crossattention rawcolumnpreviousrowattention blockattn transposeblockattn prevblockattn fulldenseattention denseattention primeprimedenseattention primeattn primeattn denseattn def fulldenseattentionlayer return fulldenseattention0 def rawcolumnpreviousrowattentionlayer return rawcolumnpreviousrowattentionlayer 3 def largeseparatedencdecwlyricslayer return largeattentionlayer 79 def encdecwithlyricslayer if layer 16 15 return primeprimedenseattentionlayer 3 return rawcolumnpreviousrowattentionlayer 3 attentionpatterns fulldenseattention fulldenseattention rawcolumnpreviousrowattention rawcolumnpreviousrowattention alternate row column and previous row attn largeseparatedencdecwlyrics largeseparatedencdecwlyrics used by large separatedencdec model with lyrics encdecwithlyrics encdecwithlyrics used by encoderdecoder model with lyrics class jukeboxpriorconfigpretrainedconfig modeltype jukeboxprior attributemap maxpositionembeddings npositions numattentionheads nhead def init self actfnquickgelu level0 alignmenthead2 alignmentlayer68 attentionmultiplier0 25 attentionpatternencdecwithlyrics attndropout0 attnresscalefalse blocks64 convresscalenone numlayers72 embdropout0 encoderconfignone encoderlossfraction0 4 hiddensize2048 initscale0 2 isencoderdecodertrue lyricvocabsize80 maskfalse maxduration600 maxnbgenres1 mergeddecodertrue metadataconditioningtrue metadatadims604 7898 minduration0 mlpmultiplier1 0 musicvocabsize2048 nctx6144 nheads2 nbrelevantlyrictokens384 resconvdepth3 resconvwidth128 resconvolutionmultiplier1 resdilationcyclenone resdilationgrowthrate1 resdownst3 2 2 resstridest2 2 2 residdropout0 samplingrate44100 spreadnone timingdims64 zerooutfalse kwargs self actfn actfn self alignmenthead alignmenthead self alignmentlayer alignmentlayer self attentionmultiplier attentionmultiplier self attentionpattern attentionpattern self attndropout attndropout self attnresscale attnresscale self blocks blocks self convresscale convresscale self numlayers numlayers self embdropout embdropout self musicvocabsize musicvocabsize if encoderconfig is not none self encoderconfig jukeboxpriorconfigencoderconfig else self encoderconfig none self encoderlossfraction encoderlossfraction self initscale initscale self isencoderdecoder isencoderdecoder self lyricvocabsize lyricvocabsize self level level self mask mask self maxduration maxduration self maxnbgenres maxnbgenres self mergeddecoder mergeddecoder self metadataconditioning metadataconditioning self metadatadims metadatadims self minduration minduration self mlpmultiplier mlpmultiplier self nctx nctx self nheads nheads self nbrelevantlyrictokens nbrelevantlyrictokens self resconvdepth resconvdepth self resconvwidth resconvwidth self resconvolutionmultiplier resconvolutionmultiplier self resdilationcycle resdilationcycle self resdilationgrowthrate resdilationgrowthrate self resdownst resdownst self resstridest resstridest self residdropout residdropout self samplingrate samplingrate self spread spread self timingdims timingdims self hiddensize hiddensize self zeroout zeroout classmethod def frompretrained cls pretrainedmodelnameorpath unionstr os pathlike level0 kwargs pretrainedconfig cls settokeninkwargskwargs configdict kwargs cls getconfigdictpretrainedmodelnameorpath kwargs get the prior config dict if we are loading from jukeboxconfig if configdict getmodeltype jukebox configdict configdictfpriorlevel if modeltype in configdict and hasattrcls modeltype and configdictmodeltype cls modeltype logger warning fyou are using a model of type configdict modeltype to instantiate a model of type fcls modeltype this is not supported for all configurations of models and can yield errors return cls fromdictconfigdict kwargs class jukeboxvqvaeconfigpretrainedconfig modeltype jukeboxvqvae def init self actfnrelu nbdiscretecodes2048 commit0 02 convinputshape1 convresscalefalse embeddim64 hopfraction0 125 0 5 0 5 levels3 lmu0 99 multipliers2 1 1 resconvdepth4 resconvwidth32 resconvolutionmultiplier1 resdilationcyclenone resdilationgrowthrate3 resdownst3 2 2 resstridest2 2 2 samplelength1058304 initscale0 2 zerooutfalse kwargs self hopfraction hopfraction self convinputshape convinputshape self samplelength samplelength vqvae parameters all used self levels levels self embeddim embeddim self nbdiscretecodes nbdiscretecodes self resconvwidth resconvwidth self resconvdepth resconvdepth self resconvolutionmultiplier resconvolutionmultiplier self resdilationgrowthrate resdilationgrowthrate self resdilationcycle resdilationcycle self multipliers multipliers self resdownst resdownst self resstridest resstridest self lmu lmu self commit commit self convresscale convresscale self actfn actfn self initscale initscale self zeroout zeroout classmethod def frompretrainedcls pretrainedmodelnameorpath unionstr os pathlike kwargs pretrainedconfig cls settokeninkwargskwargs configdict kwargs cls getconfigdictpretrainedmodelnameorpath kwargs get the text config dict if we are loading from clipconfig if configdict getmodeltype jukebox configdict configdictvqvaeconfig if modeltype in configdict and hasattrcls modeltype and configdictmodeltype cls modeltype logger warning fyou are using a model of type configdict modeltype to instantiate a model of type fcls modeltype this is not supported for all configurations of models and can yield errors return cls fromdictconfigdict kwargs class jukeboxconfigpretrainedconfig modeltype jukebox def init self vqvaeconfignone priorconfiglistnone nbpriors3 samplingrate44100 timingdims64 minduration0 maxduration600 0 maxnbgenres5 metadataconditioningtrue kwargs if vqvaeconfig is none vqvaeconfig logger infovqvaeconfig is none initializing the jukeboxvqvae with default values self vqvaeconfig jukeboxvqvaeconfigvqvaeconfig if priorconfiglist is not none self priorconfigs jukeboxpriorconfigpriorconfig for priorconfig in priorconfiglist else self priorconfigs for prioridx in rangenbpriors priorconfig kwargs popfpriorprioridx none if priorconfig is none priorconfig logger info fpriorprioridx s config is none initializing the jukeboxpriorconfig list with default values self priorconfigs appendjukeboxpriorconfigpriorconfig self hopfraction self vqvaeconfig hopfraction self nbpriors nbpriors metadata conditioning self maxnbgenres maxnbgenres self samplingrate samplingrate self timingdims timingdims self minduration minduration self maxduration maxduration self metadataconditioning metadataconditioning super initkwargs classmethod def fromconfigscls priorconfigs listjukeboxpriorconfig vqvaeconfig jukeboxvqvaeconfig kwargs r 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 override the default todict to apply todict to the list of prior configs coding utf 8 2022 the openai team s and 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 jukebox configuration alternate row column and previous row attn used by large separated_enc_dec model with lyrics used by encoder_decoder model with lyrics 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 get the prior config dict if we are loading from jukeboxconfig 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 vqvae parameters all used get the text config dict if we are loading from clipconfig 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 metadata conditioning 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 override the default to_dict to apply to_dict to the list of prior configs
import os
from typing import List, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


logger = logging.get_logger(__name__)

JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    "openai/jukebox-5b-lyrics": "https://huggingface.co/openai/jukebox-5b-lyrics/blob/main/config.json",
    "openai/jukebox-1b-lyrics": "https://huggingface.co/openai/jukebox-1b-lyrics/blob/main/config.json",
}

_LARGE_ATTENTION = [
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "block_attn",
    "transpose_block_attn",
    "prev_block_attn",
    "cross_attention",
]
_RawColumnPreviousRowAttention = ["block_attn", "transpose_block_attn", "prev_block_attn"]
_FullDenseAttention = ["dense_attention"]
_PrimePrimeDenseAttention = ["prime_attn", "prime_attn", "dense_attn"]


def full_dense_attention(layer):
    return _FullDenseAttention[0]


def raw_column_previous_row_attention(layer):
    return _RawColumnPreviousRowAttention[layer % 3]


def large_separated_enc_dec_w_lyrics(layer):
    return _LARGE_ATTENTION[layer % 79]


def enc_dec_with_lyrics(layer):
    if layer % 16 == 15:
        return _PrimePrimeDenseAttention[layer % 3]
    return _RawColumnPreviousRowAttention[layer % 3]


ATTENTION_PATTERNS = {
    "full_dense_attention": full_dense_attention,
    "raw_column_previous_row_attention": raw_column_previous_row_attention,  
    "large_separated_enc_dec_w_lyrics": large_separated_enc_dec_w_lyrics,  
    "enc_dec_with_lyrics": enc_dec_with_lyrics,  
}


class JukeboxPriorConfig(PretrainedConfig):
    

    model_type = "jukebox_prior"
    attribute_map = {
        "max_position_embeddings": "n_positions",
        "num_attention_heads": "n_head",
    }

    def __init__(
        self,
        act_fn="quick_gelu",
        level=0,
        alignment_head=2,
        alignment_layer=68,
        attention_multiplier=0.25,
        attention_pattern="enc_dec_with_lyrics",
        attn_dropout=0,
        attn_res_scale=False,
        blocks=64,
        conv_res_scale=None,
        num_layers=72,
        emb_dropout=0,
        encoder_config=None,
        encoder_loss_fraction=0.4,
        hidden_size=2048,
        init_scale=0.2,
        is_encoder_decoder=True,
        lyric_vocab_size=80,
        mask=False,
        max_duration=600,
        max_nb_genres=1,
        merged_decoder=True,
        metadata_conditioning=True,
        metadata_dims=[604, 7898],
        min_duration=0,
        mlp_multiplier=1.0,
        music_vocab_size=2048,
        n_ctx=6144,
        n_heads=2,
        nb_relevant_lyric_tokens=384,
        res_conv_depth=3,
        res_conv_width=128,
        res_convolution_multiplier=1,
        res_dilation_cycle=None,
        res_dilation_growth_rate=1,
        res_downs_t=[3, 2, 2],
        res_strides_t=[2, 2, 2],
        resid_dropout=0,
        sampling_rate=44100,
        spread=None,
        timing_dims=64,
        zero_out=False,
        **kwargs,
    ):
        self.act_fn = act_fn
        self.alignment_head = alignment_head
        self.alignment_layer = alignment_layer
        self.attention_multiplier = attention_multiplier
        self.attention_pattern = attention_pattern
        self.attn_dropout = attn_dropout
        self.attn_res_scale = attn_res_scale
        self.blocks = blocks
        self.conv_res_scale = conv_res_scale
        self.num_layers = num_layers
        self.emb_dropout = emb_dropout
        self.music_vocab_size = music_vocab_size
        if encoder_config is not None:
            self.encoder_config = JukeboxPriorConfig(**encoder_config)
        else:
            self.encoder_config = None
        self.encoder_loss_fraction = encoder_loss_fraction
        self.init_scale = init_scale
        self.is_encoder_decoder = is_encoder_decoder
        self.lyric_vocab_size = lyric_vocab_size
        self.level = level
        self.mask = mask
        self.max_duration = max_duration
        self.max_nb_genres = max_nb_genres
        self.merged_decoder = merged_decoder
        self.metadata_conditioning = metadata_conditioning
        self.metadata_dims = metadata_dims
        self.min_duration = min_duration
        self.mlp_multiplier = mlp_multiplier
        self.n_ctx = n_ctx
        self.n_heads = n_heads
        self.nb_relevant_lyric_tokens = nb_relevant_lyric_tokens
        self.res_conv_depth = res_conv_depth
        self.res_conv_width = res_conv_width
        self.res_convolution_multiplier = res_convolution_multiplier
        self.res_dilation_cycle = res_dilation_cycle
        self.res_dilation_growth_rate = res_dilation_growth_rate
        self.res_downs_t = res_downs_t
        self.res_strides_t = res_strides_t
        self.resid_dropout = resid_dropout
        self.sampling_rate = sampling_rate
        self.spread = spread
        self.timing_dims = timing_dims
        self.hidden_size = hidden_size
        self.zero_out = zero_out

    @classmethod
    def from_pretrained(
        cls, pretrained_model_name_or_path: Union[str, os.PathLike], level=0, **kwargs
    ) -> "PretrainedConfig":
        cls._set_token_in_kwargs(kwargs)

        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)

        
        if config_dict.get("model_type") == "jukebox":
            config_dict = config_dict[f"prior_{level}"]

        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
            logger.warning(
                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
            )

        return cls.from_dict(config_dict, **kwargs)


class JukeboxVQVAEConfig(PretrainedConfig):
    

    model_type = "jukebox_vqvae"

    def __init__(
        self,
        act_fn="relu",
        nb_discrete_codes=2048,
        commit=0.02,
        conv_input_shape=1,
        conv_res_scale=False,
        embed_dim=64,
        hop_fraction=[0.125, 0.5, 0.5],
        levels=3,
        lmu=0.99,
        multipliers=[2, 1, 1],
        res_conv_depth=4,
        res_conv_width=32,
        res_convolution_multiplier=1,
        res_dilation_cycle=None,
        res_dilation_growth_rate=3,
        res_downs_t=[3, 2, 2],
        res_strides_t=[2, 2, 2],
        sample_length=1058304,
        init_scale=0.2,
        zero_out=False,
        **kwargs,
    ):
        self.hop_fraction = hop_fraction
        self.conv_input_shape = conv_input_shape
        self.sample_length = sample_length

        
        self.levels = levels
        self.embed_dim = embed_dim
        self.nb_discrete_codes = nb_discrete_codes
        self.res_conv_width = res_conv_width
        self.res_conv_depth = res_conv_depth
        self.res_convolution_multiplier = res_convolution_multiplier
        self.res_dilation_growth_rate = res_dilation_growth_rate
        self.res_dilation_cycle = res_dilation_cycle
        self.multipliers = multipliers
        self.res_downs_t = res_downs_t
        self.res_strides_t = res_strides_t
        self.lmu = lmu
        self.commit = commit
        self.conv_res_scale = conv_res_scale
        self.act_fn = act_fn
        self.init_scale = init_scale
        self.zero_out = zero_out

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
        cls._set_token_in_kwargs(kwargs)

        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)

        
        if config_dict.get("model_type") == "jukebox":
            config_dict = config_dict["vqvae_config"]

        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
            logger.warning(
                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
            )

        return cls.from_dict(config_dict, **kwargs)


class JukeboxConfig(PretrainedConfig):
    

    model_type = "jukebox"

    def __init__(
        self,
        vqvae_config=None,
        prior_config_list=None,
        nb_priors=3,
        sampling_rate=44100,
        timing_dims=64,
        min_duration=0,
        max_duration=600.0,
        max_nb_genres=5,
        metadata_conditioning=True,
        **kwargs,
    ):
        if vqvae_config is None:
            vqvae_config = {}
            logger.info("vqvae_config is None. initializing the JukeboxVQVAE with default values.")

        self.vqvae_config = JukeboxVQVAEConfig(**vqvae_config)
        if prior_config_list is not None:
            self.prior_configs = [JukeboxPriorConfig(**prior_config) for prior_config in prior_config_list]
        else:
            self.prior_configs = []
            for prior_idx in range(nb_priors):
                prior_config = kwargs.pop(f"prior_{prior_idx}", None)
                if prior_config is None:
                    prior_config = {}
                    logger.info(
                        f"prior_{prior_idx}'s  config is None. Initializing the JukeboxPriorConfig list with default"
                        " values."
                    )
                self.prior_configs.append(JukeboxPriorConfig(**prior_config))

        self.hop_fraction = self.vqvae_config.hop_fraction

        self.nb_priors = nb_priors

        
        self.max_nb_genres = max_nb_genres
        self.sampling_rate = sampling_rate
        self.timing_dims = timing_dims
        self.min_duration = min_duration
        self.max_duration = max_duration
        self.metadata_conditioning = metadata_conditioning

        super().__init__(**kwargs)

    @classmethod
    def from_configs(cls, prior_configs: List[JukeboxPriorConfig], vqvae_config: JukeboxVQVAEConfig, **kwargs):
        r
        prior_config_list = [config.to_dict() for config in prior_configs]
        return cls(prior_config_list=prior_config_list, vqvae_config_dict=vqvae_config.to_dict(), **kwargs)

    def to_dict(self):
        
        result = super().to_dict()
        result["prior_config_list"] = [config.to_dict() for config in result.pop("prior_configs")]
        return result