Add gradient estimation sampler

Author: SkutteOleg

README.md

@@ -328,7 +328,7 @@ arguments:
   --skip-layer-start START           SLG enabling point: (default: 0.01)
   --skip-layer-end END               SLG disabling point: (default: 0.2)
   --scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)
-  --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}
+  --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, gradient_estimation}
                                      sampling method (default: "euler" for Flux/SD3/Wan, "euler_a" otherwise)
   --timestep-shift N                 shift timestep for NitroFusion models, default: 0, recommended N for NitroSD-Realism around 250 and 500 for NitroSD-Vibrant
   --steps  STEPS                     number of sample steps (default: 20)
@@ -342,7 +342,7 @@ arguments:
   --high-noise-skip-layer-start      (high noise) SLG enabling point: (default: 0.01)
   --high-noise-skip-layer-end END    (high noise) SLG disabling point: (default: 0.2)
   --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)
-  --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}
+  --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, gradient_estimation}
                                      (high noise) sampling method (default: "euler_a")
   --high-noise-steps  STEPS          (high noise) number of sample steps (default: -1 = auto)
                                      SLG will be enabled at step int([STEPS]*[START]) and disabled at int([STEPS]*[END])

denoiser.hpp

@@ -1462,7 +1462,50 @@ static void sample_k_diffusion(sample_method_t method,
                 }
             }
         } break;
+        case GRADIENT_ESTIMATION: {
+            // See: https://github.com/crowsonkb/k-diffusion/pull/114
+            struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
+            struct ggml_tensor* d     = ggml_dup_tensor(work_ctx, x);
+            struct ggml_tensor* x_2   = ggml_dup_tensor(work_ctx, x);
+
+            for (int i = 0; i < steps; i++) {
+                float sigma = sigmas[i];
+
+                ggml_tensor* denoised = model(x, sigma, i + 1);
 
+                float* vec_d        = (float*)d->data;
+                float* vec_x        = (float*)x->data;
+                float* vec_x_2      = (float*)x_2->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 sigma_2 = sigmas[i] * 0.5;
+
+                ggml_tensor_set_f32_randn(noise, rng);
+
+                for (int j = 0; j < ggml_nelements(x); j++) {
+                    vec_x_2[j] = vec_x[j] + sigma_2 * ((float*)noise->data)[j];
+                }
+
+                ggml_tensor* denoised_2 = model(x_2, sigma, i + 1);
+                float* vec_denoised_2   = (float*)denoised_2->data;
+
+                // d = (x - denoised) / sigma
+                for (int j = 0; j < ggml_nelements(d); j++) {
+                    float vec_d_2 = (vec_x_2[j] - vec_denoised_2[j]) / sigma;
+                    vec_d[j] += (vec_d_2 - vec_d[j]) * 1.f / (sigma_2 * sigma_2) * ((float*)noise->data)[j];
+                }
+
+                float dt = sigmas[i + 1] - sigma;
+                // x = x + d * dt
+                for (int j = 0; j < ggml_nelements(x); j++) {
+                    vec_x[j] = vec_x[j] + vec_d[j] * dt;
+                }
+            }
+        } break;
         default:
             LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);
             abort();

examples/cli/main.cpp

@@ -247,7 +247,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --skip-layer-start START           SLG enabling point: (default: 0.01)\n");
     printf("  --skip-layer-end END               SLG disabling point: (default: 0.2)\n");
     printf("  --scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)\n");
-    printf("  --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}\n");
+    printf("  --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, gradient_estimation}\n");
     printf("                                     sampling method (default: \"euler\" for Flux/SD3/Wan, \"euler_a\" otherwise)\n");
     printf("  --timestep-shift N                 shift timestep for NitroFusion models, default: 0, recommended N for NitroSD-Realism around 250 and 500 for NitroSD-Vibrant\n");
     printf("  --steps  STEPS                     number of sample steps (default: 20)\n");
@@ -261,7 +261,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --high-noise-skip-layer-start      (high noise) SLG enabling point: (default: 0.01)\n");
     printf("  --high-noise-skip-layer-end END    (high noise) SLG disabling point: (default: 0.2)\n");
     printf("  --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)\n");
-    printf("  --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}\n");
+    printf("  --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, gradient_estimation}\n");
     printf("                                     (high noise) sampling method (default: \"euler_a\")\n");
     printf("  --high-noise-steps  STEPS          (high noise) number of sample steps (default: -1 = auto)\n");
     printf("                                     SLG will be enabled at step int([STEPS]*[START]) and disabled at int([STEPS]*[END])\n");

stable-diffusion.cpp

@@ -1612,6 +1612,7 @@ const char* sample_method_to_str[] = {
     "ddim_trailing",
     "tcd",
     "euler_a",
+    "gradient_estimation",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {

stable-diffusion.h

@@ -48,6 +48,7 @@ enum sample_method_t {
     DDIM_TRAILING,
     TCD,
     EULER_A,
+    GRADIENT_ESTIMATION,
     SAMPLE_METHOD_COUNT
 };