I

RealCustom/schedulers/base.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import torch
+from dataclasses import dataclass
+from abc import ABC
+from typing import Optional, Union, List
+
+
+@dataclass
+class SchedulerConversionOutput:
+    pred_epsilon: torch.Tensor
+    pred_original_sample: torch.Tensor
+    pred_velocity: torch.Tensor
+
+
+@dataclass
+class SchedulerStepOutput:
+    prev_sample: torch.Tensor
+    pred_original_sample: Optional[torch.Tensor] = None
+
+
+class Scheduler(ABC):
+    prediction_types = ["epsilon", "sample", "v_prediction"]
+    timesteps_types = ["leading", "linspace", "trailing"]
+
+    def __init__(
+        self,
+        num_train_timesteps: int,
+        num_inference_timesteps: int,
+        betas: torch.Tensor,
+        inference_timesteps: Union[str, List[int]] = "trailing",
+        set_alpha_to_one: bool = True,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: torch.dtype = torch.float32
+    ):
+        assert num_train_timesteps > 0
+        assert num_train_timesteps >= num_inference_timesteps
+        assert num_train_timesteps == betas.size(0)
+        assert betas.ndim == 1
+
+        self.device = device or betas.device
+        self.dtype = dtype
+
+        self.num_train_timesteps = num_train_timesteps
+        self.num_inference_timesteps = num_inference_timesteps
+
+        self.betas = betas.to(device=device, dtype=dtype)
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        self.final_alpha_cumprod = torch.tensor(1.0, device=self.device, dtype=self.dtype) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        if isinstance(inference_timesteps, list):
+            # If user defines a custom inference timestep, directly assign it.
+            assert len(inference_timesteps) == num_inference_timesteps
+            self.timesteps = torch.tensor(inference_timesteps, device=self.device, dtype=torch.int)
+        elif inference_timesteps == "trailing":
+            # Example 20 steps: [999, 949, 899, 849, 799, 749, 699, 649, 599, 549, 499, 449, 399, 349, 299, 249, 199, 149,  99,  49]
+            self.timesteps = torch.arange(num_train_timesteps - 1, -1, -num_train_timesteps / num_inference_timesteps, device=self.device).round().int()
+        elif inference_timesteps == "linspace":
+            # Example 20 steps: [999, 946, 894, 841, 789, 736, 684, 631, 578, 526, 473, 421, 368, 315, 263, 210, 158, 105,  53,   0]
+            self.timesteps = torch.linspace(0, num_train_timesteps - 1, num_inference_timesteps, device=self.device).round().int().flip(0)
+        elif inference_timesteps == "leading":
+            # Original SD and DDIM paper may have a bug: <https://github.com/huggingface/diffusers/issues/2585>
+            # The inference timestep does not start from 999.
+            # Example 20 steps: [950, 900, 850, 800, 750, 700, 650, 600, 550, 500, 450, 400, 350, 300, 250, 200, 150, 100,  50,   0]
+            self.timesteps = torch.arange(0, num_train_timesteps, num_train_timesteps // num_inference_timesteps, device=self.device, dtype=torch.int).flip(0)
+        else:
+            raise NotImplementedError
+
+    def reset(self):
+        pass
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: Union[torch.Tensor, int],
+    ) -> torch.Tensor:
+        alpha_prod_t = self.alphas_cumprod[timesteps].reshape(-1, *([1] * (original_samples.ndim - 1)))
+        return alpha_prod_t ** (0.5) * original_samples + (1 - alpha_prod_t) ** (0.5) * noise
+
+    def convert_output(
+        self,
+        model_output: torch.Tensor,
+        model_output_type: str,
+        sample: torch.Tensor,
+        timesteps: Union[torch.Tensor, int]
+    ) -> SchedulerConversionOutput:
+        assert model_output_type in self.prediction_types
+
+        alpha_prod_t = self.alphas_cumprod[timesteps].reshape(-1, *([1] * (sample.ndim - 1)))
+        beta_prod_t = 1 - alpha_prod_t
+
+        if model_output_type == "epsilon":
+            pred_epsilon = model_output
+            pred_original_sample = (sample - beta_prod_t ** (0.5) * pred_epsilon) / alpha_prod_t ** (0.5)
+            pred_velocity = alpha_prod_t ** (0.5) * pred_epsilon - (1 - alpha_prod_t) ** (0.5) * pred_original_sample
+        elif model_output_type == "sample":
+            pred_original_sample = model_output
+            pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
+            pred_velocity = alpha_prod_t ** (0.5) * pred_epsilon - (1 - alpha_prod_t) ** (0.5) * pred_original_sample
+        elif model_output_type == "v_prediction":
+            pred_velocity = model_output
+            pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+            pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
+        else:
+            raise ValueError("Unknown prediction type")
+
+        return SchedulerConversionOutput(
+            pred_epsilon=pred_epsilon,
+            pred_original_sample=pred_original_sample,
+            pred_velocity=pred_velocity)
+
+    def get_velocity(
+        self,
+        sample: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor
+    ) -> torch.FloatTensor:
+        alpha_prod_t = self.alphas_cumprod[timesteps].reshape(-1, *([1] * (sample.ndim - 1)))
+        return alpha_prod_t ** (0.5) * noise - (1 - alpha_prod_t) ** (0.5) * sample

