feat: add Euler, Heun and DPM++ (2M) samplers (#50)

* Add Euler sampler

* Add Heun sampler

* Add DPM++ (2M) sampler

* Add modified DPM++ (2M) "v2" sampler.

This was proposed in a issue discussion of the stable diffusion webui,
at https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457
and apparently works around overstepping of the DPM++ (2M) method with
small step counts.

The parameter is called dpmpp2mv2 here.

* match code style

---------

Co-authored-by: Urs Ganse <urs@nerd2nerd.org>
Co-authored-by: leejet <leejet714@gmail.com>

README.md

@@ -20,6 +20,10 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) style tokenizer (not all the features, only token weighting for now)
 - Sampling method
     - `Euler A`
+    - `Euler`
+    - `Heun`
+    - `DPM++ 2M`
+    - [`DPM++ 2M v2`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457)
 - Cross-platform reproducibility (`--rng cuda`, consistent with the `stable-diffusion-webui GPU RNG`)
 - Supported platforms
     - Linux
@@ -125,8 +129,10 @@ arguments:
                                      1.0 corresponds to full destruction of information in init image
   -H, --height H                     image height, in pixel space (default: 512)
   -W, --width W                      image width, in pixel space (default: 512)
-  --sample-method SAMPLE_METHOD      sample method (default: "eular a")
+  --sampling-method {euler, euler_a, heun, dpm++2m, dpm++2mv2}
+                                     sampling method (default: "euler_a")
   --steps  STEPS                     number of sample steps (default: 20)
