File tree Expand file tree Collapse file tree 1 file changed +8
-1
lines changed
Expand file tree Collapse file tree 1 file changed +8
-1
lines changed Original file line number Diff line number Diff 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" :
You can’t perform that action at this time.
0 commit comments