from diffusers import EDMEulerScheduler
import torch

scheduler = EDMEulerScheduler(
    sigma_min=0.002,
    sigma_max=80,
    sigma_data=0.5,
    sigma_schedule="karras",
    num_train_timesteps=1000,
    prediction_type="epsilon",
    rho=7.0,
    final_sigmas_type="sigma_min",
)
num_inference_steps = 35
sigmas = torch.arange(num_inference_steps + 1) / num_inference_steps
scheduler.set_timesteps(sigmas=sigmas)
print(scheduler.sigmas)