Skip to content

GitLab

"object_detection/matchers/argmax_matcher_test.py" did not exist on "60c3ed2e34efbe428ae04ab99ccebd795187c12f"
Commit 0db67215 authored by mrfakename's avatar mrfakename
Browse files

pip package

parent f7cfcd21
Show whitespace changes
Inline Side-by-side
Showing with 1159 additions and 734 deletions
.gitignore
View file @ 0db67215
~__pycache__/
__pycache__/
*.py[cod]
*$py.class
......@@ -168,3 +168,8 @@ cython_debug/
# PyPI configuration file
.pypirc
.DS_Store
*.wav
Inference.ipynb
View file @ 0db67215
This source diff could not be displayed because it is too large. You can view the blob instead.
README.md
View file @ 0db67215
<h1 align="center">✨
<br/>
TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization
<br/>
✨✨✨
</h1>
# TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization
<div align="center">
<img src="assets/tf_teaser.png" alt="TangoFlux" width="1000" />
<br/>
[![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-Huggingface_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)
[![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)
</div>
## Quickstart on Google Colab
## 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)
| Colab |
| --- |
[![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)
## Overall Pipeline
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.
![cover-photo](assets/tangoflux.png)
🚀 **TangoFlux can generate 44.1kHz stereo audio up to 30 seconds in ~3 seconds on a single A40 GPU.**
🚀 **TangoFlux can generate up to 30 seconds long 44.1kHz stereo audios in about 3 seconds.**
## Installation
## Training TangoFlux
We use the accelerate package from HuggingFace for multi-gpu training. Run accelerate config from terminal and set up your run configuration by the answering the questions asked. We have placed the default accelerator config in the `configs` folder. Please specify the path to your training files in the configs/tangoflux_config.yaml. A sample of train.json and val.json has been provided. Replace them with your own audio.
```bash
pip install tangoflux
```
## 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:
`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'
tangoflux-demo
```
To perform DPO training, modify the training files such that each data point contains a "chosen","reject","caption" and "duration". Please specify the path to your training files in the configs/tangoflux_config.yaml. An example has been provided in train_dpo.json. Replace them with your own audio.
### CLI
Use the CLI to generate audio from text.
```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'
tangoflux "Hammer slowly hitting the wooden table" output.wav --duration 10 --steps 50
```
## Inference with TangoFlux
Download the TangoFlux model and generate audio from a text prompt.
TangoFlux can generate audios up to 30 second long through passing in a duration variable in the `model.generate` function. Please note that duration should be strictly greather than 1 and lesser than 30.
### Python API
```python
import torchaudio
from tangoflux import TangoFluxInference
from IPython.display import Audio
model = TangoFluxInference(name='declare-lab/TangoFlux')
audio = model.generate('Hammer slowly hitting the wooden table', steps=50, duration=10)
Audio(data=audio, rate=44100)
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'
```
Our evaluation shows that inferring with 50 steps yield the best results. A CFG scale of 3.5, 4, and 4.5 yield simliar quality output.
For faster inference, consider setting steps to 25 that yield similar audio quality.
## Evaluation Scripts
......@@ -77,15 +93,13 @@ This key comparison metrics include:
All the inference times are observed on the same A40 GPU. The counts of trainable parameters are reported in the **\#Params** column.
| 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 |
| 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** |
| **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 |
## Citation
......@@ -100,3 +114,7 @@ All the inference times are observed on the same A40 GPU. The counts of trainabl
url={https://arxiv.org/abs/2412.21037},
}
```
## License
TangoFlux is licensed under the MIT License. See the `LICENSE` file for more details.
\ No newline at end of file
data/.gitkeep 0 → 100644
View file @ 0db67215
*.wav
\ No newline at end of file
inference.py 0 → 100644
View file @ 0db67215
import torchaudio
from tangoflux import TangoFluxInference
model = TangoFluxInference(name="declare-lab/TangoFlux")
audio = model.generate("Hammer slowly hitting the wooden table", steps=50, duration=10)
torchaudio.save("output.wav", audio, sample_rate=44100)
replicate_demo/predict.py
View file @ 0db67215
......@@ -10,11 +10,13 @@ from diffusers import AutoencoderOobleck
import soundfile as sf
from safetensors.torch import load_file
from huggingface_hub import snapshot_download
from src.model import TangoFlux
from tangoflux.model import TangoFlux
from tangoflux import TangoFluxInference
MODEL_CACHE = "model_cache"
MODEL_URL = "https://weights.replicate.delivery/default/declare-lab/TangoFlux/model_cache.tar"
MODEL_URL = (
"https://weights.replicate.delivery/default/declare-lab/TangoFlux/model_cache.tar"
)
class CachedTangoFluxInference(TangoFluxInference):
......
setup.py 0 → 100644
View file @ 0db67215
from setuptools import setup
setup(
name="tangoflux",
description="TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching",
version="0.1.0",
packages=["tangoflux"],
install_requires=[
"torch==2.4.0",
"torchaudio==2.4.0",
"torchlibrosa==0.1.0",
"torchvision==0.19.0",
"transformers==4.44.0",
"diffusers==0.30.0",
"accelerate==0.34.2",
"datasets==2.21.0",
"librosa",
"tqdm",
"wandb",
"click",
"gradio",
"torchaudio",
],
entry_points={
"console_scripts": [
"tangoflux=tangoflux.cli:main",
"tangoflux-demo=tangoflux.demo:main",
],
},
)
tangoflux.py → tangoflux/__init__.py
View file @ 0db67215
from diffusers import AutoencoderOobleck
import torch
from transformers import T5EncoderModel,T5TokenizerFast
from transformers import T5EncoderModel, T5TokenizerFast
from diffusers import FluxTransformer2DModel
from torch import nn
from typing import List
......@@ -9,10 +9,10 @@ from diffusers.training_utils import compute_density_for_timestep_sampling
import copy
import torch.nn.functional as F
import numpy as np
from src.model import TangoFlux
from tangoflux.model import TangoFlux
from huggingface_hub import snapshot_download
from tqdm import tqdm
from typing import Optional,Union,List
from typing import Optional, Union, List
from datasets import load_dataset, Audio
from math import pi
import json
......@@ -23,8 +23,11 @@ from safetensors.torch import load_file
class TangoFluxInference:
def __init__(self,name='declare-lab/TangoFlux',device="cuda"):
def __init__(
self,
name="declare-lab/TangoFlux",
device="cuda" if torch.cuda.is_available() else "cpu",
):
self.vae = AutoencoderOobleck()
......@@ -33,29 +36,25 @@ class TangoFluxInference:
self.vae.load_state_dict(vae_weights)
weights = load_file("{}/tangoflux.safetensors".format(paths))
with open('{}/config.json'.format(paths),'r') as f:
with open("{}/config.json".format(paths), "r") as f:
config = json.load(f)
self.model = TangoFlux(config)
self.model.load_state_dict(weights,strict=False)
self.model.load_state_dict(weights, strict=False)
# _IncompatibleKeys(missing_keys=['text_encoder.encoder.embed_tokens.weight'], unexpected_keys=[]) this behaviour is expected
self.vae.to(device)
self.model.to(device)
def generate(self,prompt,steps=25,duration=10,guidance_scale=4.5):
def generate(self, prompt, steps=25, duration=10, guidance_scale=4.5):
with torch.no_grad():
latents = self.model.inference_flow(prompt,
latents = self.model.inference_flow(
prompt,
duration=duration,
num_inference_steps=steps,
guidance_scale=guidance_scale)
guidance_scale=guidance_scale,
)
wave = self.vae.decode(latents.transpose(2,1)).sample.cpu()[0]
wave = self.vae.decode(latents.transpose(2, 1)).sample.cpu()[0]
waveform_end = int(duration * self.vae.config.sampling_rate)
wave = wave[:, :waveform_end]
return wave
tangoflux/cli.py 0 → 100644
View file @ 0db67215
import click
import torchaudio
from tangoflux import TangoFluxInference
@click.command()
@click.argument('prompt')
@click.argument('output_file')
@click.option('--duration', default=10, type=int, help='Duration in seconds (1-30)')
@click.option('--steps', default=50, type=int, help='Number of inference steps (10-100)')
def main(prompt: str, output_file: str, duration: int, steps: int):
"""Generate audio from text using TangoFlux.
Args:
prompt: Text description of the audio to generate
output_file: Path to save the generated audio file
duration: Duration of generated audio in seconds (default: 10)
steps: Number of inference steps (default: 50)
"""
if not 1 <= duration <= 30:
raise click.BadParameter('Duration must be between 1 and 30 seconds')
if not 10 <= steps <= 100:
raise click.BadParameter('Steps must be between 10 and 100')
model = TangoFluxInference(name="declare-lab/TangoFlux")
audio = model.generate(prompt, steps=steps, duration=duration)
torchaudio.save(output_file, audio, sample_rate=44100)
if __name__ == '__main__':
main()
tangoflux/demo.py 0 → 100644
View file @ 0db67215
import gradio as gr
import torchaudio
import click
import tempfile
from tangoflux import TangoFluxInference
model = TangoFluxInference(name="declare-lab/TangoFlux")
def generate_audio(prompt, duration, steps):
audio = model.generate(prompt, steps=steps, duration=duration)
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
torchaudio.save(f.name, audio, sample_rate=44100)
return f.name
examples = [
["Hammer slowly hitting the wooden table", 10, 50],
["Gentle rain falling on a tin roof", 15, 50],
["Wind chimes tinkling in a light breeze", 10, 50],
["Rhythmic wooden table tapping overlaid with steady water pouring sound", 10, 50],
]
with gr.Blocks(title="TangoFlux Text-to-Audio Generation") as demo:
gr.Markdown("# TangoFlux Text-to-Audio Generation")
gr.Markdown("Generate audio from text descriptions using TangoFlux")
with gr.Row():
with gr.Column():
prompt = gr.Textbox(
label="Text Prompt", placeholder="Enter your audio description..."
)
duration = gr.Slider(
minimum=1, maximum=30, value=10, step=1, label="Duration (seconds)"
)
steps = gr.Slider(
minimum=10, maximum=100, value=50, step=10, label="Number of Steps"
)
generate_btn = gr.Button("Generate Audio")
with gr.Column():
audio_output = gr.Audio(label="Generated Audio")
generate_btn.click(
fn=generate_audio, inputs=[prompt, duration, steps], outputs=audio_output
)
gr.Examples(
examples=examples,
inputs=[prompt, duration, steps],
outputs=audio_output,
fn=generate_audio,
)
@click.command()
@click.option('--host', default='127.0.0.1', help='Host to bind to')
@click.option('--port', default=None, help='Port to bind to')
@click.option('--share', is_flag=True, help='Enable sharing via Gradio')
def main(host, port, share):
demo.queue().launch(server_name=host, server_port=port, share=share)
if __name__ == "__main__":
main()
src/model.py → tangoflux/model.py
View file @ 0db67215
from transformers import T5EncoderModel,T5TokenizerFast
from transformers import T5EncoderModel, T5TokenizerFast
import torch
from diffusers import FluxTransformer2DModel
from torch import nn
import random
from typing import List
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.training_utils import compute_density_for_timestep_sampling
......@@ -11,19 +11,16 @@ import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
from typing import Optional,Union,List
from typing import Optional, Union, List
from datasets import load_dataset, Audio
from math import pi
import inspect
import yaml
class StableAudioPositionalEmbedding(nn.Module):
"""Used for continuous time
Adapted from stable audio open.
Adapted from Stable Audio Open.
"""
def __init__(self, dim: int):
......@@ -39,6 +36,7 @@ class StableAudioPositionalEmbedding(nn.Module):
fouriered = torch.cat((times, fouriered), dim=-1)
return fouriered
class DurationEmbedder(nn.Module):
"""
A simple linear projection model to map numbers to a latent space.
......@@ -81,7 +79,9 @@ class DurationEmbedder(nn.Module):
):
floats = floats.clamp(self.min_value, self.max_value)
normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value)
normalized_floats = (floats - self.min_value) / (
self.max_value - self.min_value
)
# Cast floats to same type as embedder
embedder_dtype = next(self.time_positional_embedding.parameters()).dtype
......@@ -103,9 +103,13 @@ def retrieve_timesteps(
):
if timesteps is not None and sigmas is not None:
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
raise ValueError(
"Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
)
if timesteps is not None:
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
accepts_timesteps = "timesteps" in set(
inspect.signature(scheduler.set_timesteps).parameters.keys()
)
if not accepts_timesteps:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
......@@ -115,7 +119,9 @@ def retrieve_timesteps(
timesteps = scheduler.timesteps
num_inference_steps = len(timesteps)
elif sigmas is not None:
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
accept_sigmas = "sigmas" in set(
inspect.signature(scheduler.set_timesteps).parameters.keys()
)
if not accept_sigmas:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
......@@ -130,30 +136,22 @@ def retrieve_timesteps(
return timesteps, num_inference_steps
class TangoFlux(nn.Module):
def __init__(self,config,initialize_reference_model=False):
def __init__(self, config, initialize_reference_model=False):
super().__init__()
self.num_layers = config.get('num_layers', 6)
self.num_single_layers = config.get('num_single_layers', 18)
self.in_channels = config.get('in_channels', 64)
self.attention_head_dim = config.get('attention_head_dim', 128)
self.joint_attention_dim = config.get('joint_attention_dim', 1024)
self.num_attention_heads = config.get('num_attention_heads', 8)
self.audio_seq_len = config.get('audio_seq_len', 645)
self.max_duration = config.get('max_duration', 30)
self.uncondition = config.get('uncondition', False)
self.text_encoder_name = config.get('text_encoder_name', "google/flan-t5-large")
self.num_layers = config.get("num_layers", 6)
self.num_single_layers = config.get("num_single_layers", 18)
self.in_channels = config.get("in_channels", 64)
self.attention_head_dim = config.get("attention_head_dim", 128)
self.joint_attention_dim = config.get("joint_attention_dim", 1024)
self.num_attention_heads = config.get("num_attention_heads", 8)
self.audio_seq_len = config.get("audio_seq_len", 645)
self.max_duration = config.get("max_duration", 30)
self.uncondition = config.get("uncondition", False)
self.text_encoder_name = config.get("text_encoder_name", "google/flan-t5-large")
self.noise_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)
self.noise_scheduler_copy = copy.deepcopy(self.noise_scheduler)
......@@ -162,9 +160,12 @@ class TangoFlux(nn.Module):
self.tokenizer = T5TokenizerFast.from_pretrained(self.text_encoder_name)
self.text_embedding_dim = self.text_encoder.config.d_model
self.fc = nn.Sequential(nn.Linear(self.text_embedding_dim,self.joint_attention_dim),nn.ReLU())
self.duration_emebdder = DurationEmbedder(self.text_embedding_dim,min_value=0,max_value=self.max_duration)
self.fc = nn.Sequential(
nn.Linear(self.text_embedding_dim, self.joint_attention_dim), nn.ReLU()
)
self.duration_emebdder = DurationEmbedder(
self.text_embedding_dim, min_value=0, max_value=self.max_duration
)
self.transformer = FluxTransformer2DModel(
in_channels=self.in_channels,
......@@ -174,20 +175,15 @@ class TangoFlux(nn.Module):
num_attention_heads=self.num_attention_heads,
joint_attention_dim=self.joint_attention_dim,
pooled_projection_dim=self.text_embedding_dim,
guidance_embeds=False)
guidance_embeds=False,
)
self.beta_dpo = 2000 ## this is used for dpo training
def get_sigmas(self,timesteps, n_dim=3, dtype=torch.float32):
def get_sigmas(self, timesteps, n_dim=3, dtype=torch.float32):
device = self.text_encoder.device
sigmas = self.noise_scheduler_copy.sigmas.to(device=device, dtype=dtype)
schedule_timesteps = self.noise_scheduler_copy.timesteps.to(device)
timesteps = timesteps.to(device)
step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
......@@ -197,14 +193,18 @@ class TangoFlux(nn.Module):
sigma = sigma.unsqueeze(-1)
return sigma
def encode_text_classifier_free(self, prompt: List[str], num_samples_per_prompt=1):
device = self.text_encoder.device
batch = self.tokenizer(
prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
prompt,
max_length=self.tokenizer.model_max_length,
padding=True,
truncation=True,
return_tensors="pt",
)
input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(
device
)
input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
with torch.no_grad():
prompt_embeds = self.text_encoder(
......@@ -219,7 +219,11 @@ class TangoFlux(nn.Module):
max_length = prompt_embeds.shape[1]
uncond_batch = self.tokenizer(
uncond_tokens, max_length=max_length, padding='max_length', truncation=True, return_tensors="pt",
uncond_tokens,
max_length=max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
uncond_input_ids = uncond_batch.input_ids.to(device)
uncond_attention_mask = uncond_batch.attention_mask.to(device)
......@@ -229,8 +233,12 @@ class TangoFlux(nn.Module):
input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
)[0]
negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)
negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(
num_samples_per_prompt, 0
)
uncond_attention_mask = uncond_attention_mask.repeat_interleave(
num_samples_per_prompt, 0
)
# For classifier free guidance, we need to do two forward passes.
# We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
......@@ -245,98 +253,114 @@ class TangoFlux(nn.Module):
def encode_text(self, prompt):
device = self.text_encoder.device
batch = self.tokenizer(
prompt, max_length=self.max_text_seq_len, padding=True, truncation=True, return_tensors="pt")
input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
prompt,
max_length=self.max_text_seq_len,
padding=True,
truncation=True,
return_tensors="pt",
)
input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(
device
)
encoder_hidden_states = self.text_encoder(
input_ids=input_ids, attention_mask=attention_mask)[0]
input_ids=input_ids, attention_mask=attention_mask
)[0]
boolean_encoder_mask = (attention_mask == 1).to(device)
return encoder_hidden_states, boolean_encoder_mask
def encode_duration(self,duration):
def encode_duration(self, duration):
return self.duration_emebdder(duration)
@torch.no_grad()
def inference_flow(self, prompt,
def inference_flow(
self,
prompt,
num_inference_steps=50,
timesteps=None,
guidance_scale=3,
duration=10,
disable_progress=False,
num_samples_per_prompt=1):
'''Only tested for single inference. Haven't test for batch inference'''
num_samples_per_prompt=1,
):
"""Only tested for single inference. Haven't test for batch inference"""
bsz = num_samples_per_prompt
device = self.transformer.device
scheduler = self.noise_scheduler
if not isinstance(prompt,list):
if not isinstance(prompt, list):
prompt = [prompt]
if not isinstance(duration,torch.Tensor):
duration = torch.tensor([duration],device=device)
if not isinstance(duration, torch.Tensor):
duration = torch.tensor([duration], device=device)
classifier_free_guidance = guidance_scale > 1.0
duration_hidden_states = self.encode_duration(duration)
if classifier_free_guidance:
bsz = 2 * num_samples_per_prompt
encoder_hidden_states, boolean_encoder_mask = self.encode_text_classifier_free(prompt, num_samples_per_prompt=num_samples_per_prompt)
duration_hidden_states = duration_hidden_states.repeat(bsz,1,1)
encoder_hidden_states, boolean_encoder_mask = (
self.encode_text_classifier_free(
prompt, num_samples_per_prompt=num_samples_per_prompt
)
)
duration_hidden_states = duration_hidden_states.repeat(bsz, 1, 1)
else:
encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt,num_samples_per_prompt=num_samples_per_prompt)
encoder_hidden_states, boolean_encoder_mask = self.encode_text(
prompt, num_samples_per_prompt=num_samples_per_prompt
)
mask_expanded = boolean_encoder_mask.unsqueeze(-1).expand_as(encoder_hidden_states)
masked_data = torch.where(mask_expanded, encoder_hidden_states, torch.tensor(float('nan')))
mask_expanded = boolean_encoder_mask.unsqueeze(-1).expand_as(
encoder_hidden_states
)
masked_data = torch.where(
mask_expanded, encoder_hidden_states, torch.tensor(float("nan"))
)
pooled = torch.nanmean(masked_data, dim=1)
pooled_projection = self.fc(pooled)
encoder_hidden_states = torch.cat([encoder_hidden_states,duration_hidden_states],dim=1) ## (bs,seq_len,dim)
encoder_hidden_states = torch.cat(
[encoder_hidden_states, duration_hidden_states], dim=1
) ## (bs,seq_len,dim)
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps,
device,
timesteps,
sigmas
scheduler, num_inference_steps, device, timesteps, sigmas
)
latents = torch.randn(num_samples_per_prompt,self.audio_seq_len,64)
latents = torch.randn(num_samples_per_prompt, self.audio_seq_len, 64)
weight_dtype = latents.dtype
progress_bar = tqdm(range(num_inference_steps), disable=disable_progress)
txt_ids = torch.zeros(bsz,encoder_hidden_states.shape[1],3).to(device)
audio_ids = torch.arange(self.audio_seq_len).unsqueeze(0).unsqueeze(-1).repeat(bsz,1,3).to(device)
txt_ids = torch.zeros(bsz, encoder_hidden_states.shape[1], 3).to(device)
audio_ids = (
torch.arange(self.audio_seq_len)
.unsqueeze(0)
.unsqueeze(-1)
.repeat(bsz, 1, 3)
.to(device)
)
timesteps = timesteps.to(device)
latents = latents.to(device)
encoder_hidden_states = encoder_hidden_states.to(device)
for i, t in enumerate(timesteps):
latents_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
latents_input = (
torch.cat([latents] * 2) if classifier_free_guidance else latents
)
noise_pred = self.transformer(
hidden_states=latents_input,
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
timestep=torch.tensor([t/1000],device=device),
guidance = None,
timestep=torch.tensor([t / 1000], device=device),
guidance=None,
pooled_projections=pooled_projection,
encoder_hidden_states=encoder_hidden_states,
txt_ids=txt_ids,
......@@ -346,42 +370,45 @@ class TangoFlux(nn.Module):
if classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
noise_pred = noise_pred_uncond + guidance_scale * (
noise_pred_text - noise_pred_uncond
)
latents = scheduler.step(noise_pred, t, latents).prev_sample
return latents
def forward(self,
latents,
prompt,
duration=torch.tensor([10]),
sft=True
):
def forward(self, latents, prompt, duration=torch.tensor([10]), sft=True):
device = latents.device
audio_seq_length = self.audio_seq_len
bsz = latents.shape[0]
encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt)
duration_hidden_states = self.encode_duration(duration)
mask_expanded = boolean_encoder_mask.unsqueeze(-1).expand_as(encoder_hidden_states)
masked_data = torch.where(mask_expanded, encoder_hidden_states, torch.tensor(float('nan')))
mask_expanded = boolean_encoder_mask.unsqueeze(-1).expand_as(
encoder_hidden_states
)
masked_data = torch.where(
mask_expanded, encoder_hidden_states, torch.tensor(float("nan"))
)
pooled = torch.nanmean(masked_data, dim=1)
pooled_projection = self.fc(pooled)
## Add duration hidden states to encoder hidden states
encoder_hidden_states = torch.cat([encoder_hidden_states,duration_hidden_states],dim=1) ## (bs,seq_len,dim)
txt_ids = torch.zeros(bsz,encoder_hidden_states.shape[1],3).to(device)
audio_ids = torch.arange(audio_seq_length).unsqueeze(0).unsqueeze(-1).repeat(bsz,1,3).to(device)
encoder_hidden_states = torch.cat(
[encoder_hidden_states, duration_hidden_states], dim=1
) ## (bs,seq_len,dim)
txt_ids = torch.zeros(bsz, encoder_hidden_states.shape[1], 3).to(device)
audio_ids = (
torch.arange(audio_seq_length)
.unsqueeze(0)
.unsqueeze(-1)
.repeat(bsz, 1, 3)
.to(device)
)
if sft:
......@@ -390,27 +417,24 @@ class TangoFlux(nn.Module):
if len(mask_indices) > 0:
encoder_hidden_states[mask_indices] = 0
noise = torch.randn_like(latents)
u = compute_density_for_timestep_sampling(
weighting_scheme='logit_normal',
weighting_scheme="logit_normal",
batch_size=bsz,
logit_mean=0,
logit_std=1,
mode_scale=None,
)
indices = (u * self.noise_scheduler_copy.config.num_train_timesteps).long()
timesteps = self.noise_scheduler_copy.timesteps[indices].to(device=latents.device)
timesteps = self.noise_scheduler_copy.timesteps[indices].to(
device=latents.device
)
sigmas = self.get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
model_pred = self.transformer(
hidden_states=noisy_model_input,
encoder_hidden_states=encoder_hidden_states,
......@@ -419,34 +443,42 @@ class TangoFlux(nn.Module):
txt_ids=txt_ids,
guidance=None,
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
timestep=timesteps/1000,
return_dict=False)[0]
timestep=timesteps / 1000,
return_dict=False,
)[0]
target = noise - latents
loss = torch.mean(
( (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
((model_pred.float() - target.float()) ** 2).reshape(
target.shape[0], -1
),
1,
)
loss = loss.mean()
raw_model_loss, raw_ref_loss,implicit_acc = 0,0,0 ## default this to 0 if doing sft
raw_model_loss, raw_ref_loss, implicit_acc = (
0,
0,
0,
) ## default this to 0 if doing sft
else:
encoder_hidden_states = encoder_hidden_states.repeat(2, 1, 1)
pooled_projection = pooled_projection.repeat(2,1)
noise = torch.randn_like(latents).chunk(2)[0].repeat(2, 1, 1) ## Have to sample same noise for preferred and rejected
pooled_projection = pooled_projection.repeat(2, 1)
noise = (
torch.randn_like(latents).chunk(2)[0].repeat(2, 1, 1)
) ## Have to sample same noise for preferred and rejected
u = compute_density_for_timestep_sampling(
weighting_scheme='logit_normal',
batch_size=bsz//2,
weighting_scheme="logit_normal",
batch_size=bsz // 2,
logit_mean=0,
logit_std=1,
mode_scale=None,
)
indices = (u * self.noise_scheduler_copy.config.num_train_timesteps).long()
timesteps = self.noise_scheduler_copy.timesteps[indices].to(device=latents.device)
timesteps = self.noise_scheduler_copy.timesteps[indices].to(
device=latents.device
)
timesteps = timesteps.repeat(2)
sigmas = self.get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
......@@ -460,17 +492,21 @@ class TangoFlux(nn.Module):
txt_ids=txt_ids,
guidance=None,
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
timestep=timesteps/1000,
return_dict=False)[0]
timestep=timesteps / 1000,
return_dict=False,
)[0]
target = noise - latents
model_losses = F.mse_loss(model_pred.float(), target.float(), reduction="none")
model_losses = model_losses.mean(dim=list(range(1, len(model_losses.shape))))
model_losses = F.mse_loss(
model_pred.float(), target.float(), reduction="none"
)
model_losses = model_losses.mean(
dim=list(range(1, len(model_losses.shape)))
)
model_losses_w, model_losses_l = model_losses.chunk(2)
model_diff = model_losses_w - model_losses_l
raw_model_loss = 0.5 * (model_losses_w.mean() + model_losses_l.mean())
with torch.no_grad():
ref_preds = self.ref_transformer(
hidden_states=noisy_model_input,
......@@ -479,32 +515,25 @@ class TangoFlux(nn.Module):
img_ids=audio_ids,
txt_ids=txt_ids,
guidance=None,
timestep=timesteps/1000,
return_dict=False)[0]
timestep=timesteps / 1000,
return_dict=False,
)[0]
ref_loss = F.mse_loss(ref_preds.float(), target.float(), reduction="none")
ref_loss = F.mse_loss(
ref_preds.float(), target.float(), reduction="none"
)
ref_loss = ref_loss.mean(dim=list(range(1, len(ref_loss.shape))))
ref_losses_w, ref_losses_l = ref_loss.chunk(2)
ref_diff = ref_losses_w - ref_losses_l
raw_ref_loss = ref_loss.mean()
scale_term = -0.5 * self.beta_dpo
inside_term = scale_term * (model_diff - ref_diff)
implicit_acc = (scale_term * (model_diff - ref_diff) > 0).sum().float() / inside_term.size(0)
implicit_acc = (
scale_term * (model_diff - ref_diff) > 0
).sum().float() / inside_term.size(0)
loss = -1 * F.logsigmoid(inside_term).mean() + model_losses_w.mean()
## raw_model_loss, raw_ref_loss, implicit_acc is used to help to analyze dpo behaviour.
return loss, raw_model_loss, raw_ref_loss, implicit_acc
\ No newline at end of file
src/train.py → tangoflux/train.py
View file @ 0db67215
......@@ -22,135 +22,169 @@ from tqdm.auto import tqdm
from transformers import SchedulerType, get_scheduler
from model import TangoFlux
from datasets import load_dataset, Audio
from utils import Text2AudioDataset,read_wav_file,pad_wav
from utils import Text2AudioDataset, read_wav_file, pad_wav
from diffusers import AutoencoderOobleck
import torchaudio
logger = get_logger(__name__)
logger = get_logger(__name__)
def parse_args():
parser = argparse.ArgumentParser(description="Rectified flow for text to audio generation task.")
parser = argparse.ArgumentParser(
description="Rectified flow for text to audio generation task."
)
parser.add_argument(
"--num_examples", type=int, default=-1,
"--num_examples",
type=int,
default=-1,
help="How many examples to use for training and validation.",
)
parser.add_argument(
"--text_column", type=str, default="captions",
"--text_column",
type=str,
default="captions",
help="The name of the column in the datasets containing the input texts.",
)
parser.add_argument(
"--audio_column", type=str, default="location",
"--audio_column",
type=str,
default="location",
help="The name of the column in the datasets containing the audio paths.",
)
parser.add_argument(
"--adam_beta1", type=float, default=0.9,
help="The beta1 parameter for the Adam optimizer."
"--adam_beta1",
type=float,
default=0.9,
help="The beta1 parameter for the Adam optimizer.",
)
parser.add_argument(
"--adam_beta2", type=float, default=0.95,
help="The beta2 parameter for the Adam optimizer."
"--adam_beta2",
type=float,
default=0.95,
help="The beta2 parameter for the Adam optimizer.",
)
parser.add_argument(
"--config", type=str, default='tangoflux_config.yaml',
"--config",
type=str,
default="tangoflux_config.yaml",
help="Config file defining the model size as well as other hyper parameter.",
)
parser.add_argument(
"--prefix", type=str, default='',
"--prefix",
type=str,
default="",
help="Add prefix in text prompts.",
)
parser.add_argument(
"--learning_rate", type=float, default=3e-5,
"--learning_rate",
type=float,
default=3e-5,
help="Initial learning rate (after the potential warmup period) to use.",
)
parser.add_argument(
"--weight_decay", type=float, default=1e-8,
help="Weight decay to use."
"--weight_decay", type=float, default=1e-8, help="Weight decay to use."
)
parser.add_argument(
"--max_train_steps", type=int, default=None,
"--max_train_steps",
type=int,
default=None,
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
)
parser.add_argument(
"--lr_scheduler_type", type=SchedulerType, default="linear",
"--lr_scheduler_type",
type=SchedulerType,
default="linear",
help="The scheduler type to use.",
choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
choices=[
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
],
)
parser.add_argument(
"--num_warmup_steps", type=int, default=0,
help="Number of steps for the warmup in the lr scheduler."
"--num_warmup_steps",
type=int,
default=0,
help="Number of steps for the warmup in the lr scheduler.",
)
parser.add_argument(
"--adam_epsilon", type=float, default=1e-08,
help="Epsilon value for the Adam optimizer"
"--adam_epsilon",
type=float,
default=1e-08,
help="Epsilon value for the Adam optimizer",
)
parser.add_argument(
"--adam_weight_decay", type=float, default=1e-2,
help="Epsilon value for the Adam optimizer"
"--adam_weight_decay",
type=float,
default=1e-2,
help="Epsilon value for the Adam optimizer",
)
parser.add_argument(
"--seed", type=int, default=None,
help="A seed for reproducible training."
"--seed", type=int, default=None, help="A seed for reproducible training."
)
parser.add_argument(
"--checkpointing_steps", type=str, default="best",
"--checkpointing_steps",
type=str,
default="best",
help="Whether the various states should be saved at the end of every 'epoch' or 'best' whenever validation loss decreases.",
)
parser.add_argument(
"--save_every", type=int, default=5,
help="Save model after every how many epochs when checkpointing_steps is set to best."
"--save_every",
type=int,
default=5,
help="Save model after every how many epochs when checkpointing_steps is set to best.",
)
parser.add_argument(
"--resume_from_checkpoint", type=str, default=None,
"--resume_from_checkpoint",
type=str,
default=None,
help="If the training should continue from a local checkpoint folder.",
)
parser.add_argument(
"--load_from_checkpoint", type=str, default=None,
"--load_from_checkpoint",
type=str,
default=None,
help="Whether to continue training from a model weight",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
accelerator_log_kwargs = {}
def load_config(config_path):
with open(config_path, 'r') as file:
with open(config_path, "r") as file:
return yaml.safe_load(file)
config = load_config(args.config)
learning_rate = float(config["training"]["learning_rate"])
num_train_epochs = int(config["training"]["num_train_epochs"])
num_warmup_steps = int(config["training"]["num_warmup_steps"])
per_device_batch_size = int(config["training"]["per_device_batch_size"])
gradient_accumulation_steps = int(config["training"]["gradient_accumulation_steps"])
output_dir = config["paths"]["output_dir"]
learning_rate = float(config['training']['learning_rate'])
num_train_epochs = int(config['training']['num_train_epochs'])
num_warmup_steps = int(config['training']['num_warmup_steps'])
per_device_batch_size = int(config['training']['per_device_batch_size'])
gradient_accumulation_steps = int(config['training']['gradient_accumulation_steps'])
output_dir = config['paths']['output_dir']
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, **accelerator_log_kwargs)
accelerator = Accelerator(
gradient_accumulation_steps=gradient_accumulation_steps,
**accelerator_log_kwargs,
)
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
......@@ -187,29 +221,34 @@ def main():
accelerator.project_configuration.automatic_checkpoint_naming = False
wandb.init(project="Text to Audio Flow matching",settings=wandb.Settings(_disable_stats=True))
wandb.init(
project="Text to Audio Flow matching",
settings=wandb.Settings(_disable_stats=True),
)
accelerator.wait_for_everyone()
# Get the datasets
data_files = {}
#if args.train_file is not None:
if config['paths']['train_file'] != '':
data_files["train"] = config['paths']['train_file']
# if args.train_file is not None:
if config["paths"]["train_file"] != "":
data_files["train"] = config["paths"]["train_file"]
# if args.validation_file is not None:
if config['paths']['val_file'] != '':
data_files["validation"] = config['paths']['val_file']
if config['paths']['test_file'] != '':
data_files["test"] = config['paths']['test_file']
if config["paths"]["val_file"] != "":
data_files["validation"] = config["paths"]["val_file"]
if config["paths"]["test_file"] != "":
data_files["test"] = config["paths"]["test_file"]
else:
data_files["test"] = config['paths']['val_file']
data_files["test"] = config["paths"]["val_file"]
extension = 'json'
extension = "json"
raw_datasets = load_dataset(extension, data_files=data_files)
text_column, audio_column = args.text_column, args.audio_column
model = TangoFlux(config=config['model'])
vae = AutoencoderOobleck.from_pretrained("stabilityai/stable-audio-open-1.0",subfolder='vae')
model = TangoFlux(config=config["model"])
vae = AutoencoderOobleck.from_pretrained(
"stabilityai/stable-audio-open-1.0", subfolder="vae"
)
## Freeze vae
for param in vae.parameters():
......@@ -223,38 +262,79 @@ def main():
prefix = args.prefix
with accelerator.main_process_first():
train_dataset = Text2AudioDataset(raw_datasets["train"], prefix, text_column, audio_column,'duration', args.num_examples)
eval_dataset = Text2AudioDataset(raw_datasets["validation"], prefix, text_column, audio_column,'duration', args.num_examples)
test_dataset = Text2AudioDataset(raw_datasets["test"], prefix, text_column, audio_column,'duration', args.num_examples)
train_dataset = Text2AudioDataset(
raw_datasets["train"],
prefix,
text_column,
audio_column,
"duration",
args.num_examples,
)
eval_dataset = Text2AudioDataset(
raw_datasets["validation"],
prefix,
text_column,
audio_column,
"duration",
args.num_examples,
)
test_dataset = Text2AudioDataset(
raw_datasets["test"],
prefix,
text_column,
audio_column,
"duration",
args.num_examples,
)
accelerator.print("Num instances in train: {}, validation: {}, test: {}".format(train_dataset.get_num_instances(), eval_dataset.get_num_instances(), test_dataset.get_num_instances()))
accelerator.print(
"Num instances in train: {}, validation: {}, test: {}".format(
train_dataset.get_num_instances(),
eval_dataset.get_num_instances(),
test_dataset.get_num_instances(),
)
)
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=config['training']['per_device_batch_size'], collate_fn=train_dataset.collate_fn)
eval_dataloader = DataLoader(eval_dataset, shuffle=True, batch_size=config['training']['per_device_batch_size'], collate_fn=eval_dataset.collate_fn)
test_dataloader = DataLoader(test_dataset, shuffle=False, batch_size=config['training']['per_device_batch_size'], collate_fn=test_dataset.collate_fn)
train_dataloader = DataLoader(
train_dataset,
shuffle=True,
batch_size=config["training"]["per_device_batch_size"],
collate_fn=train_dataset.collate_fn,
)
eval_dataloader = DataLoader(
eval_dataset,
shuffle=True,
batch_size=config["training"]["per_device_batch_size"],
collate_fn=eval_dataset.collate_fn,
)
test_dataloader = DataLoader(
test_dataset,
shuffle=False,
batch_size=config["training"]["per_device_batch_size"],
collate_fn=test_dataset.collate_fn,
)
# Optimizer
optimizer_parameters = list(model.transformer.parameters())+list(model.fc.parameters())
num_trainable_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
optimizer_parameters = list(model.transformer.parameters()) + list(
model.fc.parameters()
)
num_trainable_parameters = sum(
p.numel() for p in model.parameters() if p.requires_grad
)
accelerator.print("Num trainable parameters: {}".format(num_trainable_parameters))
if args.load_from_checkpoint:
from safetensors.torch import load_file
w1 = load_file(args.load_from_checkpoint)
model.load_state_dict(w1,strict=False)
model.load_state_dict(w1, strict=False)
logger.info("Weights loaded from{}".format(args.load_from_checkpoint))
optimizer = torch.optim.AdamW(
optimizer_parameters, lr=learning_rate,
optimizer_parameters,
lr=learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
......@@ -262,17 +342,20 @@ def main():
# Scheduler and math around the number of training steps.
overrode_max_train_steps = False
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
num_update_steps_per_epoch = math.ceil(
len(train_dataloader) / gradient_accumulation_steps
)
if args.max_train_steps is None:
args.max_train_steps = num_train_epochs * num_update_steps_per_epoch
overrode_max_train_steps = True
lr_scheduler = get_scheduler(
name=args.lr_scheduler_type,
optimizer=optimizer,
num_warmup_steps=num_warmup_steps * gradient_accumulation_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * gradient_accumulation_steps
num_warmup_steps=num_warmup_steps
* gradient_accumulation_steps
* accelerator.num_processes,
num_training_steps=args.max_train_steps * gradient_accumulation_steps,
)
# Prepare everything with our `accelerator`.
......@@ -284,9 +367,10 @@ def main():
train_dataloader, eval_dataloader, test_dataloader
)
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
num_update_steps_per_epoch = math.ceil(
len(train_dataloader) / gradient_accumulation_steps
)
if overrode_max_train_steps:
args.max_train_steps = num_train_epochs * num_update_steps_per_epoch
# Afterwards we recalculate our number of training epochs
......@@ -300,35 +384,37 @@ def main():
# We need to initialize the trackers we use, and also store our configuration.
# The trackers initializes automatically on the main process.
# Train!
total_batch_size = per_device_batch_size * accelerator.num_processes * gradient_accumulation_steps
total_batch_size = (
per_device_batch_size * accelerator.num_processes * gradient_accumulation_steps
)
logger.info("***** Running training *****")
logger.info(f" Num examples = {len(train_dataset)}")
logger.info(f" Num Epochs = {num_train_epochs}")
logger.info(f" Instantaneous batch size per device = {per_device_batch_size}")
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
logger.info(
f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
)
logger.info(f" Gradient Accumulation steps = {gradient_accumulation_steps}")
logger.info(f" Total optimization steps = {args.max_train_steps}")
# Only show the progress bar once on each machine.
progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
progress_bar = tqdm(
range(args.max_train_steps), disable=not accelerator.is_local_main_process
)
completed_steps = 0
starting_epoch = 0
# Potentially load in the weights and states from a previous save
resume_from_checkpoint = config['paths']['resume_from_checkpoint']
if resume_from_checkpoint!= '':
resume_from_checkpoint = config["paths"]["resume_from_checkpoint"]
if resume_from_checkpoint != "":
accelerator.load_state(resume_from_checkpoint)
accelerator.print(f"Resumed from local checkpoint: {resume_from_checkpoint}")
# Duration of the audio clips in seconds
best_loss = np.inf
length = config['training']['max_audio_duration']
length = config["training"]["max_audio_duration"]
for epoch in range(starting_epoch, num_train_epochs):
model.train()
......@@ -345,27 +431,32 @@ def main():
for audio_path in audios:
wav = read_wav_file(audio_path,length) ## Only read the first 30 seconds of audio
if wav.shape[0] == 1 : ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2,1)
wav = read_wav_file(
audio_path, length
) ## Only read the first 30 seconds of audio
if (
wav.shape[0] == 1
): ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2, 1)
audio_list.append(wav)
audio_input = torch.stack(audio_list,dim=0)
audio_input = torch.stack(audio_list, dim=0)
audio_input = audio_input.to(device)
unwrapped_vae = accelerator.unwrap_model(vae)
duration = torch.tensor(duration,device=device)
duration = torch.clamp(duration, max=length) ## clamp duration to max audio length
duration = torch.tensor(duration, device=device)
duration = torch.clamp(
duration, max=length
) ## clamp duration to max audio length
audio_latent = unwrapped_vae.encode(audio_input).latent_dist.sample()
audio_latent = audio_latent.transpose(1,2) ## Tranpose to (bsz, seq_len, channel)
audio_latent = unwrapped_vae.encode(
audio_input
).latent_dist.sample()
audio_latent = audio_latent.transpose(
1, 2
) ## Tranpose to (bsz, seq_len, channel)
loss, _, _,_ = model(audio_latent, text ,duration=duration)
loss, _, _, _ = model(audio_latent, text, duration=duration)
total_loss += loss.detach().float()
accelerator.backward(loss)
......@@ -376,10 +467,6 @@ def main():
optimizer.step()
lr_scheduler.step()
if completed_steps % 10 == 0 and accelerator.is_main_process:
total_norm = 0.0
......@@ -388,20 +475,21 @@ def main():
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
logger.info(f"Step {completed_steps}, Loss: {loss.item()}, Grad Norm: {total_norm}")
total_norm = total_norm**0.5
logger.info(
f"Step {completed_steps}, Loss: {loss.item()}, Grad Norm: {total_norm}"
)
lr = lr_scheduler.get_last_lr()[0]
result = {
"train_loss": loss.item(),
"grad_norm": total_norm,
"learning_rate": lr
"learning_rate": lr,
}
# result["val_loss"] = round(total_val_loss.item()/len(eval_dataloader), 4)
wandb.log(result, step=completed_steps)
# Checks if the accelerator has performed an optimization step behind the scenes
if isinstance(checkpointing_steps, int):
......@@ -415,7 +503,9 @@ def main():
break
model.eval()
eval_progress_bar = tqdm(range(len(eval_dataloader)), disable=not accelerator.is_local_main_process)
eval_progress_bar = tqdm(
range(len(eval_dataloader)), disable=not accelerator.is_local_main_process
)
for step, batch in enumerate(eval_dataloader):
with accelerator.accumulate(model) and torch.no_grad():
device = model.device
......@@ -424,47 +514,49 @@ def main():
audio_list = []
for audio_path in audios:
wav = read_wav_file(audio_path,length) ## make sure none of audio exceed 30 sec
if wav.shape[0] == 1 : ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2,1)
wav = read_wav_file(
audio_path, length
) ## make sure none of audio exceed 30 sec
if (
wav.shape[0] == 1
): ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2, 1)
audio_list.append(wav)
audio_input = torch.stack(audio_list,dim=0)
audio_input = torch.stack(audio_list, dim=0)
audio_input = audio_input.to(device)
duration = torch.tensor(duration,device=device)
duration = torch.tensor(duration, device=device)
unwrapped_vae = accelerator.unwrap_model(vae)
audio_latent = unwrapped_vae.encode(audio_input).latent_dist.sample()
audio_latent = audio_latent.transpose(1,2) ## Tranpose to (bsz, seq_len, channel)
audio_latent = audio_latent.transpose(
1, 2
) ## Tranpose to (bsz, seq_len, channel)
val_loss,_, _,_ = model(audio_latent, text , duration=duration)
val_loss, _, _, _ = model(audio_latent, text, duration=duration)
total_val_loss += val_loss.detach().float()
eval_progress_bar.update(1)
if accelerator.is_main_process:
result = {}
result["epoch"] = float(epoch+1)
result["epoch"] = float(epoch + 1)
result["epoch/train_loss"] = round(total_loss.item()/len(train_dataloader), 4)
result["epoch/val_loss"] = round(total_val_loss.item()/len(eval_dataloader), 4)
result["epoch/train_loss"] = round(
total_loss.item() / len(train_dataloader), 4
)
result["epoch/val_loss"] = round(
total_val_loss.item() / len(eval_dataloader), 4
)
wandb.log(result,step=completed_steps)
wandb.log(result, step=completed_steps)
result_string = "Epoch: {}, Loss Train: {}, Val: {}\n".format(epoch, result["epoch/train_loss"],result["epoch/val_loss"])
result_string = "Epoch: {}, Loss Train: {}, Val: {}\n".format(
epoch, result["epoch/train_loss"], result["epoch/val_loss"]
)
accelerator.print(result_string)
with open("{}/summary.jsonl".format(output_dir), "a") as f:
f.write(json.dumps(result) + "\n\n")
......@@ -482,10 +574,14 @@ def main():
accelerator.save_state("{}/{}".format(output_dir, "best"))
if (epoch + 1) % args.save_every == 0:
accelerator.save_state("{}/{}".format(output_dir, "epoch_" + str(epoch+1)))
accelerator.save_state(
"{}/{}".format(output_dir, "epoch_" + str(epoch + 1))
)
if accelerator.is_main_process and args.checkpointing_steps == "epoch":
accelerator.save_state("{}/{}".format(output_dir, "epoch_" + str(epoch+1)))
accelerator.save_state(
"{}/{}".format(output_dir, "epoch_" + str(epoch + 1))
)
if __name__ == "__main__":
......
src/train_dpo.py → tangoflux/train_dpo.py
View file @ 0db67215
......@@ -5,6 +5,7 @@ import logging
import math
import os
import yaml
# from tqdm import tqdm
import copy
from pathlib import Path
......@@ -22,9 +23,9 @@ from datasets import load_dataset
from torch.utils.data import Dataset, DataLoader
from tqdm.auto import tqdm
from transformers import SchedulerType, get_scheduler
from src.model import TangoFlux
from tangoflux.model import TangoFlux
from datasets import load_dataset, Audio
from src.utils import Text2AudioDataset,read_wav_file,DPOText2AudioDataset
from tangoflux.utils import Text2AudioDataset, read_wav_file, DPOText2AudioDataset
from diffusers import AutoencoderOobleck
import torchaudio
......@@ -32,84 +33,119 @@ import torchaudio
logger = get_logger(__name__)
def parse_args():
parser = argparse.ArgumentParser(description="Rectified flow for text to audio generation task.")
parser = argparse.ArgumentParser(
description="Rectified flow for text to audio generation task."
)
parser.add_argument(
"--num_examples", type=int, default=-1,
"--num_examples",
type=int,
default=-1,
help="How many examples to use for training and validation.",
)
parser.add_argument(
"--text_column", type=str, default="captions",
"--text_column",
type=str,
default="captions",
help="The name of the column in the datasets containing the input texts.",
)
parser.add_argument(
"--audio_column", type=str, default="location",
"--audio_column",
type=str,
default="location",
help="The name of the column in the datasets containing the audio paths.",
)
parser.add_argument(
"--adam_beta1", type=float, default=0.9,
help="The beta1 parameter for the Adam optimizer."
"--adam_beta1",
type=float,
default=0.9,
help="The beta1 parameter for the Adam optimizer.",
)
parser.add_argument(
"--adam_beta2", type=float, default=0.95,
help="The beta2 parameter for the Adam optimizer."
"--adam_beta2",
type=float,
default=0.95,
help="The beta2 parameter for the Adam optimizer.",
)
parser.add_argument(
"--config", type=str, default='tangoflux_config.yaml',
"--config",
type=str,
default="tangoflux_config.yaml",
help="Config file defining the model size.",
)
parser.add_argument(
"--weight_decay", type=float, default=1e-8,
help="Weight decay to use."
"--weight_decay", type=float, default=1e-8, help="Weight decay to use."
)
parser.add_argument(
"--max_train_steps", type=int, default=None,
"--max_train_steps",
type=int,
default=None,
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
)
parser.add_argument(
"--lr_scheduler_type", type=SchedulerType, default="linear",
"--lr_scheduler_type",
type=SchedulerType,
default="linear",
help="The scheduler type to use.",
choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
choices=[
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
],
)
parser.add_argument(
"--num_warmup_steps", type=int, default=0,
help="Number of steps for the warmup in the lr scheduler."
"--num_warmup_steps",
type=int,
default=0,
help="Number of steps for the warmup in the lr scheduler.",
)
parser.add_argument(
"--adam_epsilon", type=float, default=1e-08,
help="Epsilon value for the Adam optimizer"
"--adam_epsilon",
type=float,
default=1e-08,
help="Epsilon value for the Adam optimizer",
)
parser.add_argument(
"--adam_weight_decay", type=float, default=1e-2,
help="Epsilon value for the Adam optimizer"
"--adam_weight_decay",
type=float,
default=1e-2,
help="Epsilon value for the Adam optimizer",
)
parser.add_argument(
"--seed", type=int, default=None,
help="A seed for reproducible training."
"--seed", type=int, default=None, help="A seed for reproducible training."
)
parser.add_argument(
"--checkpointing_steps", type=str, default="best",
"--checkpointing_steps",
type=str,
default="best",
help="Whether the various states should be saved at the end of every 'epoch' or 'best' whenever validation loss decreases.",
)
parser.add_argument(
"--save_every", type=int, default=5,
help="Save model after every how many epochs when checkpointing_steps is set to best."
"--save_every",
type=int,
default=5,
help="Save model after every how many epochs when checkpointing_steps is set to best.",
)
parser.add_argument(
"--resume_from_checkpoint", type=str, default=None,
"--resume_from_checkpoint",
type=str,
default=None,
help="If the training should continue from a local checkpoint folder.",
)
parser.add_argument(
"--report_to", type=str, default="all",
"--report_to",
type=str,
default="all",
help=(
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
' `"wandb"`, `"comet_ml"` and `"clearml"`. Use `"all"` (default) to report to all integrations.'
......@@ -117,59 +153,57 @@ def parse_args():
),
)
parser.add_argument(
"--load_from_checkpoint", type=str, default=None,
"--load_from_checkpoint",
type=str,
default=None,
help="Whether to continue training from a model weight",
)
parser.add_argument(
"--audio_length", type=float, default=30,
"--audio_length",
type=float,
default=30,
help="Audio duration",
)
args = parser.parse_args()
# Sanity checks
#if args.train_file is None and args.validation_file is None:
# if args.train_file is None and args.validation_file is None:
# raise ValueError("Need a training/validation file.")
#else:
# else:
# if args.train_file is not None:
# extension = args.train_file.split(".")[-1]
# assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
#if args.validation_file is not None:
# if args.validation_file is not None:
# extension = args.validation_file.split(".")[-1]
# assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
return args
def main():
args = parse_args()
accelerator_log_kwargs = {}
def load_config(config_path):
with open(config_path, 'r') as file:
with open(config_path, "r") as file:
return yaml.safe_load(file)
config = load_config(args.config)
learning_rate = float(config["training"]["learning_rate"])
num_train_epochs = int(config["training"]["num_train_epochs"])
num_warmup_steps = int(config["training"]["num_warmup_steps"])
per_device_batch_size = int(config["training"]["per_device_batch_size"])
gradient_accumulation_steps = int(config["training"]["gradient_accumulation_steps"])
output_dir = config["paths"]["output_dir"]
learning_rate = float(config['training']['learning_rate'])
num_train_epochs = int(config['training']['num_train_epochs'])
num_warmup_steps = int(config['training']['num_warmup_steps'])
per_device_batch_size = int(config['training']['per_device_batch_size'])
gradient_accumulation_steps = int(config['training']['gradient_accumulation_steps'])
output_dir = config['paths']['output_dir']
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, **accelerator_log_kwargs)
accelerator = Accelerator(
gradient_accumulation_steps=gradient_accumulation_steps,
**accelerator_log_kwargs,
)
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
......@@ -206,31 +240,36 @@ def main():
accelerator.project_configuration.automatic_checkpoint_naming = False
wandb.init(project="Text to Audio Flow matching DPO",settings=wandb.Settings(_disable_stats=True))
wandb.init(
project="Text to Audio Flow matching DPO",
settings=wandb.Settings(_disable_stats=True),
)
accelerator.wait_for_everyone()
# Get the datasets
data_files = {}
#if args.train_file is not None:
if config['paths']['train_file'] != '':
data_files["train"] = config['paths']['train_file']
# if args.train_file is not None:
if config["paths"]["train_file"] != "":
data_files["train"] = config["paths"]["train_file"]
# if args.validation_file is not None:
if config['paths']['val_file'] != '':
data_files["validation"] = config['paths']['val_file']
if config['paths']['test_file'] != '':
data_files["test"] = config['paths']['test_file']
if config["paths"]["val_file"] != "":
data_files["validation"] = config["paths"]["val_file"]
if config["paths"]["test_file"] != "":
data_files["test"] = config["paths"]["test_file"]
else:
data_files["test"] = config['paths']['val_file']
data_files["test"] = config["paths"]["val_file"]
extension = 'json'
train_dataset = load_dataset(extension,data_files=data_files['train'])
data_files.pop('train')
extension = "json"
train_dataset = load_dataset(extension, data_files=data_files["train"])
data_files.pop("train")
raw_datasets = load_dataset(extension, data_files=data_files)
text_column, audio_column = args.text_column, args.audio_column
model = TangoFlux(config=config['model'],initialize_reference_model=True)
vae = AutoencoderOobleck.from_pretrained("stabilityai/stable-audio-open-1.0",subfolder='vae')
model = TangoFlux(config=config["model"], initialize_reference_model=True)
vae = AutoencoderOobleck.from_pretrained(
"stabilityai/stable-audio-open-1.0", subfolder="vae"
)
## Freeze vae
for param in vae.parameters():
......@@ -244,35 +283,79 @@ def main():
prefix = ""
with accelerator.main_process_first():
train_dataset = DPOText2AudioDataset(train_dataset["train"], prefix, text_column, 'chosen','reject','duration', args.num_examples)
eval_dataset = Text2AudioDataset(raw_datasets["validation"], prefix, text_column, audio_column,'duration', args.num_examples)
test_dataset = Text2AudioDataset(raw_datasets["test"], prefix, text_column, audio_column,'duration', args.num_examples)
train_dataset = DPOText2AudioDataset(
train_dataset["train"],
prefix,
text_column,
"chosen",
"reject",
"duration",
args.num_examples,
)
eval_dataset = Text2AudioDataset(
raw_datasets["validation"],
prefix,
text_column,
audio_column,
"duration",
args.num_examples,
)
test_dataset = Text2AudioDataset(
raw_datasets["test"],
prefix,
text_column,
audio_column,
"duration",
args.num_examples,
)
accelerator.print("Num instances in train: {}, validation: {}, test: {}".format(train_dataset.get_num_instances(), eval_dataset.get_num_instances(), test_dataset.get_num_instances()))
accelerator.print(
"Num instances in train: {}, validation: {}, test: {}".format(
train_dataset.get_num_instances(),
eval_dataset.get_num_instances(),
test_dataset.get_num_instances(),
)
)
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=config['training']['per_device_batch_size'], collate_fn=train_dataset.collate_fn)
eval_dataloader = DataLoader(eval_dataset, shuffle=True, batch_size=config['training']['per_device_batch_size'], collate_fn=eval_dataset.collate_fn)
test_dataloader = DataLoader(test_dataset, shuffle=False, batch_size=config['training']['per_device_batch_size'], collate_fn=test_dataset.collate_fn)
train_dataloader = DataLoader(
train_dataset,
shuffle=True,
batch_size=config["training"]["per_device_batch_size"],
collate_fn=train_dataset.collate_fn,
)
eval_dataloader = DataLoader(
eval_dataset,
shuffle=True,
batch_size=config["training"]["per_device_batch_size"],
collate_fn=eval_dataset.collate_fn,
)
test_dataloader = DataLoader(
test_dataset,
shuffle=False,
batch_size=config["training"]["per_device_batch_size"],
collate_fn=test_dataset.collate_fn,
)
# Optimizer
optimizer_parameters = list(model.transformer.parameters())+list(model.fc.parameters())
num_trainable_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
optimizer_parameters = list(model.transformer.parameters()) + list(
model.fc.parameters()
)
num_trainable_parameters = sum(
p.numel() for p in model.parameters() if p.requires_grad
)
accelerator.print("Num trainable parameters: {}".format(num_trainable_parameters))
if args.load_from_checkpoint:
from safetensors.torch import load_file
w1 = load_file(args.load_from_checkpoint)
model.load_state_dict(w1,strict=False)
model.load_state_dict(w1, strict=False)
logger.info("Weights loaded from{}".format(args.load_from_checkpoint))
import copy
model.ref_transformer = copy.deepcopy(model.transformer)
model.ref_transformer.requires_grad_ = False
model.ref_transformer.eval()
......@@ -280,7 +363,9 @@ def main():
param.requires_grad = False
@torch.no_grad()
def initialize_or_update_ref_transformer(model, accelerator: Accelerator,alpha=0.5):
def initialize_or_update_ref_transformer(
model, accelerator: Accelerator, alpha=0.5
):
"""
Initializes or updates ref_transformer as alpha * ref + 1-alpha * transformer.
......@@ -296,32 +381,31 @@ def main():
# Unwrap the model to access the original underlying model
unwrapped_model = accelerator.unwrap_model(model)
with torch.no_grad():
for ref_param, model_param in zip(unwrapped_model.ref_transformer.parameters(),
unwrapped_model.transformer.parameters()):
average_param = alpha * ref_param.data + (1-alpha) * model_param.data
for ref_param, model_param in zip(
unwrapped_model.ref_transformer.parameters(),
unwrapped_model.transformer.parameters(),
):
average_param = alpha * ref_param.data + (1 - alpha) * model_param.data
ref_param.data.copy_(average_param)
unwrapped_model.ref_transformer.eval()
unwrapped_model.ref_transformer.requires_grad_= False
unwrapped_model.ref_transformer.requires_grad_ = False
for param in unwrapped_model.ref_transformer.parameters():
param.requires_grad = False
return model
model.ref_transformer = copy.deepcopy(model.transformer)
model.ref_transformer.requires_grad_ = False
model.ref_transformer.eval()
for param in model.ref_transformer.parameters():
param.requires_grad = False
optimizer = torch.optim.AdamW(
optimizer_parameters, lr=learning_rate,
optimizer_parameters,
lr=learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
......@@ -329,17 +413,20 @@ def main():
# Scheduler and math around the number of training steps.
overrode_max_train_steps = False
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
num_update_steps_per_epoch = math.ceil(
len(train_dataloader) / gradient_accumulation_steps
)
if args.max_train_steps is None:
args.max_train_steps = num_train_epochs * num_update_steps_per_epoch
overrode_max_train_steps = True
lr_scheduler = get_scheduler(
name=args.lr_scheduler_type,
optimizer=optimizer,
num_warmup_steps=num_warmup_steps * gradient_accumulation_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * gradient_accumulation_steps
num_warmup_steps=num_warmup_steps
* gradient_accumulation_steps
* accelerator.num_processes,
num_training_steps=args.max_train_steps * gradient_accumulation_steps,
)
# Prepare everything with our `accelerator`.
......@@ -351,9 +438,10 @@ def main():
train_dataloader, eval_dataloader, test_dataloader
)
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
num_update_steps_per_epoch = math.ceil(
len(train_dataloader) / gradient_accumulation_steps
)
if overrode_max_train_steps:
args.max_train_steps = num_train_epochs * num_update_steps_per_epoch
# Afterwards we recalculate our number of training epochs
......@@ -367,96 +455,108 @@ def main():
# We need to initialize the trackers we use, and also store our configuration.
# The trackers initializes automatically on the main process.
# Train!
total_batch_size = per_device_batch_size * accelerator.num_processes * gradient_accumulation_steps
total_batch_size = (
per_device_batch_size * accelerator.num_processes * gradient_accumulation_steps
)
logger.info("***** Running training *****")
logger.info(f" Num examples = {len(train_dataset)}")
logger.info(f" Num Epochs = {num_train_epochs}")
logger.info(f" Instantaneous batch size per device = {per_device_batch_size}")
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
logger.info(
f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
)
logger.info(f" Gradient Accumulation steps = {gradient_accumulation_steps}")
logger.info(f" Total optimization steps = {args.max_train_steps}")
# Only show the progress bar once on each machine.
progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
progress_bar = tqdm(
range(args.max_train_steps), disable=not accelerator.is_local_main_process
)
completed_steps = 0
starting_epoch = 0
# Potentially load in the weights and states from a previous save
resume_from_checkpoint = config['paths']['resume_from_checkpoint']
if resume_from_checkpoint!= '':
resume_from_checkpoint = config["paths"]["resume_from_checkpoint"]
if resume_from_checkpoint != "":
accelerator.load_state(resume_from_checkpoint)
accelerator.print(f"Resumed from local checkpoint: {resume_from_checkpoint}")
# Duration of the audio clips in seconds
best_loss = np.inf
length = config['training']['max_audio_duration']
length = config["training"]["max_audio_duration"]
for epoch in range(starting_epoch, num_train_epochs):
model.train()
total_loss, total_val_loss = 0, 0
for step, batch in enumerate(train_dataloader):
optimizer.zero_grad()
with accelerator.accumulate(model):
optimizer.zero_grad()
device = accelerator.device
text, audio_w,audio_l, duration, _ = batch
text, audio_w, audio_l, duration, _ = batch
with torch.no_grad():
audio_list_w = []
audio_list_l = []
for audio_path in audio_w:
wav = read_wav_file(audio_path,length) ## Only read the first 30 seconds of audio
if wav.shape[0] == 1 : ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2,1)
wav = read_wav_file(
audio_path, length
) ## Only read the first 30 seconds of audio
if (
wav.shape[0] == 1
): ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2, 1)
audio_list_w.append(wav)
for audio_path in audio_l:
wav = read_wav_file(audio_path,length) ## Only read the first 30 seconds of audio
if wav.shape[0] == 1 : ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2,1)
wav = read_wav_file(
audio_path, length
) ## Only read the first 30 seconds of audio
if (
wav.shape[0] == 1
): ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2, 1)
audio_list_l.append(wav)
audio_input_w = torch.stack(audio_list_w,dim=0).to(device)
audio_input_l = torch.stack(audio_list_l,dim=0).to(device)
#audio_input_ = audio_input.to(device)
audio_input_w = torch.stack(audio_list_w, dim=0).to(device)
audio_input_l = torch.stack(audio_list_l, dim=0).to(device)
# audio_input_ = audio_input.to(device)
unwrapped_vae = accelerator.unwrap_model(vae)
duration = torch.tensor(duration,device=device)
duration = torch.clamp(duration, max=length) ## max duration is 30 sec
audio_latent_w = unwrapped_vae.encode(audio_input_w).latent_dist.sample()
audio_latent_l = unwrapped_vae.encode(audio_input_l).latent_dist.sample()
audio_latent = torch.cat((audio_latent_w,audio_latent_l),dim=0)
audio_latent = audio_latent.transpose(1,2) ## Tranpose to (bsz, seq_len, channel)
loss, raw_model_loss, raw_ref_loss,implicit_acc = model(audio_latent, text ,duration=duration,sft=False)
duration = torch.tensor(duration, device=device)
duration = torch.clamp(
duration, max=length
) ## max duration is 30 sec
audio_latent_w = unwrapped_vae.encode(
audio_input_w
).latent_dist.sample()
audio_latent_l = unwrapped_vae.encode(
audio_input_l
).latent_dist.sample()
audio_latent = torch.cat((audio_latent_w, audio_latent_l), dim=0)
audio_latent = audio_latent.transpose(
1, 2
) ## Tranpose to (bsz, seq_len, channel)
loss, raw_model_loss, raw_ref_loss, implicit_acc = model(
audio_latent, text, duration=duration, sft=False
)
total_loss += loss.detach().float()
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
#if accelerator.sync_gradients:
# if accelerator.sync_gradients:
if accelerator.sync_gradients:
#accelerator.clip_grad_value_(model.parameters(),1.0)
# accelerator.clip_grad_value_(model.parameters(),1.0)
progress_bar.update(1)
completed_steps += 1
if completed_steps % 10 == 0 and accelerator.is_main_process:
total_norm = 0.0
......@@ -465,8 +565,10 @@ def main():
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
logger.info(f"Step {completed_steps}, Loss: {loss.item()}, Grad Norm: {total_norm}")
total_norm = total_norm**0.5
logger.info(
f"Step {completed_steps}, Loss: {loss.item()}, Grad Norm: {total_norm}"
)
lr = lr_scheduler.get_last_lr()[0]
......@@ -474,17 +576,14 @@ def main():
"train_loss": loss.item(),
"grad_norm": total_norm,
"learning_rate": lr,
'raw_model_loss':raw_model_loss,
'raw_ref_loss': raw_ref_loss,
'implicit_acc':implicit_acc
"raw_model_loss": raw_model_loss,
"raw_ref_loss": raw_ref_loss,
"implicit_acc": implicit_acc,
}
# result["val_loss"] = round(total_val_loss.item()/len(eval_dataloader), 4)
wandb.log(result, step=completed_steps)
# Checks if the accelerator has performed an optimization step behind the scenes
if isinstance(checkpointing_steps, int):
......@@ -497,9 +596,10 @@ def main():
if completed_steps >= args.max_train_steps:
break
model.eval()
eval_progress_bar = tqdm(range(len(eval_dataloader)), disable=not accelerator.is_local_main_process)
eval_progress_bar = tqdm(
range(len(eval_dataloader)), disable=not accelerator.is_local_main_process
)
for step, batch in enumerate(eval_dataloader):
with accelerator.accumulate(model) and torch.no_grad():
device = model.device
......@@ -508,60 +608,65 @@ def main():
audio_list = []
for audio_path in audios:
wav = read_wav_file(audio_path,length) ## Only read the first 30 seconds of audio
if wav.shape[0] == 1 : ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2,1)
wav = read_wav_file(
audio_path, length
) ## Only read the first 30 seconds of audio
if (
wav.shape[0] == 1
): ## If this audio is mono, we repeat the channel so it become "fake stereo"
wav = wav.repeat(2, 1)
audio_list.append(wav)
audio_input = torch.stack(audio_list,dim=0)
audio_input = torch.stack(audio_list, dim=0)
audio_input = audio_input.to(device)
duration = torch.tensor(duration,device=device)
duration = torch.tensor(duration, device=device)
unwrapped_vae = accelerator.unwrap_model(vae)
audio_latent = unwrapped_vae.encode(audio_input).latent_dist.sample()
audio_latent = audio_latent.transpose(1,2) ## Tranpose to (bsz, seq_len, channel)
audio_latent = audio_latent.transpose(
1, 2
) ## Tranpose to (bsz, seq_len, channel)
val_loss, _, _, _ = model(audio_latent, text , duration=duration,sft=True)
val_loss, _, _, _ = model(
audio_latent, text, duration=duration, sft=True
)
total_val_loss += val_loss.detach().float()
eval_progress_bar.update(1)
if accelerator.is_main_process:
result = {}
result["epoch"] = float(epoch+1)
result["epoch/train_loss"] = round(total_loss.item()/len(train_dataloader), 4)
result["epoch/val_loss"] = round(total_val_loss.item()/len(eval_dataloader), 4)
wandb.log(result,step=completed_steps)
result["epoch"] = float(epoch + 1)
result["epoch/train_loss"] = round(
total_loss.item() / len(train_dataloader), 4
)
result["epoch/val_loss"] = round(
total_val_loss.item() / len(eval_dataloader), 4
)
wandb.log(result, step=completed_steps)
with open("{}/summary.jsonl".format(output_dir), "a") as f:
f.write(json.dumps(result) + "\n\n")
logger.info(result)
save_checkpoint= True
save_checkpoint = True
accelerator.wait_for_everyone()
if accelerator.is_main_process and args.checkpointing_steps == "best":
if save_checkpoint:
accelerator.save_state("{}/{}".format(output_dir, "best"))
if (epoch + 1) % args.save_every == 0:
accelerator.save_state("{}/{}".format(output_dir, "epoch_" + str(epoch+1)))
accelerator.save_state(
"{}/{}".format(output_dir, "epoch_" + str(epoch + 1))
)
if accelerator.is_main_process and args.checkpointing_steps == "epoch":
accelerator.save_state("{}/{}".format(output_dir, "epoch_" + str(epoch+1)))
accelerator.save_state(
"{}/{}".format(output_dir, "epoch_" + str(epoch + 1))
)
if __name__ == "__main__":
......
src/utils.py → tangoflux/utils.py
View file @ 0db67215
......@@ -31,8 +31,6 @@ def pad_wav(waveform, segment_length):
return waveform
def read_wav_file(filename, duration_sec):
info = torchaudio.info(filename)
sample_rate = info.sample_rate
......@@ -40,29 +38,38 @@ def read_wav_file(filename, duration_sec):
# Calculate the number of frames corresponding to the desired duration
num_frames = int(sample_rate * duration_sec)
waveform, sr = torchaudio.load(filename,num_frames=num_frames) # Faster!!!
waveform, sr = torchaudio.load(filename, num_frames=num_frames) # Faster!!!
if waveform.shape[0] == 2 : ## Stereo audio
if waveform.shape[0] == 2: ## Stereo audio
resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=44100)
resampled_waveform = resampler(waveform)
#print(resampled_waveform.shape)
padded_left = pad_wav(resampled_waveform[0], int(44100*duration_sec)) ## We pad left and right seperately
padded_right = pad_wav(resampled_waveform[1], int(44100*duration_sec))
# print(resampled_waveform.shape)
padded_left = pad_wav(
resampled_waveform[0], int(44100 * duration_sec)
) ## We pad left and right seperately
padded_right = pad_wav(resampled_waveform[1], int(44100 * duration_sec))
return torch.stack([padded_left,padded_right])
return torch.stack([padded_left, padded_right])
else:
waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=44100)[0]
waveform = pad_wav(waveform, int(44100*duration_sec)).unsqueeze(0)
waveform = torchaudio.functional.resample(
waveform, orig_freq=sr, new_freq=44100
)[0]
waveform = pad_wav(waveform, int(44100 * duration_sec)).unsqueeze(0)
return waveform
class DPOText2AudioDataset(Dataset):
def __init__(self, dataset, prefix, text_column, audio_w_column, audio_l_column, duration, num_examples=-1):
def __init__(
self,
dataset,
prefix,
text_column,
audio_w_column,
audio_l_column,
duration,
num_examples=-1,
):
inputs = list(dataset[text_column])
self.inputs = [prefix + inp for inp in inputs]
......@@ -72,11 +79,18 @@ class DPOText2AudioDataset(Dataset):
self.indices = list(range(len(self.inputs)))
self.mapper = {}
for index, audio_w, audio_l, duration, text in zip(self.indices, self.audios_w,self.audios_l,self.durations,inputs):
for index, audio_w, audio_l, duration, text in zip(
self.indices, self.audios_w, self.audios_l, self.durations, inputs
):
self.mapper[index] = [audio_w, audio_l, duration, text]
if num_examples != -1:
self.inputs, self.audios_w, self.audios_l, self.durations = self.inputs[:num_examples], self.audios_w[:num_examples], self.audios_l[:num_examples], self.durations[:num_examples]
self.inputs, self.audios_w, self.audios_l, self.durations = (
self.inputs[:num_examples],
self.audios_w[:num_examples],
self.audios_l[:num_examples],
self.durations[:num_examples],
)
self.indices = self.indices[:num_examples]
def __len__(self):
......@@ -86,15 +100,24 @@ class DPOText2AudioDataset(Dataset):
return len(self.inputs)
def __getitem__(self, index):
s1, s2, s3, s4, s5 = self.inputs[index], self.audios_w[index], self.audios_l[index], self.durations[index], self.indices[index]
s1, s2, s3, s4, s5 = (
self.inputs[index],
self.audios_w[index],
self.audios_l[index],
self.durations[index],
self.indices[index],
)
return s1, s2, s3, s4, s5
def collate_fn(self, data):
dat = pd.DataFrame(data)
return [dat[i].tolist() for i in dat]
class Text2AudioDataset(Dataset):
def __init__(self, dataset, prefix, text_column, audio_column, duration, num_examples=-1):
def __init__(
self, dataset, prefix, text_column, audio_column, duration, num_examples=-1
):
inputs = list(dataset[text_column])
self.inputs = [prefix + inp for inp in inputs]
......@@ -103,11 +126,17 @@ class Text2AudioDataset(Dataset):
self.indices = list(range(len(self.inputs)))
self.mapper = {}
for index, audio, duration,text in zip(self.indices, self.audios, self.durations,inputs):
self.mapper[index] = [audio, text,duration]
for index, audio, duration, text in zip(
self.indices, self.audios, self.durations, inputs
):
self.mapper[index] = [audio, text, duration]
if num_examples != -1:
self.inputs, self.audios, self.durations = self.inputs[:num_examples], self.audios[:num_examples], self.durations[:num_examples]
self.inputs, self.audios, self.durations = (
self.inputs[:num_examples],
self.audios[:num_examples],
self.durations[:num_examples],
)
self.indices = self.indices[:num_examples]
def __len__(self):
......@@ -117,7 +146,12 @@ class Text2AudioDataset(Dataset):
return len(self.inputs)
def __getitem__(self, index):
s1, s2, s3, s4 = self.inputs[index], self.audios[index], self.durations[index], self.indices[index]
s1, s2, s3, s4 = (
self.inputs[index],
self.audios[index],
self.durations[index],
self.indices[index],
)
return s1, s2, s3, s4
def collate_fn(self, data):
......
train.sh
View file @ 0db67215
CUDA_VISISBLE_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'
\ No newline at end of file
CUDA_VISISBLE_DEVICES=0,1 accelerate launch --config_file='configs/accelerator_config.yaml' tangoflux/train.py --checkpointing_steps="best" --save_every=5 --config='configs/tangoflux_config.yaml'
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment