scheduling_flow_unipc_multistep

class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
    """
    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
methods the library implements for all schedulers such as loading and saving.

Args:
    num_train_timesteps (`int`, defaults to 1000):
        The number of diffusion steps to train the model.
    solver_order (`int`, default `2`):
        The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
        due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
        unconditional sampling.
    prediction_type (`str`, defaults to "flow_prediction"):
        Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
        the flow of the diffusion process.
    thresholding (`bool`, defaults to `False`):
        Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
        as Stable Diffusion.
    dynamic_thresholding_ratio (`float`, defaults to 0.995):
        The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
    sample_max_value (`float`, defaults to 1.0):
        The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
    predict_x0 (`bool`, defaults to `True`):
        Whether to use the updating algorithm on the predicted x0.
    solver_type (`str`, default `bh2`):
        Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
        otherwise.
    lower_order_final (`bool`, default `True`):
        Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
        stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
    disable_corrector (`list`, default `[]`):
        Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
        and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
        usually disabled during the first few steps.
    solver_p (`SchedulerMixin`, default `None`):
        Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
    use_karras_sigmas (`bool`, *optional*, defaults to `False`):
        Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
        the sigmas are determined according to a sequence of noise levels {σi}.
    use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
        Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
    timestep_spacing (`str`, defaults to `"linspace"`):
        The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
        Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
    steps_offset (`int`, defaults to 0):
        An offset added to the inference steps, as required by some model families.
    final_sigmas_type (`str`, defaults to `"zero"`):
        The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
        sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
"""

_compatibles = [e.name for e in KarrasDiffusionSchedulers]
order = 1

@register_to_config
def __init__(
        self,
        num_train_timesteps: int = 1000,
        solver_order: int = 2,
        prediction_type: str = "flow_prediction",
        shift: float | None = 1.0,
        use_dynamic_shifting=False,
        thresholding: bool = False,
        dynamic_thresholding_ratio: float = 0.995,
        sample_max_value: float = 1.0,
        predict_x0: bool = True,
        solver_type: str = "bh2",
        lower_order_final: bool = True,
        disable_corrector: tuple = (),
        solver_p: SchedulerMixin = None,
        timestep_spacing: str = "linspace",
        steps_offset: int = 0,
        final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
        **kwargs):

    if solver_type not in ["bh1", "bh2"]:
        if solver_type in ["midpoint", "heun", "logrho"]:
            self.register_to_config(solver_type="bh2")
        else:
            raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")

    self.predict_x0 = predict_x0
    # setable values
    self.num_inference_steps: int | None = None
    alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy()
    sigmas = 1.0 - alphas
    sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)

    # Needed when final_sigmas_type == "sigma_min" (kept for compatibility).
    self.alphas_cumprod = torch.from_numpy(alphas).to(dtype=torch.float32)

    if not use_dynamic_shifting:
        # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
        assert shift is not None
        sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore

    self.sigmas = sigmas
    self.timesteps = sigmas * num_train_timesteps
    self.num_train_timesteps = num_train_timesteps

    self.model_outputs = [None] * solver_order
    self.timestep_list: list[Any | None] = [None] * solver_order
    self.lower_order_nums = 0
    self.disable_corrector = list(disable_corrector)
    self.solver_p = solver_p
    self.last_sample = None
    self._step_index: int | None = None
    self._begin_index: int | None = None

    self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
    self.sigma_min = self.sigmas[-1].item()
    self.sigma_max = self.sigmas[0].item()

    BaseScheduler.__init__(self)

@property
def step_index(self):
    """
    The index counter for current timestep. It will increase 1 after each scheduler step.
    """
    return self._step_index

@property
def begin_index(self):
    """
    The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
    """
    return self._begin_index

def set_shift(self, shift: float) -> None:
    self.config.shift = shift

# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
def set_begin_index(self, begin_index: int = 0):
    """
    Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

    Args:
        begin_index (`int`):
            The begin index for the scheduler.
    """
    self._begin_index = begin_index

# Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
def set_timesteps(
    self,
    num_inference_steps: int | None = None,
    device: str | torch.device = None,
    sigmas: list[float] | None = None,
    mu: float | None | None = None,
    shift: float | None | None = None,
    use_karras_sigmas: bool | None = None,
    use_kerras_sigma: bool | None = None,
):
    """
    Sets the discrete timesteps used for the diffusion chain (to be run before inference).
    Args:
        num_inference_steps (`int`):
            Total number of the spacing of the time steps.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
    """

    if self.config.use_dynamic_shifting and mu is None:
        raise ValueError(" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")

    # Cosmos official uses `use_kerras_sigma=True` and a specific EDM sigma schedule.
    # Some external code uses the misspelling `use_kerras_sigma`; support both.
    if use_karras_sigmas is None and use_kerras_sigma is not None:
        use_karras_sigmas = use_kerras_sigma

    if use_karras_sigmas:
        # Force to use the exact sigma used in official EDM sampler:
        # sigma_max=200, sigma_min=0.01, rho=7
        sigma_max = 200.0
        sigma_min = 0.01
        rho = 7.0
        # Match the official Cosmos implementation: Karras/EDM schedule with
        # `num_inference_steps + 1` points, then `final_sigmas_type="zero"`
        # appends the terminal sigma (0.0).
        ramp = np.arange(num_inference_steps + 1, dtype=np.float32) / float(num_inference_steps)
        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
        # Convert EDM sigma to flow-matching sigma in [0, 1).
        sigmas = sigmas / (1.0 + sigmas)
    else:
        if sigmas is None:
            assert num_inference_steps is not None
            sigmas = np.linspace(self.sigma_max, self.sigma_min,
                                 num_inference_steps + 1).copy()[:-1]  # pyright: ignore

    if self.config.use_dynamic_shifting:
        assert mu is not None
        sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
    else:
        if not use_karras_sigmas:
            if shift is None:
                shift = self.config.shift
            assert isinstance(sigmas, np.ndarray)
            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore

    if self.config.final_sigmas_type == "sigma_min":
        sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0])**0.5
    elif self.config.final_sigmas_type == "zero":
        sigma_last = 0
    else:
        raise ValueError(
            f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}")

    timesteps = sigmas * self.config.num_train_timesteps
    sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)  # pyright: ignore

    self.sigmas = torch.from_numpy(sigmas)
    self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)

    self.num_inference_steps = len(timesteps)

    self.model_outputs = [
        None,
    ] * self.config.solver_order
    self.lower_order_nums = 0
    self.last_sample = None
    if self.solver_p:
        self.solver_p.set_timesteps(self.num_inference_steps, device=device)

    # add an index counter for schedulers that allow duplicated timesteps
    self._step_index = None
    self._begin_index = None
    self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication

# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
    """
    "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
    prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
    s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
    pixels from saturation at each step. We find that dynamic thresholding results in significantly better
    photorealism as well as better image-text alignment, especially when using very large guidance weights."

    https://arxiv.org/abs/2205.11487
    """
    dtype = sample.dtype
    batch_size, channels, *remaining_dims = sample.shape

    if dtype not in (torch.float32, torch.float64):
        sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half

    # Flatten sample for doing quantile calculation along each image
    sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))

    abs_sample = sample.abs()  # "a certain percentile absolute pixel value"

    s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
    s = torch.clamp(
        s, min=1,
        max=self.config.sample_max_value)  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
    s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
    sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"

    sample = sample.reshape(batch_size, channels, *remaining_dims)
    sample = sample.to(dtype)

    return sample

# Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
def _sigma_to_t(self, sigma):
    return sigma * self.config.num_train_timesteps

def _sigma_to_alpha_sigma_t(self, sigma) -> tuple[Any, Any]:
    return 1 - sigma, sigma

# Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
    return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)

def convert_model_output(
    self,
    model_output: torch.Tensor,
    *args,
    sample: torch.Tensor = None,
    **kwargs,
) -> torch.Tensor:
    r"""
    Convert the model output to the corresponding type the UniPC algorithm needs.

    Args:
        model_output (`torch.Tensor`):
            The direct output from the learned diffusion model.
        timestep (`int`):
            The current discrete timestep in the diffusion chain.
        sample (`torch.Tensor`):
            A current instance of a sample created by the diffusion process.

    Returns:
        `torch.Tensor`:
            The converted model output.
    """
    timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
    if sample is None:
        if len(args) > 1:
            sample = args[1]
        else:
            raise ValueError("missing `sample` as a required keyword argument")
    if timestep is not None:
        deprecate(
            "timesteps",
            "1.0.0",
            "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
        )

    sigma = self.sigmas[self.step_index]
    alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

    if self.predict_x0:
        if self.config.prediction_type == "flow_prediction":
            sigma_t = self.sigmas[self.step_index]
            x0_pred = sample - sigma_t * model_output
        else:
            raise ValueError(
                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
                " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler.")

        if self.config.thresholding:
            x0_pred = self._threshold_sample(x0_pred)

        return x0_pred
    else:
        if self.config.prediction_type == "flow_prediction":
            sigma_t = self.sigmas[self.step_index]
            epsilon = sample - (1 - sigma_t) * model_output
        else:
            raise ValueError(
                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
                " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler.")

        if self.config.thresholding:
            sigma_t = self.sigmas[self.step_index]
            x0_pred = sample - sigma_t * model_output
            x0_pred = self._threshold_sample(x0_pred)
            epsilon = model_output + x0_pred

        return epsilon

def multistep_uni_p_bh_update(
    self,
    model_output: torch.Tensor,
    *args,
    sample: torch.Tensor = None,
    order: int | None = None,  # pyright: ignore
    **kwargs,
) -> torch.Tensor:
    """
    One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.

    Args:
        model_output (`torch.Tensor`):
            The direct output from the learned diffusion model at the current timestep.
        prev_timestep (`int`):
            The previous discrete timestep in the diffusion chain.
        sample (`torch.Tensor`):
            A current instance of a sample created by the diffusion process.
        order (`int`):
            The order of UniP at this timestep (corresponds to the *p* in UniPC-p).

    Returns:
        `torch.Tensor`:
            The sample tensor at the previous timestep.
    """
    prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
    if sample is None:
        if len(args) > 1:
            sample = args[1]
        else:
            raise ValueError(" missing `sample` as a required keyword argument")
    if order is None:
        if len(args) > 2:
            order = args[2]
        else:
            raise ValueError(" missing `order` as a required keyword argument")
    if prev_timestep is not None:
        deprecate(
            "prev_timestep",
            "1.0.0",
            "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
        )
    model_output_list = self.model_outputs

    s0 = self.timestep_list[-1]
    m0 = model_output_list[-1]
    x = sample

    if self.solver_p:
        x_t = self.solver_p.step(model_output, s0, x).prev_sample
        return x_t

    sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]  # pyright: ignore
    alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
    alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

    # Numerical safety
    eps = 1e-12
    lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(torch.clamp(sigma_t, min=eps))
    lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(torch.clamp(sigma_s0, min=eps))

    h = lambda_t - lambda_s0
    device = sample.device

    rks = []
    D1s: list[Any] | None = []
    for i in range(1, order):
        si = self.step_index - i  # pyright: ignore
        mi = model_output_list[-(i + 1)]
        alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
        lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(torch.clamp(sigma_si, min=eps))
        rk = (lambda_si - lambda_s0) / h
        rks.append(rk)
        assert mi is not None
        D1s.append((mi - m0) / rk)  # pyright: ignore

    rks.append(1.0)
    rks = torch.tensor(rks, device=device)

    R = []
    b = []

    hh = -h if self.predict_x0 else h
    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 self.config.solver_type == "bh1":
        B_h = hh
    elif self.config.solver_type == "bh2":
        B_h = torch.expm1(hh)
    else:
        raise NotImplementedError()

    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=device)

    if D1s is not None and len(D1s) > 0:
        D1s = torch.stack(D1s, dim=1)  # (B, K)
        # for order 2, we use a simplified version
        if order == 2:
            rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
        else:
            assert isinstance(R, torch.Tensor)
            rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
    else:
        D1s = None

    if self.predict_x0:
        x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
        pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0  # pyright: ignore
        x_t = x_t_ - alpha_t * B_h * pred_res
    else:
        x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
        pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0  # pyright: ignore
        x_t = x_t_ - sigma_t * B_h * pred_res

    x_t = x_t.to(x.dtype)
    return x_t

def multistep_uni_c_bh_update(
    self,
    this_model_output: torch.Tensor,
    *args,
    last_sample: torch.Tensor = None,
    this_sample: torch.Tensor = None,
    order: int | None = None,  # pyright: ignore
    **kwargs,
) -> torch.Tensor:
    """
    One step for the UniC (B(h) version).

    Args:
        this_model_output (`torch.Tensor`):
            The model outputs at `x_t`.
        this_timestep (`int`):
            The current timestep `t`.
        last_sample (`torch.Tensor`):
            The generated sample before the last predictor `x_{t-1}`.
        this_sample (`torch.Tensor`):
            The generated sample after the last predictor `x_{t}`.
        order (`int`):
            The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.

    Returns:
        `torch.Tensor`:
            The corrected sample tensor at the current timestep.
    """
    this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
    if last_sample is None:
        if len(args) > 1:
            last_sample = args[1]
        else:
            raise ValueError(" missing`last_sample` as a required keyword argument")
    if this_sample is None:
        if len(args) > 2:
            this_sample = args[2]
        else:
            raise ValueError(" missing`this_sample` as a required keyword argument")
    if order is None:
        if len(args) > 3:
            order = args[3]
        else:
            raise ValueError(" missing`order` as a required keyword argument")
    if this_timestep is not None:
        deprecate(
            "this_timestep",
            "1.0.0",
            "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
        )

    model_output_list = self.model_outputs

    m0 = model_output_list[-1]
    x = last_sample
    x_t = this_sample
    model_t = this_model_output

    sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]  # pyright: ignore
    alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
    alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

    eps = 1e-12
    lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(torch.clamp(sigma_t, min=eps))
    lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(torch.clamp(sigma_s0, min=eps))

    h = lambda_t - lambda_s0
    device = this_sample.device

    rks = []
    D1s: list[Any] | None = []
    for i in range(1, order):
        si = self.step_index - (i + 1)  # pyright: ignore
        mi = model_output_list[-(i + 1)]
        alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
        lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(torch.clamp(sigma_si, min=eps))
        rk = (lambda_si - lambda_s0) / h
        rks.append(rk)
        assert mi is not None
        D1s.append((mi - m0) / rk)  # pyright: ignore

    rks.append(1.0)
    rks = torch.tensor(rks, device=device)

    R = []
    b = []

    hh = -h if self.predict_x0 else h
    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 self.config.solver_type == "bh1":
        B_h = hh
    elif self.config.solver_type == "bh2":
        B_h = torch.expm1(hh)
    else:
        raise NotImplementedError()

    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=device)

    D1s = torch.stack(D1s, dim=1) if D1s is not None and len(D1s) > 0 else None

    # for order 1, we use a simplified version
    if order == 1:
        rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
    else:
        rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)

    if self.predict_x0:
        x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
        corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0
        D1_t = model_t - m0
        x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
    else:
        x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
        corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0
        D1_t = model_t - m0
        x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
    x_t = x_t.to(x.dtype)
    return x_t

def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
    if schedule_timesteps is None:
        schedule_timesteps = self.timesteps

    indices = (schedule_timesteps == timestep).nonzero()

    # The sigma index that is taken for the **very** first `step`
    # is always the second index (or the last index if there is only 1)
    # This way we can ensure we don't accidentally skip a sigma in
    # case we start in the middle of the denoising schedule (e.g. for image-to-image)
    pos = 1 if len(indices) > 1 else 0
    step_index: int = indices[pos].item()

    return step_index

# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
def _init_step_index(self, timestep) -> None:
    """
    Initialize the step_index counter for the scheduler.
    """

    if self.begin_index is None:
        if isinstance(timestep, torch.Tensor):
            timestep = timestep.to(self.timesteps.device)
        self._step_index = self.index_for_timestep(timestep)
    else:
        self._step_index = self._begin_index

def step(self,
         model_output: torch.Tensor,
         timestep: int | torch.Tensor,
         sample: torch.Tensor,
         return_dict: bool = True,
         generator=None) -> SchedulerOutput | tuple:
    """
    Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
    the multistep UniPC.

    Args:
        model_output (`torch.Tensor`):
            The direct output from learned diffusion model.
        timestep (`int`):
            The current discrete timestep in the diffusion chain.
        sample (`torch.Tensor`):
            A current instance of a sample created by the diffusion process.
        return_dict (`bool`):
            Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.

    Returns:
        [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
            If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
            tuple is returned where the first element is the sample tensor.

    """
    if self.num_inference_steps is None:
        raise ValueError(
            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler")

    if self.step_index is None:
        self._init_step_index(timestep)

    use_corrector = (
        self.step_index > 0 and self.step_index - 1 not in self.disable_corrector
        and self.last_sample is not None  # pyright: ignore
    )

    model_output_convert = self.convert_model_output(model_output, sample=sample)

    if use_corrector:
        sample = self.multistep_uni_c_bh_update(
            this_model_output=model_output_convert,
            last_sample=self.last_sample,
            this_sample=sample,
            order=self.this_order,
        )

    for i in range(self.config.solver_order - 1):
        self.model_outputs[i] = self.model_outputs[i + 1]
        self.timestep_list[i] = self.timestep_list[i + 1]

    self.model_outputs[-1] = model_output_convert
    self.timestep_list[-1] = timestep  # pyright: ignore

    if self.config.lower_order_final:
        this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)  # pyright: ignore
    else:
        this_order = self.config.solver_order

    self.this_order: int = min(this_order, self.lower_order_nums + 1)  # warmup for multistep
    assert self.this_order > 0

    self.last_sample = sample
    prev_sample = self.multistep_uni_p_bh_update(
        model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
        sample=sample,
        order=self.this_order,
    )

    if self.lower_order_nums < self.config.solver_order:
        self.lower_order_nums += 1

    # upon completion increase step index by one
    assert self._step_index is not None
    self._step_index += 1  # pyright: ignore

    if not return_dict:
        return (prev_sample, )

    return SchedulerOutput(prev_sample=prev_sample)

def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
    """
    Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
    current timestep.

    Args:
        sample (`torch.Tensor`):
            The input sample.

    Returns:
        `torch.Tensor`:
            A scaled input sample.
    """
    return sample

# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
def add_noise(
    self,
    original_samples: torch.Tensor,
    noise: torch.Tensor,
    timesteps: torch.IntTensor,
) -> torch.Tensor:
    # Make sure sigmas and timesteps have the same device and dtype as original_samples
    sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
    if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
        # mps does not support float64
        schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
        timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
    else:
        schedule_timesteps = self.timesteps.to(original_samples.device)
        timesteps = timesteps.to(original_samples.device)

    # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
    if self.begin_index is None:
        step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
    elif self.step_index is not None:
        # add_noise is called after first denoising step (for inpainting)
        step_indices = [self.step_index] * timesteps.shape[0]
    else:
        # add noise is called before first denoising step to create initial latent(img2img)
        step_indices = [self.begin_index] * timesteps.shape[0]

    sigma = sigmas[step_indices].flatten()
    while len(sigma.shape) < len(original_samples.shape):
        sigma = sigma.unsqueeze(-1)

    alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
    noisy_samples = alpha_t * original_samples + sigma_t * noise
    return noisy_samples

def __len__(self):
    return self.config.num_train_timesteps</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.step_index
step_index
The index counter for current timestep. It will increase 1 after each scheduler step.
Source
def step_index(self):
    """
    The index counter for current timestep. It will increase 1 after each scheduler step.
    """
    return self._step_index
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.begin_index
begin_index
The index for the first timestep. It should be set from pipeline with set_begin_index method.
Source
def begin_index(self):
    """
    The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
    """
    return self._begin_index
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.set_shift
set_shift(shift: float) -> None
Source
def set_shift(self, shift: float) -> None:
    self.config.shift = shift
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.set_begin_index
set_begin_index(begin_index: int = 0)
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
Parameters

begin_index: `int`
The begin index for the scheduler.

