src/transformers/models/jukebox/modeling_jukebox.py/JukeboxPrior/__init__
class JukeboxPrior:    def __init__(self, config: JukeboxPriorConfig, level=None, nb_priors=3, vqvae_encoder=None, vqvae_decoder=None):
        super().__init__(config)
        # Passing functions instead of the vqvae module to avoid getting params, only used in the
        # forward loop
        self.vqvae_encoder = vqvae_encoder
        self.vqvae_decoder = vqvae_decoder

        self.levels = nb_priors
        self.level = level if level is not None else config.level

        self.base_model_prefix = f"priors.{self.level}"
        self._keys_to_ignore_on_load_unexpected += [r"priors.[^%d]." % self.level]

        self.n_ctx = config.n_ctx

        self.lyric_conditioning = config.nb_relevant_lyric_tokens > 0
        self.nb_relevant_lyric_tokens = config.nb_relevant_lyric_tokens
        self.encoder_loss_fraction = config.encoder_loss_fraction

        # Audio conditioning : conditioning on music tokens (either from audio or from previous levels or both)
        self.audio_conditioning = self.level != 0
        self.cond_level = self.level - 1
        if self.audio_conditioning:
            self.conditioner_blocks = JukeboxMusicTokenConditioner(config, self.level)

        # metadata conditioning : contioning on timing, genres, and artist
        self.metadata_conditioning = config.metadata_conditioning
        if self.metadata_conditioning:
            self.metadata_embedding = JukeboxLabelConditioner(config, include_time_signal=not self.audio_conditioning)

        # define encoder-decoder or encoder and decoder
        self.is_encoder_decoder = config.is_encoder_decoder
        if config.is_encoder_decoder:
            # encoder-decoder transformer
            self.input_shapes = [config.nb_relevant_lyric_tokens, config.n_ctx]
            self.embed_dim_shift = [0, config.lyric_vocab_size]
            self.width = config.hidden_size

            self.nb_relevant_lyric_tokens = config.nb_relevant_lyric_tokens

            self.prior = JukeboxConditionalAutoregressive(
                config,
                n_ctx=config.nb_relevant_lyric_tokens + config.n_ctx,
                embed_dim=config.lyric_vocab_size + config.music_vocab_size,
                audio_conditioning=(self.audio_conditioning or self.metadata_conditioning),
                metadata_conditioning=True,
            )

        else:
            # Separate encoder-decoder transformer
            encoder_config = config.encoder_config

            if self.nb_relevant_lyric_tokens != 0 and self.lyric_conditioning:
                self.lyric_acts_width = encoder_config.hidden_size
                self.encoder_width = config.hidden_size
                self.encoder_dim = config.lyric_vocab_size
                self.encoder = JukeboxConditionalAutoregressive(
                    encoder_config,
                    n_ctx=self.nb_relevant_lyric_tokens,
                    embed_dim=self.encoder_dim,
                    audio_conditioning=False,
                    metadata_conditioning=False,
                    is_encoder=True,
                )
                self.encoder.proj_in = JukeboxConv1D(encoder_config.hidden_size, config.hidden_size)
                self.encoder.final_layer_norm = JukeboxLayerNorm(config.hidden_size)
                self.encoder.lm_head = nn.Linear(config.hidden_size, config.lyric_vocab_size, bias=False)
            else:
                self.nb_relevant_lyric_tokens = 0

            # decoder model on the tokens
            self.prior = JukeboxConditionalAutoregressive(
                config,
                audio_conditioning=(self.audio_conditioning or self.metadata_conditioning),
                metadata_conditioning=self.metadata_conditioning,
            )

        self.next_token_prediction_loss_dims = config.n_ctx
        self.total_loss_dims = self.nb_relevant_lyric_tokens + self.next_token_prediction_loss_dims

        self.downsamples = [stride**down for stride, down in zip(config.res_strides_t, config.res_downs_t)]
        self.cond_downsample = self.downsamples[self.level] if self.level != 0 else None
        self.raw_to_tokens = np.prod(self.downsamples[: nb_priors - self.level])
        self.sample_length = self.n_ctx * self.raw_to_tokens

        logger.info(
            f"Level:{self.level}, Cond downsample:{self.cond_downsample}, Raw to tokens:{self.raw_to_tokens}, Sample"
            f" length:{self.sample_length}"
        )