    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,
    ):
        """
        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.

        num_steps = num_steps + 1
        step_indices = torch.arange(num_steps, device=latents.device)
        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)
            d_cur = (x_hat - denoised) / t_hat
            x_next = x_hat + (t_next - t_hat) * d_cur

            # add Heun’s 2nd order method
            # if i < num_steps - 1:
            #     t = torch.zeros((x_cur.shape[0], 1, 1), device=x_cur.device)
            #     t[:, 0, 0] = t_next
            #     #t_next = t
            #     denoised = self.EDMPrecond(x_next, t, cond, self.denoise_fn, nonpadding)
            #     d_prime = (x_next - denoised) / t_next
            #     x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime)

        return x_next
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]
models/svc/comosvc/comosvc.py