# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

import math

import torch
import torch.nn.functional as F
from tqdm import tqdm


class NoiseScheduleVP:
    def __init__(
        self,
        schedule="discrete",
        betas=None,
        alphas_cumprod=None,
        continuous_beta_0=0.1,
        continuous_beta_1=20.0,
        dtype=torch.float32,
    ):
        """Thanks to DPM-Solver for their code base"""
        r"""Create a wrapper class for the forward SDE (VP type).
        ***
        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
        ***
        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
            log_alpha_t = self.marginal_log_mean_coeff(t)
            sigma_t = self.marginal_std(t)
            lambda_t = self.marginal_lambda(t)
        Moreover, as lambda(t) is an invertible function, we also support its inverse function:
            t = self.inverse_lambda(lambda_t)
        ===============================================================
        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
        1. For discrete-time DPMs:
            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
                t_i = (i + 1) / N
            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
            Args:
                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
            Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
                and
                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
        2. For continuous-time DPMs:
            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
            schedule are the default settings in DDPM and improved-DDPM:
            Args:
                beta_min: A `float` number. The smallest beta for the linear schedule.
                beta_max: A `float` number. The largest beta for the linear schedule.
                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
                T: A `float` number. The ending time of the forward process.
        ===============================================================
        Args:
            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
                    'linear' or 'cosine' for continuous-time DPMs.
        Returns:
            A wrapper object of the forward SDE (VP type).

        ===============================================================
        Example:
        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
        >>> ns = NoiseScheduleVP('discrete', betas=betas)
        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
        # For continuous-time DPMs (VPSDE), linear schedule:
        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
        """

        if schedule not in ["discrete", "linear", "cosine"]:
            raise ValueError(
                "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
                    schedule
                )
            )

        self.schedule = schedule
        if schedule == "discrete":
            if betas is not None:
                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
            else:
                assert alphas_cumprod is not None
                log_alphas = 0.5 * torch.log(alphas_cumprod)
            self.total_N = len(log_alphas)
            self.T = 1.0
            self.t_array = torch.linspace(0.0, 1.0, self.total_N + 1)[1:].reshape((1, -1)).to(dtype=dtype)
            self.log_alpha_array = log_alphas.reshape(
                (
                    1,
                    -1,
                )
            ).to(dtype=dtype)
        else:
            self.total_N = 1000
            self.beta_0 = continuous_beta_0
            self.beta_1 = continuous_beta_1
            self.cosine_s = 0.008
            self.cosine_beta_max = 999.0
            self.cosine_t_max = (
                math.atan(self.cosine_beta_max * (1.0 + self.cosine_s) / math.pi)
                * 2.0
                * (1.0 + self.cosine_s)
                / math.pi
                - self.cosine_s
            )
            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1.0 + self.cosine_s) * math.pi / 2.0))
            self.schedule = schedule
            if schedule == "cosine":
                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
                self.T = 0.9946
            else:
                self.T = 1.0

    def marginal_log_mean_coeff(self, t):
        """
        Compute log(alpha_t) of a given continuous-time label t in [0, T].
        """
        if self.schedule == "discrete":
            return interpolate_fn(
                t.reshape((-1, 1)), self.t_array.to(t.device), self.log_alpha_array.to(t.device)
            ).reshape(-1)
        elif self.schedule == "linear":
            return -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
        elif self.schedule == "cosine":
            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1.0 + self.cosine_s) * math.pi / 2.0))
            log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
            return log_alpha_t

    def marginal_alpha(self, t):
        """
        Compute alpha_t of a given continuous-time label t in [0, T].
        """
        return torch.exp(self.marginal_log_mean_coeff(t))

    def marginal_std(self, t):
        """
        Compute sigma_t of a given continuous-time label t in [0, T].
        """
        return torch.sqrt(1.0 - torch.exp(2.0 * self.marginal_log_mean_coeff(t)))

    def marginal_lambda(self, t):
        """
        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
        """
        log_mean_coeff = self.marginal_log_mean_coeff(t)
        log_std = 0.5 * torch.log(1.0 - torch.exp(2.0 * log_mean_coeff))
        return log_mean_coeff - log_std

    def inverse_lambda(self, lamb):
        """
        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
        """
        if self.schedule == "linear":
            tmp = 2.0 * (self.beta_1 - self.beta_0) * torch.logaddexp(-2.0 * lamb, torch.zeros((1,)).to(lamb))
            Delta = self.beta_0**2 + tmp
            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
        elif self.schedule == "discrete":
            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2.0 * lamb)
            t = interpolate_fn(
                log_alpha.reshape((-1, 1)),
                torch.flip(self.log_alpha_array.to(lamb.device), [1]),
                torch.flip(self.t_array.to(lamb.device), [1]),
            )
            return t.reshape((-1,))
        else:
            log_alpha = -0.5 * torch.logaddexp(-2.0 * lamb, torch.zeros((1,)).to(lamb))
            t_fn = (
                lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0))
                * 2.0
                * (1.0 + self.cosine_s)
                / math.pi
                - self.cosine_s
            )
            t = t_fn(log_alpha)
            return t

    def edm_sigma(self, t):
        return self.marginal_std(t) / self.marginal_alpha(t)

    def edm_inverse_sigma(self, edmsigma):
        alpha = 1 / (edmsigma**2 + 1).sqrt()
        sigma = alpha * edmsigma
        lambda_t = torch.log(alpha / sigma)
        t = self.inverse_lambda(lambda_t)
        return t


