Jukebox

## [](#overview)Overview

The Jukebox model was proposed in [Jukebox: A generative model for music](https://arxiv.org/pdf/2005.00341.pdf) by Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, Ilya Sutskever. It introduces a generative music model which can produce minute long samples that can be conditioned on an artist, genres and lyrics.

The abstract from the paper is the following:

_We introduce Jukebox, a model that generates music with singing in the raw audio domain. We tackle the long context of raw audio using a multiscale VQ-VAE to compress it to discrete codes, and modeling those using autoregressive Transformers. We show that the combined model at scale can generate high-fidelity and diverse songs with coherence up to multiple minutes. We can condition on artist and genre to steer the musical and vocal style, and on unaligned lyrics to make the singing more controllable. We are releasing thousands of non cherry-picked samples, along with model weights and code._

As shown on the following figure, Jukebox is made of 3 `priors` which are decoder only models. They follow the architecture described in [Generating Long Sequences with Sparse Transformers](https://arxiv.org/abs/1904.10509), modified to support longer context length. First, a autoencoder is used to encode the text lyrics. Next, the first (also called `top_prior`) prior attends to the last hidden states extracted from the lyrics encoder. The priors are linked to the previous priors respectively via an `AudioConditionner` module. The`AudioConditioner` upsamples the outputs of the previous prior to raw tokens at a certain audio frame per second resolution. The metadata such as _artist, genre and timing_ are passed to each prior, in the form of a start token and positionnal embedding for the timing data. The hidden states are mapped to the closest codebook vector from the VQVAE in order to convert them to raw audio.

![JukeboxModel](https://gist.githubusercontent.com/ArthurZucker/92c1acaae62ebf1b6a951710bdd8b6af/raw/c9c517bf4eff61393f6c7dec9366ef02bdd059a3/jukebox.svg)

Tips:

-   This model only supports inference. This is for a few reasons, mostly because it requires a crazy amount of memory to train. Feel free to open a PR and add what’s missing to have a full integration with the hugging face traineer!
-   This model is very slow, and takes 8h to generate a minute long audio using the 5b top prior on a V100 GPU. In order automaticallay handle the device on which the model should execute, use `accelerate`.
-   Contrary to the paper, the order of the priors goes from `0` to `1` as it felt more intuitive : we sample starting from `0`.
-   Primed sampling (conditionning the sampling on raw audio) requires more memory than ancestral sampling and should be used with `fp16` set to `True`.

This model was contributed by [Arthur Zucker](https://huggingface.co/ArthurZ). The original code can be found [here](https://github.com/openai/jukebox).

## [](#transformers.JukeboxConfig)JukeboxConfig

### class transformers.JukeboxConfig

[](#transformers.JukeboxConfig)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/configuration_jukebox.py#L504)

( vqvae\_config = Noneprior\_config\_list = Nonenb\_priors = 3sampling\_rate = 44100timing\_dims = 64min\_duration = 0max\_duration = 600.0max\_nb\_genres = 5metadata\_conditioning = True\*\*kwargs )

This is the configuration class to store the configuration of a [JukeboxModel](/docs/transformers/v4.30.0/en/model_doc/jukebox#transformers.JukeboxModel).

Configuration objects inherit from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) and can be used to control the model outputs. Read the documentation from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.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.

Example:

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

>>> 
>>> configuration = JukeboxConfig()

>>> 
>>> model = JukeboxModel(configuration)

>>> 
>>> configuration = model.config```

#### from\_configs

[](#transformers.JukeboxConfig.from_configs)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/configuration_jukebox.py#L608)

( prior\_configs: typing.List\[transformers.models.jukebox.configuration\_jukebox.JukeboxPriorConfig\]vqvae\_config: JukeboxVQVAEConfig\*\*kwargs ) → [JukeboxConfig](/docs/transformers/v4.30.0/en/model_doc/jukebox#transformers.JukeboxConfig)

An instance of a configuration object

Instantiate a [JukeboxConfig](/docs/transformers/v4.30.0/en/model_doc/jukebox#transformers.JukeboxConfig) (or a derived class) from clip text model configuration and clip vision model configuration.

#### to\_dict

[](#transformers.JukeboxConfig.to_dict)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/configuration_jukebox.py#L620)

( ) → `Dict[str, any]`

Dictionary of all the attributes that make up this configuration instance,

Serializes this instance to a Python dictionary. Override the default [to\_dict()](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig.to_dict).

## [](#transformers.JukeboxPriorConfig)JukeboxPriorConfig

### class transformers.JukeboxPriorConfig

[](#transformers.JukeboxPriorConfig)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/configuration_jukebox.py#L144)

( act\_fn = 'quick\_gelu'level = 0alignment\_head = 2alignment\_layer = 68attention\_multiplier = 0.25attention\_pattern = 'enc\_dec\_with\_lyrics'attn\_dropout = 0attn\_res\_scale = Falseblocks = 64conv\_res\_scale = Nonenum\_layers = 72emb\_dropout = 0encoder\_config = Noneencoder\_loss\_fraction = 0.4hidden\_size = 2048init\_scale = 0.2is\_encoder\_decoder = Truelyric\_vocab\_size = 80mask = Falsemax\_duration = 600max\_nb\_genres = 1merged\_decoder = Truemetadata\_conditioning = Truemetadata\_dims = \[604, 7898\]min\_duration = 0mlp\_multiplier = 1.0music\_vocab\_size = 2048n\_ctx = 6144n\_heads = 2nb\_relevant\_lyric\_tokens = 384res\_conv\_depth = 3res\_conv\_width = 128res\_convolution\_multiplier = 1res\_dilation\_cycle = Noneres\_dilation\_growth\_rate = 1res\_downs\_t = \[3, 2, 2\]res\_strides\_t = \[2, 2, 2\]resid\_dropout = 0sampling\_rate = 44100spread = Nonetiming\_dims = 64zero\_out = False\*\*kwargs )

This is the configuration class to store the configuration of a [JukeboxPrior](/docs/transformers/v4.30.0/en/model_doc/jukebox#transformers.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](https://huggingface.co/openai/jukebox) -1b-lyrics) architecture.

Configuration objects inherit from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) and can be used to control the model outputs. Read the documentation from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) for more information.

#### to\_dict

[](#transformers.JukeboxPriorConfig.to_dict)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/configuration_jukebox.py#L370)

( ) → `Dict[str, any]`

Dictionary of all the attributes that make up this configuration instance,

Serializes this instance to a Python dictionary. Override the default [to\_dict()](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig.to_dict).

## [](#transformers.JukeboxVQVAEConfig)JukeboxVQVAEConfig

### class transformers.JukeboxVQVAEConfig

[](#transformers.JukeboxVQVAEConfig)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/configuration_jukebox.py#L383)

( act\_fn = 'relu'nb\_discrete\_codes = 2048commit = 0.02conv\_input\_shape = 1conv\_res\_scale = Falseembed\_dim = 64hop\_fraction = \[0.125, 0.5, 0.5\]levels = 3lmu = 0.99multipliers = \[2, 1, 1\]res\_conv\_depth = 4res\_conv\_width = 32res\_convolution\_multiplier = 1res\_dilation\_cycle = Noneres\_dilation\_growth\_rate = 3res\_downs\_t = \[3, 2, 2\]res\_strides\_t = \[2, 2, 2\]sample\_length = 1058304init\_scale = 0.2zero\_out = False\*\*kwargs )

This is the configuration class to store the configuration of a [JukeboxVQVAE](/docs/transformers/v4.30.0/en/model_doc/jukebox#transformers.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](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) and can be used to control the model outputs. Read the documentation from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) for more information.

## [](#transformers.JukeboxTokenizer)JukeboxTokenizer

### class transformers.JukeboxTokenizer

[](#transformers.JukeboxTokenizer)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/tokenization_jukebox.py#L59)

( artists\_filegenres\_filelyrics\_fileversion = \['v3', 'v2', 'v2'\]max\_n\_lyric\_tokens = 512n\_genres = 5unk\_token = '
'\*\*kwargs )

Constructs a Jukebox tokenizer. Jukebox can be conditioned on 3 different inputs :

-   Artists, unique ids are associated to each artist from the provided dictionary.
-   Genres, unique ids are associated to each genre from the provided dictionary.
-   Lyrics, character based tokenization. Must be initialized with the list of characters that are inside the vocabulary.

This tokenizer does not require training. It should be able to process a different number of inputs: as the conditioning of the model can be done on the three different queries. If None is provided, defaults values will be used.:

Depending on the number of genres on which the model should be conditioned (`n_genres`).

```
>>> from transformers import JukeboxTokenizer

>>> tokenizer = JukeboxTokenizer.from_pretrained("openai/jukebox-1b-lyrics")
>>> tokenizer("Alan Jackson", "Country Rock", "old town road")["input_ids"]
[tensor([[   0,    0,    0, 6785,  546,   41,   38,   30,   76,   46,   41,   49,
           40,   76,   44,   41,   27,   30]]), tensor([[  0,   0,   0, 145,   0]]), tensor([[  0,   0,   0, 145,   0]])]```

You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance.

If nothing is provided, the genres and the artist will either be selected randomly or set to None

This tokenizer inherits from [PreTrainedTokenizer](/docs/transformers/v4.30.0/en/main_classes/tokenizer#transformers.PreTrainedTokenizer) which contains most of the main methods. Users should refer to: this superclass for more information regarding those methods.

However the code does not allow that and only supports composing from various genres.

#### save\_vocabulary

[](#transformers.JukeboxTokenizer.save_vocabulary)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/tokenization_jukebox.py#L372)

( save\_directory: strfilename\_prefix: typing.Optional\[str\] = None )

Parameters

-   [](#transformers.JukeboxTokenizer.save_vocabulary.save_directory)**save\_directory** (`str`) — A path to the directory where to saved. It will be created if it doesn’t exist.
-   [](#transformers.JukeboxTokenizer.save_vocabulary.filename_prefix)**filename\_prefix** (`Optional[str]`, _optional_) — A prefix to add to the names of the files saved by the tokenizer.

Saves the tokenizer’s vocabulary dictionary to the provided save\_directory.

## [](#transformers.JukeboxModel)JukeboxModel

### class transformers.JukeboxModel

[](#transformers.JukeboxModel)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2302)

( config )

Parameters

-   [](#transformers.JukeboxModel.config)**config** (`JukeboxConfig`) — 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 [from\_pretrained()](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method to load the model weights.

The bare JUKEBOX Model used for music generation. 4 sampling techniques are supported : `primed_sample`, `upsample`, `continue_sample` and `ancestral_sample`. It does not have a `forward` method as the training is not end to end. If you want to fine-tune the model, it is recommended to use the `JukeboxPrior` class and train each prior individually.

This model inherits from [PreTrainedModel](/docs/transformers/v4.30.0/en/main_classes/model#transformers.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.

#### ancestral\_sample

[](#transformers.JukeboxModel.ancestral_sample)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2576)

( labelsn\_samples = 1\*\*sampling\_kwargs )

Parameters

-   [](#transformers.JukeboxModel.ancestral_sample.labels)**labels** (`List[torch.LongTensor]`) — List of length `n_sample`, and shape `(self.levels, 4 + self.config.max_nb_genre + lyric_sequence_length)` metadata such as `artist_id`, `genre_id` and the full list of lyric tokens which are used to condition the generation.
-   [](#transformers.JukeboxModel.ancestral_sample.n_samples)**n\_samples** (`int`, _optional_, default to 1) — Number of samples to be generated in parallel.

Generates music tokens based on the provided `labels. Will start at the desired prior level and automatically upsample the sequence. If you want to create the audio, you should call` model.decode(tokens)\`, which will use the VQ-VAE decoder to convert the music tokens to raw audio.

Example:

```
>>> from transformers import AutoTokenizer, JukeboxModel, set_seed

>>> model = JukeboxModel.from_pretrained("openai/jukebox-1b-lyrics", min_duration=0).eval()
>>> tokenizer = AutoTokenizer.from_pretrained("openai/jukebox-1b-lyrics")

>>> lyrics = "Hey, are you awake? Can you talk to me?"
>>> artist = "Zac Brown Band"
>>> genre = "Country"
>>> metas = tokenizer(artist=artist, genres=genre, lyrics=lyrics)
>>> set_seed(0)
>>> music_tokens = model.ancestral_sample(metas.input_ids, sample_length=400)

>>> with torch.no_grad():
...     model.decode(music_tokens)[:, :10].squeeze(-1)
tensor([[-0.0219, -0.0679, -0.1050, -0.1203, -0.1271, -0.0936, -0.0396, -0.0405,
    -0.0818, -0.0697]])```

#### primed\_sample

[](#transformers.JukeboxModel.primed_sample)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2652)

( raw\_audiolabels\*\*sampling\_kwargs )

Parameters

-   [](#transformers.JukeboxModel.primed_sample.raw_audio)**raw\_audio** (`List[torch.Tensor]` of length `n_samples` ) — A list of raw audio that will be used as conditioning information for each samples that will be generated.
-   [](#transformers.JukeboxModel.primed_sample.labels)**labels** (`List[torch.LongTensor]` of length `n_sample`, and shape `(self.levels, self.config.max_nb_genre + lyric_sequence_length)` — List of metadata such as `artist_id`, `genre_id` and the full list of lyric tokens which are used to condition the generation.
-   [](#transformers.JukeboxModel.primed_sample.sampling_kwargs)**sampling\_kwargs** (`Dict[Any]`) — Various additional sampling arguments that are used by the `_sample` function. A detail list of the arguments can bee seen in the `_sample` function documentation.

Generate a raw audio conditioned on the provided `raw_audio` which is used as conditioning at each of the generation levels. The audio is encoded to music tokens using the 3 levels of the VQ-VAE. These tokens are used: as conditioning for each level, which means that no ancestral sampling is required.

#### continue\_sample

[](#transformers.JukeboxModel.continue_sample)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2622)

( music\_tokenslabels\*\*sampling\_kwargs )

Parameters

-   [](#transformers.JukeboxModel.continue_sample.music_tokens)**music\_tokens** (`List[torch.LongTensor]` of length `self.levels` ) — A sequence of music tokens which will be used as context to continue the sampling process. Should have `self.levels` tensors, each corresponding to the generation at a certain level.
-   [](#transformers.JukeboxModel.continue_sample.labels)**labels** (`List[torch.LongTensor]` of length `n_sample`, and shape `(self.levels, self.config.max_nb_genre + lyric_sequence_length)` — List of metadata such as `artist_id`, `genre_id` and the full list of lyric tokens which are used to condition the generation.
-   [](#transformers.JukeboxModel.continue_sample.sampling_kwargs)**sampling\_kwargs** (`Dict[Any]`) — Various additional sampling arguments that are used by the `_sample` function. A detail list of the arguments can bee seen in the `_sample` function documentation.

Generates a continuation of the previously generated tokens.

#### upsample

[](#transformers.JukeboxModel.upsample)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2637)

( music\_tokenslabels\*\*sampling\_kwargs )

Parameters

-   [](#transformers.JukeboxModel.upsample.music_tokens)**music\_tokens** (`List[torch.LongTensor]` of length `self.levels` ) — A sequence of music tokens which will be used as context to continue the sampling process. Should have `self.levels` tensors, each corresponding to the generation at a certain level.
-   [](#transformers.JukeboxModel.upsample.labels)**labels** (`List[torch.LongTensor]` of length `n_sample`, and shape `(self.levels, self.config.max_nb_genre + lyric_sequence_length)` — List of metadata such as `artist_id`, `genre_id` and the full list of lyric tokens which are used to condition the generation.
-   [](#transformers.JukeboxModel.upsample.sampling_kwargs)**sampling\_kwargs** (`Dict[Any]`) — Various additional sampling arguments that are used by the `_sample` function. A detail list of the arguments can bee seen in the `_sample` function documentation.

Upsamples a sequence of music tokens using the prior at level `level`.

#### \_sample

[](#transformers.JukeboxModel._sample)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2437)

( music\_tokenslabelssample\_levelsmetas = Nonechunk\_size = 32sampling\_temperature = 0.98lower\_batch\_size = 16max\_batch\_size = 16sample\_length\_in\_seconds = 24compute\_alignments = Falsesample\_tokens = Noneoffset = 0save\_results = Truesample\_length = None )

Core sampling function used to generate music tokens. Iterates over the provided list of levels, while saving the generated raw audio at each step.

Returns: torch.Tensor

Example:

```
>>> from transformers import AutoTokenizer, JukeboxModel, set_seed
>>> import torch

>>> metas = dict(artist="Zac Brown Band", genres="Country", lyrics="I met a traveller from an antique land")
>>> tokenizer = AutoTokenizer.from_pretrained("openai/jukebox-1b-lyrics")
>>> model = JukeboxModel.from_pretrained("openai/jukebox-1b-lyrics", min_duration=0).eval()

>>> labels = tokenizer(**metas)["input_ids"]
>>> set_seed(0)
>>> zs = [torch.zeros(1, 0, dtype=torch.long) for _ in range(3)]
>>> zs = model._sample(zs, labels, [0], sample_length=40 * model.priors[0].raw_to_tokens, save_results=False)
>>> zs[0]
tensor([[1853, 1369, 1150, 1869, 1379, 1789,  519,  710, 1306, 1100, 1229,  519,
      353, 1306, 1379, 1053,  519,  653, 1631, 1467, 1229, 1229,   10, 1647,
     1254, 1229, 1306, 1528, 1789,  216, 1631, 1434,  653,  475, 1150, 1528,
     1804,  541, 1804, 1434]])```

## [](#transformers.JukeboxPrior)JukeboxPrior

### class transformers.JukeboxPrior

[](#transformers.JukeboxPrior)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L1770)

( config: JukeboxPriorConfiglevel = Nonenb\_priors = 3vqvae\_encoder = Nonevqvae\_decoder = None )

Parameters

-   [](#transformers.JukeboxPrior.config)**config** (`JukeboxPriorConfig`) — 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 [from\_pretrained()](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method to load the model weights.
-   [](#transformers.JukeboxPrior.level)**level** (`int`, _optional_) — Current level of the Prior. Should be in range `[0,nb_priors]`.
-   [](#transformers.JukeboxPrior.nb_priors)**nb\_priors** (`int`, _optional_, defaults to 3) — Total number of priors.
-   [](#transformers.JukeboxPrior.vqvae_encoder)**vqvae\_encoder** (`Callable`, _optional_) — Encoding method of the VQVAE encoder used in the forward pass of the model. Passing functions instead of the vqvae module to avoid getting the parameters.
-   [](#transformers.JukeboxPrior.vqvae_decoder)**vqvae\_decoder** (`Callable`, _optional_) — Decoding method of the VQVAE decoder used in the forward pass of the model. Passing functions instead of the vqvae module to avoid getting the parameters.

The JukeboxPrior class, which is a wrapper around the various conditioning and the transformer. JukeboxPrior can be seen as language models trained on music. They model the next `music token` prediction task. If a (lyric) `encoderù is defined, it also models the` next character\` prediction on the lyrics. Can be conditionned on timing, artist, genre, lyrics and codes from lower-levels Priors.

#### sample

[](#transformers.JukeboxPrior.sample)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2062)

( n\_samplesmusic\_tokens = Nonemusic\_tokens\_conds = Nonemetadata = Nonetemp = 1.0top\_k = 0top\_p = 0.0chunk\_size = Nonesample\_tokens = None )

Parameters

-   [](#transformers.JukeboxPrior.sample.n_samples)**n\_samples** (`int`) — Number of samples to generate.
-   [](#transformers.JukeboxPrior.sample.music_tokens)**music\_tokens** (`List[torch.LongTensor]`, _optional_) — Previously gemerated tokens at the current level. Used as context for the generation.
-   [](#transformers.JukeboxPrior.sample.music_tokens_conds)**music\_tokens\_conds** (`List[torch.FloatTensor]`, _optional_) — Upper-level music tokens generated by the previous prior model. Is `None` if the generation is not conditionned on the upper-level tokens.
-   [](#transformers.JukeboxPrior.sample.metadata)**metadata** (`List[torch.LongTensor]`, _optional_) — List containing the metatdata tensor with the artist, genre and the lyric tokens.
-   [](#transformers.JukeboxPrior.sample.temp)**temp** (`float`, _optional_, defaults to 1.0) — Sampling temperature.
-   [](#transformers.JukeboxPrior.sample.top_k)**top\_k** (`int`, _optional_, defaults to 0) — Top k probabilities used for filtering.
-   [](#transformers.JukeboxPrior.sample.top_p)**top\_p** (`float`, _optional_, defaults to 0.0) — Top p probabilities used for filtering.
-   [](#transformers.JukeboxPrior.sample.chunk_size)**chunk\_size** (`int`, _optional_) — Size of the chunks used to prepare the cache of the transformer.
-   [](#transformers.JukeboxPrior.sample.sample_tokens)**sample\_tokens** (`int`, _optional_) — Number of tokens to sample.

Ancestral/Prime sampling a window of tokens using the provided conditioning and metadatas.

#### forward

[](#transformers.JukeboxPrior.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L2230)

( hidden\_states: Tensormetadata: typing.Optional\[typing.List\[torch.LongTensor\]\]decode: typing.Optional\[bool\] = Falseget\_preds: typing.Optional\[bool\] = False )

Parameters

-   [](#transformers.JukeboxPrior.forward.hidden_states)**hidden\_states** (`torch.Tensor`) — Hidden states which should be raw audio
-   [](#transformers.JukeboxPrior.forward.metadata)**metadata** (`List[torch.LongTensor]`, _optional_) — List containing the metadata conditioning tensorwith the lyric and the metadata tokens.
-   [](#transformers.JukeboxPrior.forward.decode)**decode** (`bool`, _optional_, defaults to `False`) — Whether or not to decode the encoded to tokens.
-   [](#transformers.JukeboxPrior.forward.get_preds)**get\_preds** (`bool`, _optional_, defaults to `False`) — Whether or not to return the actual predicitons of the model.

Encode the hidden states using the `vqvae` encoder, and then predicts the next token in the `forward_tokens` function. The loss is the sum of the `encoder` loss and the `decoder` loss.

## [](#transformers.JukeboxVQVAE)JukeboxVQVAE

### class transformers.JukeboxVQVAE

[](#transformers.JukeboxVQVAE)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L602)

( config: JukeboxVQVAEConfig )

Parameters

-   [](#transformers.JukeboxVQVAE.config)**config** (`JukeboxConfig`) — 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 [from\_pretrained()](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method to load the model weights.

The Hierarchical VQ-VAE model used in Jukebox. This model follows the Hierarchical VQVAE paper from [Will Williams, Sam Ringer, Tom Ash, John Hughes, David MacLeod, Jamie Dougherty](https://arxiv.org/abs/2002.08111).

This model inherits from [PreTrainedModel](/docs/transformers/v4.30.0/en/main_classes/model#transformers.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.

#### forward

[](#transformers.JukeboxVQVAE.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L740)

( raw\_audio: FloatTensor )

Parameters

-   [](#transformers.JukeboxVQVAE.forward.raw_audio)**raw\_audio** (`torch.FloatTensor`) — Audio input which will be encoded and decoded.

Forward pass of the VQ-VAE, encodes the `raw_audio` to latent states, which are then decoded for each level. The commit loss, which ensure that the encoder’s computed embeddings are close to the codebook vectors, is computed.

Example:

```
>>> from transformers import JukeboxVQVAE, set_seed
>>> import torch

>>> model = JukeboxVQVAE.from_pretrained("openai/jukebox-1b-lyrics").eval()
>>> set_seed(0)
>>> zs = [torch.randint(100, (4, 1))]
>>> model.decode(zs).shape
torch.Size([4, 8, 1])```

#### encode

[](#transformers.JukeboxVQVAE.encode)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L710)

( input\_audiostart\_level = 0end\_level = Nonebs\_chunks = 1 )

Parameters

-   [](#transformers.JukeboxVQVAE.encode.input_audio)**input\_audio** (`torch.Tensor`) — Raw audio which will be encoded to its discrete representation using the codebook. The closest `code` form the codebook will be computed for each sequence of samples.
-   [](#transformers.JukeboxVQVAE.encode.start_level)**start\_level** (`int`, _optional_, defaults to 0) — Level at which the encoding process will start. Default to 0.
-   [](#transformers.JukeboxVQVAE.encode.end_level)**end\_level** (`int`, _optional_) — Level at which the encoding process will start. Default to None.
-   [](#transformers.JukeboxVQVAE.encode.bs_chunks)**bs\_chunks** (int, _optional_, defaults to 1) — Number of chunks of raw audio to process at the same time.

Transforms the `input_audio` to a discrete representation made out of `music_tokens`.

#### decode

[](#transformers.JukeboxVQVAE.decode)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/jukebox/modeling_jukebox.py#L674)

( music\_tokensstart\_level = 0end\_level = Nonebs\_chunks = 1 )

Parameters

-   [](#transformers.JukeboxVQVAE.decode.music_tokens)**music\_tokens** (`torch.LongTensor`) — Tensor of music tokens which will be decoded to raw audio by using the codebook. Each music token should be an index to a corresponding `code` vector in the codebook.
-   [](#transformers.JukeboxVQVAE.decode.start_level)**start\_level** (`int`, _optional_) — Level at which the decoding process will start. Default to 0.
-   [](#transformers.JukeboxVQVAE.decode.end_level)**end\_level** (`int`, _optional_) — Level at which the decoding process will start. Default to None.
-   [](#transformers.JukeboxVQVAE.decode.bs_chunks)**bs\_chunks** (int, _optional_) — Number of chunks to process at the same time.

Transforms the input `music_tokens` to their `raw_audio` representation.