--- title: scheduling_flow_unipc_multistep ---

fastvideo.models.schedulers.scheduling_flow_unipc_multistep.BaseScheduler

BaseScheduler(*args, **kwargs)

fastvideo.models.schedulers.scheduling_flow_unipc_multistep.BaseScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.models.schedulers.scheduling_flow_unipc_multistep.BaseScheduler.set_timesteps

set_timesteps(*args, **kwargs) -> None

fastvideo.models.schedulers.scheduling_flow_unipc_multistep.BaseScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor

fastvideo.models.schedulers.scheduling_flow_unipc_multistep.FlowUniPCMultistepScheduler

FlowUniPCMultistepScheduler(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', **kwargs)

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.

Parameters

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 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.

Source
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

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.
    """
    self._begin_index = begin_index

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(" 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

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.
    """
    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

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.
    """
    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

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_{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

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'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

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 [`~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)

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.
    """
    return sample

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 = [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