NeuralSeq/modules/diff/diffusion.py

import math
import random
from functools import partial
from inspect import isfunction
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from tqdm import tqdm
from einops import rearrange

from modules.fastspeech.fs2 import FastSpeech2
from modules.diffsinger_midi.fs2 import FastSpeech2MIDI
from utils.hparams import hparams



def exists(x):
 return x is not None


def default(val, d):
 if exists(val):
 return val
 return d() if isfunction(d) else d


def cycle(dl):
 while True:
 for data in dl:
 yield data


def num_to_groups(num, divisor):
 groups = num // divisor
 remainder = num % divisor
 arr = [divisor] * groups
 if remainder > 0:
 arr.append(remainder)
 return arr


class Residual(nn.Module):
 def __init__(self, fn):
 super().__init__()
 self.fn = fn

 def forward(self, x, *args, **kwargs):
 return self.fn(x, *args, **kwargs) + x


class SinusoidalPosEmb(nn.Module):
 def __init__(self, dim):
 super().__init__()
 self.dim = dim

 def forward(self, x):
 device = x.device
 half_dim = self.dim // 2
 emb = math.log(10000) / (half_dim - 1)
 emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
 emb = x[:, None] * emb[None, :]
 emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
 return emb


class Mish(nn.Module):
 def forward(self, x):
 return x * torch.tanh(F.softplus(x))


class Upsample(nn.Module):
 def __init__(self, dim):
 super().__init__()
 self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1)

 def forward(self, x):
 return self.conv(x)


class Downsample(nn.Module):
 def __init__(self, dim):
 super().__init__()
 self.conv = nn.Conv2d(dim, dim, 3, 2, 1)

 def forward(self, x):
 return self.conv(x)


class Rezero(nn.Module):
 def __init__(self, fn):
 super().__init__()
 self.fn = fn
 self.g = nn.Parameter(torch.zeros(1))

 def forward(self, x):
 return self.fn(x) * self.g


# building block modules

class Block(nn.Module):
 def __init__(self, dim, dim_out, groups=8):
 super().__init__()
 self.block = nn.Sequential(
 nn.Conv2d(dim, dim_out, 3, padding=1),
 nn.GroupNorm(groups, dim_out),
 Mish()
 )

 def forward(self, x):
 return self.block(x)


class ResnetBlock(nn.Module):
 def __init__(self, dim, dim_out, *, time_emb_dim, groups=8):
 super().__init__()
 self.mlp = nn.Sequential(
 Mish(),
 nn.Linear(time_emb_dim, dim_out)
 )

 self.block1 = Block(dim, dim_out)
 self.block2 = Block(dim_out, dim_out)
 self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()

 def forward(self, x, time_emb):
 h = self.block1(x)
 h += self.mlp(time_emb)[:, :, None, None]
 h = self.block2(h)
 return h + self.res_conv(x)


class LinearAttention(nn.Module):
 def __init__(self, dim, heads=4, dim_head=32):
 super().__init__()
 self.heads = heads
 hidden_dim = dim_head * heads
 self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
 self.to_out = nn.Conv2d(hidden_dim, dim, 1)

 def forward(self, x):
 b, c, h, w = x.shape
 qkv = self.to_qkv(x)
 q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads=self.heads, qkv=3)
 k = k.softmax(dim=-1)
 context = torch.einsum('bhdn,bhen->bhde', k, v)
 out = torch.einsum('bhde,bhdn->bhen', context, q)
 out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
 return self.to_out(out)


# gaussian diffusion trainer class

def extract(a, t, x_shape):
 b, *_ = t.shape
 out = a.gather(-1, t)
 return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def noise_like(shape, device, repeat=False):
 repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
 noise = lambda: torch.randn(shape, device=device)
 return repeat_noise() if repeat else noise()


def cosine_beta_schedule(timesteps, s=0.008):
 """
 cosine schedule
 as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
 """
 steps = timesteps + 1
 x = np.linspace(0, steps, steps)
 alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
 alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
 betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
 return np.clip(betas, a_min=0, a_max=0.999)


