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.

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)
 
 

README.md

@@ -22,9 +22,13 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
     - `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
@@ -69,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
@@ -77,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
     ```
@@ -129,7 +134,7 @@ 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)
-  --sampling-method {euler, euler_a, heun, dpm++2m, dpm++2mv2}
+  --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)
@@ -193,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)

examples/main.cpp

@@ -77,14 +77,19 @@ const char* sample_method_str[] = {
     "euler_a",
     "euler",
     "heun",
+    "dpm2",
+    "dpm++2s_a",
     "dpm++2m",
-    "dpm++2mv2"};
+    "dpm++2mv2",
+    "lcm",
+};
 
 // Names of the sigma schedule overrides, same order as Schedule in stable-diffusion.h
 const char* schedule_str[] = {
     "default",
     "discrete",
-    "karras"};
+    "karras"
+};
 
 struct Option {
     int n_threads = -1;
@@ -144,7 +149,7 @@ 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("  --sampling-method {euler, euler_a, heun, dpm++2m, dpm++2mv2}\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: cuda)\n");
@@ -363,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);
@@ -435,7 +452,24 @@ 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;

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);

models/.gitignore

@@ -1,4 +1,5 @@
 *.bin
 *.ckpt
 *.safetensor
+*.safetensors
 *.log

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):

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)

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);
         }

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;
     }

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),
@@ -2705,13 +2707,13 @@ struct DiscreteSchedule : SigmaSchedule {
         if (n == 0) {
             return result;
         } else if (n == 1) {
-            result.push_back(t_to_sigma(t_max));
+            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 (int i = 0; i < n; ++i) {
+        for (uint32_t i = 0; i < n; ++i) {
             float t = t_max - step * i;
             result.push_back(t_to_sigma(t));
         }
@@ -2724,17 +2726,17 @@ 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.;
+        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. / rho));
-        float max_inv_rho = pow(sigma_max, (1. / rho));
-        for (int i = 0; i < n; i++) {
+        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.) * (min_inv_rho - max_inv_rho), rho);
+            result[i] = pow(max_inv_rho + (float)i / ((float)n - 1.f) * (min_inv_rho - max_inv_rho), rho);
         }
         result[n] = 0.;
         return result;
@@ -3155,6 +3157,8 @@ 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);
 
+        struct ggml_cplan cplan;
+
         size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
@@ -3179,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",
@@ -3212,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);
@@ -3257,6 +3261,7 @@ class StableDiffusionGGML {
                                                             true);
         std::vector<int>& tokens = tokens_and_weights.first;
         std::vector<float>& weights = tokens_and_weights.second;
+        struct ggml_cplan cplan;
         size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
@@ -3278,8 +3283,8 @@ class StableDiffusionGGML {
 
             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_cplan cplan = ggml_graph_plan(&cond_graph, n_threads);
+            struct ggml_cgraph* cond_graph = ggml_build_forward_ctx(ctx, hidden_states);
+            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);
@@ -3390,6 +3395,7 @@ 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 = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
@@ -3417,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",
@@ -3449,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);
@@ -3706,6 +3714,137 @@ class StableDiffusionGGML {
                     }
                 }
             } 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");
@@ -3723,7 +3862,7 @@ class StableDiffusionGGML {
                     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 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;
@@ -3737,7 +3876,7 @@ class StableDiffusionGGML {
                         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];
+                            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;
                         }
                     }
@@ -3771,7 +3910,7 @@ class StableDiffusionGGML {
 
                     if (i == 0 || sigmas[i + 1] == 0) {
                         // Simpler step for the edge cases
-                        float b = exp(-h) - 1.;
+                        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];
                         }
@@ -3780,10 +3919,10 @@ class StableDiffusionGGML {
                         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.;
+                        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. + 1. / (2. * r)) * vec_denoised[j] - (1. / (2. * r)) * vec_old_denoised[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;
                         }
                     }
@@ -3794,6 +3933,44 @@ class StableDiffusionGGML {
                     }
                 }
             } 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);
@@ -3830,6 +4007,7 @@ 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 = 10 * 1024 * 1024;  // 10MB
@@ -3850,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",
@@ -3874,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);
@@ -3952,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;
@@ -3979,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",
@@ -4003,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);

stable-diffusion.h

@@ -20,8 +20,11 @@ enum SampleMethod {
     EULER_A,
     EULER,
     HEUN,
+    DPM2,
+    DPMPP2S_A,
     DPMPP2M,
     DPMPP2Mv2,
+    LCM,
     N_SAMPLE_METHODS
 };