chore: import upstream snapshot with attribution

invokeai/backend/rectified_flow/er_sde_scheduler.py

@@ -0,0 +1,601 @@
+"""ER-SDE (Extended Reverse-time SDE) ``diffusers`` scheduler.
+
+Implements the multistep Taylor-expansion solver from:
+
+    Cui, Q., Zhang, X., Lu, Z., & Liao, Q. (2023).
+    Elucidating the solution space of extended reverse-time SDE
+    for diffusion models. arXiv:2309.06169.
+    https://arxiv.org/abs/2309.06169
+
+Reference implementation (MIT-licensed):
+    https://github.com/QinpengCui/ER-SDE-Solver/blob/main/er_sde_solver.py
+
+This scheduler unifies two regimes under a single API:
+
+* **VP-SDE** (``use_flow_sigmas=False``) — Stable Diffusion / SDXL style models
+  with epsilon, x0, or v prediction. Uses the standard
+  ``alpha_t = 1 / sqrt(1 + sigma^2), sigma_t = sigma * alpha_t`` parameterization
+  and ports ``vp_*_order_*`` from the reference impl.
+* **Rectified flow / flow matching** (``use_flow_sigmas=True``) — FLUX, Z-Image,
+  Anima style models with flow_prediction. Uses ``alpha_t = 1 - sigma, sigma_t = sigma``
+  and the rectified-flow integral helpers defined locally (``_fn``,
+  ``_integral_one_over_fn``, ``_integral_lam_minus_curr_over_fn``).
+
+The rectified-flow integral helpers are kept local so this class is self-contained.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+# Number of sample points for the left Riemann sums approximating the
+# Taylor-extension integrals. Matches the reference impl's nums_intergrate=100.
+_INTEGRAL_NUM_POINTS = 100
+
+
+def _fn(x: float) -> float:
+    """ER-SDE noise-scale function ``SDE_5`` (paper appendix A.8).
+
+    Mirrors ``customized_func(..., func_type=7)`` in the reference impl —
+    the variant the paper recommends and tests for fast (~20 NFE) sampling.
+    """
+    return x * (math.exp(x**0.3) + 10.0)
+
+
+def _integral_one_over_fn(lambda_next: float, lambda_curr: float) -> float:
+    """Left Riemann sum of int_{lambda_next}^{lambda_curr} 1/_fn(lam) dlam.
+
+    Precondition: ``lambda_next > 0``. The integrand has a logarithmic singularity
+    at ``lam = 0`` (``_fn(0) = 0``); callers must skip this when ``sigma_next == 0``.
+    """
+    delta = lambda_curr - lambda_next
+    if delta <= 0:
+        return 0.0
+    step = delta / _INTEGRAL_NUM_POINTS
+    total = 0.0
+    for k in range(_INTEGRAL_NUM_POINTS):
+        lam = lambda_next + k * step
+        total += step / _fn(lam)
+    return total
+
+
+def _integral_lam_minus_curr_over_fn(lambda_next: float, lambda_curr: float) -> float:
+    """Left Riemann sum of int_{lambda_next}^{lambda_curr} (lam - lambda_curr)/_fn(lam) dlam.
+
+    Precondition: ``lambda_next > 0``. Same singularity at ``lam = 0`` as
+    :func:`_integral_one_over_fn`.
+    """
+    delta = lambda_curr - lambda_next
+    if delta <= 0:
+        return 0.0
+    step = delta / _INTEGRAL_NUM_POINTS
+    total = 0.0
+    for k in range(_INTEGRAL_NUM_POINTS):
+        lam = lambda_next + k * step
+        total += step * (lam - lambda_curr) / _fn(lam)
+    return total
+
+
+class ERSDEScheduler(SchedulerMixin, ConfigMixin):
+    """``diffusers`` scheduler for the ER-SDE multistep solver.
+
+    See module docstring for paper / reference-impl citations.
+
+    Args:
+        num_train_timesteps: Number of diffusion steps used during training.
+        beta_start: VP-SDE beta schedule start (ignored when ``use_flow_sigmas=True``).
+        beta_end: VP-SDE beta schedule end (ignored when ``use_flow_sigmas=True``).
+        beta_schedule: ``"linear"``, ``"scaled_linear"``, or ``"squaredcos_cap_v2"``.
+        trained_betas: Override betas with a pre-computed schedule.
+        prediction_type: ``"epsilon"``, ``"v_prediction"``, or ``"flow_prediction"``.
+        solver_order: Multistep order (1, 2, or 3). The solver auto-warms from order 1.
+        use_flow_sigmas: If True, use the rectified-flow parameterization
+            (``alpha_t = 1 - sigma``); else VP-SDE.
+        flow_shift: Sigma shift applied to the default flow schedule.
+        stochastic: If True, inject noise (full ER-SDE). If False, deterministic
+            ODE companion — same Taylor expansion with the noise term zeroed.
+        sigma_one_tolerance: Boundary tolerance for the ``sigma = 1`` limit
+            (rectified-flow only). Numerically paranoid; keep small.
+        timestep_spacing: ``"linspace"``, ``"leading"``, or ``"trailing"``.
+        steps_offset: Offset added to ``"leading"`` timesteps.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.00085,
+        beta_end: float = 0.012,
+        beta_schedule: str = "scaled_linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        prediction_type: str = "epsilon",
+        solver_order: int = 3,
+        use_flow_sigmas: bool = False,
+        flow_shift: float = 1.0,
+        stochastic: bool = True,
+        sigma_one_tolerance: float = 1e-6,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+    ):
+        if prediction_type not in ("epsilon", "v_prediction", "flow_prediction"):
+            raise ValueError(
+                f"prediction_type must be one of 'epsilon', 'v_prediction', 'flow_prediction', got {prediction_type!r}"
+            )
+        if solver_order not in (1, 2, 3):
+            raise ValueError(f"solver_order must be 1, 2, or 3, got {solver_order}")
+        if prediction_type == "flow_prediction" and not use_flow_sigmas:
+            # Not strictly invalid, but almost certainly a misconfiguration.
+            raise ValueError("prediction_type='flow_prediction' requires use_flow_sigmas=True (rectified-flow regime).")
+
+        # VP-SDE noise schedule (only used when use_flow_sigmas=False).
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule.
+            betas = []
+            for i in range(num_train_timesteps):
+                t1 = i / num_train_timesteps
+                t2 = (i + 1) / num_train_timesteps
+                a1 = math.cos((t1 + 0.008) / 1.008 * math.pi / 2) ** 2
+                a2 = math.cos((t2 + 0.008) / 1.008 * math.pi / 2) ** 2
+                betas.append(min(1 - a2 / a1, 0.999))
+            self.betas = torch.tensor(betas, dtype=torch.float32)
+        else:
+            raise NotImplementedError(f"beta_schedule {beta_schedule!r} is not implemented for ERSDEScheduler")
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # Default sigmas (VP-SDE form). Overwritten in set_timesteps.
+        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+
+        # Standard deviation of initial noise distribution (per Euler convention).
+        self.init_noise_sigma = 1.0
+
+        self.num_inference_steps: Optional[int] = None
+        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        self.timesteps = torch.from_numpy(timesteps)
+
+        # Multistep history. ``model_outputs`` stores x0 predictions; ``_sigma_history``
+        # stores the sigma at which each prediction was made. Both are FIFO with
+        # length == solver_order. Slot ``-1`` is the most recent.
+        self.model_outputs: List[Optional[torch.Tensor]] = [None] * solver_order
+        self._sigma_history: List[Optional[float]] = [None] * solver_order
+        self.lower_order_nums = 0
+        self._step_index: Optional[int] = None
+        self._begin_index: Optional[int] = None
+        self.sigmas = self.sigmas.to("cpu")
+
+    # ---- Index plumbing (mirrors DPM++) ---------------------------------------
+
+    @property
+    def step_index(self) -> Optional[int]:
+        return self._step_index
+
+    @property
+    def begin_index(self) -> Optional[int]:
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        self._begin_index = begin_index
+
+    def index_for_timestep(
+        self,
+        timestep: Union[int, torch.Tensor],
+        schedule_timesteps: Optional[torch.Tensor] = None,
+    ) -> int:
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+        index_candidates = (schedule_timesteps == timestep).nonzero()
+        if len(index_candidates) == 0:
+            return len(self.timesteps) - 1
+        # On the very first step, prefer the second match if duplicated, so
+        # img2img doesn't accidentally skip a sigma.
+        if len(index_candidates) > 1:
+            return index_candidates[1].item()
+        return index_candidates[0].item()
+
+    def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
+        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
+
+    # ---- Timestep / sigma scheduling ------------------------------------------
+
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        sigmas: Optional[Union[List[float], np.ndarray, torch.Tensor]] = None,
+        timesteps: Optional[List[int]] = None,
+    ) -> None:
+        """Set the discrete timesteps used for inference.
+
+        Exactly one of ``num_inference_steps``, ``timesteps``, or ``sigmas`` must
+        be provided. The ``sigmas`` form (mirroring :class:`EulerDiscreteScheduler`)
+        lets Anima/FLUX/Z-Image inject pre-shifted sigma schedules directly.
+        """
+        n_set = sum(x is not None for x in (num_inference_steps, timesteps, sigmas))
+        if n_set != 1:
+            raise ValueError("Must pass exactly one of `num_inference_steps`, `timesteps`, or `sigmas`.")
+
+        if sigmas is not None:
+            if isinstance(sigmas, torch.Tensor):
+                sigmas_np = sigmas.detach().cpu().numpy().astype(np.float32)
+            else:
+                sigmas_np = np.array(sigmas, dtype=np.float32)
+            num_inference_steps = len(sigmas_np) - 1
+            # Timesteps in the rectified-flow / Anima convention scale sigma to t.
+            # For VP-SDE this approximation is wrong but timesteps are only used
+            # for indexing; the algebra runs entirely off self.sigmas.
+            timesteps_np = (sigmas_np[:-1] * self.config.num_train_timesteps).astype(np.float32)
+        elif timesteps is not None:
+            timesteps_np = np.array(timesteps, dtype=np.float32)
+            num_inference_steps = len(timesteps_np)
+            sigmas_np = self._sigmas_for_timesteps(timesteps_np)
+        else:
+            assert num_inference_steps is not None
+            timesteps_np = self._default_timesteps(num_inference_steps)
+            sigmas_np = self._sigmas_for_timesteps(timesteps_np)
+
+        self.num_inference_steps = num_inference_steps
+        self.sigmas = torch.from_numpy(sigmas_np.astype(np.float32))
+        self.timesteps = torch.from_numpy(timesteps_np.astype(np.float32)).to(device=device)
+
+        # Reset multistep state.
+        self.model_outputs = [None] * self.config.solver_order
+        self._sigma_history = [None] * self.config.solver_order
+        self.lower_order_nums = 0
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")
+
+    def _default_timesteps(self, num_inference_steps: int) -> np.ndarray:
+        """Standard linspace/leading/trailing schedule (VP-SDE timesteps)."""
+        if self.config.timestep_spacing == "linspace":
+            timesteps = (
+                np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1)
+                .round()[::-1][:-1]
+                .copy()
+                .astype(np.float32)
+            )
+        elif self.config.timestep_spacing == "leading":
+            step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1)
+            timesteps = (
+                (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.float32)
+            )
+            timesteps += self.config.steps_offset
+        elif self.config.timestep_spacing == "trailing":
+            step_ratio = self.config.num_train_timesteps / num_inference_steps
+            timesteps = np.arange(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.float32)
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"timestep_spacing {self.config.timestep_spacing!r} must be one of 'linspace', 'leading', 'trailing'"
+            )
+        return timesteps
+
+    def _sigmas_for_timesteps(self, timesteps_np: np.ndarray) -> np.ndarray:
+        """Build the sigma schedule (with terminal 0 appended) for given timesteps."""
+        if self.config.use_flow_sigmas:
+            # Rectified-flow sigmas in [0, 1], time-shifted per Anima/FLUX convention.
+            num_inference_steps = len(timesteps_np)
+            alphas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)
+            sigmas = 1.0 - alphas
+            shift = self.config.flow_shift
+            sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy()
+            # Terminal sigma is exactly 0.
+            return np.concatenate([sigmas, [0.0]]).astype(np.float32)
+
+        # VP-SDE: interpolate against the train sigmas using timestep indexing.
+        train_sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        sigmas = np.interp(timesteps_np, np.arange(0, len(train_sigmas)), train_sigmas)
+        return np.concatenate([sigmas, [0.0]]).astype(np.float32)
+
+    # ---- Math helpers ---------------------------------------------------------
+
+    def _sigma_to_alpha_sigma_t(self, sigma: float) -> Tuple[float, float]:
+        """Map ``sigma`` to ``(alpha_t, sigma_t)``.
+
+        Rectified flow: ``alpha_t = 1 - sigma, sigma_t = sigma``.
+        VP-SDE: ``alpha_t = 1 / sqrt(1 + sigma^2), sigma_t = sigma * alpha_t``.
+        """
+        if self.config.use_flow_sigmas:
+            return 1.0 - sigma, sigma
+        alpha_t = 1.0 / math.sqrt(1.0 + sigma * sigma)
+        return alpha_t, sigma * alpha_t
+
+    @staticmethod
+    def _lambda(alpha_t: float, sigma_t: float) -> float:
+        """ER-SDE ``lambda = sigma_t / alpha_t`` — the noise-to-signal ratio.
+
+        This matches the reference impl's ``lambdas = sigmas / alphas`` in both
+        VP and rectified-flow regimes (see ``vp_*_order_*`` in
+        ``https://github.com/QinpengCui/ER-SDE-Solver``). For VP-SDE this equals
+        the stored sigma; for rectified flow it equals ``sigma / (1 - sigma)``.
+        Diverges at ``sigma_t = alpha_t = 0`` (rectified flow at sigma=1) — the
+        boundary branch in :meth:`_first_order_update` handles that case.
+        """
+        if alpha_t == 0.0:
+            return float("inf")
+        return sigma_t / alpha_t
+
+    # ---- Model output conversion ----------------------------------------------
+
+    def _convert_model_output(self, model_output: torch.Tensor, sample: torch.Tensor) -> torch.Tensor:
+        """Convert raw model output to an ``x0`` prediction at the current sigma."""
+        sigma = float(self.sigmas[self.step_index].item())
+        if self.config.prediction_type == "flow_prediction":
+            # v = (x - x0) / sigma  =>  x0 = x - sigma * v
+            return sample - sigma * model_output
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        if self.config.prediction_type == "epsilon":
+            return (sample - sigma_t * model_output) / alpha_t
+        if self.config.prediction_type == "v_prediction":
+            return alpha_t * sample - sigma_t * model_output
+        raise ValueError(f"Unsupported prediction_type {self.config.prediction_type!r}")
+
+    # ---- Order-N updates -------------------------------------------------------
+
+    def _first_order_update(
+        self,
+        x0: torch.Tensor,
+        sample: torch.Tensor,
+        sigma_curr: float,
+        sigma_next: float,
+        noise: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Order-1 ER-SDE step (ports ``vp_1_order`` / ``er_sde_rf_step`` order-1 branch)."""
+        # Rectified-flow boundary: sigma_curr ~= 1 means alpha_curr ~= 0 so lambda diverges.
+        # Closed-form limit (er_sde.py:136-142): x_next = (1 - sigma_next) * x0 + sigma_next * noise.
+        if self.config.use_flow_sigmas and 1.0 - sigma_curr < self.config.sigma_one_tolerance:
+            x_next = (1.0 - sigma_next) * x0
+            if self.config.stochastic and noise is not None and sigma_next > 0.0:
+                x_next = x_next + sigma_next * noise
+            return x_next
+
+        alpha_curr, sigma_curr_t = self._sigma_to_alpha_sigma_t(sigma_curr)
+        alpha_next, sigma_next_t = self._sigma_to_alpha_sigma_t(sigma_next)
+
+        # Reference impl uses lambda = sigma_t / alpha_t in both VP and flow regimes.
+        lambda_curr = self._lambda(alpha_curr, sigma_curr_t)
+        # At the terminal step, sigma_next == 0 so lambda_next == 0 and fn_next == 0.
+        lambda_next = self._lambda(alpha_next, sigma_next_t) if sigma_next_t > 0.0 else 0.0
+
+        fn_curr = _fn(lambda_curr)
+        fn_next = _fn(lambda_next)
+        r_fn = fn_next / fn_curr if fn_curr != 0.0 else 0.0
+        r_alphas = alpha_next / alpha_curr
+
+        # Stochastic noise std (paper appendix eq. for ER-SDE_5 variance).
+        # ``inner`` can underflow to tiny negatives by roundoff; clip.
+        inner = lambda_next**2 - lambda_curr**2 * r_fn**2
+        if inner < 0.0:
+            inner = 0.0
+        noise_std = math.sqrt(inner) * alpha_next
+
+        x_next = r_alphas * r_fn * sample + alpha_next * (1.0 - r_fn) * x0
+        if self.config.stochastic and noise is not None and sigma_next > 0.0:
+            x_next = x_next + noise_std * noise
+        return x_next
+
+    def _second_order_update(
+        self,
+        sample: torch.Tensor,
+        sigma_curr: float,
+        sigma_next: float,
+        noise: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Order-2 ER-SDE step (ports ``vp_2_order_taylor``)."""
+        x0 = self.model_outputs[-1]
+        old_x0 = self.model_outputs[-2]
+        sigma_prev_curr = self._sigma_history[-2]
+        assert x0 is not None and old_x0 is not None and sigma_prev_curr is not None
+
+        # If the previous step used the sigma=1 closed-form limit, the finite-difference
+        # derivative across that boundary is meaningless — fall back to order 1.
+        if self.config.use_flow_sigmas and 1.0 - sigma_prev_curr < self.config.sigma_one_tolerance:
+            return self._first_order_update(x0, sample, sigma_curr, sigma_next, noise)
+
+        # Order-1 base.
+        x_next = self._first_order_update(x0, sample, sigma_curr, sigma_next, noise)
+
+        # Skip the higher-order term at the terminal step — the integral helpers diverge
+        # at lambda = 0 (sigma = 0), see _integral_one_over_fn docstring.
+        if sigma_next <= 0.0:
+            return x_next
+
+        alpha_curr, sigma_curr_t = self._sigma_to_alpha_sigma_t(sigma_curr)
+        alpha_next, sigma_next_t = self._sigma_to_alpha_sigma_t(sigma_next)
+        alpha_prev, sigma_prev_t = self._sigma_to_alpha_sigma_t(sigma_prev_curr)
+        lambda_curr = self._lambda(alpha_curr, sigma_curr_t)
+        lambda_next = self._lambda(alpha_next, sigma_next_t)
+        lambda_prev = self._lambda(alpha_prev, sigma_prev_t)
+
+        denom = lambda_curr - lambda_prev
+        if denom == 0.0:
+            return x_next
+        d_x0 = (x0 - old_x0) / denom
+
+        fn_next = _fn(lambda_next)
+        s_int = _integral_one_over_fn(lambda_next, lambda_curr)
+        x_next = x_next + alpha_next * (lambda_next - lambda_curr + s_int * fn_next) * d_x0
+        return x_next
+
+    def _third_order_update(
+        self,
+        sample: torch.Tensor,
+        sigma_curr: float,
+        sigma_next: float,
+        noise: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Order-3 ER-SDE step (ports ``vp_3_order_taylor``)."""
+        x0 = self.model_outputs[-1]
+        old_x0 = self.model_outputs[-2]
+        old_old_x0 = self.model_outputs[-3]
+        sigma_prev_curr = self._sigma_history[-2]
+        sigma_prev_prev = self._sigma_history[-3]
+        assert (
+            x0 is not None
+            and old_x0 is not None
+            and old_old_x0 is not None
+            and sigma_prev_curr is not None
+            and sigma_prev_prev is not None
+        )
+
+        # If any sigma in the lookback hits the boundary, fall back to order 2.
+        if self.config.use_flow_sigmas and (
+            1.0 - sigma_prev_curr < self.config.sigma_one_tolerance
+            or 1.0 - sigma_prev_prev < self.config.sigma_one_tolerance
+        ):
+            return self._second_order_update(sample, sigma_curr, sigma_next, noise)
+
+        # Order-2 base.
+        x_next = self._second_order_update(sample, sigma_curr, sigma_next, noise)
+
+        if sigma_next <= 0.0:
+            return x_next
+
+        alpha_curr, sigma_curr_t = self._sigma_to_alpha_sigma_t(sigma_curr)
+        alpha_next, sigma_next_t = self._sigma_to_alpha_sigma_t(sigma_next)
+        alpha_prev, sigma_prev_t = self._sigma_to_alpha_sigma_t(sigma_prev_curr)
+        alpha_pprev, sigma_pprev_t = self._sigma_to_alpha_sigma_t(sigma_prev_prev)
+        lambda_curr = self._lambda(alpha_curr, sigma_curr_t)
+        lambda_next = self._lambda(alpha_next, sigma_next_t)
+        lambda_prev = self._lambda(alpha_prev, sigma_prev_t)
+        lambda_pprev = self._lambda(alpha_pprev, sigma_pprev_t)
+
+        denom_d = lambda_curr - lambda_prev
+        denom_d_prev = lambda_prev - lambda_pprev
+        denom_dd = lambda_curr - lambda_pprev
+        if denom_d == 0.0 or denom_d_prev == 0.0 or denom_dd == 0.0:
+            return x_next
+
+        d_x0 = (x0 - old_x0) / denom_d
+        old_d_x0 = (old_x0 - old_old_x0) / denom_d_prev
+        dd_x0 = 2.0 * (d_x0 - old_d_x0) / denom_dd
+
+        fn_next = _fn(lambda_next)
+        s_d_int = _integral_lam_minus_curr_over_fn(lambda_next, lambda_curr)
+        x_next = x_next + alpha_next * ((lambda_next - lambda_curr) ** 2 / 2.0 + s_d_int * fn_next) * dd_x0
+        return x_next
+
+    # ---- Public step ----------------------------------------------------------
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: Optional[Union[int, torch.Tensor]] = None
+    ) -> torch.Tensor:
+        """No-op (matches ``FlowMatchEulerDiscreteScheduler``)."""
+        return sample
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SchedulerOutput, Tuple]:
+        """Predict the sample at the next timestep using one ER-SDE step."""
+        if self.num_inference_steps is None:
+            raise ValueError("num_inference_steps is None — call `set_timesteps` before calling `step`.")
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        sigma_curr = float(self.sigmas[self.step_index].item())
+        sigma_next = float(self.sigmas[self.step_index + 1].item())
+
+        # 1. Convert model output to x0 prediction.
+        x0 = self._convert_model_output(model_output, sample)
+
+        # 2. FIFO-shift the multistep history. New entry goes in slot -1.
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self._sigma_history[i] = self._sigma_history[i + 1]
+        self.model_outputs[-1] = x0
+        self._sigma_history[-1] = sigma_curr
+
+        # 3. Sample noise (only when stochastic and not at terminal step).
+        if self.config.stochastic and sigma_next > 0.0:
+            noise = randn_tensor(
+                model_output.shape,
+                generator=generator,
+                device=model_output.device,
+                dtype=model_output.dtype,
+            )
+        else:
+            noise = None
+
+        # 4. Dispatch by available history.
+        if self.config.solver_order == 1 or self.lower_order_nums < 1:
+            prev_sample = self._first_order_update(x0, sample, sigma_curr, sigma_next, noise)
+        elif self.config.solver_order == 2 or self.lower_order_nums < 2:
+            prev_sample = self._second_order_update(sample, sigma_curr, sigma_next, noise)
+        else:
+            prev_sample = self._third_order_update(sample, sigma_curr, sigma_next, noise)
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # 5. Advance step index.
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    # ---- Forward noising (training / img2img) ---------------------------------
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward-noise ``original_samples`` at the given timesteps (img2img style)."""
+        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+            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)
+
+        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:
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            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)
+
+        if self.config.use_flow_sigmas:
+            alpha_t = 1.0 - sigma
+            sigma_t = sigma
+        else:
+            alpha_t = 1.0 / torch.sqrt(1.0 + sigma * sigma)
+            sigma_t = sigma * alpha_t
+        return alpha_t * original_samples + sigma_t * noise
+
+    def __len__(self) -> int:
+        return self.config.num_train_timesteps

