def stochastic_adams_moulton_update(
        self,
        this_model_output: torch.Tensor,
        *args,
        last_sample: torch.Tensor,
        last_noise: torch.Tensor,
        this_sample: torch.Tensor,
        order: int,
        tau: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        """
        One step for the SA-Corrector.

        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 order of SA-Corrector at this step.

        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 last_noise is None:
            if len(args) > 2:
                last_noise = args[2]
            else:
                raise ValueError("missing `last_noise` as a required keyword argument")
        if this_sample is None:
            if len(args) > 3:
                this_sample = args[3]
            else:
                raise ValueError("missing `this_sample` as a required keyword argument")
        if order is None:
            if len(args) > 4:
                order = args[4]
            else:
                raise ValueError("missing `order` as a required keyword argument")
        if tau is None:
            if len(args) > 5:
                tau = args[5]
            else:
                raise ValueError("missing `tau` 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
        sigma_t, sigma_s0 = (
            self.sigmas[self.step_index],
            self.sigmas[self.step_index - 1],
        )
        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
        gradient_part = torch.zeros_like(this_sample)
        h = lambda_t - lambda_s0
        lambda_list = []
        for i in range(order):
            si = self.step_index - i
            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
            lambda_list.append(lambda_si)

        model_prev_list = model_output_list + [this_model_output]

        gradient_coefficients = self.get_coefficients_fn(order, lambda_s0, lambda_t, lambda_list, tau)

        x = last_sample

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

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

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

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

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

        One step for the SA-Corrector.

        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 order of SA-Corrector at this step.

        Returns:
            `torch.Tensor`:
                The corrected sample tensor at the current timestep.
        
src/diffusers/schedulers/scheduling_sasolver.py