fix: fix schedulers when sigma_min gets too close to zero

Author: daniandtheweb

denoiser.hpp

@@ -347,7 +347,11 @@ struct SmoothStepScheduler : SigmaScheduler {
     }
 };
 
-// KL Optimal: https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608
+/*
+* KL Optimal:
+*     Original work from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608.
+*     Implemented using https://github.com/comfyanonymous/ComfyUI/pull/6206 as a reference.
+*/
 struct KLOptimalScheduler : SigmaScheduler {
     std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
         std::vector<float> sigmas;
@@ -367,39 +371,15 @@ struct KLOptimalScheduler : SigmaScheduler {
         float alpha_min = std::atan(sigma_min);
         float alpha_max = std::atan(sigma_max);
 
-        if (sigma_min <= 1e-6f) {
-            /* sigma_min = 0:
-             * implemented using
-             * https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608
-             * as a reference.
-             */
-
-            LOG_INFO("KL Optimal with sigma_min close to zero");
-            for (uint32_t i = 0; i <= n; ++i) {
-
-                float t = static_cast<float>(i) / static_cast<float>(n);
-                float angle = t * alpha_min + (1.0f - t) * alpha_max;
-                sigmas.push_back(std::tan(angle));
-            }
-
-        } else {
-            /* sigma_min != 0:
-             * implemented using
-             * https://github.com/comfyanonymous/ComfyUI/pull/6206
-             * as a reference.
-             */
-
-            LOG_INFO("KL Optimal with sigma_min greater than zero");
-            for (uint32_t i = 0; i < n; ++i) {
-
-                float t = static_cast<float>(i) / static_cast<float>(n - 1);
-                float angle = t * alpha_min + (1.0f - t) * alpha_max;
-                sigmas.push_back(std::tan(angle));
-            }
+        for (uint32_t i = 0; i < n; ++i) {
 
-            sigmas.push_back(0.0f);
+            float t = static_cast<float>(i) / static_cast<float>(n - 1);
+            float angle = t * alpha_min + (1.0f - t) * alpha_max;
+            sigmas.push_back(std::tan(angle));
         }
 
+        sigmas.push_back(0.0f);
+
         return sigmas;
     }
 };
@@ -481,6 +461,10 @@ struct CompVisDenoiser : public Denoiser {
     }
 
     float sigma_to_t(float sigma) override {
+        if (sigma <= 1e-6f) {
+            return (float)(TIMESTEPS - 1);
+        }
+
         float log_sigma = std::log(sigma);
         std::vector<float> dists;
         dists.reserve(TIMESTEPS);
@@ -756,8 +740,12 @@ static bool sample_k_diffusion(sample_method_t method,
                     float* vec_x        = (float*)x->data;
                     float* vec_denoised = (float*)denoised->data;
 
-                    for (int i = 0; i < ggml_nelements(d); i++) {
-                        vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
+                    if (sigma < 1e-6f) {
+                        ggml_set_f32(d, 0.0f);
+                    } else {
+                        for (int i = 0; i < ggml_nelements(d); i++) {
+                            vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
+                        }
                     }
                 }