class GaussianDiffusion(nn.Module):
 def __init__(self, phone_encoder, out_dims, denoise_fn,
 timesteps=1000, loss_type='l1', betas=None, spec_min=None, spec_max=None):
 super().__init__()
 self.denoise_fn = denoise_fn
 if hparams.get('use_midi') is not None and hparams['use_midi']:
 self.fs2 = FastSpeech2MIDI(phone_encoder, out_dims)
 else:
 self.fs2 = FastSpeech2(phone_encoder, out_dims)
 self.fs2.decoder = None
 self.mel_bins = out_dims

 if exists(betas):
 betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
 else:
 betas = cosine_beta_schedule(timesteps)

 alphas = 1. - betas
 alphas_cumprod = np.cumprod(alphas, axis=0)
 alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])

 timesteps, = betas.shape
 self.num_timesteps = int(timesteps)
 self.loss_type = loss_type

 to_torch = partial(torch.tensor, dtype=torch.float32)

 self.register_buffer('betas', to_torch(betas))
 self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
 self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))

 # calculations for diffusion q(x_t | x_{t-1}) and others
 self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
 self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
 self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
 self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
 self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))

 # calculations for posterior q(x_{t-1} | x_t, x_0)
 posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
 # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
 self.register_buffer('posterior_variance', to_torch(posterior_variance))
 # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
 self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
 self.register_buffer('posterior_mean_coef1', to_torch(
 betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
 self.register_buffer('posterior_mean_coef2', to_torch(
 (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))

 self.register_buffer('spec_min', torch.FloatTensor(spec_min)[None, None, :hparams['keep_bins']])
 self.register_buffer('spec_max', torch.FloatTensor(spec_max)[None, None, :hparams['keep_bins']])

 def q_mean_variance(self, x_start, t):
 mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
 variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
 log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
 return mean, variance, log_variance

 def predict_start_from_noise(self, x_t, t, noise):
 return (
 extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
 extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
 )

 def q_posterior(self, x_start, x_t, t):
 posterior_mean = (
 extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
 extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
 )
 posterior_variance = extract(self.posterior_variance, t, x_t.shape)
 posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
 return posterior_mean, posterior_variance, posterior_log_variance_clipped

 def p_mean_variance(self, x, t, cond, clip_denoised: bool):
 noise_pred = self.denoise_fn(x, t, cond=cond)
 x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)

 if clip_denoised:
 x_recon.clamp_(-1., 1.)

 model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
 return model_mean, posterior_variance, posterior_log_variance

 @torch.no_grad()
 def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
 b, *_, device = *x.shape, x.device
 model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond, clip_denoised=clip_denoised)
 noise = noise_like(x.shape, device, repeat_noise)
 # no noise when t == 0
 nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
 return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise

 def q_sample(self, x_start, t, noise=None):
 noise = default(noise, lambda: torch.randn_like(x_start))
 return (
 extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
 extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
 )

 def p_losses(self, x_start, t, cond, noise=None, nonpadding=None):
 noise = default(noise, lambda: torch.randn_like(x_start))

 x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
 x_recon = self.denoise_fn(x_noisy, t, cond)

 if self.loss_type == 'l1':
 if nonpadding is not None:
 loss = ((noise - x_recon).abs() * nonpadding.unsqueeze(1)).mean()
 else:
 # print('are you sure w/o nonpadding?')
 loss = (noise - x_recon).abs().mean()

 elif self.loss_type == 'l2':
 loss = F.mse_loss(noise, x_recon)
 else:
 raise NotImplementedError()

 return loss

 def forward(self, txt_tokens, mel2ph=None, spk_embed=None,
 ref_mels=None, f0=None, uv=None, energy=None, infer=False):
 b, *_, device = *txt_tokens.shape, txt_tokens.device
 ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
 skip_decoder=True, infer=infer)
 cond = ret['decoder_inp'].transpose(1, 2)
 if not infer:
 t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
 x = ref_mels
 x = self.norm_spec(x)
 x = x.transpose(1, 2)[:, None, :, :] # [B, 1, M, T]
 nonpadding = (mel2ph != 0).float()
 ret['diff_loss'] = self.p_losses(x, t, cond, nonpadding=nonpadding)
 else:
 t = self.num_timesteps
 shape = (cond.shape[0], 1, self.mel_bins, cond.shape[2])
 x = torch.randn(shape, device=device)
 for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
 x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
 x = x[:, 0].transpose(1, 2)
 ret['mel_out'] = self.denorm_spec(x)

 return ret

 def norm_spec(self, x):
 return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1

 def denorm_spec(self, x):
 return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min

 def cwt2f0_norm(self, cwt_spec, mean, std, mel2ph):
 return self.fs2.cwt2f0_norm(cwt_spec, mean, std, mel2ph)

 def out2mel(self, x):
 return x