Merge pull request #4745 from vladmandic/feat/er-sde-solver

add ER-SDE solver scheduler

modules/schedulers/scheduler_ersde.py

@@ -0,0 +1,552 @@
+# Copyright 2024 Qinpeng Cui and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ER-SDE Solver: Extended Reverse-time SDE for diffusion sampling
+# Based on https://github.com/QinpengCui/ER-SDE-Solver (WACV 2025)
+# Ported to diffusers-compatible scheduler for SD.Next
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+@dataclass
+class ERSDESchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+    pred_original_sample: Optional[torch.Tensor] = None
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999, alpha_transform_type="cosine"):
+    if alpha_transform_type == "cosine":
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+    elif alpha_transform_type == "exp":
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+    else:
+        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+def rescale_zero_terminal_snr(betas):
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+    alphas_bar = alphas_bar_sqrt ** 2
+    alphas = alphas_bar[1:] / alphas_bar[:-1]
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+    return betas
+
+
+def customized_func(x, func_type=7):
+    """Noise scaling function for ER-SDE. Controls the stochasticity of sampling.
+    func_type=1 gives ODE (deterministic), func_type=7 is the paper's recommended default."""
+    if func_type == 1:  # ODE
+        return x
+    elif func_type == 2:  # Original SDE
+        return x ** 2
+    elif func_type == 3:  # SDE_1
+        eta = 0.0
+        return x * (torch.exp(x ** (eta - 1) / (eta - 1)) + 10)
+    elif func_type == 4:  # SDE_2
+        return x ** 1.5
+    elif func_type == 5:  # SDE_3
+        return x ** 2.5
+    elif func_type == 6:  # SDE_4
+        return x ** 0.9 * torch.log10(1 + 100 * x ** 1.5)
+    elif func_type == 7:  # SDE_5 (default)
+        return x * (torch.exp(x ** 0.3) + 10)
+    else:
+        raise ValueError(f"func_type must be 1-7, got {func_type}")
+
+
+class ERSDEScheduler(SchedulerMixin, ConfigMixin):
+    """ER-SDE (Extended Reverse-time SDE) Solver for diffusion models.
+
+    Supports VP-type models (SD 1.5, SDXL) and flow-matching models (SD3, Flux)
+    with 1st, 2nd, and 3rd order Taylor expansion methods.
+    """
+
+    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 = 1,
+        func_type: int = 7,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        clip_sample: bool = False,
+        clip_sample_range: float = 1.0,
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        rescale_betas_zero_snr: bool = False,
+        set_alpha_to_one: bool = True,
+        lower_order_final: bool = True,
+        use_karras_sigmas: bool = False,
+        num_integration_steps: int = 100,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+        base_shift: float = 0.5,
+        max_shift: float = 1.15,
+    ):
+        if solver_order not in (1, 2, 3):
+            raise ValueError(f"solver_order must be 1, 2, or 3, got {solver_order}")
+        if func_type < 1 or func_type > 7:
+            raise ValueError(f"func_type must be 1-7, got {func_type}")
+
+        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":
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        self.init_noise_sigma = 1.0
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+
+        # ER-SDE state
+        self._step_index = None
+        self._begin_index = None
+        self._is_flow = False
+        self.old_x0 = None
+        self.old_d_x0 = None
+        self.prev_lambda = None
+        self.prev_prev_lambda = None
+        self.lower_order_nums = 0
+        # Flow-matching arrays (populated in set_timesteps for flow mode)
+        self.sigmas = None
+        self._flow_alphas = None
+        self._flow_sigmas = None
+        self._flow_lambdas = None
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        return self._begin_index
+
+    def set_begin_index(self, begin_index=0):
+        self._begin_index = begin_index
+
+    def _init_step_index(self, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.to(self.timesteps.device)
+        index_candidates = (self.timesteps == timestep).nonzero()
+        if len(index_candidates) == 0:
+            step_index = len(self.timesteps) - 1
+        elif len(index_candidates) > 1:
+            step_index = index_candidates[1].item()
+        else:
+            step_index = index_candidates[0].item()
+        self._step_index = step_index
+
+    def _reset_state(self):
+        self._step_index = None
+        self._begin_index = None
+        self.old_x0 = None
+        self.old_d_x0 = None
+        self.prev_lambda = None
+        self.prev_prev_lambda = None
+        self.lower_order_nums = 0
+
+    def _setup_flow(self, sigmas, device, mu=None):
+        """Common flow-matching schedule setup from a sigmas tensor.
+
+        Uses alpha=1 so the VP update formula r_alpha*r_fn*x + alpha*(1-r_fn)*x0
+        reduces to the correct flow-matching form r_f*x + (1-r_f)*x0, and lambda=sigma
+        keeps the customized noise function in a well-behaved range on [0,1].
+        """
+        self._is_flow = True
+        if self.config.use_dynamic_shifting:
+            if mu is None:
+                raise ValueError("mu must be provided when use_dynamic_shifting=True")
+            sigmas = self._time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
+        flow_sigmas = sigmas.clamp(min=1e-8, max=1.0 - 1e-8)
+        flow_alphas = torch.ones_like(flow_sigmas)
+        flow_lambdas = flow_sigmas.clone()
+        self.timesteps = (sigmas * self.config.num_train_timesteps).to(device=device)
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device, dtype=sigmas.dtype)])
+        self._flow_alphas = torch.cat([flow_alphas, torch.ones(1, device=device, dtype=torch.float64)])
+        self._flow_sigmas = torch.cat([flow_sigmas, torch.zeros(1, device=device, dtype=torch.float64)])
+        self._flow_lambdas = torch.cat([flow_lambdas, torch.zeros(1, device=device, dtype=torch.float64)])
+
+    def set_timesteps(self, num_inference_steps: Optional[int] = None, device: Union[str, torch.device] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, mu: Optional[float] = None):
+        if sigmas is not None:
+            # Flow-matching path: sigmas provided externally
+            sigmas = np.array(sigmas, dtype=np.float64) if not isinstance(sigmas, np.ndarray) else sigmas.astype(np.float64)
+            self.num_inference_steps = len(sigmas)
+            sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device)
+            self._setup_flow(sigmas, device, mu)
+        elif self.config.prediction_type == "flow_prediction":
+            # Flow-matching path: compute flow schedule internally
+            self.num_inference_steps = num_inference_steps
+            sigmas = np.linspace(1.0, 1.0 / self.config.num_train_timesteps, num_inference_steps)
+            sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device)
+            self._setup_flow(sigmas, device, mu)
+        elif timesteps is not None:
+            # Custom timesteps path
+            self._is_flow = False
+            self.num_inference_steps = len(timesteps)
+            timesteps = np.array(timesteps, dtype=np.int64)
+            self.timesteps = torch.from_numpy(timesteps).to(device)
+            sigmas_arr = []
+            for t in timesteps:
+                acp = self.alphas_cumprod[t].item()
+                sigmas_arr.append(((1 - acp) / max(acp, 1e-8)) ** 0.5)
+            sigmas_arr.append(0.0)
+            self.sigmas = torch.tensor(sigmas_arr, dtype=torch.float64, device=device)
+            self._flow_alphas = None
+            self._flow_sigmas = None
+            self._flow_lambdas = None
+        else:
+            # VP path: standard diffusion models
+            self._is_flow = False
+            self.num_inference_steps = num_inference_steps
+
+            if self.config.timestep_spacing == "linspace":
+                timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps).round()[::-1].copy().astype(np.int64)
+            elif self.config.timestep_spacing == "leading":
+                step_ratio = self.config.num_train_timesteps // num_inference_steps
+                timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+                timesteps += self.config.steps_offset
+            elif self.config.timestep_spacing == "trailing":
+                step_ratio = self.config.num_train_timesteps / num_inference_steps
+                timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+                timesteps -= 1
+            else:
+                raise ValueError(f"{self.config.timestep_spacing} is not supported")
+
+            self.timesteps = torch.from_numpy(timesteps).to(device)
+            sigmas_arr = []
+            for t in timesteps:
+                acp = self.alphas_cumprod[t].item()
+                sigmas_arr.append(((1 - acp) / max(acp, 1e-8)) ** 0.5)
+            sigmas_arr.append(0.0)
+            self.sigmas = torch.tensor(sigmas_arr, dtype=torch.float64, device=device)
+            self._flow_alphas = None
+            self._flow_sigmas = None
+            self._flow_lambdas = None
+
+        self._reset_state()
+
+    def _time_shift(self, mu, sigma, t):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def _get_alpha_sigma_lambda(self, step_idx):
+        """Get alpha, sigma, lambda values for a given step index."""
+        if self._is_flow:
+            alpha = self._flow_alphas[step_idx]
+            sigma = self._flow_sigmas[step_idx]
+            lam = self._flow_lambdas[step_idx]
+        else:
+            t = self.timesteps[step_idx] if step_idx < len(self.timesteps) else 0
+            if step_idx >= len(self.timesteps):
+                # Past the last timestep: fully denoised
+                alpha = torch.tensor(1.0, dtype=torch.float64)
+                sigma = torch.tensor(0.0, dtype=torch.float64)
+                lam = torch.tensor(0.0, dtype=torch.float64)
+            else:
+                acp = self.alphas_cumprod[t].to(torch.float64)
+                alpha = acp.sqrt()
+                sigma = (1.0 - acp).sqrt()
+                lam = sigma / alpha.clamp(min=1e-8)
+        return alpha, sigma, lam
+
+    def _numerical_clip(self, x, eps=1e-6):
+        """Clip near-zero values to prevent negative variance from float errors."""
+        if isinstance(x, torch.Tensor):
+            return torch.where(torch.abs(x) < eps, torch.zeros_like(x), x)
+        return 0.0 if abs(x) < eps else x
+
+    @staticmethod
+    def _safe_div(numerator, denominator, eps=1e-10):
+        """Division with sign-preserving zero protection for lambda differences."""
+        if denominator.abs() < eps:
+            return torch.zeros_like(numerator)
+        return numerator / denominator
+
+    def _compute_fn(self, x):
+        """Evaluate the noise scaling function on a scalar or tensor."""
+        if isinstance(x, (int, float)):
+            x = torch.tensor(x, dtype=torch.float64)
+        return customized_func(x, self.config.func_type)
+
+    def _compute_integral(self, lambda_next, lambda_curr):
+        """Numerical integration of 1/fn(lambda) over [lambda_next, lambda_curr] for 2nd-order correction."""
+        N = self.config.num_integration_steps
+        delta = (lambda_curr - lambda_next)
+        if abs(delta.item()) < 1e-10:
+            return torch.tensor(0.0, dtype=torch.float64, device=lambda_curr.device)
+        indices = lambda_next + torch.arange(1, N + 1, dtype=torch.float64, device=lambda_curr.device) / N * delta
+        fn_vals = torch.stack([self._compute_fn(idx) for idx in indices])
+        fn_vals = fn_vals.clamp(min=1e-10)
+        s_int = (1.0 / fn_vals * delta / N).sum()
+        return s_int
+
+    def _compute_derivative_integral(self, lambda_next, lambda_curr):
+        """Numerical integration of (lambda - lambda_curr)/fn(lambda) over [lambda_next, lambda_curr] for 3rd-order correction."""
+        N = self.config.num_integration_steps
+        delta = (lambda_curr - lambda_next)
+        if abs(delta.item()) < 1e-10:
+            return torch.tensor(0.0, dtype=torch.float64, device=lambda_curr.device)
+        indices = lambda_next + torch.arange(1, N + 1, dtype=torch.float64, device=lambda_curr.device) / N * delta
+        fn_vals = torch.stack([self._compute_fn(idx) for idx in indices])
+        fn_vals = fn_vals.clamp(min=1e-10)
+        s_d_int = ((indices - lambda_curr) / fn_vals * delta / N).sum()
+        return s_d_int
+
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+        abs_sample = sample.abs()
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(s, min=1, max=self.config.sample_max_value)
+        s = s.unsqueeze(1)
+        sample = torch.clamp(sample, -s, s) / s
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+        return sample
+
+    def _predict_x0(self, model_output, sample, alpha, sigma):
+        """Convert model output to x0 prediction based on prediction_type."""
+        prediction_type = self.config.prediction_type
+        if prediction_type == "epsilon":
+            x0 = (sample - sigma * model_output) / alpha.clamp(min=1e-8)
+        elif prediction_type == "sample":
+            x0 = model_output
+        elif prediction_type == "v_prediction":
+            x0 = alpha * sample - sigma * model_output
+        elif prediction_type == "flow_prediction":
+            x0 = sample - sigma * model_output
+        else:
+            raise ValueError(f"prediction_type {prediction_type} must be one of epsilon, sample, v_prediction, flow_prediction")
+        return x0
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+        return sample
+
+    def scale_noise(self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], noise: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
+        """Forward process for flow-matching models (img2img noise addition)."""
+        if self.sigmas is None:
+            return sample
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+        if self._begin_index is None:
+            step_indices = [self._index_for_timestep(t, schedule_timesteps) for t in timestep]
+        elif self._step_index is not None:
+            step_indices = [self._step_index] * timestep.shape[0]
+        else:
+            step_indices = [self._begin_index] * timestep.shape[0]
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+        sample = sigma * noise + (1.0 - sigma) * sample
+        return sample
+
+    def _index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+        index_candidates = (schedule_timesteps == timestep).nonzero()
+        if len(index_candidates) == 0:
+            return len(schedule_timesteps) - 1
+        return index_candidates[0].item()
+
+    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[ERSDESchedulerOutput, Tuple]:
+        if self.num_inference_steps is None:
+            raise ValueError("Number of inference steps is None, run set_timesteps first")
+
+        if self._step_index is None:
+            self._init_step_index(timestep)
+
+        dtype = sample.dtype
+        device = sample.device
+
+        # Get alpha, sigma, lambda for current and next step
+        alpha_curr, sigma_curr, lambda_curr = self._get_alpha_sigma_lambda(self._step_index)
+        alpha_next, sigma_next, lambda_next = self._get_alpha_sigma_lambda(self._step_index + 1)
+
+        alpha_curr = alpha_curr.to(device)
+        sigma_curr = sigma_curr.to(device)
+        lambda_curr = lambda_curr.to(device)
+        alpha_next = alpha_next.to(device)
+        sigma_next = sigma_next.to(device)
+        lambda_next = lambda_next.to(device)
+
+        # Predict x0 from model output
+        x0 = self._predict_x0(model_output.to(torch.float64), sample.to(torch.float64), alpha_curr, sigma_curr)
+
+        # Apply thresholding or clipping
+        if self.config.thresholding:
+            x0 = self._threshold_sample(x0)
+        elif self.config.clip_sample:
+            x0 = x0.clamp(-self.config.clip_sample_range, self.config.clip_sample_range)
+
+        # Determine effective solver order for this step
+        max_order = self.config.solver_order
+        effective_order = min(max_order, self.lower_order_nums + 1)
+
+        # Lower order for final step
+        is_last_step = (self._step_index == len(self.timesteps) - 1)
+        if self.config.lower_order_final and is_last_step:
+            effective_order = 1
+
+        # When next sigma is 0 (fully denoised), use 1st order and skip noise
+        at_final_sigma = sigma_next.abs() < 1e-8
+        if at_final_sigma:
+            effective_order = 1
+
+        # Compute noise scaling function ratio
+        fn_lambda_curr = self._compute_fn(lambda_curr)
+        fn_lambda_next = self._compute_fn(lambda_next)
+        # Avoid division by zero
+        if fn_lambda_curr.abs() < 1e-10:
+            r_fn = torch.tensor(1.0, dtype=torch.float64, device=device)
+        else:
+            r_fn = fn_lambda_next / fn_lambda_curr
+
+        r_alpha = alpha_next / alpha_curr.clamp(min=1e-8)
+
+        # Compute stochastic noise standard deviation
+        noise_var = self._numerical_clip(lambda_next ** 2 - lambda_curr ** 2 * r_fn ** 2)
+        if isinstance(noise_var, torch.Tensor):
+            noise_std = alpha_next * torch.sqrt(noise_var.clamp(min=0))
+        else:
+            noise_std = alpha_next * math.sqrt(max(noise_var, 0))
+
+        # Generate noise
+        if at_final_sigma or noise_std.abs() < 1e-10:
+            noise_term = torch.zeros_like(sample, dtype=torch.float64)
+        else:
+            noise = randn_tensor(sample.shape, generator=generator, device=device, dtype=torch.float64)
+            noise_term = noise_std * noise
+
+        sample_f64 = sample.to(torch.float64)
+
+        # ER-SDE VP update
+        if effective_order == 1:
+            # 1st order: x_next = r_alpha * r_fn * x + alpha_next * (1 - r_fn) * x0 + noise
+            prev_sample = r_alpha * r_fn * sample_f64 + alpha_next * (1.0 - r_fn) * x0 + noise_term
+
+        elif effective_order == 2:
+            # 2nd order: add first derivative correction
+            d_x0 = self._safe_div(x0 - self.old_x0, lambda_curr - self.prev_lambda)
+            s_int = self._compute_integral(lambda_next, lambda_curr)
+            s_int = s_int.to(device)
+            delta_lambda = lambda_next - lambda_curr
+            correction = alpha_next * (delta_lambda + s_int * fn_lambda_next) * d_x0
+            prev_sample = r_alpha * r_fn * sample_f64 + alpha_next * (1.0 - r_fn) * x0 + correction + noise_term
+
+        elif effective_order == 3:
+            # 3rd order: add second derivative correction
+            d_x0 = self._safe_div(x0 - self.old_x0, lambda_curr - self.prev_lambda)
+            dd_x0 = self._safe_div(2.0 * (d_x0 - self.old_d_x0), lambda_curr - self.prev_prev_lambda)
+            s_int = self._compute_integral(lambda_next, lambda_curr).to(device)
+            s_d_int = self._compute_derivative_integral(lambda_next, lambda_curr).to(device)
+            delta_lambda = lambda_next - lambda_curr
+            correction_1 = alpha_next * (delta_lambda + s_int * fn_lambda_next) * d_x0
+            correction_2 = alpha_next * (delta_lambda ** 2 / 2.0 + s_d_int * fn_lambda_next) * dd_x0
+            prev_sample = r_alpha * r_fn * sample_f64 + alpha_next * (1.0 - r_fn) * x0 + correction_1 + correction_2 + noise_term
+
+        # Update state for next step
+        # Compute d_x0 for buffer (needed by next step's 3rd order)
+        if self.old_x0 is not None and (lambda_curr - self.prev_lambda).abs() > 1e-10:
+            current_d_x0 = (x0 - self.old_x0) / (lambda_curr - self.prev_lambda)
+        else:
+            current_d_x0 = None
+
+        self.prev_prev_lambda = self.prev_lambda
+        self.prev_lambda = lambda_curr
+        self.old_d_x0 = current_d_x0
+        self.old_x0 = x0
+        self.lower_order_nums = min(self.lower_order_nums + 1, self.config.solver_order)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample.to(dtype),)
+        return ERSDESchedulerOutput(prev_sample=prev_sample.to(dtype), pred_original_sample=x0.to(dtype))
+
+    def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
+        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

