Awesome text to music generation (TTM) : QA-MDT

Official Pytorch Implementation

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

I recommend anyone to listen to our demo, even under the clutter of tabs in Musiccaps, we still perform well

Down the main checkpoint of our QA-MDT model from 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:

Requirements

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:

Training

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.

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:

Create a Proto File

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


syntax = "proto2";

message Datum_all {
 repeated float wav_file = 1;
 required string caption_original = 2;
 repeated string caption_generated = 3;
 required float mos = 4;
}

Generate Python Bindings

Run the following command in your terminal to generate Python bindings:
```
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.

Code for Preparing a toy LMDB Dataset

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

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')

Input your generated lmdb path and its corresponding key file path into the config
Start your training

Inference

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
WeChat: 19524292801

Citation

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}
}

ivcylc / qa-mdt

readme