soapboxguy/MusicGen

1# EnCodec: High Fidelity Neural Audio Compression2 3AudioCraft provides the training code for EnCodec, a state-of-the-art deep learning4based audio codec supporting both mono stereo audio, presented in the5[High Fidelity Neural Audio Compression][arxiv] paper.6Check out our [sample page][encodec_samples].7 8## Original EnCodec models9 10The EnCodec models presented in High Fidelity Neural Audio Compression can be accessed11and used with the [EnCodec repository](https://github.com/facebookresearch/encodec).12 13**Note**: We do not guarantee compatibility between the AudioCraft and EnCodec codebases14and released checkpoints at this stage.15 16 17## Installation18 19Please follow the AudioCraft installation instructions from the [README](../README.md).20 21 22## Training23 24The [CompressionSolver](../audiocraft/solvers/compression.py) implements the audio reconstruction25task to train an EnCodec model. Specifically, it trains an encoder-decoder with a quantization26bottleneck - a SEANet encoder-decoder with Residual Vector Quantization bottleneck for EnCodec -27using a combination of objective and perceptual losses in the forms of discriminators.28 29The default configuration matches a causal EnCodec training with at a single bandwidth.30 31### Example configuration and grids32 33We provide sample configuration and grids for training EnCodec models.34 35The compression configuration are defined in36[config/solver/compression](../config/solver/compression).37 38The example grids are available at39[audiocraft/grids/compression](../audiocraft/grids/compression).40 41```shell42# base causal encodec on monophonic audio sampled at 24 khz43dora grid compression.encodec_base_24khz44# encodec model used for MusicGen on monophonic audio sampled at 32 khz45dora grid compression.encodec_musicgen_32khz46```47 48### Training and valid stages49 50The model is trained using a combination of objective and perceptual losses.51More specifically, EnCodec is trained with the MS-STFT discriminator along with52objective losses through the use of a loss balancer to effectively weight53the different losses, in an intuitive manner.54 55### Evaluation stage56 57Evaluations metrics for audio generation:58* SI-SNR: Scale-Invariant Signal-to-Noise Ratio.59* ViSQOL: Virtual Speech Quality Objective Listener.60 61Note: Path to the ViSQOL binary (compiled with bazel) needs to be provided in62order to run the ViSQOL metric on the reference and degraded signals.63The metric is disabled by default.64Please refer to the [metrics documentation](../METRICS.md) to learn more.65 66### Generation stage67 68The generation stage consists in generating the reconstructed audio from samples69with the current model. The number of samples generated and the batch size used are70controlled by the `dataset.generate` configuration. The output path and audio formats71are defined in the generate stage configuration.72 73```shell74# generate samples every 5 epoch75dora run solver=compression/encodec_base_24khz generate.every=576# run with a different dset77dora run solver=compression/encodec_base_24khz generate.path=<PATH_IN_DORA_XP_FOLDER>78# limit the number of samples or use a different batch size79dora grid solver=compression/encodec_base_24khz dataset.generate.num_samples=10 dataset.generate.batch_size=480```81 82### Playing with the model83 84Once you have a model trained, it is possible to get the entire solver, or just85the trained model with the following functions:86 87```python88from audiocraft.solvers import CompressionSolver89 90# If you trained a custom model with signature SIG.91model = CompressionSolver.model_from_checkpoint('//sig/SIG')92# If you want to get one of the pretrained models with the `//pretrained/` prefix.93model = CompressionSolver.model_from_checkpoint('//pretrained/facebook/encodec_32khz')94# Or load from a custom checkpoint path95model = CompressionSolver.model_from_checkpoint('/my_checkpoints/foo/bar/checkpoint.th')96 97 98# If you only want to use a pretrained model, you can also directly get it99# from the CompressionModel base model class.100from audiocraft.models import CompressionModel101 102# Here do not put the `//pretrained/` prefix!103model = CompressionModel.get_pretrained('facebook/encodec_32khz')104model = CompressionModel.get_pretrained('dac_44khz')105 106# Finally, you can also retrieve the full Solver object, with its dataloader etc.107from audiocraft import train108from pathlib import Path109import logging110import os111import sys112 113# uncomment the following line if you want some detailed logs when loading a Solver.114logging.basicConfig(stream=sys.stderr, level=logging.INFO)115# You must always run the following function from the root directory.116os.chdir(Path(train.__file__).parent.parent)117 118 119# You can also get the full solver (only for your own experiments).120# You can provide some overrides to the parameters to make things more convenient.121solver = train.get_solver_from_sig('SIG', {'device': 'cpu', 'dataset': {'batch_size': 8}})122solver.model123solver.dataloaders124```125 126### Importing / Exporting models127 128At the moment we do not have a definitive workflow for exporting EnCodec models, for129instance to Hugging Face (HF). We are working on supporting automatic convertion between130AudioCraft and Hugging Face implementations.131 132We still have some support for fine tuning an EnCodec model coming from HF in AudioCraft,133using for instance `continue_from=//pretrained/facebook/encodec_32k`.134 135An AudioCraft checkpoint can be exported in a more compact format (excluding the optimizer etc.)136using `audiocraft.utils.export.export_encodec`. For instance, you could run137 138```python139from audiocraft.utils import export140from audiocraft import train141xp = train.main.get_xp_from_sig('SIG')142export.export_encodec(143    xp.folder / 'checkpoint.th',144    '/checkpoints/my_audio_lm/compression_state_dict.bin')145 146 147from audiocraft.models import CompressionModel148model = CompressionModel.get_pretrained('/checkpoints/my_audio_lm/compression_state_dict.bin')149 150from audiocraft.solvers import CompressionSolver151# The two are strictly equivalent, but this function supports also loading from non already exported models.152model = CompressionSolver.model_from_checkpoint('//pretrained//checkpoints/my_audio_lm/compression_state_dict.bin')153```154 155We will see then how to use this model as a tokenizer for MusicGen/Audio gen in the156[MusicGen documentation](./MUSICGEN.md).157 158### Learn more159 160Learn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).161 162 163## Citation164```165@article{defossez2022highfi,166  title={High Fidelity Neural Audio Compression},167  author={Défossez, Alexandre and Copet, Jade and Synnaeve, Gabriel and Adi, Yossi},168  journal={arXiv preprint arXiv:2210.13438},169  year={2022}170}171```172 173 174## License175 176See license information in the [README](../README.md).177 178[arxiv]: https://arxiv.org/abs/2210.13438179[encodec_samples]: https://ai.honu.io/papers/encodec/samples.html180