soapboxguy/MusicGen

1# AudioCraft conditioning modules2 3AudioCraft provides a4[modular implementation of conditioning modules](../audiocraft/modules/conditioners.py)5that can be used with the language model to condition the generation.6The codebase was developed in order to easily extend the set of modules7currently supported to easily develop new ways of controlling the generation.8 9 10## Conditioning methods11 12For now, we support 3 main types of conditioning within AudioCraft:13* Text-based conditioning methods14* Waveform-based conditioning methods15* Joint embedding conditioning methods for text and audio projected in a shared latent space.16 17The Language Model relies on 2 core components that handle processing information:18* The `ConditionProvider` class, that maps metadata to processed conditions leveraging19all the defined conditioners for the given task.20* The `ConditionFuser` class, that takes preprocessed conditions and properly fuse the21conditioning embedding to the language model inputs following a given fusing strategy.22 23Different conditioners (for text, waveform, joint embeddings...) are provided as torch24modules in AudioCraft and are used internally in the language model to process the25conditioning signals and feed them to the language model.26 27 28## Core concepts29 30### Conditioners31 32The `BaseConditioner` torch module is the base implementation for all conditioners in audiocraft.33 34Each conditioner is expected to implement 2 methods:35* The `tokenize` method that is used as a preprocessing method that contains all processing36that can lead to synchronization points (e.g. BPE tokenization with transfer to the GPU).37The output of the tokenize method will then be used to feed the forward method.38* The `forward` method that takes the output of the tokenize method and contains the core computation39to obtain the conditioning embedding along with a mask indicating valid indices (e.g. padding tokens).40 41### ConditionProvider42 43The ConditionProvider prepares and provides conditions given a dictionary of conditioners.44 45Conditioners are specified as a dictionary of attributes and the corresponding conditioner46providing the processing logic for the given attribute.47 48Similarly to the conditioners, the condition provider works in two steps to avoid sychronization points:49* A `tokenize` method that takes a list of conditioning attributes for the batch,50and run all tokenize steps for the set of conditioners.51* A `forward` method that takes the output of the tokenize step and run all the forward steps52for the set of conditioners.53 54The list of conditioning attributes is passed as a list of `ConditioningAttributes`55that is presented just below.56 57### ConditionFuser58 59Once all conditioning signals have been extracted and processed by the `ConditionProvider`60as dense embeddings, they remain to be passed to the language model along with the original61language model inputs.62 63The `ConditionFuser` handles specifically the logic to combine the different conditions64to the actual model input, supporting different strategies to combine them.65 66One can therefore define different strategies to combine or fuse the condition to the input, in particular:67* Prepending the conditioning signal to the input with the `prepend` strategy,68* Summing the conditioning signal to the input with the `sum` strategy,69* Combining the conditioning relying on a cross-attention mechanism with the `cross` strategy,70* Using input interpolation with the `input_interpolate` strategy.71 72### SegmentWithAttributes and ConditioningAttributes: From metadata to conditions73 74The `ConditioningAttributes` dataclass is the base class for metadata75containing all attributes used for conditioning the language model.76 77It currently supports the following types of attributes:78* Text conditioning attributes: Dictionary of textual attributes used for text-conditioning.79* Wav conditioning attributes: Dictionary of waveform attributes used for waveform-based80conditioning such as the chroma conditioning.81* JointEmbed conditioning attributes: Dictionary of text and waveform attributes82that are expected to be represented in a shared latent space.83 84These different types of attributes are the attributes that are processed85by the different conditioners.86 87`ConditioningAttributes` are extracted from metadata loaded along the audio in the datasets,88provided that the metadata used by the dataset implements the `SegmentWithAttributes` abstraction.89 90All metadata-enabled datasets to use for conditioning in AudioCraft inherits91the [`audiocraft.data.info_dataset.InfoAudioDataset`](../audiocraft/data/info_audio_dataset.py) class92and the corresponding metadata inherits and implements the `SegmentWithAttributes` abstraction.93Refer to the [`audiocraft.data.music_dataset.MusicAudioDataset`](../audiocraft/data/music_dataset.py)94class as an example.95 96 97## Available conditioners98 99### Text conditioners100 101All text conditioners are expected to inherit from the `TextConditioner` class.102 103AudioCraft currently provides two text conditioners:104* The `LUTConditioner` that relies on look-up-table of embeddings learned at train time,105and relying on either no tokenizer or a spacy tokenizer. This conditioner is particularly106useful for simple experiments and categorical labels.107* The `T5Conditioner` that relies on a108[pre-trained T5 model](https://huggingface.co/docs/transformers/model_doc/t5)109frozen or fine-tuned at train time to extract the text embeddings.110 111### Waveform conditioners112 113All waveform conditioners are expected to inherit from the `WaveformConditioner` class and114consists of conditioning method that takes a waveform as input. The waveform conditioner115must implement the logic to extract the embedding from the waveform and define the downsampling116factor from the waveform to the resulting embedding.117 118The `ChromaStemConditioner` conditioner is a waveform conditioner for the chroma features119conditioning used by MusicGen. It takes a given waveform, extract relevant stems for melody120(namely all non drums and bass stems) using a121[pre-trained Demucs model](https://github.com/facebookresearch/demucs)122and then extract the chromagram bins from the remaining mix of stems.123 124### Joint embeddings conditioners125 126We finally provide support for conditioning based on joint text and audio embeddings through127the `JointEmbeddingConditioner` class and the `CLAPEmbeddingConditioner` that implements such128a conditioning method relying on a [pretrained CLAP model](https://github.com/LAION-AI/CLAP).129 130## Classifier Free Guidance131 132We provide a Classifier Free Guidance implementation in AudioCraft. With the classifier free133guidance dropout, all attributes are dropped with the same probability.134 135## Attribute Dropout136 137We further provide an attribute dropout strategy. Unlike the classifier free guidance dropout,138the attribute dropout drops given attributes with a defined probability, allowing the model139not to expect all conditioning signals to be provided at once.140 141## Faster computation of conditions142 143Conditioners that require some heavy computation on the waveform can be cached, in particular144the `ChromaStemConditioner` or `CLAPEmbeddingConditioner`. You just need to provide the145`cache_path` parameter to them. We recommend running dummy jobs for filling up the cache quickly.146An example is provied in the [musicgen.musicgen_melody_32khz grid](../audiocraft/grids/musicgen/musicgen_melody_32khz.py).