Source
def set_begin_index(self, begin_index: int = 0):
    """
    Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
Args:
    begin_index (`int`):
        The begin index for the scheduler.
&quot;&quot;&quot;
self._begin_index = begin_index</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.set_timesteps
set_timesteps(num_inference_steps: int | None = None, device: str | torch.device = None, sigmas: list[float] | None = None, mu: float | None | None = None, shift: float | None | None = None, use_karras_sigmas: bool | None = None, use_kerras_sigma: bool | None = None)
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
Args:
    num_inference_steps (int):
        Total number of the spacing of the time steps.
    device (str or torch.device, optional):
        The device to which the timesteps should be moved to. If None, the timesteps are not moved.
Source
def set_timesteps(
    self,
    num_inference_steps: int | None = None,
    device: str | torch.device = None,
    sigmas: list[float] | None = None,
    mu: float | None | None = None,
    shift: float | None | None = None,
    use_karras_sigmas: bool | None = None,
    use_kerras_sigma: bool | None = None,
):
    """
    Sets the discrete timesteps used for the diffusion chain (to be run before inference).
    Args:
        num_inference_steps (`int`):
            Total number of the spacing of the time steps.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
    """
if self.config.use_dynamic_shifting and mu is None:
    raise ValueError(&quot; you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`&quot;)

# Cosmos official uses `use_kerras_sigma=True` and a specific EDM sigma schedule.
# Some external code uses the misspelling `use_kerras_sigma`; support both.
if use_karras_sigmas is None and use_kerras_sigma is not None:
    use_karras_sigmas = use_kerras_sigma

if use_karras_sigmas:
    # Force to use the exact sigma used in official EDM sampler:
    # sigma_max=200, sigma_min=0.01, rho=7
    sigma_max = 200.0
    sigma_min = 0.01
    rho = 7.0
    # Match the official Cosmos implementation: Karras/EDM schedule with
    # `num_inference_steps + 1` points, then `final_sigmas_type=&quot;zero&quot;`
    # appends the terminal sigma (0.0).
    ramp = np.arange(num_inference_steps + 1, dtype=np.float32) / float(num_inference_steps)
    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
    # Convert EDM sigma to flow-matching sigma in &#91;0, 1).
    sigmas = sigmas / (1.0 + sigmas)
else:
    if sigmas is None:
        assert num_inference_steps is not None
        sigmas = np.linspace(self.sigma_max, self.sigma_min,
                             num_inference_steps + 1).copy()&#91;:-1&#93;  # pyright: ignore

if self.config.use_dynamic_shifting:
    assert mu is not None
    sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
else:
    if not use_karras_sigmas:
        if shift is None:
            shift = self.config.shift
        assert isinstance(sigmas, np.ndarray)
        sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore

if self.config.final_sigmas_type == &quot;sigma_min&quot;:
    sigma_last = ((1 - self.alphas_cumprod&#91;0&#93;) / self.alphas_cumprod&#91;0&#93;)**0.5
elif self.config.final_sigmas_type == &quot;zero&quot;:
    sigma_last = 0
else:
    raise ValueError(
        f&quot;`final_sigmas_type` must be one of &#x27;zero&#x27;, or &#x27;sigma_min&#x27;, but got &#123;self.config.final_sigmas_type&#125;&quot;)

timesteps = sigmas * self.config.num_train_timesteps
sigmas = np.concatenate(&#91;sigmas, &#91;sigma_last&#93;&#93;).astype(np.float32)  # pyright: ignore

self.sigmas = torch.from_numpy(sigmas)
self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)

self.num_inference_steps = len(timesteps)

self.model_outputs = &#91;
    None,
&#93; * self.config.solver_order
self.lower_order_nums = 0
self.last_sample = None
if self.solver_p:
    self.solver_p.set_timesteps(self.num_inference_steps, device=device)

# add an index counter for schedulers that allow duplicated timesteps
self._step_index = None
self._begin_index = None
self.sigmas = self.sigmas.to(&quot;cpu&quot;)  # to avoid too much CPU/GPU communication</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.time_shift
time_shift(mu: float, sigma: float, t: torch.Tensor)
Source
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
    return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.convert_model_output
convert_model_output(model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs) -> torch.Tensor
Convert the model output to the corresponding type the UniPC algorithm needs.
Parameters

model_output: `torch.Tensor`
The direct output from the learned diffusion model.
timestep: `int`
The current discrete timestep in the diffusion chain.
sample: `torch.Tensor`
A current instance of a sample created by the diffusion process.

Returns

torch.Tensor
torch.Tensor:
The converted model output.

Source
def convert_model_output(
    self,
    model_output: torch.Tensor,
    *args,
    sample: torch.Tensor = None,
    **kwargs,
) -> torch.Tensor:
    r"""
    Convert the model output to the corresponding type the UniPC algorithm needs.
Args:
    model_output (`torch.Tensor`):
        The direct output from the learned diffusion model.
    timestep (`int`):
        The current discrete timestep in the diffusion chain.
    sample (`torch.Tensor`):
        A current instance of a sample created by the diffusion process.

Returns:
    `torch.Tensor`:
        The converted model output.
&quot;&quot;&quot;
timestep = args&#91;0&#93; if len(args) &gt; 0 else kwargs.pop(&quot;timestep&quot;, None)
if sample is None:
    if len(args) &gt; 1:
        sample = args&#91;1&#93;
    else:
        raise ValueError(&quot;missing `sample` as a required keyword argument&quot;)
if timestep is not None:
    deprecate(
        &quot;timesteps&quot;,
        &quot;1.0.0&quot;,
        &quot;Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`&quot;,
    )

sigma = self.sigmas&#91;self.step_index&#93;
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

if self.predict_x0:
    if self.config.prediction_type == &quot;flow_prediction&quot;:
        sigma_t = self.sigmas&#91;self.step_index&#93;
        x0_pred = sample - sigma_t * model_output
    else:
        raise ValueError(
            f&quot;prediction_type given as &#123;self.config.prediction_type&#125; must be one of `epsilon`, `sample`,&quot;
            &quot; `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler.&quot;)

    if self.config.thresholding:
        x0_pred = self._threshold_sample(x0_pred)

    return x0_pred