invokeai/backend/rectified_flow/rectified_flow_inpaint_extension.py

@@ -0,0 +1,58 @@
+import torch
+
+
+def assert_broadcastable(*shapes):
+    try:
+        torch.broadcast_shapes(*shapes)
+    except RuntimeError as e:
+        raise AssertionError(f"Shapes {shapes} are not broadcastable.") from e
+
+
+class RectifiedFlowInpaintExtension:
+    """A class for managing inpainting with rectified flow models (e.g. FLUX, SD3, CogView4)."""
+
+    def __init__(self, init_latents: torch.Tensor, inpaint_mask: torch.Tensor, noise: torch.Tensor):
+        """Initialize InpaintExtension.
+
+        Args:
+            init_latents (torch.Tensor): The initial latents (i.e. un-noised at timestep 0). In 'packed' format.
+            inpaint_mask (torch.Tensor): A mask specifying which elements to inpaint. Range [0, 1]. Values of 1 will be
+                re-generated. Values of 0 will remain unchanged. Values between 0 and 1 can be used to blend the
+                inpainted region with the background. In 'packed' format.
+            noise (torch.Tensor): The noise tensor used to noise the init_latents. In 'packed' format.
+        """
+        assert_broadcastable(init_latents.shape, inpaint_mask.shape, noise.shape)
+
+        self._init_latents = init_latents
+        self._inpaint_mask = inpaint_mask
+        self._noise = noise
+
+    def _apply_mask_gradient_adjustment(self, t_prev: float) -> torch.Tensor:
+        """Applies inpaint mask gradient adjustment and returns the inpaint mask to be used at the current timestep."""
+        # As we progress through the denoising process, we promote gradient regions of the mask to have a full weight of
+        # 1.0. This helps to produce more coherent seams around the inpainted region.
+
+        # We use a small epsilon to avoid any potential issues with floating point precision.
+        eps = 1e-4
+        mask = torch.where(self._inpaint_mask >= t_prev + eps, 1.0, 0.0).to(
+            dtype=self._inpaint_mask.dtype, device=self._inpaint_mask.device
+        )
+
+        return mask
+
+    def merge_intermediate_latents_with_init_latents(
+        self, intermediate_latents: torch.Tensor, t_prev: float
+    ) -> torch.Tensor:
+        """Merge the intermediate latents with the initial latents for the current timestep using the inpaint mask. I.e.
+        update the intermediate latents to keep the regions that are not being inpainted on the correct noise
+        trajectory.
+
+        This function should be called after each denoising step.
+        """
+        mask = self._apply_mask_gradient_adjustment(t_prev)
+
+        # Noise the init latents for the current timestep.
+        noised_init_latents = self._noise * t_prev + (1.0 - t_prev) * self._init_latents
+
+        # Merge the intermediate latents with the noised_init_latents using the inpaint_mask.
+        return intermediate_latents * mask + noised_init_latents * (1.0 - mask)