feat: add lcm sampler support

This referenced an issue discussion of the stable-diffusion-webui at https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13952, which may not be too perfect.
chore: clear the msvc compilation warning
2026-07-24 03:40:53 -05:00 · 2023-11-17 22:53:46 +08:00 · 2023-10-28 20:55:24 +08:00 · 2023-10-28 20:19:15 +08:00 · 2023-10-27 00:50:23 +08:00 · 2023-10-24 22:04:00 +08:00
12 changed files with 1062 additions and 151 deletions
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -24,6 +24,7 @@ endif()
 # general
 #option(SD_BUILD_TESTS                "sd: build tests"    ${SD_STANDALONE})
 option(SD_BUILD_EXAMPLES             "sd: build examples" ${SD_STANDALONE})
+option(BUILD_SHARED_LIBS             "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)


--- a/README.md
+++ b/README.md
@@ -20,7 +20,15 @@ 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`
+    - `DPM2`
+    - `DPM++ 2M`
+    - [`DPM++ 2M v2`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457)
+    - `DPM++ 2S a`
+    - [`LCM`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13952)
 - Cross-platform reproducibility (`--rng cuda`, consistent with the `stable-diffusion-webui GPU RNG`)
+- Embedds generation parameters into png output as webui-compatible text string
 - Supported platforms
    - Linux
    - Mac OS
@@ -65,7 +73,7 @@ git submodule update
    ```shell
    curl -L -O https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt
    # curl -L -O https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
-    # curl -L -o https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-nonema-pruned.safetensors
+    # curl -L -O https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-nonema-pruned.safetensors
    ```

 - convert weights to ggml model format
@@ -73,6 +81,7 @@ git submodule update
    ```shell
    cd models
    pip install -r requirements.txt
+    # (optional) python convert_diffusers_to_original_stable_diffusion.py --model_path [path to diffusers weights] --checkpoint_path [path to weights]
    python convert.py [path to weights] --out_type [output precision]
    # For example, python convert.py sd-v1-4.ckpt --out_type f16
    ```
@@ -125,8 +134,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, lcm}
+                                     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
 ```
@@ -187,3 +198,4 @@ docker run -v /path/to/models:/models -v /path/to/output/:/output sd [args...]
 - [stable-diffusion-stability-ai](https://github.com/Stability-AI/stablediffusion)
 - [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
 - [k-diffusion](https://github.com/crowsonkb/k-diffusion)
+- [latent-consistency-model](https://github.com/luosiallen/latent-consistency-model)
--- a/examples/main.cpp
+++ b/examples/main.cpp
@@ -72,6 +72,25 @@ 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",
+    "dpm2",
+    "dpm++2s_a",
+    "dpm++2m",
+    "dpm++2mv2",
+    "lcm",
+};
+
+// 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;
@@ -83,10 +102,11 @@ 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;
-    RNGType rng_type = STD_DEFAULT_RNG;
+    RNGType rng_type = CUDA_RNG;
    int64_t seed = 42;
    bool verbose = false;

@@ -102,7 +122,8 @@ 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("    rng:             %s\n", rng_type_to_str[rng_type]);
@@ -128,10 +149,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, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, lcm}\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: std_default)\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");
 }

@@ -228,12 +251,46 @@ void parse_args(int argc, const char* argv[], Option* opt) {
                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::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);
@@ -285,13 +342,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);
    }

@@ -311,6 +368,18 @@ void parse_args(int argc, const char* argv[], Option* opt) {
    }
 }

+std::string basename(const std::string& path) {
+    size_t pos = path.find_last_of('/');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    pos = path.find_last_of('\\');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    return path;
+}
+
 int main(int argc, const char* argv[]) {
    Option opt;
    parse_args(argc, argv, &opt);
@@ -337,13 +406,13 @@ 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;
        }
@@ -351,7 +420,7 @@ int main(int argc, const char* argv[]) {
    }

    StableDiffusion sd(opt.n_threads, vae_decode_only, true, opt.rng_type);
-    if (!sd.load_from_file(opt.model_path)) {
+    if (!sd.load_from_file(opt.model_path, opt.schedule)) {
        return 1;
    }

@@ -383,8 +452,25 @@ int main(int argc, const char* argv[]) {
        return 1;
    }

-    stbi_write_png(opt.output_path.c_str(), opt.w, opt.h, 3, img.data(), 0);
+    std::string parameter_string = opt.prompt + "\n";
+    if (opt.negative_prompt.size() != 0) {
+        parameter_string += "Negative prompt: " + opt.negative_prompt + "\n";
+    }
+    parameter_string += "Steps: " + std::to_string(opt.sample_steps) + ", ";
+    parameter_string += "CFG scale: " + std::to_string(opt.cfg_scale) + ", ";
+    parameter_string += "Seed: " + std::to_string(opt.seed) + ", ";
+    parameter_string += "Size: " + std::to_string(opt.w) + "x" + std::to_string(opt.h) + ", ";
+    parameter_string += "Model: " + basename(opt.model_path) + ", ";
+    parameter_string += "RNG: " + std::string(rng_type_to_str[opt.rng_type]) + ", ";
+    parameter_string += "Sampler: " + std::string(sample_method_str[opt.sample_method]);
+    if (opt.schedule == KARRAS) {
+        parameter_string += " karras";
+    }
+    parameter_string += ", ";
+    parameter_string += "Version: stable-diffusion.cpp";
+
+    stbi_write_png(opt.output_path.c_str(), opt.w, opt.h, 3, img.data(), 0, parameter_string.c_str());
    printf("save result image to '%s'\n", opt.output_path.c_str());

    return 0;
-}
+}
--- a/examples/stb_image_write.h
+++ b/examples/stb_image_write.h
@@ -173,7 +173,7 @@ STBIWDEF int stbi_write_force_png_filter;
 #endif

 #ifndef STBI_WRITE_NO_STDIO
