| import torch |
| from torch.utils.data import DataLoader, Dataset |
| import torchaudio |
| import torchvision.transforms as tvt |
| from denoising_diffusion_pytorch.classifier_free_guidance import Unet, GaussianDiffusion |
| import glob |
| import torch.nn as nn |
| import time, math |
| from PIL import Image |
| from diffusers import Mel |
| import sys |
| import torchaudio |
| import librosa |
| import matplotlib.pyplot as plt |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| args = sys.argv[1:] |
|
|
| class Audio(Dataset): |
| def __init__(self, folder): |
| |
| self.waveforms = [] |
| self.labels = [] |
| print("Loading files...") |
| for file in glob.iglob(folder + '/**/*.wav', recursive=True): |
| self.labels.append(int(file.split('/')[-1][0])) |
| waveform, _ = torchaudio.load(file) |
| |
| self.waveforms.append(waveform) |
| |
| def __len__(self): |
| return len(self.waveforms) |
|
|
| def __getitem__(self, index): |
| return self.waveforms[index], self.labels[index] |
|
|
|
|
| image_size = 256 |
| if len(args) >= 1: |
| image_size = int(args[0]) |
|
|
| MEL = Mel(x_res=image_size, y_res=image_size) |
| img_to_tensor = tvt.PILToTensor() |
|
|
| def collate(batch): |
| spectros = [] |
| labels = [] |
| for waveform, label in batch: |
| MEL.load_audio(raw_audio=waveform[0]) |
| for slice in range(MEL.get_number_of_slices()): |
| spectro = MEL.audio_slice_to_image(slice) |
| spectro = img_to_tensor(spectro) / 255.0 |
| |
| |
| |
| |
| spectros.append(spectro) |
| labels.append(label) |
|
|
| spectros = torch.stack(spectros) |
| labels = torch.tensor(labels) |
| |
| return spectros.to(device), labels.to(device) |
|
|
|
|
| def initialize(scheduler = None, batch_size=32): |
| model = Unet( |
| dim = 64, |
| num_classes=10, |
| dim_mults=(1, 2, 4, 8), |
| channels=1 |
| ) |
| diffusion = GaussianDiffusion( |
| model, |
| image_size=image_size, |
| timesteps=1000, |
| loss_type = 'l2', |
| objective='pred_x0', |
| |
| ) |
| diffusion.to(device) |
| |
| optim = torch.optim.AdamW(model.parameters(), lr=1e-4, eps=1e-8) |
| if scheduler: |
| scheduler = torch.optim.lr_scheduler.CyclicLR(optim, base_lr=1e-5, max_lr=1e-3, mode="exp_range", cycle_momentum=False) |
| return diffusion, optim, scheduler |
|
|
| def timeSince(since): |
| now = time.time() |
| s = now - since |
| m = math.floor(s / 60) |
| s -= m * 60 |
| return '%dm %ds' % (m, s) |
|
|
| start = time.time() |
|
|
| def train(model, optim, train_dl, batch_size=32, epochs=5, scheduler = None): |
| size = len(train_dl.dataset) |
| model.train() |
| losses = [] |
| |
| for e in range(epochs): |
| batch_loss, batch_counts = 0, 0 |
| for step, batch in enumerate(train_dl): |
| model.zero_grad() |
| batch_counts += 1 |
| spectros, labels = batch |
| loss = model(spectros, classes=labels) |
| |
| batch_loss += loss.item() |
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1) |
| optim.step() |
| if scheduler is not None: |
| scheduler.step() |
| |
| if (step % 100 == 0 and step != 0) or (step == len(train_dl) - 1): |
| to_print = f"{e + 1:^7} | {step:^7} | {batch_loss / batch_counts:^12.6f} | {timeSince(start)} | {step*batch_size:>5d}/{size:>5d}" |
| print(to_print) |
| losses.append(batch_loss) |
| batch_loss, batch_counts = 0, 0 |
|
|
| labels = torch.randint(0,9,(8, )).to(device) |
| print(labels) |
| samples = model.sample(labels) |
| for i, sample in enumerate(samples): |
| im = Image.fromarray(sample[0].cpu().numpy() * 255).convert('L') |
| audio = torch.tensor([MEL.image_to_audio(im)]) |
| torchaudio.save(f"audio/sample{e}_{i}_{labels[i]}.wav", audio, 48000) |
| im.save(f"images/sample{e}_{i}_{labels[i]}.jpg") |
| return losses |
|
|
| if __name__ == "__main__": |
| num_epochs = 10 |
| if len(args) >= 2: |
| num_epochs = int(args[1]) |
|
|
| batch_size = 32 |
| if len(args) >= 3: |
| batch_size = int(args[2]) |
|
|
| print(image_size, num_epochs, batch_size) |
| model, optim, scheduler = initialize(scheduler=True, batch_size=batch_size) |
| train_data = Audio("AudioMNIST/data") |
| print("Done Loading") |
| train_dl = DataLoader(train_data, batch_size, True, collate_fn=collate) |
| train(model, optim, train_dl, batch_size, num_epochs, scheduler) |
| torch.save(model.state_dict(), "diffusion_condition_model.pt") |
