feat: add option to switch the sigma schedule (#51)

Concretely, this allows switching to the "Karras" schedule from the
Karras et al 2022 paper, equivalent to the samplers marked as "Karras"
in the AUTOMATIC1111 WebUI. This choice is in principle orthogonal to
the sampler choice and can be given independently.

README.md

@@ -20,6 +20,11 @@ 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
     - Mac OS
@@ -35,8 +40,6 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
 - [ ] LoRA support
 - [ ] k-quants support
-- [ ] Cross-platform reproducibility (perhaps ensuring consistency with the original SD)
-- [ ] Adapting to more weight formats
 
 ## Usage
 
@@ -126,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

@@ -67,6 +67,25 @@ int32_t get_num_physical_cores() {
     return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
 }
 
+const char* rng_type_to_str[] = {
+    "std_default",
+    "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"};
+
+// Names of the sigma schedule overrides, same order as Schedule in stable-diffusion.h
+const char* schedule_str[] = {
+    "default",
+    "discrete",
+    "karras"};
+
 struct Option {
     int n_threads = -1;
     std::string mode = TXT2IMG;
@@ -78,10 +97,12 @@ struct Option {
     float cfg_scale = 7.0f;
     int w = 512;
     int h = 512;
-    SampleMethod sample_method = EULAR_A;
+    SampleMethod sample_method = EULER_A;
+    Schedule schedule = DEFAULT;
     int sample_steps = 20;
     float strength = 0.75f;
-    int seed = 42;
+    RNGType rng_type = CUDA_RNG;
+    int64_t seed = 42;
     bool verbose = false;
 
     void print() {
@@ -96,10 +117,12 @@ 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("    schedule:        %s\n", schedule_str[schedule]);
         printf("    sample_steps:    %d\n", sample_steps);
         printf("    strength:        %.2f\n", strength);
-        printf("    seed:            %d\n", seed);
+        printf("    rng:             %s\n", rng_type_to_str[rng_type]);
+        printf("    seed:            %ld\n", seed);
     }
 };
 
@@ -121,9 +144,12 @@ 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");
+    printf("  --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -206,12 +232,60 @@ void parse_args(int argc, const char* argv[], Option* opt) {
                 break;
             }
             opt->sample_steps = std::stoi(argv[i]);
+        } else if (arg == "--rng") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            std::string rng_type_str = argv[i];
+            if (rng_type_str == "std_default") {
+                opt->rng_type = STD_DEFAULT_RNG;
+            } else if (rng_type_str == "cuda") {
+                opt->rng_type = CUDA_RNG;
+            } else {
+                invalid_arg = true;
+                break;
+            }
+        } else if (arg == "--schedule") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            const char* schedule_selected = argv[i];
+            int schedule_found = -1;
+            for (int d = 0; d < N_SCHEDULES; d++) {
+                if (!strcmp(schedule_selected, schedule_str[d])) {
+                    schedule_found = d;
+                }
+            }
+            if (schedule_found == -1) {
+                invalid_arg = true;
+                break;
+            }
+            opt->schedule = (Schedule)schedule_found;
         } else if (arg == "-s" || arg == "--seed") {
             if (++i >= argc) {
                 invalid_arg = true;
                 break;
             }
-            opt->seed = std::stoi(argv[i]);
+            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);
@@ -263,13 +337,13 @@ void parse_args(int argc, const char* argv[], Option* opt) {
         exit(1);
     }
 
