scheduling_flow_unipc_multistep ¶

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

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

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

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)