Repository files navigation

Paper accepted at IJCAI-25

(Trying to support audio-to-audio generation is in my todo list, such that you can input music tracks and the LDM will help merge them~)

  Paper accepted at ICCV-25

Setting up is super easy! Just follow the instructions below:

pip install -r gradio/requirements.txtpython gradio/gradio_app.py

We have to admit that the Unet architecture still has some probability advantage in subjective musicality, but this is not measured in the metric.And, we did have some models that were better on the metric, or trained for longer, but we observed that the models generally became less musicality after training too long, so we picked a model that was moderate on the metric as an open source sample. If you need more models (extreme metric pursuit or extreme musically pursuit, please contact me)

without any fancy design, just a quality injection, and enjoy your beautiful music

Down the main checkpoint of our QA-MDT model fromhttps://huggingface.co/lichang0928/QA-MDT

For chinese users, you can also download your checkpoint through following link:

https://pan.baidu.com/s/1N0XqVxtF_x9I7fWb07LPqw?pwd=9nkt

This repository provides an implementation of QA-MDT, integrating state-of-the-art models for music generation. The code and methods are based on the following repositories:

• AudioLDM
• PixArt-alpha
• MDT
• AudioMAE
• Open-Sora

Python 3.10qamdt.yaml

Before training, you need to download extra ckpts needed in ./audioldm_train/config/mos_as_token/qa_mdt.yaml and offset_pretrained_checkpoints.json

Noted that: All above checkpoints can be downloaded from:

flan-t5-large

clap_music

roberta-base

others

sh run.sh

Our model is already well-pretrained. If you wish to retrain or fine-tune it, you can choose to use or not use our QA strategy. We offer several training strategies:

• MDT w.o quality token:PixArt_MDT
• MDT with quality token:Pixart_MDT_MOS_AS_TOKEN
• DiT:PixArt_Slow
• U-net w / w.o quality prefix:you can just follow AudioLDM and make your dataset as illustrated in our paper (method part)

To train or fine-tune, simply change"Your_Class" inaudioldm_train.modules.diffusionmodules.PixArt.Your_Class in ourconfig file.

you can also try modifying the patch size, overlap size for your best performance and computing resources trade off (see our Appendix in arXiv paper)

We use theLMDB dataset format for training. You can modify the dataloader according to your own training needs.

If you'd like to follow our process (though we don't recommend it, as it can be complex), here's how you can create a toy LMDB dataset:

1. Create a Proto File

   First, create a file nameddatum_all.proto with the following content:

   syntax="proto2";messageDatum_all {repeatedfloatwav_file=1;requiredstringcaption_original=2;repeatedstringcaption_generated=3;requiredfloatmos=4;}

2. Generate Python Bindings

(Your protoc version should be 3.4, and you can download ithere)

Run the following command in your terminal to generate Python bindings:

protoc --python_out=./ datum_all.proto

This will create a file calleddatum_all_pb2.py. We have also provided this file in our datasets folder, and you can check if it matches the one you generated.Never attempt to modify this file, as doing so could cause errors.

3. Code for Preparing a toy LMDB Dataset

The following Python script demonstrates how to prepare your dataset in the LMDB format:

importtorchimportosimportlmdbimporttimeimportnumpyasnpimportlibrosaimportosimportsoundfileassfimportiofromdatum_all_pb2importDatum_allasDatum_outdevice='cpu'count=0total_hours=0# Define pathslmdb_file='/disk1/changli/toy_lmdb'toy_path='/disk1/changli/audioset'lmdb_key=os.path.join(lmdb_file,'data_key.key')# Open LMDB environmentenv=lmdb.open(lmdb_file,map_size=1e12)txn=env.begin(write=True)final_keys= []def_resample_load_librosa(path:str,sample_rate:int,downmix_to_mono:bool,**kwargs):"""Load and resample audio using librosa."""src,sr=librosa.load(path,sr=sample_rate,mono=downmix_to_mono,**kwargs)returnsrcstart_time=time.time()# Walk through the dataset directoryforroot,_,filesinos.walk(toy_path):forfileinfiles:audio_path=os.path.join(root,file)key_tmp=audio_path.replace('/','_')audio=_resample_load_librosa(audio_path,16000,True)# Create a new Datum objectdatum=Datum_out()datum.wav_file.extend(audio)datum.caption_original='audio'.encode()datum.caption_generated.append('audio'.encode())datum.mos=-1# Write to LMDBtxn.put(key_tmp.encode(),datum.SerializeToString())final_keys.append(key_tmp)count+=1total_hours+=1.00/60/10ifcount%1==0:elapsed_time=time.time()-start_timeprint(f'{count} files written, time:{elapsed_time:.2f}s')txn.commit()txn=env.begin(write=True)# Finalize transactiontry:total_time=time.time()-start_timeprint(f'Packing completed:{count} files written, total_hours:{total_hours:.2f}, time:{total_time:.2f}s')txn.commit()except:passenv.close()# Save the LMDB keyswithopen(lmdb_key,'w')asf:forkeyinfinal_keys:f.write(key+'\n')

otherwise, you can just use the dataloader inAudioLDM

4. Input your generated lmdb path and its corresponding key file path into the config

5. Start your training

sh infer/infer.sh# you may change the infer.sh for witch quality level you want to infer# defaultly, it should be set to 5 which represent highest quality# Additionally, it may be useful to change the prompt with text prefix "high quality",# which match the training process and may further improve performance

This is the first time I open source such a project, the code, the organization, the open source may not be perfect.If you have any questions about our model, code and datasets, feel free to contact me via below links, and I'm looking forward to any suggestions:

• Email:lc_lca@mail.ustc.edu.cn
• WeChat: 19524292801

I will try my best to provide more projects on music and audio in the future ~

If you find this project useful, please consider citing:

@article{li2024quality,  title={Quality-aware Masked Diffusion Transformer for Enhanced Music Generation},  author={Li, Chang and Wang, Ruoyu and Liu, Lijuan and Du, Jun and Sun, Yixuan and Guo, Zilu and Zhang, Zhenrong and Jiang, Yuan},  journal={arXiv preprint arXiv:2405.15863},  year={2024}}