+  --rng {std_default, cuda}          RNG (default: cuda)
   -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
   -v, --verbose                      print extra info
 ```

examples/main.cpp

@@ -72,6 +72,14 @@ const char* rng_type_to_str[] = {
     "cuda",
 };
 
+// Names of the sampler method, same order as enum SampleMethod in stable-diffusion.h
+const char* sample_method_str[] = {
+    "euler_a",
+    "euler",
+    "heun",
+    "dpm++2m",
+    "dpm++2mv2"};
+
 struct Option {
     int n_threads = -1;
     std::string mode = TXT2IMG;
@@ -83,7 +91,7 @@ struct Option {
     float cfg_scale = 7.0f;
     int w = 512;
     int h = 512;
-    SampleMethod sample_method = EULAR_A;
+    SampleMethod sample_method = EULER_A;
     int sample_steps = 20;
     float strength = 0.75f;
     RNGType rng_type = CUDA_RNG;
@@ -102,7 +110,7 @@ struct Option {
         printf("    cfg_scale:       %.2f\n", cfg_scale);
         printf("    width:           %d\n", w);
         printf("    height:          %d\n", h);
-        printf("    sample_method:   %s\n", "eular a");
+        printf("    sample_method:   %s\n", sample_method_str[sample_method]);
         printf("    sample_steps:    %d\n", sample_steps);
         printf("    strength:        %.2f\n", strength);
         printf("    rng:             %s\n", rng_type_to_str[rng_type]);
@@ -128,7 +136,8 @@ void print_usage(int argc, const char* argv[]) {
     printf("                                     1.0 corresponds to full destruction of information in init image\n");
     printf("  -H, --height H                     image height, in pixel space (default: 512)\n");
     printf("  -W, --width W                      image width, in pixel space (default: 512)\n");
-    printf("  --sample-method SAMPLE_METHOD      sample method (default: \"eular a\")\n");
+    printf("  --sampling-method {euler, euler_a, heun, dpm++2m, dpm++2mv2}\n");
+    printf("                                     sampling method (default: \"euler_a\")\n");
     printf("  --steps  STEPS                     number of sample steps (default: 20)\n");
     printf("  --rng {std_default, cuda}          RNG (default: cuda)\n");
     printf("  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)\n");
@@ -234,6 +243,23 @@ void parse_args(int argc, const char* argv[], Option* opt) {
                 break;
             }
             opt->seed = std::stoll(argv[i]);
+        } else if (arg == "--sampling-method") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            const char* sample_method_selected = argv[i];
+            int sample_method_found = -1;
+            for (int m = 0; m < N_SAMPLE_METHODS; m++) {
+                if (!strcmp(sample_method_selected, sample_method_str[m])) {
+                    sample_method_found = m;
+                }
+            }
+            if (sample_method_found == -1) {
+                invalid_arg = true;
+                break;
+            }
+            opt->sample_method = (SampleMethod)sample_method_found;
         } else if (arg == "-h" || arg == "--help") {
             print_usage(argc, argv);
             exit(0);

stable-diffusion.cpp

@@ -3502,16 +3502,20 @@ class StableDiffusionGGML {
             ggml_graph_print(&diffusion_graph);
 #endif
             int64_t t1 = ggml_time_ms();
+            if (step > 0) {
                 LOG_INFO("step %d sampling completed, taking %.2fs", step, (t1 - t0) * 1.0f / 1000);
                 LOG_DEBUG("diffusion graph use %.2fMB runtime memory: static %.2fMB, dynamic %.2fMB",
                           (ctx_size + ggml_curr_max_dynamic_size()) * 1.0f / 1024 / 1024,
                           ctx_size * 1.0f / 1024 / 1024,
                           ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024);
                 LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size());
+            }
         };
 
         // sample_euler_ancestral
-        {
+        switch (method) {
+            case EULER_A: {
+                LOG_INFO("sampling using Euler A method");
                 ggml_set_dynamic(ctx, false);
                 struct ggml_tensor* noise = ggml_dup_tensor(ctx, x);
                 struct ggml_tensor* d = ggml_dup_tensor(ctx, x);
@@ -3565,6 +3569,188 @@ class StableDiffusionGGML {
                         }
                     }
                 }
+            } break;
+            case EULER:  // Implemented without any sigma churn
+            {
+                LOG_INFO("sampling using Euler method");
+                ggml_set_dynamic(ctx, false);
+                struct ggml_tensor* d = ggml_dup_tensor(ctx, x);
+                ggml_set_dynamic(ctx, params.dynamic);
+
+                for (int i = 0; i < steps; i++) {
+                    float sigma = sigmas[i];
+
+                    // denoise
+                    denoise(x, sigma, i + 1);
+
+                    // d = (x - denoised) / sigma
+                    {
+                        float* vec_d = (float*)d->data;
+                        float* vec_x = (float*)x->data;
+                        float* vec_denoised = (float*)denoised->data;
+
+                        for (int j = 0; j < ggml_nelements(d); j++) {
+                            vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigma;
+                        }
+                    }
+
+                    float dt = sigmas[i + 1] - sigma;
+                    // x = x + d * dt
+                    {
+                        float* vec_d = (float*)d->data;
+                        float* vec_x = (float*)x->data;
+
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = vec_x[j] + vec_d[j] * dt;
+                        }
+                    }
+                }
+            } break;
+            case HEUN: {
+                LOG_INFO("sampling using Heun method");
+                ggml_set_dynamic(ctx, false);
+                struct ggml_tensor* d = ggml_dup_tensor(ctx, x);
+                struct ggml_tensor* x2 = ggml_dup_tensor(ctx, x);
+                ggml_set_dynamic(ctx, params.dynamic);
+
+                for (int i = 0; i < steps; i++) {
+                    // denoise
+                    denoise(x, sigmas[i], -(i + 1));
+
+                    // d = (x - denoised) / sigma
+                    {
+                        float* vec_d = (float*)d->data;
+                        float* vec_x = (float*)x->data;
+                        float* vec_denoised = (float*)denoised->data;
+
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i];
+                        }
+                    }
+
+                    float dt = sigmas[i + 1] - sigmas[i];
+                    if (sigmas[i + 1] == 0) {
+                        // Euler step
+                        // x = x + d * dt
+                        float* vec_d = (float*)d->data;
+                        float* vec_x = (float*)x->data;
+
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = vec_x[j] + vec_d[j] * dt;
+                        }
+                    } else {
+                        // Heun step
+                        float* vec_d = (float*)d->data;
+                        float* vec_d2 = (float*)d->data;
+                        float* vec_x = (float*)x->data;
+                        float* vec_x2 = (float*)x2->data;
+
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x2[j] = vec_x[j] + vec_d[j] * dt;
+                        }
+
+                        denoise(x2, sigmas[i + 1], i + 1);
+                        float* vec_denoised = (float*)denoised->data;
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            float d2 = (vec_x2[j] - vec_denoised[j]) / sigmas[i + 1];
+                            vec_d[j] = (vec_d[j] + d2) / 2;
+                            vec_x[j] = vec_x[j] + vec_d[j] * dt;
+                        }
+                    }
+                }
+            } break;
+            case DPMPP2M:  // DPM++ (2M) from Karras et al (2022)
+            {
+                LOG_INFO("sampling using DPM++ (2M) method");
+                ggml_set_dynamic(ctx, false);
+                struct ggml_tensor* old_denoised = ggml_dup_tensor(ctx, x);
+                ggml_set_dynamic(ctx, params.dynamic);
+
+                auto t_fn = [](float sigma) -> float { return -log(sigma); };
+
+                for (int i = 0; i < steps; i++) {
+                    // denoise
+                    denoise(x, sigmas[i], i + 1);
+
+                    float t = t_fn(sigmas[i]);
+                    float t_next = t_fn(sigmas[i + 1]);
+                    float h = t_next - t;
+                    float a = sigmas[i + 1] / sigmas[i];
+                    float b = exp(-h) - 1.;
+                    float* vec_x = (float*)x->data;
+                    float* vec_denoised = (float*)denoised->data;
+                    float* vec_old_denoised = (float*)old_denoised->data;
+
+                    if (i == 0 || sigmas[i + 1] == 0) {
+                        // Simpler step for the edge cases
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = a * vec_x[j] - b * vec_denoised[j];
+                        }
+                    } else {
+                        float h_last = t - t_fn(sigmas[i - 1]);
+                        float r = h_last / h;
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            float denoised_d = (1. + 1. / (2. * r)) * vec_denoised[j] - (1. / (2. * r)) * vec_old_denoised[j];
+                            vec_x[j] = a * vec_x[j] - b * denoised_d;
+                        }
+                    }
+
+                    // old_denoised = denoised
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_old_denoised[j] = vec_denoised[j];
+                    }
+                }
+            } break;
+            case DPMPP2Mv2:  // Modified DPM++ (2M) from https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457
+            {
+                LOG_INFO("sampling using modified DPM++ (2M) method");
+                ggml_set_dynamic(ctx, false);
+                struct ggml_tensor* old_denoised = ggml_dup_tensor(ctx, x);
+                ggml_set_dynamic(ctx, params.dynamic);
+
+                auto t_fn = [](float sigma) -> float { return -log(sigma); };
+
+                for (int i = 0; i < steps; i++) {
+                    // denoise
+                    denoise(x, sigmas[i], i + 1);
+
+                    float t = t_fn(sigmas[i]);
+                    float t_next = t_fn(sigmas[i + 1]);
+                    float h = t_next - t;
+                    float a = sigmas[i + 1] / sigmas[i];
+                    float* vec_x = (float*)x->data;
+                    float* vec_denoised = (float*)denoised->data;
+                    float* vec_old_denoised = (float*)old_denoised->data;
+
+                    if (i == 0 || sigmas[i + 1] == 0) {
+                        // Simpler step for the edge cases
+                        float b = exp(-h) - 1.;
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = a * vec_x[j] - b * vec_denoised[j];
+                        }
+                    } else {
+                        float h_last = t - t_fn(sigmas[i - 1]);
+                        float h_min = std::min(h_last, h);
+                        float h_max = std::max(h_last, h);
+                        float r = h_max / h_min;
+                        float h_d = (h_max + h_min) / 2.;
+                        float b = exp(-h_d) - 1.;
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            float denoised_d = (1. + 1. / (2. * r)) * vec_denoised[j] - (1. / (2. * r)) * vec_old_denoised[j];
+                            vec_x[j] = a * vec_x[j] - b * denoised_d;
+                        }
+                    }
+
+                    // old_denoised = denoised
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_old_denoised[j] = vec_denoised[j];
+                    }
+                }
+            } break;
+
+            default:
+                LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);
+                abort();
         }
 
         size_t rt_mem_size = ctx_size + ggml_curr_max_dynamic_size();

stable-diffusion.h

@@ -17,7 +17,12 @@ enum RNGType {
 };
 
 enum SampleMethod {
-    EULAR_A,
+    EULER_A,
+    EULER,
+    HEUN,
+    DPMPP2M,
+    DPMPP2Mv2,
+    N_SAMPLE_METHODS
 };
 
 class StableDiffusionGGML;