import math import torch from einops import rearrange from model.base import BaseModule class Mish(BaseModule): def forward(self, x): return x * torch.tanh(torch.nn.functional.softplus(x)) class Upsample(BaseModule): def __init__(self, dim): super(Upsample, self).__init__() self.conv = torch.nn.ConvTranspose2d(dim, dim, 4, 2, 1) def forward(self, x): return self.conv(x) class Downsample(BaseModule): def __init__(self, dim): super(Downsample, self).__init__() self.conv = torch.nn.Conv2d(dim, dim, 3, 2, 1) # kernel=3, stride=2, padding=1. def forward(self, x): return self.conv(x) class Rezero(BaseModule): def __init__(self, fn): super(Rezero, self).__init__() self.fn = fn self.g = torch.nn.Parameter(torch.zeros(1)) def forward(self, x): return self.fn(x) * self.g class Block(BaseModule): def __init__(self, dim, dim_out, groups=8): super(Block, self).__init__() self.block = torch.nn.Sequential(torch.nn.Conv2d(dim, dim_out, 3, padding=1), torch.nn.GroupNorm( groups, dim_out), Mish()) def forward(self, x, mask): output = self.block(x * mask) return output * mask class ResnetBlock(BaseModule): def __init__(self, dim, dim_out, time_emb_dim, groups=8): super(ResnetBlock, self).__init__() self.mlp = torch.nn.Sequential(Mish(), torch.nn.Linear(time_emb_dim, dim_out)) self.block1 = Block(dim, dim_out, groups=groups) self.block2 = Block(dim_out, dim_out, groups=groups) if dim != dim_out: self.res_conv = torch.nn.Conv2d(dim, dim_out, 1) else: self.res_conv = torch.nn.Identity() def forward(self, x, mask, time_emb): h = self.block1(x, mask) h += self.mlp(time_emb).unsqueeze(-1).unsqueeze(-1) h = self.block2(h, mask) output = h + self.res_conv(x * mask) return output class LinearAttention(BaseModule): def __init__(self, dim, heads=4, dim_head=32): super(LinearAttention, self).__init__() self.heads = heads hidden_dim = dim_head * heads self.to_qkv = torch.nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) # NOTE: 1x1 conv self.to_out = torch.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) class Residual(BaseModule): def __init__(self, fn): super(Residual, self).__init__() self.fn = fn def forward(self, x, *args, **kwargs): output = self.fn(x, *args, **kwargs) + x return output class SinusoidalPosEmb(BaseModule): def __init__(self, dim): super(SinusoidalPosEmb, self).__init__() self.dim = dim def forward(self, x, scale=1000): device = x.device half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb class GradLogPEstimator2d(BaseModule): def __init__(self, dim, dim_mults=(1, 2, 4), groups=8, spk_emb_dim=64, n_feats=80, pe_scale=1000): super(GradLogPEstimator2d, self).__init__() self.dim = dim self.dim_mults = dim_mults self.groups = groups self.spk_emb_dim = spk_emb_dim self.pe_scale = pe_scale self.spk_mlp = torch.nn.Sequential(torch.nn.Linear(spk_emb_dim, spk_emb_dim * 4), Mish(), torch.nn.Linear(spk_emb_dim * 4, n_feats)) self.time_pos_emb = SinusoidalPosEmb(dim) self.mlp = torch.nn.Sequential(torch.nn.Linear(dim, dim * 4), Mish(), torch.nn.Linear(dim * 4, dim)) dims = [3, *map(lambda m: dim * m, dim_mults)] in_out = list(zip(dims[:-1], dims[1:])) self.downs = torch.nn.ModuleList([]) self.ups = torch.nn.ModuleList([]) num_resolutions = len(in_out) for ind, (dim_in, dim_out) in enumerate(in_out): is_last = ind >= (num_resolutions - 1) self.downs.append(torch.nn.ModuleList([ ResnetBlock(dim_in, dim_out, time_emb_dim=dim), ResnetBlock(dim_out, dim_out, time_emb_dim=dim), Residual(Rezero(LinearAttention(dim_out))), Downsample(dim_out) if not is_last else torch.nn.Identity()])) mid_dim = dims[-1] self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim) self.mid_attn = Residual(Rezero(LinearAttention(mid_dim))) self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim) for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): self.ups.append(torch.nn.ModuleList([ ResnetBlock(dim_out * 2, dim_in, time_emb_dim=dim), ResnetBlock(dim_in, dim_in, time_emb_dim=dim), Residual(Rezero(LinearAttention(dim_in))), Upsample(dim_in)])) self.final_block = Block(dim, dim) self.final_conv = torch.nn.Conv2d(dim, 1, 1) def forward(self, x, mask, mu, t, spk=None): # x, mu: [B, 80, L], t: [B, ], mask: [B, 1, L] if not isinstance(spk, type(None)): s = self.spk_mlp(spk) t = self.time_pos_emb(t, scale=self.pe_scale) t = self.mlp(t) # [B, 64] s = s.unsqueeze(-1).repeat(1, 1, x.shape[-1]) x = torch.stack([mu, x, s], 1) # [B, 3, 80, L] mask = mask.unsqueeze(1) # [B, 1, 1, L] hiddens = [] masks = [mask] for resnet1, resnet2, attn, downsample in self.downs: mask_down = masks[-1] x = resnet1(x, mask_down, t) # [B, 64, 80, L] x = resnet2(x, mask_down, t) x = attn(x) hiddens.append(x) x = downsample(x * mask_down) masks.append(mask_down[:, :, :, ::2]) masks = masks[:-1] mask_mid = masks[-1] x = self.mid_block1(x, mask_mid, t) x = self.mid_attn(x) x = self.mid_block2(x, mask_mid, t) for resnet1, resnet2, attn, upsample in self.ups: mask_up = masks.pop() x = torch.cat((x, hiddens.pop()), dim=1) x = resnet1(x, mask_up, t) x = resnet2(x, mask_up, t) x = attn(x) x = upsample(x * mask_up) x = self.final_block(x, mask) output = self.final_conv(x * mask) return (output * mask).squeeze(1) def get_noise(t, beta_init, beta_term, cumulative=False): if cumulative: noise = beta_init*t + 0.5*(beta_term - beta_init)*(t**2) else: noise = beta_init + (beta_term - beta_init)*t return noise class Diffusion(BaseModule): def __init__(self, n_feats, dim, spk_emb_dim=64, beta_min=0.05, beta_max=20, pe_scale=1000): super(Diffusion, self).__init__() self.n_feats = n_feats self.dim = dim # self.n_spks = n_spks self.spk_emb_dim = spk_emb_dim self.beta_min = beta_min self.beta_max = beta_max self.pe_scale = pe_scale self.estimator = GradLogPEstimator2d(dim, spk_emb_dim=spk_emb_dim, pe_scale=pe_scale, n_feats=n_feats) def forward_diffusion(self, x0, mask, mu, t): time = t.unsqueeze(-1).unsqueeze(-1) cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True) # it is actually the integral of beta mean = x0*torch.exp(-0.5*cum_noise) + mu*(1.0 - torch.exp(-0.5*cum_noise)) variance = 1.0 - torch.exp(-cum_noise) z = torch.randn(x0.shape, dtype=x0.dtype, device=x0.device, requires_grad=False) xt = mean + z * torch.sqrt(variance) return xt * mask, z * mask @torch.no_grad() def reverse_diffusion(self, z, mask, mu, n_timesteps, stoc=False, spk=None, use_classifier_free=False, classifier_free_guidance=3.0, dummy_spk=None): # emo need to be merged by spk # looks like a plain Euler-Maruyama method h = 1.0 / n_timesteps xt = z * mask for i in range(n_timesteps): t = (1.0 - (i + 0.5)*h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) time = t.unsqueeze(-1).unsqueeze(-1) noise_t = get_noise(time, self.beta_min, self.beta_max, cumulative=False) if not use_classifier_free: if stoc: # adds stochastic term dxt_det = 0.5 * (mu - xt) - self.estimator(xt, mask, mu, t, spk) dxt_det = dxt_det * noise_t * h dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device, requires_grad=False) dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h) dxt = dxt_det + dxt_stoc else: dxt = 0.5 * (mu - xt - self.estimator(xt, mask, mu, t, spk)) dxt = dxt * noise_t * h xt = (xt - dxt) * mask else: if stoc: # adds stochastic term score_estimate = (1 + classifier_free_guidance) * self.estimator(xt, mask, mu, t, spk) \ - classifier_free_guidance * self.estimator(xt, mask, mu, t, dummy_spk) dxt_det = 0.5 * (mu - xt) - score_estimate dxt_det = dxt_det * noise_t * h dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device, requires_grad=False) dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h) dxt = dxt_det + dxt_stoc else: score_estimate = (1 + classifier_free_guidance) * self.estimator(xt, mask, mu, t, spk) \ - classifier_free_guidance * self.estimator(xt, mask, mu, t, dummy_spk) dxt = 0.5 * (mu - xt - score_estimate) dxt = dxt * noise_t * h xt = (xt - dxt) * mask return xt @torch.no_grad() def forward(self, z, mask, mu, n_timesteps, stoc=False, spk=None, use_classifier_free=False, classifier_free_guidance=3.0, dummy_spk=None ): return self.reverse_diffusion(z, mask, mu, n_timesteps, stoc, spk, use_classifier_free, classifier_free_guidance, dummy_spk) def loss_t(self, x0, mask, mu, t, spk=None): xt, z = self.forward_diffusion(x0, mask, mu, t) # z is sampled from N(0, I) time = t.unsqueeze(-1).unsqueeze(-1) cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True) noise_estimation = self.estimator(xt, mask, mu, t, spk) noise_estimation *= torch.sqrt(1.0 - torch.exp(-cum_noise)) # multiply by lambda which is set to be variance # actually multiplied by sqrt(lambda), but not lambda # NOTE: here use a trick to put lambda into L2 norm so that don't divide z with std. loss = torch.sum((noise_estimation + z)**2) / (torch.sum(mask)*self.n_feats) return loss, xt def compute_loss(self, x0, mask, mu, spk=None, offset=1e-5): t = torch.rand(x0.shape[0], dtype=x0.dtype, device=x0.device, requires_grad=False) t = torch.clamp(t, offset, 1.0 - offset) return self.loss_t(x0, mask, mu, t, spk) def classifier_decode(self, z, mask, mu, n_timesteps, stoc=False, spk=None, classifier_func=None, guidance=1.0, control_emo=None, classifier_type="conformer"): # control_emo should be [B, ] tensor h = 1.0 / n_timesteps xt = z * mask for i in range(n_timesteps): t = (1.0 - (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) time = t.unsqueeze(-1).unsqueeze(-1) noise_t = get_noise(time, self.beta_min, self.beta_max, cumulative=False) # =========== classifier part ============== xt = xt.detach() xt.requires_grad_(True) if classifier_type == 'CNN-with-time': logits = classifier_func(xt.transpose(1, 2), mu.transpose(1, 2), (mask == 1.0).squeeze(1), t=t) else: logits = classifier_func(xt.transpose(1, 2), mu.transpose(1, 2), (mask == 1.0).squeeze(1)) if classifier_type == 'conformer': # [B, C] probs = torch.log_softmax(logits, dim=-1) # [B, C] elif classifier_type == 'CNN' or classifier_type == 'CNN-with-time' : probs_every_place = torch.softmax(logits, dim=-1) # [B, T', C] probs_mean = torch.mean(probs_every_place, dim=1) # [B, C] probs = torch.log(probs_mean) else: raise NotImplementedError control_emo_probs = probs[torch.arange(len(control_emo)).to(control_emo.device), control_emo] control_emo_probs.sum().backward(retain_graph=True) # NOTE: sum is to treat all the components as the same weight. xt_grad = xt.grad # ========================================== if stoc: # adds stochastic term dxt_det = 0.5 * (mu - xt) - self.estimator(xt, mask, mu, t, spk) - guidance * xt_grad dxt_det = dxt_det * noise_t * h dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device, requires_grad=False) dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h) dxt = dxt_det + dxt_stoc else: dxt = 0.5 * (mu - xt - self.estimator(xt, mask, mu, t, spk) - guidance * xt_grad) dxt = dxt * noise_t * h xt = (xt - dxt) * mask return xt def classifier_decode_DPS(self, z, mask, mu, n_timesteps, stoc=False, spk=None, classifier_func=None, guidance=1.0, control_emo=None, classifier_type="conformer"): # control_emo should be [B, ] tensor h = 1.0 / n_timesteps xt = z * mask for i in range(n_timesteps): t = (1.0 - (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) time = t.unsqueeze(-1).unsqueeze(-1) noise_t = get_noise(time, self.beta_min, self.beta_max, cumulative=False) beta_integral_t = get_noise(time, self.beta_min, self.beta_max, cumulative=True) bar_alpha_t = math.exp(-beta_integral_t) # =========== classifier part ============== xt = xt.detach() xt.requires_grad_(True) score_estimate = self.estimator(xt, mask, mu, t, spk) x0_hat = (xt + (1-bar_alpha_t) * score_estimate) / math.sqrt(bar_alpha_t) if classifier_type == 'CNN-with-time': raise NotImplementedError else: logits = classifier_func(x0_hat.transpose(1, 2), mu.transpose(1, 2), (mask == 1.0).squeeze(1)) if classifier_type == 'conformer': # [B, C] probs = torch.log_softmax(logits, dim=-1) # [B, C] elif classifier_type == 'CNN': probs_every_place = torch.softmax(logits, dim=-1) # [B, T', C] probs_mean = torch.mean(probs_every_place, dim=1) # [B, C] probs_mean = probs_mean + 10E-10 # NOTE: at the first few steps, x0 may be very large. Then the classifier output logits will also have extreme value range. # probs = torch.log(probs_mean) else: raise NotImplementedError control_emo_probs = probs[torch.arange(len(control_emo)).to(control_emo.device), control_emo] control_emo_probs.sum().backward(retain_graph=True) # NOTE: sum is to treat all the components as the same weight. xt_grad = xt.grad # ========================================== if stoc: # adds stochastic term dxt_det = 0.5 * (mu - xt) - score_estimate - guidance * xt_grad dxt_det = dxt_det * noise_t * h dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device, requires_grad=False) dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h) dxt = dxt_det + dxt_stoc else: dxt = 0.5 * (mu - xt - score_estimate - guidance * xt_grad) dxt = dxt * noise_t * h xt = (xt - dxt) * mask return xt def classifier_decode_mixture(self, z, mask, mu, n_timesteps, stoc=False, spk=None, classifier_func=None, guidance=1.0, control_emo1=None,control_emo2=None, emo1_weight=None, classifier_type="conformer"): # control_emo should be [B, ] tensor h = 1.0 / n_timesteps xt = z * mask for i in range(n_timesteps): t = (1.0 - (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) time = t.unsqueeze(-1).unsqueeze(-1) noise_t = get_noise(time, self.beta_min, self.beta_max, cumulative=False) # =========== classifier part ============== xt = xt.detach() xt.requires_grad_(True) if classifier_type == 'CNN-with-time': logits = classifier_func(xt.transpose(1, 2), mu.transpose(1, 2), (mask == 1.0).squeeze(1), t=t) else: logits = classifier_func(xt.transpose(1, 2), mu.transpose(1, 2), (mask == 1.0).squeeze(1)) if classifier_type == 'conformer': # [B, C] probs = torch.log_softmax(logits, dim=-1) # [B, C] elif classifier_type == 'CNN' or classifier_type == 'CNN-with-time' : probs_every_place = torch.softmax(logits, dim=-1) # [B, T', C] probs_mean = torch.mean(probs_every_place, dim=1) # [B, C] probs = torch.log(probs_mean) else: raise NotImplementedError control_emo_probs1 = probs[torch.arange(len(control_emo1)).to(control_emo1.device), control_emo1] control_emo_probs2 = probs[torch.arange(len(control_emo2)).to(control_emo2.device), control_emo2] control_emo_probs = control_emo_probs1 * emo1_weight + control_emo_probs2 * (1-emo1_weight) # interpolate control_emo_probs.sum().backward(retain_graph=True) # NOTE: sum is to treat all the components as the same weight. xt_grad = xt.grad # ========================================== if stoc: # adds stochastic term dxt_det = 0.5 * (mu - xt) - self.estimator(xt, mask, mu, t, spk) - guidance * xt_grad dxt_det = dxt_det * noise_t * h dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device, requires_grad=False) dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h) dxt = dxt_det + dxt_stoc else: dxt = 0.5 * (mu - xt - self.estimator(xt, mask, mu, t, spk) - guidance * xt_grad) dxt = dxt * noise_t * h xt = (xt - dxt) * mask return xt def classifier_decode_mixture_DPS(self, z, mask, mu, n_timesteps, stoc=False, spk=None, classifier_func=None, guidance=1.0, control_emo1=None,control_emo2=None, emo1_weight=None, classifier_type="conformer"): # control_emo should be [B, ] tensor h = 1.0 / n_timesteps xt = z * mask for i in range(n_timesteps): t = (1.0 - (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) time = t.unsqueeze(-1).unsqueeze(-1) noise_t = get_noise(time, self.beta_min, self.beta_max, cumulative=False) beta_integral_t = get_noise(time, self.beta_min, self.beta_max, cumulative=True) bar_alpha_t = math.exp(-beta_integral_t) # =========== classifier part ============== xt = xt.detach() xt.requires_grad_(True) score_estimate = self.estimator(xt, mask, mu, t, spk) x0_hat = (xt + (1 - bar_alpha_t) * score_estimate) / math.sqrt(bar_alpha_t) if classifier_type == 'CNN-with-time': raise NotImplementedError else: logits = classifier_func(x0_hat.transpose(1, 2), mu.transpose(1, 2), (mask == 1.0).squeeze(1)) if classifier_type == 'conformer': # [B, C] probs = torch.log_softmax(logits, dim=-1) # [B, C] elif classifier_type == 'CNN' or classifier_type == 'CNN-with-time' : probs_every_place = torch.softmax(logits, dim=-1) # [B, T', C] probs_mean = torch.mean(probs_every_place, dim=1) # [B, C] probs_mean = probs_mean + 10E-10 probs = torch.log(probs_mean) else: raise NotImplementedError control_emo_probs1 = probs[torch.arange(len(control_emo1)).to(control_emo1.device), control_emo1] control_emo_probs2 = probs[torch.arange(len(control_emo2)).to(control_emo2.device), control_emo2] control_emo_probs = control_emo_probs1 * emo1_weight + control_emo_probs2 * (1-emo1_weight) # interpolate control_emo_probs.sum().backward(retain_graph=True) # NOTE: sum is to treat all the components as the same weight. xt_grad = xt.grad # ========================================== if stoc: # adds stochastic term dxt_det = 0.5 * (mu - xt) - score_estimate - guidance * xt_grad dxt_det = dxt_det * noise_t * h dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device, requires_grad=False) dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h) dxt = dxt_det + dxt_stoc else: dxt = 0.5 * (mu - xt - score_estimate - guidance * xt_grad) dxt = dxt * noise_t * h xt = (xt - dxt) * mask return xt