diffsynth-engine 0.0.0__py3-none-any.whl

diffsynth_engine/__init__.py

@@ -0,0 +1,28 @@
+from .pipelines import (
+    FluxImagePipeline,
+    SDXLImagePipeline,
+    SDImagePipeline,
+    WanVideoPipeline,
+    FluxModelConfig,
+    SDXLModelConfig,
+    SDModelConfig,
+    WanModelConfig,
+)
+from .utils.download import fetch_model, fetch_modelscope_model, fetch_civitai_model
+from .utils.video import load_video, save_video
+
+__all__ = [
+    "FluxImagePipeline",
+    "SDXLImagePipeline",
+    "SDImagePipeline",
+    "WanVideoPipeline",
+    "FluxModelConfig",
+    "SDXLModelConfig",
+    "SDModelConfig",
+    "WanModelConfig",
+    "fetch_model",
+    "fetch_modelscope_model",
+    "fetch_civitai_model",
+    "load_video",
+    "save_video",
+]

diffsynth_engine/algorithm/noise_scheduler/__init__.py

@@ -0,0 +1,21 @@
+from .stable_diffusion.linear import ScaledLinearScheduler
+from .stable_diffusion.beta import BetaScheduler
+from .stable_diffusion.karras import KarrasScheduler
+from .stable_diffusion.exponential import ExponentialScheduler
+from .stable_diffusion.ddim import DDIMScheduler
+from .stable_diffusion.sgm_uniform import SGMUniformScheduler
+from .flow_match.recifited_flow import RecifitedFlowScheduler
+from .flow_match.flow_ddim import FlowDDIMScheduler
+from .flow_match.flow_beta import FlowBetaScheduler
+
+__all__ = [
+    "ScaledLinearScheduler",
+    "BetaScheduler",
+    "KarrasScheduler",
+    "ExponentialScheduler",
+    "DDIMScheduler",
+    "SGMUniformScheduler",
+    "RecifitedFlowScheduler",
+    "FlowDDIMScheduler",
+    "FlowBetaScheduler",
+]

diffsynth_engine/algorithm/noise_scheduler/base_scheduler.py

@@ -0,0 +1,10 @@
+import torch
+
+
+def append_zero(x):
+    return torch.cat([x, x.new_zeros([1])])
+
+
+class BaseScheduler:
+    def schedule(self, num_inference_steps: int):
+        raise NotImplementedError()

diffsynth_engine/algorithm/noise_scheduler/flow_match/__init__.py

@@ -0,0 +1,5 @@
+from .recifited_flow import RecifitedFlowScheduler
+from .flow_ddim import FlowDDIMScheduler
+from .flow_beta import FlowBetaScheduler
+
+__all__ = ["RecifitedFlowScheduler", "FlowDDIMScheduler", "FlowBetaScheduler"]

diffsynth_engine/algorithm/noise_scheduler/flow_match/flow_beta.py

@@ -0,0 +1,28 @@
+import torch
+import numpy as np
+import scipy.stats as stats
+
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+from diffsynth_engine.algorithm.noise_scheduler.flow_match.recifited_flow import RecifitedFlowScheduler
+
+
+class FlowBetaScheduler(RecifitedFlowScheduler):
+    def __init__(self):
+        super().__init__()
+        self.alpha = 0.6
+        self.beta = 0.6
+
+    def schedule(self, num_inference_steps: int, mu: float | None = None, sigmas: torch.Tensor | None = None):
+        pseudo_timestep_range = 10000
+        inner_sigmas = torch.arange(1, pseudo_timestep_range + 1, 1) / pseudo_timestep_range
+        inner_sigmas = self._time_shift(mu, 1.0, inner_sigmas)
+        sigma_min = inner_sigmas[0]
+        sigma_max = inner_sigmas[-1]
+
+        timesteps = 1 - np.linspace(0, 1, num_inference_steps)
+        timesteps = [stats.beta.ppf(x, self.alpha, self.beta) for x in timesteps]
+        sigmas = [sigma_min + (x * (sigma_max - sigma_min)) for x in timesteps]
+        sigmas = torch.FloatTensor(sigmas)
+        timesteps = self._sigma_to_t(sigmas)
+        sigmas = append_zero(sigmas)
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/flow_match/flow_ddim.py

