File tree Expand file tree Collapse file tree 1 file changed +6
-2
lines changed Expand file tree Collapse file tree 1 file changed +6
-2
lines changed Original file line number Diff line number Diff line change @@ -399,12 +399,16 @@ def set_timesteps(
399399 if self .config .use_karras_sigmas :
400400 sigmas = np .flip (sigmas ).copy ()
401401 sigmas = self ._convert_to_karras (in_sigmas = sigmas , num_inference_steps = num_inference_steps )
402- timesteps = np .array ([self ._sigma_to_t (sigma , log_sigmas ) for sigma in sigmas ]).round ()
402+ timesteps = np .array ([self ._sigma_to_t (sigma , log_sigmas ) for sigma in sigmas ])
403+ if self .config .beta_schedule != "squaredcos_cap_v2" :
404+ timesteps = timesteps .round ()
403405 elif self .config .use_lu_lambdas :
404406 lambdas = np .flip (log_sigmas .copy ())
405407 lambdas = self ._convert_to_lu (in_lambdas = lambdas , num_inference_steps = num_inference_steps )
406408 sigmas = np .exp (lambdas )
407- timesteps = np .array ([self ._sigma_to_t (sigma , log_sigmas ) for sigma in sigmas ]).round ()
409+ timesteps = np .array ([self ._sigma_to_t (sigma , log_sigmas ) for sigma in sigmas ])
410+ if self .config .beta_schedule != "squaredcos_cap_v2" :
411+ timesteps = timesteps .round ()
408412 elif self .config .use_exponential_sigmas :
409413 sigmas = np .flip (sigmas ).copy ()
410414 sigmas = self ._convert_to_exponential (in_sigmas = sigmas , num_inference_steps = num_inference_steps )
You can’t perform that action at this time.
0 commit comments