From 3a215deff23d28c06c8de98423c12628b8ce6326 Mon Sep 17 00:00:00 2001 From: drhead <1313496+drhead@users.noreply.github.com> Date: Sun, 28 Apr 2024 00:15:58 -0400 Subject: [PATCH] vectorize kl-optimal sigma calculation Co-authored-by: mamei16 --- modules/sd_schedulers.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 10ae4e081..99a6f7be2 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -34,9 +34,8 @@ def sgm_uniform(n, sigma_min, sigma_max, inner_model, device): def kl_optimal(n, sigma_min, sigma_max, device): alpha_min = torch.arctan(torch.tensor(sigma_min, device=device)) alpha_max = torch.arctan(torch.tensor(sigma_max, device=device)) - sigmas = torch.empty((n+1,), device=device) - for i in range(n+1): - sigmas[i] = torch.tan((i/n) * alpha_min + (1.0-i/n) * alpha_max) + step_indices = torch.arange(n + 1, device=device) + sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max) return sigmas