def model_wrapper(
    model,
    noise_schedule,
    model_type="noise",
    model_kwargs={},
    guidance_type="uncond",
    condition=None,
    unconditional_condition=None,
    guidance_scale=1.0,
    classifier_fn=None,
    classifier_kwargs={},
):
    """Thanks to DPM-Solver for their code base"""
    """Create a wrapper function for the noise prediction model.
    SA-Solver needs to solve the continuous-time diffusion SDEs. For DPMs trained on discrete-time labels, we need to
    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
    We support four types of the diffusion model by setting `model_type`:
        1. "noise": noise prediction model. (Trained by predicting noise).
        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
        3. "v": velocity prediction model. (Trained by predicting the velocity).
            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
                arXiv preprint arXiv:2202.00512 (2022).
            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
                arXiv preprint arXiv:2210.02303 (2022).

        4. "score": marginal score function. (Trained by denoising score matching).
            Note that the score function and the noise prediction model follows a simple relationship:
            ```
                noise(x_t, t) = -sigma_t * score(x_t, t)
            ```
    We support three types of guided sampling by DPMs by setting `guidance_type`:
        1. "uncond": unconditional sampling by DPMs.
            The input `model` has the following format:
            ``
                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
            ``
        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
            The input `model` has the following format:
            ``
                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
            ``
            The input `classifier_fn` has the following format:
            ``
                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
            ``
            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
            The input `model` has the following format:
            ``
                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
            ``
            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
                arXiv preprint arXiv:2207.12598 (2022).

    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
    or continuous-time labels (i.e. epsilon to T).
    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
    ``
        def model_fn(x, t_continuous) -> noise:
            t_input = get_model_input_time(t_continuous)
            return noise_pred(model, x, t_input, **model_kwargs)
    ``
    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for SA-Solver.
    ===============================================================
    Args:
        model: A diffusion model with the corresponding format described above.
        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
        model_type: A `str`. The parameterization type of the diffusion model.
                    "noise" or "x_start" or "v" or "score".
        model_kwargs: A `dict`. A dict for the other inputs of the model function.
        guidance_type: A `str`. The type of the guidance for sampling.
                    "uncond" or "classifier" or "classifier-free".
        condition: A pytorch tensor. The condition for the guided sampling.
                    Only used for "classifier" or "classifier-free" guidance type.
        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
                    Only used for "classifier-free" guidance type.
        guidance_scale: A `float`. The scale for the guided sampling.
        classifier_fn: A classifier function. Only used for the classifier guidance.
        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
    Returns:
        A noise prediction model that accepts the noised data and the continuous time as the inputs.
    """

    def get_model_input_time(t_continuous):
        """
        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
        For continuous-time DPMs, we just use `t_continuous`.
        """
        if noise_schedule.schedule == "discrete":
            return (t_continuous - 1.0 / noise_schedule.total_N) * 1000.0
        else:
            return t_continuous

    def noise_pred_fn(x, t_continuous, cond=None):
        t_input = get_model_input_time(t_continuous)
        if cond is None:
            output = model(x, t_input, **model_kwargs)
        else:
            output = model(x, t_input, cond, **model_kwargs)
        if model_type == "noise":
            return output
        elif model_type == "x_start":
            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
            return (x - alpha_t[0] * output) / sigma_t[0]
        elif model_type == "v":
            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
            return alpha_t[0] * output + sigma_t[0] * x
        elif model_type == "score":
            sigma_t = noise_schedule.marginal_std(t_continuous)
            return -sigma_t[0] * output

    def cond_grad_fn(x, t_input):
        """
        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
        """
        with torch.enable_grad():
            x_in = x.detach().requires_grad_(True)
            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
            return torch.autograd.grad(log_prob.sum(), x_in)[0]

    def model_fn(x, t_continuous):
        """
        The noise predicition model function that is used for DPM-Solver.
        """
        if guidance_type == "uncond":
            return noise_pred_fn(x, t_continuous)
        elif guidance_type == "classifier":
            assert classifier_fn is not None
            t_input = get_model_input_time(t_continuous)
            cond_grad = cond_grad_fn(x, t_input)
            sigma_t = noise_schedule.marginal_std(t_continuous)
            noise = noise_pred_fn(x, t_continuous)
            return noise - guidance_scale * sigma_t * cond_grad
        elif guidance_type == "classifier-free":
            if guidance_scale == 1.0 or unconditional_condition is None:
                return noise_pred_fn(x, t_continuous, cond=condition)
            else:
                x_in = torch.cat([x] * 2)
                t_in = torch.cat([t_continuous] * 2)
                c_in = torch.cat([unconditional_condition, condition])
                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
                return noise_uncond + guidance_scale * (noise - noise_uncond)

    assert model_type in ["noise", "x_start", "v", "score"]
    assert guidance_type in ["uncond", "classifier", "classifier-free"]
    return model_fn


