Skip to content

Commit 94d2b41

Browse files
UniPCMultistep: support use_karras_sigmas=True and use_flow_sigmas=True
1 parent 7e0ee06 commit 94d2b41

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

src/diffusers/schedulers/scheduling_unipc_multistep.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,8 @@ def __init__(
217217
rescale_betas_zero_snr: bool = False,
218218
use_dynamic_shifting: bool = False,
219219
time_shift_type: Literal["exponential"] = "exponential",
220+
sigma_min: Optional[float] = None,
221+
sigma_max: Optional[float] = None,
220222
) -> None:
221223
if self.config.use_beta_sigmas and not is_scipy_available():
222224
raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
@@ -350,7 +352,12 @@ def set_timesteps(
350352
log_sigmas = np.log(sigmas)
351353
sigmas = np.flip(sigmas).copy()
352354
sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
353-
timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
355+
if self.config.use_flow_sigmas:
356+
sigmas = sigmas / (sigmas + 1)
357+
timesteps = (sigmas * self.config.num_train_timesteps).copy()
358+
else:
359+
timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
360+
354361
if self.config.final_sigmas_type == "sigma_min":
355362
sigma_last = sigmas[-1]
356363
elif self.config.final_sigmas_type == "zero":

0 commit comments

Comments
 (0)