passing sigmas beta, exponential and karras

huggingface · Nov 23, 2024 · 6be2301 · 6be2301
1 parent e8c9f9d
commit 6be2301
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py
@@ -428,6 +428,15 @@ def set_timesteps(
                 else:
                     sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
 
+            if self.config.use_karras_sigmas:
+                sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps)
+
+            elif self.config.use_exponential_sigmas:
+                sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+            elif self.config.use_beta_sigmas:
+                sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
             if self.config.invert_sigmas:
                 sigmas = 1.0 - sigmas
             timesteps = sigmas * self.config.num_train_timesteps