justyoung/DiffSinger

1import math2import random3from functools import partial4from inspect import isfunction5from pathlib import Path6import numpy as np7import torch8import torch.nn.functional as F9from torch import nn10from tqdm import tqdm11from einops import rearrange12 13from modules.fastspeech.fs2 import FastSpeech214from modules.diffsinger_midi.fs2 import FastSpeech2MIDI15from utils.hparams import hparams16 17 18 19def exists(x):20    return x is not None21 22 23def default(val, d):24    if exists(val):25        return val26    return d() if isfunction(d) else d27 28 29def cycle(dl):30    while True:31        for data in dl:32            yield data33 34 35def num_to_groups(num, divisor):36    groups = num // divisor37    remainder = num % divisor38    arr = [divisor] * groups39    if remainder > 0:40        arr.append(remainder)41    return arr42 43 44class Residual(nn.Module):45    def __init__(self, fn):46        super().__init__()47        self.fn = fn48 49    def forward(self, x, *args, **kwargs):50        return self.fn(x, *args, **kwargs) + x51 52 53class SinusoidalPosEmb(nn.Module):54    def __init__(self, dim):55        super().__init__()56        self.dim = dim57 58    def forward(self, x):59        device = x.device60        half_dim = self.dim // 261        emb = math.log(10000) / (half_dim - 1)62        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)63        emb = x[:, None] * emb[None, :]64        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)65        return emb66 67 68class Mish(nn.Module):69    def forward(self, x):70        return x * torch.tanh(F.softplus(x))71 72 73class Upsample(nn.Module):74    def __init__(self, dim):75        super().__init__()76        self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1)77 78    def forward(self, x):79        return self.conv(x)80 81 82class Downsample(nn.Module):83    def __init__(self, dim):84        super().__init__()85        self.conv = nn.Conv2d(dim, dim, 3, 2, 1)86 87    def forward(self, x):88        return self.conv(x)89 90 91class Rezero(nn.Module):92    def __init__(self, fn):93        super().__init__()94        self.fn = fn95        self.g = nn.Parameter(torch.zeros(1))96 97    def forward(self, x):98        return self.fn(x) * self.g99 100 101# building block modules102 103class Block(nn.Module):104    def __init__(self, dim, dim_out, groups=8):105        super().__init__()106        self.block = nn.Sequential(107            nn.Conv2d(dim, dim_out, 3, padding=1),108            nn.GroupNorm(groups, dim_out),109            Mish()110        )111 112    def forward(self, x):113        return self.block(x)114 115 116class ResnetBlock(nn.Module):117    def __init__(self, dim, dim_out, *, time_emb_dim, groups=8):118        super().__init__()119        self.mlp = nn.Sequential(120            Mish(),121            nn.Linear(time_emb_dim, dim_out)122        )123 124        self.block1 = Block(dim, dim_out)125        self.block2 = Block(dim_out, dim_out)126        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()127 128    def forward(self, x, time_emb):129        h = self.block1(x)130        h += self.mlp(time_emb)[:, :, None, None]131        h = self.block2(h)132        return h + self.res_conv(x)133 134 135class LinearAttention(nn.Module):136    def __init__(self, dim, heads=4, dim_head=32):137        super().__init__()138        self.heads = heads139        hidden_dim = dim_head * heads140        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)141        self.to_out = nn.Conv2d(hidden_dim, dim, 1)142 143    def forward(self, x):144        b, c, h, w = x.shape145        qkv = self.to_qkv(x)146        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads=self.heads, qkv=3)147        k = k.softmax(dim=-1)148        context = torch.einsum('bhdn,bhen->bhde', k, v)149        out = torch.einsum('bhde,bhdn->bhen', context, q)150        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)151        return self.to_out(out)152 153 154# gaussian diffusion trainer class155 156def extract(a, t, x_shape):157    b, *_ = t.shape158    out = a.gather(-1, t)159    return out.reshape(b, *((1,) * (len(x_shape) - 1)))160 161 162def noise_like(shape, device, repeat=False):163    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))164    noise = lambda: torch.randn(shape, device=device)165    return repeat_noise() if repeat else noise()166 167 168def cosine_beta_schedule(timesteps, s=0.008):169    """170    cosine schedule171    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ172    """173    steps = timesteps + 1174    x = np.linspace(0, steps, steps)175    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2176    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]177    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])178    return np.clip(betas, a_min=0, a_max=0.999)179 180 181class GaussianDiffusion(nn.Module):182    def __init__(self, phone_encoder, out_dims, denoise_fn,183                 timesteps=1000, loss_type='l1', betas=None, spec_min=None, spec_max=None):184        super().__init__()185        self.denoise_fn = denoise_fn186        if hparams.get('use_midi') is not None and hparams['use_midi']:187            self.fs2 = FastSpeech2MIDI(phone_encoder, out_dims)188        else:189            self.fs2 = FastSpeech2(phone_encoder, out_dims)190        self.fs2.decoder = None191        self.mel_bins = out_dims192 193        if exists(betas):194            betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas195        else:196            betas = cosine_beta_schedule(timesteps)197 198        alphas = 1. - betas199        alphas_cumprod = np.cumprod(alphas, axis=0)200        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])201 202        timesteps, = betas.shape203        self.num_timesteps = int(timesteps)204        self.loss_type = loss_type205 206        to_torch = partial(torch.tensor, dtype=torch.float32)207 208        self.register_buffer('betas', to_torch(betas))209        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))210        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))211 212        # calculations for diffusion q(x_t | x_{t-1}) and others213        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))214        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))215        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))216        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))217        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))218 219        # calculations for posterior q(x_{t-1} | x_t, x_0)220        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)221        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)222        self.register_buffer('posterior_variance', to_torch(posterior_variance))223        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain224        self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))225        self.register_buffer('posterior_mean_coef1', to_torch(226            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))227        self.register_buffer('posterior_mean_coef2', to_torch(228            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))229 230        self.register_buffer('spec_min', torch.FloatTensor(spec_min)[None, None, :hparams['keep_bins']])231        self.register_buffer('spec_max', torch.FloatTensor(spec_max)[None, None, :hparams['keep_bins']])232 233    def q_mean_variance(self, x_start, t):234        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start235        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)236        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)237        return mean, variance, log_variance238 239    def predict_start_from_noise(self, x_t, t, noise):240        return (241                extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -242                extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise243        )244 245    def q_posterior(self, x_start, x_t, t):246        posterior_mean = (247                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +248                extract(self.posterior_mean_coef2, t, x_t.shape) * x_t249        )250        posterior_variance = extract(self.posterior_variance, t, x_t.shape)251        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)252        return posterior_mean, posterior_variance, posterior_log_variance_clipped253 254    def p_mean_variance(self, x, t, cond, clip_denoised: bool):255        noise_pred = self.denoise_fn(x, t, cond=cond)256        x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)257 258        if clip_denoised:259            x_recon.clamp_(-1., 1.)260 261        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)262        return model_mean, posterior_variance, posterior_log_variance263 264    @torch.no_grad()265    def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):266        b, *_, device = *x.shape, x.device267        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond, clip_denoised=clip_denoised)268        noise = noise_like(x.shape, device, repeat_noise)269        # no noise when t == 0270        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))271        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise272 273    def q_sample(self, x_start, t, noise=None):274        noise = default(noise, lambda: torch.randn_like(x_start))275        return (276                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +277                extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise278        )279 280    def p_losses(self, x_start, t, cond, noise=None, nonpadding=None):281        noise = default(noise, lambda: torch.randn_like(x_start))282 283        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)284        x_recon = self.denoise_fn(x_noisy, t, cond)285 286        if self.loss_type == 'l1':287            if nonpadding is not None:288                loss = ((noise - x_recon).abs() * nonpadding.unsqueeze(1)).mean()289            else:290                # print('are you sure w/o nonpadding?')291                loss = (noise - x_recon).abs().mean()292 293        elif self.loss_type == 'l2':294            loss = F.mse_loss(noise, x_recon)295        else:296            raise NotImplementedError()297 298        return loss299 300    def forward(self, txt_tokens, mel2ph=None, spk_embed=None,301                ref_mels=None, f0=None, uv=None, energy=None, infer=False):302        b, *_, device = *txt_tokens.shape, txt_tokens.device303        ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,304                       skip_decoder=True, infer=infer)305        cond = ret['decoder_inp'].transpose(1, 2)306        if not infer:307            t = torch.randint(0, self.num_timesteps, (b,), device=device).long()308            x = ref_mels309            x = self.norm_spec(x)310            x = x.transpose(1, 2)[:, None, :, :]  # [B, 1, M, T]311            nonpadding = (mel2ph != 0).float()312            ret['diff_loss'] = self.p_losses(x, t, cond, nonpadding=nonpadding)313        else:314            t = self.num_timesteps315            shape = (cond.shape[0], 1, self.mel_bins, cond.shape[2])316            x = torch.randn(shape, device=device)317            for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):318                x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)319            x = x[:, 0].transpose(1, 2)320            ret['mel_out'] = self.denorm_spec(x)321 322        return ret323 324    def norm_spec(self, x):325        return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1326 327    def denorm_spec(self, x):328        return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min329 330    def cwt2f0_norm(self, cwt_spec, mean, std, mel2ph):331        return self.fs2.cwt2f0_norm(cwt_spec, mean, std, mel2ph)332 333    def out2mel(self, x):334        return x335