modules/sd_samplers_diffusers.py

@@ -62,6 +62,7 @@ try:
     from modules.schedulers.scheduler_unipc_flowmatch import FlowUniPCMultistepScheduler # pylint: disable=ungrouped-imports
     from modules.schedulers.scheduler_flashflow import FlashFlowMatchEulerDiscreteScheduler # pylint: disable=ungrouped-imports
     from modules.schedulers.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports
+    from modules.schedulers.scheduler_ersde import ERSDEScheduler # pylint: disable=ungrouped-imports
 except Exception as e:
     log.error(f'Sampler import: version={diffusers.__version__} error: {e}')
     if os.environ.get('SD_SAMPLER_DEBUG', None) is not None:
@@ -154,6 +155,13 @@ config.update({
     'UFOGen': { },
     'BDIA DDIM': { 'clip_sample': False, 'set_alpha_to_one': True, 'steps_offset': 0, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'leading', 'rescale_betas_zero_snr': False, 'thresholding': False, 'gamma': 1.0 },
 
+    'ER-SDE': { 'solver_order': 1, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True },
+    'ER-SDE 2M': { 'solver_order': 2, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True },
+    'ER-SDE 3M': { 'solver_order': 3, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True },
+    'ER-SDE FlowMatch': { 'solver_order': 1, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
+    'ER-SDE 2M FlowMatch': { 'solver_order': 2, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
+    'ER-SDE 3M FlowMatch': { 'solver_order': 3, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
+
     'PNDM': { 'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0, 'timestep_spacing': 'linspace' },
     'IPNDM': { },
     'DDPM': { 'variance_type': "fixed_small", 'clip_sample': False, 'thresholding': False, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False },
@@ -277,6 +285,12 @@ samplers_data_diffusers = [
 
     SamplerData('VDM Solver', lambda model: DiffusionSampler('VDM Solver', VDMScheduler, model), [], {}),
     SamplerData('BDIA DDIM', lambda model: DiffusionSampler('BDIA DDIM g=0', BDIA_DDIMScheduler, model), [], {}),
+    SamplerData('ER-SDE', lambda model: DiffusionSampler('ER-SDE', ERSDEScheduler, model), [], {}),
+    SamplerData('ER-SDE 2M', lambda model: DiffusionSampler('ER-SDE 2M', ERSDEScheduler, model), [], {}),
+    SamplerData('ER-SDE 3M', lambda model: DiffusionSampler('ER-SDE 3M', ERSDEScheduler, model), [], {}),
+    SamplerData('ER-SDE FlowMatch', lambda model: DiffusionSampler('ER-SDE FlowMatch', ERSDEScheduler, model), [], {}),
+    SamplerData('ER-SDE 2M FlowMatch', lambda model: DiffusionSampler('ER-SDE 2M FlowMatch', ERSDEScheduler, model), [], {}),
+    SamplerData('ER-SDE 3M FlowMatch', lambda model: DiffusionSampler('ER-SDE 3M FlowMatch', ERSDEScheduler, model), [], {}),
     SamplerData('LCM', lambda model: DiffusionSampler('LCM', LCMScheduler, model), [], {}),
     SamplerData('LCM FlowMatch', lambda model: DiffusionSampler('LCM FlowMatch', FlowMatchLCMScheduler, model), [], {}),
     SamplerData('TCD', lambda model: DiffusionSampler('TCD', TCDScheduler, model), [], {}),