README.md 6.13 KB
Newer Older
mrfakename's avatar
mrfakename committed
1
# TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization 
Soujanya Poria's avatar
Soujanya Poria committed
2
3

<div align="center">
Soujanya Poria's avatar
Soujanya Poria committed
4
  <img src="assets/tf_teaser.png" alt="TangoFlux" width="1000" />
Soujanya Poria's avatar
Soujanya Poria committed
5
6
7

<br/>

mrfakename's avatar
mrfakename committed
8
[![arXiv](https://img.shields.io/badge/Read_the_Paper-blue?link=https%3A%2F%2Fopenreview.net%2Fattachment%3Fid%3DtpJPlFTyxd%26name%3Dpdf)](https://arxiv.org/abs/2412.21037) [![Static Badge](https://img.shields.io/badge/TangoFlux-Huggingface-violet?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fdeclare-lab%2FTangoFlux)](https://huggingface.co/declare-lab/TangoFlux) [![Static Badge](https://img.shields.io/badge/Demos-declare--lab-brightred?style=flat)](https://tangoflux.github.io/) [![Static Badge](https://img.shields.io/badge/TangoFlux-Hugging_Face_Space-8A2BE2?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fdeclare-lab%2FTangoFlux)](https://huggingface.co/spaces/declare-lab/TangoFlux) [![Static Badge](https://img.shields.io/badge/TangoFlux_Dataset-Huggingface-red?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fdatasets%2Fdeclare-lab%2FTangoFlux)](https://huggingface.co/datasets/declare-lab/CRPO) [![Replicate](https://replicate.com/chenxwh/tangoflux/badge)](https://replicate.com/chenxwh/tangoflux)
Soujanya Poria's avatar
Soujanya Poria committed
9
10
11

</div>

mrfakename's avatar
mrfakename committed
12
13
14
## Demos

[![Hugging Face Space](https://img.shields.io/badge/Hugging_Face_Space-TangoFlux-blue?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fdeclare-lab%2FTangoFlux)](https://huggingface.co/spaces/declare-lab/TangoFlux)
Chia-Yu Hung's avatar
Chia-Yu Hung committed
15

mrfakename's avatar
mrfakename committed
16
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1j__4fl_BlaVS_225M34d-EKxsVDJPRiR?usp=sharing)
Chia-Yu Hung's avatar
Chia-Yu Hung committed
17

Soujanya Poria's avatar
Soujanya Poria committed
18
## Overall Pipeline
mrfakename's avatar
mrfakename committed
19

Navonil Majumder's avatar
Navonil Majumder committed
20
TangoFlux consists of FluxTransformer blocks, which are Diffusion Transformers (DiT) and Multimodal Diffusion Transformers (MMDiT) conditioned on a textual prompt and a duration embedding to generate a 44.1kHz audio up to 30 seconds long. TangoFlux learns a rectified flow trajectory to an audio latent representation encoded by a variational autoencoder (VAE). TangoFlux training pipeline consists of three stages: pre-training, fine-tuning, and preference optimization with CRPO. CRPO, particularly, iteratively generates new synthetic data and constructs preference pairs for preference optimization using DPO loss for flow matching.
Soujanya Poria's avatar
Soujanya Poria committed
21

Soujanya Poria's avatar
Soujanya Poria committed
22
![cover-photo](assets/tangoflux.png)
Soujanya Poria's avatar
Soujanya Poria committed
23

mrfakename's avatar
mrfakename committed
24
25
26
🚀 **TangoFlux can generate 44.1kHz stereo audio up to 30 seconds in ~3 seconds on a single A40 GPU.**

## Installation
Soujanya Poria's avatar
Soujanya Poria committed
27

mrfakename's avatar
mrfakename committed
28
29
30
```bash
pip install tangoflux
```
Soujanya Poria's avatar
Soujanya Poria committed
31

mrfakename's avatar
mrfakename committed
32
33
34
35
36
37
38
## Inference

TangoFlux can generate audio up to 30 seconds long. You must pass a duration to the `model.generate` function when using the Python API. Please note that duration should be between 1 and 30.

### Web Interface

Run the following command to start the web interface:
Chia-Yu Hung's avatar
Chia-Yu Hung committed
39

Navonil Majumder's avatar
Navonil Majumder committed
40
```bash
mrfakename's avatar
mrfakename committed
41
tangoflux-demo
Chia-Yu Hung's avatar
Chia-Yu Hung committed
42
```
Chia-Yu Hung's avatar
Chia-Yu Hung committed
43

mrfakename's avatar
mrfakename committed
44
45
46
### CLI

Use the CLI to generate audio from text.
Chia-Yu Hung's avatar
Chia-Yu Hung committed
47
48

```bash
mrfakename's avatar
mrfakename committed
49
tangoflux "Hammer slowly hitting the wooden table" output.wav --duration 10 --steps 50
Chia-Yu Hung's avatar
Chia-Yu Hung committed
50
```
mrfakename's avatar
mrfakename committed
51
52
53

### Python API

Chia-Yu Hung's avatar
Chia-Yu Hung committed
54
55
```python
import torchaudio
hungchiayu1's avatar
updates  
hungchiayu1 committed
56
from tangoflux import TangoFluxInference
Chia-Yu Hung's avatar
Chia-Yu Hung committed
57

hungchiayu1's avatar
updates  
hungchiayu1 committed
58
59
model = TangoFluxInference(name='declare-lab/TangoFlux')
audio = model.generate('Hammer slowly hitting the wooden table', steps=50, duration=10)
Chia-Yu Hung's avatar
Chia-Yu Hung committed
60

mrfakename's avatar
mrfakename committed
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
torchaudio.save('output.wav', audio, 44100)
```

Our evaluation shows that inference with 50 steps yields the best results. A CFG scale of 3.5, 4, and 4.5 yield similar quality output. Inference with 25 steps yields similar audio quality at a faster speed.

## Training

We use the `accelerate` package from Hugging Face for multi-GPU training. Run `accelerate config` to setup your run configuration. The default accelerate config is in the `configs` folder. Please specify the path to your training files in the `configs/tangoflux_config.yaml`. Samples of `train.json` and `val.json` have been provided. Replace them with your own audio.

`tangoflux_config.yaml` defines the training file paths and model hyperparameters:

```bash
CUDA_VISIBLE_DEVICES=0,1 accelerate launch --config_file='configs/accelerator_config.yaml' src/train.py   --checkpointing_steps="best" --save_every=5 --config='configs/tangoflux_config.yaml'
```

To perform DPO training, modify the training files such that each data point contains "chosen", "reject", "caption" and "duration" fields. Please specify the path to your training files in `configs/tangoflux_config.yaml`. An example has been provided in `train_dpo.json`. Replace it with your own audio.

```bash
CUDA_VISIBLE_DEVICES=0,1 accelerate launch --config_file='configs/accelerator_config.yaml' src/train_dpo.py   --checkpointing_steps="best" --save_every=5 --config='configs/tangoflux_config.yaml'
hungchiayu1's avatar
updates  
hungchiayu1 committed
80
```
Soujanya Poria's avatar
Soujanya Poria committed
81
82
83

## Evaluation Scripts

Navonil Majumder's avatar
Navonil Majumder committed
84
## TangoFlux vs. Other Audio Generation Models
Soujanya Poria's avatar
Soujanya Poria committed
85

Navonil Majumder's avatar
Navonil Majumder committed
86
This key comparison metrics include:
Soujanya Poria's avatar
Soujanya Poria committed
87
88

- **Output Length**: Represents the duration of the generated audio.
Navonil Majumder's avatar
Navonil Majumder committed
89
- **FD**<sub>openl3</sub>: Fréchet Distance.
Soujanya Poria's avatar
Soujanya Poria committed
90
91
92
93
- **KL**<sub>passt</sub>: KL divergence.
- **CLAP**<sub>score</sub>: Alignment score.


Navonil Majumder's avatar
Navonil Majumder committed
94
All the inference times are observed on the same A40 GPU. The counts of trainable parameters are reported in the **\#Params** column.
Soujanya Poria's avatar
Soujanya Poria committed
95

mrfakename's avatar
mrfakename committed
96
97
98
99
100
101
102
| Model | Params | Duration | Steps | FD<sub>openl3</sub> ↓ | KL<sub>passt</sub> ↓ | CLAP<sub>score</sub> ↑ | IS ↑ | Inference Time (s) |
|---|---|---|---|---|---|---|---|---|
| **AudioLDM 2 (Large)** | 712M | 10 sec | 200 | 108.3 | 1.81 | 0.419 | 7.9 | 24.8 |
| **Stable Audio Open** | 1056M | 47 sec | 100 | 89.2 | 2.58 | 0.291 | 9.9 | 8.6 |
| **Tango 2** | 866M | 10 sec | 200 | 108.4 | 1.11 | 0.447 | 9.0 | 22.8 |
| **TangoFlux (Base)** | 515M | 30 sec | 50 | 80.2 | 1.22 | 0.431 | 11.7 | 3.7 |
| **TangoFlux** | 515M | 30 sec | 50 | 75.1 | 1.15 | 0.480 | 12.2 | 3.7 |
Soujanya Poria's avatar
Soujanya Poria committed
103
104
105

## Citation

Soujanya Poria's avatar
Soujanya Poria committed
106
107
108
109
110
111
112
113
114
```bibtex
@misc{hung2024tangofluxsuperfastfaithful,
      title={TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization}, 
      author={Chia-Yu Hung and Navonil Majumder and Zhifeng Kong and Ambuj Mehrish and Rafael Valle and Bryan Catanzaro and Soujanya Poria},
      year={2024},
      eprint={2412.21037},
      archivePrefix={arXiv},
      primaryClass={cs.SD},
      url={https://arxiv.org/abs/2412.21037}, 
Soujanya Poria's avatar
Soujanya Poria committed
115
}
Soujanya Poria's avatar
Soujanya Poria committed
116
```
mrfakename's avatar
mrfakename committed
117
118
119
120

## License

TangoFlux is licensed under the MIT License. See the `LICENSE` file for more details.