else:
    if self.config.prediction_type == &quot;flow_prediction&quot;:
        sigma_t = self.sigmas&#91;self.step_index&#93;
        epsilon = sample - (1 - sigma_t) * model_output
    else:
        raise ValueError(
            f&quot;prediction_type given as &#123;self.config.prediction_type&#125; must be one of `epsilon`, `sample`,&quot;
            &quot; `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler.&quot;)

    if self.config.thresholding:
        sigma_t = self.sigmas&#91;self.step_index&#93;
        x0_pred = sample - sigma_t * model_output
        x0_pred = self._threshold_sample(x0_pred)
        epsilon = model_output + x0_pred

    return epsilon</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.multistep_uni_p_bh_update
multistep_uni_p_bh_update(model_output: torch.Tensor, *args, sample: torch.Tensor = None, order: int | None = None, **kwargs) -> torch.Tensor
One step for the UniP (B(h) version). Alternatively, self.solver_p is used if is specified.
Parameters

model_output: `torch.Tensor`
The direct output from the learned diffusion model at the current timestep.
prev_timestep: `int`
The previous discrete timestep in the diffusion chain.
sample: `torch.Tensor`
A current instance of a sample created by the diffusion process.
order: `int`
The order of UniP at this timestep (corresponds to the p in UniPC-p).

Returns

torch.Tensor
torch.Tensor:
The sample tensor at the previous timestep.

Source
def multistep_uni_p_bh_update(
    self,
    model_output: torch.Tensor,
    *args,
    sample: torch.Tensor = None,
    order: int | None = None,  # pyright: ignore
    **kwargs,
) -> torch.Tensor:
    """
    One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
Args:
    model_output (`torch.Tensor`):
        The direct output from the learned diffusion model at the current timestep.
    prev_timestep (`int`):
        The previous discrete timestep in the diffusion chain.
    sample (`torch.Tensor`):
        A current instance of a sample created by the diffusion process.
    order (`int`):
        The order of UniP at this timestep (corresponds to the *p* in UniPC-p).

Returns:
    `torch.Tensor`:
        The sample tensor at the previous timestep.
&quot;&quot;&quot;
prev_timestep = args&#91;0&#93; if len(args) &gt; 0 else kwargs.pop(&quot;prev_timestep&quot;, None)
if sample is None:
    if len(args) &gt; 1:
        sample = args&#91;1&#93;
    else:
        raise ValueError(&quot; missing `sample` as a required keyword argument&quot;)
if order is None:
    if len(args) &gt; 2:
        order = args&#91;2&#93;
    else:
        raise ValueError(&quot; missing `order` as a required keyword argument&quot;)
if prev_timestep is not None:
    deprecate(
        &quot;prev_timestep&quot;,
        &quot;1.0.0&quot;,
        &quot;Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`&quot;,
    )
model_output_list = self.model_outputs

s0 = self.timestep_list&#91;-1&#93;
m0 = model_output_list&#91;-1&#93;
x = sample

if self.solver_p:
    x_t = self.solver_p.step(model_output, s0, x).prev_sample
    return x_t

sigma_t, sigma_s0 = self.sigmas&#91;self.step_index + 1&#93;, self.sigmas&#91;self.step_index&#93;  # pyright: ignore
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

# Numerical safety
eps = 1e-12
lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(torch.clamp(sigma_t, min=eps))
lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(torch.clamp(sigma_s0, min=eps))

h = lambda_t - lambda_s0
device = sample.device

rks = &#91;&#93;
D1s: list&#91;Any&#93; | None = &#91;&#93;
for i in range(1, order):
    si = self.step_index - i  # pyright: ignore
    mi = model_output_list&#91;-(i + 1)&#93;
    alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas&#91;si&#93;)
    lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(torch.clamp(sigma_si, min=eps))
    rk = (lambda_si - lambda_s0) / h
    rks.append(rk)
    assert mi is not None
    D1s.append((mi - m0) / rk)  # pyright: ignore

rks.append(1.0)
rks = torch.tensor(rks, device=device)

R = &#91;&#93;
b = &#91;&#93;

hh = -h if self.predict_x0 else h
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 self.config.solver_type == &quot;bh1&quot;:
    B_h = hh
elif self.config.solver_type == &quot;bh2&quot;:
    B_h = torch.expm1(hh)
else:
    raise NotImplementedError()

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=device)

