src/transformers/models/jukebox/modeling_jukebox.py/JukeboxConditionalAutoregressive/primed_sample
class JukeboxConditionalAutoregressive:    def primed_sample(
        self,
        n_samples,
        lyric_and_music_tokens,
        audio_conditioning=None,
        metadata_conditioning=None,
        last_encoder_hidden_states=None,
        temp=1.0,
        top_k=0,
        top_p=0.0,
        get_preds=False,
        chunk_size=None,
        sample_tokens=None,
    ):
        if sample_tokens is None:
            sample_tokens = self.n_ctx
        # Preprocess.
        batch_size = lyric_and_music_tokens.shape[0]
        with torch.no_grad():
            lyric_and_music_tokens = lyric_and_music_tokens.view(batch_size, -1).long()

        sampled_audio = torch.split(lyric_and_music_tokens, 1, dim=1)
        sampled_audio = list(sampled_audio)

        if not self.audio_conditioning:
            audio_conditioning = torch.zeros(
                (n_samples, 1, self.width), dtype=self.transformer._attn_mods[0].mlp.c_fc.weight.dtype
            ).to(lyric_and_music_tokens.device)

        with torch.no_grad():
            if get_preds:
                preds = []

            # Fill up key/value cache for past context by runing forward pass.
            # We do so in chunks instead of doing the whole past in one forward pass to reduce max memory usage.
            if chunk_size is None:
                chunk_size = len(sampled_audio)
            chunk_sizes = self.split_chunks(len(sampled_audio), chunk_size)
            x_primes = []
            start = 0
            token = None

            for current_chunk_size in tqdm(chunk_sizes, desc="Preparing past key value", leave=False):
                sampled_audio_prime, conds_prime = [], []
                for sample_t in range(start, start + current_chunk_size):
                    x_prime, cond_prime = self.get_emb(
                        sample_t, n_samples, token, audio_conditioning, metadata_conditioning
                    )
                    token = sampled_audio[sample_t]
                    sampled_audio_prime.append(x_prime)
                    conds_prime.append(cond_prime)
                start = start + current_chunk_size
                x_prime, cond_prime = torch.cat(sampled_audio_prime, dim=1), torch.cat(conds_prime, dim=1)
                del sampled_audio_prime
                del conds_prime
                if not get_preds:
                    del cond_prime
                x_prime = self.transformer(x_prime, last_encoder_hidden_states=last_encoder_hidden_states, sample=True)

                if get_preds:
                    if self.add_cond_after_transformer:
                        x_prime = x_prime + cond_prime
                    del cond_prime
                    x_primes.append(x_prime)
                else:
                    del x_prime

            if get_preds:
                x_prime = torch.cat(x_primes, dim=1)
                x_prime = self.fc_proj_out(x_prime)  # Predictions
                preds.append(x_prime)

            # the input of the encoder and decoder can be merged into (lyrics, music tokens)
            input_tokens = sampled_audio[-1]

            itererator = tqdm(
                range(len(sampled_audio), sample_tokens),
                desc=f"Sampling {len(range(len(sampled_audio), sample_tokens))} music tokens",
                leave=False,
            )
            for sample_t in itererator:
                hidden_states, cond = self.get_emb(
                    sample_t, n_samples, input_tokens, audio_conditioning, metadata_conditioning
                )

                hidden_states = self.transformer(
                    hidden_states, last_encoder_hidden_states=last_encoder_hidden_states, sample=True
                )
                if self.add_cond_after_transformer:
                    hidden_states = hidden_states + cond
                hidden_states = self.fc_proj_out(hidden_states)  # Predictions
                if get_preds:
                    preds.append(hidden_states)
                # Adjust logits
                hidden_states = hidden_states / temp
                hidden_states = filter_logits(hidden_states, top_k=top_k, top_p=top_p)
                # only music tokens are sampled
                music_tokens = torch.distributions.Categorical(logits=hidden_states).sample()
                sampled_audio.append(music_tokens.clone())
                input_tokens = music_tokens

            del input_tokens, music_tokens
            self.transformer.del_cache()

            music_tokens = torch.cat(sampled_audio, dim=1)
            if get_preds:
                preds = torch.cat(preds, dim=1)
        if get_preds:
            return music_tokens, preds
        else:
            return music_tokens