OpenMusic / qa_mdt /readme.md
jadechoghari's picture
add model
9b9e0ee verified
# Awesome text to music generation (TTM) !!! : QA-MDT
## Official Pytorch Implementation
## I recommend anyone to listen to our demo, even under the clutter of tabs in Musiccaps, we still perform well
**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**
<a href="https://arxiv.org/pdf/2405.15863"><img src="https://img.shields.io/static/v1?label=Paper&message=arXiv.2405&color=red&logo=arXiv"></a> &ensp;
<a href="https://qa-mdt.github.io/"><img src="https://img.shields.io/static/v1?label=Demo&message=QA-MDT&color=black&logo=github.io"></a> &ensp;
<a href="https://huggingface.co/lichang0928/QA-MDT"><img src="https://img.shields.io/static/v1?label=ckpts&message=huggingface&color=yellow&logo=huggingface.co"></a> &ensp;
Down the main checkpoint of our QA-MDT model from [https://huggingface.co/lichang0928/QA-MDT](https://huggingface.co/lichang0928/QA-MDT)
For chinese users, you can also download your checkpoint through following link:
https://pan.baidu.com/s/1pkLnQhbNeFjKRadXUy_7Iw?pwd=v9dd
## Overview
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](https://github.com/haoheliu/AudioLDM-training-finetuning)
- [PixArt-alpha](https://github.com/PixArt-alpha/PixArt-alpha)
- [MDT](https://github.com/sail-sg/MDT)
- [AudioMAE](https://github.com/facebookresearch/AudioMAE)
- [Open-Sora](https://github.com/hpcaitech/Open-Sora)
## Requirements
```bash
Python 3.10
qamdt.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](https://huggingface.co/google/flan-t5-large)
[clap_music](https://huggingface.co/lukewys/laion_clap/blob/main/music_speech_audioset_epoch_15_esc_89.98.pt)
[roberta-base](https://huggingface.co/FacebookAI/roberta-base)
[others](https://drive.google.com/file/d/1T6EnuAHIc8ioeZ9kB1OZ_WGgwXAVGOZS/view)
## Training
```bash
sh run.sh
```
### How to Prepare for Training or Fine-tuning
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"` in `audioldm_train.modules.diffusionmodules.PixArt.Your_Class` in our [config file](https://github.com/ivcylc/qa-mdt/blob/main/audioldm_train/config/mos_as_token/qa_mdt.yaml).
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)
#### How to Prepare Your Dataset for Training or Fine-tuning
We use the **LMDB** 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 named `datum_all.proto` with the following content:
```proto
syntax = "proto2";
message Datum_all {
repeated float wav_file = 1;
required string caption_original = 2;
repeated string caption_generated = 3;
required float mos = 4;
}
2. **Generate Python Bindings**
Run the following command in your terminal to generate Python bindings:
```bash
protoc --python_out=./ datum_all.proto
```
This will create a file called **datum_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:
```python
import torch
import os
import lmdb
import time
import numpy as np
import librosa
import os
import soundfile as sf
import io
from datum_all_pb2 import Datum_all as Datum_out
device = 'cpu'
count = 0
total_hours = 0
# Define paths
lmdb_file = '/disk1/changli/toy_lmdb'
toy_path = '/disk1/changli/audioset'
lmdb_key = os.path.join(lmdb_file, 'data_key.key')
# Open LMDB environment
env = 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)
return src
start_time = time.time()
# Walk through the dataset directory
for root, _, files in os.walk(toy_path):
for file in files:
audio_path = os.path.join(root, file)
key_tmp = audio_path.replace('/', '_')
audio = _resample_load_librosa(audio_path, 16000, True)
# Create a new Datum object
datum = Datum_out()
datum.wav_file.extend(audio)
datum.caption_original = 'audio'.encode()
datum.caption_generated.append('audio'.encode())
datum.mos = -1
# Write to LMDB
txn.put(key_tmp.encode(), datum.SerializeToString())
final_keys.append(key_tmp)
count += 1
total_hours += 1.00 / 60 / 10
if count % 1 == 0:
elapsed_time = time.time() - start_time
print(f'{count} files written, time: {elapsed_time:.2f}s')
txn.commit()
txn = env.begin(write=True)
# Finalize transaction
try:
total_time = time.time() - start_time
print(f'Packing completed: {count} files written, total_hours: {total_hours:.2f}, time: {total_time:.2f}s')
txn.commit()
except:
pass
env.close()
# Save the LMDB keys
with open(lmdb_key, 'w') as f:
for key in final_keys:
f.write(key + '\n')
```
4. **Input your generated lmdb path and its corresponding key file path into the config**
5. **Start your training**
## Inference
()
```bash
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
# Additionly, it may be useful to change the prompt with text prefix "high quality",
# which match the training process and may further improve performance
```
## Contact
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](mailto:lc_lca@mail.ustc.edu.cn)
- **WeChat**: 19524292801
## Citation
If you find this project useful, please consider citing:
```bash
@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}
}
```