audiolm-pytorch

AudioLM - Pytorch

Implementation of AudioLM, a Language Modeling Approach to Audio Generation out of Google Research, in Pytorch

It also extends the work for conditioning with classifier free guidance with T5. This allows for one to do text-to-audio or TTS, not offered in the paper. Yes, this means VALL-E can be trained from this repository. It is essentially the same.

Please join if you are interested in replicating this work in the open

This repository now also contains a MIT licensed version of SoundStream. It is also compatible with EnCodec, which is also MIT-licensed at the time of writing.

Update: AudioLM was essentially used to 'solve' music generation in the new MusicLM

In the future, this movie clip would no longer make any sense. You would just prompt an AI instead.

Appreciation

• Stability.ai for the generous sponsorship to work and open source cutting edge artificial intelligence research

• Huggingface for their amazing accelerate and transformers libraries

• MetaAI for Fairseq and the liberal license

• @eonglints and Joseph for offering their professional advice and expertise as well as pull requests!

• @djqualia, @yigityu, @inspirit, and @BlackFox1197 for helping with the debugging of soundstream

• Allen and LWprogramming for reviewing the code and submitting bug fixes!

• Ilya for finding an issue with multi-scale discriminator downsampling and for soundstream trainer improvements

• Andrey for identifying a missing loss in soundstream and guiding me through the proper mel spectrogram hyperparameters

• Alejandro and Ilya for sharing their results with training soundstream, and for working through a few issues with the local attention positional embeddings

• LWprogramming for adding Encodec compatibility!

• LWprogramming for finding an issue with handling of the EOS token when sampling from the FineTransformer!

• @YoungloLee for identifying a big bug in the 1d causal convolution for soundstream related to padding not accounting for strides!

• Hayden for pointing out some discrepancies in the multi-scale discriminator for Soundstream

Install

$ pip install audiolm-pytorch

Usage

SoundStream & Encodec

There are two options for the neural codec. If you want to use the pretrained 24kHz Encodec, just create an Encodec object as follows:

from audiolm_pytorch import EncodecWrapper
encodec = EncodecWrapper()
# Now you can use the encodec variable in the same way you'd use the soundstream variables below.

Otherwise, to stay more true to the original paper, you can use SoundStream. First, SoundStream needs to be trained on a large corpus of audio data

…

Your trained SoundStream can then be used as a generic tokenizer for audio

audio = torch.randn(1, 512 * 320)

codes = soundstream.tokenize(audio)

# you can now train anything with the codebook ids

recon_audio_from_codes = soundstream.decode_from_codebook_indices(codes)

# sanity check

assert torch.allclose(
    recon_audio_from_codes,
    soundstream(audio, return_recons_only = True)
)

You can also use soundstreams that are specific to AudioLM and MusicLM by importing AudioLMSoundStream and MusicLMSoundStream respectively

from audiolm_pytorch import AudioLMSoundStream, MusicLMSoundStream

soundstream = AudioLMSoundStream(...) # say you want the hyperparameters as in Audio LM paper

# rest is the same as above

As of version 0.17.0, you can now invoke the class method on SoundStream to load from checkpoint files, without having to remember your configurations.

from audiolm_pytorch import SoundStream

soundstream = SoundStream.init_and_load_from('./path/to/checkpoint.pt')

To use Weights & Biases tracking, first set use_wandb_tracking = True on the SoundStreamTrainer, then do the following

trainer = SoundStreamTrainer(
    soundstream,
    ...,
    use_wandb_tracking = True
)

# wrap .train() with contextmanager, specifying project and run name

with trainer.wandb_tracker(project = 'soundstream', run = 'baseline'):
    trainer.train()

Hierarchical Transformers

Then three separate transformers (SemanticTransformer, CoarseTransformer, FineTransformer) need to be trained

ex. SemanticTransformer

…

ex. CoarseTransformer

…

ex. FineTransformer

…

All together now

from audiolm_pytorch import AudioLM

audiolm = AudioLM(
    wav2vec = wav2vec,
    codec = soundstream,
    semantic_transformer = semantic_transformer,
    coarse_transformer = coarse_transformer,
    fine_transformer = fine_transformer
)

generated_wav = audiolm(batch_size = 1)

# or with priming

generated_wav_with_prime = audiolm(prime_wave = torch.randn(1, 320 * 8))

# or with text condition, if given

generated_wav_with_text_condition = audiolm(text = ['chirping of birds and the distant echos of bells'])

Text Conditioned Audio Synthesis

Update: Looks like this will work, given 'VALL-E'

ex. Semantic Transformer

…

Multi-GPU

Because all the trainer classes uses Accelerator, you can easily do multi gpu training by using the accelerate command as so