-    if (opt->w <= 0 || opt->w % 32 != 0) {
-        fprintf(stderr, "error: the width must be a multiple of 32\n");
+    if (opt->w <= 0 || opt->w % 64 != 0) {
+        fprintf(stderr, "error: the width must be a multiple of 64\n");
         exit(1);
     }
 
-    if (opt->h <= 0 || opt->h % 32 != 0) {
-        fprintf(stderr, "error: the height must be a multiple of 32\n");
+    if (opt->h <= 0 || opt->h % 64 != 0) {
+        fprintf(stderr, "error: the height must be a multiple of 64\n");
         exit(1);
     }
 
@@ -315,21 +389,21 @@ int main(int argc, const char* argv[]) {
             free(img_data);
             return 1;
         }
-        if (opt.w <= 0 || opt.w % 32 != 0) {
-            fprintf(stderr, "error: the width of image must be a multiple of 32\n");
+        if (opt.w <= 0 || opt.w % 64 != 0) {
+            fprintf(stderr, "error: the width of image must be a multiple of 64\n");
             free(img_data);
             return 1;
         }
-        if (opt.h <= 0 || opt.h % 32 != 0) {
-            fprintf(stderr, "error: the height of image must be a multiple of 32\n");
+        if (opt.h <= 0 || opt.h % 64 != 0) {
+            fprintf(stderr, "error: the height of image must be a multiple of 64\n");
             free(img_data);
             return 1;
         }
         init_img.assign(img_data, img_data + (opt.w * opt.h * c));
     }
 
-    StableDiffusion sd(opt.n_threads, vae_decode_only, true);
-    if (!sd.load_from_file(opt.model_path)) {
+    StableDiffusion sd(opt.n_threads, vae_decode_only, true, opt.rng_type);
+    if (!sd.load_from_file(opt.model_path, opt.schedule)) {
         return 1;
     }
 
@@ -365,4 +439,4 @@ int main(int argc, const char* argv[]) {
     printf("save result image to '%s'\n", opt.output_path.c_str());
 
     return 0;
-}
+}

rng.h

@@ -0,0 +1,35 @@
+#ifndef __RNG_H__
+#define __RNG_H__
+
+#include <random>
+#include <vector>
+
+class RNG {
+   public:
+    virtual void manual_seed(uint64_t seed) = 0;
+    virtual std::vector<float> randn(uint32_t n) = 0;
+};
+
+class STDDefaultRNG : public RNG {
+   private:
+    std::default_random_engine generator;
+
+   public:
+    void manual_seed(uint64_t seed) {
+        generator.seed(seed);
+    }
+
+    std::vector<float> randn(uint32_t n) {
+        std::vector<float> result;
+        float mean = 0.0;
+        float stddev = 1.0;
+        std::normal_distribution<float> distribution(mean, stddev);
+        for (int i = 0; i < n; i++) {
+            float random_number = distribution(generator);
+            result.push_back(random_number);
+        }
+        return result;
+    }
+};
+
+#endif  // __RNG_H__

rng_philox.h

@@ -0,0 +1,125 @@
+#ifndef __RNG_PHILOX_H__
+#define __RNG_PHILOX_H__
+
+#include <cmath>
+#include <vector>
+
+#include "rng.h"
+
+// RNG imitiating torch cuda randn on CPU.
+// Port from: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/modules/rng_philox.py
+class PhiloxRNG : public RNG {
+   private:
+    uint64_t seed;
+    uint32_t offset;
+
+   private:
+    std::vector<uint32_t> philox_m = {0xD2511F53, 0xCD9E8D57};
+    std::vector<uint32_t> philox_w = {0x9E3779B9, 0xBB67AE85};
+    float two_pow32_inv = 2.3283064e-10;
+    float two_pow32_inv_2pi = 2.3283064e-10 * 6.2831855;
+
+    std::vector<uint32_t> uint32(uint64_t x) {
+        std::vector<uint32_t> result(2);
+        result[0] = static_cast<uint32_t>(x & 0xFFFFFFFF);
+        result[1] = static_cast<uint32_t>(x >> 32);
+        return result;
+    }
+
+    std::vector<std::vector<uint32_t>> uint32(const std::vector<uint64_t>& x) {
+        int N = x.size();
+        std::vector<std::vector<uint32_t>> result(2, std::vector<uint32_t>(N));
+
+        for (int i = 0; i < N; ++i) {
+            result[0][i] = static_cast<uint32_t>(x[i] & 0xFFFFFFFF);
+            result[1][i] = static_cast<uint32_t>(x[i] >> 32);
+        }
+
+        return result;
+    }
+
+    //  A single round of the Philox 4x32 random number generator.
+    void philox4_round(std::vector<std::vector<uint32_t>>& counter,
+                       const std::vector<std::vector<uint32_t>>& key) {
+        uint32_t N = counter[0].size();
+        for (uint32_t i = 0; i < N; i++) {
+            std::vector<uint32_t> v1 = uint32(static_cast<uint64_t>(counter[0][i]) * static_cast<uint64_t>(philox_m[0]));
+            std::vector<uint32_t> v2 = uint32(static_cast<uint64_t>(counter[2][i]) * static_cast<uint64_t>(philox_m[1]));
+
+            counter[0][i] = v2[1] ^ counter[1][i] ^ key[0][i];
+            counter[1][i] = v2[0];
+            counter[2][i] = v1[1] ^ counter[3][i] ^ key[1][i];
+            counter[3][i] = v1[0];
+        }
+    }
+
+    // Generates 32-bit random numbers using the Philox 4x32 random number generator.
+    // Parameters:
+    //     counter : A 4xN array of 32-bit integers representing the counter values (offset into generation).
+    //     key : A 2xN array of 32-bit integers representing the key values (seed).
+    //     rounds : The number of rounds to perform.
+    // Returns:
+    //     std::vector<std::vector<uint32_t>>: A 4xN array of 32-bit integers containing the generated random numbers.
+    std::vector<std::vector<uint32_t>> philox4_32(std::vector<std::vector<uint32_t>>& counter,
+                                                  std::vector<std::vector<uint32_t>>& key,
+                                                  int rounds = 10) {
+        uint32_t N = counter[0].size();
+        for (int i = 0; i < rounds - 1; ++i) {
+            philox4_round(counter, key);
+
+            for (uint32_t j = 0; j < N; ++j) {
+                key[0][j] += philox_w[0];
+                key[1][j] += philox_w[1];
+            }
+        }
+
+        philox4_round(counter, key);
+        return counter;
+    }
+
+    float box_muller(float x, float y) {
+        float u = x * two_pow32_inv + two_pow32_inv / 2;
+        float v = y * two_pow32_inv_2pi + two_pow32_inv_2pi / 2;
+
+        float s = sqrt(-2.0 * log(u));
+
+        float r1 = s * sin(v);
+        return r1;
+    }
+
+   public:
+    PhiloxRNG(uint64_t seed = 0) {
+        this->seed = seed;
+        this->offset = 0;
+    }
+
+    void manual_seed(uint64_t seed) {
+        this->seed = seed;
+        this->offset = 0;
+    }
+
+    std::vector<float> randn(uint32_t n) {
+        std::vector<std::vector<uint32_t>> counter(4, std::vector<uint32_t>(n, 0));
+        for (uint32_t i = 0; i < n; i++) {
+            counter[0][i] = this->offset;
+        }
+
+        for (uint32_t i = 0; i < n; i++) {
+            counter[2][i] = i;
+        }
+        this->offset += 1;
+
+        std::vector<uint64_t> key(n, this->seed);
+        std::vector<std::vector<uint32_t>> key_uint32 = uint32(key);
+
+        std::vector<std::vector<uint32_t>> g = philox4_32(counter, key_uint32);
+
+        std::vector<float> result;
+        for (int i = 0; i < n; ++i) {
+            result.push_back(box_muller(g[0][i], g[1][i]));
+        }
+        return result;
+    }
+};
+
+#endif  // __RNG_PHILOX_H__

stable-diffusion.cpp

@@ -14,6 +14,8 @@
 #include <vector>
 
 #include "ggml/ggml.h"
+#include "rng.h"
+#include "rng_philox.h"
 #include "stable-diffusion.h"
 
 static SDLogLevel log_level = SDLogLevel::INFO;
@@ -117,19 +119,11 @@ ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_pa
     return tensor;
 }
 
-static std::default_random_engine generator;
-
-void set_random_seed(int seed) {
-    generator.seed(seed);
-}
-
-void ggml_tensor_set_f32_randn(struct ggml_tensor* tensor) {
-    float mean = 0.0;
-    float stddev = 1.0;
-    std::normal_distribution<float> distribution(mean, stddev);
-    for (int i = 0; i < ggml_nelements(tensor); i++) {
-        float random_number = distribution(generator);
-        ggml_set_f32_1d(tensor, i, random_number);
+void ggml_tensor_set_f32_randn(struct ggml_tensor* tensor, std::shared_ptr<RNG> rng) {
+    uint32_t n = ggml_nelements(tensor);
+    std::vector<float> random_numbers = rng->randn(n);
+    for (int i = 0; i < n; i++) {
+        ggml_set_f32_1d(tensor, i, random_numbers[i]);
     }
 }
 
@@ -2660,23 +2654,12 @@ struct AutoEncoderKL {
 
 // Ref: https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/external.py
 
-struct DiscreteSchedule {
+struct SigmaSchedule {
     float alphas_cumprod[TIMESTEPS];
     float sigmas[TIMESTEPS];
     float log_sigmas[TIMESTEPS];
 
-    std::vector<float> get_sigmas(int n) {
-        std::vector<float> result;
-
-        int t_max = TIMESTEPS - 1;
-        float step = static_cast<float>(t_max) / static_cast<float>(n - 1);
-        for (int i = 0; i < n; ++i) {
-            float t = t_max - step * i;
-            result.push_back(t_to_sigma(t));
-        }
-        result.push_back(0);
-        return result;
-    }
+    virtual std::vector<float> get_sigmas(uint32_t n) = 0;
 
     float sigma_to_t(float sigma) {
         float log_sigma = std::log(sigma);
@@ -2711,11 +2694,59 @@ struct DiscreteSchedule {
         float log_sigma = (1.0f - w) * log_sigmas[low_idx] + w * log_sigmas[high_idx];
         return std::exp(log_sigma);
     }
+};
 
+struct DiscreteSchedule : SigmaSchedule {
+    std::vector<float> get_sigmas(uint32_t n) {
+        std::vector<float> result;
+
+        int t_max = TIMESTEPS - 1;
+
+        if (n == 0) {
+            return result;
+        } else if (n == 1) {
+            result.push_back(t_to_sigma(t_max));
+            result.push_back(0);
+            return result;
+        }
+
+        float step = static_cast<float>(t_max) / static_cast<float>(n - 1);
+        for (int i = 0; i < n; ++i) {
+            float t = t_max - step * i;
+            result.push_back(t_to_sigma(t));
+        }
+        result.push_back(0);
+        return result;
+    }
+};
+
+struct KarrasSchedule : SigmaSchedule {
+    std::vector<float> get_sigmas(uint32_t n) {
+        // These *COULD* be function arguments here,
+        // but does anybody ever bother to touch them?
+        float sigma_min = 0.1;
+        float sigma_max = 10.;
+        float rho = 7.;
+
+        std::vector<float> result(n + 1);
+
+        float min_inv_rho = pow(sigma_min, (1. / rho));
+        float max_inv_rho = pow(sigma_max, (1. / rho));
+        for (int i = 0; i < n; i++) {
+            // Eq. (5) from Karras et al 2022
+            result[i] = pow(max_inv_rho + (float)i / ((float)n - 1.) * (min_inv_rho - max_inv_rho), rho);
+        }
+        result[n] = 0.;
+        return result;
+    }
+};
+
+struct Denoiser {
+    std::shared_ptr<SigmaSchedule> schedule = std::make_shared<DiscreteSchedule>();
     virtual std::vector<float> get_scalings(float sigma) = 0;
 };
 
-struct CompVisDenoiser : public DiscreteSchedule {
+struct CompVisDenoiser : public Denoiser {
     float sigma_data = 1.0f;
 
     std::vector<float> get_scalings(float sigma) {
@@ -2725,7 +2756,7 @@ struct CompVisDenoiser : public DiscreteSchedule {
     }
 };
 
-struct CompVisVDenoiser : public DiscreteSchedule {
+struct CompVisVDenoiser : public Denoiser {
     float sigma_data = 1.0f;
 
     std::vector<float> get_scalings(float sigma) {
@@ -2747,6 +2778,8 @@ class StableDiffusionGGML {
     bool dynamic = true;
     bool vae_decode_only = false;
     bool free_params_immediately = false;
+
+    std::shared_ptr<RNG> rng = std::make_shared<STDDefaultRNG>();
     int32_t ftype = 1;
     int n_threads = -1;
     float scale_factor = 0.18215f;
@@ -2759,17 +2792,23 @@ class StableDiffusionGGML {
     UNetModel diffusion_model;
     AutoEncoderKL first_stage_model;
 
-    std::shared_ptr<DiscreteSchedule> denoiser = std::make_shared<CompVisDenoiser>();
+    std::shared_ptr<Denoiser> denoiser = std::make_shared<CompVisDenoiser>();
 
     StableDiffusionGGML() = default;
 
     StableDiffusionGGML(int n_threads,
                         bool vae_decode_only,
-                        bool free_params_immediately)
+                        bool free_params_immediately,
+                        RNGType rng_type)
         : n_threads(n_threads),
           vae_decode_only(vae_decode_only),
           free_params_immediately(free_params_immediately) {
         first_stage_model.decode_only = vae_decode_only;
+        if (rng_type == STD_DEFAULT_RNG) {
+            rng = std::make_shared<STDDefaultRNG>();
+        } else if (rng_type == CUDA_RNG) {
+            rng = std::make_shared<PhiloxRNG>();
+        }
     }
 
     ~StableDiffusionGGML() {
@@ -2787,7 +2826,7 @@ class StableDiffusionGGML {
         }
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, Schedule schedule) {
         LOG_INFO("loading model from '%s'", file_path.c_str());
 
         std::ifstream file(file_path, std::ios::binary);
@@ -3027,6 +3066,9 @@ class StableDiffusionGGML {
             }
             bool some_tensor_not_init = false;
             for (auto pair : tensors) {
+                if (pair.first.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos) {
+                    continue;
+                }
                 if (tensor_names_in_file.find(pair.first) == tensor_names_in_file.end()) {
                     LOG_ERROR("tensor '%s' not in model file", pair.first.c_str());
                     some_tensor_not_init = true;
@@ -3079,10 +3121,29 @@ class StableDiffusionGGML {
             LOG_INFO("running in eps-prediction mode");
         }
 
+        if (schedule != DEFAULT) {
+            switch (schedule) {
+                case DISCRETE:
+                    LOG_INFO("running with discrete schedule");
+                    denoiser->schedule = std::make_shared<DiscreteSchedule>();
+                    break;
+                case KARRAS:
+                    LOG_INFO("running with Karras schedule");
+                    denoiser->schedule = std::make_shared<KarrasSchedule>();
+                    break;
+                case DEFAULT:
+                    // Don't touch anything.
+                    break;
+                default:
+                    LOG_ERROR("Unknown schedule %i", schedule);
+                    abort();
+            }
+        }
+
         for (int i = 0; i < TIMESTEPS; i++) {
-            denoiser->alphas_cumprod[i] = alphas_cumprod[i];
-            denoiser->sigmas[i] = std::sqrt((1 - denoiser->alphas_cumprod[i]) / denoiser->alphas_cumprod[i]);
-            denoiser->log_sigmas[i] = std::log(denoiser->sigmas[i]);
+            denoiser->schedule->alphas_cumprod[i] = alphas_cumprod[i];
+            denoiser->schedule->sigmas[i] = std::sqrt((1 - denoiser->schedule->alphas_cumprod[i]) / denoiser->schedule->alphas_cumprod[i]);
+            denoiser->schedule->log_sigmas[i] = std::log(denoiser->schedule->sigmas[i]);
         }
 
         return true;
@@ -3094,7 +3155,7 @@ class StableDiffusionGGML {
         struct ggml_tensor* c = ggml_new_tensor_4d(res_ctx, GGML_TYPE_F32, 1024, 2, 1, 1);
         ggml_set_f32(c, 0.5);
 
-        size_t ctx_size = 1 * 1024 * 1024;  // 1MB
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
             struct ggml_init_params params;
@@ -3117,8 +3178,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* out = diffusion_model.forward(ctx, x_t, NULL, c, t_emb);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-            struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+            struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+            struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("diffusion context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3150,8 +3211,8 @@ class StableDiffusionGGML {
         struct ggml_tensor* out = diffusion_model.forward(ctx, x_t, NULL, c, t_emb);
         ggml_hold_dynamic_tensor(out);
 
-        struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-        struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+        struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+        struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
         ggml_set_dynamic(ctx, false);
         struct ggml_tensor* buf = ggml_new_tensor_1d(ctx, GGML_TYPE_I8, cplan.work_size);
@@ -3160,7 +3221,7 @@ class StableDiffusionGGML {
         cplan.work_data = (uint8_t*)buf->data;
 
         int64_t t0 = ggml_time_ms();
-        ggml_graph_compute(&diffusion_graph, &cplan);
+        ggml_graph_compute(diffusion_graph, &cplan);
 
         double result = 0.f;
 
@@ -3196,7 +3257,7 @@ class StableDiffusionGGML {
                                                             true);
         std::vector<int>& tokens = tokens_and_weights.first;
         std::vector<float>& weights = tokens_and_weights.second;
-        size_t ctx_size = 1 * 1024 * 1024;  // 1MB
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
             struct ggml_init_params params;
@@ -3246,14 +3307,14 @@ class StableDiffusionGGML {
         ggml_set_dynamic(ctx, params.dynamic);
 
         struct ggml_tensor* hidden_states = cond_stage_model.text_model.forward(ctx, input_ids);
-        struct ggml_cgraph cond_graph = ggml_build_forward(hidden_states);
+        struct ggml_cgraph* cond_graph = ggml_build_forward_ctx(ctx, hidden_states);
         LOG_DEBUG("building condition graph completed: %d nodes, %d leafs",
-                  cond_graph.n_nodes, cond_graph.n_leafs);
+                  cond_graph->n_nodes, cond_graph->n_leafs);
 
         memcpy(input_ids->data, tokens.data(), tokens.size() * ggml_element_size(input_ids));
 
         int64_t t0 = ggml_time_ms();
-        ggml_graph_compute_with_ctx(ctx, &cond_graph, n_threads);
+        ggml_graph_compute_with_ctx(ctx, cond_graph, n_threads);
         int64_t t1 = ggml_time_ms();
         LOG_DEBUG("computing condition graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
 
@@ -3330,7 +3391,7 @@ class StableDiffusionGGML {
         struct ggml_tensor* x = ggml_dup_tensor(res_ctx, x_t);
         copy_ggml_tensor(x, x_t);
 
-        size_t ctx_size = 1 * 1024 * 1024;  // 1MB
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
             struct ggml_init_params params;
@@ -3355,8 +3416,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* out = diffusion_model.forward(ctx, noised_input, NULL, context, t_emb);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-            struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+            struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+            struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("diffusion context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3388,8 +3449,8 @@ class StableDiffusionGGML {
         struct ggml_tensor* out = diffusion_model.forward(ctx, noised_input, NULL, context, t_emb);
         ggml_hold_dynamic_tensor(out);
 
-        struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-        struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+        struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+        struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
         ggml_set_dynamic(ctx, false);
         struct ggml_tensor* buf = ggml_new_tensor_1d(ctx, GGML_TYPE_I8, cplan.work_size);
@@ -3431,7 +3492,7 @@ class StableDiffusionGGML {
                 c_in = scaling[1];
             }
 
-            float t = denoiser->sigma_to_t(sigma);
+            float t = denoiser->schedule->sigma_to_t(sigma);
             ggml_set_f32(timesteps, t);
             set_timestep_embedding(timesteps, t_emb, diffusion_model.model_channels);
 
@@ -3447,12 +3508,12 @@ class StableDiffusionGGML {
             if (cfg_scale != 1.0 && uc != NULL) {
                 // uncond
                 copy_ggml_tensor(context, uc);
-                ggml_graph_compute(&diffusion_graph, &cplan);
+                ggml_graph_compute(diffusion_graph, &cplan);
                 copy_ggml_tensor(out_uncond, out);
 
                 // cond
                 copy_ggml_tensor(context, c);
-                ggml_graph_compute(&diffusion_graph, &cplan);
+                ggml_graph_compute(diffusion_graph, &cplan);
 
                 out_cond = out;
 
@@ -3469,7 +3530,7 @@ class StableDiffusionGGML {
             } else {
                 // cond
                 copy_ggml_tensor(context, c);
-                ggml_graph_compute(&diffusion_graph, &cplan);
+                ggml_graph_compute(diffusion_graph, &cplan);
             }
 
             // v = out, eps = out
@@ -3488,69 +3549,255 @@ class StableDiffusionGGML {
             ggml_graph_print(&diffusion_graph);
 #endif
             int64_t t1 = ggml_time_ms();
-            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());
+            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
-        {
-            ggml_set_dynamic(ctx, false);
-            struct ggml_tensor* noise = ggml_dup_tensor(ctx, x);
-            struct ggml_tensor* d = ggml_dup_tensor(ctx, x);
-            ggml_set_dynamic(ctx, params.dynamic);
+        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);
+                ggml_set_dynamic(ctx, params.dynamic);
 
-            for (int i = 0; i < steps; i++) {
-                float sigma = sigmas[i];
+                for (int i = 0; i < steps; i++) {
+                    float sigma = sigmas[i];
 
-                // denoise
-                denoise(x, sigma, i + 1);
+                    // 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 i = 0; i < ggml_nelements(d); i++) {
-                        vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
-                    }
-                }
-
-                // get_ancestral_step
-                float sigma_up = std::min(sigmas[i + 1],
-                                          std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i])));
-                float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up);
-
-                // Euler method
-                float dt = sigma_down - sigmas[i];
-                // x = x + d * dt
-                {
-                    float* vec_d = (float*)d->data;
-                    float* vec_x = (float*)x->data;
-
-                    for (int i = 0; i < ggml_nelements(x); i++) {
-                        vec_x[i] = vec_x[i] + vec_d[i] * dt;
-                    }
-                }
-
-                if (sigmas[i + 1] > 0) {
-                    // x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-                    ggml_tensor_set_f32_randn(noise);
-                    // noise = load_tensor_from_file(res_ctx, "./rand" + std::to_string(i+1) + ".bin");
+                    // d = (x - denoised) / sigma
                     {
+                        float* vec_d = (float*)d->data;
+                        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;
+                        }
+                    }
+
+                    // get_ancestral_step
+                    float sigma_up = std::min(sigmas[i + 1],
+                                              std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i])));
+                    float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up);
+
+                    // Euler method
+                    float dt = sigma_down - sigmas[i];
+                    // x = x + d * dt
+                    {
+                        float* vec_d = (float*)d->data;
                         float* vec_x = (float*)x->data;
-                        float* vec_noise = (float*)noise->data;
 
                         for (int i = 0; i < ggml_nelements(x); i++) {
-                            vec_x[i] = vec_x[i] + vec_noise[i] * sigma_up;
+                            vec_x[i] = vec_x[i] + vec_d[i] * dt;
+                        }
+                    }
+
+                    if (sigmas[i + 1] > 0) {
+                        // x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+                        ggml_tensor_set_f32_randn(noise, rng);
+                        // noise = load_tensor_from_file(res_ctx, "./rand" + std::to_string(i+1) + ".bin");
+                        {
+                            float* vec_x = (float*)x->data;
+                            float* vec_noise = (float*)noise->data;
+
+                            for (int i = 0; i < ggml_nelements(x); i++) {
+                                vec_x[i] = vec_x[i] + vec_noise[i] * sigma_up;
+                            }
                         }
                     }
                 }
-            }
+            } 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();
@@ -3585,7 +3832,7 @@ class StableDiffusionGGML {
         struct ggml_tensor* result = NULL;
 
         // calculate the amount of memory required
-        size_t ctx_size = 1 * 1024 * 1024;
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         {
             struct ggml_init_params params;
             params.mem_size = ctx_size;
@@ -3602,8 +3849,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* moments = first_stage_model.encode(ctx, x);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph vae_graph = ggml_build_forward(moments);
-            struct ggml_cplan cplan = ggml_graph_plan(&vae_graph, n_threads);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments);
+            struct ggml_cplan cplan = ggml_graph_plan(vae_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("vae context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3627,10 +3874,10 @@ class StableDiffusionGGML {
             }
 
             struct ggml_tensor* moments = first_stage_model.encode(ctx, x);
-            struct ggml_cgraph vae_graph = ggml_build_forward(moments);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments);
 
             int64_t t0 = ggml_time_ms();
-            ggml_graph_compute_with_ctx(ctx, &vae_graph, n_threads);
+            ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads);
             int64_t t1 = ggml_time_ms();
 
 #ifdef GGML_PERF
@@ -3674,7 +3921,7 @@ class StableDiffusionGGML {
         ggml_tensor* latent = ggml_new_tensor_4d(res_ctx, moments->type, moments->ne[0],
                                                  moments->ne[1], moments->ne[2] / 2, moments->ne[3]);
         struct ggml_tensor* noise = ggml_dup_tensor(res_ctx, latent);
-        ggml_tensor_set_f32_randn(noise);
+        ggml_tensor_set_f32_randn(noise, rng);
         // noise = load_tensor_from_file(res_ctx, "noise.bin");
         {
             float mean = 0;
@@ -3714,7 +3961,7 @@ class StableDiffusionGGML {
         }
 
         // calculate the amount of memory required
-        size_t ctx_size = 1 * 1024 * 1024;
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         {
             struct ggml_init_params params;
             params.mem_size = ctx_size;
@@ -3731,8 +3978,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* img = first_stage_model.decoder.forward(ctx, z);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph vae_graph = ggml_build_forward(img);
-            struct ggml_cplan cplan = ggml_graph_plan(&vae_graph, n_threads);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img);
+            struct ggml_cplan cplan = ggml_graph_plan(vae_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("vae context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3756,10 +4003,10 @@ class StableDiffusionGGML {
             }
 
             struct ggml_tensor* img = first_stage_model.decode(ctx, z);
-            struct ggml_cgraph vae_graph = ggml_build_forward(img);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img);
 
             int64_t t0 = ggml_time_ms();
-            ggml_graph_compute_with_ctx(ctx, &vae_graph, n_threads);
+            ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads);
             int64_t t1 = ggml_time_ms();
 
 #ifdef GGML_PERF
@@ -3802,14 +4049,16 @@ class StableDiffusionGGML {
 
 StableDiffusion::StableDiffusion(int n_threads,
                                  bool vae_decode_only,
-                                 bool free_params_immediately) {
+                                 bool free_params_immediately,
+                                 RNGType rng_type) {
     sd = std::make_shared<StableDiffusionGGML>(n_threads,
                                                vae_decode_only,
-                                               free_params_immediately);
+                                               free_params_immediately,
+                                               rng_type);
 }
 
-bool StableDiffusion::load_from_file(const std::string& file_path) {
-    return sd->load_from_file(file_path);
+bool StableDiffusion::load_from_file(const std::string& file_path, Schedule s) {
+    return sd->load_from_file(file_path, s);
 }
 
 std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
@@ -3819,10 +4068,11 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
                                               int height,
                                               SampleMethod sample_method,
                                               int sample_steps,
-                                              int seed) {
+                                              int64_t seed) {
     std::vector<uint8_t> result;
     struct ggml_init_params params;
     params.mem_size = static_cast<size_t>(10 * 1024) * 1024;  // 10M
+    params.mem_size += width * height * 3 * sizeof(float) * 2;
     params.mem_buffer = NULL;
     params.no_alloc = false;
     params.dynamic = false;
@@ -3835,7 +4085,7 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
     if (seed < 0) {
         seed = (int)time(NULL);
     }
-    set_random_seed(seed);
+    sd->rng->manual_seed(seed);
 
     int64_t t0 = ggml_time_ms();
     ggml_tensor* c = sd->get_learned_condition(ctx, prompt);
@@ -3856,9 +4106,9 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
     int W = width / 8;
     int H = height / 8;
     struct ggml_tensor* x_t = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, W, H, C, 1);
-    ggml_tensor_set_f32_randn(x_t);
+    ggml_tensor_set_f32_randn(x_t, sd->rng);
 
-    std::vector<float> sigmas = sd->denoiser->get_sigmas(sample_steps);
+    std::vector<float> sigmas = sd->denoiser->schedule->get_sigmas(sample_steps);
 
     LOG_INFO("start sampling");
     struct ggml_tensor* x_0 = sd->sample(ctx, x_t, c, uc, cfg_scale, sample_method, sigmas);
@@ -3907,14 +4157,14 @@ std::vector<uint8_t> StableDiffusion::img2img(const std::vector<uint8_t>& init_i
                                               SampleMethod sample_method,
                                               int sample_steps,
                                               float strength,
-                                              int seed) {
+                                              int64_t seed) {
     std::vector<uint8_t> result;
     if (init_img_vec.size() != width * height * 3) {
         return result;
     }
     LOG_INFO("img2img %dx%d", width, height);
 
-    std::vector<float> sigmas = sd->denoiser->get_sigmas(sample_steps);
+    std::vector<float> sigmas = sd->denoiser->schedule->get_sigmas(sample_steps);
     size_t t_enc = static_cast<size_t>(sample_steps * strength);
     LOG_INFO("target t_enc is %zu steps", t_enc);
     std::vector<float> sigma_sched;
@@ -3935,7 +4185,7 @@ std::vector<uint8_t> StableDiffusion::img2img(const std::vector<uint8_t>& init_i
     if (seed < 0) {
         seed = (int)time(NULL);
     }
-    set_random_seed(seed);
+    sd->rng->manual_seed(seed);
 
     ggml_tensor* init_img = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, height, 3, 1);
     image_vec_to_ggml(init_img_vec, init_img);

stable-diffusion.h

@@ -4,15 +4,32 @@
 #include <memory>
 #include <vector>
 
-enum class SDLogLevel {
+enum SDLogLevel {
     DEBUG,
     INFO,
     WARN,
     ERROR
 };
 
+enum RNGType {
+    STD_DEFAULT_RNG,
+    CUDA_RNG
+};
+
 enum SampleMethod {
-    EULAR_A,
+    EULER_A,
+    EULER,
+    HEUN,
+    DPMPP2M,
+    DPMPP2Mv2,
+    N_SAMPLE_METHODS
+};
+
+enum Schedule {
+    DEFAULT,
+    DISCRETE,
+    KARRAS,
+    N_SCHEDULES
 };
 
 class StableDiffusionGGML;
@@ -24,8 +41,9 @@ class StableDiffusion {
    public:
     StableDiffusion(int n_threads = -1,
                     bool vae_decode_only = false,
-                    bool free_params_immediately = false);
-    bool load_from_file(const std::string& file_path);
+                    bool free_params_immediately = false,
+                    RNGType rng_type = STD_DEFAULT_RNG);
+    bool load_from_file(const std::string& file_path, Schedule d = DEFAULT);
     std::vector<uint8_t> txt2img(
         const std::string& prompt,
         const std::string& negative_prompt,
@@ -34,7 +52,7 @@ class StableDiffusion {
         int height,
         SampleMethod sample_method,
         int sample_steps,
-        int seed);
+        int64_t seed);
     std::vector<uint8_t> img2img(
         const std::vector<uint8_t>& init_img,
         const std::string& prompt,
@@ -45,7 +63,7 @@ class StableDiffusion {
         SampleMethod sample_method,
         int sample_steps,
         float strength,
-        int seed);
+        int64_t seed);
 };
 
 void set_sd_log_level(SDLogLevel level);