if D1s is not None and len(D1s) &gt; 0:
    D1s = torch.stack(D1s, dim=1)  # (B, K)
    # for order 2, we use a simplified version
    if order == 2:
        rhos_p = torch.tensor(&#91;0.5&#93;, dtype=x.dtype, device=device)
    else:
        assert isinstance(R, torch.Tensor)
        rhos_p = torch.linalg.solve(R&#91;:-1, :-1&#93;, b&#91;:-1&#93;).to(device).to(x.dtype)
else:
    D1s = None

if self.predict_x0:
    x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
    pred_res = torch.einsum(&quot;k,bkc...-&gt;bc...&quot;, rhos_p, D1s) if D1s is not None else 0  # pyright: ignore
    x_t = x_t_ - alpha_t * B_h * pred_res
else:
    x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
    pred_res = torch.einsum(&quot;k,bkc...-&gt;bc...&quot;, rhos_p, D1s) if D1s is not None else 0  # pyright: ignore
    x_t = x_t_ - sigma_t * B_h * pred_res

x_t = x_t.to(x.dtype)
return x_t</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.multistep_uni_c_bh_update
multistep_uni_c_bh_update(this_model_output: torch.Tensor, *args, last_sample: torch.Tensor = None, this_sample: torch.Tensor = None, order: int | None = None, **kwargs) -> torch.Tensor
One step for the UniC (B(h) version).
Parameters

this_model_output: `torch.Tensor`
The model outputs at x_t.
this_timestep: `int`
The current timestep t.
last_sample: `torch.Tensor`
The generated sample before the last predictor x_{t-1}.
this_sample: `torch.Tensor`
The generated sample after the last predictor x_{t}.
order: `int`
The p of UniC-p at this step. The effective order of accuracy should be order + 1.

Returns

torch.Tensor
torch.Tensor:
The corrected sample tensor at the current timestep.

Source
def multistep_uni_c_bh_update(
    self,
    this_model_output: torch.Tensor,
    *args,
    last_sample: torch.Tensor = None,
    this_sample: torch.Tensor = None,
    order: int | None = None,  # pyright: ignore
    **kwargs,
) -> torch.Tensor:
    """
    One step for the UniC (B(h) version).
Args:
    this_model_output (`torch.Tensor`):
        The model outputs at `x_t`.
    this_timestep (`int`):
        The current timestep `t`.
    last_sample (`torch.Tensor`):
        The generated sample before the last predictor `x_&#123;t-1&#125;`.
    this_sample (`torch.Tensor`):
        The generated sample after the last predictor `x_&#123;t&#125;`.
    order (`int`):
        The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.

Returns:
    `torch.Tensor`:
        The corrected sample tensor at the current timestep.
&quot;&quot;&quot;
this_timestep = args&#91;0&#93; if len(args) &gt; 0 else kwargs.pop(&quot;this_timestep&quot;, None)
if last_sample is None:
    if len(args) &gt; 1:
        last_sample = args&#91;1&#93;
    else:
        raise ValueError(&quot; missing`last_sample` as a required keyword argument&quot;)
if this_sample is None:
    if len(args) &gt; 2:
        this_sample = args&#91;2&#93;
    else:
        raise ValueError(&quot; missing`this_sample` as a required keyword argument&quot;)
if order is None:
    if len(args) &gt; 3:
        order = args&#91;3&#93;
    else:
        raise ValueError(&quot; missing`order` as a required keyword argument&quot;)
if this_timestep is not None:
    deprecate(
        &quot;this_timestep&quot;,
        &quot;1.0.0&quot;,
        &quot;Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`&quot;,
    )

model_output_list = self.model_outputs

m0 = model_output_list&#91;-1&#93;
x = last_sample
x_t = this_sample
model_t = this_model_output

sigma_t, sigma_s0 = self.sigmas&#91;self.step_index&#93;, self.sigmas&#91;self.step_index - 1&#93;  # pyright: ignore
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

eps = 1e-12
lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(torch.clamp(sigma_t, min=eps))
lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(torch.clamp(sigma_s0, min=eps))

h = lambda_t - lambda_s0
device = this_sample.device

rks = &#91;&#93;
D1s: list&#91;Any&#93; | None = &#91;&#93;
for i in range(1, order):
    si = self.step_index - (i + 1)  # pyright: ignore
    mi = model_output_list&#91;-(i + 1)&#93;
    alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas&#91;si&#93;)
    lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(torch.clamp(sigma_si, min=eps))
    rk = (lambda_si - lambda_s0) / h
    rks.append(rk)
    assert mi is not None
    D1s.append((mi - m0) / rk)  # pyright: ignore

rks.append(1.0)
rks = torch.tensor(rks, device=device)

R = &#91;&#93;
b = &#91;&#93;

hh = -h if self.predict_x0 else h
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 self.config.solver_type == &quot;bh1&quot;:
    B_h = hh
elif self.config.solver_type == &quot;bh2&quot;:
    B_h = torch.expm1(hh)
else:
    raise NotImplementedError()

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=device)

D1s = torch.stack(D1s, dim=1) if D1s is not None and len(D1s) &gt; 0 else None

# for order 1, we use a simplified version
if order == 1:
    rhos_c = torch.tensor(&#91;0.5&#93;, dtype=x.dtype, device=device)
else:
    rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)

if self.predict_x0:
    x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
    corr_res = torch.einsum(&quot;k,bkc...-&gt;bc...&quot;, rhos_c&#91;:-1&#93;, D1s) if D1s is not None else 0
    D1_t = model_t - m0
    x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c&#91;-1&#93; * D1_t)
else:
    x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
    corr_res = torch.einsum(&quot;k,bkc...-&gt;bc...&quot;, rhos_c&#91;:-1&#93;, D1s) if D1s is not None else 0
    D1_t = model_t - m0
    x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c&#91;-1&#93; * D1_t)
x_t = x_t.to(x.dtype)
return x_t</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.index_for_timestep
index_for_timestep(timestep, schedule_timesteps = None) -> int
Source
def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
    if schedule_timesteps is None:
        schedule_timesteps = self.timesteps
indices = (schedule_timesteps == timestep).nonzero()

# The sigma index that is taken for the **very** first `step`
# is always the second index (or the last index if there is only 1)
# This way we can ensure we don&#x27;t accidentally skip a sigma in
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
pos = 1 if len(indices) &gt; 1 else 0
step_index: int = indices&#91;pos&#93;.item()

return step_index</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.step
step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, generator = None) -> SchedulerOutput | tuple
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep UniPC.
Parameters

model_output: `torch.Tensor`
The direct output from learned diffusion model.
timestep: `int`
The current discrete timestep in the diffusion chain.
sample: `torch.Tensor`
A current instance of a sample created by the diffusion process.
return_dict: `bool`
Whether or not to return a [~schedulers.scheduling_utils.SchedulerOutput] or tuple.

Returns

SchedulerOutput | tuple
[~schedulers.scheduling_utils.SchedulerOutput] or tuple:
If return_dict is True, [~schedulers.scheduling_utils.SchedulerOutput] is returned, otherwise a
tuple is returned where the first element is the sample tensor.

Source
def step(self,
         model_output: torch.Tensor,
         timestep: int | torch.Tensor,
         sample: torch.Tensor,
         return_dict: bool = True,
         generator=None) -> SchedulerOutput | tuple:
    """
    Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
    the multistep UniPC.
Args:
    model_output (`torch.Tensor`):
        The direct output from learned diffusion model.
    timestep (`int`):
        The current discrete timestep in the diffusion chain.
    sample (`torch.Tensor`):
        A current instance of a sample created by the diffusion process.
    return_dict (`bool`):
        Whether or not to return a &#91;`~schedulers.scheduling_utils.SchedulerOutput`&#93; or `tuple`.

Returns:
    &#91;`~schedulers.scheduling_utils.SchedulerOutput`&#93; or `tuple`:
        If return_dict is `True`, &#91;`~schedulers.scheduling_utils.SchedulerOutput`&#93; is returned, otherwise a
        tuple is returned where the first element is the sample tensor.

&quot;&quot;&quot;
if self.num_inference_steps is None:
    raise ValueError(
        &quot;Number of inference steps is &#x27;None&#x27;, you need to run &#x27;set_timesteps&#x27; after creating the scheduler&quot;)

if self.step_index is None:
    self._init_step_index(timestep)

use_corrector = (
    self.step_index &gt; 0 and self.step_index - 1 not in self.disable_corrector
    and self.last_sample is not None  # pyright: ignore
)

model_output_convert = self.convert_model_output(model_output, sample=sample)

if use_corrector:
    sample = self.multistep_uni_c_bh_update(
        this_model_output=model_output_convert,
        last_sample=self.last_sample,
        this_sample=sample,
        order=self.this_order,
    )

for i in range(self.config.solver_order - 1):
    self.model_outputs&#91;i&#93; = self.model_outputs&#91;i + 1&#93;
    self.timestep_list&#91;i&#93; = self.timestep_list&#91;i + 1&#93;

self.model_outputs&#91;-1&#93; = model_output_convert
self.timestep_list&#91;-1&#93; = timestep  # pyright: ignore

if self.config.lower_order_final:
    this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)  # pyright: ignore
else:
    this_order = self.config.solver_order

self.this_order: int = min(this_order, self.lower_order_nums + 1)  # warmup for multistep
assert self.this_order &gt; 0

self.last_sample = sample
prev_sample = self.multistep_uni_p_bh_update(
    model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
    sample=sample,
    order=self.this_order,
)

if self.lower_order_nums &lt; self.config.solver_order:
    self.lower_order_nums += 1

# upon completion increase step index by one
assert self._step_index is not None
self._step_index += 1  # pyright: ignore

if not return_dict:
    return (prev_sample, )

return SchedulerOutput(prev_sample=prev_sample)</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.scale_model_input
scale_model_input(sample: torch.Tensor, *args, **kwargs) -> torch.Tensor
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
Parameters

sample: `torch.Tensor`
The input sample.

Returns

torch.Tensor
torch.Tensor:
A scaled input sample.

Source
def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
    """
    Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
    current timestep.
Args:
    sample (`torch.Tensor`):
        The input sample.

Returns:
    `torch.Tensor`:
        A scaled input sample.
&quot;&quot;&quot;
return sample</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler.add_noise
add_noise(original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor
Source
def add_noise(
    self,
    original_samples: torch.Tensor,
    noise: torch.Tensor,
    timesteps: torch.IntTensor,
) -> torch.Tensor:
    # Make sure sigmas and timesteps have the same device and dtype as original_samples
    sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
    if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
        # mps does not support float64
        schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
        timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
    else:
        schedule_timesteps = self.timesteps.to(original_samples.device)
        timesteps = timesteps.to(original_samples.device)
# begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
if self.begin_index is None:
    step_indices = &#91;self.index_for_timestep(t, schedule_timesteps) for t in timesteps&#93;
elif self.step_index is not None:
    # add_noise is called after first denoising step (for inpainting)
    step_indices = &#91;self.step_index&#93; * timesteps.shape&#91;0&#93;
else:
    # add noise is called before first denoising step to create initial latent(img2img)
    step_indices = &#91;self.begin_index&#93; * timesteps.shape&#91;0&#93;

sigma = sigmas&#91;step_indices&#93;.flatten()
while len(sigma.shape) &lt; len(original_samples.shape):
    sigma = sigma.unsqueeze(-1)

alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
noisy_samples = alpha_t * original_samples + sigma_t * noise
return noisy_samples</code></pre></details>