@@ -0,0 +1,25 @@
+import torch
+
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+from diffsynth_engine.algorithm.noise_scheduler.flow_match.recifited_flow import RecifitedFlowScheduler
+
+
+class FlowDDIMScheduler(RecifitedFlowScheduler):
+    def __init__(self, shift=1.0, num_train_timesteps=1000, use_dynamic_shifting=False):
+        super().__init__(shift, num_train_timesteps, use_dynamic_shifting)
+        self.pseudo_timestep_range = 10000
+
+    def schedule(self, num_inference_steps: int, mu: float | None = None, sigmas: torch.Tensor | None = None):
+        inner_sigmas = torch.arange(1, self.pseudo_timestep_range + 1, 1) / self.pseudo_timestep_range
+        inner_sigmas = self._time_shift(mu, 1.0, inner_sigmas)
+        sigmas = []
+        ss = max(len(inner_sigmas) // num_inference_steps, 1)
+        for i in range(1, len(inner_sigmas), ss):
+            sigmas.append(float(inner_sigmas[i]))
+        sigmas = sigmas[::-1]
+        sigmas = torch.FloatTensor(sigmas)
+
+        timesteps = self._sigma_to_t(sigmas)
+        sigmas = append_zero(sigmas)
+
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/flow_match/recifited_flow.py

@@ -0,0 +1,50 @@
+import torch
+import math
+
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero, BaseScheduler
+
+
+class RecifitedFlowScheduler(BaseScheduler):
+    def __init__(
+        self,
+        shift=1.0,
+        sigma_min=0.001,
+        sigma_max=1.0,
+        num_train_timesteps=1000,
+        use_dynamic_shifting=False,
+    ):
+        self.shift = shift
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.num_train_timesteps = num_train_timesteps
+        self.use_dynamic_shifting = use_dynamic_shifting
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.num_train_timesteps
+
+    def _t_to_sigma(self, t):
+        return t / self.num_train_timesteps
+
+    def _time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def _shift_sigma(self, sigma: torch.Tensor, shift: float):
+        return shift * sigma / (1 + (shift - 1) * sigma)
+
+    def schedule(
+        self,
+        num_inference_steps: int,
+        mu: float | None = None,
+        sigma_min: float | None = None,
+        sigma_max: float | None = None,
+    ):
+        sigma_min = self.sigma_min if sigma_min is None else sigma_min
+        sigma_max = self.sigma_max if sigma_max is None else sigma_max
+        sigmas = torch.linspace(sigma_max, sigma_min, num_inference_steps)
+        if self.use_dynamic_shifting:
+            sigmas = self._time_shift(mu, 1.0, sigmas)  # FLUX
+        else:
+            sigmas = self._shift_sigma(sigmas, self.shift)
+        timesteps = sigmas * self.num_train_timesteps
+        sigmas = append_zero(sigmas)
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/stable_diffusion/beta.py

@@ -0,0 +1,26 @@
+import torch
+import numpy as np
+import scipy.stats as stats
+
+from diffsynth_engine.algorithm.noise_scheduler.stable_diffusion.linear import ScaledLinearScheduler
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+
+
+class BetaScheduler(ScaledLinearScheduler):
+    """
+    Implemented based on: https://arxiv.org/abs/2407.12173
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.alpha = 0.6
+        self.beta = 0.6
+
+    def schedule(self, num_inference_steps: int):
+        timesteps = 1 - np.linspace(0, 1, num_inference_steps)
+        timesteps = [stats.beta.ppf(x, self.alpha, self.beta) for x in timesteps]
+        sigmas = [self.sigma_min + (x * (self.sigma_max - self.sigma_min)) for x in timesteps]
+        sigmas = torch.FloatTensor(sigmas).to(self.device)
+        timesteps = self.sigma_to_t(sigmas)
+        sigmas = append_zero(sigmas)
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/stable_diffusion/ddim.py

@@ -0,0 +1,25 @@
+import torch
+
+from diffsynth_engine.algorithm.noise_scheduler.stable_diffusion.linear import ScaledLinearScheduler
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+
+
+class DDIMScheduler(ScaledLinearScheduler):
+    """
+    Implemented based on: https://arxiv.org/pdf/2010.02502.pdf
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def schedule(self, num_inference_steps: int):
+        inner_sigmas = self.get_sigmas()
+        sigmas = []
+        ss = max(len(inner_sigmas) // num_inference_steps, 1)
+        for i in range(1, len(inner_sigmas), ss):
+            sigmas.append(float(inner_sigmas[i]))
+        sigmas = sigmas[::-1]
+        sigmas = torch.FloatTensor(sigmas).to(self.device)
+        timesteps = self.sigma_to_t(sigmas)
+        sigmas = append_zero(sigmas)
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/stable_diffusion/exponential.py

@@ -0,0 +1,19 @@
+import torch
+import math
+
+from diffsynth_engine.algorithm.noise_scheduler.stable_diffusion.linear import ScaledLinearScheduler
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+
+
+class ExponentialScheduler(ScaledLinearScheduler):
+    def __init__(self):
+        super().__init__()
+
+    def schedule(self, num_inference_steps: int):
+        """Constructs an exponential noise schedule."""
+        sigmas = torch.linspace(
+            math.log(self.sigma_max), math.log(self.sigma_min), num_inference_steps, device=self.device
+        ).exp()
+        timesteps = self.sigma_to_t(sigmas)
+        sigmas = append_zero(sigmas)
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/stable_diffusion/karras.py

@@ -0,0 +1,21 @@
+import torch
+
+from diffsynth_engine.algorithm.noise_scheduler.stable_diffusion.linear import ScaledLinearScheduler
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+
+
+class KarrasScheduler(ScaledLinearScheduler):
+    def __init__(self):
+        super().__init__()
+        self.rho = 7.0
+        self.device = "cpu"
+
+    def schedule(self, num_inference_steps: int):
+        """Constructs the noise schedule of Karras et al. (2022)."""
+        ramp = torch.linspace(0, 1, num_inference_steps, device=self.device)
+        min_inv_rho = self.sigma_min ** (1 / self.rho)
+        max_inv_rho = self.sigma_max ** (1 / self.rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** self.rho
+        timesteps = self.sigma_to_t(sigmas)
+        sigmas = append_zero(sigmas).to(self.device)
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/stable_diffusion/linear.py

@@ -0,0 +1,77 @@
+import torch
+
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import BaseScheduler, append_zero
+
+
+def linear_beta_schedule(beta_start: float = 0.00085, beta_end: float = 0.0120, num_train_steps: int = 1000):
+    """
+    DDPM Schedule
+    """
+    return torch.linspace(beta_start, beta_end, num_train_steps)
+
+
+def scaled_linear_beta_schedule(beta_start: float = 0.00085, beta_end: float = 0.0120, num_train_steps: int = 1000):
+    """
+    Stable Diffusion Schedule
+    """
+    return torch.linspace(beta_start**0.5, beta_end**0.5, num_train_steps) ** 2
+
+
+class ScaledLinearScheduler(BaseScheduler):
+    def __init__(self):
+        self.device = "cpu"
+        self.num_train_steps = 1000
+        self.beta_start = 0.00085
+        self.beta_end = 0.0120
+        self.sigmas = self.get_sigmas()
+        self.log_sigmas = self.sigmas.log()
+
+    @property
+    def sigma_min(self):
+        return self.sigmas[0]
+
+    @property
+    def sigma_max(self):
+        return self.sigmas[-1]
+
+    def get_sigmas(self):
+        # Stable Diffusion Sigmas
+        # len(sigmas) == 1000, sigma_min=sigmas[0] == 0.0292, sigma_max=sigmas[-1] == 14.6146
+        betas = scaled_linear_beta_schedule(
+            beta_start=self.beta_start, beta_end=self.beta_end, num_train_steps=self.num_train_steps
+        )
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+        return sigmas
+
+    def sigma_to_t(self, sigma):
+        """
+        找到sigma.log()在self.log_sigmas中的位置(low和high), 进行加权插值得到t
+        """
+        log_sigma = sigma.log()
+        dists = log_sigma - self.log_sigmas[:, None]
+        low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.log_sigmas.shape[0] - 2)
+        high_idx = low_idx + 1
+        low, high = self.log_sigmas[low_idx], self.log_sigmas[high_idx]
+        w = (low - log_sigma) / (low - high)
+        w = w.clamp(0, 1)
+        t = (1 - w) * low_idx + w * high_idx
+        return t.view(sigma.shape)
+
+    def t_to_sigma(self, t):
+        """
+        对t进行floor和ceil, 得到low_idx和high_idx, 计算对应位置的log_sigma, 进行加权插值并exp得到sigma
+        """
+        t = t.float()
+        low_idx, high_idx, w = t.floor().long(), t.ceil().long(), t.frac()
+        log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
+        return log_sigma.exp()
+
+    def schedule(self, num_inference_steps: int):
+        """
+        Uniformly sample timesteps for inference
+        """
+        timesteps = torch.linspace(self.num_train_steps - 1, 0, num_inference_steps, device=self.sigmas.device)
+        sigmas = append_zero(self.t_to_sigma(timesteps))
+        return sigmas, timesteps

diffsynth_engine/algorithm/noise_scheduler/stable_diffusion/sgm_uniform.py

@@ -0,0 +1,17 @@
+import torch
+
+from diffsynth_engine.algorithm.noise_scheduler.stable_diffusion.linear import ScaledLinearScheduler
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero
+
+
+class SGMUniformScheduler(ScaledLinearScheduler):
+    def __init__(self):
+        super().__init__()
+
+    def schedule(self, num_inference_steps: int):
+        # suppose sigma_min and sigma_max is default value
+        timesteps = torch.linspace(999, 0, num_inference_steps + 1)[:-1]
+        sigmas = [self.t_to_sigma(timestep) for timestep in timesteps]
+        sigmas = torch.FloatTensor(sigmas).to(self.device)
+        sigmas = append_zero(sigmas)
+        return sigmas, timesteps

diffsynth_engine/algorithm/sampler/__init__.py

@@ -0,0 +1,19 @@
+from .stable_diffusion.ddpm import DDPMSampler
+from .stable_diffusion.euler import EulerSampler
+from .stable_diffusion.euler_ancestral import EulerAncestralSampler
+from .stable_diffusion.dpmpp_2m import DPMSolverPlusPlus2MSampler
+from .stable_diffusion.dpmpp_2m_sde import DPMSolverPlusPlus2MSDESampler
+from .stable_diffusion.dpmpp_3m_sde import DPMSolverPlusPlus3MSDESampler
+from .stable_diffusion.deis import DEISSampler
+from .flow_match.flow_match_euler import FlowMatchEulerSampler
+
+__all__ = [
+    "DDPMSampler",
+    "EulerSampler",
+    "EulerAncestralSampler",
+    "DPMSolverPlusPlus2MSampler",
+    "DPMSolverPlusPlus2MSDESampler",
+    "DPMSolverPlusPlus3MSDESampler",
+    "DEISSampler",
+    "FlowMatchEulerSampler",
+]

diffsynth_engine/algorithm/sampler/flow_match/flow_match_euler.py

@@ -0,0 +1,22 @@
+import torch
+
+
+class FlowMatchEulerSampler:
+    def initialize(self, init_latents, timesteps, sigmas, mask=None):
+        self.init_latents = init_latents
+        self.timesteps = timesteps
+        self.sigmas = sigmas
+        self.mask = mask
+
+    def step(self, latents, model_outputs, i):
+        if self.mask is not None:
+            model_outputs = model_outputs * self.mask + self.init_latents * (1 - self.mask)
+
+        dt = self.sigmas[i + 1] - self.sigmas[i]
+        latents = latents.to(dtype=torch.float32)
+        latents = latents + model_outputs * dt
+        latents = latents.to(dtype=model_outputs.dtype)
+        return latents
+
+    def add_noise(self, latents, noise, sigma):
+        return (1 - sigma) * latents + noise * sigma

diffsynth_engine/algorithm/sampler/stable_diffusion/brownian_tree.py

@@ -0,0 +1,54 @@
+import torch
+import torchsde
+
+
+class BatchedBrownianTree:
+    """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
+
+    def __init__(self, x, t0, t1, seed=None, **kwargs):
+        t0, t1, self.sign = self.sort(t0, t1)
+        w0 = kwargs.get("w0", torch.zeros_like(x))
+        if seed is None:
+            seed = torch.randint(0, 2**63 - 1, []).item()
+        self.batched = True
+        try:
+            assert len(seed) == x.shape[0]
+            w0 = w0[0]
+        except TypeError:
+            seed = [seed]
+            self.batched = False
+        self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
+
+    @staticmethod
+    def sort(a, b):
+        return (a, b, 1) if a < b else (b, a, -1)
+
+    def __call__(self, t0, t1):
+        t0, t1, sign = self.sort(t0, t1)
+        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
+        return w if self.batched else w[0]
+
+
+class BrownianTreeNoiseSampler:
+    """A noise sampler backed by a torchsde.BrownianTree.
+
+    Args:
+        x (Tensor): The tensor whose shape, device and dtype to use to generate
+            random samples.
+        sigma_min (float): The low end of the valid interval.
+        sigma_max (float): The high end of the valid interval.
+        seed (int or List[int]): The random seed. If a list of seeds is
+            supplied instead of a single integer, then the noise sampler will
+            use one BrownianTree per batch item, each with its own seed.
+        transform (callable): A function that maps sigma to the sampler's
+            internal timestep.
+    """
+
+    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x):
+        self.transform = transform
+        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
+        self.tree = BatchedBrownianTree(x, t0, t1, seed)
+
+    def __call__(self, sigma, sigma_next):
+        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
+        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()

diffsynth_engine/algorithm/sampler/stable_diffusion/ddpm.py

@@ -0,0 +1,32 @@
+import torch
+from .epsilon import EpsilonSampler
+
+
+class DDPMSampler(EpsilonSampler):
+    def _step_function(self, x, sigma, sigma_prev, noise):
+        alpha_cumprod = 1 / ((sigma * sigma) + 1)
+        alpha_cumprod_prev = 1 / ((sigma_prev * sigma_prev) + 1)
+        alpha = alpha_cumprod / alpha_cumprod_prev
+
+        mu = (1.0 / alpha) ** 0.5 * (x - (1 - alpha) * noise / (1 - alpha_cumprod) ** 0.5)
+        if sigma_prev > 0:
+            # Caution: this randn tensor needs to be controlled by `torch.manual_seed`.
+            mu += ((1 - alpha) * (1.0 - alpha_cumprod_prev) / (1.0 - alpha_cumprod)) ** 0.5 * torch.randn_like(x)
+        return mu
+
+    def step(self, latents, model_outputs, i):
+        sigma = self.sigmas[i]
+        sigma_next = self.sigmas[i + 1]
+        latents = self._scaling(sigma, latents)
+
+        denoised = self._to_denoised(sigma, model_outputs, latents)
+        latents = self._step_function(
+            latents / (1.0 + sigma**2.0) ** 0.5, sigma, sigma_next, (latents - denoised) / sigma
+        )
+
+        latents *= (1.0 + sigma_next**2.0) ** 0.5
+
+        return self._unscaling(self.sigmas[i + 1], latents)
+
+    def step2(self, latents, model_outputs, i):
+        return self._step_function(latents, self.sigmas[i], self.sigmas[i + 1], model_outputs)

diffsynth_engine/algorithm/sampler/stable_diffusion/deis.py

@@ -0,0 +1,125 @@
+import torch
+
+from diffsynth_engine.algorithm.sampler.stable_diffusion.epsilon import EpsilonSampler
+
+
+class DEISSampler(EpsilonSampler):
+    """
+    According to the implementation of the webui forge, deis_mode only supports tab and rhoab.
+    """
+
+    def initialize(self, init_latents, timesteps, sigmas, mask):
+        super().initialize(init_latents, timesteps, sigmas, mask)
+        self.max_order = 3
+        self.sigmas = sigmas
+        self.timesteps = timesteps
+        self.lower_order_nums = 0
+        self.coeff_list = get_deis_coeff_list(self.sigmas, self.max_order)
+        self.coeff_buffer = []
+
+    def step(self, latents, model_outputs, i):
+        s, s_next = self.sigmas[i], self.sigmas[i + 1]
+        denoised = latents - model_outputs * s
+
+        d = (latents - denoised) / s
+        order = min(self.max_order, i + 1)
+        if self.sigmas[i + 1] <= 0:
+            order = 1
+        if order == 1:
+            x_next = latents + (s_next - s) * d
+        elif order == 2:
+            coeff, coeff_prev1 = self.coeff_list[i]
+            x_next = latents + coeff * d + coeff_prev1 * self.coeff_buffer[-1]
+        elif order == 3:
+            coeff, coeff_prev1, coeff_prev2 = self.coeff_list[i]
+            x_next = latents + coeff * d + coeff_prev1 * self.coeff_buffer[-1] + coeff_prev2 * self.coeff_buffer[-2]
+        elif order == 4:
+            coeff, coeff_prev1, coeff_prev2, coeff_prev3 = self.coeff_list[i]
+            x_next = (
+                latents
+                + coeff * d
+                + coeff_prev1 * self.coeff_buffer[-1]
+                + coeff_prev2 * self.coeff_buffer[-2]
+                + coeff_prev3 * self.coeff_buffer[-3]
+            )
+
+        if len(self.coeff_buffer) == self.max_order - 1:
+            for k in range(self.max_order - 2):
+                self.coeff_buffer[k] = self.coeff_buffer[k + 1]
+            self.coeff_buffer[-1] = d
+        else:
+            self.coeff_buffer.append(d)
+        return x_next
+
+
+# Taken from: https://github.com/zju-pi/diff-sampler/blob/main/gits-main/solver_utils.py
+# under Apache 2 license
+# A pytorch reimplementation of DEIS (https://github.com/qsh-zh/deis).
+#############################
+### Utils for DEIS solver ###
+#############################
+# ----------------------------------------------------------------------------
+# Transfer from the input time (sigma) used in EDM to that (t) used in DEIS.
+
+
+def vp_sigma_inv(beta_d, beta_min, sigma):
+    return ((beta_min**2 + 2 * beta_d * (sigma**2 + 1).log()).sqrt() - beta_min) / beta_d
+
+
+def edm2t(edm_steps, epsilon_s=1e-3, sigma_min=0.002, sigma_max=80):
+    vp_beta_d = (
+        2
+        * (torch.log(torch.tensor(sigma_min) ** 2 + 1) / epsilon_s - torch.log(torch.tensor(sigma_max) ** 2 + 1))
+        / (epsilon_s - 1)
+    )
+    vp_beta_min = torch.log(torch.tensor(sigma_max) ** 2 + 1) - 0.5 * vp_beta_d
+    t_steps = vp_sigma_inv(vp_beta_d, vp_beta_min, edm_steps)
+    return t_steps, vp_beta_min, vp_beta_d + vp_beta_min
+
+
+def cal_poly(prev_t, j, taus):
+    poly = 1
+    for k in range(prev_t.shape[0]):
+        if k == j:
+            continue
+        poly *= (taus - prev_t[k]) / (prev_t[j] - prev_t[k])
+    return poly
+
+
+def t2alpha_fn(beta_0, beta_1, t):
+    return torch.exp(-0.5 * t**2 * (beta_1 - beta_0) - t * beta_0)
+
+
+def cal_intergrand(beta_0, beta_1, taus):
+    with torch.inference_mode(mode=False):
+        taus = taus.clone()
+        beta_0 = beta_0.clone()
+        beta_1 = beta_1.clone()
+        with torch.enable_grad():
+            taus.requires_grad_(True)
+            alpha = t2alpha_fn(beta_0, beta_1, taus)
+            log_alpha = alpha.log()
+            log_alpha.sum().backward()
+            d_log_alpha_dtau = taus.grad
+    integrand = -0.5 * d_log_alpha_dtau / torch.sqrt(alpha * (1 - alpha))
+    return integrand
+
+
+def get_deis_coeff_list(t_steps, max_order, N=10000):
+    t_steps, beta_0, beta_1 = edm2t(t_steps)
+    C = []
+    for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
+        order = min(i + 1, max_order)
+        if order == 1:
+            C.append([])
+        else:
+            taus = torch.linspace(t_cur, t_next, N).to(t_next.device)
+            dtau = (t_next - t_cur) / N
+            prev_t = t_steps[[i - k for k in range(order)]]
+            coeff_temp = []
+            integrand = cal_intergrand(beta_0, beta_1, taus)
+            for j in range(order):
+                poly = cal_poly(prev_t, j, taus)
+                coeff_temp.append(torch.sum(integrand * poly) * dtau)
+            C.append(coeff_temp)
+    return C

diffsynth_engine/algorithm/sampler/stable_diffusion/dpmpp_2m.py

@@ -0,0 +1,29 @@
+from diffsynth_engine.algorithm.sampler.stable_diffusion.epsilon import EpsilonSampler
+
+
+class DPMSolverPlusPlus2MSampler(EpsilonSampler):
+    """
+    DPM Solver++ 2M sampler
+    """
+
+    def initialize(self, init_latents, timesteps, sigmas, mask):
+        super().initialize(init_latents, timesteps, sigmas, mask)
+        self.old_denoised = None
+
+    def step(self, latents, model_outputs, i):
+        s_prev, s, s_next = self.sigmas[i - 1], self.sigmas[i], self.sigmas[i + 1]
+        t_prev, t, t_next = self._sigma_to_t(s_prev), self._sigma_to_t(s), self._sigma_to_t(s_next)
+        h = t_next - t
+        x = self._scaling(latents, s)
+        denoised = self._to_denoised(s, model_outputs, x)
+        if self.old_denoised is None or s_next == 0:
+            self.old_denoised = denoised
+            return (s_next / s) * x - (-h).expm1() * denoised
+        h_last = t - t_prev
+        r = h_last / h
+        denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * self.old_denoised
+        x = (s_next / s) * x - (-h).expm1() * denoised_d
+        return self._unscaling(x, s_next)
+
+    def _sigma_to_t(self, sigma):
+        return sigma.log().neg()