Spaces:
Sleeping
Sleeping
| # Copyright (c) 2023 Amphion. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """Adapted from https://github.com/zhenye234/CoMoSpeech""" | |
| import torch | |
| import torch.nn as nn | |
| import copy | |
| import numpy as np | |
| import math | |
| from tqdm.auto import tqdm | |
| from utils.ssim import SSIM | |
| from models.svc.transformer.conformer import Conformer, BaseModule | |
| from models.svc.diffusion.diffusion_wrapper import DiffusionWrapper | |
| from models.svc.comosvc.utils import slice_segments, rand_ids_segments | |
| class Consistency(nn.Module): | |
| def __init__(self, cfg, distill=False): | |
| super().__init__() | |
| self.cfg = cfg | |
| # self.denoise_fn = GradLogPEstimator2d(96) | |
| self.denoise_fn = DiffusionWrapper(self.cfg) | |
| self.cfg = cfg.model.comosvc | |
| self.teacher = not distill | |
| self.P_mean = self.cfg.P_mean | |
| self.P_std = self.cfg.P_std | |
| self.sigma_data = self.cfg.sigma_data | |
| self.sigma_min = self.cfg.sigma_min | |
| self.sigma_max = self.cfg.sigma_max | |
| self.rho = self.cfg.rho | |
| self.N = self.cfg.n_timesteps | |
| self.ssim_loss = SSIM() | |
| # Time step discretization | |
| step_indices = torch.arange(self.N) | |
| # karras boundaries formula | |
| t_steps = ( | |
| self.sigma_min ** (1 / self.rho) | |
| + step_indices | |
| / (self.N - 1) | |
| * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** (1 / self.rho)) | |
| ) ** self.rho | |
| self.t_steps = torch.cat( | |
| [torch.zeros_like(t_steps[:1]), self.round_sigma(t_steps)] | |
| ) | |
| def init_consistency_training(self): | |
| self.denoise_fn_ema = copy.deepcopy(self.denoise_fn) | |
| self.denoise_fn_pretrained = copy.deepcopy(self.denoise_fn) | |
| def EDMPrecond(self, x, sigma, cond, denoise_fn, mask, spk=None): | |
| """ | |
| karras diffusion reverse process | |
| Args: | |
| x: noisy mel-spectrogram [B x n_mel x L] | |
| sigma: noise level [B x 1 x 1] | |
| cond: output of conformer encoder [B x n_mel x L] | |
| denoise_fn: denoiser neural network e.g. DilatedCNN | |
| mask: mask of padded frames [B x n_mel x L] | |
| Returns: | |
| denoised mel-spectrogram [B x n_mel x L] | |
| """ | |
| sigma = sigma.reshape(-1, 1, 1) | |
| c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2) | |
| c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2).sqrt() | |
| c_in = 1 / (self.sigma_data**2 + sigma**2).sqrt() | |
| c_noise = sigma.log() / 4 | |
| x_in = c_in * x | |
| x_in = x_in.transpose(1, 2) | |
| x = x.transpose(1, 2) | |
| cond = cond.transpose(1, 2) | |
| F_x = denoise_fn(x_in, c_noise.squeeze(), cond) | |
| # F_x = denoise_fn((c_in * x), mask, cond, c_noise.flatten()) | |
| D_x = c_skip * x + c_out * (F_x) | |
| D_x = D_x.transpose(1, 2) | |
| return D_x | |
| def EDMLoss(self, x_start, cond, mask): | |
| """ | |
| compute loss for EDM model | |
| Args: | |
| x_start: ground truth mel-spectrogram [B x n_mel x L] | |
| cond: output of conformer encoder [B x n_mel x L] | |
| mask: mask of padded frames [B x n_mel x L] | |
| """ | |
| rnd_normal = torch.randn([x_start.shape[0], 1, 1], device=x_start.device) | |
| sigma = (rnd_normal * self.P_std + self.P_mean).exp() | |
| weight = (sigma**2 + self.sigma_data**2) / (sigma * self.sigma_data) ** 2 | |
| # follow Grad-TTS, start from Gaussian noise with mean cond and std I | |
| noise = (torch.randn_like(x_start) + cond) * sigma | |
| D_yn = self.EDMPrecond(x_start + noise, sigma, cond, self.denoise_fn, mask) | |
| loss = weight * ((D_yn - x_start) ** 2) | |
| loss = torch.sum(loss * mask) / torch.sum(mask) | |
| return loss | |
| def round_sigma(self, sigma): | |
| return torch.as_tensor(sigma) | |
| def edm_sampler( | |
| self, | |
| latents, | |
| cond, | |
| nonpadding, | |
| num_steps=50, | |
| sigma_min=0.002, | |
| sigma_max=80, | |
| rho=7, | |
| S_churn=0, | |
| S_min=0, | |
| S_max=float("inf"), | |
| S_noise=1, | |
| # S_churn=40 ,S_min=0.05,S_max=50,S_noise=1.003,# S_churn=0, S_min=0, S_max=float('inf'), S_noise=1, | |
| # S_churn=30 ,S_min=0.01,S_max=30,S_noise=1.007, | |
| # S_churn=30 ,S_min=0.01,S_max=1,S_noise=1.007, | |
| # S_churn=80 ,S_min=0.05,S_max=50,S_noise=1.003, | |
| ): | |
| """ | |
| karras diffusion sampler | |
| Args: | |
| latents: noisy mel-spectrogram [B x n_mel x L] | |
| cond: output of conformer encoder [B x n_mel x L] | |
| nonpadding: mask of padded frames [B x n_mel x L] | |
| num_steps: number of steps for diffusion inference | |
| Returns: | |
| denoised mel-spectrogram [B x n_mel x L] | |
| """ | |
| # Time step discretization. | |
| step_indices = torch.arange(num_steps, device=latents.device) | |
| num_steps = num_steps + 1 | |
| t_steps = ( | |
| sigma_max ** (1 / rho) | |
| + step_indices | |
| / (num_steps - 1) | |
| * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho)) | |
| ) ** rho | |
| t_steps = torch.cat([self.round_sigma(t_steps), torch.zeros_like(t_steps[:1])]) | |
| # Main sampling loop. | |
| x_next = latents * t_steps[0] | |
| # wrap in tqdm for progress bar | |
| bar = tqdm(enumerate(zip(t_steps[:-1], t_steps[1:]))) | |
| for i, (t_cur, t_next) in bar: | |
| x_cur = x_next | |
| # Increase noise temporarily. | |
| gamma = ( | |
| min(S_churn / num_steps, np.sqrt(2) - 1) | |
| if S_min <= t_cur <= S_max | |
| else 0 | |
| ) | |
| t_hat = self.round_sigma(t_cur + gamma * t_cur) | |
| t = torch.zeros((x_cur.shape[0], 1, 1), device=x_cur.device) | |
| t[:, 0, 0] = t_hat | |
| t_hat = t | |
| x_hat = x_cur + ( | |
| t_hat**2 - t_cur**2 | |
| ).sqrt() * S_noise * torch.randn_like(x_cur) | |
| # Euler step. | |
| denoised = self.EDMPrecond(x_hat, t_hat, cond, self.denoise_fn, nonpadding) | |
| d_cur = (x_hat - denoised) / t_hat | |
| x_next = x_hat + (t_next - t_hat) * d_cur | |
| return x_next | |
| def CTLoss_D(self, y, cond, mask): | |
| """ | |
| compute loss for consistency distillation | |
| Args: | |
| y: ground truth mel-spectrogram [B x n_mel x L] | |
| cond: output of conformer encoder [B x n_mel x L] | |
| mask: mask of padded frames [B x n_mel x L] | |
| """ | |
| with torch.no_grad(): | |
| mu = 0.95 | |
| for p, ema_p in zip( | |
| self.denoise_fn.parameters(), self.denoise_fn_ema.parameters() | |
| ): | |
| ema_p.mul_(mu).add_(p, alpha=1 - mu) | |
| n = torch.randint(1, self.N, (y.shape[0],)) | |
| z = torch.randn_like(y) + cond | |
| tn_1 = self.t_steps[n + 1].reshape(-1, 1, 1).to(y.device) | |
| f_theta = self.EDMPrecond(y + tn_1 * z, tn_1, cond, self.denoise_fn, mask) | |
| with torch.no_grad(): | |
| tn = self.t_steps[n].reshape(-1, 1, 1).to(y.device) | |
| # euler step | |
| x_hat = y + tn_1 * z | |
| denoised = self.EDMPrecond( | |
| x_hat, tn_1, cond, self.denoise_fn_pretrained, mask | |
| ) | |
| d_cur = (x_hat - denoised) / tn_1 | |
| y_tn = x_hat + (tn - tn_1) * d_cur | |
| f_theta_ema = self.EDMPrecond(y_tn, tn, cond, self.denoise_fn_ema, mask) | |
| # loss = (f_theta - f_theta_ema.detach()) ** 2 | |
| # loss = torch.sum(loss * mask) / torch.sum(mask) | |
| loss = self.ssim_loss(f_theta, f_theta_ema.detach()) | |
| loss = torch.sum(loss * mask) / torch.sum(mask) | |
| return loss | |
| def get_t_steps(self, N): | |
| N = N + 1 | |
| step_indices = torch.arange(N) # , device=latents.device) | |
| t_steps = ( | |
| self.sigma_min ** (1 / self.rho) | |
| + step_indices | |
| / (N - 1) | |
| * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** (1 / self.rho)) | |
| ) ** self.rho | |
| return t_steps.flip(0) | |
| def CT_sampler(self, latents, cond, nonpadding, t_steps=1): | |
| """ | |
| consistency distillation sampler | |
| Args: | |
| latents: noisy mel-spectrogram [B x n_mel x L] | |
| cond: output of conformer encoder [B x n_mel x L] | |
| nonpadding: mask of padded frames [B x n_mel x L] | |
| t_steps: number of steps for diffusion inference | |
| Returns: | |
| denoised mel-spectrogram [B x n_mel x L] | |
| """ | |
| # one-step | |
| if t_steps == 1: | |
| t_steps = [80] | |
| # multi-step | |
| else: | |
| t_steps = self.get_t_steps(t_steps) | |
| t_steps = torch.as_tensor(t_steps).to(latents.device) | |
| latents = latents * t_steps[0] | |
| _t = torch.zeros((latents.shape[0], 1, 1), device=latents.device) | |
| _t[:, 0, 0] = t_steps | |
| x = self.EDMPrecond(latents, _t, cond, self.denoise_fn_ema, nonpadding) | |
| for t in t_steps[1:-1]: | |
| z = torch.randn_like(x) + cond | |
| x_tn = x + (t**2 - self.sigma_min**2).sqrt() * z | |
| _t = torch.zeros((x.shape[0], 1, 1), device=x.device) | |
| _t[:, 0, 0] = t | |
| t = _t | |
| print(t) | |
| x = self.EDMPrecond(x_tn, t, cond, self.denoise_fn_ema, nonpadding) | |
| return x | |
| def forward(self, x, nonpadding, cond, t_steps=1, infer=False): | |
| """ | |
| calculate loss or sample mel-spectrogram | |
| Args: | |
| x: | |
| training: ground truth mel-spectrogram [B x n_mel x L] | |
| inference: output of encoder [B x n_mel x L] | |
| """ | |
| if self.teacher: # teacher model -- karras diffusion | |
| if not infer: | |
| loss = self.EDMLoss(x, cond, nonpadding) | |
| return loss | |
| else: | |
| shape = (cond.shape[0], self.cfg.n_mel, cond.shape[2]) | |
| x = torch.randn(shape, device=x.device) + cond | |
| x = self.edm_sampler(x, cond, nonpadding, t_steps) | |
| return x | |
| else: # Consistency distillation | |
| if not infer: | |
| loss = self.CTLoss_D(x, cond, nonpadding) | |
| return loss | |
| else: | |
| shape = (cond.shape[0], self.cfg.n_mel, cond.shape[2]) | |
| x = torch.randn(shape, device=x.device) + cond | |
| x = self.CT_sampler(x, cond, nonpadding, t_steps=1) | |
| return x | |
| class ComoSVC(BaseModule): | |
| def __init__(self, cfg): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.cfg.model.comosvc.n_mel = self.cfg.preprocess.n_mel | |
| self.distill = self.cfg.model.comosvc.distill | |
| self.encoder = Conformer(self.cfg.model.comosvc) | |
| self.decoder = Consistency(self.cfg, distill=self.distill) | |
| self.ssim_loss = SSIM() | |
| def forward(self, x_mask, x, n_timesteps, temperature=1.0): | |
| """ | |
| Generates mel-spectrogram from pitch, content vector, energy. Returns: | |
| 1. encoder outputs (from conformer) | |
| 2. decoder outputs (from diffusion-based decoder) | |
| Args: | |
| x_mask : mask of padded frames in mel-spectrogram. [B x L x n_mel] | |
| x : output of encoder framework. [B x L x d_condition] | |
| n_timesteps : number of steps to use for reverse diffusion in decoder. | |
| temperature : controls variance of terminal distribution. | |
| """ | |
| # Get encoder_outputs `mu_x` | |
| mu_x = self.encoder(x, x_mask) | |
| encoder_outputs = mu_x | |
| mu_x = mu_x.transpose(1, 2) | |
| x_mask = x_mask.transpose(1, 2) | |
| # Generate sample by performing reverse dynamics | |
| decoder_outputs = self.decoder( | |
| mu_x, x_mask, mu_x, t_steps=n_timesteps, infer=True | |
| ) | |
| decoder_outputs = decoder_outputs.transpose(1, 2) | |
| return encoder_outputs, decoder_outputs | |
| def compute_loss(self, x_mask, x, mel, out_size=None, skip_diff=False): | |
| """ | |
| Computes 2 losses: | |
| 1. prior loss: loss between mel-spectrogram and encoder outputs. | |
| 2. diffusion loss: loss between gaussian noise and its reconstruction by diffusion-based decoder. | |
| Args: | |
| x_mask : mask of padded frames in mel-spectrogram. [B x L x n_mel] | |
| x : output of encoder framework. [B x L x d_condition] | |
| mel : ground truth mel-spectrogram. [B x L x n_mel] | |
| """ | |
| mu_x = self.encoder(x, x_mask) | |
| # prior loss | |
| prior_loss = torch.sum( | |
| 0.5 * ((mel - mu_x) ** 2 + math.log(2 * math.pi)) * x_mask | |
| ) | |
| prior_loss = prior_loss / (torch.sum(x_mask) * self.cfg.model.comosvc.n_mel) | |
| # ssim loss | |
| ssim_loss = self.ssim_loss(mu_x, mel) | |
| ssim_loss = torch.sum(ssim_loss * x_mask) / torch.sum(x_mask) | |
| x_mask = x_mask.transpose(1, 2) | |
| mu_x = mu_x.transpose(1, 2) | |
| mel = mel.transpose(1, 2) | |
| if not self.distill and skip_diff: | |
| diff_loss = prior_loss.clone() | |
| diff_loss.fill_(0) | |
| # Cut a small segment of mel-spectrogram in order to increase batch size | |
| else: | |
| if self.distill: | |
| mu_y = mu_x.detach() | |
| else: | |
| mu_y = mu_x | |
| mask_y = x_mask | |
| diff_loss = self.decoder(mel, mask_y, mu_y, infer=False) | |
| return ssim_loss, prior_loss, diff_loss | |