feat: added exponential scheduler

daniandtheweb · daniandtheweb · commit 032f6651f1f2 · 2024-08-14T15:55:07.000+02:00
diff --git a/denoiser.hpp b/denoiser.hpp
@@ -40,6 +40,27 @@ struct DiscreteSchedule : SigmaSchedule {
     }
 };
 
+struct ExponentialSchedule : SigmaSchedule {
+    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;
+
+        // Calculate step size
+        float log_sigma_min = std::log(sigma_min);
+        float log_sigma_max = std::log(sigma_max);
+        float step = (log_sigma_max - log_sigma_min) / (n - 1);
+
+        // Fill sigmas with exponential values
+        for (uint32_t i = 0; i < n; ++i) {
+            float sigma = std::exp(log_sigma_max - step * i);
+            sigmas.push_back(sigma);
+        }
+
+        sigmas.push_back(0);
+
+        return sigmas;
+    }
+};
+
 /*
 https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
 */
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -45,6 +45,7 @@ const char* schedule_str[] = {
     "default",
     "discrete",
     "karras",
+    "exponential",
     "ays",
 };
 
@@ -193,7 +194,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --rng {std_default, cuda}          RNG (default: cuda)\n");
     printf("  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)\n");
     printf("  -b, --batch-count COUNT            number of images to generate.\n");
-    printf("  --schedule {discrete, karras, ays} Denoiser sigma schedule (default: discrete)\n");
+    printf("  --schedule {discrete, karras, exponential, ays} Denoiser sigma schedule (default: discrete)\n");
     printf("  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n");
     printf("                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n");
     printf("  --vae-tiling                       process vae in tiles to reduce memory usage\n");
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -4,6 +4,7 @@
 #include "rng.hpp"
 #include "rng_philox.hpp"
 #include "stable-diffusion.h"
+#include <memory>
 #include "util.h"
 
 #include "conditioner.hpp"
@@ -456,6 +457,10 @@ class StableDiffusionGGML {
                     LOG_INFO("running with Karras schedule");
                     denoiser->schedule = std::make_shared<KarrasSchedule>();
                     break;
+                case EXPONENTIAL:
+                    LOG_INFO("running exponential schedule");
+                    denoiser->schedule = std::make_shared<ExponentialSchedule>();
+                    break;
                 case AYS:
                     LOG_INFO("Running with Align-Your-Steps schedule");
                     denoiser->schedule          = std::make_shared<AYSSchedule>();
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -49,6 +49,7 @@ enum schedule_t {
     DEFAULT,
     DISCRETE,
     KARRAS,
+    EXPONENTIAL,
     AYS,
     N_SCHEDULES
 };
