import torch from tqdm import tqdm class CFM(torch.nn.Module): def __init__( self, estimator: torch.nn.Module, ): super().__init__() self.sigma_min = 1e-6 self.estimator = estimator self.in_channels = estimator.in_channels self.criterion = torch.nn.L1Loss() @torch.inference_mode() def inference(self, mu: torch.Tensor, x_lens: torch.Tensor, prompt: torch.Tensor, style: torch.Tensor, n_timesteps=10, temperature=1.0, inference_cfg_rate=[0.5, 0.5], random_voice=False, ): """Forward diffusion Args: mu (torch.Tensor): output of encoder shape: (batch_size, n_feats, mel_timesteps) x_lens (torch.Tensor): length of each mel-spectrogram shape: (batch_size,) prompt (torch.Tensor): prompt shape: (batch_size, n_feats, prompt_len) style (torch.Tensor): style shape: (batch_size, style_dim) n_timesteps (int): number of diffusion steps temperature (float, optional): temperature for scaling noise. Defaults to 1.0. inference_cfg_rate (float, optional): Classifier-Free Guidance inference introduced in VoiceBox. Defaults to 0.5. Returns: sample: generated mel-spectrogram shape: (batch_size, n_feats, mel_timesteps) """ B, T = mu.size(0), mu.size(1) z = torch.randn([B, self.in_channels, T], device=mu.device) * temperature t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device) t_span = t_span + (-1) * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span) return self.solve_euler(z, x_lens, prompt, mu, style, t_span, inference_cfg_rate, random_voice) def solve_euler(self, x, x_lens, prompt, mu, style, t_span, inference_cfg_rate=[0.5, 0.5], random_voice=False,): """ Fixed euler solver for ODEs. Args: x (torch.Tensor): random noise t_span (torch.Tensor): n_timesteps interpolated shape: (n_timesteps + 1,) mu (torch.Tensor): output of encoder shape: (batch_size, n_feats, mel_timesteps) x_lens (torch.Tensor): length of each mel-spectrogram shape: (batch_size,) prompt (torch.Tensor): prompt shape: (batch_size, n_feats, prompt_len) style (torch.Tensor): style shape: (batch_size, style_dim) inference_cfg_rate (float, optional): Classifier-Free Guidance inference introduced in VoiceBox. Defaults to 0.5. sway_sampling (bool, optional): Sway sampling. Defaults to False. amo_sampling (bool, optional): AMO sampling. Defaults to False. """ t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] # apply prompt prompt_len = prompt.size(-1) prompt_x = torch.zeros_like(x) prompt_x[..., :prompt_len] = prompt[..., :prompt_len] x[..., :prompt_len] = 0 for step in tqdm(range(1, len(t_span))): if random_voice: cfg_dphi_dt = self.estimator( torch.cat([x, x], dim=0), torch.cat([torch.zeros_like(prompt_x), torch.zeros_like(prompt_x)], dim=0), torch.cat([x_lens, x_lens], dim=0), torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0), torch.cat([torch.zeros_like(style), torch.zeros_like(style)], dim=0), torch.cat([mu, torch.zeros_like(mu)], dim=0), ) cond_txt, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2] dphi_dt = ((1.0 + inference_cfg_rate[0]) * cond_txt - inference_cfg_rate[0] * uncond) elif all(i == 0 for i in inference_cfg_rate): dphi_dt = self.estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu) elif inference_cfg_rate[0] == 0: # Classifier-Free Guidance inference introduced in VoiceBox cfg_dphi_dt = self.estimator( torch.cat([x, x], dim=0), torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0), torch.cat([x_lens, x_lens], dim=0), torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0), torch.cat([style, torch.zeros_like(style)], dim=0), torch.cat([mu, mu], dim=0), ) cond_txt_spk, cond_txt = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2] dphi_dt = ((1.0 + inference_cfg_rate[1]) * cond_txt_spk - inference_cfg_rate[1] * cond_txt) elif inference_cfg_rate[1] == 0: cfg_dphi_dt = self.estimator( torch.cat([x, x], dim=0), torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0), torch.cat([x_lens, x_lens], dim=0), torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0), torch.cat([style, torch.zeros_like(style)], dim=0), torch.cat([mu, torch.zeros_like(mu)], dim=0), ) cond_txt_spk, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2] dphi_dt = ((1.0 + inference_cfg_rate[0]) * cond_txt_spk - inference_cfg_rate[0] * uncond) else: # Multi-condition Classifier-Free Guidance inference introduced in MegaTTS3 cfg_dphi_dt = self.estimator( torch.cat([x, x, x], dim=0), torch.cat([prompt_x, torch.zeros_like(prompt_x), torch.zeros_like(prompt_x)], dim=0), torch.cat([x_lens, x_lens, x_lens], dim=0), torch.cat([t.unsqueeze(0), t.unsqueeze(0), t.unsqueeze(0)], dim=0), torch.cat([style, torch.zeros_like(style), torch.zeros_like(style)], dim=0), torch.cat([mu, mu, torch.zeros_like(mu)], dim=0), ) cond_txt_spk, cond_txt, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2], cfg_dphi_dt[2:3] dphi_dt = (1.0 + inference_cfg_rate[0] + inference_cfg_rate[1]) * cond_txt_spk - \ inference_cfg_rate[0] * uncond - inference_cfg_rate[1] * cond_txt x = x + dt * dphi_dt t = t + dt if step < len(t_span) - 1: dt = t_span[step + 1] - t x[:, :, :prompt_len] = 0 return x def forward(self, x1, x_lens, prompt_lens, mu, style): """Computes diffusion loss Args: x1 (torch.Tensor): Target shape: (batch_size, n_feats, mel_timesteps) mask (torch.Tensor): target mask shape: (batch_size, 1, mel_timesteps) mu (torch.Tensor): output of encoder shape: (batch_size, n_feats, mel_timesteps) spks (torch.Tensor, optional): speaker embedding. Defaults to None. shape: (batch_size, spk_emb_dim) Returns: loss: conditional flow matching loss y: conditional flow shape: (batch_size, n_feats, mel_timesteps) """ b, _, t = x1.shape # random timestep t = torch.rand([b, 1, 1], device=mu.device, dtype=x1.dtype) # sample noise p(x_0) z = torch.randn_like(x1) y = (1 - (1 - self.sigma_min) * t) * z + t * x1 u = x1 - (1 - self.sigma_min) * z prompt = torch.zeros_like(x1) for bib in range(b): prompt[bib, :, :prompt_lens[bib]] = x1[bib, :, :prompt_lens[bib]] # range covered by prompt are set to 0 y[bib, :, :prompt_lens[bib]] = 0 estimator_out = self.estimator(y, prompt, x_lens, t.squeeze(), style, mu) loss = 0 for bib in range(b): loss += self.criterion(estimator_out[bib, :, prompt_lens[bib]:x_lens[bib]], u[bib, :, prompt_lens[bib]:x_lens[bib]]) loss /= b return loss