# based on https://github.com/huggingface/diffusers/blob/main/examples/train_unconditional.py
import argparse
import os
import torch
import torch.nn.functional as F
from accelerate import Accelerator
from accelerate.logging import get_logger
from datasets import load_from_disk, load_dataset
from diffusers import (DiffusionPipeline, DDPMScheduler, UNet2DModel,
DDIMScheduler, AutoencoderKL)
from diffusers.hub_utils import init_git_repo, push_to_hub
from diffusers.optimization import get_scheduler
from diffusers.training_utils import EMAModel
from torchvision.transforms import (
Compose,
Normalize,
ToTensor,
)
import numpy as np
from tqdm.auto import tqdm
from librosa.util import normalize
from audiodiffusion.mel import Mel
from audiodiffusion import LatentAudioDiffusionPipeline, AudioDiffusionPipeline
logger = get_logger(__name__)
def main(args):
output_dir = os.environ.get("SM_MODEL_DIR", None) or args.output_dir
logging_dir = os.path.join(output_dir, args.logging_dir)
accelerator = Accelerator(
gradient_accumulation_steps=args.gradient_accumulation_steps,
mixed_precision=args.mixed_precision,
log_with="tensorboard",
logging_dir=logging_dir,
)
if args.dataset_name is not None:
if os.path.exists(args.dataset_name):
dataset = load_from_disk(args.dataset_name,
args.dataset_config_name)["train"]
else:
dataset = load_dataset(
args.dataset_name,
args.dataset_config_name,
cache_dir=args.cache_dir,
use_auth_token=True if args.use_auth_token else None,
split="train",
)
else:
dataset = load_dataset(
"imagefolder",
data_dir=args.train_data_dir,
cache_dir=args.cache_dir,
split="train",
)
# Determine image resolution
resolution = dataset[0]['image'].height, dataset[0]['image'].width
augmentations = Compose([
ToTensor(),
Normalize([0.5], [0.5]),
])
def transforms(examples):
if args.vae is not None and vqvae.config['in_channels'] == 3:
images = [
augmentations(image.convert('RGB'))
for image in examples["image"]
]
else:
images = [augmentations(image) for image in examples["image"]]
return {"input": images}
dataset.set_transform(transforms)
train_dataloader = torch.utils.data.DataLoader(
dataset, batch_size=args.train_batch_size, shuffle=True)
vqvae = None
if args.vae is not None:
try:
vqvae = AutoencoderKL.from_pretrained(args.vae)
except EnvironmentError:
vqvae = LatentAudioDiffusionPipeline.from_pretrained(
args.vae).vqvae
# Determine latent resolution
with torch.no_grad():
latent_resolution = vqvae.encode(
torch.zeros((1, 1) +
resolution)).latent_dist.sample().shape[2:]
if args.from_pretrained is not None:
pipeline = {
'LatentAudioDiffusionPipeline': LatentAudioDiffusionPipeline,
'AudioDiffusionPipeline': AudioDiffusionPipeline
}.get(
DiffusionPipeline.get_config_dict(
args.from_pretrained)['_class_name'], AudioDiffusionPipeline)
pipeline = pipeline.from_pretrained(args.from_pretrained)
model = pipeline.unet
if hasattr(pipeline, 'vqvae'):
vqvae = pipeline.vqvae
else:
model = UNet2DModel(
sample_size=resolution if vqvae is None else latent_resolution,
in_channels=1
if vqvae is None else vqvae.config['latent_channels'],
out_channels=1
if vqvae is None else vqvae.config['latent_channels'],
layers_per_block=2,
block_out_channels=(128, 128, 256, 256, 512, 512),
down_block_types=(
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"AttnDownBlock2D",
"DownBlock2D",
),
up_block_types=(
"UpBlock2D",
"AttnUpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
),
)
if args.scheduler == "ddpm":
noise_scheduler = DDPMScheduler(
num_train_timesteps=args.num_train_steps)
else:
noise_scheduler = DDIMScheduler(
num_train_timesteps=args.num_train_steps)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps,
num_training_steps=(len(train_dataloader) * args.num_epochs) //
args.gradient_accumulation_steps,
)
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
model, optimizer, train_dataloader, lr_scheduler)
ema_model = EMAModel(
getattr(model, "module", model),
inv_gamma=args.ema_inv_gamma,
power=args.ema_power,
max_value=args.ema_max_decay,
)
if args.push_to_hub:
repo = init_git_repo(args, at_init=True)
if accelerator.is_main_process:
run = os.path.split(__file__)[-1].split(".")[0]
accelerator.init_trackers(run)
mel = Mel(x_res=resolution[1],
y_res=resolution[0],
hop_length=args.hop_length)
global_step = 0
for epoch in range(args.num_epochs):
progress_bar = tqdm(total=len(train_dataloader),
disable=not accelerator.is_local_main_process)
progress_bar.set_description(f"Epoch {epoch}")
if epoch < args.start_epoch:
for step in range(len(train_dataloader)):
optimizer.step()
lr_scheduler.step()
progress_bar.update(1)
global_step += 1
if epoch == args.start_epoch - 1 and args.use_ema:
ema_model.optimization_step = global_step
continue
model.train()
for step, batch in enumerate(train_dataloader):
clean_images = batch["input"]
if vqvae is not None:
vqvae.to(clean_images.device)
with torch.no_grad():
clean_images = vqvae.encode(
clean_images).latent_dist.sample()
# Scale latent images to ensure approximately unit variance
clean_images = clean_images * 0.18215
# Sample noise that we'll add to the images
noise = torch.randn(clean_images.shape).to(clean_images.device)
bsz = clean_images.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(
0,
noise_scheduler.num_train_timesteps,
(bsz, ),
device=clean_images.device,
).long()
# Add noise to the clean images according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_images = noise_scheduler.add_noise(clean_images, noise,
timesteps)
with accelerator.accumulate(model):
# Predict the noise residual
noise_pred = model(noisy_images, timesteps)["sample"]
loss = F.mse_loss(noise_pred, noise)
accelerator.backward(loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
lr_scheduler.step()
if args.use_ema:
ema_model.step(model)
optimizer.zero_grad()
progress_bar.update(1)
global_step += 1
logs = {
"loss": loss.detach().item(),
"lr": lr_scheduler.get_last_lr()[0],
"step": global_step,
}
if args.use_ema:
logs["ema_decay"] = ema_model.decay
progress_bar.set_postfix(**logs)
accelerator.log(logs, step=global_step)
progress_bar.close()
accelerator.wait_for_everyone()
# Generate sample images for visual inspection
if accelerator.is_main_process:
if (
epoch + 1
) % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
if vqvae is not None:
pipeline = LatentAudioDiffusionPipeline(
unet=accelerator.unwrap_model(
ema_model.averaged_model if args.use_ema else model
),
vqvae=vqvae,
scheduler=noise_scheduler)
else:
pipeline = AudioDiffusionPipeline(
unet=accelerator.unwrap_model(
ema_model.averaged_model if args.use_ema else model
),
scheduler=noise_scheduler,
)
# save the model
if args.push_to_hub:
try:
push_to_hub(
args,
pipeline,
repo,
commit_message=f"Epoch {epoch}",
blocking=False,
)
except NameError: # current version of diffusers has a little bug
pass
else:
pipeline.save_pretrained(output_dir)
if (epoch + 1) % args.save_images_epochs == 0:
generator = torch.manual_seed(42)
# run pipeline in inference (sample random noise and denoise)
images, (sample_rate, audios) = pipeline(
mel=mel,
generator=generator,
batch_size=args.eval_batch_size,
)
# denormalize the images and save to tensorboard
images = np.array([
np.frombuffer(image.tobytes(), dtype="uint8").reshape(
(len(image.getbands()), image.height, image.width))
for image in images
])
accelerator.trackers[0].writer.add_images(
"test_samples", images, epoch)
for _, audio in enumerate(audios):
accelerator.trackers[0].writer.add_audio(
f"test_audio_{_}",
normalize(audio),
epoch,
sample_rate=sample_rate,
)
accelerator.wait_for_everyone()
accelerator.end_training()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Simple example of a training script.")
parser.add_argument("--local_rank", type=int, default=-1)
parser.add_argument("--dataset_name", type=str, default=None)
parser.add_argument("--dataset_config_name", type=str, default=None)
parser.add_argument(
"--train_data_dir",
type=str,
default=None,
help="A folder containing the training data.",
)
parser.add_argument("--output_dir", type=str, default="ddpm-model-64")
parser.add_argument("--overwrite_output_dir", type=bool, default=False)
parser.add_argument("--cache_dir", type=str, default=None)
parser.add_argument("--train_batch_size", type=int, default=16)
parser.add_argument("--eval_batch_size", type=int, default=16)
parser.add_argument("--num_epochs", type=int, default=100)
parser.add_argument("--save_images_epochs", type=int, default=10)
parser.add_argument("--save_model_epochs", type=int, default=10)
parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
parser.add_argument("--learning_rate", type=float, default=1e-4)
parser.add_argument("--lr_scheduler", type=str, default="cosine")
parser.add_argument("--lr_warmup_steps", type=int, default=500)
parser.add_argument("--adam_beta1", type=float, default=0.95)
parser.add_argument("--adam_beta2", type=float, default=0.999)
parser.add_argument("--adam_weight_decay", type=float, default=1e-6)
parser.add_argument("--adam_epsilon", type=float, default=1e-08)
parser.add_argument("--use_ema", type=bool, default=True)
parser.add_argument("--ema_inv_gamma", type=float, default=1.0)
parser.add_argument("--ema_power", type=float, default=3 / 4)
parser.add_argument("--ema_max_decay", type=float, default=0.9999)
parser.add_argument("--push_to_hub", type=bool, default=False)
parser.add_argument("--use_auth_token", type=bool, default=False)
parser.add_argument("--hub_token", type=str, default=None)
parser.add_argument("--hub_model_id", type=str, default=None)
parser.add_argument("--hub_private_repo", type=bool, default=False)
parser.add_argument("--logging_dir", type=str, default="logs")
parser.add_argument(
"--mixed_precision",
type=str,
default="no",
choices=["no", "fp16", "bf16"],
help=(
"Whether to use mixed precision. Choose"
"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
"and an Nvidia Ampere GPU."),
)
parser.add_argument("--hop_length", type=int, default=512)
parser.add_argument("--from_pretrained", type=str, default=None)
parser.add_argument("--start_epoch", type=int, default=0)
parser.add_argument("--num_train_steps", type=int, default=1000)
parser.add_argument("--scheduler",
type=str,
default="ddpm",
help="ddpm or ddim")
parser.add_argument("--vae",
type=str,
default=None,
help="pretrained VAE model for latent diffusion")
args = parser.parse_args()
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
if env_local_rank != -1 and env_local_rank != args.local_rank:
args.local_rank = env_local_rank
if args.dataset_name is None and args.train_data_dir is None:
raise ValueError(
"You must specify either a dataset name from the hub or a train data directory."
)
if args.dataset_name is not None and args.dataset_name == args.hub_model_id:
raise ValueError(
"The local dataset name must be different from the hub model id.")
main(args)