RealCustom/schedulers/ddim.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from .base import *
+
+
+class DDIMScheduler(Scheduler):
+    def step(
+        self,
+        model_output: torch.Tensor,
+        model_output_type: str,
+        timestep: Union[torch.Tensor, int],
+        sample: torch.Tensor,
+        eta: float = 0.0,
+        clip_sample: bool = False,
+        dynamic_threshold: Optional[float] = None,
+        variance_noise: Optional[torch.Tensor] = None,
+    ) -> SchedulerStepOutput:
+        # 1. get previous step value (t-1)
+        if not isinstance(timestep, torch.Tensor):
+            timestep = torch.tensor(timestep, device=self.device, dtype=torch.int)
+
+        idx = timestep.reshape(-1, 1).eq(self.timesteps.reshape(1, -1)).nonzero()[:, 1]
+        prev_timestep = self.timesteps[idx.add(1).clamp_max(self.num_inference_timesteps - 1)]
+
+        # 2. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[timestep].reshape(-1, *([1] * (sample.ndim - 1)))
+        alpha_prod_t_prev = torch.where(idx < self.num_inference_timesteps - 1, self.alphas_cumprod[prev_timestep], self.final_alpha_cumprod).reshape(-1, *([1] * (sample.ndim - 1)))
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        # 3. compute predicted original sample from predicted noise also called
+        model_output_conversion = self.convert_output(model_output, model_output_type, sample, timestep)
+        pred_original_sample = model_output_conversion.pred_original_sample
+        pred_epsilon = model_output_conversion.pred_epsilon
+
+        # 4. Clip or threshold "predicted x_0"
+        if clip_sample:
+            pred_original_sample = torch.clamp(pred_original_sample, -1, 1)
+            pred_epsilon = self.convert_output(pred_original_sample, "sample", sample, timestep).pred_epsilon
+
+        if dynamic_threshold is not None:
+            # Dynamic thresholding in https://arxiv.org/abs/2205.11487
+            dynamic_max_val = pred_original_sample \
+                .flatten(1) \
+                .abs() \
+                .float() \
+                .quantile(dynamic_threshold, dim=1) \
+                .type_as(pred_original_sample) \
+                .clamp_min(1) \
+                .view(-1, *([1] * (pred_original_sample.ndim - 1)))
+            pred_original_sample = pred_original_sample.clamp(-dynamic_max_val, dynamic_max_val) / dynamic_max_val
+            pred_epsilon = self.convert_output(pred_original_sample, "sample", sample, timestep).pred_epsilon
+
+        # 5. compute variance: "sigma_t(η)" -> see formula (16) from https://arxiv.org/pdf/2010.02502.pdf
+        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
+        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+        std_dev_t = eta * variance ** (0.5)
+
+        # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon
+
+        # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
+
+        # 8. add "random noise" if needed.
+        if eta > 0:
+            if variance_noise is None:
+                variance_noise = torch.randn_like(model_output)
+            prev_sample = prev_sample + std_dev_t * variance_noise
+
+        return SchedulerStepOutput(
+            prev_sample=prev_sample,
+            pred_original_sample=pred_original_sample)

