From 40c5eb26596bf20c5f3c6450e03336e2d1ca4e2e Mon Sep 17 00:00:00 2001 From: hlky Date: Mon, 23 Sep 2024 15:21:50 +0100 Subject: [PATCH 1/3] exponential sigmas --- .../schedulers/scheduling_euler_discrete.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 46e0e6baef81..92ec6adeaf12 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -158,6 +158,8 @@ class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin): use_karras_sigmas (`bool`, *optional*, defaults to `False`): Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. timestep_spacing (`str`, defaults to `"linspace"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. @@ -186,6 +188,7 @@ def __init__( prediction_type: str = "epsilon", interpolation_type: str = "linear", use_karras_sigmas: Optional[bool] = False, + use_exponential_sigmas: Optional[bool] = False, sigma_min: Optional[float] = None, sigma_max: Optional[float] = None, timestep_spacing: str = "linspace", @@ -235,6 +238,7 @@ def __init__( self.is_scale_input_called = False self.use_karras_sigmas = use_karras_sigmas + self.use_exponential_sigmas = use_exponential_sigmas self._step_index = None self._begin_index = None @@ -332,6 +336,8 @@ def set_timesteps( raise ValueError("Can only pass one of `num_inference_steps` or `timesteps` or `sigmas`.") if timesteps is not None and self.config.use_karras_sigmas: raise ValueError("Cannot set `timesteps` with `config.use_karras_sigmas = True`.") + if timesteps is not None and self.config.use_exponential_sigmas: + raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") if ( timesteps is not None and self.config.timestep_type == "continuous" @@ -396,6 +402,10 @@ def set_timesteps( sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) + if self.config.use_exponential_sigmas: + sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) + timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) + if self.config.final_sigmas_type == "sigma_min": sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 elif self.config.final_sigmas_type == "zero": @@ -468,6 +478,28 @@ def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> to sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho return sigmas + # Copied from https://github.com/crowsonkb/k-diffusion/blob/686dbad0f39640ea25c8a8c6a6e56bb40eacefa2/k_diffusion/sampling.py#L26 + def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: + """Constructs an exponential noise schedule.""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps).exp() + return sigmas + def index_for_timestep(self, timestep, schedule_timesteps=None): if schedule_timesteps is None: schedule_timesteps = self.timesteps From e04194ce8765bfc360d1c40aba52e3ffd707f47b Mon Sep 17 00:00:00 2001 From: hlky Date: Mon, 23 Sep 2024 23:52:33 +0100 Subject: [PATCH 2/3] Apply suggestions from code review Co-authored-by: YiYi Xu --- src/diffusers/schedulers/scheduling_euler_discrete.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 92ec6adeaf12..4e11f7da3f74 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -338,6 +338,8 @@ def set_timesteps( raise ValueError("Cannot set `timesteps` with `config.use_karras_sigmas = True`.") if timesteps is not None and self.config.use_exponential_sigmas: raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") + if self.config.use_exponential_sigmas and self.config.use_karras_sigmas: + raise ValueErrror("Cannot set both `config.use_exponential_sigmas = True` and config.use_karras_sigmas = True`") if ( timesteps is not None and self.config.timestep_type == "continuous" @@ -402,7 +404,7 @@ def set_timesteps( sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.use_exponential_sigmas: + elif self.config.use_exponential_sigmas: sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) From db54cc0dd562d00dc39ab94275de0cd68db49111 Mon Sep 17 00:00:00 2001 From: hlky Date: Mon, 23 Sep 2024 23:56:52 +0100 Subject: [PATCH 3/3] make style --- src/diffusers/schedulers/scheduling_euler_discrete.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 4e11f7da3f74..2b74558d3cb7 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -339,7 +339,9 @@ def set_timesteps( if timesteps is not None and self.config.use_exponential_sigmas: raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") if self.config.use_exponential_sigmas and self.config.use_karras_sigmas: - raise ValueErrror("Cannot set both `config.use_exponential_sigmas = True` and config.use_karras_sigmas = True`") + raise ValueError( + "Cannot set both `config.use_exponential_sigmas = True` and config.use_karras_sigmas = True`" + ) if ( timesteps is not None and self.config.timestep_type == "continuous"