fix: avoid issues when sigma_min is close to 0

Author: daniandtheweb

denoiser.hpp

@@ -347,35 +347,37 @@ struct SmoothStepScheduler : SigmaScheduler {
     }
 };
 
-// Implementation adapted from 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;
 
         if (n == 0) {
             return sigmas;
         }
+
         if (n == 1) {
             sigmas.push_back(sigma_max);
             sigmas.push_back(0.0f);
             return sigmas;
         }
 
+        sigmas.reserve(n + 1);
+
         float alpha_min = std::atan(sigma_min);
         float alpha_max = std::atan(sigma_max);
 
         for (uint32_t i = 0; i < n; ++i) {
-            // t goes from 0.0 to 1.0
-            float t = static_cast<float>(i) / static_cast<float>(n - 1);
 
-            // Interpolate in the angle domain
+            float t = static_cast<float>(i) / static_cast<float>(n - 1);
             float angle = t * alpha_min + (1.0f - t) * alpha_max;
-
-            // Convert back to sigma
             sigmas.push_back(std::tan(angle));
         }
 
-        // Append the final zero to sigma
         sigmas.push_back(0.0f);
 
         return sigmas;
@@ -459,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);
@@ -734,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;
+                        }
                     }
                 }