import os import pytorch_lightning as pl import torch import torch.nn.functional as F import torch.nn as nn import matplotlib.pyplot as plt from tqdm import tqdm from sklearn.metrics import precision_recall_fscore_support import numpy as np from scipy.stats import pearsonr import matplotlib.animation as animation from mpl_toolkits.axes_grid1 import make_axes_locatable MIN_MIDI = 21 MAX_MIDI = 108 HOP_LENGTH = 160 SAMPLE_RATE = 16000 import partitura as pt from utils import animate_sampling, get_batch_slice, plot_codec, p_codec_scale, tensor_pair_swap from renderer import Renderer # from model.utils import Normalization def linear_beta_schedule(beta_start, beta_end, timesteps): return torch.linspace(beta_start, beta_end, timesteps) def q_sample(x_start, t, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, noise=None): """ x_start: x0 (B, 1, T, F) t: timestep information (B,) """ # sqrt_alphas is mean of the Gaussian N() - for the forward process distribution sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[t] # extract the value of \bar{\alpha} at time=t # one minus alphas is variance of the Gaussian N() sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[t] # boardcasting into correct shape sqrt_alphas_cumprod_t = sqrt_alphas_cumprod_t[:, None, None, None].to(x_start.device) sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod_t[:, None, None, None].to(x_start.device) # scale down the input, and scale up the noise as time increases? return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise def extract_x0(x_t, epsilon, t, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod): """ x_t: The output from q_sample epsilon: The noise predicted from the model t: timestep information """ # sqrt_alphas is mean of the Gaussian N() sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[t] # extract the value of \bar{\alpha} at time=t sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[t] sqrt_alphas_cumprod_t = sqrt_alphas_cumprod_t[:, None, None, None].to(x_t.device) sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod_t[:, None, None, None].to(x_t.device) # obtaining x0 like the one in audioLDM? pred_x0 = 1.0 / sqrt_alphas_cumprod_t * x_t - epsilon / sqrt_one_minus_alphas_cumprod_t # pred_x0 = (x_t - sqrt_one_minus_alphas_cumprod_t * epsilon) / sqrt_alphas_cumprod_t # original - should be wrong - but it's the same as the labml?? return torch.clamp(pred_x0, -1.0, 1.0) class RollDiffusion(pl.LightningModule): def __init__(self, lr, timesteps, loss_type, beta_start, ): super().__init__() self.save_hyperparameters() # define beta schedule # beta is variance self.betas = linear_beta_schedule(timesteps=timesteps) # define alphas - reparameterization alphas = 1. - self.betas alphas_cumprod = torch.cumprod(alphas, axis=0) # alpha bar alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) self.sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # calculations for diffusion q(x_t | x_{t-1}) and others self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod) # calculations for posterior q(x_{t-1} | x_t, x_0) self.posterior_variance = self.betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) def training_step(self, batch, batch_idx): batch_size = batch["frame"].shape[0] batch = batch["frame"].unsqueeze(1) device = batch.device # Algorithm 1 line 3: sample t uniformally for every example in the batch t = torch.randint(0, self.hparams.timesteps, (batch_size,), device=device).long() loss = self.p_losses(batch, t, self.sqrt_alphas_cumprod, self.sqrt_one_minus_alphas_cumprod, loss_type=self.hparams.loss_type) self.log("Train/loss", loss) return loss def validation_step(self, batch, batch_idx): batch_size = batch["frame"].shape[0] batch = batch["frame"].unsqueeze(1) device = batch.device t = torch.randint(0, self.hparams.timesteps, (batch_size,), device=device).long() loss = self.p_losses(batch, t, self.sqrt_alphas_cumprod, self.sqrt_one_minus_alphas_cumprod, loss_type=self.hparams.loss_type) self.log("Val/loss", loss) def test_step(self, batch, batch_idx): batch_size = batch["frame"].shape[0] batch = batch["frame"].unsqueeze(1) device = batch.device t = torch.randint(0, self.hparams.timesteps, (batch_size,), device=device).long() loss = self.p_losses(batch, t, self.sqrt_alphas_cumprod, self.sqrt_one_minus_alphas_cumprod, loss_type=self.hparams.loss_type) self.log("Test/loss", loss) def predict_step(self, batch, batch_idx): # inference code # Unwrapping TensorDataset (list) # It is a pure noise img = batch[0] b = img.shape[0] # extracting batchsize device=img.device imgs = [] for i in tqdm(reversed(range(0, self.hparams.timesteps)), desc='sampling loop time step', total=self.hparams.timesteps): img = self.p_sample( img, i) img_npy = img.cpu().numpy() if (i+1)%10==0: for idx, j in enumerate(img_npy): # j (1, T, F) fig, ax = plt.subplots(1,1) ax.imshow(j[0].T, aspect='auto', origin='lower') self.logger.experiment.add_figure( f"sample_{idx}", fig, global_step=self.hparams.timesteps-i) # self.hparams.timesteps-i is used because slide bar won't show # if global step starts from self.hparams.timesteps imgs.append(img_npy) torch.save(imgs, 'imgs.pt') def p_losses(self, x_start, t, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, noise=None, loss_type="l1"): if noise is None: noise = torch.randn_like(x_start) x_noisy = q_sample(x_start=x_start, t=t, sqrt_alphas_cumprod=sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod=sqrt_one_minus_alphas_cumprod, noise=noise) predicted_noise = self(x_noisy, t) if loss_type == 'l1': loss = F.l1_loss(noise, predicted_noise) elif loss_type == 'l2': loss = F.mse_loss(noise, predicted_noise) elif loss_type == "huber": loss = F.smooth_l1_loss(noise, predicted_noise) else: raise NotImplementedError() return loss def p_sample(self, x, t_index): # x is Guassian noise # extracting coefficients at time t betas_t = self.betas[t_index] sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t_index] sqrt_recip_alphas_t = self.sqrt_recip_alphas[t_index] # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean model_mean = sqrt_recip_alphas_t * ( x - betas_t * self(x, t_tensor) / sqrt_one_minus_alphas_cumprod_t ) if t_index == 0: return model_mean else: # posterior_variance_t = extract(self.posterior_variance, t, x.shape) posterior_variance_t = self.posterior_variance[t_index] noise = torch.randn_like(x) # Algorithm 2 line 4: return model_mean + torch.sqrt(posterior_variance_t) * noise def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr) # scheduler = TriStageLRSchedule(optimizer, # [1e-8, self.hparams.lr, 1e-8], # [0.2,0.6,0.2], # max_update=len(self.train_dataloader.dataloader)*self.trainer.max_epochs) # scheduler = MultiStepLR(optimizer, [1,3,5,7,9], gamma=0.1, last_epoch=-1, verbose=False) # return [optimizer], [{"scheduler":scheduler, "interval": "step"}] return [optimizer] class CodecDiffusion(pl.LightningModule): def __init__(self, lr, timesteps, loss_type, loss_keys, beta_start, beta_end, training, sampling, samples_root, debug=False, generation_filter=0.0, **kwargs ): super().__init__() self.save_hyperparameters() # define beta schedule (beta is variance) self.betas = linear_beta_schedule(beta_start, beta_end, timesteps=timesteps) # define alphas alphas = 1. - self.betas alphas_cumprod = torch.cumprod(alphas, axis=0) alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) self.sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # calculations for diffusion q(x_t | x_{t-1}) and others self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod) # calculations for posterior q(x_{t-1} | x_t, x_0) self.posterior_variance = self.betas * (1. - alphas_cumprod_prev) / (1- alphas_cumprod) self.inner_loop = tqdm(range(self.hparams.timesteps), desc='sampling loop time step') self.reverse_diffusion = getattr(self, sampling['type']) # self.reverse_diffusion = getattr(self, sampling.type) self.alphas = alphas def training_step(self, batch, batch_idx): losses, tensors = self.step(batch, batch_idx) # calculating total loss based on keys give total_loss = 0 for k in self.hparams.loss_keys: total_loss += losses[k] self.log(f"Train/{k}", losses[k]) return total_loss def validation_step(self, batch, batch_idx): losses, tensors = self.step(batch, batch_idx) total_loss = 0 for k in self.hparams.loss_keys: total_loss += losses[k] self.log(f"Val/{k}", losses[k]) # sample a fraction (1 percent) of the testing set randgen = torch.rand(1)[0] if randgen <= self.hparams.valid_fraction: sampled_loss, tc_fig, tempo_vel_loss, tempo_vel_cor = self.predict(batch, batch_idx) self.logger.log_image(key=f"Val/tempo_curves", images=[tc_fig]) self.log(f"Val/sampled_loss", sampled_loss) self.log(f"Val/tempo_vel_loss", tempo_vel_loss) self.log(f"Val/tempo_vel_cor", tempo_vel_cor) def test_step(self, batch, batch_idx): # evaluation if batch_idx < self.hparams.eval_starting_epoch: return _, _, _, _ = self.predict(batch, batch_idx, evaluate=True) def predict(self, batch, batch_idx, save_animation=False, evaluate=False): batch_label = batch batch_source = dict([(k, tensor_pair_swap(v)) for k, v in batch.items()]) # rescale batch_source['p_codec'] = p_codec_scale(batch_source['p_codec'], self.hparams.dataset_means, self.hparams.dataset_stds) batch_label['p_codec'] = p_codec_scale(batch_label['p_codec'], self.hparams.dataset_means, self.hparams.dataset_stds) save_root = f'{self.hparams.samples_root}/{self.logger._name}/epoch={self.current_epoch}/batch={batch_idx}' c_configs = self.set_eval_c_config() for c_config in c_configs: save_root_ = save_root + f"/{c_config}" start_noise, sample_steps, c_codec = self.set_predict_start(batch_source, batch_label, c_config=c_config) pred_list = self.p_sample(batch_label, start_noise=start_noise, sample_steps=sample_steps, c_codec=c_codec) # noise_list: [(pred_t, t), ..., (pred_0, 0)] pcodec_pred, _ = pred_list[-1] # (B, 1, T, F) batch_label_codec = batch_label['p_codec'].unsqueeze(1).cpu() N = len(batch_label['score_path']) os.makedirs(save_root_, exist_ok=True) # save the prediction samples (of one batch) and render. np.save(f'{save_root_}/test_sample.npy', pcodec_pred) # rescale the predictions and labels back to normal pcodec_pred = p_codec_scale(pcodec_pred, self.hparams.dataset_means, self.hparams.dataset_stds) tempo_vel_loss, tempo_vel_cor = 0, 0 if len(pcodec_pred) % 2 != 1: tc_fig, tempo_vel_loss, tempo_vel_cor = self.render_batch( pcodec_pred, batch_source, batch_label, save_root_, evaluate=evaluate) if save_animation: t_list = torch.arange(1, self.hparams.timesteps, 20).flip(0) if t_list[-1] != self.hparams.timesteps: t_list = torch.cat((t_list, torch.tensor([self.hparams.timesteps])), 0) fig, axes = plt.subplots(2 * N, 1, figsize=(24, 4 * N)) ax_flat = axes.flatten() caxs = [] for ax in axes.flatten(): div = make_axes_locatable(ax) caxs.append(div.append_axes('right', '5%', '5%')) ani = animation.FuncAnimation(fig, animate_sampling, frames=tqdm(t_list, desc='Animating'), fargs=(fig, ax_flat, caxs, noise_list, self.hparams.timesteps), interval=500, blit=False, repeat_delay=1000) ani.save(f'{save_root}/animation.gif', dpi=80, writer='imagemagick') # plot the comparison between prediction and label fig, ax = plt.subplots(2 * N, 1, figsize=(24, 4 * N)) for idx, (pred, label) in enumerate(zip(pcodec_pred, batch_label_codec)): plot_codec(pred, label, ax[idx*2], ax[idx*2+1], fig) self.logger.log_image(key=f"Val/pred_label", images=[fig]) plt.savefig(f"{save_root}/pred_label.png") sampled_loss = self.p_losses(batch_label_codec, torch.tensor(pcodec_pred), loss_type='l2') return sampled_loss, tc_fig, tempo_vel_loss, tempo_vel_cor def inference_one(self, batch): sample_steps = self.hparams.timesteps - 1 pred_list = self.p_sample(batch, sample_steps=sample_steps) pcodec_pred, _ = pred_list[-1] # Extract the final prediction pcodec_pred_processed = p_codec_scale(pcodec_pred, self.hparams.dataset_means, self.hparams.dataset_stds) return pcodec_pred_processed def step(self, batch, batch_idx): p_codec, s_codec, c_codec = batch['p_codec'], batch['s_codec'], batch['c_codec'] if self.hparams.drop_c_con: c_codec = torch.zeros_like(c_codec) batch_size = p_codec.shape[0] device = p_codec.device p_codec = p_codec.unsqueeze(1) # (B, 1, T, F) t = torch.randint(0, self.hparams.timesteps, (batch_size,), device=device).long().cpu() # more diverse sampling noise = torch.randn_like(p_codec) if self.hparams.training['target'] == "transfer": # invert each pair noise = tensor_pair_swap(p_codec) # in transfer context, c_codec is the difference between the two that's being transfered. c_codec = tensor_pair_swap(c_codec) - c_codec # (tgt - src) x_t = q_sample( # sampling noise at time t x_start=p_codec, t=t, sqrt_alphas_cumprod=self.sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod, noise=noise) # train to predict the noise, loss is computed for the noise if self.hparams.training['mode'] == 'epsilon': epsilon_pred, _ = self(x_t, s_codec, c_codec, t) # predict the noise N(0, 1) diffusion_loss = self.p_losses(noise, epsilon_pred, loss_type=self.hparams.loss_type) pred_p_codec = extract_x0( x_t, epsilon_pred, t, sqrt_alphas_cumprod=self.sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod) # train to predict p_codec elif self.hparams.training['mode'] == 'x_0': pred_p_codec, _ = self(x_t, s_codec, c_codec, t) # pred_p_codec = p_codec_scale(pred_p_codec, self.hparams.dataset_means, self.hparams.dataset_stds) diffusion_loss = self.p_losses(p_codec, pred_p_codec, loss_type=self.hparams.loss_type) elif self.hparams.training['mode'] == 'ex_0': epsilon_pred, _ = self(x_t, s_codec, c_codec, t) # predict the noise N(0, 1) pred_p_codec = extract_x0( x_t, epsilon_pred, t, sqrt_alphas_cumprod=self.sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod) diffusion_loss = self.p_losses(p_codec, pred_p_codec, loss_type=self.hparams.loss_type) else: raise ValueError(f"training mode {self.training.mode} is not supported. Please either use 'x_0' or 'epsilon'.") losses = {} # compute tempo_vel_loss for the prediced pcodec, for bp and velocity rows. # tempo_vel_loss = self.p_losses(p_codec[:, :, :, :2], pred_p_codec[:, :, :, :2], loss_type='l2') # only tempo and velocity weight_sum = sum(self.hparams.loss_weight) for idx, name in enumerate(['tempo', 'velocity', 'timing', 'duration', 'pedal']): losses[f"recon_{name}_loss"] = (self.p_losses(p_codec[..., idx], pred_p_codec[..., idx], loss_type='l2') * self.hparams.loss_weight[idx] / weight_sum * self.hparams.recon_loss_weight) losses["diffusion_loss"] = diffusion_loss tensors = { "pred_pcodec": pred_p_codec, "label_pcodec": p_codec } return losses, tensors def p_sample(self, batch, start_noise=None, sample_steps=None, c_codec=None): p_codec = batch['p_codec'].unsqueeze(1) s_codec = batch['s_codec'] if type(c_codec) == type(None): c_codec = batch['c_codec'] # Algorithm 1 line 3: sample t uniformally for every example in the batch self.inner_loop.refresh() self.inner_loop.reset() if type(start_noise) != type(None): noise = start_noise else: noise = torch.randn_like(p_codec.float()).to(self.device) noise_list = [] noise_list.append((noise, self.hparams.timesteps)) if type(sample_steps) == type(None): sample_steps = self.hparams.timesteps for t_index in reversed(range(0, sample_steps)): noise, _ = self.reverse_diffusion(noise, s_codec, c_codec, t_index) # cfdg_ddpm_x0 or cfdg_ddpm noise_npy = noise.detach().cpu().numpy() # self.hparams.timesteps-i is used because slide bar won't show # if global step starts from self.hparams.timesteps noise_list.append((noise_npy, t_index)) self.inner_loop.update() return noise_list def p_losses(self, label, prediction, loss_type="l1"): if loss_type == 'l1': loss = F.l1_loss(label, prediction) elif loss_type == 'l2': loss = F.mse_loss(label, prediction) elif loss_type == "huber": loss = F.smooth_l1_loss(label, prediction) else: raise NotImplementedError() return loss def ddpm(self, x, waveform, t_index): # x is Guassian noise # extracting coefficients at time t betas_t = self.betas[t_index] sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t_index] sqrt_recip_alphas_t = self.sqrt_recip_alphas[t_index] # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean epsilon, spec = self(x, waveform, t_tensor) model_mean = sqrt_recip_alphas_t * ( x - betas_t * epsilon / sqrt_one_minus_alphas_cumprod_t ) if t_index == 0: return model_mean, spec else: # posterior_variance_t = extract(self.posterior_variance, t, x.shape) posterior_variance_t = self.posterior_variance[t_index] noise = torch.randn_like(x) # Algorithm 2 line 4: return (model_mean + torch.sqrt(posterior_variance_t) * noise), spec def ddpm_x0(self, x, waveform, t_index): # x is x_t, when t=T it is pure Gaussian # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean x0_pred, spec = self(x, waveform, t_tensor) if t_index == 0: sigma = (1/self.sqrt_one_minus_alphas_cumprod[t_index]) * ( torch.sqrt(1-self.alphas[t_index])) model_mean = x0_pred / self.sqrt_alphas_cumprod[t_index] else: sigma = (self.sqrt_one_minus_alphas_cumprod[t_index-1]/self.sqrt_one_minus_alphas_cumprod[t_index]) * ( torch.sqrt(1-self.alphas[t_index])) model_mean = (self.sqrt_alphas_cumprod[t_index-1]) * x0_pred + ( torch.sqrt(1 - self.sqrt_alphas_cumprod[t_index-1]**2 - sigma**2) * ( x-self.sqrt_alphas_cumprod[t_index]* x0_pred)/self.sqrt_one_minus_alphas_cumprod[t_index]) + ( sigma * torch.randn_like(x)) return model_mean, spec def ddim_x0(self, x, waveform, t_index): # x is x_t, when t=T it is pure Gaussian # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean x0_pred, spec = self(x, waveform, t_tensor) if t_index == 0: sigma = 0 model_mean = x0_pred / self.sqrt_alphas_cumprod[t_index] else: sigma = 0 model_mean = (self.sqrt_alphas_cumprod[t_index-1]) * x0_pred + ( torch.sqrt(1 - self.sqrt_alphas_cumprod[t_index-1]**2 - sigma**2) * ( x-self.sqrt_alphas_cumprod[t_index]* x0_pred)/self.sqrt_one_minus_alphas_cumprod[t_index]) + ( sigma * torch.randn_like(x)) return model_mean, spec def ddim(self, x, waveform, t_index): # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean epsilon, spec = self(x, waveform, t_tensor) if t_index == 0: model_mean = (x - self.sqrt_one_minus_alphas_cumprod[t_index] * epsilon) / self.sqrt_alphas_cumprod[t_index] else: model_mean = (self.sqrt_alphas_cumprod[t_index-1]) * ( (x - self.sqrt_one_minus_alphas_cumprod[t_index] * epsilon) / self.sqrt_alphas_cumprod[t_index]) + ( self.sqrt_one_minus_alphas_cumprod[t_index-1] * epsilon) return model_mean, spec def ddim2ddpm(self, x, waveform, t_index): # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean epsilon, spec = self(x, waveform, t_tensor) if t_index == 0: model_mean = (x - self.sqrt_one_minus_alphas_cumprod[t_index] * epsilon) / self.sqrt_alphas_cumprod[t_index] else: sigma = (self.sqrt_one_minus_alphas_cumprod[t_index-1]/self.sqrt_one_minus_alphas_cumprod[t_index]) * ( torch.sqrt(1-self.alphas[t_index])) model_mean = (self.sqrt_alphas_cumprod[t_index-1]) * ( (x - self.sqrt_one_minus_alphas_cumprod[t_index] * epsilon) / self.sqrt_alphas_cumprod[t_index]) + ( torch.sqrt(1 - self.sqrt_alphas_cumprod[t_index-1]**2 - sigma**2) * epsilon) + sigma * torch.randn_like(x) return model_mean, spec def cfdg_ddpm(self, x, s_codec, c_codec, t_index): # x is Guassian noise # extracting coefficients at time t betas_t = self.betas[t_index] sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t_index] sqrt_recip_alphas_t = self.sqrt_recip_alphas[t_index] # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Use our model (noise predictor) to predict the mean epsilon_c, cond = self(x, s_codec, c_codec, t_tensor) epsilon_0, _ = self(x, torch.zeros_like(s_codec), torch.zeros_like(c_codec), t_tensor) epsilon = self.hparams.sampling['w'] * epsilon_c + (1 - self.hparams.sampling['w']) * epsilon_0 # Equation 11 in the paper model_mean = sqrt_recip_alphas_t * ( x - betas_t * epsilon / sqrt_one_minus_alphas_cumprod_t ) if t_index == 0: return model_mean, cond else: # posterior_variance_t = extract(self.posterior_variance, t, x.shape) posterior_variance_t = self.posterior_variance[t_index] noise = torch.randn_like(x) variance = torch.sqrt(posterior_variance_t) * noise # From the diffusion-inspired training strategy paper, remove the deviation parameter since noise is deterministic if self.hparams.training['target'] == "transfer": return model_mean, cond return (model_mean + variance), cond def cfdg_ddpm_x0(self, x, condition, t_index): # x is x_t, when t=T it is pure Gaussian # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Use model (noise & x0 predictor) to predict the mean and weight the conditioned & unconditioned if self.hparams.training.mode == "ex_0" or self.hparams.training.mode == "epsilon": epsilon_c, cond = self(x, condition, t_tensor) epsilon_0, _ = self(x, torch.zeros_like(condition), t_tensor) epsilon = self.hparams.sampling['w'] * epsilon_c + (1 - self.hparams.sampling['w']) * epsilon_0 x0_pred = extract_x0(x, epsilon, t_tensor, sqrt_alphas_cumprod=self.sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod) elif self.hparams.training.mode == "x0": x0_pred_c, cond = self(x, condition, t_tensor) x0_pred_0, _ = self(x, torch.zeros_like(condition), t_tensor, sampling=True) # if sampling = True, the input condition will be overwritten # wait... is this really weighting??? # x0_pred = (1 + self.hparams.sampling.w) * x0_pred_c - self.hparams.sampling.w * x0_pred_0 x0_pred = self.hparams.sampling['w'] * x0_pred_c + (1 - self.hparams.sampling['w']) * x0_pred_0 if t_index == 0: sigma = (1/self.sqrt_one_minus_alphas_cumprod[t_index]) * ( torch.sqrt(1-self.alphas[t_index])) model_mean = x0_pred / self.sqrt_alphas_cumprod[t_index] else: # sigma is the noise magnitude (variance to scale noise) sigma = (self.sqrt_one_minus_alphas_cumprod[t_index - 1] / self.sqrt_one_minus_alphas_cumprod[t_index]) * ( torch.sqrt(1 - self.alphas[t_index])) if self.hparams.training.mode == "x0": # I modified the division to multiplication according to the inverse of DDPM paper eq.11 epsilon = (x - self.sqrt_alphas_cumprod[t_index] * x0_pred) * self.sqrt_one_minus_alphas_cumprod[t_index] model_mean = (self.sqrt_alphas_cumprod[t_index - 1]) * x0_pred + ( torch.sqrt(1 - self.sqrt_alphas_cumprod[t_index - 1] ** 2 - sigma ** 2) * epsilon) + ( # epsilon sigma * torch.randn_like(x)) return model_mean, cond def cfdg_ddim_x0(self, x, waveform, t_index): # x is x_t, when t=T it is pure Gaussian # boardcasting t_index into a tensor t_tensor = torch.tensor(t_index).repeat(x.shape[0]).to(x.device) # Equation 11 in the paper # Use our model (noise predictor) to predict the mean x0_pred_c, spec = self(x, waveform, t_tensor) x0_pred_0, _ = self(x, torch.zeros_like(waveform), t_tensor) x0_pred = (1 + self.hparams.sampling.w) * x0_pred_c - self.hparams.sampling.w * x0_pred_0 # x0_pred = x0_pred_c # x0_pred = x0_pred_0 if t_index == 0: sigma = 0 model_mean = x0_pred / self.sqrt_alphas_cumprod[t_index] else: sigma = 0 model_mean = (self.sqrt_alphas_cumprod[t_index-1]) * x0_pred + ( torch.sqrt(1 - self.sqrt_alphas_cumprod[t_index - 1] ** 2 - sigma ** 2) * ( x - self.sqrt_alphas_cumprod[t_index] * x0_pred) / self.sqrt_one_minus_alphas_cumprod[t_index]) + ( sigma * torch.randn_like(x)) return model_mean, spec def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr) # scheduler = TriStageLRSchedule(optimizer, # [1e-8, self.hparams.lr, 1e-8], # [0.2,0.6,0.2], # max_update=len(self.train_dataloader.dataloader)*self.trainer.max_epochs) # scheduler = MultiStepLR(optimizer, [1,3,5,7,9], gamma=0.1, last_epoch=-1, verbose=False) # return [optimizer], [{"scheduler":scheduler, "interval": "step"}] return [optimizer] def animate_sampling(self, t_idx, fig, ax_flat, caxs, noise_list): # noise_list: Tuple of (x_t, t), (x_t-1, t-1), ... (x_0, 0) # x_t (B, 1, T, F) # clearing figures to prevent slow down in each iteration.d fig.canvas.draw() for idx in range(len(noise_list[0][0])): # visualize only 4 piano rolls in a batch ax_flat[idx].cla() ax_flat[4+idx].cla() caxs[idx].cla() caxs[4+idx].cla() # roll_pred (1, T, F) im1 = ax_flat[idx].imshow(noise_list[0][0][idx][0].detach().T.cpu(), aspect='auto', origin='lower') im2 = ax_flat[4+idx].imshow(noise_list[1 + self.hparams.timesteps - t_idx][0][idx][0].T, aspect='auto', origin='lower') fig.colorbar(im1, cax=caxs[idx]) fig.colorbar(im2, cax=caxs[4+idx]) fig.suptitle(f't={t_idx}') row1_txt = ax_flat[0].text(-400,45,f'Gaussian N(0,1)') row2_txt = ax_flat[4].text(-300,45,'x_{t-1}') def render_batch(self, pcodec_pred, batch_source, batch_label, save_root, evaluate=False): B = len(pcodec_pred) fig, ax = plt.subplots(int(B/2), 4, figsize=(24, 3*B)) tempo_vel_loss, tempo_vel_cor = 0, 0 for idx in range(B): renderer = Renderer(save_root, pcodec_pred[idx], source_data=get_batch_slice(batch_source, idx), label_data=get_batch_slice(batch_label, idx), with_source=self.hparams.transfer, idx=idx, B=B) tvl, tvc = renderer.render_sample(save_sourcelabel=True) tempo_vel_loss += tvl tempo_vel_cor += tvc if renderer.success: renderer.plot_curves(ax) if evaluate: renderer.save_performance_features(save_source=False, save_label=True) renderer.save_pf_distribution() plt.savefig(f"{save_root}/tempo_curves.png") return fig, tempo_vel_loss / B, tempo_vel_cor / B def set_predict_start(self, batch_source, batch_label, c_config=""): """set up the starting point of inference given the configuration. - transfer in training: start from source codec, conditioned on label c_codec - source c_codec - transfer in inference: starting from source codec + noise (75%), conditioned on label c_codec - no transfer: starting from noise, conditioned on label c_codec [or other adjustable c_codec] c_config: e.g. "melodiousness_more", changes the parameter by half or double """ batch_source_codec = batch_source['p_codec'] batch_source_codec = batch_source_codec.unsqueeze(1) # (B, 1, T, F) sample_steps = self.hparams.timesteps - 1 c_codec = batch_label['c_codec'] if "transfer" not in self.hparams.training.target: # pure noise noise = torch.randn_like(batch_source_codec) if self.hparams.transfer: # only transfer in inference sample_steps = int((self.hparams.timesteps - 1) * self.hparams.sample_steps_frac) # steps for noisify the source start_noise = q_sample(x_start=batch_source_codec, t=torch.tensor([sample_steps] * int(batch_source['p_codec'].shape[0])), sqrt_alphas_cumprod=self.sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod, noise=noise) else: start_noise = None else: # transfer in training start_noise = batch_source_codec c_codec = batch_label['c_codec'] - batch_source['c_codec'] if c_config != "": c, control = c_config.split("/") c_idx = self.c_type.index(c) if control == "more": c_codec[c_idx] *= 2 if control == "less": c_codec[c_idx] *= 0.5 # normal is just itself return start_noise, sample_steps, c_codec def set_eval_c_config(self): # set the c_config for oracle evaluation if self.hparams.condition_eval: self.c_type = ["melodiousness", "articulation", "rhythm complexity", "rhythm stability", "dissonance", "tonal stability", "minorness"] c_configs = [f"{c}_{control}" for c in self.c_type for control in ["less", "normal", "more"]] else: c_configs = [""] return c_configs