uni_pc.py

from dpm_solver_v3 import NoiseScheduleVP, model_wrapper
import torch
import torch.nn.functional as F
import math
import numpy as np
import os

class UniPC:
    def __init__(
        self,
        noise_schedule,
        steps=10,
        t_start=None,
        t_end=None,
        skip_type="customed_time_karras",
        degenerated=False,
        use_afs = False,
        denoise_to_zero=False,
        need_fp16_discrete_method = False,
        ultilize_vae_in_fp16 = False,
        is_high_resoulution = True,
        device="cuda",
    ):
        self.device = device
        self.model = None
        self.noise_schedule = noise_schedule
        self.steps = steps if not use_afs else steps + 1
        self.use_afs = use_afs
        self.ultilize_vae_in_fp16 = ultilize_vae_in_fp16
        self.need_fp16_discrete_method = need_fp16_discrete_method
        t_0 = 1.0 / self.noise_schedule.total_N if t_end is None else t_end
        t_T = self.noise_schedule.T if t_start is None else t_start
        self.is_high_resolution = is_high_resoulution
        assert (
            t_0 > 0 and t_T > 0
        ), "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"

        
        # precompute timesteps
        if skip_type == "logSNR" or skip_type == "time_uniform" or skip_type == "time_quadratic" or skip_type == "customed_time_karras":
            self.timesteps = self.get_time_steps(skip_type
                                                 , t_T=t_T
                                                 , t_0=t_0
                                                 , N=steps
                                                 , device=device,denoise_to_zero=denoise_to_zero
                                                 , is_high_resolution=self.is_high_resolution)
        else:
            raise ValueError(f"Unsupported timestep strategy {skip_type}")
        self.lambda_T = self.timesteps[0].cpu().item()
        self.lambda_0 = self.timesteps[-1].cpu().item()

        # print("Time steps", self.timesteps)
        # print("LogSNR steps", self.noise_schedule.marginal_lambda(self.timesteps))

        # store high-order exponential coefficients (lazy)
        self.exp_coeffs = {}

    def noise_prediction_fn(self, x, t):
        """
        Return the noise prediction model.
        """
        return self.model(x, t)

    def append_zero(self, x):
        return torch.cat([x, x.new_zeros([1])])

    def get_sigmas_karras(self, n, sigma_min, sigma_max, rho=7., device='cpu', need_append_zero=True):
        """Constructs the noise schedule of Karras et al. (2022)."""
        ramp = torch.linspace(0, 1, n)
        min_inv_rho = sigma_min ** (1 / rho)
        max_inv_rho = sigma_max ** (1 / rho)
        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
        return self.append_zero(sigmas).to(device) if need_append_zero else sigmas.to(device)
    
    def sigma_to_t(self, sigma, quantize=None):
        quantize = False
        log_sigma = sigma.log()
        dists = log_sigma - self.noise_schedule.log_sigmas[:, None]
        if quantize:
            return dists.abs().argmin(dim=0).view(sigma.shape)
        low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.noise_schedule.log_sigmas.shape[0] - 2)
        high_idx = low_idx + 1
        low, high = self.noise_schedule.log_sigmas[low_idx], self.noise_schedule.log_sigmas[high_idx]
        w = (low - log_sigma) / (low - high)
        w = w.clamp(0, 1)
        t = (1 - w) * low_idx + w * high_idx
        return t.view(sigma.shape)

    def get_time_steps(self, skip_type, t_T, t_0, N, device, denoise_to_zero=False, is_high_resolution=True):
        """Compute the intermediate time steps for sampling.

        Args:
            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
                - 'logSNR': uniform logSNR for the time steps.
                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
            t_T: A `float`. The starting time of the sampling (default is T).
            t_0: A `float`. The ending time of the sampling (default is epsilon).
            N: A `int`. The total number of the spacing of the time steps.
            device: A torch device.
        Returns:
            A pytorch tensor of the time steps, with the shape (N + 1,).
        """
        if skip_type == "logSNR":
            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
            return self.noise_schedule.inverse_lambda(logSNR_steps)
        elif skip_type == "time_uniform":
            return torch.linspace(t_T, t_0, N + 1).to(device)
        elif skip_type == "time_quadratic":
            t_order = 2
            t = torch.linspace(t_T ** (1.0 / t_order), t_0 ** (1.0 / t_order), N + 1).pow(t_order).to(device)
            return t
        elif skip_type == "customed_time_karras" and is_high_resolution:
            sigma_T = self.noise_schedule.sigmas[-1].cpu().item()
            sigma_0 = self.noise_schedule.sigmas[0].cpu().item()
            if N == 8:
                sigmas = self.get_sigmas_karras(12, sigma_0, sigma_T,rho=12.0, device=device)
                if not self.need_fp16_discrete_method:
                    ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[10])
                    ct = self.get_sigmas_karras(9, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                    sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                    real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
                else:
                    sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0, device=device)
                    ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                    ct = self.get_sigmas_karras(8, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                    sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                    tmp_t = [self.noise_schedule.sigma_to_t(sigma).to('cpu') for sigma in sigmas_ct]
                    real_ct = [ t / 999 for t in tmp_t]
            elif N == 5:
                sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0, device=device)
                if not self.need_fp16_discrete_method:
                    sigmas = self.get_sigmas_karras(12, sigma_0, sigma_T,rho=12.0, device=device)
                    ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[9])
                    ct = self.get_sigmas_karras(6, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                    sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                    real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
                else:
                    ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                    ct = self.get_sigmas_karras(5, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                    sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                    real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
            elif N == 6:
                sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0, device=device)
                if not self.need_fp16_discrete_method:
                    sigmas = self.get_sigmas_karras(12, sigma_0, sigma_T,rho=12.0, device=device)
                    ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[10])
                    ct = self.get_sigmas_karras(7, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                    sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                    real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
                else:
                    if denoise_to_zero:
                        ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                        ct = self.get_sigmas_karras(6, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                        sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                        real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
                        real_ct.append(torch.tensor(t_0).to(dtype=real_ct[-1].dtype,device='cpu'))
                    else:
                        sigmas = self.get_sigmas_karras(12, sigma_0, sigma_T, rho=7.0, device=device)
                        ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[7])
                        ct = self.get_sigmas_karras(7, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                        sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                        real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
            elif N == 7:
                sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0, device=device)
                if not self.need_fp16_discrete_method:
                    ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                    ct = self.get_sigmas_karras(8, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                    sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                    real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
                else:
                    if denoise_to_zero:
                        ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                        ct = self.get_sigmas_karras(7, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                        sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                        real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
                        real_ct.append(torch.tensor(t_0).to(dtype=real_ct[-1].dtype,device='cpu'))
            # if denoise_to_zero:
            #     real_ct.append(torch.tensor(t_0).to(dtype=real_ct[-1].dtype,device='cpu'))
            
            if self.use_afs:
                tmp_t = (real_ct[0] + real_ct[1]) / 2
                real_ct.insert(1, tmp_t)
            none_k_ct = torch.from_numpy(np.array(real_ct)).to(device) 
            return none_k_ct#real_ct
        elif skip_type == "customed_time_karras" and not is_high_resolution:
            sigma_T = self.noise_schedule.sigmas[-1].cpu().item()
            sigma_0 = self.noise_schedule.sigmas[0].cpu().item()
            if N == 8:
                sigmas = self.get_sigmas_karras(12, sigma_0, sigma_T, rho=7.0, device=device)
                ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[9])
                ct = self.get_sigmas_karras(9, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
            elif N == 5:
                sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0, device=device)
                ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                ct = self.get_sigmas_karras(6, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
            elif N == 6:
                sigmas = self.sigmas = self.get_sigmas_karras(8, sigma_0, sigma_T, rho=5.0, device=device)
                ct_start, ct_end = self.noise_schedule.sigma_to_t(sigmas[0]), self.sigma_to_t(sigmas[6])
                ct = self.get_sigmas_karras(7, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
                sigmas_ct = self.noise_schedule.get_special_sigmas_with_timesteps(ct).to(device=device)
                real_ct = [self.noise_schedule.sigma_to_t(sigma).to('cpu') / 999 for sigma in sigmas_ct]
            none_k_ct = torch.from_numpy(np.array(real_ct)).to(device) 
            return none_k_ct#real_ct
        else:
            raise ValueError(
                "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type)
            )


    def multistep_uni_pc_update(self, x, model_prev_list:list, t_prev_list: list, t, order, **kwargs):
        if len(model_prev_list) == 0 or len(t_prev_list) == 0:
            return None, None
        if len(t.shape) == 0:
            t = t.view(-1)
        if True:#'bh' in self.variant:
            return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
        else:
            # assert self.variant == 'vary_coeff'
            return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)

    def multistep_uni_pc_sde_update(self, x, model_prev_list:list, t_prev_list: list, t, order, level = 1.0, **kwargs):
        if len(model_prev_list) == 0 or len(t_prev_list) == 0:
            return None, None
        if len(t.shape) == 0:
            t = t.view(-1)
        if True:#'bh' in self.variant:
            return self.multistep_uni_pc_bh_sde_update(x, model_prev_list, t_prev_list, t, level=level, order= order, **kwargs)
        else:
            # assert self.variant == 'vary_coeff'
            return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)

    def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
        # print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
        ns = self.noise_schedule
        assert order <= len(model_prev_list)
        dims = x.dim()

        # first compute rks
        t_prev_0 = t_prev_list[-1]
        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
        lambda_t = ns.marginal_lambda(t)
        model_prev_0 = model_prev_list[-1]
        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
        alpha_t = torch.exp(log_alpha_t)

        h = lambda_t - lambda_prev_0

        rks = []
        D1s = []
        for i in range(1, order):
            t_prev_i = t_prev_list[-(i + 1)]
            model_prev_i = model_prev_list[-(i + 1)]
            lambda_prev_i = ns.marginal_lambda(t_prev_i)
            rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
            rks.append(rk)
            D1s.append((model_prev_i - model_prev_0) / rk)

        rks.append(1.)
        rks = torch.tensor(rks, device=x.device)

        R = []
        b = []

        hh = h[0]
        h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
        h_phi_k = h_phi_1 / hh - 1

        factorial_i = 1

        if True:
            B_h = hh
        else:
            B_h = torch.expm1(hh)
            
        for i in range(1, order + 1):
            R.append(torch.pow(rks, i - 1))
            b.append(h_phi_k * factorial_i / B_h)
            factorial_i *= (i + 1)
            h_phi_k = h_phi_k / hh - 1 / factorial_i 

        R = torch.stack(R)
        b = torch.tensor(b, device=x.device)

        # now predictor
        use_predictor = len(D1s) > 0 and x_t is None
        if len(D1s) > 0:
            D1s = torch.stack(D1s, dim=1) # (B, K)
            if x_t is None:
                # for order 2, we use a simplified version
                if order == 2:
                    rhos_p = torch.tensor([0.5], device=b.device)
                else:
                    rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
        else:
            D1s = None

        if use_corrector:
            # print('using corrector')
            # for order 1, we use a simplified version
            if order == 1:
                rhos_c = torch.tensor([0.5], device=b.device)
            else:
                rhos_c = torch.linalg.solve(R, b)

        model_t = None
        
        x_t_ = (
            expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
            - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
        )
        if x_t is None:
            if use_predictor:
                pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
            else:
                pred_res = 0
            x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * pred_res

        if use_corrector:
            model_t = self.noise_prediction_fn(x_t, t)
            if D1s is not None:
                corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
            else:
                corr_res = 0
            D1_t = (model_t - model_prev_0)
            x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)

        return x_t, model_t
    
    def multistep_uni_pc_bh_sde_update(self, x, model_prev_list, t_prev_list, t, order, level = 0, x_t=None, use_corrector=True):
        # print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
        ns = self.noise_schedule
        assert order <= len(model_prev_list)
        dims = x.dim()

        # first compute rks
        t_prev_0 = t_prev_list[-1]
        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
        lambda_t = ns.marginal_lambda(t)
        model_prev_0 = model_prev_list[-1]
        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
        alpha_t = torch.exp(log_alpha_t)

        h = lambda_t - lambda_prev_0
        z = torch.randn(x.shape, device=self.device)
        z = sigma_t * torch.sqrt(torch.expm1(2.0 * h[0])) * z

        rks = []
        D1s = []
        for i in range(1, order):
            t_prev_i = t_prev_list[-(i + 1)]
            model_prev_i = model_prev_list[-(i + 1)]
            lambda_prev_i = ns.marginal_lambda(t_prev_i)
            rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
            rks.append(rk)
            D1s.append((model_prev_i - model_prev_0) / rk)

        rks.append(1.)
        rks = torch.tensor(rks, device=x.device)

        R = []
        b = []

        hh = h[0]
        h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
        h_phi_k = h_phi_1 / hh - 1

        factorial_i = 1

        if True:
            B_h = hh
        else:
            B_h = torch.expm1(hh)
            
        for i in range(1, order + 1):
            R.append(torch.pow(rks, i - 1))
            b.append(h_phi_k * factorial_i / B_h)
            factorial_i *= (i + 1)
            h_phi_k = h_phi_k / hh - 1 / factorial_i 

        R = torch.stack(R)
        b = torch.tensor(b, device=x.device)

        # now predictor
        use_predictor = len(D1s) > 0 and x_t is None
        if len(D1s) > 0:
            D1s = torch.stack(D1s, dim=1) # (B, K)
            if x_t is None:
                # for order 2, we use a simplified version
                if order == 2:
                    rhos_p = torch.tensor([0.5], device=b.device)
                else:
                    rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
        else:
            D1s = None

        if use_corrector:
            # print('using corrector')
            # for order 1, we use a simplified version
            if order == 1:
                rhos_c = torch.tensor([0.5], device=b.device)
            else:
                rhos_c = torch.linalg.solve(R, b)

        model_t = None
        
        x_t_ = (
            expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
            - expand_dims(sigma_t * h_phi_1, dims) * (1 + level) * model_prev_0
        )
        if x_t is None:
            if use_predictor:
                pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
            else:
                pred_res = 0
            
            x_t_p = (
                expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
                - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
            )
            x_t = x_t_p - expand_dims(sigma_t * B_h, dims) * pred_res

        if use_corrector:
            model_t = self.noise_prediction_fn(x_t, t)
            if D1s is not None:
                corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
            else:
                corr_res = 0
            D1_t = (model_t - model_prev_0)
            x_t = x_t_ - (1 + level)  * expand_dims(sigma_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t) + z * level

        return x_t, model_t
    
    
    def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
        # print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
        ns = self.noise_schedule
        assert order <= len(model_prev_list)
        dims = x.dim()
        # first compute rks
        t_prev_0 = t_prev_list[-1]
        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
        lambda_t = ns.marginal_lambda(t)
        model_prev_0 = model_prev_list[-1]
        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
        log_alpha_t = ns.marginal_log_mean_coeff(t)
        alpha_t = torch.exp(log_alpha_t)

        h = lambda_t - lambda_prev_0

        rks = []
        D1s = []
        for i in range(1, order):
            t_prev_i = t_prev_list[-(i + 1)]
            model_prev_i = model_prev_list[-(i + 1)]
            lambda_prev_i = ns.marginal_lambda(t_prev_i)
            rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
            rks.append(rk)
            D1s.append((model_prev_i - model_prev_0) / rk)

        rks.append(1.)
        rks = torch.tensor(rks, device=x.device)

        K = len(rks)
        # build C matrix
        C = []

        col = torch.ones_like(rks)
        for k in range(1, K + 1):
            C.append(col)
            col = col * rks / (k + 1) 
        C = torch.stack(C, dim=1)

        if len(D1s) > 0:
            D1s = torch.stack(D1s, dim=1) # (B, K)
            C_inv_p = torch.linalg.inv(C[:-1, :-1])
            A_p = C_inv_p

        if use_corrector:
            # print('using corrector')
            C_inv = torch.linalg.inv(C)
            A_c = C_inv

        hh = h
        h_phi_1 = torch.expm1(hh)
        h_phi_ks = []
        factorial_k = 1
        h_phi_k = h_phi_1
        for k in range(1, K + 2):
            h_phi_ks.append(h_phi_k)
            h_phi_k = h_phi_k / hh - 1 / factorial_k
            factorial_k *= (k + 1)

        model_t = None
        if True:
            log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
            x_t_ = (
                expand_dims((torch.exp(log_alpha_t - log_alpha_prev_0)),dims) * x
                - expand_dims((sigma_t * h_phi_1),dims) * model_prev_0
            )
            # now predictor
            x_t = x_t_
            if len(D1s) > 0:
                # compute the residuals for predictor
                for k in range(K - 1):
                    x_t = x_t - expand_dims(sigma_t * h_phi_ks[k + 1],dims) * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
            # now corrector
            if use_corrector:
                model_t = self.noise_prediction_fn(x_t, t)
                D1_t = (model_t - model_prev_0)
                x_t = x_t_
                k = 0
                for k in range(K - 1):
                    x_t = x_t - expand_dims(sigma_t * h_phi_ks[k + 1],dims) * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
                x_t = x_t - expand_dims(sigma_t * h_phi_ks[K],dims) * (D1_t * A_c[k][-1])
        return x_t, model_t

    def sample(
        self,
        x,
        model_fn,
        order,
        use_corrector,
        lower_order_final,
        start_free_u_step=None,
        free_u_apply_callback=None,
        free_u_stop_callback=None,
        npnet_x = None,
        npnet_scale = None,
        half=False,
        return_intermediate=False,
    ):
        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
        steps = self.steps
        vec_t = self.timesteps[0].expand((x.shape[0]))
        if free_u_stop_callback is not None:
            free_u_stop_callback()
        if start_free_u_step is not None and 0 == start_free_u_step and free_u_apply_callback is not None:
            free_u_apply_callback()
            has_called_free_u = True
        if not self.use_afs:
            fir_output = self.noise_prediction_fn(x, vec_t)
        else:
            fir_output = x  # ultilize npnet there in the future
            if npnet_x is not None and npnet_scale is not None:
                fir_output = npnet_x 
                # fir_output = fir_output - npnet_scale * (npnet_out - fir_output) #guidance_scale * (noise - noise_uncond)
                x = fir_output.clone().detach().to(fir_output.device)
            

        model_prev_list = [fir_output]
        full_cache = [fir_output]
        t_prev_list = [vec_t]
        has_called_free_u = False
        for init_order in range(1, order):
            if start_free_u_step is not None and init_order == start_free_u_step and free_u_apply_callback is not None and (not has_called_free_u):
                free_u_apply_callback()
                has_called_free_u = True
            vec_t = self.timesteps[init_order].expand(x.shape[0])
            x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
            if model_x is None:
                model_x = self.noise_prediction_fn(x, vec_t)
                x = model_x.clone().detach().to(torch.float32).to(model_x.device)
            full_cache.append(x)
            model_prev_list.append(model_x)
            t_prev_list.append(vec_t)
        
        for step in range(order, steps + 1):
            if start_free_u_step is not None and step == start_free_u_step and free_u_apply_callback is not None and (not has_called_free_u):
                free_u_apply_callback()
            vec_t = self.timesteps[step].expand(x.shape[0])
            if lower_order_final:
                step_order = min(order, steps + 1 - step)
            else:
                step_order = order
            # print('this step order:', step_order)
            if step == steps:
                # print('do not run corrector at the last step')
                use_corrector = False
            else:
                use_corrector = True
            x, model_x =  self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
            for i in range(order - 1):
                t_prev_list[i] = t_prev_list[i + 1]
                model_prev_list[i] = model_prev_list[i + 1]
            t_prev_list[-1] = vec_t
                    # We do not need to evaluate the final model value.
            full_cache.append(x)
            if step < steps:
                if model_x is None:
                    model_x = self.noise_prediction_fn(x, vec_t)
                model_prev_list[-1] = model_x
        return x, full_cache
    def sample_mix(
        self,
        x,
        model_fn,
        order,
        use_corrector,
        lower_order_final,
        start_free_u_step=None,
        free_u_apply_callback=None,
        free_u_stop_callback=None,
        noise_level = 0.1,
        half=False,
        return_intermediate=False,
    ):
        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
        steps = self.steps
        vec_t = self.timesteps[0].expand((x.shape[0]))
        fir_output = self.noise_prediction_fn(x, vec_t)
        model_prev_list = [fir_output]
        full_cache = [fir_output]
        t_prev_list = [vec_t]
        has_called_free_u = False
        if free_u_stop_callback is not None:
            free_u_stop_callback()
        for init_order in range(1, order):
            if start_free_u_step is not None and init_order == start_free_u_step and free_u_apply_callback is not None:
                free_u_apply_callback()
                has_called_free_u = True
            vec_t = self.timesteps[init_order].expand(x.shape[0])
            if start_free_u_step is not None and init_order >= start_free_u_step and free_u_apply_callback is not None:
                x, model_x = self.multistep_uni_pc_sde_update(x
                                                              , model_prev_list
                                                              , t_prev_list
                                                              , vec_t
                                                              , init_order
                                                              , use_corrector=True
                                                              ,level=noise_level)
            else:
                x, model_x = self.multistep_uni_pc_sde_update(x
                                                              , model_prev_list
                                                              , t_prev_list
                                                              , vec_t
                                                              , init_order
                                                              , use_corrector=True
                                                              ,level=0.0)
            if model_x is None:
                model_x = self.noise_prediction_fn(x, vec_t)
                x = model_x.clone().detach().to(torch.float32).to(model_x.device)
            full_cache.append(x)
            model_prev_list.append(model_x)
            t_prev_list.append(vec_t)
        
        if free_u_stop_callback is not None:
            free_u_stop_callback()
        for step in range(order, steps + 1):
            if start_free_u_step is not None and step == start_free_u_step and free_u_apply_callback is not None and (not has_called_free_u):
                free_u_apply_callback()
            vec_t = self.timesteps[step].expand(x.shape[0])
            if lower_order_final:
                step_order = min(order, steps + 1 - step)
            else:
                step_order = order
            # print('this step order:', step_order)
            if step == steps:
                # print('do not run corrector at the last step')
                use_corrector = False
            else:
                use_corrector = True
            if start_free_u_step is not None and step >= start_free_u_step and free_u_apply_callback is not None:
                x, model_x =  self.multistep_uni_pc_sde_update(x
                                                               , model_prev_list
                                                               , t_prev_list
                                                               , vec_t
                                                               , step_order
                                                               , use_corrector=use_corrector
                                                               , level=noise_level)
            else:
                x, model_x =  self.multistep_uni_pc_sde_update(x
                                                               , model_prev_list
                                                               , t_prev_list
                                                               , vec_t
                                                               , step_order
                                                               , use_corrector=use_corrector
                                                               , level=0.0)
            for i in range(order - 1):
                t_prev_list[i] = t_prev_list[i + 1]
                model_prev_list[i] = model_prev_list[i + 1]
            t_prev_list[-1] = vec_t
                    # We do not need to evaluate the final model value.
            full_cache.append(x)
            if step < steps:
                if model_x is None:
                    model_x = self.noise_prediction_fn(x, vec_t)
                model_prev_list[-1] = model_x
        return x, full_cache
        
        


#############################################################
# other utility functions
#############################################################

def interpolate_fn(x, xp, yp):
    """
    A piecewise linear function y = f(x), using xp and yp as keypoints.
    We implement f(x) in a differentiable way (i.e. applicable for autograd).
    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)

    Args:
        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
        yp: PyTorch tensor with shape [C, K].
    Returns:
        The function values f(x), with shape [N, C].
    """
    N, K = x.shape[0], xp.shape[1]
    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
    x_idx = torch.argmin(x_indices, dim=2)
    cand_start_idx = x_idx - 1
    start_idx = torch.where(
        torch.eq(x_idx, 0),
        torch.tensor(1, device=x.device),
        torch.where(
            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
        ),
    )
    end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
    start_idx2 = torch.where(
        torch.eq(x_idx, 0),
        torch.tensor(0, device=x.device),
        torch.where(
            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
        ),
    )
    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
    start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
    end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
    return cand


def expand_dims(v, dims):
    """
    Expand the tensor `v` to the dim `dims`.

    Args:
        `v`: a PyTorch tensor with shape [N].
        `dim`: a `int`.
    Returns:
        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
    """
    return v[(...,) + (None,)*(dims - 1)]