class SASolver:
    def __init__(
        self,
        model_fn,
        noise_schedule,
        algorithm_type="data_prediction",
        correcting_x0_fn=None,
        correcting_xt_fn=None,
        thresholding_max_val=1.0,
        dynamic_thresholding_ratio=0.995,
    ):
        """
        Construct a SA-Solver
        The default value for algorithm_type is "data_prediction" and we recommend not to change it to
        "noise_prediction". For details, please see Appendix A.2.4 in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        """

        self.model = lambda x, t: model_fn(x, t.expand(x.shape[0]))
        self.noise_schedule = noise_schedule
        assert algorithm_type in ["data_prediction", "noise_prediction"]

        if correcting_x0_fn == "dynamic_thresholding":
            self.correcting_x0_fn = self.dynamic_thresholding_fn
        else:
            self.correcting_x0_fn = correcting_x0_fn

        self.correcting_xt_fn = correcting_xt_fn
        self.dynamic_thresholding_ratio = dynamic_thresholding_ratio
        self.thresholding_max_val = thresholding_max_val

        self.predict_x0 = algorithm_type == "data_prediction"

        self.sigma_min = float(self.noise_schedule.edm_sigma(torch.tensor([1e-3])))
        self.sigma_max = float(self.noise_schedule.edm_sigma(torch.tensor([1])))

    def dynamic_thresholding_fn(self, x0, t=None):
        """
        The dynamic thresholding method.
        """
        dims = x0.dim()
        p = self.dynamic_thresholding_ratio
        s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
        s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
        x0 = torch.clamp(x0, -s, s) / s
        return x0

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

    def data_prediction_fn(self, x, t):
        """
        Return the data prediction model (with corrector).
        """
        noise = self.noise_prediction_fn(x, t)
        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
        x0 = (x - sigma_t * noise) / alpha_t
        if self.correcting_x0_fn is not None:
            x0 = self.correcting_x0_fn(x0)
        return x0

    def model_fn(self, x, t):
        """
        Convert the model to the noise prediction model or the data prediction model.
        """

        if self.predict_x0:
            return self.data_prediction_fn(x, t)
        else:
            return self.noise_prediction_fn(x, t)

    def get_time_steps(self, skip_type, t_T, t_0, N, order, device):
        """Compute the intermediate time steps for sampling."""
        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 = lambda_T + torch.linspace(
                torch.tensor(0.0).cpu().item(), (lambda_0 - lambda_T).cpu().item() ** (1.0 / order), N + 1
            ).pow(order).to(device)
            return self.noise_schedule.inverse_lambda(logSNR_steps)
        elif skip_type == "time":
            t = torch.linspace(t_T ** (1.0 / order), t_0 ** (1.0 / order), N + 1).pow(order).to(device)
            return t
        elif skip_type == "karras":
            sigma_min = max(0.002, self.sigma_min)
            sigma_max = min(80, self.sigma_max)
            sigma_steps = torch.linspace(sigma_max ** (1.0 / 7), sigma_min ** (1.0 / 7), N + 1).pow(7).to(device)
            t = self.noise_schedule.edm_inverse_sigma(sigma_steps)
            return t
        else:
            raise ValueError(f"Unsupported skip_type {skip_type}, need to be 'logSNR' or 'time' or 'karras'")

    def denoise_to_zero_fn(self, x, s):
        """
        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
        """
        return self.data_prediction_fn(x, s)

    def get_coefficients_exponential_negative(self, order, interval_start, interval_end):
        """
        Calculate the integral of exp(-x) * x^order dx from interval_start to interval_end
        For calculating the coefficient of gradient terms after the lagrange interpolation,
        see Eq.(15) and Eq.(18) in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        For noise_prediction formula.
        """
        assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3"

        if order == 0:
            return torch.exp(-interval_end) * (torch.exp(interval_end - interval_start) - 1)
        elif order == 1:
            return torch.exp(-interval_end) * (
                (interval_start + 1) * torch.exp(interval_end - interval_start) - (interval_end + 1)
            )
        elif order == 2:
            return torch.exp(-interval_end) * (
                (interval_start**2 + 2 * interval_start + 2) * torch.exp(interval_end - interval_start)
                - (interval_end**2 + 2 * interval_end + 2)
            )
        elif order == 3:
            return torch.exp(-interval_end) * (
                (interval_start**3 + 3 * interval_start**2 + 6 * interval_start + 6)
                * torch.exp(interval_end - interval_start)
                - (interval_end**3 + 3 * interval_end**2 + 6 * interval_end + 6)
            )

    def get_coefficients_exponential_positive(self, order, interval_start, interval_end, tau):
        """
        Calculate the integral of exp(x(1+tau^2)) * x^order dx from interval_start to interval_end
        For calculating the coefficient of gradient terms after the lagrange interpolation,
        see Eq.(15) and Eq.(18) in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        For data_prediction formula.
        """
        assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3"

        # after change of variable(cov)
        interval_end_cov = (1 + tau**2) * interval_end
        interval_start_cov = (1 + tau**2) * interval_start

        if order == 0:
            return (
                torch.exp(interval_end_cov) * (1 - torch.exp(-(interval_end_cov - interval_start_cov))) / (1 + tau**2)
            )
        elif order == 1:
            return (
                torch.exp(interval_end_cov)
                * (
                    (interval_end_cov - 1)
                    - (interval_start_cov - 1) * torch.exp(-(interval_end_cov - interval_start_cov))
                )
                / ((1 + tau**2) ** 2)
            )
        elif order == 2:
            return (
                torch.exp(interval_end_cov)
                * (
                    (interval_end_cov**2 - 2 * interval_end_cov + 2)
                    - (interval_start_cov**2 - 2 * interval_start_cov + 2)
                    * torch.exp(-(interval_end_cov - interval_start_cov))
                )
                / ((1 + tau**2) ** 3)
            )
        elif order == 3:
            return (
                torch.exp(interval_end_cov)
                * (
                    (interval_end_cov**3 - 3 * interval_end_cov**2 + 6 * interval_end_cov - 6)
                    - (interval_start_cov**3 - 3 * interval_start_cov**2 + 6 * interval_start_cov - 6)
                    * torch.exp(-(interval_end_cov - interval_start_cov))
                )
                / ((1 + tau**2) ** 4)
            )

    def lagrange_polynomial_coefficient(self, order, lambda_list):
        """
        Calculate the coefficient of lagrange polynomial
        For lagrange interpolation
        """
        assert order in [0, 1, 2, 3]
        assert order == len(lambda_list) - 1
        if order == 0:
            return [[1]]
        elif order == 1:
            return [
                [1 / (lambda_list[0] - lambda_list[1]), -lambda_list[1] / (lambda_list[0] - lambda_list[1])],
                [1 / (lambda_list[1] - lambda_list[0]), -lambda_list[0] / (lambda_list[1] - lambda_list[0])],
            ]
        elif order == 2:
            denominator1 = (lambda_list[0] - lambda_list[1]) * (lambda_list[0] - lambda_list[2])
            denominator2 = (lambda_list[1] - lambda_list[0]) * (lambda_list[1] - lambda_list[2])
            denominator3 = (lambda_list[2] - lambda_list[0]) * (lambda_list[2] - lambda_list[1])
            return [
                [
                    1 / denominator1,
                    (-lambda_list[1] - lambda_list[2]) / denominator1,
                    lambda_list[1] * lambda_list[2] / denominator1,
                ],
                [
                    1 / denominator2,
                    (-lambda_list[0] - lambda_list[2]) / denominator2,
                    lambda_list[0] * lambda_list[2] / denominator2,
                ],
                [
                    1 / denominator3,
                    (-lambda_list[0] - lambda_list[1]) / denominator3,
                    lambda_list[0] * lambda_list[1] / denominator3,
                ],
            ]
        elif order == 3:
            denominator1 = (
                (lambda_list[0] - lambda_list[1])
                * (lambda_list[0] - lambda_list[2])
                * (lambda_list[0] - lambda_list[3])
            )
            denominator2 = (
                (lambda_list[1] - lambda_list[0])
                * (lambda_list[1] - lambda_list[2])
                * (lambda_list[1] - lambda_list[3])
            )
            denominator3 = (
                (lambda_list[2] - lambda_list[0])
                * (lambda_list[2] - lambda_list[1])
                * (lambda_list[2] - lambda_list[3])
            )
            denominator4 = (
                (lambda_list[3] - lambda_list[0])
                * (lambda_list[3] - lambda_list[1])
                * (lambda_list[3] - lambda_list[2])
            )
            return [
                [
                    1 / denominator1,
                    (-lambda_list[1] - lambda_list[2] - lambda_list[3]) / denominator1,
                    (
                        lambda_list[1] * lambda_list[2]
                        + lambda_list[1] * lambda_list[3]
                        + lambda_list[2] * lambda_list[3]
                    )
                    / denominator1,
                    (-lambda_list[1] * lambda_list[2] * lambda_list[3]) / denominator1,
                ],
                [
                    1 / denominator2,
                    (-lambda_list[0] - lambda_list[2] - lambda_list[3]) / denominator2,
                    (
                        lambda_list[0] * lambda_list[2]
                        + lambda_list[0] * lambda_list[3]
                        + lambda_list[2] * lambda_list[3]
                    )
                    / denominator2,
                    (-lambda_list[0] * lambda_list[2] * lambda_list[3]) / denominator2,
                ],
                [
                    1 / denominator3,
                    (-lambda_list[0] - lambda_list[1] - lambda_list[3]) / denominator3,
                    (
                        lambda_list[0] * lambda_list[1]
                        + lambda_list[0] * lambda_list[3]
                        + lambda_list[1] * lambda_list[3]
                    )
                    / denominator3,
                    (-lambda_list[0] * lambda_list[1] * lambda_list[3]) / denominator3,
                ],
                [
                    1 / denominator4,
                    (-lambda_list[0] - lambda_list[1] - lambda_list[2]) / denominator4,
                    (
                        lambda_list[0] * lambda_list[1]
                        + lambda_list[0] * lambda_list[2]
                        + lambda_list[1] * lambda_list[2]
                    )
                    / denominator4,
                    (-lambda_list[0] * lambda_list[1] * lambda_list[2]) / denominator4,
                ],
            ]

    def get_coefficients_fn(self, order, interval_start, interval_end, lambda_list, tau):
        """
        Calculate the coefficient of gradients.
        """
        assert order in [1, 2, 3, 4]
        assert order == len(lambda_list), "the length of lambda list must be equal to the order"
        coefficients = []
        lagrange_coefficient = self.lagrange_polynomial_coefficient(order - 1, lambda_list)
        for i in range(order):
            coefficient = 0
            for j in range(order):
                if self.predict_x0:
                    coefficient += lagrange_coefficient[i][j] * self.get_coefficients_exponential_positive(
                        order - 1 - j, interval_start, interval_end, tau
                    )
                else:
                    coefficient += lagrange_coefficient[i][j] * self.get_coefficients_exponential_negative(
                        order - 1 - j, interval_start, interval_end
                    )
            coefficients.append(coefficient)
        assert len(coefficients) == order, "the length of coefficients does not match the order"
        return coefficients

    def adams_bashforth_update(self, order, x, tau, model_prev_list, t_prev_list, noise, t):
        """
        SA-Predictor, without the "rescaling" trick in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        """
        assert order in [1, 2, 3, 4], "order of stochastic adams bashforth method is only supported for 1, 2, 3 and 4"

        # get noise schedule
        ns = self.noise_schedule
        alpha_t = ns.marginal_alpha(t)
        sigma_t = ns.marginal_std(t)
        lambda_t = ns.marginal_lambda(t)
        alpha_prev = ns.marginal_alpha(t_prev_list[-1])
        sigma_prev = ns.marginal_std(t_prev_list[-1])
        gradient_part = torch.zeros_like(x)
        h = lambda_t - ns.marginal_lambda(t_prev_list[-1])
        lambda_list = []
        for i in range(order):
            lambda_list.append(ns.marginal_lambda(t_prev_list[-(i + 1)]))
        gradient_coefficients = self.get_coefficients_fn(
            order, ns.marginal_lambda(t_prev_list[-1]), lambda_t, lambda_list, tau
        )

        for i in range(order):
            if self.predict_x0:
                gradient_part += (
                    (1 + tau**2)
                    * sigma_t
                    * torch.exp(-(tau**2) * lambda_t)
                    * gradient_coefficients[i]
                    * model_prev_list[-(i + 1)]
                )
            else:
                gradient_part += -(1 + tau**2) * alpha_t * gradient_coefficients[i] * model_prev_list[-(i + 1)]

        if self.predict_x0:
            noise_part = sigma_t * torch.sqrt(1 - torch.exp(-2 * tau**2 * h)) * noise
        else:
            noise_part = tau * sigma_t * torch.sqrt(torch.exp(2 * h) - 1) * noise

        if self.predict_x0:
            x_t = torch.exp(-(tau**2) * h) * (sigma_t / sigma_prev) * x + gradient_part + noise_part
        else:
            x_t = (alpha_t / alpha_prev) * x + gradient_part + noise_part

        return x_t

    def adams_moulton_update(self, order, x, tau, model_prev_list, t_prev_list, noise, t):
        """
        SA-Corrector, without the "rescaling" trick in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        """

        assert order in [1, 2, 3, 4], "order of stochastic adams bashforth method is only supported for 1, 2, 3 and 4"

        # get noise schedule
        ns = self.noise_schedule
        alpha_t = ns.marginal_alpha(t)
        sigma_t = ns.marginal_std(t)
        lambda_t = ns.marginal_lambda(t)
        alpha_prev = ns.marginal_alpha(t_prev_list[-1])
        sigma_prev = ns.marginal_std(t_prev_list[-1])
        gradient_part = torch.zeros_like(x)
        h = lambda_t - ns.marginal_lambda(t_prev_list[-1])
        lambda_list = []
        t_list = t_prev_list + [t]
        for i in range(order):
            lambda_list.append(ns.marginal_lambda(t_list[-(i + 1)]))
        gradient_coefficients = self.get_coefficients_fn(
            order, ns.marginal_lambda(t_prev_list[-1]), lambda_t, lambda_list, tau
        )

        for i in range(order):
            if self.predict_x0:
                gradient_part += (
                    (1 + tau**2)
                    * sigma_t
                    * torch.exp(-(tau**2) * lambda_t)
                    * gradient_coefficients[i]
                    * model_prev_list[-(i + 1)]
                )
            else:
                gradient_part += -(1 + tau**2) * alpha_t * gradient_coefficients[i] * model_prev_list[-(i + 1)]

        if self.predict_x0:
            noise_part = sigma_t * torch.sqrt(1 - torch.exp(-2 * tau**2 * h)) * noise
        else:
            noise_part = tau * sigma_t * torch.sqrt(torch.exp(2 * h) - 1) * noise

        if self.predict_x0:
            x_t = torch.exp(-(tau**2) * h) * (sigma_t / sigma_prev) * x + gradient_part + noise_part
        else:
            x_t = (alpha_t / alpha_prev) * x + gradient_part + noise_part

        return x_t

    def adams_bashforth_update_few_steps(self, order, x, tau, model_prev_list, t_prev_list, noise, t):
        """
        SA-Predictor, with the "rescaling" trick in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        """

        assert order in [1, 2, 3, 4], "order of stochastic adams bashforth method is only supported for 1, 2, 3 and 4"

        # get noise schedule
        ns = self.noise_schedule
        alpha_t = ns.marginal_alpha(t)
        sigma_t = ns.marginal_std(t)
        lambda_t = ns.marginal_lambda(t)
        alpha_prev = ns.marginal_alpha(t_prev_list[-1])
        sigma_prev = ns.marginal_std(t_prev_list[-1])
        gradient_part = torch.zeros_like(x)
        h = lambda_t - ns.marginal_lambda(t_prev_list[-1])
        lambda_list = []
        for i in range(order):
            lambda_list.append(ns.marginal_lambda(t_prev_list[-(i + 1)]))
        gradient_coefficients = self.get_coefficients_fn(
            order, ns.marginal_lambda(t_prev_list[-1]), lambda_t, lambda_list, tau
        )

        if self.predict_x0:
            if (
                order == 2
            ):  ## if order = 2 we do a modification that does not influence the convergence order similar to unipc. Note: This is used only for few steps sampling.
                # The added term is O(h^3). Empirically we find it will slightly improve the image quality.
                # ODE case
                # gradient_coefficients[0] += 1.0 * torch.exp(lambda_t) * (h ** 2 / 2 - (h - 1 + torch.exp(-h))) / (ns.marginal_lambda(t_prev_list[-1]) - ns.marginal_lambda(t_prev_list[-2]))
                # gradient_coefficients[1] -= 1.0 * torch.exp(lambda_t) * (h ** 2 / 2 - (h - 1 + torch.exp(-h))) / (ns.marginal_lambda(t_prev_list[-1]) - ns.marginal_lambda(t_prev_list[-2]))
                gradient_coefficients[0] += (
                    1.0
                    * torch.exp((1 + tau**2) * lambda_t)
                    * (h**2 / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2))
                    / (ns.marginal_lambda(t_prev_list[-1]) - ns.marginal_lambda(t_prev_list[-2]))
                )
                gradient_coefficients[1] -= (
                    1.0
                    * torch.exp((1 + tau**2) * lambda_t)
                    * (h**2 / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2))
                    / (ns.marginal_lambda(t_prev_list[-1]) - ns.marginal_lambda(t_prev_list[-2]))
                )

        for i in range(order):
            if self.predict_x0:
                gradient_part += (
                    (1 + tau**2)
                    * sigma_t
                    * torch.exp(-(tau**2) * lambda_t)
                    * gradient_coefficients[i]
                    * model_prev_list[-(i + 1)]
                )
            else:
                gradient_part += -(1 + tau**2) * alpha_t * gradient_coefficients[i] * model_prev_list[-(i + 1)]

        if self.predict_x0:
            noise_part = sigma_t * torch.sqrt(1 - torch.exp(-2 * tau**2 * h)) * noise
        else:
            noise_part = tau * sigma_t * torch.sqrt(torch.exp(2 * h) - 1) * noise

        if self.predict_x0:
            x_t = torch.exp(-(tau**2) * h) * (sigma_t / sigma_prev) * x + gradient_part + noise_part
        else:
            x_t = (alpha_t / alpha_prev) * x + gradient_part + noise_part

        return x_t

    def adams_moulton_update_few_steps(self, order, x, tau, model_prev_list, t_prev_list, noise, t):
        """
        SA-Corrector, without the "rescaling" trick in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        """

        assert order in [1, 2, 3, 4], "order of stochastic adams bashforth method is only supported for 1, 2, 3 and 4"

        # get noise schedule
        ns = self.noise_schedule
        alpha_t = ns.marginal_alpha(t)
        sigma_t = ns.marginal_std(t)
        lambda_t = ns.marginal_lambda(t)
        alpha_prev = ns.marginal_alpha(t_prev_list[-1])
        sigma_prev = ns.marginal_std(t_prev_list[-1])
        gradient_part = torch.zeros_like(x)
        h = lambda_t - ns.marginal_lambda(t_prev_list[-1])
        lambda_list = []
        t_list = t_prev_list + [t]
        for i in range(order):
            lambda_list.append(ns.marginal_lambda(t_list[-(i + 1)]))
        gradient_coefficients = self.get_coefficients_fn(
            order, ns.marginal_lambda(t_prev_list[-1]), lambda_t, lambda_list, tau
        )

        if self.predict_x0:
            if (
                order == 2
            ):  ## if order = 2 we do a modification that does not influence the convergence order similar to UniPC. Note: This is used only for few steps sampling.
                # The added term is O(h^3). Empirically we find it will slightly improve the image quality.
                # ODE case
                # gradient_coefficients[0] += 1.0 * torch.exp(lambda_t) * (h / 2 - (h - 1 + torch.exp(-h)) / h)
                # gradient_coefficients[1] -= 1.0 * torch.exp(lambda_t) * (h / 2 - (h - 1 + torch.exp(-h)) / h)
                gradient_coefficients[0] += (
                    1.0
                    * torch.exp((1 + tau**2) * lambda_t)
                    * (h / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2 * h))
                )
                gradient_coefficients[1] -= (
                    1.0
                    * torch.exp((1 + tau**2) * lambda_t)
                    * (h / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2 * h))
                )

        for i in range(order):
            if self.predict_x0:
                gradient_part += (
                    (1 + tau**2)
                    * sigma_t
                    * torch.exp(-(tau**2) * lambda_t)
                    * gradient_coefficients[i]
                    * model_prev_list[-(i + 1)]
                )
            else:
                gradient_part += -(1 + tau**2) * alpha_t * gradient_coefficients[i] * model_prev_list[-(i + 1)]

        if self.predict_x0:
            noise_part = sigma_t * torch.sqrt(1 - torch.exp(-2 * tau**2 * h)) * noise
        else:
            noise_part = tau * sigma_t * torch.sqrt(torch.exp(2 * h) - 1) * noise

        if self.predict_x0:
            x_t = torch.exp(-(tau**2) * h) * (sigma_t / sigma_prev) * x + gradient_part + noise_part
        else:
            x_t = (alpha_t / alpha_prev) * x + gradient_part + noise_part

        return x_t

    def sample_few_steps(
        self,
        x,
        tau,
        steps=5,
        t_start=None,
        t_end=None,
        skip_type="time",
        skip_order=1,
        predictor_order=3,
        corrector_order=4,
        pc_mode="PEC",
        return_intermediate=False,
    ):
        """
        For the PC-mode, please refer to the wiki page
        https://en.wikipedia.org/wiki/Predictor%E2%80%93corrector_method#PEC_mode_and_PECE_mode
        'PEC' needs one model evaluation per step while 'PECE' needs two model evaluations
        We recommend use pc_mode='PEC' for NFEs is limited. 'PECE' mode is only for test with sufficient NFEs.
        """

        skip_first_step = False
        skip_final_step = True
        lower_order_final = True
        denoise_to_zero = False

        assert pc_mode in ["PEC", "PECE"], "Predictor-corrector mode only supports PEC and PECE"
        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
        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"

        device = x.device
        intermediates = []
        with torch.no_grad():
            assert steps >= max(predictor_order, corrector_order - 1)
            timesteps = self.get_time_steps(
                skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, order=skip_order, device=device
            )
            assert timesteps.shape[0] - 1 == steps
            # Init the initial values.
            step = 0
            t = timesteps[step]
            noise = torch.randn_like(x)
            t_prev_list = [t]
            # do not evaluate if skip_first_step
            if skip_first_step:
                if self.predict_x0:
                    alpha_t = self.noise_schedule.marginal_alpha(t)
                    sigma_t = self.noise_schedule.marginal_std(t)
                    model_prev_list = [(1 - sigma_t) / alpha_t * x]
                else:
                    model_prev_list = [x]
            else:
                model_prev_list = [self.model_fn(x, t)]

            if self.correcting_xt_fn is not None:
                x = self.correcting_xt_fn(x, t, step)
            if return_intermediate:
                intermediates.append(x)

            # determine the first several values
            for step in tqdm(range(1, max(predictor_order, corrector_order - 1))):

                t = timesteps[step]
                predictor_order_used = min(predictor_order, step)
                corrector_order_used = min(corrector_order, step + 1)
                noise = torch.randn_like(x)
                # predictor step
                x_p = self.adams_bashforth_update_few_steps(
                    order=predictor_order_used,
                    x=x,
                    tau=tau(t),
                    model_prev_list=model_prev_list,
                    t_prev_list=t_prev_list,
                    noise=noise,
                    t=t,
                )
                # evaluation step
                model_x = self.model_fn(x_p, t)

                # update model_list
                model_prev_list.append(model_x)
                # corrector step
                if corrector_order > 0:
                    x = self.adams_moulton_update_few_steps(
                        order=corrector_order_used,
                        x=x,
                        tau=tau(t),
                        model_prev_list=model_prev_list,
                        t_prev_list=t_prev_list,
                        noise=noise,
                        t=t,
                    )
                else:
                    x = x_p

                # evaluation step if correction and mode = pece
                if corrector_order > 0:
                    if pc_mode == "PECE":
                        model_x = self.model_fn(x, t)
                        del model_prev_list[-1]
                        model_prev_list.append(model_x)

                if self.correcting_xt_fn is not None:
                    x = self.correcting_xt_fn(x, t, step)
                if return_intermediate:
                    intermediates.append(x)

                t_prev_list.append(t)

            for step in tqdm(range(max(predictor_order, corrector_order - 1), steps + 1)):
                if lower_order_final:
                    predictor_order_used = min(predictor_order, steps - step + 1)
                    corrector_order_used = min(corrector_order, steps - step + 2)

                else:
                    predictor_order_used = predictor_order
                    corrector_order_used = corrector_order
                t = timesteps[step]
                noise = torch.randn_like(x)

                # predictor step
                if skip_final_step and step == steps and not denoise_to_zero:
                    x_p = self.adams_bashforth_update_few_steps(
                        order=predictor_order_used,
                        x=x,
                        tau=0,
                        model_prev_list=model_prev_list,
                        t_prev_list=t_prev_list,
                        noise=noise,
                        t=t,
                    )
                else:
                    x_p = self.adams_bashforth_update_few_steps(
                        order=predictor_order_used,
                        x=x,
                        tau=tau(t),
                        model_prev_list=model_prev_list,
                        t_prev_list=t_prev_list,
                        noise=noise,
                        t=t,
                    )

                # evaluation step
                # do not evaluate if skip_final_step and step = steps
                if not skip_final_step or step < steps:
                    model_x = self.model_fn(x_p, t)

                # update model_list
                # do not update if skip_final_step and step = steps
                if not skip_final_step or step < steps:
                    model_prev_list.append(model_x)

                # corrector step
                # do not correct if skip_final_step and step = steps
                if corrector_order > 0:
                    if not skip_final_step or step < steps:
                        x = self.adams_moulton_update_few_steps(
                            order=corrector_order_used,
                            x=x,
                            tau=tau(t),
                            model_prev_list=model_prev_list,
                            t_prev_list=t_prev_list,
                            noise=noise,
                            t=t,
                        )
                    else:
                        x = x_p
                else:
                    x = x_p

                # evaluation step if mode = pece and step != steps
                if corrector_order > 0:
                    if pc_mode == "PECE" and step < steps:
                        model_x = self.model_fn(x, t)
                        del model_prev_list[-1]
                        model_prev_list.append(model_x)

                if self.correcting_xt_fn is not None:
                    x = self.correcting_xt_fn(x, t, step)
                if return_intermediate:
                    intermediates.append(x)

                t_prev_list.append(t)
                del model_prev_list[0]

            if denoise_to_zero:
                t = torch.ones((1,)).to(device) * t_0
                x = self.denoise_to_zero_fn(x, t)
                if self.correcting_xt_fn is not None:
                    x = self.correcting_xt_fn(x, t, step + 1)
                if return_intermediate:
                    intermediates.append(x)
        if return_intermediate:
            return x, intermediates
        else:
            return x

    def sample_more_steps(
        self,
        x,
        tau,
        steps=20,
        t_start=None,
        t_end=None,
        skip_type="time",
        skip_order=1,
        predictor_order=3,
        corrector_order=4,
        pc_mode="PEC",
        return_intermediate=False,
    ):
        """
        For the PC-mode, please refer to the wiki page
        https://en.wikipedia.org/wiki/Predictor%E2%80%93corrector_method#PEC_mode_and_PECE_mode
        'PEC' needs one model evaluation per step while 'PECE' needs two model evaluations
        We recommend use pc_mode='PEC' for NFEs is limited. 'PECE' mode is only for test with sufficient NFEs.
        """

        skip_first_step = False
        skip_final_step = False
        lower_order_final = True
        denoise_to_zero = True

        assert pc_mode in ["PEC", "PECE"], "Predictor-corrector mode only supports PEC and PECE"
        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
        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"

        device = x.device
        intermediates = []
        with torch.no_grad():
            assert steps >= max(predictor_order, corrector_order - 1)
            timesteps = self.get_time_steps(
                skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, order=skip_order, device=device
            )
            assert timesteps.shape[0] - 1 == steps
            # Init the initial values.
            step = 0
            t = timesteps[step]
            noise = torch.randn_like(x)
            t_prev_list = [t]
            # do not evaluate if skip_first_step
            if skip_first_step:
                if self.predict_x0:
                    alpha_t = self.noise_schedule.marginal_alpha(t)
                    sigma_t = self.noise_schedule.marginal_std(t)
                    model_prev_list = [(1 - sigma_t) / alpha_t * x]
                else:
                    model_prev_list = [x]
            else:
                model_prev_list = [self.model_fn(x, t)]

            if self.correcting_xt_fn is not None:
                x = self.correcting_xt_fn(x, t, step)
            if return_intermediate:
                intermediates.append(x)

            # determine the first several values
            for step in tqdm(range(1, max(predictor_order, corrector_order - 1))):

                t = timesteps[step]
                predictor_order_used = min(predictor_order, step)
                corrector_order_used = min(corrector_order, step + 1)
                noise = torch.randn_like(x)
                # predictor step
                x_p = self.adams_bashforth_update(
                    order=predictor_order_used,
                    x=x,
                    tau=tau(t),
                    model_prev_list=model_prev_list,
                    t_prev_list=t_prev_list,
                    noise=noise,
                    t=t,
                )
                # evaluation step
                model_x = self.model_fn(x_p, t)

                # update model_list
                model_prev_list.append(model_x)
                # corrector step
                if corrector_order > 0:
                    x = self.adams_moulton_update(
                        order=corrector_order_used,
                        x=x,
                        tau=tau(t),
                        model_prev_list=model_prev_list,
                        t_prev_list=t_prev_list,
                        noise=noise,
                        t=t,
                    )
                else:
                    x = x_p

                # evaluation step if mode = pece
                if corrector_order > 0:
                    if pc_mode == "PECE":
                        model_x = self.model_fn(x, t)
                        del model_prev_list[-1]
                        model_prev_list.append(model_x)
                if self.correcting_xt_fn is not None:
                    x = self.correcting_xt_fn(x, t, step)
                if return_intermediate:
                    intermediates.append(x)

                t_prev_list.append(t)

            for step in tqdm(range(max(predictor_order, corrector_order - 1), steps + 1)):
                if lower_order_final:
                    predictor_order_used = min(predictor_order, steps - step + 1)
                    corrector_order_used = min(corrector_order, steps - step + 2)

                else:
                    predictor_order_used = predictor_order
                    corrector_order_used = corrector_order
                t = timesteps[step]
                noise = torch.randn_like(x)

                # predictor step
                if skip_final_step and step == steps and not denoise_to_zero:
                    x_p = self.adams_bashforth_update(
                        order=predictor_order_used,
                        x=x,
                        tau=0,
                        model_prev_list=model_prev_list,
                        t_prev_list=t_prev_list,
                        noise=noise,
                        t=t,
                    )
                else:
                    x_p = self.adams_bashforth_update(
                        order=predictor_order_used,
                        x=x,
                        tau=tau(t),
                        model_prev_list=model_prev_list,
                        t_prev_list=t_prev_list,
                        noise=noise,
                        t=t,
                    )

                # evaluation step
                # do not evaluate if skip_final_step and step = steps
                if not skip_final_step or step < steps:
                    model_x = self.model_fn(x_p, t)

                # update model_list
                # do not update if skip_final_step and step = steps
                if not skip_final_step or step < steps:
                    model_prev_list.append(model_x)

                # corrector step
                # do not correct if skip_final_step and step = steps
                if corrector_order > 0:
                    if not skip_final_step or step < steps:
                        x = self.adams_moulton_update(
                            order=corrector_order_used,
                            x=x,
                            tau=tau(t),
                            model_prev_list=model_prev_list,
                            t_prev_list=t_prev_list,
                            noise=noise,
                            t=t,
                        )
                    else:
                        x = x_p
                else:
                    x = x_p

                # evaluation step if mode = pece and step != steps
                if corrector_order > 0:
                    if pc_mode == "PECE" and step < steps:
                        model_x = self.model_fn(x, t)
                        del model_prev_list[-1]
                        model_prev_list.append(model_x)

                if self.correcting_xt_fn is not None:
                    x = self.correcting_xt_fn(x, t, step)
                if return_intermediate:
                    intermediates.append(x)

                t_prev_list.append(t)
                del model_prev_list[0]

            if denoise_to_zero:
                t = torch.ones((1,)).to(device) * t_0
                x = self.denoise_to_zero_fn(x, t)
                if self.correcting_xt_fn is not None:
                    x = self.correcting_xt_fn(x, t, step + 1)
                if return_intermediate:
                    intermediates.append(x)
        if return_intermediate:
            return x, intermediates
        else:
            return x

    def sample(
        self,
        mode,
        x,
        tau,
        steps,
        t_start=None,
        t_end=None,
        skip_type="time",
        skip_order=1,
        predictor_order=3,
        corrector_order=4,
        pc_mode="PEC",
        return_intermediate=False,
    ):
        """
        For the PC-mode, please refer to the wiki page
        https://en.wikipedia.org/wiki/Predictor%E2%80%93corrector_method#PEC_mode_and_PECE_mode
        'PEC' needs one model evaluation per step while 'PECE' needs two model evaluations
        We recommend use pc_mode='PEC' for NFEs is limited. 'PECE' mode is only for test with sufficient NFEs.

        'few_steps' mode is recommended. The differences between 'few_steps' and 'more_steps' are as below:
        1) 'few_steps' do not correct at final step and do not denoise to zero, while 'more_steps' do these two.
        Thus the NFEs for 'few_steps' = steps, NFEs for 'more_steps' = steps + 2
        For most of the experiments and tasks, we find these two operations do not have much help to sample quality.
        2) 'few_steps' use a rescaling trick as in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf
        We find it will slightly improve the sample quality especially in few steps.
        """
        assert mode in ["few_steps", "more_steps"], "mode must be either 'few_steps' or 'more_steps'"
        if mode == "few_steps":
            return self.sample_few_steps(
                x=x,
                tau=tau,
                steps=steps,
                t_start=t_start,
                t_end=t_end,
                skip_type=skip_type,
                skip_order=skip_order,
                predictor_order=predictor_order,
                corrector_order=corrector_order,
                pc_mode=pc_mode,
                return_intermediate=return_intermediate,
            )
        else:
            return self.sample_more_steps(
                x=x,
                tau=tau,
                steps=steps,
                t_start=t_start,
                t_end=t_end,
                skip_type=skip_type,
                skip_order=skip_order,
                predictor_order=predictor_order,
                corrector_order=corrector_order,
                pc_mode=pc_mode,
                return_intermediate=return_intermediate,
            )


#############################################################
# 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)]