At the project root

$ accelerate config

Then, in the same directory

$ accelerate launch train.py

Todo

• complete CoarseTransformer

• use fairseq vq-wav2vec for embeddings

• add conditioning

• add classifier free guidance

• add unique consecutive for

• incorporate ability to use hubert intermediate features as semantic tokens, recommended by eonglints

• accommodate variable lengthed audio, bring in eos token

• make sure unique consecutive works with coarse transformer

• pretty printing all discriminator losses to log

• handle when generating semantic tokens, that last logits may not be necessarily the last in the sequence given unique consecutive processing

• complete sampling code for both Coarse and Fine Transformers, which will be tricky

• make sure full inference with or without prompting works on the AudioLM class

• complete full training code for soundstream, taking care of discriminator training

• add efficient gradient penalty for discriminators for soundstream

• wire up sample hz from sound dataset -> transformers, and have proper resampling within during training - think about whether to allow for dataset to have sound files of varying or enforce same sample hz

• full transformer training code for all three transformers

• refactor so semantic transformer has a wrapper to that handles unique consecutives as well as wav to hubert or vq-wav2vec

• simply not self attend to eos token on the prompting side (semantic for coarse transformer, coarse for fine transformer)

• add structured dropout from forgetful causal masking, far better than traditional dropouts

• figure out how to suppress logging in fairseq

• assert that all three transformers passed into audiolm is compatible

• allow for specialized relative positional embeddings in fine transformer based on absolute matching positions of quantizers between coarse and fine

• allow for grouped residual vq in soundstream (use GroupedResidualVQ from vector-quantize-pytorch lib), from hifi-codec

• add flash attention with NoPE

• accept prime wave in AudioLM as a path to an audio file, and auto resample for semantic vs acoustic

• add key / value caching to all transformers, speeding up inference

• design a hierarchical coarse and fine transformer

• investigate spec decoding, first test in x-transformers, then port over if applicable

• redo the positional embeddings in the presence of groups in residual vq

• test with speech synthesis for starters

• cli tool, something like audiolm generate <wav.file | text> and save generated wav file to local directory

• return a list of waves in the case of variable lengthed audio

• just take care of the edge case in coarse transformer text conditioned training, where the raw wave is resampled at different frequencies. autodetermine how to route based on length

Citations

@inproceedings{Borsos2022AudioLMAL,
  title  = {AudioLM: a Language Modeling Approach to Audio Generation},
  author = {Zal{\'a}n Borsos and Rapha{\"e}l Marinier and Damien Vincent and Eugene Kharitonov and Olivier Pietquin and Matthew Sharifi and Olivier Teboul and David Grangier and Marco Tagliasacchi and Neil Zeghidour},
  year   = {2022}
}

@misc{https://doi.org/10.48550/arxiv.2107.03312,
  title  = {SoundStream: An End-to-End Neural Audio Codec},
  author = {Zeghidour, Neil and Luebs, Alejandro and Omran, Ahmed and Skoglund, Jan and Tagliasacchi, Marco},
  publisher = {arXiv},
  url    = {https://arxiv.org/abs/2107.03312},
  year   = {2021}
}

@misc{shazeer2020glu,
    title   = {GLU Variants Improve Transformer},
    author  = {Noam Shazeer},
    year    = {2020},
    url     = {https://arxiv.org/abs/2002.05202}
}

@article{Shazeer2019FastTD,
    title   = {Fast Transformer Decoding: One Write-Head is All You Need},
    author  = {Noam M. Shazeer},
    journal = {ArXiv},
    year    = {2019},
    volume  = {abs/1911.02150}
}

@article{Ho2022ClassifierFreeDG,
    title   = {Classifier-Free Diffusion Guidance},
    author  = {Jonathan Ho},
    journal = {ArXiv},
    year    = {2022},
    volume  = {abs/2207.12598}
}

@misc{crowson2022,
    author  = {Katherine Crowson},
    url     = {https://twitter.com/rivershavewings}
}

@misc{ding2021cogview,
    title   = {CogView: Mastering Text-to-Image Generation via Transformers},
    author  = {Ming Ding and Zhuoyi Yang and Wenyi Hong and Wendi Zheng and Chang Zhou and Da Yin and Junyang Lin and Xu Zou and Zhou Shao and Hongxia Yang and Jie Tang},
    year    = {2021},
    eprint  = {2105.13290},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}

@article{Liu2022FCMFC,
    title   = {FCM: Forgetful Causal Masking Makes Causal Language Models Better Zero-Shot Learners},
    author  = {Hao Liu and Xinyang Geng and Lisa Lee and Igor Mordatch and Sergey Levine and Sharan Narang and P. Ab