RealCustom/schedulers/dpm_m.py

@@ -0,0 +1,412 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023 TSAIL Team 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.
+
+# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
+
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+
+from .base import *
+
+
+class DPMSolverMultistepScheduler(Scheduler):
+    def __init__(
+        self,
+        # Generic scheduler settings
+        num_inference_timesteps: int,
+        betas: torch.Tensor,
+        num_train_timesteps: int = 1000,
+        inference_timesteps: Union[str, List[str]] = "trailing",
+        set_alpha_to_one: bool = True,
+        device: Optional[torch.device] = None,
+        dtype: torch.dtype = torch.float32,
+        # DPM scheduler settings
+        solver_order: int = 2,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        use_karras_sigmas: bool = False,
+    ):
+        super().__init__(
+            num_train_timesteps=num_train_timesteps,
+            num_inference_timesteps=num_inference_timesteps,
+            betas=betas,
+            inference_timesteps=inference_timesteps,
+            set_alpha_to_one=set_alpha_to_one,
+            device=device,
+            dtype=dtype,
+        )
+
+        self.solver_order = solver_order
+        self.solver_type = solver_type
+        self.lower_order_final = lower_order_final
+        self.algorithm_type = algorithm_type
+
+        # Currently we only support VP-type noise schedule
+        self.alpha_t = torch.sqrt(self.alphas_cumprod)
+        self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
+        self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
+
+        sigmas = torch.sqrt((1 - self.alphas_cumprod) / self.alphas_cumprod)
+        if use_karras_sigmas:
+            log_sigmas = torch.log(sigmas)
+            sigmas = self._convert_to_karras(
+                in_sigmas=sigmas, num_inference_timesteps=num_inference_timesteps)
+            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas)
+                                 for sigma in sigmas]).round()
+            timesteps = np.flip(timesteps).copy().astype(np.int64)
+            self.timesteps = torch.from_numpy(timesteps).to(device)
+            sigmas = torch.from_numpy(sigmas).to(device)
+        self.sigmas = sigmas
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # settings for DPM-Solver
+        if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++", "deis"]:
+            raise NotImplementedError(
+                f"{algorithm_type} does is not implemented for {self.__class__}")
+
+        if solver_type not in ["midpoint", "heun", "logrho", "bh1", "bh2"]:
+            raise NotImplementedError(
+                f"{solver_type} does is not implemented for {self.__class__}")
+
+        # setable values
+        self.reset()
+
+    def reset(self):
+        self.model_outputs = [None] * self.solver_order
+        self.lower_order_nums = 0
+
+    def _sigma_to_t(self, sigma, log_sigmas):
+        # get log sigma
+        log_sigma = np.log(sigma)
+
+        # get distribution
+        dists = log_sigma - log_sigmas[:, np.newaxis]
+
+        # get sigmas range
+        low_idx = np.cumsum((dists >= 0), axis=0).argmax(
+            axis=0).clip(max=log_sigmas.shape[0] - 2)
+        high_idx = low_idx + 1
+
+        low = log_sigmas[low_idx]
+        high = log_sigmas[high_idx]
+
+        # interpolate sigmas
+        w = (low - log_sigma) / (low - high)
+        w = np.clip(w, 0, 1)
+
+        # transform interpolation to time range
+        t = (1 - w) * low_idx + w * high_idx
+        t = t.reshape(sigma.shape)
+        return t
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.FloatTensor, num_inference_timesteps) -> torch.FloatTensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        sigma_min: float = in_sigmas[-1].item()
+        sigma_max: float = in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_timesteps)
+        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
+        return sigmas
+
+    def dpm_solver_first_order_update(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        One step for the first-order DPM-Solver (equivalent to DDIM).
+
+        See https://arxiv.org/abs/2206.00927 for the detailed derivation.
+
+        Args:
+            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
+            timestep (`int`): current discrete timestep in the diffusion chain.
+            prev_timestep (`int`): previous discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                current instance of sample being created by diffusion process.
+
+        Returns:
+            `torch.FloatTensor`: the sample tensor at the previous timestep.
+        """
+        lambda_t, lambda_s = self.lambda_t[prev_timestep], self.lambda_t[timestep]
+        alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep]
+        sigma_t, sigma_s = self.sigma_t[prev_timestep], self.sigma_t[timestep]
+        h = lambda_t - lambda_s
+        if self.algorithm_type == "dpmsolver++":
+            x_t = (sigma_t / sigma_s) * sample - \
+                (alpha_t * (torch.exp(-h) - 1.0)) * model_output
+        elif self.algorithm_type == "dpmsolver":
+            x_t = (alpha_t / alpha_s) * sample - \
+                (sigma_t * (torch.exp(h) - 1.0)) * model_output
+        elif self.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            x_t = (
+                (sigma_t / sigma_s * torch.exp(-h)) * sample
+                + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output
+                + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+            )
+        elif self.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            x_t = (
+                (alpha_t / alpha_s) * sample
+                - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output
+                + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+            )
+        return x_t
+
+    def multistep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.FloatTensor],
+        timestep_list: List[int],
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        One step for the second-order multistep DPM-Solver.
+
+        Args:
+            model_output_list (`List[torch.FloatTensor]`):
+                direct outputs from learned diffusion model at current and latter timesteps.
+            timestep (`int`): current and latter discrete timestep in the diffusion chain.
+            prev_timestep (`int`): previous discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                current instance of sample being created by diffusion process.
+
+        Returns:
+            `torch.FloatTensor`: the sample tensor at the previous timestep.
+        """
+        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
+        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
+        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+        if self.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                )
+        elif self.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                )
+        elif self.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+        elif self.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+        return x_t
+
+    def multistep_dpm_solver_third_order_update(
+        self,
+        model_output_list: List[torch.FloatTensor],
+        timestep_list: List[int],
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        """
+        One step for the third-order multistep DPM-Solver.
+
+        Args:
+            model_output_list (`List[torch.FloatTensor]`):
+                direct outputs from learned diffusion model at current and latter timesteps.
+            timestep (`int`): current and latter discrete timestep in the diffusion chain.
+            prev_timestep (`int`): previous discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                current instance of sample being created by diffusion process.
+
+        Returns:
+            `torch.FloatTensor`: the sample tensor at the previous timestep.
+        """
+        t, s0, s1, s2 = prev_timestep, timestep_list[-1], timestep_list[-2], timestep_list[-3]
+        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
+        lambda_t, lambda_s0, lambda_s1, lambda_s2 = (
+            self.lambda_t[t],
+            self.lambda_t[s0],
+            self.lambda_t[s1],
+            self.lambda_t[s2],
+        )
+        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
+        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
+        r0, r1 = h_0 / h, h_1 / h
+        D0 = m0
+        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+        if self.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (sigma_t / sigma_s0) * sample
+                - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2
+            )
+        elif self.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (alpha_t / alpha_s0) * sample
+                - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2
+            )
+        return x_t
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        model_output_type: str,
+        timestep: int,
+        sample: torch.FloatTensor,
+    ) -> SchedulerStepOutput:
+        """
+        Step function propagating the sample with the multistep DPM-Solver.
+
+        Args:
+            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
+            timestep (`int`): current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                current instance of sample being created by diffusion process.
+            return_dict (`bool`): option for returning tuple rather than SchedulerOutput class
+
+        Returns:
+            [`~scheduling_utils.SchedulerOutput`] or `tuple`: [`~scheduling_utils.SchedulerOutput`] if `return_dict` is
+            True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.
+
+        """
+        if self.num_inference_timesteps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.to(self.timesteps.device)
+        step_index = (self.timesteps == timestep).nonzero()
+        if len(step_index) == 0:
+            step_index = len(self.timesteps) - 1
+        else:
+            step_index = step_index.item()
+        prev_timestep = 0 if step_index == len(
+            self.timesteps) - 1 else self.timesteps[step_index + 1]
+        lower_order_final = (
+            (step_index == len(self.timesteps) -
+             1) and self.lower_order_final and len(self.timesteps) < 15
+        )
+        lower_order_second = (
+            (step_index == len(self.timesteps) -
+             2) and self.lower_order_final and len(self.timesteps) < 15
+        )
+
+        model_output_convert = self.convert_output(
+            model_output, model_output_type=model_output_type, sample=sample, timesteps=timestep)
+        # DPM-Solver++ needs to solve an integral of the data prediction model.
+        if self.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
+            model_output = model_output_convert.pred_original_sample
+        # DPM-Solver needs to solve an integral of the noise prediction model.
+        elif self.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            model_output = model_output_convert.pred_epsilon
+
+        for i in range(self.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+        self.model_outputs[-1] = model_output
+
+        if self.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
+            noise = torch.randn_like(
+                model_output, device=model_output.device, dtype=model_output.dtype)
+        else:
+            noise = None
+
+        if self.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
+            prev_sample = self.dpm_solver_first_order_update(
+                model_output, timestep, prev_timestep, sample, noise=noise
+            )
+        elif self.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
+            timestep_list = [self.timesteps[step_index - 1], timestep]
+            prev_sample = self.multistep_dpm_solver_second_order_update(
+                self.model_outputs, timestep_list, prev_timestep, sample, noise=noise
+            )
+        else:
+            timestep_list = [self.timesteps[step_index - 2],
+                             self.timesteps[step_index - 1], timestep]
+            prev_sample = self.multistep_dpm_solver_third_order_update(
+                self.model_outputs, timestep_list, prev_timestep, sample
+            )
+
+        if self.lower_order_nums < self.solver_order:
+            self.lower_order_nums += 1
+
+        return SchedulerStepOutput(prev_sample=prev_sample, pred_original_sample=model_output_convert.pred_original_sample)

RealCustom/schedulers/dpm_s.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from .base import *
+
+
+class DPMSolverSingleStepScheduler(Scheduler):
+    def __init__(
+        self,
+        # Generic scheduler settings
+        num_train_timesteps: int,
+        num_inference_timesteps: int,
+        betas: torch.Tensor,
+        inference_timesteps: Union[str, List[int]] = "trailing",
+        set_alpha_to_one: bool = True,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: torch.dtype = torch.float32,
+        # DPM scheduler settings
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        solver_order: int = 2,
+        lower_order_final: bool = True,
+    ):
+        super().__init__(
+            num_train_timesteps=num_train_timesteps,
+            num_inference_timesteps=num_inference_timesteps,
+            betas=betas,
+            inference_timesteps=inference_timesteps,
+            set_alpha_to_one=set_alpha_to_one,
+            device=device,
+            dtype=dtype,
+        )
+
+        self.solver_order = solver_order
+        self.solver_type = solver_type
+        self.lower_order_final = lower_order_final
+        self.algorithm_type = algorithm_type
+
+        self.alpha_t = torch.sqrt(self.alphas_cumprod)
+        self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
+        self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
+
+        self.reset()
+
+    def reset(self):
+        self.model_outputs = [None] * self.solver_order
+        self.sample = None
+        self.order_list = self.get_order_list()
+        self.last_step_index = None
+
+    def get_order_list(self) -> List[int]:
+        steps = self.num_inference_timesteps
+        order = self.solver_order
+        # First step must be order 1
+        # Second step must be order 1 in case of terminal zero SNR
+        orders = [1] + [(i % order) + 1 for i in range(steps - 1)] + [1]
+        # Last step should be order 1 for better quality.
+        if self.lower_order_final:
+            orders[-1] = 1
+        return orders
+
+    def dpm_solver_first_order_update(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        lambda_t, lambda_s = self.lambda_t[prev_timestep], self.lambda_t[timestep]
+        alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep]
+        sigma_t, sigma_s = self.sigma_t[prev_timestep], self.sigma_t[timestep]
+        h = lambda_t - lambda_s
+        if self.algorithm_type == "dpmsolver++":
+            x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
+        elif self.algorithm_type == "dpmsolver":
+            x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output
+        return x_t
+
+    def singlestep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.FloatTensor],
+        timestep_list: List[int],
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
+        alpha_t, alpha_s1 = self.alpha_t[t], self.alpha_t[s1]
+        sigma_t, sigma_s1 = self.sigma_t[t], self.sigma_t[s1]
+        h, h_0 = lambda_t - lambda_s1, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m1, (1.0 / r0) * (m0 - m1)
+        if self.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s1) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s1) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                )
+        elif self.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s1) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s1) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                )
+        return x_t
+
+    def singlestep_dpm_solver_third_order_update(
+        self,
+        model_output_list: List[torch.FloatTensor],
+        timestep_list: List[int],
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        t, s0, s1, s2 = prev_timestep, timestep_list[-1], timestep_list[-2], timestep_list[-3]
+        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
+        lambda_t, lambda_s0, lambda_s1, lambda_s2 = (
+            self.lambda_t[t],
+            self.lambda_t[s0],
+            self.lambda_t[s1],
+            self.lambda_t[s2],
+        )
+        alpha_t, alpha_s2 = self.alpha_t[t], self.alpha_t[s2]
+        sigma_t, sigma_s2 = self.sigma_t[t], self.sigma_t[s2]
+        h, h_0, h_1 = lambda_t - lambda_s2, lambda_s0 - lambda_s2, lambda_s1 - lambda_s2
+        r0, r1 = h_0 / h, h_1 / h
+        D0 = m2
+        D1_0, D1_1 = (1.0 / r1) * (m1 - m2), (1.0 / r0) * (m0 - m2)
+        D1 = (r0 * D1_0 - r1 * D1_1) / (r0 - r1)
+        D2 = 2.0 * (D1_1 - D1_0) / (r0 - r1)
+        if self.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s2) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1_1
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s2) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                    - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2
+                )
+        elif self.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s2) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1_1
+                )
+            elif self.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s2) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                    - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2
+                )
+        return x_t
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        model_output_type: str,
+        timestep: int,
+        sample: torch.FloatTensor,
+    ) -> SchedulerStepOutput:
+
+        step_index = (self.timesteps == timestep).nonzero().item()
+
+        # Check if this step is the follow-up of the previous step.
+        # If not, then we reset and treat it as a new run.
+        if self.last_step_index and self.last_step_index != step_index - 1:
+            self.reset()
+        self.last_step_index = step_index
+
+        prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1]
+        model_output_convert = self.convert_output(model_output, model_output_type, sample, timestep)
+
+        if self.algorithm_type == "dpmsolver++":
+            model_output = model_output_convert.pred_original_sample
+        else:
+            model_output = model_output_convert.pred_epsilon
+
+        for i in range(self.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+        self.model_outputs[-1] = model_output
+
+        order = self.order_list[step_index]
+
+        #  For img2img denoising might start with order>1 which is not possible
+        #  In this case make sure that the first two steps are both order=1
+        while self.model_outputs[-order] is None:
+            order -= 1
+
+        # For single-step solvers, we use the initial value at each time with order = 1.
+        if order == 1:
+            self.sample = sample
+
+        timestep_list = [self.timesteps[step_index - i] for i in range(order - 1, 0, -1)] + [timestep]
+
+        if order == 1:
+            prev_sample = self.dpm_solver_first_order_update(self.model_outputs[-1], timestep_list[-1], prev_timestep, self.sample)
+        elif order == 2:
+            prev_sample = self.singlestep_dpm_solver_second_order_update(self.model_outputs, timestep_list, prev_timestep, self.sample)
+        elif order == 3:
+            prev_sample = self.singlestep_dpm_solver_third_order_update(self.model_outputs, timestep_list, prev_timestep, self.sample)
+        else:
+            raise NotImplementedError
+
+        return SchedulerStepOutput(
+            prev_sample=prev_sample,
+            pred_original_sample=model_output_convert.pred_original_sample
+        )

RealCustom/schedulers/utils.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import math
+from typing import List
+
+import torch
+
+
+def get_betas(name: str, num_steps: int = 1000, shift_snr: float = 1, terminal_pure_noise: bool = False):
+    # Get betas
+    max_beta = 1 if terminal_pure_noise else 0.999
+    if name == "squared_linear":
+        betas = torch.linspace(0.00085**0.5, 0.012**0.5, num_steps) ** 2
+    elif name == "cosine":
+        betas = get_cosine_betas(num_steps, max_beta=max_beta)
+    elif name == "alphas_cumprod_linear":
+        betas = get_alphas_cumprod_linear_betas(num_steps, max_beta=max_beta)
+    elif name == "sigmoid":
+        betas = get_sigmoid_betas(num_steps, max_beta=max_beta, square=True, slop=0.7)
+    else:
+        raise NotImplementedError
+
+    # Shift snr
+    betas = shift_betas_by_snr_factor(betas, shift_snr)
+
+    # Ensure terminal pure noise
+    # Only non-cosine schedule uses rescale, cosine schedule can directly set max_beta=1 to ensure temrinal pure noise.
+    if name == "squared_linear" and terminal_pure_noise:
+        betas = rescale_betas_to_ensure_terminal_pure_noise(betas)
+
+    return betas
+
+
+def validate_betas(betas: List[float]) -> bool:
+    """
+    Validate betas is monotonic and within 0 to 1 range, i.e. 0 < beta_{t-1} < beta_{t} <= 1
+
+    Args:
+        betas (List[float]): betas
+
+    Returns:
+        bool: True if betas is correct
+    """
+    return all(b1 < b2 for b1, b2 in zip(betas, betas[1:])) and betas[0] > 0 and betas[-1] <= 1
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar_fn, max_beta=0.999):
+    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))
+    if not validate_betas(betas):
+        import logging
+        logging.warning("No feasible betas for given alpha bar")
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+def get_cosine_betas(num_diffusion_timesteps, max_beta=0.999) -> torch.Tensor:
+    def alpha_bar_fn(time_step):
+        return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2
+    return betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar_fn, max_beta)
+
+
+def get_sigmoid_betas(num_diffusion_timesteps, max_beta, square=False, slop=1):
+    def alpha_bar_fn(t):
+        def sigmoid(x):
+            return 1 / (1 + math.exp(-x * slop))
+        s = 6  # (-6, 6) from geodiff
+        vb = sigmoid(-s)
+        ve = sigmoid(s)
+        return ((sigmoid(s - t * 2 * s) - vb) / (ve - vb))**(1 if not square else 2)
+    return betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar_fn, max_beta)
+
+
+def get_alphas_cumprod_linear_betas(num_diffusion_timesteps, max_beta):
+    def alpha_bar_fn(t):
+        return 1 - t
+    return betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar_fn, max_beta=max_beta)
+
+
+def shift_betas_by_snr_factor(betas: torch.Tensor, factor: float) -> torch.Tensor:
+    if factor == 1.0:
+        return betas
+    # Convert betas to snr
+    alphas = 1 - betas
+    alphas_cumprod = alphas.cumprod(dim=0)
+    snr = alphas_cumprod / (1 - alphas_cumprod)
+    # Shift snr
+    snr *= factor
+    # Convert snr to betas
+    alphas_cumprod = snr / (1 + snr)
+    alphas = torch.cat(
+        [alphas_cumprod[0:1], alphas_cumprod[1:] / alphas_cumprod[:-1]])
+    betas = 1 - alphas
+    return betas
+
+
+def rescale_betas_to_ensure_terminal_pure_noise(betas: torch.Tensor) -> torch.Tensor:
+    # Convert betas to alphas_cumprod_sqrt
+    alphas = 1 - betas
+    alphas_cumprod = alphas.cumprod(0)
+    alphas_cumprod_sqrt = alphas_cumprod.sqrt()
+    # Rescale alphas_cumprod_sqrt such that alphas_cumprod_sqrt[0] remains unchanged but alphas_cumprod_sqrt[-1] = 0
+    alphas_cumprod_sqrt = (alphas_cumprod_sqrt - alphas_cumprod_sqrt[-1]) / (
+        alphas_cumprod_sqrt[0] - alphas_cumprod_sqrt[-1]) * alphas_cumprod_sqrt[0]
+    # Convert alphas_cumprod_sqrt to betas
+    alphas_cumprod = alphas_cumprod_sqrt ** 2
+    alphas = torch.cat(
+        [alphas_cumprod[0:1], alphas_cumprod[1:] / alphas_cumprod[:-1]])
+    betas = 1 - alphas
+    return betas