-STBIWDEF int stbi_write_png(char const *filename, int w, int h, int comp, const void  *data, int stride_in_bytes);
+STBIWDEF int stbi_write_png(char const *filename, int w, int h, int comp, const void  *data, int stride_in_bytes, const char* parameters = NULL);
 STBIWDEF int stbi_write_bmp(char const *filename, int w, int h, int comp, const void  *data);
 STBIWDEF int stbi_write_tga(char const *filename, int w, int h, int comp, const void  *data);
 STBIWDEF int stbi_write_hdr(char const *filename, int w, int h, int comp, const float *data);
@@ -1125,9 +1125,10 @@ static void stbiw__encode_png_line(unsigned char *pixels, int stride_bytes, int
   }
 }

-STBIWDEF unsigned char *stbi_write_png_to_mem(const unsigned char *pixels, int stride_bytes, int x, int y, int n, int *out_len)
+STBIWDEF unsigned char *stbi_write_png_to_mem(const unsigned char *pixels, int stride_bytes, int x, int y, int n, int *out_len, const char* parameters)
 {
   int force_filter = stbi_write_force_png_filter;
+   int param_length = 0;
   int ctype[5] = { -1, 0, 4, 2, 6 };
   unsigned char sig[8] = { 137,80,78,71,13,10,26,10 };
   unsigned char *out,*o, *filt, *zlib;
@@ -1177,10 +1178,15 @@ STBIWDEF unsigned char *stbi_write_png_to_mem(const unsigned char *pixels, int s
   STBIW_FREE(filt);
   if (!zlib) return 0;

+   if(parameters != NULL) {
+      param_length = strlen(parameters);
+      param_length += strlen("parameters") + 1; // For the name and the null-byte
+   }
+
   // each tag requires 12 bytes of overhead
-   out = (unsigned char *) STBIW_MALLOC(8 + 12+13 + 12+zlen + 12);
+   out = (unsigned char *) STBIW_MALLOC(8 + 12+13 + 12+zlen + 12 + ((parameters)?(param_length+12):0));
   if (!out) return 0;
-   *out_len = 8 + 12+13 + 12+zlen + 12;
+   *out_len = 8 + 12+13 + 12+zlen + 12 + ((parameters)?(param_length+12):0);

   o=out;
   STBIW_MEMMOVE(o,sig,8); o+= 8;
@@ -1195,6 +1201,17 @@ STBIWDEF unsigned char *stbi_write_png_to_mem(const unsigned char *pixels, int s
   *o++ = 0;
   stbiw__wpcrc(&o,13);

+   if(parameters != NULL) {
+      stbiw__wp32(o, param_length);
+      stbiw__wptag(o, "tEXt");
+      STBIW_MEMMOVE(o, "parameters", strlen("parameters"));
+      o+=strlen("parameters");
+      *o++ = 0; // Null pyte separator
+      STBIW_MEMMOVE(o, parameters, strlen(parameters));
+      o+=strlen(parameters);
+      stbiw__wpcrc(&o, param_length);
+   }
+
   stbiw__wp32(o, zlen);
   stbiw__wptag(o, "IDAT");
   STBIW_MEMMOVE(o, zlib, zlen);
@@ -1212,11 +1229,11 @@ STBIWDEF unsigned char *stbi_write_png_to_mem(const unsigned char *pixels, int s
 }

 #ifndef STBI_WRITE_NO_STDIO
-STBIWDEF int stbi_write_png(char const *filename, int x, int y, int comp, const void *data, int stride_bytes)
+STBIWDEF int stbi_write_png(char const *filename, int x, int y, int comp, const void *data, int stride_bytes, const char* parameters)
 {
   FILE *f;
   int len;
-   unsigned char *png = stbi_write_png_to_mem((const unsigned char *) data, stride_bytes, x, y, comp, &len);
+   unsigned char *png = stbi_write_png_to_mem((const unsigned char *) data, stride_bytes, x, y, comp, &len, parameters);
   if (png == NULL) return 0;

   f = stbiw__fopen(filename, "wb");
@@ -1231,7 +1248,7 @@ STBIWDEF int stbi_write_png(char const *filename, int x, int y, int comp, const
 STBIWDEF int stbi_write_png_to_func(stbi_write_func *func, void *context, int x, int y, int comp, const void *data, int stride_bytes)
 {
   int len;
-   unsigned char *png = stbi_write_png_to_mem((const unsigned char *) data, stride_bytes, x, y, comp, &len);
+   unsigned char *png = stbi_write_png_to_mem((const unsigned char *) data, stride_bytes, x, y, comp, &len, NULL);
   if (png == NULL) return 0;
   func(context, png, len);
   STBIW_FREE(png);
--- a/2
+++ b/2
--- a/models/.gitignore
+++ b/models/.gitignore
@@ -1,4 +1,5 @@
 *.bin
 *.ckpt
 *.safetensor
+*.safetensors
 *.log
--- a/models/convert.py
+++ b/models/convert.py
@@ -179,9 +179,9 @@ def preprocess(state_dict):
        state_dict["alphas_cumprod"] = alphas_cumprod
    
    new_state_dict = {}
-    for name in state_dict.keys():
+    for name, w in state_dict.items():
        # ignore unused tensors
-        if not isinstance(state_dict[name], torch.Tensor):
+        if not isinstance(w, torch.Tensor):
            continue
        skip = False
        for unused_tensor in unused_tensors:
@@ -190,13 +190,25 @@ def preprocess(state_dict):
                break
        if skip:
            continue
-        
+
+        # # convert BF16 to FP16
+        if w.dtype == torch.bfloat16:
+            w = w.to(torch.float16)
+
        # convert open_clip to hf CLIPTextModel (for SD2.x)
        open_clip_to_hf_clip_model = {
            "cond_stage_model.model.ln_final.bias": "cond_stage_model.transformer.text_model.final_layer_norm.bias",
            "cond_stage_model.model.ln_final.weight": "cond_stage_model.transformer.text_model.final_layer_norm.weight",
            "cond_stage_model.model.positional_embedding": "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
            "cond_stage_model.model.token_embedding.weight": "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight",
+            "first_stage_model.decoder.mid.attn_1.to_k.bias": "first_stage_model.decoder.mid.attn_1.k.bias",
+            "first_stage_model.decoder.mid.attn_1.to_k.weight": "first_stage_model.decoder.mid.attn_1.k.weight",
+            "first_stage_model.decoder.mid.attn_1.to_out.0.bias": "first_stage_model.decoder.mid.attn_1.proj_out.bias",
+            "first_stage_model.decoder.mid.attn_1.to_out.0.weight": "first_stage_model.decoder.mid.attn_1.proj_out.weight",
+            "first_stage_model.decoder.mid.attn_1.to_q.bias": "first_stage_model.decoder.mid.attn_1.q.bias",
+            "first_stage_model.decoder.mid.attn_1.to_q.weight": "first_stage_model.decoder.mid.attn_1.q.weight",
+            "first_stage_model.decoder.mid.attn_1.to_v.bias": "first_stage_model.decoder.mid.attn_1.v.bias",
+            "first_stage_model.decoder.mid.attn_1.to_v.weight": "first_stage_model.decoder.mid.attn_1.v.weight",
        }
        open_clip_to_hk_clip_resblock = {
            "attn.out_proj.bias": "self_attn.out_proj.bias",
@@ -214,22 +226,19 @@ def preprocess(state_dict):
        hf_clip_resblock_prefix = "cond_stage_model.transformer.text_model.encoder.layers."
        if name in open_clip_to_hf_clip_model:
            new_name = open_clip_to_hf_clip_model[name]
-            new_state_dict[new_name] = state_dict[name]
            print(f"preprocess {name} => {new_name}")
-            continue
+            name = new_name
        if name.startswith(open_clip_resblock_prefix):
            remain = name[len(open_clip_resblock_prefix):]
            idx = remain.split(".")[0]
            suffix = remain[len(idx)+1:]
            if suffix == "attn.in_proj_weight":
-                w = state_dict[name]
                w_q, w_k, w_v = w.chunk(3)
                for new_suffix, new_w in zip(["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], [w_q, w_k, w_v]):
                    new_name = hf_clip_resblock_prefix + idx + "." + new_suffix
                    new_state_dict[new_name] = new_w
                    print(f"preprocess {name}{w.size()} => {new_name}{new_w.size()}")
            elif suffix == "attn.in_proj_bias":
-                w = state_dict[name]
                w_q, w_k, w_v = w.chunk(3)
                for new_suffix, new_w in zip(["self_attn.q_proj.bias", "self_attn.k_proj.bias", "self_attn.v_proj.bias"], [w_q, w_k, w_v]):
                    new_name = hf_clip_resblock_prefix + idx + "." + new_suffix
@@ -238,20 +247,27 @@ def preprocess(state_dict):
            else:
                new_suffix = open_clip_to_hk_clip_resblock[suffix]
                new_name = hf_clip_resblock_prefix + idx + "." + new_suffix
-                new_state_dict[new_name] = state_dict[name]
+                new_state_dict[new_name] = w
                print(f"preprocess {name} => {new_name}")
            continue
        
        # convert unet transformer linear to conv2d 1x1
        if name.startswith("model.diffusion_model.") and (name.endswith("proj_in.weight") or name.endswith("proj_out.weight")):
-            w = state_dict[name]
-            if len(state_dict[name].shape) == 2:
+            if len(w.shape) == 2:
                new_w = w.unsqueeze(2).unsqueeze(3)
                new_state_dict[name] = new_w
                print(f"preprocess {name} {w.size()} => {name} {new_w.size()}")
                continue

-        new_state_dict[name] = state_dict[name]
+        # convert vae attn block linear to conv2d 1x1
+        if name.startswith("first_stage_model.") and "attn_1" in name:
+            if len(w.shape) == 2:
+                new_w = w.unsqueeze(2).unsqueeze(3)
+                new_state_dict[name] = new_w
+                print(f"preprocess {name} {w.size()} => {name} {new_w.size()}")
+                continue
+
+        new_state_dict[name] = w
    return new_state_dict

 def convert(model_path, out_type = None, out_file=None):
--- a/models/convert_diffusers_to_original_stable_diffusion.py
+++ b/models/convert_diffusers_to_original_stable_diffusion.py
@@ -0,0 +1,335 @@
+# Copy from https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
+# LICENSE: https://github.com/huggingface/diffusers/blob/main/LICENSE
+# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
+# *Only* converts the UNet, VAE, and Text Encoder.
+# Does not convert optimizer state or any other thing.
+
+import argparse
+import os.path as osp
+import re
+
+import torch
+from safetensors.torch import load_file, save_file
+
+
+# =================#
+# UNet Conversion #
+# =================#
+
+unet_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+    ("input_blocks.0.0.weight", "conv_in.weight"),
+    ("input_blocks.0.0.bias", "conv_in.bias"),
+    ("out.0.weight", "conv_norm_out.weight"),
+    ("out.0.bias", "conv_norm_out.bias"),
+    ("out.2.weight", "conv_out.weight"),
+    ("out.2.bias", "conv_out.bias"),
+]
+
+unet_conversion_map_resnet = [
+    # (stable-diffusion, HF Diffusers)
+    ("in_layers.0", "norm1"),
+    ("in_layers.2", "conv1"),
+    ("out_layers.0", "norm2"),
+    ("out_layers.3", "conv2"),
+    ("emb_layers.1", "time_emb_proj"),
+    ("skip_connection", "conv_shortcut"),
+]
+
+unet_conversion_map_layer = []
+# hardcoded number of downblocks and resnets/attentions...
+# would need smarter logic for other networks.
+for i in range(4):
+    # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+        hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+        if i < 3:
+            # no attention layers in down_blocks.3
+            hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+        # loop over resnets/attentions for upblocks
+        hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+        sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+        if i > 0:
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+        # no downsample in down_blocks.3
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+        sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        # no upsample in up_blocks.3
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+hf_mid_atn_prefix = "mid_block.attentions.0."
+sd_mid_atn_prefix = "middle_block.1."
+unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+def convert_unet_state_dict(unet_state_dict):
+    # buyer beware: this is a *brittle* function,
+    # and correct output requires that all of these pieces interact in
+    # the exact order in which I have arranged them.
+    mapping = {k: k for k in unet_state_dict.keys()}
+    for sd_name, hf_name in unet_conversion_map:
+        mapping[hf_name] = sd_name
+    for k, v in mapping.items():
+        if "resnets" in k:
+            for sd_part, hf_part in unet_conversion_map_resnet:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    for k, v in mapping.items():
+        for sd_part, hf_part in unet_conversion_map_layer:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+    return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+vae_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("nin_shortcut", "conv_shortcut"),
+    ("norm_out", "conv_norm_out"),
+    ("mid.attn_1.", "mid_block.attentions.0."),
+]
+
+for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+        hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+        sd_down_prefix = f"encoder.down.{i}.block.{j}."
+        vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+        sd_downsample_prefix = f"down.{i}.downsample."
+        vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"up.{3-i}.upsample."
+        vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+        hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+        sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+        vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+# this part accounts for mid blocks in both the encoder and the decoder
+for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+vae_conversion_map_attn = [
+    # (stable-diffusion, HF Diffusers)
+    ("norm.", "group_norm."),
+    ("q.", "query."),
+    ("k.", "key."),
+    ("v.", "value."),
+    ("proj_out.", "proj_attn."),
+]
+
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+    return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+    mapping = {k: k for k in vae_state_dict.keys()}
+    for k, v in mapping.items():
+        for sd_part, hf_part in vae_conversion_map:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    for k, v in mapping.items():
+        if "attentions" in k:
+            for sd_part, hf_part in vae_conversion_map_attn:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+    weights_to_convert = ["q", "k", "v", "proj_out"]
+    for k, v in new_state_dict.items():
+        for weight_name in weights_to_convert:
+            if f"mid.attn_1.{weight_name}.weight" in k:
+                print(f"Reshaping {k} for SD format")
+                new_state_dict[k] = reshape_weight_for_sd(v)
+    return new_state_dict
+
+
+# =========================#
+# Text Encoder Conversion #
+# =========================#
+
+
+textenc_conversion_lst = [
+    # (stable-diffusion, HF Diffusers)
+    ("resblocks.", "text_model.encoder.layers."),
+    ("ln_1", "layer_norm1"),
+    ("ln_2", "layer_norm2"),
+    (".c_fc.", ".fc1."),
+    (".c_proj.", ".fc2."),
+    (".attn", ".self_attn"),
+    ("ln_final.", "transformer.text_model.final_layer_norm."),
+    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
+    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
+]
+protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
+textenc_pattern = re.compile("|".join(protected.keys()))
+
+# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
+code2idx = {"q": 0, "k": 1, "v": 2}
+
+
+def convert_text_enc_state_dict_v20(text_enc_dict):
+    new_state_dict = {}
+    capture_qkv_weight = {}
+    capture_qkv_bias = {}
+    for k, v in text_enc_dict.items():
+        if (
+            k.endswith(".self_attn.q_proj.weight")
+            or k.endswith(".self_attn.k_proj.weight")
+            or k.endswith(".self_attn.v_proj.weight")
+        ):
+            k_pre = k[: -len(".q_proj.weight")]
+            k_code = k[-len("q_proj.weight")]
+            if k_pre not in capture_qkv_weight:
+                capture_qkv_weight[k_pre] = [None, None, None]
+            capture_qkv_weight[k_pre][code2idx[k_code]] = v
+            continue
+
+        if (
+            k.endswith(".self_attn.q_proj.bias")
+            or k.endswith(".self_attn.k_proj.bias")
+            or k.endswith(".self_attn.v_proj.bias")
+        ):
+            k_pre = k[: -len(".q_proj.bias")]
+            k_code = k[-len("q_proj.bias")]
+            if k_pre not in capture_qkv_bias:
+                capture_qkv_bias[k_pre] = [None, None, None]
+            capture_qkv_bias[k_pre][code2idx[k_code]] = v
+            continue
+
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
+        new_state_dict[relabelled_key] = v
+
+    for k_pre, tensors in capture_qkv_weight.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
+
+    for k_pre, tensors in capture_qkv_bias.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
+
+    return new_state_dict
+
+
+def convert_text_enc_state_dict(text_enc_dict):
+    return text_enc_dict
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
+    parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
+    parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
+    parser.add_argument(
+        "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
+    )
+
+    args = parser.parse_args()
+
+    assert args.model_path is not None, "Must provide a model path!"
+
+    assert args.checkpoint_path is not None, "Must provide a checkpoint path!"
+
+    # Path for safetensors
+    unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
+    vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
+    text_enc_path = osp.join(args.model_path, "text_encoder", "model.safetensors")
+
+    # Load models from safetensors if it exists, if it doesn't pytorch
+    if osp.exists(unet_path):
+        unet_state_dict = load_file(unet_path, device="cpu")
+    else:
+        unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
+        unet_state_dict = torch.load(unet_path, map_location="cpu")
+
+    if osp.exists(vae_path):
+        vae_state_dict = load_file(vae_path, device="cpu")
+    else:
+        vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
+        vae_state_dict = torch.load(vae_path, map_location="cpu")
+
+    if osp.exists(text_enc_path):
+        text_enc_dict = load_file(text_enc_path, device="cpu")
+    else:
+        text_enc_path = osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
+        text_enc_dict = torch.load(text_enc_path, map_location="cpu")
+
+    # Convert the UNet model
+    unet_state_dict = convert_unet_state_dict(unet_state_dict)
+    unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
+
+    # Convert the VAE model
+    vae_state_dict = convert_vae_state_dict(vae_state_dict)
+    vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
+
+    # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
+    is_v20_model = "text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict
+
+    if is_v20_model:
+        # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
+        text_enc_dict = {"transformer." + k: v for k, v in text_enc_dict.items()}
+        text_enc_dict = convert_text_enc_state_dict_v20(text_enc_dict)
+        text_enc_dict = {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
+    else:
+        text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
+        text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
+
+    # Put together new checkpoint
+    state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
+    if args.half:
+        state_dict = {k: v.half() for k, v in state_dict.items()}
+
+    if args.use_safetensors:
+        save_file(state_dict, args.checkpoint_path)
+    else:
+        state_dict = {"state_dict": state_dict}
+        torch.save(state_dict, args.checkpoint_path)
--- a/rng.h
+++ b/rng.h
@@ -16,7 +16,7 @@ class STDDefaultRNG : public RNG {

   public:
    void manual_seed(uint64_t seed) {
-        generator.seed(seed);
+        generator.seed((unsigned int)seed);
    }

    std::vector<float> randn(uint32_t n) {
@@ -24,7 +24,7 @@ class STDDefaultRNG : public RNG {
        float mean = 0.0;
        float stddev = 1.0;
        std::normal_distribution<float> distribution(mean, stddev);
-        for (int i = 0; i < n; i++) {
+        for (uint32_t i = 0; i < n; i++) {
            float random_number = distribution(generator);
            result.push_back(random_number);
        }
--- a/rng_philox.h
+++ b/rng_philox.h
@@ -16,8 +16,8 @@ class PhiloxRNG : public RNG {
   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;
+    float two_pow32_inv = 2.3283064e-10f;
+    float two_pow32_inv_2pi = 2.3283064e-10f * 6.2831855f;

    std::vector<uint32_t> uint32(uint64_t x) {
        std::vector<uint32_t> result(2);
@@ -27,10 +27,10 @@ class PhiloxRNG : public RNG {
    }

    std::vector<std::vector<uint32_t>> uint32(const std::vector<uint64_t>& x) {
-        int N = x.size();
+        uint32_t N = (uint32_t)x.size();
        std::vector<std::vector<uint32_t>> result(2, std::vector<uint32_t>(N));

-        for (int i = 0; i < N; ++i) {
+        for (uint32_t 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);
        }
@@ -41,7 +41,7 @@ class PhiloxRNG : public RNG {
    //  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();
+        uint32_t N = (uint32_t)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]));
@@ -63,7 +63,7 @@ class PhiloxRNG : public RNG {
    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();
+        uint32_t N = (uint32_t)counter[0].size();
        for (int i = 0; i < rounds - 1; ++i) {
            philox4_round(counter, key);

@@ -81,7 +81,7 @@ class PhiloxRNG : public RNG {
        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 s = sqrt(-2.0f * log(u));

        float r1 = s * sin(v);
        return r1;
@@ -115,8 +115,8 @@ class PhiloxRNG : public RNG {
        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]));
+        for (uint32_t i = 0; i < n; ++i) {
+            result.push_back(box_muller((float)g[0][i], (float)g[1][i]));
        }
        return result;
    }
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -18,6 +18,8 @@
 #include "rng_philox.h"
 #include "stable-diffusion.h"

+#define EPS 1e-05f
+
 static SDLogLevel log_level = SDLogLevel::INFO;

 #define __FILENAME__ "stable-diffusion.cpp"
@@ -120,9 +122,9 @@ ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_pa
 }

 void ggml_tensor_set_f32_randn(struct ggml_tensor* tensor, std::shared_ptr<RNG> rng) {
-    uint32_t n = ggml_nelements(tensor);
+    uint32_t n = (uint32_t)ggml_nelements(tensor);
    std::vector<float> random_numbers = rng->randn(n);
-    for (int i = 0; i < n; i++) {
+    for (uint32_t i = 0; i < n; i++) {
        ggml_set_f32_1d(tensor, i, random_numbers[i]);
    }
 }
@@ -436,12 +438,12 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
        std::string weight = m[1];

        if (text == "(") {
-            round_brackets.push_back(res.size());
+            round_brackets.push_back((int)res.size());
        } else if (text == "[") {
-            square_brackets.push_back(res.size());
+            square_brackets.push_back((int)res.size());
        } else if (!weight.empty()) {
            if (!round_brackets.empty()) {
-                multiply_range(round_brackets.back(), std::stod(weight));
+                multiply_range(round_brackets.back(), std::stof(weight));
                round_brackets.pop_back();
            }
        } else if (text == ")" && !round_brackets.empty()) {
@@ -586,7 +588,7 @@ struct ResidualAttentionBlock {

        // layer norm 1
        {
-            x = ggml_norm(ctx, x);
+            x = ggml_norm(ctx, x, EPS);
            x = ggml_add(ctx,
                         ggml_mul(ctx, ggml_repeat(ctx, ln1_w, x), x),
                         ggml_repeat(ctx, ln1_b, x));
@@ -636,7 +638,7 @@ struct ResidualAttentionBlock {

        // layer norm 2
        {
-            x = ggml_norm(ctx, x);
+            x = ggml_norm(ctx, x, EPS);

            x = ggml_add(ctx, ggml_mul(ctx, ggml_repeat(ctx, ln2_w, x), x),
                         ggml_repeat(ctx, ln2_b, x));
@@ -766,7 +768,7 @@ struct CLIPTextModel {

        // final layer norm
        {
-            x = ggml_norm(ctx, x);
+            x = ggml_norm(ctx, x, EPS);

            x = ggml_add(ctx, ggml_mul(ctx, ggml_repeat(ctx, final_ln_w, x), x),
                         ggml_repeat(ctx, final_ln_b, x));
@@ -1200,7 +1202,7 @@ struct SpatialTransformer {
            // layer norm 1
            {
                x = ggml_reshape_2d(ctx, x, c, w * h * n);
-                x = ggml_norm(ctx, x);
+                x = ggml_norm(ctx, x, EPS);
                x = ggml_add(ctx,
                             ggml_mul(ctx,
                                      ggml_repeat(ctx, transformer.norm1_w, x),
@@ -1248,7 +1250,7 @@ struct SpatialTransformer {

            // layer norm 2
            {
-                x = ggml_norm(ctx, x);
+                x = ggml_norm(ctx, x, EPS);
                x = ggml_add(ctx,
                             ggml_mul(ctx,
                                      ggml_repeat(ctx, transformer.norm2_w, x), x),
@@ -1299,7 +1301,7 @@ struct SpatialTransformer {
            // layer norm 3
            {
                x = ggml_reshape_2d(ctx, x, c, h * w * n);  // [N * h * w, in_channels]
-                x = ggml_norm(ctx, x);
+                x = ggml_norm(ctx, x, EPS);
                x = ggml_add(ctx,
                             ggml_mul(ctx,
                                      ggml_repeat(ctx, transformer.norm3_w, x), x),
@@ -2654,23 +2656,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);
@@ -2705,11 +2696,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((float)t_max));
+            result.push_back(0);
+            return result;
+        }
+
+        float step = static_cast<float>(t_max) / static_cast<float>(n - 1);
+        for (uint32_t 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.1f;
+        float sigma_max = 10.f;
+        float rho = 7.f;
+
+        std::vector<float> result(n + 1);
+
+        float min_inv_rho = pow(sigma_min, (1.f / rho));
+        float max_inv_rho = pow(sigma_max, (1.f / rho));
+        for (uint32_t i = 0; i < n; i++) {
+            // Eq. (5) from Karras et al 2022
+            result[i] = pow(max_inv_rho + (float)i / ((float)n - 1.f) * (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) {
@@ -2719,7 +2758,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) {
@@ -2755,7 +2794,7 @@ 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;

@@ -2789,7 +2828,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);
@@ -3084,10 +3123,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;
@@ -3099,7 +3157,9 @@ 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
+        struct ggml_cplan cplan;
+
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
        // calculate the amount of memory required
        {
            struct ggml_init_params params;
@@ -3123,7 +3183,7 @@ class StableDiffusionGGML {
            ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);

            struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
-            struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
+            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",
@@ -3156,7 +3216,7 @@ class StableDiffusionGGML {
        ggml_hold_dynamic_tensor(out);

        struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
-        struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
+        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);
@@ -3201,7 +3261,8 @@ 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
+        struct ggml_cplan cplan;
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
        // calculate the amount of memory required
        {
            struct ggml_init_params params;
@@ -3223,7 +3284,7 @@ class StableDiffusionGGML {
            struct ggml_tensor* hidden_states = cond_stage_model.text_model.forward(ctx, input_ids);

            struct ggml_cgraph* cond_graph = ggml_build_forward_ctx(ctx, hidden_states);
-            struct ggml_cplan cplan = ggml_graph_plan(cond_graph, n_threads);
+            cplan = ggml_graph_plan(cond_graph, n_threads);
            ctx_size += cplan.work_size;

            ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
@@ -3334,8 +3395,9 @@ class StableDiffusionGGML {
        // print_ggml_tensor(x_t);
        struct ggml_tensor* x = ggml_dup_tensor(res_ctx, x_t);
        copy_ggml_tensor(x, x_t);
+        struct ggml_cplan cplan;

-        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;
@@ -3361,7 +3423,7 @@ class StableDiffusionGGML {
            ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);

            struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
-            struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
+            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",
@@ -3393,8 +3455,10 @@ 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_ctx(ctx, out);
-        struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
+        struct ggml_cgraph* diffusion_graph = ggml_new_graph(ctx);
+        diffusion_graph->order = GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT;
+        ggml_build_forward_expand(diffusion_graph, out);
+        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);
@@ -3436,7 +3500,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);

@@ -3493,69 +3557,424 @@ 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, rng);
-                    // 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 DPM2: {
+                LOG_INFO("sampling using DPM2 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];
+                        }
+                    }
+
+                    if (sigmas[i + 1] == 0) {
+                        // Euler step
+                        // x = x + d * dt
+                        float dt = sigmas[i + 1] - sigmas[i];
+                        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 {
+                        // DPM-Solver-2
+                        float sigma_mid = exp(0.5f * (log(sigmas[i]) + log(sigmas[i + 1])));
+                        float dt_1 = sigma_mid - sigmas[i];
+                        float dt_2 = sigmas[i + 1] - sigmas[i];
+
+                        float* vec_d = (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_1;
+                        }
+
+                        denoise(x2, sigma_mid, 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]) / sigma_mid;
+                            vec_x[j] = vec_x[j] + d2 * dt_2;
+                        }
+                    }
+                }
+
+            } break;
+            case DPMPP2S_A: {
+                LOG_INFO("sampling using DPM++ (2s) 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);
+                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);
+
+                    // 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);
+                    auto t_fn = [](float sigma) -> float { return -log(sigma); };
+                    auto sigma_fn = [](float t) -> float { return exp(-t); };
+
+                    if (sigma_down == 0) {
+                        // Euler step
+                        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]) / sigmas[i];
+                        }
+
+                        // TODO: If sigma_down == 0, isn't this wrong?
+                        // But
+                        // https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L525
+                        // has this exactly the same way.
+                        float dt = sigma_down - sigmas[i];
+                        for (int j = 0; j < ggml_nelements(d); j++) {
+                            vec_x[j] = vec_x[j] + vec_d[j] * dt;
+                        }
+                    } else {
+                        // DPM-Solver++(2S)
+                        float t = t_fn(sigmas[i]);
+                        float t_next = t_fn(sigma_down);
+                        float h = t_next - t;
+                        float s = t + 0.5f * h;
+
+                        float* vec_d = (float*)d->data;
+                        float* vec_x = (float*)x->data;
+                        float* vec_x2 = (float*)x2->data;
+                        float* vec_denoised = (float*)denoised->data;
+
+                        // First half-step
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x2[j] = (sigma_fn(s) / sigma_fn(t)) * vec_x[j] - (exp(-h * 0.5f) - 1) * vec_denoised[j];
+                        }
+
+                        denoise(x2, sigmas[i + 1], i + 1);
+
+                        // Second half-step
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = (sigma_fn(t_next) / sigma_fn(t)) * vec_x[j] - (exp(-h) - 1) * vec_denoised[j];
+                        }
+                    }
+
+                    // Noise addition
+                    if (sigmas[i + 1] > 0) {
+                        ggml_tensor_set_f32_randn(noise, rng);
+                        {
+                            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 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.f;
+                    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.f + 1.f / (2.f * r)) * vec_denoised[j] - (1.f / (2.f * 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.f;
+                        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.f;
+                        float b = exp(-h_d) - 1.f;
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            float denoised_d = (1.f + 1.f / (2.f * r)) * vec_denoised[j] - (1.f / (2.f * 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 LCM:  // Latent Consistency Models
+            {
+                LOG_INFO("sampling using LCM 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];
+
+                    // denoise
+                    denoise(x, sigma, i + 1);
+
+                    // x = denoised
+                    {
+                        float* vec_x = (float*)x->data;
+                        float* vec_denoised = (float*)denoised->data;
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = vec_denoised[j];
+                        }
+                    }
+
+                    if (sigmas[i + 1] > 0) {
+                        // x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
+                        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 j = 0; j < ggml_nelements(x); j++) {
+                                vec_x[j] = vec_x[j] + sigmas[i + 1] * vec_noise[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();
@@ -3588,9 +4007,10 @@ class StableDiffusionGGML {
        int64_t W = x->ne[0];
        int64_t H = x->ne[1];
        struct ggml_tensor* result = NULL;
+        struct ggml_cplan cplan;

        // 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;
@@ -3608,7 +4028,7 @@ class StableDiffusionGGML {
            ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);

            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments);
-            struct ggml_cplan cplan = ggml_graph_plan(vae_graph, n_threads);
+            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",
@@ -3632,7 +4052,10 @@ class StableDiffusionGGML {
            }

            struct ggml_tensor* moments = first_stage_model.encode(ctx, x);
-            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments);
+
+            struct ggml_cgraph* vae_graph = ggml_new_graph(ctx);
+            vae_graph->order = GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT;
+            ggml_build_forward_expand(vae_graph, moments);

            int64_t t0 = ggml_time_ms();
            ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads);
@@ -3710,6 +4133,7 @@ class StableDiffusionGGML {
        int64_t W = z->ne[0];
        int64_t H = z->ne[1];
        struct ggml_tensor* result_img = NULL;
+        struct ggml_cplan cplan;

        {
            float* vec = (float*)z->data;
@@ -3719,7 +4143,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;
@@ -3737,7 +4161,7 @@ class StableDiffusionGGML {
            ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);

            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img);
-            struct ggml_cplan cplan = ggml_graph_plan(vae_graph, n_threads);
+            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",
@@ -3761,7 +4185,10 @@ class StableDiffusionGGML {
            }

            struct ggml_tensor* img = first_stage_model.decode(ctx, z);
-            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img);
+
+            struct ggml_cgraph* vae_graph = ggml_new_graph(ctx);
+            vae_graph->order = GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT;
+            ggml_build_forward_expand(vae_graph, img);

            int64_t t0 = ggml_time_ms();
            ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads);
@@ -3815,8 +4242,8 @@ StableDiffusion::StableDiffusion(int n_threads,
                                               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,
@@ -3830,6 +4257,7 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
    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;
@@ -3865,7 +4293,7 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
    struct ggml_tensor* x_t = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, W, H, C, 1);
    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);
@@ -3921,7 +4349,7 @@ std::vector<uint8_t> StableDiffusion::img2img(const std::vector<uint8_t>& init_i
    }
    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;
--- a/stable-diffusion.h
+++ b/stable-diffusion.h
@@ -17,7 +17,22 @@ enum RNGType {
 };

 enum SampleMethod {
-    EULAR_A,
+    EULER_A,
+    EULER,
+    HEUN,
+    DPM2,
+    DPMPP2S_A,
+    DPMPP2M,
+    DPMPP2Mv2,
+    LCM,
+    N_SAMPLE_METHODS
+};
+
+enum Schedule {
+    DEFAULT,
+    DISCRETE,
+    KARRAS,
+    N_SCHEDULES
 };

 class StableDiffusionGGML;
@@ -31,7 +46,7 @@ class StableDiffusion {
                    bool vae_decode_only = false,
                    bool free_params_immediately = false,
                    RNGType rng_type = STD_DEFAULT_RNG);
-    bool load_from_file(const std::string& file_path);
+    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,
Author	SHA1	Message	Date
leejet	536f3af672	feat: add lcm sampler support This referenced an issue discussion of the stable-diffusion-webui at https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13952, which may not be too perfect.	2023-11-17 22:53:46 +08:00
leejet	3bf1665885	chore: clear the msvc compilation warning	2023-10-28 20:55:24 +08:00
leejet	3001c23f7d	perf: change ggml graph eval order to RIGHT_TO_LEFT to optimize memory usage	2023-10-28 20:19:15 +08:00
leejet	ed374983f3	fix: set eps of ggml_norm(LayerNorm) to 1e-5	2023-10-27 00:50:23 +08:00
leejet	4c96185fcc	fix: update ggml to avoid insufficient memory error on macOS	2023-10-24 22:04:00 +08:00
leejet	fbd18e1059	fix: avoid stack overflow on MSVC	2023-10-23 21:10:46 +08:00
leejet	09cab2a2ae	chore: set default BUILD_SHARED_LIBS to OFF	2023-10-22 14:59:03 +08:00
leejet	69e54ace14	sync: update ggml	2023-10-22 14:11:06 +08:00
Robert Bledsaw	29a56f2e98	docs: update README.md (#71 ) fixed type in curl command.	2023-10-22 13:03:32 +08:00
Urs Ganse	afec5051cf	feat: write generation parameter exif data into output png (#57 ) * Write generation parameter exif data into output pngs. This adds prompt, negative prompt (if nonempty) and other generation parameters to the output file as a tEXt PNG block, in the same format as AUTOMATIC1111 webui does. In order to keep everything free of external library dependencies, I have somewhat dirtily hacked this into the stb_image_write implementation. * Mention png text data in README.md, include "karras" in sampler text * add Steps/Model/RNG to parameter string --------- Co-authored-by: leejet <leejet714@gmail.com>	2023-09-18 21:09:15 +08:00
leejet	bd62138751	feat: adapt to more weight formats	2023-09-13 00:07:16 +08:00
Urs Ganse	3a25179d52	feat: add DPM2 and DPM++(2s) a samplers (#56 ) * Add DPM2 sampler. * Add DPM++ (2s) a sampler. * Update README.md with added samplers --------- Co-authored-by: leejet <leejet714@gmail.com>	2023-09-12 23:02:09 +08:00
Urs Ganse	968fbf02aa	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.	2023-09-09 00:02:07 +08:00
Urs Ganse	b6899e8fc2	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>	2023-09-08 23:47:28 +08:00
leejet	b85b236b13	feat: set default rng to cuda	2023-09-04 21:46:54 +08:00
leejet	34a118d407	fix: avoid coredump when steps == 1	2023-09-04 21:44:38 +08:00
leejet	f6ff06fcb7	fix: avoid coredump when generating large image	2023-09-04 21:37:46 +08:00
leejet	cf38e238d4	fix: width and height should be a multiple of 64	2023-09-04 20:49:52 +08:00
leejet	b247581782	fix: insufficient memory error on macOS	2023-09-04 03:50:42 +08:00
leejet	bb3f19cb40	fix: increase ctx_size	2023-09-04 03:45:43 +08:00