style: format code

* add controlnet to pipeline

* add cli params

* control strength cli param

* cli param keep controlnet in cpu

* add Textual Inversion

* add canny preprocessor

* refactor: change ggml_type_sizef to ggml_row_size

* process hint once time

* ignore the embedding name case

---------

Co-authored-by: leejet <leejet714@gmail.com>

.clang-format

@@ -3,7 +3,6 @@ UseTab: Never
 IndentWidth: 4
 TabWidth: 4
 AllowShortIfStatementsOnASingleLine: false
-IndentCaseLabels: false
 ColumnLimit: 0
 AccessModifierOffset: -4
 NamespaceIndentation: All

.gitignore

@@ -10,5 +10,4 @@ test/
 *.gguf
 output*.png
 models*
-!taesd-model.gguf
 *.log

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "ggml"]
     path = ggml
-	url = https://github.com/leejet/ggml.git
+	url = https://github.com/ggerganov/ggml.git

CMakeLists.txt

@@ -25,9 +25,9 @@ endif()
 #option(SD_BUILD_TESTS                "sd: build tests"    ${SD_STANDALONE})
 option(SD_BUILD_EXAMPLES             "sd: build examples" ${SD_STANDALONE})
 option(SD_CUBLAS                     "sd: cuda backend" OFF)
+option(SD_HIPBLAS                    "sd: rocm backend" OFF)
 option(SD_METAL                      "sd: metal backend" OFF)
 option(SD_FLASH_ATTN                 "sd: use flash attention for x4 less memory usage" OFF)
-option(SD_FAST_SOFTMAX               "sd: x1.5 faster softmax, indeterministic (sometimes, same seed don't generate same image), cuda only" OFF)
 option(BUILD_SHARED_LIBS             "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
@@ -35,9 +35,6 @@ if(SD_CUBLAS)
 	message("Use CUBLAS as backend stable-diffusion")
     set(GGML_CUBLAS ON)
     add_definitions(-DSD_USE_CUBLAS)
-    if(SD_FAST_SOFTMAX)
-        set(GGML_CUDA_FAST_SOFTMAX ON)
-    endif()
 endif()
 
 if(SD_METAL)
@@ -46,22 +43,45 @@ if(SD_METAL)
     add_definitions(-DSD_USE_METAL)
 endif()
 
+if (SD_HIPBLAS)
+    message("Use HIPBLAS as backend stable-diffusion")
+    set(GGML_HIPBLAS ON)
+    add_definitions(-DSD_USE_CUBLAS)
+    if(SD_FAST_SOFTMAX)
+        set(GGML_CUDA_FAST_SOFTMAX ON)
+    endif()
+endif ()
+
 if(SD_FLASH_ATTN)
     message("Use Flash Attention for memory optimization")
     add_definitions(-DSD_USE_FLASH_ATTENTION)
 endif()
 
+set(SD_LIB stable-diffusion)
+
+add_library(${SD_LIB} stable-diffusion.h stable-diffusion.cpp model.h model.cpp util.h util.cpp upscaler.cpp
+             ggml_extend.hpp clip.hpp common.hpp unet.hpp tae.hpp esrgan.hpp lora.hpp denoiser.hpp rng.hpp rng_philox.hpp)
+
+if(BUILD_SHARED_LIBS)
+    message("Build shared library")
+    add_definitions(-DSD_BUILD_SHARED_LIB)
+    target_compile_definitions(${SD_LIB} PRIVATE -DSD_BUILD_DLL)
+    set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+else()
+    message("Build static library")
+endif()
+
 
 set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
+
+# see https://github.com/ggerganov/ggml/pull/682
+add_definitions(-DGGML_MAX_NAME=128)
+
 # deps
 add_subdirectory(ggml)
 
 add_subdirectory(thirdparty)
 
-set(SD_LIB stable-diffusion)
-
-add_library(${SD_LIB} stable-diffusion.h stable-diffusion.cpp model.h model.cpp util.h util.cpp upscaler.cpp
-             ggml_extend.hpp clip.hpp common.hpp unet.hpp tae.hpp esrgan.hpp lora.hpp denoiser.hpp rng.hpp rng_philox.hpp)
 target_link_libraries(${SD_LIB} PUBLIC ggml zip)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
 target_compile_features(${SD_LIB} PUBLIC cxx_std_11)

README.md

@@ -31,6 +31,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - Faster and memory efficient latent decoding with [TAESD](https://github.com/madebyollin/taesd)
 - Upscale images generated with [ESRGAN](https://github.com/xinntao/Real-ESRGAN)
 - VAE tiling processing for reduce memory usage
+- Control Net support with SD 1.5
 - Sampling method
     - `Euler A`
     - `Euler`
@@ -53,9 +54,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - [ ] More sampling methods
 - [ ] Make inference faster
     - The current implementation of ggml_conv_2d is slow and has high memory usage
-    - Implement Winograd Convolution 2D for 3x3 kernel filtering
 - [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
-- [ ] Implement Textual Inversion (embeddings)
 - [ ] Implement Inpainting support
 - [ ] k-quants support
 
@@ -117,6 +116,17 @@ cmake .. -DSD_CUBLAS=ON
 cmake --build . --config Release
 ```
 
+##### Using HipBLAS
+This provides BLAS acceleration using the ROCm cores of your AMD GPU. Make sure to have the ROCm toolkit installed.
+
+Windows User Refer to [docs/hipBLAS_on_Windows.md](docs%2FhipBLAS_on_Windows.md) for a comprehensive guide.
+
+```
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
+cmake --build . --config Release
+```
+
+
 ##### Using Metal
 
 Using Metal makes the computation run on the GPU. Currently, there are some issues with Metal when performing operations on very large matrices, making it highly inefficient at the moment. Performance improvements are expected in the near future.
@@ -126,7 +136,7 @@ cmake .. -DSD_METAL=ON
 cmake --build . --config Release
 ```
 
-### Using Flash Attention
+##### Using Flash Attention
 
 Enabling flash attention reduces memory usage by at least 400 MB. At the moment, it is not supported when CUBLAS is enabled because the kernel implementation is missing.
 
@@ -142,22 +152,26 @@ usage: ./bin/sd [arguments]
 
 arguments:
   -h, --help                         show this help message and exit
-  -M, --mode [txt2img or img2img]    generation mode (default: txt2img)
+  -M, --mode [MODEL]                 run mode (txt2img or img2img or convert, default: txt2img)
   -t, --threads N                    number of threads to use during computation (default: -1).
                                      If threads <= 0, then threads will be set to the number of CPU physical cores
   -m, --model [MODEL]                path to model
   --vae [VAE]                        path to vae
   --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
+  --control-net [CONTROL_PATH]       path to control net model
+  --embd-dir [EMBEDDING_PATH]        path to embeddings.
   --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.
   --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
                                      If not specified, the default is the type of the weight file.
   --lora-model-dir [DIR]             lora model directory
   -i, --init-img [IMAGE]             path to the input image, required by img2img
+  --control-image [IMAGE]            path to image condition, control net
   -o, --output OUTPUT                path to write result image to (default: ./output.png)
   -p, --prompt [PROMPT]              the prompt to render
   -n, --negative-prompt PROMPT       the negative prompt (default: "")
   --cfg-scale SCALE                  unconditional guidance scale: (default: 7.0)
   --strength STRENGTH                strength for noising/unnoising (default: 0.75)
+  --control-strength STRENGTH        strength to apply Control Net (default: 0.9)
                                      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)
@@ -168,8 +182,10 @@ arguments:
   -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
   -b, --batch-count COUNT            number of images to generate.
   --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)
-  --clip-skip N                      number of layers to skip of clip model (default: 0)
+  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)
+                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
   --vae-tiling                       process vae in tiles to reduce memory usage
+  --control-net-cpu                  keep controlnet in cpu (for low vram)
   -v, --verbose                      print extra info
 ```
 
@@ -183,6 +199,16 @@ You can specify the model weight type using the `--type` parameter. The weights
 - `q5_0` or `q5_1` for 5-bit integer quantization
 - `q4_0` or `q4_1` for 4-bit integer quantization
 
+#### Convert to GGUF
+
+You can also convert weights in the formats `ckpt/safetensors/diffusers` to gguf and perform quantization in advance, avoiding the need for quantization every time you load them.
+
+For example:
+
+```sh
+./bin/sd -M convert -m ../models/v1-5-pruned-emaonly.safetensors -o  ../models/v1-5-pruned-emaonly.q8_0.gguf -v --type q8_0
+```
+
 #### txt2img example
 
 ```sh
@@ -240,7 +266,7 @@ Here's a simple example:
 | ----  |----    |
 | ![](./assets/without_lcm.png) |![](./assets/with_lcm.png)  |
 
-## Using TAESD to faster decoding
+#### Using TAESD to faster decoding
 
 You can use TAESD to accelerate the decoding of latent images by following these steps:
 
@@ -258,7 +284,7 @@ curl -L -O https://huggingface.co/madebyollin/taesd/blob/main/diffusion_pytorch_
 sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --taesd ../models/diffusion_pytorch_model.safetensors
 ```
 
-## Using ESRGAN to upscale results
+#### Using ESRGAN to upscale results
 
 You can use ESRGAN to upscale the generated images. At the moment, only the [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) model is supported. Support for more models of this architecture will be added soon.
 

clip.hpp

@@ -2,6 +2,7 @@
 #define __CLIP_HPP__
 
 #include "ggml_extend.hpp"
+#include "model.h"
 
 /*================================================== CLIPTokenizer ===================================================*/
 
@@ -67,6 +68,9 @@ std::vector<std::pair<int, std::u32string>> bytes_to_unicode() {
 }
 
 // Ref: https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py
+
+typedef std::function<bool(std::string&, std::vector<int32_t>&)> on_new_token_cb_t;
+
 class CLIPTokenizer {
 private:
     SDVersion version = VERSION_1_x;
@@ -234,8 +238,11 @@ public:
         return result;
     }
 
-    std::vector<int> tokenize(std::string text, size_t max_length = 0, bool padding = false) {
-        std::vector<int32_t> tokens = encode(text);
+    std::vector<int> tokenize(std::string text,
+                              on_new_token_cb_t on_new_token_cb,
+                              size_t max_length = 0,
+                              bool padding      = false) {
+        std::vector<int32_t> tokens = encode(text, on_new_token_cb);
         tokens.insert(tokens.begin(), BOS_TOKEN_ID);
         if (max_length > 0) {
             if (tokens.size() > max_length - 1) {
@@ -255,7 +262,7 @@ public:
         return tokens;
     }
 
-    std::vector<int> encode(std::string text) {
+    std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb) {
         std::string original_text = text;
         std::vector<int32_t> bpe_tokens;
         text = whitespace_clean(text);
@@ -268,6 +275,10 @@ public:
         std::string str = text;
         std::vector<std::string> token_strs;
         while (std::regex_search(str, matches, pat)) {
+            bool skip = on_new_token_cb(str, bpe_tokens);
+            if (skip) {
+                continue;
+            }
             for (auto& token : matches) {
                 std::string token_str = token.str();
                 std::u32string utf32_token;
@@ -443,17 +454,14 @@ struct ResidualAttentionBlock {
     struct ggml_tensor* ln2_w;  // [hidden_size, ]
     struct ggml_tensor* ln2_b;  // [hidden_size, ]
 
-    struct ggml_tensor* attn_scale;  // [hidden_size, ]
-
     size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
-        mem_size += 4 * hidden_size * hidden_size * ggml_type_sizef(wtype);        // q_w/k_w/v_w/out_w
-        mem_size += 8 * hidden_size * ggml_type_sizef(GGML_TYPE_F32);              // q_b/k_b/v_b/out_b/ln1_w/ln1_b/ln2_w/ln2_b
-        mem_size += 2 * hidden_size * intermediate_size * ggml_type_sizef(wtype);  // fc1_w/fc2_w
-        mem_size += intermediate_size * ggml_type_sizef(GGML_TYPE_F32);            // fc1_b
-        mem_size += hidden_size * ggml_type_sizef(GGML_TYPE_F32);                  // fc2_b
-        mem_size += ggml_type_sizef(GGML_TYPE_F32);                                // attn_scale
-        return static_cast<size_t>(mem_size);
+        size_t mem_size = 0;
+        mem_size += 4 * ggml_row_size(wtype, hidden_size * hidden_size);        // q_w/k_w/v_w/out_w
+        mem_size += 8 * ggml_row_size(GGML_TYPE_F32, hidden_size);              // q_b/k_b/v_b/out_b/ln1_w/ln1_b/ln2_w/ln2_b
+        mem_size += 2 * ggml_row_size(wtype, hidden_size * intermediate_size);  // fc1_w/fc2_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, intermediate_size);            // fc1_b
+        mem_size += ggml_row_size(GGML_TYPE_F32, hidden_size);                  // fc2_b
+        return mem_size;
     }
 
     void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
@@ -478,11 +486,6 @@ struct ResidualAttentionBlock {
 
         ln2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);
         ln2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);
-
-        attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, attn_scale);
-        float scale = 1.0f / sqrt((float)d_model);
-        ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));
     }
 
     void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
@@ -521,7 +524,7 @@ struct ResidualAttentionBlock {
         // self-attention
         {
             struct ggml_tensor* q = ggml_nn_linear(ctx, x, q_w, q_b);
-            q                     = ggml_scale_inplace(ctx, q, attn_scale);
+            q                     = ggml_scale_inplace(ctx, q, 1.0f / sqrt((float)d_model));
             q                     = ggml_reshape_4d(ctx, q, d_model, n_head, n_token, N);   // [N, n_token, n_head, d_model]
             q                     = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));       // [N, n_head, n_token, d_model]
             q                     = ggml_reshape_3d(ctx, q, d_model, n_token, n_head * N);  // [N * n_head, n_token, d_model]
@@ -604,6 +607,7 @@ struct CLIPTextModel {
     struct ggml_tensor* position_ids;
     struct ggml_tensor* token_embed_weight;
     struct ggml_tensor* position_embed_weight;
+    struct ggml_tensor* token_embed_custom;
 
     // transformer
     std::vector<ResidualAttentionBlock> resblocks;
@@ -611,6 +615,9 @@ struct CLIPTextModel {
     struct ggml_tensor* final_ln_b;
 
     struct ggml_tensor* text_projection;
+    std::string embd_dir;
+    int32_t num_custom_embeddings = 0;
+    std::vector<std::string> readed_embeddings;
 
     CLIPTextModel(CLIPVersion version = OPENAI_CLIP_VIT_L_14,
                   int clip_skip       = -1,
@@ -649,18 +656,21 @@ struct CLIPTextModel {
     }
 
     size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
-        mem_size += hidden_size * max_position_embeddings * ggml_type_sizef(GGML_TYPE_I32);  // position_ids
-        mem_size += hidden_size * vocab_size * ggml_type_sizef(wtype);                       // token_embed_weight
-        mem_size += hidden_size * max_position_embeddings * ggml_type_sizef(wtype);          // position_embed_weight
+        size_t mem_size = 0;
+        mem_size += ggml_row_size(GGML_TYPE_I32, hidden_size * max_position_embeddings);  // position_ids
+        mem_size += ggml_row_size(wtype, hidden_size * vocab_size);                       // token_embed_weight
+        mem_size += ggml_row_size(wtype, hidden_size * max_position_embeddings);          // position_embed_weight
+        if (version == OPENAI_CLIP_VIT_L_14) {
+            mem_size += ggml_row_size(wtype, hidden_size * max_position_embeddings);  // token_embed_custom
+        }
         for (int i = 0; i < num_hidden_layers; i++) {
             mem_size += resblocks[i].calculate_mem_size(wtype);
         }
-        mem_size += 2 * hidden_size * ggml_type_sizef(GGML_TYPE_F32);  // final_ln_w/b
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, hidden_size);  // final_ln_w/b
         if (version == OPEN_CLIP_VIT_BIGG_14) {
-            mem_size += hidden_size * projection_dim * ggml_type_sizef(GGML_TYPE_F32);  // text_projection
+            mem_size += ggml_row_size(GGML_TYPE_F32, hidden_size * projection_dim);  // text_projection
         }
-        return static_cast<size_t>(mem_size);
+        return mem_size;
     }
 
     void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
@@ -677,14 +687,48 @@ struct CLIPTextModel {
         }
     }
 
-    struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor* input_ids, size_t max_token_idx = 0, bool return_pooled = false) {
+    bool load_embedding(std::string embd_name, std::string embd_path, std::vector<int32_t>& bpe_tokens) {
+        // the order matters
+        ModelLoader model_loader;
+        if (!model_loader.init_from_file(embd_path)) {
+            LOG_ERROR("embedding '%s' failed", embd_name.c_str());
+            return false;
+        }
+        struct ggml_init_params params;
+        params.mem_size               = 32 * 1024;  // max for custom embeddings 32 KB
+        params.mem_buffer             = NULL;
+        params.no_alloc               = false;
+        struct ggml_context* embd_ctx = ggml_init(params);
+        struct ggml_tensor* embd      = NULL;
+        auto on_load                  = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
+            if (tensor_storage.ne[0] != hidden_size) {
+                LOG_DEBUG("embedding wrong hidden size, got %i, expected %i", tensor_storage.ne[0], hidden_size);
+                return false;
+            }
+            embd        = ggml_new_tensor_2d(embd_ctx, token_embed_weight->type, hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
+            *dst_tensor = embd;
+            return true;
+        };
+        model_loader.load_tensors(on_load, NULL);
+        ggml_backend_tensor_set(token_embed_custom, embd->data, num_custom_embeddings * hidden_size * ggml_type_size(token_embed_custom->type), ggml_nbytes(embd));
+        readed_embeddings.push_back(embd_name);
+        for (int i = 0; i < embd->ne[1]; i++) {
+            bpe_tokens.push_back(vocab_size + num_custom_embeddings);
+            // LOG_DEBUG("new custom token: %i", vocab_size + num_custom_embeddings);
+            num_custom_embeddings++;
+        }
+        LOG_DEBUG("embedding '%s' applied, custom embeddings: %i", embd_name.c_str(), num_custom_embeddings);
+        return true;
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor* input_ids, struct ggml_tensor* tkn_embeddings, uint32_t max_token_idx = 0, bool return_pooled = false) {
         // input_ids: [N, n_token]
         GGML_ASSERT(input_ids->ne[0] <= position_ids->ne[0]);
 
         // token_embedding + position_embedding
         struct ggml_tensor* x;
         x = ggml_add(ctx0,
-                     ggml_get_rows(ctx0, token_embed_weight, input_ids),
+                     ggml_get_rows(ctx0, tkn_embeddings == NULL ? token_embed_weight : tkn_embeddings, input_ids),
                      ggml_get_rows(ctx0,
                                    position_embed_weight,
                                    ggml_view_1d(ctx0, position_ids, input_ids->ne[0], 0)));  // [N, n_token, hidden_size]
@@ -730,6 +774,10 @@ struct CLIPTextModel {
 
         final_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);
 
+        if (version == OPENAI_CLIP_VIT_L_14) {
+            token_embed_custom = ggml_new_tensor_2d(ctx, wtype, hidden_size, max_position_embeddings);
+        }
+
         if (version == OPEN_CLIP_VIT_BIGG_14) {
             text_projection = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, projection_dim, hidden_size);
         }
@@ -812,11 +860,11 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
         }
     }
 
-    struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor* input_ids, struct ggml_tensor* input_ids2, size_t max_token_idx = 0, bool return_pooled = false) {
+    struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor* input_ids, struct ggml_tensor* input_ids2, struct ggml_tensor* embeddings, size_t max_token_idx = 0, bool return_pooled = false) {
         if (return_pooled) {
-            return text_model2.forward(ctx0, input_ids2, max_token_idx, return_pooled);
+            return text_model2.forward(ctx0, input_ids2, NULL, max_token_idx, return_pooled);
         }
-        auto hidden_states = text_model.forward(ctx0, input_ids);  // [N, n_token, hidden_size]
+        auto hidden_states = text_model.forward(ctx0, input_ids, embeddings);  // [N, n_token, hidden_size]
         // LOG_DEBUG("hidden_states: %d %d %d %d %d", hidden_states->n_dims, hidden_states->ne[0], hidden_states->ne[1], hidden_states->ne[2], hidden_states->ne[3]);
         if (version == VERSION_XL) {
             hidden_states = ggml_reshape_4d(ctx0,
@@ -827,13 +875,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
                                             hidden_states->ne[3]);
             hidden_states = ggml_cont(ctx0, ggml_permute(ctx0, hidden_states, 2, 0, 1, 3));
 
-            auto hidden_states2 = text_model2.forward(ctx0, input_ids2);  // [N, n_token, hidden_size2]
+            auto hidden_states2 = text_model2.forward(ctx0, input_ids2, NULL);  // [N, n_token, hidden_size2]
             hidden_states2      = ggml_reshape_4d(ctx0,
-                                             hidden_states2,
-                                             hidden_states2->ne[0],
-                                             hidden_states2->ne[1],
-                                             hidden_states2->ne[2],
-                                             hidden_states2->ne[3]);
+                                                  hidden_states2,
+                                                  hidden_states2->ne[0],
+                                                  hidden_states2->ne[1],
+                                                  hidden_states2->ne[2],
+                                                  hidden_states2->ne[3]);
             hidden_states2      = ggml_cont(ctx0, ggml_permute(ctx0, hidden_states2, 2, 0, 1, 3));
 
             hidden_states = ggml_concat(ctx0, hidden_states, hidden_states2);  // [N, n_token, hidden_size + hidden_size2]
@@ -864,12 +912,36 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
             LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
         }
 
+        auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
+            size_t word_end       = str.find(",");
+            std::string embd_name = word_end == std::string::npos ? str : str.substr(0, word_end);
+            embd_name             = trim(embd_name);
+            std::string embd_path = get_full_path(text_model.embd_dir, embd_name + ".pt");
+            if (embd_path.size() == 0) {
+                embd_path = get_full_path(text_model.embd_dir, embd_name + ".ckpt");
+            }
+            if (embd_path.size() == 0) {
+                embd_path = get_full_path(text_model.embd_dir, embd_name + ".safetensors");
+            }
+            if (embd_path.size() > 0) {
+                if (text_model.load_embedding(embd_name, embd_path, bpe_tokens)) {
+                    if (word_end != std::string::npos) {
+                        str = str.substr(word_end);
+                    } else {
+                        str = "";
+                    }
+                    return true;
+                }
+            }
+            return false;
+        };
+
         std::vector<int> tokens;
         std::vector<float> weights;
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = tokenizer.encode(curr_text);
+            std::vector<int> curr_tokens = tokenizer.encode(curr_text, on_new_token_cb);
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
@@ -958,7 +1030,26 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
             }
         }
 
-        struct ggml_tensor* hidden_states = forward(ctx0, input_ids, input_ids2, max_token_idx, return_pooled);
+        struct ggml_tensor* embeddings = NULL;
+
+        if (text_model.num_custom_embeddings > 0 && version != VERSION_XL) {
+            embeddings = ggml_new_tensor_2d(ctx0, wtype, text_model.hidden_size, text_model.vocab_size + text_model.num_custom_embeddings /* custom placeholder */);
+            ggml_allocr_alloc(allocr, embeddings);
+            if (!ggml_allocr_is_measure(allocr)) {
+                // really bad, there is memory inflexibility (this is for host<->device memory conflicts)
+                void* freeze_data = malloc(ggml_nbytes(text_model.token_embed_weight));
+                ggml_backend_tensor_get_and_sync(backend, text_model.token_embed_weight, freeze_data, 0, ggml_nbytes(text_model.token_embed_weight));
+                ggml_backend_tensor_set(embeddings, freeze_data, 0, ggml_nbytes(text_model.token_embed_weight));
+                free(freeze_data);
+                // concatenate custom embeddings
+                void* custom_data = malloc(ggml_nbytes(text_model.token_embed_custom));
+                ggml_backend_tensor_get_and_sync(backend, text_model.token_embed_custom, custom_data, 0, ggml_nbytes(text_model.token_embed_custom));
+                ggml_backend_tensor_set(embeddings, custom_data, ggml_nbytes(text_model.token_embed_weight), text_model.num_custom_embeddings * text_model.hidden_size * ggml_type_size(wtype));
+                free(custom_data);
+            }
+        }
+
+        struct ggml_tensor* hidden_states = forward(ctx0, input_ids, input_ids2, embeddings, max_token_idx, return_pooled);
 
         ggml_build_forward_expand(gf, hidden_states);
         ggml_free(ctx0);

common.hpp

@@ -15,10 +15,10 @@ struct DownSample {
     bool vae_downsample = false;
 
     size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
-        mem_size += out_channels * channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // op_w
-        mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32);                     // op_b
-        return static_cast<size_t>(mem_size);
+        size_t mem_size = 0;
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * channels * 3 * 3);  // op_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                     // op_b
+        return mem_size;
     }
 
     void init_params(struct ggml_context* ctx, ggml_type wtype) {
@@ -59,10 +59,10 @@ struct UpSample {
     struct ggml_tensor* conv_b;  // [out_channels,]
 
     size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
-        mem_size += out_channels * channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // op_w
-        mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32);                     // op_b
-        return static_cast<size_t>(mem_size);
+        size_t mem_size = 0;
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * channels * 3 * 3);  // op_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                     // op_b
+        return mem_size;
     }
 
     void init_params(struct ggml_context* ctx, ggml_type wtype) {
@@ -83,4 +83,461 @@ struct UpSample {
     }
 };
 
+struct ResBlock {
+    // network hparams
+    int channels;      // model_channels * (1, 1, 1, 2, 2, 4, 4, 4)
+    int emb_channels;  // time_embed_dim
+    int out_channels;  // mult * model_channels
+
+    // network params
+    // in_layers
+    struct ggml_tensor* in_layer_0_w;  // [channels, ]
+    struct ggml_tensor* in_layer_0_b;  // [channels, ]
+    // in_layer_1 is nn.SILU()
+    struct ggml_tensor* in_layer_2_w;  // [out_channels, channels, 3, 3]
+    struct ggml_tensor* in_layer_2_b;  // [out_channels, ]
+
+    // emb_layers
+    // emb_layer_0 is nn.SILU()
+    struct ggml_tensor* emb_layer_1_w;  // [out_channels, emb_channels]
+    struct ggml_tensor* emb_layer_1_b;  // [out_channels, ]
+
+    // out_layers
+    struct ggml_tensor* out_layer_0_w;  // [out_channels, ]
+    struct ggml_tensor* out_layer_0_b;  // [out_channels, ]
+    // out_layer_1 is nn.SILU()
+    // out_layer_2 is nn.Dropout(), p = 0 for inference
+    struct ggml_tensor* out_layer_3_w;  // [out_channels, out_channels, 3, 3]
+    struct ggml_tensor* out_layer_3_b;  // [out_channels, ]
+
+    // skip connection, only if out_channels != channels
+    struct ggml_tensor* skip_w;  // [out_channels, channels, 1, 1]
+    struct ggml_tensor* skip_b;  // [out_channels, ]
+
+    size_t calculate_mem_size(ggml_type wtype) {
+        size_t mem_size = 0;
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, channels);                         // in_layer_0_w/b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * channels * 3 * 3);      // in_layer_2_w
+        mem_size += 5 * ggml_row_size(GGML_TYPE_F32, out_channels);                     // in_layer_2_b/emb_layer_1_b/out_layer_0_w/out_layer_0_b/out_layer_3_b
+        mem_size += ggml_row_size(wtype, out_channels * emb_channels);                  // emb_layer_1_w
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * out_channels * 3 * 3);  // out_layer_3_w
+
+        if (out_channels != channels) {
+            mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * channels * 1 * 1);  // skip_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                     // skip_b
+        }
+        return mem_size;
+    }
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        in_layer_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+        in_layer_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+        in_layer_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels);
+        in_layer_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+
+        emb_layer_1_w = ggml_new_tensor_2d(ctx, wtype, emb_channels, out_channels);
+        emb_layer_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+
+        out_layer_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        out_layer_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        out_layer_3_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels);
+        out_layer_3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+
+        if (out_channels != channels) {
+            skip_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, channels, out_channels);
+            skip_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "in_layers.0.weight"] = in_layer_0_w;
+        tensors[prefix + "in_layers.0.bias"]   = in_layer_0_b;
+        tensors[prefix + "in_layers.2.weight"] = in_layer_2_w;
+        tensors[prefix + "in_layers.2.bias"]   = in_layer_2_b;
+
+        tensors[prefix + "emb_layers.1.weight"] = emb_layer_1_w;
+        tensors[prefix + "emb_layers.1.bias"]   = emb_layer_1_b;
+
+        tensors[prefix + "out_layers.0.weight"] = out_layer_0_w;
+        tensors[prefix + "out_layers.0.bias"]   = out_layer_0_b;
+        tensors[prefix + "out_layers.3.weight"] = out_layer_3_w;
+        tensors[prefix + "out_layers.3.bias"]   = out_layer_3_b;
+
+        if (out_channels != channels) {
+            tensors[prefix + "skip_connection.weight"] = skip_w;
+            tensors[prefix + "skip_connection.bias"]   = skip_b;
+        }
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* emb) {
+        // x: [N, channels, h, w]
+        // emb: [N, emb_channels]
+
+        // in_layers
+        auto h = ggml_nn_group_norm(ctx, x, in_layer_0_w, in_layer_0_b);
+        h      = ggml_silu_inplace(ctx, h);
+        h      = ggml_nn_conv_2d(ctx, h, in_layer_2_w, in_layer_2_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
+
+        // emb_layers
+        auto emb_out = ggml_silu(ctx, emb);
+        emb_out      = ggml_nn_linear(ctx, emb_out, emb_layer_1_w, emb_layer_1_b);           // [N, out_channels]
+        emb_out      = ggml_reshape_4d(ctx, emb_out, 1, 1, emb_out->ne[0], emb_out->ne[1]);  // [N, out_channels, 1, 1]
+
+        // out_layers
+        h = ggml_add(ctx, h, emb_out);
+        h = ggml_nn_group_norm(ctx, h, out_layer_0_w, out_layer_0_b);
+        h = ggml_silu_inplace(ctx, h);
+
+        // dropout, skip for inference
+
+        h = ggml_nn_conv_2d(ctx, h, out_layer_3_w, out_layer_3_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
+
+        // skip connection
+        if (out_channels != channels) {
+            x = ggml_nn_conv_2d(ctx, x, skip_w, skip_b);  // [N, out_channels, h, w]
+        }
+
+        h = ggml_add(ctx, h, x);
+        return h;  // [N, out_channels, h, w]
+    }
+};
+
+struct SpatialTransformer {
+    int in_channels;        // mult * model_channels
+    int n_head;             // num_heads
+    int d_head;             // in_channels // n_heads
+    int depth       = 1;    // 1
+    int context_dim = 768;  // hidden_size, 1024 for VERSION_2_x
+
+    // group norm
+    struct ggml_tensor* norm_w;  // [in_channels,]
+    struct ggml_tensor* norm_b;  // [in_channels,]
+
+    // proj_in
+    struct ggml_tensor* proj_in_w;  // [in_channels, in_channels, 1, 1]
+    struct ggml_tensor* proj_in_b;  // [in_channels,]
+
+    // transformer
+    struct Transformer {
+        // layer norm 1
+        struct ggml_tensor* norm1_w;  // [in_channels, ]
+        struct ggml_tensor* norm1_b;  // [in_channels, ]
+
+        // attn1
+        struct ggml_tensor* attn1_q_w;  // [in_channels, in_channels]
+        struct ggml_tensor* attn1_k_w;  // [in_channels, in_channels]
+        struct ggml_tensor* attn1_v_w;  // [in_channels, in_channels]
+
+        struct ggml_tensor* attn1_out_w;  // [in_channels, in_channels]
+        struct ggml_tensor* attn1_out_b;  // [in_channels, ]
+
+        // layer norm 2
+        struct ggml_tensor* norm2_w;  // [in_channels, ]
+        struct ggml_tensor* norm2_b;  // [in_channels, ]
+
+        // attn2
+        struct ggml_tensor* attn2_q_w;  // [in_channels, in_channels]
+        struct ggml_tensor* attn2_k_w;  // [in_channels, context_dim]
+        struct ggml_tensor* attn2_v_w;  // [in_channels, context_dim]
+
+        struct ggml_tensor* attn2_out_w;  // [in_channels, in_channels]
+        struct ggml_tensor* attn2_out_b;  // [in_channels, ]
+
+        // layer norm 3
+        struct ggml_tensor* norm3_w;  // [in_channels, ]
+        struct ggml_tensor* norm3_b;  // [in_channels, ]
+
+        // ff
+        struct ggml_tensor* ff_0_proj_w;  // [in_channels * 4 * 2, in_channels]
+        struct ggml_tensor* ff_0_proj_b;  // [in_channels * 4 * 2]
+
+        struct ggml_tensor* ff_2_w;  // [in_channels, in_channels * 4]
+        struct ggml_tensor* ff_2_b;  // [in_channels,]
+    };
+
+    std::vector<Transformer> transformers;
+
+    // proj_out
+    struct ggml_tensor* proj_out_w;  // [in_channels, in_channels, 1, 1]
+    struct ggml_tensor* proj_out_b;  // [in_channels,]
+
+    SpatialTransformer(int depth = 1)
+        : depth(depth) {
+        transformers.resize(depth);
+    }
+
+    int get_num_tensors() {
+        return depth * 20 + 7;
+    }
+
+    size_t calculate_mem_size(ggml_type wtype) {
+        size_t mem_size = 0;
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, in_channels);                        // norm_w/norm_b
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F16, in_channels * in_channels * 1 * 1);  // proj_in_w/proj_out_w
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, in_channels);                        // proj_in_b/proj_out_b
+
+        // transformer
+        for (auto& transformer : transformers) {
+            mem_size += 6 * ggml_row_size(GGML_TYPE_F32, in_channels);            // norm1-3_w/b
+            mem_size += 6 * ggml_row_size(wtype, in_channels * in_channels);      // attn1_q/k/v/out_w attn2_q/out_w
+            mem_size += 2 * ggml_row_size(wtype, in_channels * context_dim);      // attn2_k/v_w
+            mem_size += ggml_row_size(wtype, in_channels * 4 * 2 * in_channels);  // ff_0_proj_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, in_channels * 4 * 2);        // ff_0_proj_b
+            mem_size += ggml_row_size(wtype, in_channels * 4 * in_channels);      // ff_2_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, in_channels);                // ff_2_b
+        }
+        return mem_size;
+    }
+
+    void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
+        norm_w    = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        norm_b    = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        proj_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
+        proj_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+        proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
+        proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+        // transformer
+        for (auto& transformer : transformers) {
+            transformer.norm1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+            transformer.norm1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+            transformer.attn1_q_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
+            transformer.attn1_k_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
+            transformer.attn1_v_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
+
+            transformer.attn1_out_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
+            transformer.attn1_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+            transformer.norm2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+            transformer.norm2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+            transformer.attn2_q_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
+            transformer.attn2_k_w = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
+            transformer.attn2_v_w = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
+
+            transformer.attn2_out_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
+            transformer.attn2_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+            transformer.norm3_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+            transformer.norm3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+            transformer.ff_0_proj_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels * 4 * 2);
+            transformer.ff_0_proj_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels * 4 * 2);
+
+            transformer.ff_2_w = ggml_new_tensor_2d(ctx, wtype, in_channels * 4, in_channels);
+            transformer.ff_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        }
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "norm.weight"]    = norm_w;
+        tensors[prefix + "norm.bias"]      = norm_b;
+        tensors[prefix + "proj_in.weight"] = proj_in_w;
+        tensors[prefix + "proj_in.bias"]   = proj_in_b;
+
+        // transformer
+        for (int i = 0; i < transformers.size(); i++) {
+            auto& transformer                                 = transformers[i];
+            std::string transformer_prefix                    = prefix + "transformer_blocks." + std::to_string(i) + ".";
+            tensors[transformer_prefix + "attn1.to_q.weight"] = transformer.attn1_q_w;
+            tensors[transformer_prefix + "attn1.to_k.weight"] = transformer.attn1_k_w;
+            tensors[transformer_prefix + "attn1.to_v.weight"] = transformer.attn1_v_w;
+
+            tensors[transformer_prefix + "attn1.to_out.0.weight"] = transformer.attn1_out_w;
+            tensors[transformer_prefix + "attn1.to_out.0.bias"]   = transformer.attn1_out_b;
+
+            tensors[transformer_prefix + "ff.net.0.proj.weight"] = transformer.ff_0_proj_w;
+            tensors[transformer_prefix + "ff.net.0.proj.bias"]   = transformer.ff_0_proj_b;
+            tensors[transformer_prefix + "ff.net.2.weight"]      = transformer.ff_2_w;
+            tensors[transformer_prefix + "ff.net.2.bias"]        = transformer.ff_2_b;
+
+            tensors[transformer_prefix + "attn2.to_q.weight"] = transformer.attn2_q_w;
+            tensors[transformer_prefix + "attn2.to_k.weight"] = transformer.attn2_k_w;
+            tensors[transformer_prefix + "attn2.to_v.weight"] = transformer.attn2_v_w;
+
+            tensors[transformer_prefix + "attn2.to_out.0.weight"] = transformer.attn2_out_w;
+            tensors[transformer_prefix + "attn2.to_out.0.bias"]   = transformer.attn2_out_b;
+
+            tensors[transformer_prefix + "norm1.weight"] = transformer.norm1_w;
+            tensors[transformer_prefix + "norm1.bias"]   = transformer.norm1_b;
+            tensors[transformer_prefix + "norm2.weight"] = transformer.norm2_w;
+            tensors[transformer_prefix + "norm2.bias"]   = transformer.norm2_b;
+            tensors[transformer_prefix + "norm3.weight"] = transformer.norm3_w;
+            tensors[transformer_prefix + "norm3.bias"]   = transformer.norm3_b;
+        }
+
+        tensors[prefix + "proj_out.weight"] = proj_out_w;
+        tensors[prefix + "proj_out.bias"]   = proj_out_b;
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* context) {
+        // x: [N, in_channels, h, w]
+        // context: [N, max_position, hidden_size(aka context_dim)]
+        auto x_in = x;
+        x         = ggml_nn_group_norm(ctx, x, norm_w, norm_b);
+        // proj_in
+        x = ggml_nn_conv_2d(ctx, x, proj_in_w, proj_in_b);  // [N, in_channels, h, w]
+
+        // transformer
+        const int64_t n            = x->ne[3];
+        const int64_t c            = x->ne[2];
+        const int64_t h            = x->ne[1];
+        const int64_t w            = x->ne[0];
+        const int64_t max_position = context->ne[1];
+        x                          = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3));  // [N, h, w, in_channels]
+
+        for (auto& transformer : transformers) {
+            auto r = x;
+            // layer norm 1
+            x = ggml_reshape_2d(ctx, x, c, w * h * n);
+            x = ggml_nn_layer_norm(ctx, x, transformer.norm1_w, transformer.norm1_b);
+
+            // self-attention
+            {
+                x                     = ggml_reshape_2d(ctx, x, c, h * w * n);        // [N * h * w, in_channels]
+                struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn1_q_w, x);  // [N * h * w, in_channels]
+#if !defined(SD_USE_FLASH_ATTENTION) || defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)
+                q = ggml_scale_inplace(ctx, q, 1.0f / sqrt((float)d_head));
+#endif
+                q = ggml_reshape_4d(ctx, q, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
+                q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));    // [N, n_head, h * w, d_head]
+                q = ggml_reshape_3d(ctx, q, d_head, h * w, n_head * n);  // [N * n_head, h * w, d_head]
+
+                struct ggml_tensor* k = ggml_mul_mat(ctx, transformer.attn1_k_w, x);         // [N * h * w, in_channels]
+                k                     = ggml_reshape_4d(ctx, k, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
+                k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));    // [N, n_head, h * w, d_head]
+                k                     = ggml_reshape_3d(ctx, k, d_head, h * w, n_head * n);  // [N * n_head, h * w, d_head]
+
+                struct ggml_tensor* v = ggml_mul_mat(ctx, transformer.attn1_v_w, x);         // [N * h * w, in_channels]
+                v                     = ggml_reshape_4d(ctx, v, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
+                v                     = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));    // [N, n_head, d_head, h * w]
+                v                     = ggml_reshape_3d(ctx, v, h * w, d_head, n_head * n);  // [N * n_head, d_head, h * w]
+
+#if defined(SD_USE_FLASH_ATTENTION) && !defined(SD_USE_CUBLAS) && !defined(SD_USE_METAL)
+                struct ggml_tensor* kqv = ggml_flash_attn(ctx, q, k, v, false);  // [N * n_head, h * w, d_head]
+#else
+                struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, h * w, h * w]
+                // kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);
+                kq = ggml_soft_max_inplace(ctx, kq);
+
+                struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, h * w, d_head]
+#endif
+                kqv = ggml_reshape_4d(ctx, kqv, d_head, h * w, n_head, n);
+                kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));  // [N, h * w, n_head, d_head]
+
+                // x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n));
+                x = ggml_reshape_2d(ctx, kqv, d_head * n_head, h * w * n);
+
+                x = ggml_nn_linear(ctx, x, transformer.attn1_out_w, transformer.attn1_out_b);
+
+                x = ggml_reshape_4d(ctx, x, c, w, h, n);
+            }
+
+            x = ggml_add(ctx, x, r);
+            r = x;
+
+            // layer norm 2
+            x = ggml_nn_layer_norm(ctx, x, transformer.norm2_w, transformer.norm2_b);
+
+            // cross-attention
+            {
+                x                     = ggml_reshape_2d(ctx, x, c, h * w * n);                                           // [N * h * w, in_channels]
+                context               = ggml_reshape_2d(ctx, context, context->ne[0], context->ne[1] * context->ne[2]);  // [N * max_position, hidden_size]
+                struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn2_q_w, x);                                     // [N * h * w, in_channels]
+#if !defined(SD_USE_FLASH_ATTENTION) || defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)
+                q = ggml_scale_inplace(ctx, q, 1.0f / sqrt((float)d_head));
+#endif
+                q = ggml_reshape_4d(ctx, q, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
+                q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));    // [N, n_head, h * w, d_head]
+                q = ggml_reshape_3d(ctx, q, d_head, h * w, n_head * n);  // [N * n_head, h * w, d_head]
+
+                struct ggml_tensor* k = ggml_mul_mat(ctx, transformer.attn2_k_w, context);          // [N * max_position, in_channels]
+                k                     = ggml_reshape_4d(ctx, k, d_head, n_head, max_position, n);   // [N, max_position, n_head, d_head]
+                k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));           // [N, n_head, max_position, d_head]
+                k                     = ggml_reshape_3d(ctx, k, d_head, max_position, n_head * n);  // [N * n_head, max_position, d_head]
+
+                struct ggml_tensor* v = ggml_mul_mat(ctx, transformer.attn2_v_w, context);          // [N * max_position, in_channels]
+                v                     = ggml_reshape_4d(ctx, v, d_head, n_head, max_position, n);   // [N, max_position, n_head, d_head]
+                v                     = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));           // [N, n_head, d_head, max_position]
+                v                     = ggml_reshape_3d(ctx, v, max_position, d_head, n_head * n);  // [N * n_head, d_head, max_position]
+#if defined(SD_USE_FLASH_ATTENTION) && !defined(SD_USE_CUBLAS) && !defined(SD_USE_METAL)
+                struct ggml_tensor* kqv = ggml_flash_attn(ctx, q, k, v, false);  // [N * n_head, h * w, d_head]
+#else
+                struct ggml_tensor* kq  = ggml_mul_mat(ctx, k, q);   // [N * n_head, h * w, max_position]
+                // kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);
+                kq = ggml_soft_max_inplace(ctx, kq);
+
+                struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, h * w, d_head]
+#endif
+                kqv = ggml_reshape_4d(ctx, kqv, d_head, h * w, n_head, n);
+                kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));
+
+                // x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n)); // [N * h * w, in_channels]
+                x = ggml_reshape_2d(ctx, kqv, d_head * n_head, h * w * n);  // [N * h * w, in_channels]
+
+                x = ggml_nn_linear(ctx, x, transformer.attn2_out_w, transformer.attn2_out_b);
+
+                x = ggml_reshape_4d(ctx, x, c, w, h, n);
+            }
+
+            x = ggml_add(ctx, x, r);
+            r = x;
+
+            // layer norm 3
+            x = ggml_reshape_2d(ctx, x, c, h * w * n);  // [N * h * w, in_channels]
+            x = ggml_nn_layer_norm(ctx, x, transformer.norm3_w, transformer.norm3_b);
+
+            // ff
+            {
+                // GEGLU
+                auto x_w    = ggml_view_2d(ctx,
+                                           transformer.ff_0_proj_w,
+                                           transformer.ff_0_proj_w->ne[0],
+                                           transformer.ff_0_proj_w->ne[1] / 2,
+                                           transformer.ff_0_proj_w->nb[1],
+                                           0);  // [in_channels * 4, in_channels]
+                auto x_b    = ggml_view_1d(ctx,
+                                           transformer.ff_0_proj_b,
+                                           transformer.ff_0_proj_b->ne[0] / 2,
+                                           0);  // [in_channels * 4, in_channels]
+                auto gate_w = ggml_view_2d(ctx,
+                                           transformer.ff_0_proj_w,
+                                           transformer.ff_0_proj_w->ne[0],
+                                           transformer.ff_0_proj_w->ne[1] / 2,
+                                           transformer.ff_0_proj_w->nb[1],
+                                           transformer.ff_0_proj_w->nb[1] * transformer.ff_0_proj_w->ne[1] / 2);  // [in_channels * 4, ]
+                auto gate_b = ggml_view_1d(ctx,
+                                           transformer.ff_0_proj_b,
+                                           transformer.ff_0_proj_b->ne[0] / 2,
+                                           transformer.ff_0_proj_b->nb[0] * transformer.ff_0_proj_b->ne[0] / 2);  // [in_channels * 4, ]
+                x           = ggml_reshape_2d(ctx, x, c, w * h * n);
+                auto x_in   = x;
+                x           = ggml_nn_linear(ctx, x_in, x_w, x_b);        // [N * h * w, in_channels * 4]
+                auto gate   = ggml_nn_linear(ctx, x_in, gate_w, gate_b);  // [N * h * w, in_channels * 4]
+
+                gate = ggml_gelu_inplace(ctx, gate);
+
+                x = ggml_mul(ctx, x, gate);  // [N * h * w, in_channels * 4]
+                // fc
+                x = ggml_nn_linear(ctx, x, transformer.ff_2_w, transformer.ff_2_b);  // [N * h * w, in_channels]
+            }
+
+            x = ggml_reshape_4d(ctx, x, c, w, h, n);  // [N, h, w, in_channels]
+
+            // residual
+            x = ggml_add(ctx, x, r);
+        }
+
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3));  // [N, in_channels, h, w]
+
+        // proj_out
+        x = ggml_nn_conv_2d(ctx, x, proj_out_w, proj_out_b);  // [N, in_channels, h, w]
+
+        x = ggml_add(ctx, x, x_in);
+        return x;
+    }
+};
+
 #endif  // __COMMON_HPP__

control.hpp

@@ -0,0 +1,697 @@
+#ifndef __CONTROL_HPP__
+#define __CONTROL_HPP__
+
+#include "common.hpp"
+#include "ggml_extend.hpp"
+#include "model.h"
+
+#define CONTROL_NET_GRAPH_SIZE 1536
+
+/*
+    =================================== ControlNet ===================================
+    Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/cldm/cldm.py
+
+*/
+
+struct CNHintBlock {
+    int hint_channels    = 3;
+    int model_channels   = 320;  // SD 1.5
+    int feat_channels[4] = {16, 32, 96, 256};
+    int num_blocks       = 3;
+    ggml_tensor* conv_first_w;  // [feat_channels[0], hint_channels, 3, 3]
+    ggml_tensor* conv_first_b;  // [feat_channels[0]]
+
+    struct hint_block {
+        ggml_tensor* conv_0_w;  // [feat_channels[idx], feat_channels[idx], 3, 3]
+        ggml_tensor* conv_0_b;  // [feat_channels[idx]]
+
+        ggml_tensor* conv_1_w;  // [feat_channels[idx + 1], feat_channels[idx], 3, 3]
+        ggml_tensor* conv_1_b;  // [feat_channels[idx + 1]]
+    };
+
+    hint_block blocks[3];
+    ggml_tensor* conv_final_w;  // [model_channels, feat_channels[3], 3, 3]
+    ggml_tensor* conv_final_b;  // [model_channels]
+
+    size_t calculate_mem_size() {
+        size_t mem_size = feat_channels[0] * hint_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_first_w
+        mem_size += feat_channels[0] * ggml_type_size(GGML_TYPE_F32);                                // conv_first_b
+        for (int i = 0; i < num_blocks; i++) {
+            mem_size += feat_channels[i] * feat_channels[i] * 3 * 3 * ggml_type_size(GGML_TYPE_F16);      // conv_0_w
+            mem_size += feat_channels[i] * ggml_type_size(GGML_TYPE_F32);                                 // conv_0_b
+            mem_size += feat_channels[i + 1] * feat_channels[i] * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_1_w
+            mem_size += feat_channels[i + 1] * ggml_type_size(GGML_TYPE_F32);                             // conv_1_b
+        }
+        mem_size += model_channels * feat_channels[3] * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_final_w
+        mem_size += model_channels * ggml_type_size(GGML_TYPE_F32);                             // conv_final_b
+        return mem_size;
+    }
+
+    void init_params(struct ggml_context* ctx) {
+        conv_first_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, hint_channels, feat_channels[0]);
+        conv_first_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, feat_channels[0]);
+
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[i].conv_0_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, feat_channels[i], feat_channels[i]);
+            blocks[i].conv_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, feat_channels[i]);
+            blocks[i].conv_1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, feat_channels[i], feat_channels[i + 1]);
+            blocks[i].conv_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, feat_channels[i + 1]);
+        }
+
+        conv_final_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, feat_channels[3], model_channels);
+        conv_final_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, model_channels);
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "input_hint_block.0.weight"] = conv_first_w;
+        tensors[prefix + "input_hint_block.0.bias"]   = conv_first_b;
+        int index                                     = 2;
+        for (int i = 0; i < num_blocks; i++) {
+            tensors[prefix + "input_hint_block." + std::to_string(index) + ".weight"] = blocks[i].conv_0_w;
+            tensors[prefix + "input_hint_block." + std::to_string(index) + ".bias"]   = blocks[i].conv_0_b;
+            index += 2;
+            tensors[prefix + "input_hint_block." + std::to_string(index) + ".weight"] = blocks[i].conv_1_w;
+            tensors[prefix + "input_hint_block." + std::to_string(index) + ".bias"]   = blocks[i].conv_1_b;
+            index += 2;
+        }
+        tensors[prefix + "input_hint_block.14.weight"] = conv_final_w;
+        tensors[prefix + "input_hint_block.14.bias"]   = conv_final_b;
+    }
+
+    struct ggml_tensor* forward(ggml_context* ctx, struct ggml_tensor* x) {
+        auto h = ggml_nn_conv_2d(ctx, x, conv_first_w, conv_first_b, 1, 1, 1, 1);
+        h      = ggml_silu_inplace(ctx, h);
+
+        auto body_h = h;
+        for (int i = 0; i < num_blocks; i++) {
+            // operations.conv_nd(dims, 16, 16, 3, padding=1)
+            body_h = ggml_nn_conv_2d(ctx, body_h, blocks[i].conv_0_w, blocks[i].conv_0_b, 1, 1, 1, 1);
+            body_h = ggml_silu_inplace(ctx, body_h);
+            // operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2)
+            body_h = ggml_nn_conv_2d(ctx, body_h, blocks[i].conv_1_w, blocks[i].conv_1_b, 2, 2, 1, 1);
+            body_h = ggml_silu_inplace(ctx, body_h);
+        }
+
+        h = ggml_nn_conv_2d(ctx, body_h, conv_final_w, conv_final_b, 1, 1, 1, 1);
+        h = ggml_silu_inplace(ctx, h);
+        return h;
+    }
+};
+
+struct CNZeroConv {
+    int channels;
+    ggml_tensor* conv_w;  // [channels, channels, 1, 1]
+    ggml_tensor* conv_b;  // [channels]
+
+    void init_params(struct ggml_context* ctx) {
+        conv_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, channels, channels);
+        conv_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+    }
+};
+
+struct ControlNet : public GGMLModule {
+    int in_channels                        = 4;
+    int model_channels                     = 320;
+    int out_channels                       = 4;
+    int num_res_blocks                     = 2;
+    std::vector<int> attention_resolutions = {4, 2, 1};
+    std::vector<int> channel_mult          = {1, 2, 4, 4};
+    std::vector<int> transformer_depth     = {1, 1, 1, 1};
+    int time_embed_dim                     = 1280;  // model_channels*4
+    int num_heads                          = 8;
+    int num_head_channels                  = -1;  // channels // num_heads
+    int context_dim                        = 768;
+    int middle_out_channel;
+    CNHintBlock input_hint_block;
+    CNZeroConv zero_convs[12];
+    int num_zero_convs = 1;
+
+    // network params
+    struct ggml_tensor* time_embed_0_w;  // [time_embed_dim, model_channels]
+    struct ggml_tensor* time_embed_0_b;  // [time_embed_dim, ]
+    // time_embed_1 is nn.SILU()
+    struct ggml_tensor* time_embed_2_w;  // [time_embed_dim, time_embed_dim]
+    struct ggml_tensor* time_embed_2_b;  // [time_embed_dim, ]
+
+    struct ggml_tensor* input_block_0_w;  // [model_channels, in_channels, 3, 3]
+    struct ggml_tensor* input_block_0_b;  // [model_channels, ]
+
+    // input_blocks
+    ResBlock input_res_blocks[4][2];
+    SpatialTransformer input_transformers[3][2];
+    DownSample input_down_samples[3];
+
+    // middle_block
+    ResBlock middle_block_0;
+    SpatialTransformer middle_block_1;
+    ResBlock middle_block_2;
+
+    struct ggml_tensor* middle_block_out_w;       // [middle_out_channel, middle_out_channel, 1, 1]
+    struct ggml_tensor* middle_block_out_b;       // [middle_out_channel, ]
+    ggml_backend_buffer_t control_buffer = NULL;  // keep control output tensors in backend memory
+    ggml_context* control_ctx            = NULL;
+    std::vector<struct ggml_tensor*> controls;  // (12 input block outputs, 1 middle block output) SD 1.5
+
+    ControlNet() {
+        name = "controlnet";
+        // input_blocks
+        std::vector<int> input_block_chans;
+        input_block_chans.push_back(model_channels);
+        int ch                 = model_channels;
+        zero_convs[0].channels = model_channels;
+        int ds                 = 1;
+
+        int len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            int mult = channel_mult[i];
+            for (int j = 0; j < num_res_blocks; j++) {
+                input_res_blocks[i][j].channels     = ch;
+                input_res_blocks[i][j].emb_channels = time_embed_dim;
+                input_res_blocks[i][j].out_channels = mult * model_channels;
+
+                ch = mult * model_channels;
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    int n_head = num_heads;
+                    int d_head = ch / num_heads;
+                    if (num_head_channels != -1) {
+                        d_head = num_head_channels;
+                        n_head = ch / d_head;
+                    }
+                    input_transformers[i][j]             = SpatialTransformer(transformer_depth[i]);
+                    input_transformers[i][j].in_channels = ch;
+                    input_transformers[i][j].n_head      = n_head;
+                    input_transformers[i][j].d_head      = d_head;
+                    input_transformers[i][j].context_dim = context_dim;
+                }
+                input_block_chans.push_back(ch);
+
+                zero_convs[num_zero_convs].channels = ch;
+                num_zero_convs++;
+            }
+            if (i != len_mults - 1) {
+                input_down_samples[i].channels     = ch;
+                input_down_samples[i].out_channels = ch;
+                input_block_chans.push_back(ch);
+
+                zero_convs[num_zero_convs].channels = ch;
+                num_zero_convs++;
+                ds *= 2;
+            }
+        }
+        GGML_ASSERT(num_zero_convs == 12);
+
+        // middle blocks
+        middle_block_0.channels     = ch;
+        middle_block_0.emb_channels = time_embed_dim;
+        middle_block_0.out_channels = ch;
+
+        int n_head = num_heads;
+        int d_head = ch / num_heads;
+        if (num_head_channels != -1) {
+            d_head = num_head_channels;
+            n_head = ch / d_head;
+        }
+        middle_block_1             = SpatialTransformer(transformer_depth[transformer_depth.size() - 1]);
+        middle_block_1.in_channels = ch;
+        middle_block_1.n_head      = n_head;
+        middle_block_1.d_head      = d_head;
+        middle_block_1.context_dim = context_dim;
+
+        middle_block_2.channels     = ch;
+        middle_block_2.emb_channels = time_embed_dim;
+        middle_block_2.out_channels = ch;
+        middle_out_channel          = ch;
+    }
+
+    size_t calculate_mem_size() {
+        size_t mem_size = 0;
+        mem_size += input_hint_block.calculate_mem_size();
+        mem_size += ggml_row_size(wtype, time_embed_dim * model_channels);  // time_embed_0_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);           // time_embed_0_b
+        mem_size += ggml_row_size(wtype, time_embed_dim * time_embed_dim);  // time_embed_2_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);           // time_embed_2_b
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, model_channels * in_channels * 3 * 3);  // input_block_0_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, model_channels);                        // input_block_0_b
+
+        // input_blocks
+        int ds        = 1;
+        int len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                mem_size += input_res_blocks[i][j].calculate_mem_size(wtype);
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    mem_size += input_transformers[i][j].calculate_mem_size(wtype);
+                }
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                mem_size += input_down_samples[i].calculate_mem_size(wtype);
+            }
+        }
+
+        for (int i = 0; i < num_zero_convs; i++) {
+            mem_size += ggml_row_size(GGML_TYPE_F16, zero_convs[i].channels * zero_convs[i].channels);
+            mem_size += ggml_row_size(GGML_TYPE_F32, zero_convs[i].channels);
+        }
+
+        // middle_block
+        mem_size += middle_block_0.calculate_mem_size(wtype);
+        mem_size += middle_block_1.calculate_mem_size(wtype);
+        mem_size += middle_block_2.calculate_mem_size(wtype);
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, middle_out_channel * middle_out_channel);  // middle_block_out_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, middle_out_channel);                       // middle_block_out_b
+
+        return mem_size;
+    }
+
+    size_t get_num_tensors() {
+        // in
+        size_t num_tensors = 6;
+
+        num_tensors += num_zero_convs * 2;
+
+        // input blocks
+        int ds        = 1;
+        int len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                num_tensors += 12;
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    num_tensors += input_transformers[i][j].get_num_tensors();
+                }
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                num_tensors += 2;
+            }
+        }
+
+        // middle blocks
+        num_tensors += 13 * 2;
+        num_tensors += middle_block_1.get_num_tensors();
+        return num_tensors;
+    }
+
+    void init_params() {
+        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
+
+        input_hint_block.init_params(params_ctx);
+
+        time_embed_0_w = ggml_new_tensor_2d(params_ctx, wtype, model_channels, time_embed_dim);
+        time_embed_0_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, time_embed_dim);
+        time_embed_2_w = ggml_new_tensor_2d(params_ctx, wtype, time_embed_dim, time_embed_dim);
+        time_embed_2_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, time_embed_dim);
+
+        // input_blocks
+        input_block_0_w = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, in_channels, model_channels);
+        input_block_0_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, model_channels);
+
+        int ds        = 1;
+        int len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                input_res_blocks[i][j].init_params(params_ctx, wtype);
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    input_transformers[i][j].init_params(params_ctx, alloc, wtype);
+                }
+            }
+            if (i != len_mults - 1) {
+                input_down_samples[i].init_params(params_ctx, wtype);
+                ds *= 2;
+            }
+        }
+
+        for (int i = 0; i < num_zero_convs; i++) {
+            zero_convs[i].init_params(params_ctx);
+        }
+
+        // middle_blocks
+        middle_block_0.init_params(params_ctx, wtype);
+        middle_block_1.init_params(params_ctx, alloc, wtype);
+        middle_block_2.init_params(params_ctx, wtype);
+
+        // middle_block_out
+        middle_block_out_w = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 1, 1, middle_out_channel, middle_out_channel);
+        middle_block_out_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, middle_out_channel);
+
+        // alloc all tensors linked to this context
+        for (struct ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != NULL; t = ggml_get_next_tensor(params_ctx, t)) {
+            if (t->data == NULL) {
+                ggml_allocr_alloc(alloc, t);
+            }
+        }
+
+        ggml_allocr_free(alloc);
+    }
+
+    bool load_from_file(const std::string& file_path, ggml_backend_t backend_, ggml_type wtype_) {
+        LOG_INFO("loading control net from '%s'", file_path.c_str());
+
+        std::map<std::string, ggml_tensor*> control_tensors;
+
+        ModelLoader model_loader;
+        if (!model_loader.init_from_file(file_path)) {
+            LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str());
+            return false;
+        }
+
+        if (!alloc_params_buffer(backend_, wtype_)) {
+            return false;
+        }
+
+        // prepare memory for the weights
+        {
+            init_params();
+            map_by_name(control_tensors, "");
+        }
+
+        std::set<std::string> tensor_names_in_file;
+
+        auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+            const std::string& name = tensor_storage.name;
+            tensor_names_in_file.insert(name);
+
+            struct ggml_tensor* real;
+            if (control_tensors.find(name) != control_tensors.end()) {
+                real = control_tensors[name];
+            } else {
+                LOG_ERROR("unknown tensor '%s' in model file", name.data());
+                return true;
+            }
+
+            if (
+                real->ne[0] != tensor_storage.ne[0] ||
+                real->ne[1] != tensor_storage.ne[1] ||
+                real->ne[2] != tensor_storage.ne[2] ||
+                real->ne[3] != tensor_storage.ne[3]) {
+                LOG_ERROR(
+                    "tensor '%s' has wrong shape in model file: "
+                    "got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
+                    name.c_str(),
+                    (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
+                    (int)real->ne[0], (int)real->ne[1], (int)real->ne[2], (int)real->ne[3]);
+                return false;
+            }
+
+            *dst_tensor = real;
+
+            return true;
+        };
+
+        bool success = model_loader.load_tensors(on_new_tensor_cb, backend);
+
+        bool some_tensor_not_init = false;
+
+        for (auto pair : control_tensors) {
+            if (tensor_names_in_file.find(pair.first) == tensor_names_in_file.end()) {
+                LOG_ERROR("tensor '%s' not in model file", pair.first.c_str());
+                some_tensor_not_init = true;
+            }
+        }
+
+        if (some_tensor_not_init) {
+            return false;
+        }
+
+        LOG_INFO("control net model loaded");
+        return success;
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        input_hint_block.map_by_name(tensors, "");
+        tensors[prefix + "time_embed.0.weight"] = time_embed_0_w;
+        tensors[prefix + "time_embed.0.bias"]   = time_embed_0_b;
+        tensors[prefix + "time_embed.2.weight"] = time_embed_2_w;
+        tensors[prefix + "time_embed.2.bias"]   = time_embed_2_b;
+
+        // input_blocks
+        tensors[prefix + "input_blocks.0.0.weight"] = input_block_0_w;
+        tensors[prefix + "input_blocks.0.0.bias"]   = input_block_0_b;
+
+        int len_mults       = channel_mult.size();
+        int input_block_idx = 0;
+        int ds              = 1;
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                input_block_idx += 1;
+                input_res_blocks[i][j].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".0.");
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    input_transformers[i][j].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".1.");
+                }
+            }
+            if (i != len_mults - 1) {
+                input_block_idx += 1;
+                input_down_samples[i].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".0.");
+                ds *= 2;
+            }
+        }
+
+        for (int i = 0; i < num_zero_convs; i++) {
+            tensors[prefix + "zero_convs." + std::to_string(i) + ".0.weight"] = zero_convs[i].conv_w;
+            tensors[prefix + "zero_convs." + std::to_string(i) + ".0.bias"]   = zero_convs[i].conv_b;
+        }
+
+        // middle_blocks
+        middle_block_0.map_by_name(tensors, prefix + "middle_block.0.");
+        middle_block_1.map_by_name(tensors, prefix + "middle_block.1.");
+        middle_block_2.map_by_name(tensors, prefix + "middle_block.2.");
+
+        tensors[prefix + "middle_block_out.0.weight"] = middle_block_out_w;
+        tensors[prefix + "middle_block_out.0.bias"]   = middle_block_out_b;
+    }
+
+    struct ggml_cgraph* build_graph_hint(struct ggml_tensor* hint) {
+        // since we are using ggml-alloc, this buffer only needs enough space to hold the ggml_tensor and ggml_cgraph structs, but not the tensor data
+        static size_t buf_size = ggml_tensor_overhead() * GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead();
+        static std::vector<uint8_t> buf(buf_size);
+
+        struct ggml_init_params params = {
+            /*.mem_size   =*/buf_size,
+            /*.mem_buffer =*/buf.data(),
+            /*.no_alloc   =*/true,  // the tensors will be allocated later by ggml_allocr_alloc_graph()
+        };
+
+        struct ggml_context* ctx0 = ggml_init(params);
+        struct ggml_cgraph* gf    = ggml_new_graph(ctx0);
+        // temporal tensors for transfer tensors from cpu to gpu if needed
+        struct ggml_tensor* hint_t = NULL;
+        // it's performing a compute, check if backend isn't cpu
+        if (!ggml_backend_is_cpu(backend)) {
+            // pass input tensors to gpu memory
+            hint_t = ggml_dup_tensor(ctx0, hint);
+            ggml_allocr_alloc(compute_allocr, hint_t);
+            // pass data to device backend
+            if (!ggml_allocr_is_measure(compute_allocr)) {
+                ggml_backend_tensor_set(hint_t, hint->data, 0, ggml_nbytes(hint));
+            }
+        } else {
+            // if it's cpu backend just pass the same tensors
+            hint_t = hint;
+        }
+        struct ggml_tensor* out = input_hint_block.forward(ctx0, hint_t);
+        ggml_build_forward_expand(gf, out);
+        ggml_free(ctx0);
+        return gf;
+    }
+
+    void process_hint(struct ggml_tensor* output, int n_threads, struct ggml_tensor* hint) {
+        // compute buffer size
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph_hint(hint);
+        };
+        GGMLModule::alloc_compute_buffer(get_graph);
+        // perform computation
+        GGMLModule::compute(get_graph, n_threads, output);
+        GGMLModule::free_compute_buffer();
+    }
+
+    void forward(struct ggml_cgraph* gf,
+                 struct ggml_context* ctx0,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* hint,
+                 struct ggml_tensor* timesteps,
+                 struct ggml_tensor* context,
+                 struct ggml_tensor* t_emb = NULL) {
+        // x: [N, in_channels, h, w]
+        // timesteps: [N, ]
+        // t_emb: [N, model_channels]
+        // context: [N, max_position, hidden_size]([N, 77, 768])
+        if (t_emb == NULL && timesteps != NULL) {
+            t_emb = new_timestep_embedding(ctx0, compute_allocr, timesteps, model_channels);  // [N, model_channels]
+        }
+
+        // time_embed = nn.Sequential
+        auto emb = ggml_nn_linear(ctx0, t_emb, time_embed_0_w, time_embed_0_b);
+        emb      = ggml_silu_inplace(ctx0, emb);
+        emb      = ggml_nn_linear(ctx0, emb, time_embed_2_w, time_embed_2_b);  // [N, time_embed_dim]
+
+        // input_blocks
+        int zero_conv_offset = 0;
+
+        // input block 0
+        struct ggml_tensor* h = ggml_nn_conv_2d(ctx0, x, input_block_0_w, input_block_0_b, 1, 1, 1, 1);  // [N, model_channels, h, w]
+        h                     = ggml_add(ctx0, h, hint);
+
+        auto h_c = ggml_nn_conv_2d(ctx0, h, zero_convs[zero_conv_offset].conv_w, zero_convs[zero_conv_offset].conv_b);
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_c, controls[zero_conv_offset]));
+        zero_conv_offset++;
+
+        // input block 1-11
+        int len_mults = channel_mult.size();
+        int ds        = 1;
+        for (int i = 0; i < len_mults; i++) {
+            int mult = channel_mult[i];
+            for (int j = 0; j < num_res_blocks; j++) {
+                h = input_res_blocks[i][j].forward(ctx0, h, emb);  // [N, mult*model_channels, h, w]
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    h = input_transformers[i][j].forward(ctx0, h, context);  // [N, mult*model_channels, h, w]
+                }
+                h_c = ggml_nn_conv_2d(ctx0, h, zero_convs[zero_conv_offset].conv_w, zero_convs[zero_conv_offset].conv_b);
+                ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_c, controls[zero_conv_offset]));
+                zero_conv_offset++;
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                h   = input_down_samples[i].forward(ctx0, h);  // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
+                h_c = ggml_nn_conv_2d(ctx0, h, zero_convs[zero_conv_offset].conv_w, zero_convs[zero_conv_offset].conv_b);
+                ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_c, controls[zero_conv_offset]));
+                zero_conv_offset++;
+            }
+        }
+        // [N, 4*model_channels, h/8, w/8]
+
+        // middle_block
+        h = middle_block_0.forward(ctx0, h, emb);      // [N, 4*model_channels, h/8, w/8]
+        h = middle_block_1.forward(ctx0, h, context);  // [N, 4*model_channels, h/8, w/8]
+        h = middle_block_2.forward(ctx0, h, emb);      // [N, 4*model_channels, h/8, w/8]
+
+        h_c = ggml_nn_conv_2d(ctx0, h, middle_block_out_w, middle_block_out_b);
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_c, controls[zero_conv_offset]));
+    }
+
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                    struct ggml_tensor* hint,
+                                    struct ggml_tensor* timesteps,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* t_emb = NULL) {
+        // since we are using ggml-alloc, this buffer only needs enough space to hold the ggml_tensor and ggml_cgraph structs, but not the tensor data
+        static size_t buf_size = ggml_tensor_overhead() * CONTROL_NET_GRAPH_SIZE + ggml_graph_overhead();
+        static std::vector<uint8_t> buf(buf_size);
+
+        struct ggml_init_params params = {
+            /*.mem_size   =*/buf_size,
+            /*.mem_buffer =*/buf.data(),
+            /*.no_alloc   =*/true,  // the tensors will be allocated later by ggml_allocr_alloc_graph()
+        };
+        // LOG_DEBUG("mem_size %u ", params.mem_size);
+
+        struct ggml_context* ctx0 = ggml_init(params);
+
+        struct ggml_cgraph* gf = ggml_new_graph_custom(ctx0, CONTROL_NET_GRAPH_SIZE, false);
+
+        // temporal tensors for transfer tensors from cpu to gpu if needed
+        struct ggml_tensor* x_t         = NULL;
+        struct ggml_tensor* hint_t      = NULL;
+        struct ggml_tensor* timesteps_t = NULL;
+        struct ggml_tensor* context_t   = NULL;
+        struct ggml_tensor* t_emb_t     = NULL;
+
+        // it's performing a compute, check if backend isn't cpu
+        if (!ggml_backend_is_cpu(backend)) {
+            // pass input tensors to gpu memory
+            x_t       = ggml_dup_tensor(ctx0, x);
+            context_t = ggml_dup_tensor(ctx0, context);
+            hint_t    = ggml_dup_tensor(ctx0, hint);
+            ggml_allocr_alloc(compute_allocr, x_t);
+            if (timesteps != NULL) {
+                timesteps_t = ggml_dup_tensor(ctx0, timesteps);
+                ggml_allocr_alloc(compute_allocr, timesteps_t);
+            }
+            ggml_allocr_alloc(compute_allocr, context_t);
+            ggml_allocr_alloc(compute_allocr, hint_t);
+            if (t_emb != NULL) {
+                t_emb_t = ggml_dup_tensor(ctx0, t_emb);
+                ggml_allocr_alloc(compute_allocr, t_emb_t);
+            }
+            // pass data to device backend
+            if (!ggml_allocr_is_measure(compute_allocr)) {
+                ggml_backend_tensor_set(x_t, x->data, 0, ggml_nbytes(x));
+                ggml_backend_tensor_set(context_t, context->data, 0, ggml_nbytes(context));
+                ggml_backend_tensor_set(hint_t, hint->data, 0, ggml_nbytes(hint));
+                if (timesteps_t != NULL) {
+                    ggml_backend_tensor_set(timesteps_t, timesteps->data, 0, ggml_nbytes(timesteps));
+                }
+                if (t_emb_t != NULL) {
+                    ggml_backend_tensor_set(t_emb_t, t_emb->data, 0, ggml_nbytes(t_emb));
+                }
+            }
+        } else {
+            // if it's cpu backend just pass the same tensors
+            x_t         = x;
+            timesteps_t = timesteps;
+            context_t   = context;
+            t_emb_t     = t_emb;
+            hint_t      = hint;
+        }
+
+        forward(gf, ctx0, x_t, hint_t, timesteps_t, context_t, t_emb_t);
+
+        ggml_free(ctx0);
+
+        return gf;
+    }
+
+    void alloc_compute_buffer(struct ggml_tensor* x,
+                              struct ggml_tensor* hint,
+                              struct ggml_tensor* context,
+                              struct ggml_tensor* t_emb = NULL) {
+        {
+            struct ggml_init_params params;
+            params.mem_size            = static_cast<size_t>(14 * ggml_tensor_overhead()) + 256;
+            params.mem_buffer          = NULL;
+            params.no_alloc            = true;
+            control_ctx                = ggml_init(params);
+            size_t control_buffer_size = 0;
+            int w = x->ne[0], h = x->ne[1], steps = 0;
+            for (int i = 0; i < (num_zero_convs + 1); i++) {
+                bool last = i == num_zero_convs;
+                int c     = last ? middle_out_channel : zero_convs[i].channels;
+                if (!last && steps == 3) {
+                    w /= 2;
+                    h /= 2;
+                    steps = 0;
+                }
+                controls.push_back(ggml_new_tensor_4d(control_ctx, GGML_TYPE_F32, w, h, c, 1));
+                control_buffer_size += ggml_nbytes(controls[i]);
+                steps++;
+            }
+            control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, backend);
+        }
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(x, hint, NULL, context, t_emb);
+        };
+        GGMLModule::alloc_compute_buffer(get_graph);
+    }
+
+    void compute(int n_threads,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* hint,
+                 struct ggml_tensor* context,
+                 struct ggml_tensor* t_emb = NULL) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(x, hint, NULL, context, t_emb);
+        };
+        GGMLModule::compute(get_graph, n_threads, NULL);
+    }
+
+    void free_compute_buffer() {
+        GGMLModule::free_compute_buffer();
+        ggml_free(control_ctx);
+        ggml_backend_buffer_free(control_buffer);
+        control_buffer = NULL;
+    }
+};
+
+#endif  // __CONTROL_HPP__

docs/hipBLAS_on_Windows.md

@@ -0,0 +1,85 @@
+# Using hipBLAS on Windows
+
+To get hipBLAS in `stable-diffusion.cpp` working on Windows, go through this guide section by section.
+
+## Build Tools for Visual Studio 2022
+
+Skip this step if you already have Build Tools installed.
+
+To install Build Tools, go to [Visual Studio Downloads](https://visualstudio.microsoft.com/vs/), download `Visual Studio 2022 and other Products` and run the installer.
+
+## CMake
+
+Skip this step if you already have CMake installed: running `cmake --version` should output `cmake version x.y.z`.
+
+Download latest `Windows x64 Installer` from [Download | CMake](https://cmake.org/download/) and run it.
+
+## ROCm
+
+Skip this step if you already have Build Tools installed.
+
+The [validation tools](https://rocm.docs.amd.com/en/latest/reference/validation_tools.html) not support on Windows. So you should confirm the Version of `ROCM` by yourself.
+
+Fortunately, `AMD` provides complete help documentation, you can use the help documentation to install [ROCM](https://rocm.docs.amd.com/en/latest/deploy/windows/quick_start.html)
+
+>**If you encounter an error, if it is [AMD ROCm Windows Installation Error 215](https://github.com/RadeonOpenCompute/ROCm/issues/2363), don't worry about this error. ROCM has been installed correctly, but the vs studio plugin installation failed, we can ignore it.**
+
+Then we must set `ROCM` as environment variables before running cmake.
+
+Usually if you install according to the official tutorial and do not modify the ROCM path, then there is a high probability that it is here `C:\Program Files\AMD\ROCm\5.5\bin`
+
+This is what I use to set the clang:
+```Commandline
+set CC=C:\Program Files\AMD\ROCm\5.5\bin\clang.exe
+set CXX=C:\Program Files\AMD\ROCm\5.5\bin\clang++.exe
+```
+
+## Ninja
+
+Skip this step if you already have Ninja installed: running `ninja --version` should output `1.11.1`.
+
+Download latest `ninja-win.zip` from [GitHub Releases Page](https://github.com/ninja-build/ninja/releases/tag/v1.11.1) and unzip. Then set as environment variables. I unzipped it in `C:\Program Files\ninja`, so I set it like this:
+
+```Commandline
+set ninja=C:\Program Files\ninja\ninja.exe
+```
+## Building stable-diffusion.cpp
+
+The thing different from the regular CPU build is `-DSD_HIPBLAS=ON` ,
+`-G "Ninja"`, `-DCMAKE_C_COMPILER=clang`, `-DCMAKE_CXX_COMPILER=clang++`, `-DAMDGPU_TARGETS=gfx1100`
+
+>**Notice**: check the `clang` and `clang++` information:
+```Commandline
+clang --version
+clang++ --version
+```
+
+If you see like this, we can continue:
+```
+clang version 17.0.0 (git@github.amd.com:Compute-Mirrors/llvm-project e3201662d21c48894f2156d302276eb1cf47c7be)
+Target: x86_64-pc-windows-msvc
+Thread model: posix
+InstalledDir: C:\Program Files\AMD\ROCm\5.5\bin
+```
+
+```
+clang version 17.0.0 (git@github.amd.com:Compute-Mirrors/llvm-project e3201662d21c48894f2156d302276eb1cf47c7be)
+Target: x86_64-pc-windows-msvc
+Thread model: posix
+InstalledDir: C:\Program Files\AMD\ROCm\5.5\bin
+```
+
+>**Notice** that the `gfx1100` is the GPU architecture of my GPU, you can change it to your GPU architecture. Click here to see your architecture [LLVM Target](https://rocm.docs.amd.com/en/latest/release/windows_support.html#windows-supported-gpus)
+
+My GPU is AMD Radeon™ RX 7900 XTX Graphics, so I set it to `gfx1100`.
+
+option:
+
+```commandline
+mkdir build
+cd build
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
+cmake --build . --config Release
+```
+
+If everything went OK, `build\bin\sd.exe` file should appear.

esrgan.hpp

@@ -91,7 +91,7 @@ struct ResidualDenseBlock {
         tensors[prefix + "conv5.bias"]   = conv5_b;
     }
 
-    ggml_tensor* forward(ggml_context* ctx, ggml_tensor* out_scale, ggml_tensor* x /* feat */) {
+    ggml_tensor* forward(ggml_context* ctx, float out_scale, ggml_tensor* x /* feat */) {
         // x1 = self.lrelu(self.conv1(x))
         ggml_tensor* x1 = ggml_nn_conv_2d(ctx, x, conv1_w, conv1_b, 1, 1, 1, 1);
         x1              = ggml_leaky_relu(ctx, x1, 0.2f, true);
@@ -161,7 +161,7 @@ struct EsrganBlock {
         }
     }
 
-    ggml_tensor* forward(ggml_context* ctx, ggml_tensor* out_scale, ggml_tensor* x) {
+    ggml_tensor* forward(ggml_context* ctx, float out_scale, ggml_tensor* x) {
         ggml_tensor* out = x;
         for (int i = 0; i < num_residual_blocks; i++) {
             // out = self.rdb...(x)
@@ -325,7 +325,7 @@ struct ESRGAN : public GGMLModule {
         tensors["conv_last.bias"]   = conv_last_b;
     }
 
-    ggml_tensor* forward(ggml_context* ctx0, ggml_tensor* out_scale, ggml_tensor* x /* feat */) {
+    ggml_tensor* forward(ggml_context* ctx0, float out_scale, ggml_tensor* x /* feat */) {
         // feat = self.conv_first(feat)
         auto h = ggml_nn_conv_2d(ctx0, x, conv_first_w, conv_first_b, 1, 1, 1, 1);
 
@@ -376,12 +376,7 @@ struct ESRGAN : public GGMLModule {
         struct ggml_cgraph* gf = ggml_new_graph(ctx0);
 
         struct ggml_tensor* x_ = NULL;
-        struct ggml_tensor* os = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(compute_allocr, os);
-        if (!ggml_allocr_is_measure(compute_allocr)) {
-            float scale = 0.2f;
-            ggml_backend_tensor_set(os, &scale, 0, sizeof(scale));
-        }
+        float out_scale        = 0.2f;
 
         // it's performing a compute, check if backend isn't cpu
         if (!ggml_backend_is_cpu(backend)) {
@@ -397,7 +392,7 @@ struct ESRGAN : public GGMLModule {
             x_ = x;
         }
 
-        struct ggml_tensor* out = forward(ctx0, os, x);
+        struct ggml_tensor* out = forward(ctx0, out_scale, x);
 
         ggml_build_forward_expand(gf, out);
         ggml_free(ctx0);

examples/cli/main.cpp

@@ -6,6 +6,7 @@
 #include <string>
 #include <vector>
 
+#include "preprocessing.hpp"
 #include "stable-diffusion.h"
 
 #define STB_IMAGE_IMPLEMENTATION
@@ -42,11 +43,13 @@ const char* schedule_str[] = {
 const char* modes_str[] = {
     "txt2img",
     "img2img",
+    "convert",
 };
 
 enum SDMode {
     TXT2IMG,
     IMG2IMG,
+    CONVERT,
     MODE_COUNT
 };
 
@@ -58,10 +61,13 @@ struct SDParams {
     std::string vae_path;
     std::string taesd_path;
     std::string esrgan_path;
+    std::string controlnet_path;
+    std::string embeddings_path;
     sd_type_t wtype = SD_TYPE_COUNT;
     std::string lora_model_dir;
     std::string output_path = "output.png";
     std::string input_path;
+    std::string control_image_path;
 
     std::string prompt;
     std::string negative_prompt;
@@ -75,24 +81,15 @@ struct SDParams {
     schedule_t schedule           = DEFAULT;
     int sample_steps              = 20;
     float strength                = 0.75f;
+    float control_strength        = 0.9f;
     rng_type_t rng_type           = CUDA_RNG;
     int64_t seed                  = 42;
     bool verbose                  = false;
     bool vae_tiling               = false;
+    bool control_net_cpu          = false;
+    bool canny_preprocess         = false;
 };
 
-static std::string sd_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;
-}
-
 void print_params(SDParams params) {
     printf("Option: \n");
     printf("    n_threads:         %d\n", params.n_threads);
@@ -102,8 +99,13 @@ void print_params(SDParams params) {
     printf("    vae_path:          %s\n", params.vae_path.c_str());
     printf("    taesd_path:        %s\n", params.taesd_path.c_str());
     printf("    esrgan_path:       %s\n", params.esrgan_path.c_str());
+    printf("    controlnet_path:   %s\n", params.controlnet_path.c_str());
+    printf("    embeddings_path:   %s\n", params.embeddings_path.c_str());
     printf("    output_path:       %s\n", params.output_path.c_str());
     printf("    init_img:          %s\n", params.input_path.c_str());
+    printf("    control_image:     %s\n", params.control_image_path.c_str());
+    printf("    controlnet cpu:    %s\n", params.control_net_cpu ? "true" : "false");
+    printf("    strength(control): %.2f\n", params.control_strength);
     printf("    prompt:            %s\n", params.prompt.c_str());
     printf("    negative_prompt:   %s\n", params.negative_prompt.c_str());
     printf("    cfg_scale:         %.2f\n", params.cfg_scale);
@@ -125,22 +127,26 @@ void print_usage(int argc, const char* argv[]) {
     printf("\n");
     printf("arguments:\n");
     printf("  -h, --help                         show this help message and exit\n");
-    printf("  -M, --mode [txt2img or img2img]    generation mode (default: txt2img)\n");
+    printf("  -M, --mode [MODEL]                 run mode (txt2img or img2img or convert, default: txt2img)\n");
     printf("  -t, --threads N                    number of threads to use during computation (default: -1).\n");
     printf("                                     If threads <= 0, then threads will be set to the number of CPU physical cores\n");
     printf("  -m, --model [MODEL]                path to model\n");
     printf("  --vae [VAE]                        path to vae\n");
     printf("  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n");
+    printf("  --control-net [CONTROL_PATH]       path to control net model\n");
+    printf("  --embd-dir [EMBEDDING_PATH]        path to embeddings.\n");
     printf("  --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n");
     printf("  --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n");
     printf("                                     If not specified, the default is the type of the weight file.\n");
     printf("  --lora-model-dir [DIR]             lora model directory\n");
     printf("  -i, --init-img [IMAGE]             path to the input image, required by img2img\n");
+    printf("  --control-image [IMAGE]            path to image condition, control net\n");
     printf("  -o, --output OUTPUT                path to write result image to (default: ./output.png)\n");
     printf("  -p, --prompt [PROMPT]              the prompt to render\n");
     printf("  -n, --negative-prompt PROMPT       the negative prompt (default: \"\")\n");
     printf("  --cfg-scale SCALE                  unconditional guidance scale: (default: 7.0)\n");
     printf("  --strength STRENGTH                strength for noising/unnoising (default: 0.75)\n");
+    printf("  --control-strength STRENGTH        strength to apply Control Net (default: 0.9)\n");
     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");
@@ -154,6 +160,8 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n");
     printf("                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n");
     printf("  --vae-tiling                       process vae in tiles to reduce memory usage\n");
+    printf("  --control-net-cpu                  keep controlnet in cpu (for low vram)\n");
+    printf("  --canny                            apply canny preprocessor (edge detection)\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -205,12 +213,24 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.taesd_path = argv[i];
+        } else if (arg == "--control-net") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.controlnet_path = argv[i];
         } else if (arg == "--upscale-model") {
             if (++i >= argc) {
                 invalid_arg = true;
                 break;
             }
             params.esrgan_path = argv[i];
+        } else if (arg == "--embd-dir") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.embeddings_path = argv[i];
         } else if (arg == "--type") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -248,6 +268,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.input_path = argv[i];
+        } else if (arg == "--control-image") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.control_image_path = argv[i];
         } else if (arg == "-o" || arg == "--output") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -278,6 +304,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.strength = std::stof(argv[i]);
+        } else if (arg == "--control-strength") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.control_strength = std::stof(argv[i]);
         } else if (arg == "-H" || arg == "--height") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -304,6 +336,10 @@ void parse_args(int argc, const char** argv, SDParams& params) {
             params.clip_skip = std::stoi(argv[i]);
         } else if (arg == "--vae-tiling") {
             params.vae_tiling = true;
+        } else if (arg == "--control-net-cpu") {
+            params.control_net_cpu = true;
+        } else if (arg == "--canny") {
+            params.canny_preprocess = true;
         } else if (arg == "-b" || arg == "--batch-count") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -384,7 +420,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         params.n_threads = get_num_physical_cores();
     }
 
-    if (params.prompt.length() == 0) {
+    if (params.mode != CONVERT && params.prompt.length() == 0) {
         fprintf(stderr, "error: the following arguments are required: prompt\n");
         print_usage(argc, argv);
         exit(1);
@@ -432,6 +468,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         srand((int)time(NULL));
         params.seed = rand();
     }
+
+    if (params.mode == CONVERT) {
+        if (params.output_path == "output.png") {
+            params.output_path = "output.gguf";
+        }
+    }
 }
 
 std::string get_image_params(SDParams params, int64_t seed) {
@@ -460,13 +502,13 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
         return;
     }
     if (level <= SD_LOG_INFO) {
-        fprintf(stdout, log);
+        fputs(log, stdout);
         fflush(stdout);
     } else {
-        fprintf(stderr, log);
+        fputs(log, stderr);
         fflush(stderr);
     }
-};
+}
 
 int main(int argc, const char* argv[]) {
     SDParams params;
@@ -479,6 +521,24 @@ int main(int argc, const char* argv[]) {
         printf("%s", sd_get_system_info());
     }
 
+    if (params.mode == CONVERT) {
+        bool success = convert(params.model_path.c_str(), params.vae_path.c_str(), params.output_path.c_str(), params.wtype);
+        if (!success) {
+            fprintf(stderr,
+                    "convert '%s'/'%s' to '%s' failed\n",
+                    params.model_path.c_str(),
+                    params.vae_path.c_str(),
+                    params.output_path.c_str());
+            return 1;
+        } else {
+            printf("convert '%s'/'%s' to '%s' success\n",
+                   params.model_path.c_str(),
+                   params.vae_path.c_str(),
+                   params.output_path.c_str());
+            return 0;
+        }
+    }
+
     bool vae_decode_only        = true;
     uint8_t* input_image_buffer = NULL;
     if (params.mode == IMG2IMG) {
@@ -510,14 +570,17 @@ int main(int argc, const char* argv[]) {
     sd_ctx_t* sd_ctx = new_sd_ctx(params.model_path.c_str(),
                                   params.vae_path.c_str(),
                                   params.taesd_path.c_str(),
+                                  params.controlnet_path.c_str(),
                                   params.lora_model_dir.c_str(),
+                                  params.embeddings_path.c_str(),
                                   vae_decode_only,
                                   params.vae_tiling,
                                   true,
                                   params.n_threads,
                                   params.wtype,
                                   params.rng_type,
-                                  params.schedule);
+                                  params.schedule,
+                                  params.control_net_cpu);
 
     if (sd_ctx == NULL) {
         printf("new_sd_ctx_t failed\n");
@@ -526,6 +589,23 @@ int main(int argc, const char* argv[]) {
 
     sd_image_t* results;
     if (params.mode == TXT2IMG) {
+        sd_image_t* control_image = NULL;
+        if (params.controlnet_path.size() > 0 && params.control_image_path.size() > 0) {
+            int c              = 0;
+            input_image_buffer = stbi_load(params.control_image_path.c_str(), &params.width, &params.height, &c, 3);
+            if (input_image_buffer == NULL) {
+                fprintf(stderr, "load image from '%s' failed\n", params.control_image_path.c_str());
+                return 1;
+            }
+            control_image = new sd_image_t{(uint32_t)params.width,
+                                           (uint32_t)params.height,
+                                           3,
+                                           input_image_buffer};
+            if (params.canny_preprocess) {  // apply preprocessor
+                LOG_INFO("Applying canny preprocessor");
+                control_image->data = preprocess_canny(control_image->data, control_image->width, control_image->height);
+            }
+        }
         results = txt2img(sd_ctx,
                           params.prompt.c_str(),
                           params.negative_prompt.c_str(),
@@ -536,7 +616,9 @@ int main(int argc, const char* argv[]) {
                           params.sample_method,
                           params.sample_steps,
                           params.seed,
-                          params.batch_count);
+                          params.batch_count,
+                          control_image,
+                          params.control_strength);
     } else {
         sd_image_t input_image = {(uint32_t)params.width,
                                   (uint32_t)params.height,
@@ -560,6 +642,7 @@ int main(int argc, const char* argv[]) {
 
     if (results == NULL) {
         printf("generate failed\n");
+        free_sd_ctx(sd_ctx);
         return 1;
     }
 
@@ -600,6 +683,8 @@ int main(int argc, const char* argv[]) {
         free(results[i].data);
         results[i].data = NULL;
     }
+    free(results);
+    free_sd_ctx(sd_ctx);
 
     return 0;
 }

ggml_extend.hpp

@@ -219,7 +219,7 @@ __STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
     for (int iy = 0; iy < height; iy++) {
         for (int ix = 0; ix < width; ix++) {
             for (int k = 0; k < channels; k++) {
-                float value = *(image_data + iy * width * channels + ix * channels + k);
+                int value = *(image_data + iy * width * channels + ix * channels + k);
                 ggml_tensor_set_f32(output, value / 255.0f, ix, iy, k);
             }
         }
@@ -449,7 +449,7 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_group_norm(struct ggml_context* ct
                                                          struct ggml_tensor* w,
                                                          struct ggml_tensor* b,
                                                          int num_groups = 32) {
-    if (x->n_dims == 4) {
+    if (ggml_n_dims(x) >= 3) {
         w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
         b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
     }
@@ -462,8 +462,12 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_group_norm(struct ggml_context* ct
 
 __STATIC_INLINE__ void ggml_backend_tensor_get_and_sync(ggml_backend_t backend, const struct ggml_tensor* tensor, void* data, size_t offset, size_t size) {
 #ifdef SD_USE_CUBLAS
-    ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
-    ggml_backend_synchronize(backend);
+    if (!ggml_backend_is_cpu(backend)) {
+        ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
+        ggml_backend_synchronize(backend);
+    } else {
+        ggml_backend_tensor_get(tensor, data, offset, size);
+    }
 #else
     ggml_backend_tensor_get(tensor, data, offset, size);
 #endif
@@ -544,7 +548,7 @@ struct GGMLModule {
     bool alloc_params_buffer(ggml_backend_t backend_, ggml_type wtype_ = GGML_TYPE_F32) {
         backend            = backend_;
         wtype              = wtype_;
-        params_buffer_size = 10 * 1024 * 1024;  // 10 MB, for padding
+        params_buffer_size = 4 * 1024 * 1024;  // 10 MB, for padding
         params_buffer_size += calculate_mem_size();
         size_t num_tensors = get_num_tensors();
 

lora.hpp

@@ -113,7 +113,7 @@ struct LoraModel : public GGMLModule {
             applied_lora_tensors.insert(scale_name);
 
             // calc_cale
-            int64_t dim       = lora_down->ne[lora_down->n_dims - 1];
+            int64_t dim       = lora_down->ne[ggml_n_dims(lora_down) - 1];
             float scale_value = 1.0f;
             if (lora_tensors.find(scale_name) != lora_tensors.end()) {
                 scale_value = ggml_backend_tensor_get_f32(lora_tensors[scale_name]);
@@ -123,17 +123,10 @@ struct LoraModel : public GGMLModule {
             }
             scale_value *= multiplier;
 
-            ggml_tensor* lora_scale = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
-
-            ggml_allocr_alloc(compute_allocr, lora_scale);
-            if (!ggml_allocr_is_measure(compute_allocr)) {
-                ggml_backend_tensor_set(lora_scale, &scale_value, 0, ggml_nbytes(lora_scale));
-            }
-
             // flat lora tensors to multiply it
-            int64_t lora_up_rows   = lora_up->ne[lora_up->n_dims - 1];
+            int64_t lora_up_rows   = lora_up->ne[ggml_n_dims(lora_up) - 1];
             lora_up                = ggml_reshape_2d(ctx0, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
-            int64_t lora_down_rows = lora_down->ne[lora_down->n_dims - 1];
+            int64_t lora_down_rows = lora_down->ne[ggml_n_dims(lora_down) - 1];
             lora_down              = ggml_reshape_2d(ctx0, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
 
             // ggml_mul_mat requires tensor b transposed
@@ -142,7 +135,7 @@ struct LoraModel : public GGMLModule {
             updown                     = ggml_cont(ctx0, ggml_transpose(ctx0, updown));
             updown                     = ggml_reshape(ctx0, updown, weight);
             GGML_ASSERT(ggml_nelements(updown) == ggml_nelements(weight));
-            updown = ggml_scale_inplace(ctx0, updown, lora_scale);
+            updown = ggml_scale_inplace(ctx0, updown, scale_value);
             ggml_tensor* final_weight;
             // if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) {
             //     final_weight = ggml_new_tensor(ctx0, GGML_TYPE_F32, weight->n_dims, weight->ne);

model.cpp

@@ -15,6 +15,8 @@
 #include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
 
+#include "stable-diffusion.h"
+
 #ifdef SD_USE_METAL
 #include "ggml-metal.h"
 #endif
@@ -87,7 +89,6 @@ const char* unused_tensors[] = {
     "model_ema.decay",
     "model_ema.num_updates",
     "model_ema.diffusion_model",
-    "control_model",
     "embedding_manager",
     "denoiser.sigmas",
 };
@@ -374,6 +375,11 @@ std::string convert_tensor_name(const std::string& name) {
         new_name = convert_open_clip_to_hf_clip(name);
     } else if (starts_with(name, "first_stage_model.decoder")) {
         new_name = convert_vae_decoder_name(name);
+    } else if (starts_with(name, "control_model.")) {  // for controlnet pth models
+        size_t pos = name.find('.');
+        if (pos != std::string::npos) {
+            new_name = name.substr(pos + 1);
+        }
     } else if (starts_with(name, "lora_")) {  // for lora
         size_t pos = name.find('.');
         if (pos != std::string::npos) {
@@ -609,7 +615,7 @@ bool is_safetensors_file(const std::string& file_path) {
     }
 
     size_t header_size_ = read_u64(header_size_buf);
-    if (header_size_ >= file_size_) {
+    if (header_size_ >= file_size_ || header_size_ <= 2) {
         return false;
     }
 
@@ -673,7 +679,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
 
         // LOG_DEBUG("%s", name.c_str());
 
-        TensorStorage tensor_storage(prefix + name, dummy->type, dummy->ne, dummy->n_dims, file_index, offset);
+        TensorStorage tensor_storage(prefix + name, dummy->type, dummy->ne, ggml_n_dims(dummy), file_index, offset);
 
         GGML_ASSERT(ggml_nbytes(dummy) == tensor_storage.nbytes());
 
@@ -1181,6 +1187,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("conditioner.embedders.1") != std::string::npos) {
             return VERSION_XL;
         }
+        if (tensor_storage.name.find("cond_stage_model.1") != std::string::npos) {
+            return VERSION_XL;
+        }
         if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
             tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
             tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
@@ -1218,7 +1227,35 @@ std::string ModelLoader::load_merges() {
     return merges_utf8_str;
 }
 
+void remove_duplicates(std::vector<TensorStorage>& vec) {
+    std::unordered_map<std::string, size_t> name_to_index_map;
+
+    for (size_t i = 0; i < vec.size(); ++i) {
+        const std::string& current_name = vec[i].name;
+        auto it                         = name_to_index_map.find(current_name);
+
+        if (it != name_to_index_map.end()) {
+            vec[it->second] = vec[i];
+        } else {
+            name_to_index_map[current_name] = i;
+        }
+    }
+
+    vec.resize(name_to_index_map.size());
+}
+
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend_t backend) {
+    std::vector<TensorStorage> processed_tensor_storages;
+    for (auto& tensor_storage : tensor_storages) {
+        // LOG_DEBUG("%s", name.c_str());
+
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+
+        preprocess_tensor(tensor_storage, processed_tensor_storages);
+    }
+    remove_duplicates(processed_tensor_storages);
     bool success = true;
     for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
         std::string file_path = file_paths_[file_index];
@@ -1276,22 +1313,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
             return true;
         };
 
-        std::vector<TensorStorage> processed_tensor_storages;
-        for (auto& tensor_storage : tensor_storages) {
+        for (auto& tensor_storage : processed_tensor_storages) {
             if (tensor_storage.file_index != file_index) {
                 continue;
             }
-
-            // LOG_DEBUG("%s", name.c_str());
-
-            if (is_unused_tensor(tensor_storage.name)) {
-                continue;
-            }
-
-            preprocess_tensor(tensor_storage, processed_tensor_storages);
-        }
-
-        for (auto& tensor_storage : processed_tensor_storages) {
             // LOG_DEBUG("%s", tensor_storage.name.c_str());
 
             ggml_tensor* dst_tensor = NULL;
@@ -1308,11 +1333,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
 
             size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
-            if (dst_tensor->buffer == NULL || ggml_backend_is_cpu(backend)
-#ifdef SD_USE_METAL
-                || ggml_backend_is_metal(backend)
-#endif
-            ) {
+            if (dst_tensor->buffer == NULL || ggml_backend_buffer_is_host(dst_tensor->buffer)) {
                 // for the CPU and Metal backend, we can copy directly into the tensor
                 if (tensor_storage.type == dst_tensor->type) {
                     GGML_ASSERT(ggml_nbytes(dst_tensor) == tensor_storage.nbytes());
@@ -1417,6 +1438,9 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
         if (pair.first.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos) {
             continue;
         }
+        if (pair.first.find("alphas_cumprod") != std::string::npos) {
+            continue;
+        }
 
         if (pair.first.find("alphas_cumprod") != std::string::npos) {
             continue;
@@ -1434,7 +1458,61 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
     return true;
 }
 
-int64_t ModelLoader::cal_mem_size(ggml_backend_t backend) {
+bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type) {
+    auto backend    = ggml_backend_cpu_init();
+    size_t mem_size = 1 * 1024 * 1024;  // for padding
+    mem_size += tensor_storages.size() * ggml_tensor_overhead();
+    mem_size += cal_mem_size(backend, type);
+    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
+    ggml_context* ggml_ctx = ggml_init({mem_size, NULL, false});
+
+    gguf_context* gguf_ctx = gguf_init_empty();
+
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+
+        ggml_type tensor_type = tensor_storage.type;
+        if (type != GGML_TYPE_COUNT) {
+            if (ggml_is_quantized(type) && tensor_storage.ne[0] % 32 != 0) {
+                tensor_type = GGML_TYPE_F16;
+            } else {
+                tensor_type = type;
+            }
+        }
+
+        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
+        if (tensor == NULL) {
+            LOG_ERROR("ggml_new_tensor failed");
+            return false;
+        }
+        ggml_set_name(tensor, name.c_str());
+
+        // LOG_DEBUG("%s %d %s %d[%d %d %d %d] %d[%d %d %d %d]", name.c_str(),
+        // ggml_nbytes(tensor), ggml_type_name(tensor_type),
+        // tensor_storage.n_dims,
+        // tensor_storage.ne[0], tensor_storage.ne[1], tensor_storage.ne[2], tensor_storage.ne[3],
+        // tensor->n_dims, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+
+        *dst_tensor = tensor;
+
+        gguf_add_tensor(gguf_ctx, tensor);
+
+        return true;
+    };
+
+    bool success = load_tensors(on_new_tensor_cb, backend);
+    ggml_backend_free(backend);
+    LOG_INFO("load tensors done");
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    if (success) {
+        gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
+    }
+    ggml_free(ggml_ctx);
+    gguf_free(gguf_ctx);
+    return success;
+}
+
+int64_t ModelLoader::cal_mem_size(ggml_backend_t backend, ggml_type type) {
     size_t alignment = 128;
     if (backend != NULL) {
         alignment = ggml_backend_get_alignment(backend);
@@ -1449,8 +1527,35 @@ int64_t ModelLoader::cal_mem_size(ggml_backend_t backend) {
     }
 
     for (auto& tensor_storage : processed_tensor_storages) {
+        ggml_type tensor_type = tensor_storage.type;
+        if (type != GGML_TYPE_COUNT) {
+            if (ggml_is_quantized(type) && tensor_storage.ne[0] % 32 != 0) {
+                tensor_type = GGML_TYPE_F16;
+            } else {
+                tensor_type = type;
+            }
+        }
+        tensor_storage.type = tensor_type;
         mem_size += tensor_storage.nbytes() + alignment;
     }
 
     return mem_size;
 }
+
+bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type) {
+    ModelLoader model_loader;
+
+    if (!model_loader.init_from_file(input_path)) {
+        LOG_ERROR("init model loader from file failed: '%s'", input_path);
+        return false;
+    }
+
+    if (vae_path != NULL && strlen(vae_path) > 0) {
+        if (!model_loader.init_from_file(vae_path, "vae.")) {
+            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
+            return false;
+        }
+    }
+    bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type);
+    return success;
+}

model.h

@@ -4,6 +4,7 @@
 #include <functional>
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -92,7 +93,6 @@ struct TensorStorage {
 };
 
 typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
-typedef std::function<void(const std::string&, int32_t)> on_new_token_cb_t;
 
 class ModelLoader {
 protected:
@@ -120,7 +120,8 @@ public:
     bool load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
                       ggml_backend_t backend,
                       std::set<std::string> ignore_tensors = {});
-    int64_t cal_mem_size(ggml_backend_t backend);
+    bool save_to_gguf_file(const std::string& file_path, ggml_type type);
+    int64_t cal_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
     ~ModelLoader() = default;
 };
 #endif  // __MODEL_H__

preprocessing.hpp

@@ -0,0 +1,227 @@
+#ifndef __PREPROCESSING_HPP__
+#define __PREPROCESSING_HPP__
+
+#include "ggml_extend.hpp"
+#define M_PI_ 3.14159265358979323846
+
+void convolve(struct ggml_tensor* input, struct ggml_tensor* output, struct ggml_tensor* kernel, int padding) {
+    struct ggml_init_params params;
+    params.mem_size                 = 20 * 1024 * 1024;  // 10
+    params.mem_buffer               = NULL;
+    params.no_alloc                 = false;
+    struct ggml_context* ctx0       = ggml_init(params);
+    struct ggml_tensor* kernel_fp16 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, kernel->ne[0], kernel->ne[1], 1, 1);
+    ggml_fp32_to_fp16_row((float*)kernel->data, (ggml_fp16_t*)kernel_fp16->data, ggml_nelements(kernel));
+    ggml_tensor* h  = ggml_conv_2d(ctx0, kernel_fp16, input, 1, 1, padding, padding, 1, 1);
+    ggml_cgraph* gf = ggml_new_graph(ctx0);
+    ggml_build_forward_expand(gf, ggml_cpy(ctx0, h, output));
+    ggml_graph_compute_with_ctx(ctx0, gf, 1);
+    ggml_free(ctx0);
+}
+
+void gaussian_kernel(struct ggml_tensor* kernel) {
+    int ks_mid   = kernel->ne[0] / 2;
+    float sigma  = 1.4f;
+    float normal = 1.f / (2.0f * M_PI_ * powf(sigma, 2.0f));
+    for (int y = 0; y < kernel->ne[0]; y++) {
+        float gx = -ks_mid + y;
+        for (int x = 0; x < kernel->ne[1]; x++) {
+            float gy = -ks_mid + x;
+            float k_ = expf(-((gx * gx + gy * gy) / (2.0f * powf(sigma, 2.0f)))) * normal;
+            ggml_tensor_set_f32(kernel, k_, x, y);
+        }
+    }
+}
+
+void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
+    for (int iy = 0; iy < rgb_img->ne[1]; iy++) {
+        for (int ix = 0; ix < rgb_img->ne[0]; ix++) {
+            float r    = ggml_tensor_get_f32(rgb_img, ix, iy);
+            float g    = ggml_tensor_get_f32(rgb_img, ix, iy, 1);
+            float b    = ggml_tensor_get_f32(rgb_img, ix, iy, 2);
+            float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b;
+            ggml_tensor_set_f32(grayscale, gray, ix, iy);
+        }
+    }
+}
+
+void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
+    int n_elements = ggml_nelements(h);
+    float* dx      = (float*)x->data;
+    float* dy      = (float*)y->data;
+    float* dh      = (float*)h->data;
+    for (int i = 0; i < n_elements; i++) {
+        dh[i] = sqrtf(dx[i] * dx[i] + dy[i] * dy[i]);
+    }
+}
+
+void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
+    int n_elements = ggml_nelements(h);
+    float* dx      = (float*)x->data;
+    float* dy      = (float*)y->data;
+    float* dh      = (float*)h->data;
+    for (int i = 0; i < n_elements; i++) {
+        dh[i] = atan2f(dy[i], dx[i]);
+    }
+}
+
+void normalize_tensor(struct ggml_tensor* g) {
+    int n_elements = ggml_nelements(g);
+    float* dg      = (float*)g->data;
+    float max      = -INFINITY;
+    for (int i = 0; i < n_elements; i++) {
+        max = dg[i] > max ? dg[i] : max;
+    }
+    max = 1.0f / max;
+    for (int i = 0; i < n_elements; i++) {
+        dg[i] *= max;
+    }
+}
+
+void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struct ggml_tensor* D) {
+    for (int iy = 1; iy < result->ne[1] - 1; iy++) {
+        for (int ix = 1; ix < result->ne[0] - 1; ix++) {
+            float angle = ggml_tensor_get_f32(D, ix, iy) * 180.0f / M_PI_;
+            angle       = angle < 0.0f ? angle += 180.0f : angle;
+            float q     = 1.0f;
+            float r     = 1.0f;
+
+            // angle 0
+            if ((0 >= angle && angle < 22.5f) || (157.5f >= angle && angle <= 180)) {
+                q = ggml_tensor_get_f32(G, ix, iy + 1);
+                r = ggml_tensor_get_f32(G, ix, iy - 1);
+            }
+            // angle 45
+            else if (22.5f >= angle && angle < 67.5f) {
+                q = ggml_tensor_get_f32(G, ix + 1, iy - 1);
+                r = ggml_tensor_get_f32(G, ix - 1, iy + 1);
+            }
+            // angle 90
+            else if (67.5f >= angle && angle < 112.5) {
+                q = ggml_tensor_get_f32(G, ix + 1, iy);
+                r = ggml_tensor_get_f32(G, ix - 1, iy);
+            }
+            // angle 135
+            else if (112.5 >= angle && angle < 157.5f) {
+                q = ggml_tensor_get_f32(G, ix - 1, iy - 1);
+                r = ggml_tensor_get_f32(G, ix + 1, iy + 1);
+            }
+
+            float cur = ggml_tensor_get_f32(G, ix, iy);
+            if ((cur >= q) && (cur >= r)) {
+                ggml_tensor_set_f32(result, cur, ix, iy);
+            } else {
+                ggml_tensor_set_f32(result, 0.0f, ix, iy);
+            }
+        }
+    }
+}
+
+void threshold_hystersis(struct ggml_tensor* img, float highThreshold, float lowThreshold, float weak, float strong) {
+    int n_elements = ggml_nelements(img);
+    float* imd     = (float*)img->data;
+    float max      = -INFINITY;
+    for (int i = 0; i < n_elements; i++) {
+        max = imd[i] > max ? imd[i] : max;
+    }
+    float ht = max * highThreshold;
+    float lt = ht * lowThreshold;
+    for (int i = 0; i < n_elements; i++) {
+        float img_v = imd[i];
+        if (img_v >= ht) {  // strong pixel
+            imd[i] = strong;
+        } else if (img_v <= ht && img_v >= lt) {  // strong pixel
+            imd[i] = weak;
+        }
+    }
+
+    for (int iy = 0; iy < img->ne[1]; iy++) {
+        for (int ix = 0; ix < img->ne[0]; ix++) {
+            if (ix >= 3 && ix <= img->ne[0] - 3 && iy >= 3 && iy <= img->ne[1] - 3) {
+                ggml_tensor_set_f32(img, ggml_tensor_get_f32(img, ix, iy), ix, iy);
+            } else {
+                ggml_tensor_set_f32(img, 0.0f, ix, iy);
+            }
+        }
+    }
+
+    // hysteresis
+    for (int iy = 1; iy < img->ne[1] - 1; iy++) {
+        for (int ix = 1; ix < img->ne[0] - 1; ix++) {
+            float imd_v = ggml_tensor_get_f32(img, ix, iy);
+            if (imd_v == weak) {
+                if (ggml_tensor_get_f32(img, ix + 1, iy - 1) == strong || ggml_tensor_get_f32(img, ix + 1, iy) == strong ||
+                    ggml_tensor_get_f32(img, ix, iy - 1) == strong || ggml_tensor_get_f32(img, ix, iy + 1) == strong ||
+                    ggml_tensor_get_f32(img, ix - 1, iy - 1) == strong || ggml_tensor_get_f32(img, ix - 1, iy) == strong) {
+                    ggml_tensor_set_f32(img, strong, ix, iy);
+                } else {
+                    ggml_tensor_set_f32(img, 0.0f, ix, iy);
+                }
+            }
+        }
+    }
+}
+
+uint8_t* preprocess_canny(uint8_t* img, int width, int height, float highThreshold = 0.08f, float lowThreshold = 0.08f, float weak = 0.8f, float strong = 1.0f, bool inverse = false) {
+    struct ggml_init_params params;
+    params.mem_size               = static_cast<size_t>(10 * 1024 * 1024);  // 10
+    params.mem_buffer             = NULL;
+    params.no_alloc               = false;
+    struct ggml_context* work_ctx = ggml_init(params);
+
+    if (!work_ctx) {
+        LOG_ERROR("ggml_init() failed");
+        return NULL;
+    }
+
+    float kX[9] = {
+        -1, 0, 1,
+        -2, 0, 2,
+        -1, 0, 1};
+
+    float kY[9] = {
+        1, 2, 1,
+        0, 0, 0,
+        -1, -2, -1};
+
+    // generate kernel
+    int kernel_size             = 5;
+    struct ggml_tensor* gkernel = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, 1);
+    struct ggml_tensor* sf_kx   = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
+    memcpy(sf_kx->data, kX, ggml_nbytes(sf_kx));
+    struct ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
+    memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
+    gaussian_kernel(gkernel);
+    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
+    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 1, 1);
+    struct ggml_tensor* iX         = ggml_dup_tensor(work_ctx, image_gray);
+    struct ggml_tensor* iY         = ggml_dup_tensor(work_ctx, image_gray);
+    struct ggml_tensor* G          = ggml_dup_tensor(work_ctx, image_gray);
+    struct ggml_tensor* tetha      = ggml_dup_tensor(work_ctx, image_gray);
+    sd_image_to_tensor(img, image);
+    grayscale(image, image_gray);
+    convolve(image_gray, image_gray, gkernel, 2);
+    convolve(image_gray, iX, sf_kx, 1);
+    convolve(image_gray, iY, sf_ky, 1);
+    prop_hypot(iX, iY, G);
+    normalize_tensor(G);
+    prop_arctan2(iX, iY, tetha);
+    non_max_supression(image_gray, G, tetha);
+    threshold_hystersis(image_gray, highThreshold, lowThreshold, weak, strong);
+    // to RGB channels
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            float gray = ggml_tensor_get_f32(image_gray, ix, iy);
+            gray       = inverse ? 1.0f - gray : gray;
+            ggml_tensor_set_f32(image, gray, ix, iy);
+            ggml_tensor_set_f32(image, gray, ix, iy, 1);
+            ggml_tensor_set_f32(image, gray, ix, iy, 2);
+        }
+    }
+    free(img);
+    uint8_t* output = sd_tensor_to_image(image);
+    ggml_free(work_ctx);
+    return output;
+}
+
+#endif  // __PREPROCESSING_HPP__

stable-diffusion.cpp

@@ -7,6 +7,7 @@
 #include "util.h"
 
 #include "clip.hpp"
+#include "control.hpp"
 #include "denoiser.hpp"
 #include "esrgan.hpp"
 #include "lora.hpp"
@@ -80,6 +81,8 @@ public:
     TinyAutoEncoder tae_first_stage;
     std::string taesd_path;
 
+    ControlNet control_net;
+
     StableDiffusionGGML() = default;
 
     StableDiffusionGGML(int n_threads,
@@ -101,17 +104,19 @@ public:
     }
 
     ~StableDiffusionGGML() {
+        ggml_backend_free(backend);
     }
 
     bool load_from_file(const std::string& model_path,
                         const std::string& vae_path,
+                        const std::string control_net_path,
+                        const std::string embeddings_path,
                         const std::string& taesd_path,
-                        bool vae_tiling,
+                        bool vae_tiling_,
                         ggml_type wtype,
-                        schedule_t schedule) {
-        this->use_tiny_autoencoder = taesd_path.size() > 0;
-        this->taesd_path           = taesd_path;
-        this->vae_tiling           = vae_tiling;
+                        schedule_t schedule,
+                        bool control_net_cpu) {
+        use_tiny_autoencoder = taesd_path.size() > 0;
 #ifdef SD_USE_CUBLAS
         LOG_DEBUG("Using CUDA backend");
         backend = ggml_backend_cuda_init(0);
@@ -136,6 +141,8 @@ public:
         LOG_INFO("loading model from '%s'", model_path.c_str());
         ModelLoader model_loader;
 
+        vae_tiling = vae_tiling_;
+
         if (!model_loader.init_from_file(model_path)) {
             LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
             return false;
@@ -185,6 +192,8 @@ public:
             return false;
         }
 
+        cond_stage_model.text_model.embd_dir = embeddings_path;
+
         ggml_type vae_type = model_data_type;
         if (version == VERSION_XL) {
             vae_type = GGML_TYPE_F32;  // avoid nan, not work...
@@ -307,8 +316,23 @@ public:
             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]);
         }
+
         LOG_DEBUG("finished loaded file");
         ggml_free(ctx);
+
+        if (control_net_path.size() > 0) {
+            ggml_backend_t cn_backend = NULL;
+            if (control_net_cpu && !ggml_backend_is_cpu(backend)) {
+                LOG_DEBUG("ControlNet: Using CPU backend");
+                cn_backend = ggml_backend_cpu_init();
+            } else {
+                cn_backend = backend;
+            }
+            if (!control_net.load_from_file(control_net_path, cn_backend, GGML_TYPE_F16 /* just f16 controlnet models */)) {
+                return false;
+            }
+        }
+
         if (use_tiny_autoencoder) {
             return tae_first_stage.load_from_file(taesd_path, backend);
         }
@@ -328,8 +352,9 @@ public:
         ggml_set_f32(timesteps, 999);
         set_timestep_embedding(timesteps, t_emb, diffusion_model.model_channels);
         struct ggml_tensor* out = ggml_dup_tensor(work_ctx, x_t);
-        diffusion_model.alloc_compute_buffer(x_t, c, t_emb);
-        diffusion_model.compute(out, n_threads, x_t, NULL, c, t_emb);
+        std::vector<struct ggml_tensor*> controls;
+        diffusion_model.alloc_compute_buffer(x_t, c, controls, t_emb);
+        diffusion_model.compute(out, n_threads, x_t, NULL, c, controls, 1.0f, t_emb);
         diffusion_model.free_compute_buffer();
 
         double result = 0.f;
@@ -510,9 +535,11 @@ public:
                         ggml_tensor* c_vector,
                         ggml_tensor* uc,
                         ggml_tensor* uc_vector,
+                        ggml_tensor* control_hint,
                         float cfg_scale,
                         sample_method_t method,
-                        const std::vector<float>& sigmas) {
+                        const std::vector<float>& sigmas,
+                        float control_strength) {
         size_t steps = sigmas.size() - 1;
         // x_t = load_tensor_from_file(work_ctx, "./rand0.bin");
         // print_ggml_tensor(x_t);
@@ -522,7 +549,14 @@ public:
         struct ggml_tensor* noised_input = ggml_dup_tensor(work_ctx, x_t);
         struct ggml_tensor* timesteps    = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 1);                                     // [N, ]
         struct ggml_tensor* t_emb        = new_timestep_embedding(work_ctx, NULL, timesteps, diffusion_model.model_channels);  // [N, model_channels]
-        diffusion_model.alloc_compute_buffer(noised_input, c, t_emb, c_vector);
+        struct ggml_tensor* guided_hint  = NULL;
+        if (control_hint != NULL) {
+            guided_hint = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, noised_input->ne[0], noised_input->ne[1], diffusion_model.model_channels, 1);
+            control_net.process_hint(guided_hint, n_threads, control_hint);
+            control_net.alloc_compute_buffer(noised_input, guided_hint, c, t_emb);
+        }
+
+        diffusion_model.alloc_compute_buffer(noised_input, c, control_net.controls, t_emb, c_vector);
 
         bool has_unconditioned = cfg_scale != 1.0 && uc != NULL;
 
@@ -572,12 +606,19 @@ public:
             ggml_tensor_scale(noised_input, c_in);
 
             // cond
-            diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c, t_emb, c_vector);
+            if (control_hint != NULL) {
+                control_net.compute(n_threads, noised_input, guided_hint, c, t_emb);
+            }
+            diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c, control_net.controls, control_strength, t_emb, c_vector);
 
             float* negative_data = NULL;
             if (has_unconditioned) {
                 // uncond
-                diffusion_model.compute(out_uncond, n_threads, noised_input, NULL, uc, t_emb, uc_vector);
+                if (control_hint != NULL) {
+                    control_net.compute(n_threads, noised_input, guided_hint, uc, t_emb);
+                }
+
+                diffusion_model.compute(out_uncond, n_threads, noised_input, NULL, uc, control_net.controls, control_strength, t_emb, uc_vector);
                 negative_data = (float*)out_uncond->data;
             }
             float* vec_denoised  = (float*)denoised->data;
@@ -626,7 +667,7 @@ public:
 
                     // 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])));
+                                                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
@@ -802,7 +843,7 @@ public:
 
                     // 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])));
+                                                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); };
@@ -986,6 +1027,7 @@ public:
                 LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);
                 abort();
         }
+        control_net.free_compute_buffer();
         diffusion_model.free_compute_buffer();
         return x;
     }
@@ -1097,14 +1139,17 @@ struct sd_ctx_t {
 sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                      const char* vae_path_c_str,
                      const char* taesd_path_c_str,
+                     const char* control_net_path_c_str,
                      const char* lora_model_dir_c_str,
+                     const char* embed_dir_c_str,
                      bool vae_decode_only,
                      bool vae_tiling,
                      bool free_params_immediately,
                      int n_threads,
                      enum sd_type_t wtype,
                      enum rng_type_t rng_type,
-                     enum schedule_t s) {
+                     enum schedule_t s,
+                     bool keep_control_net_cpu) {
     sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
     if (sd_ctx == NULL) {
         return NULL;
@@ -1112,6 +1157,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
     std::string model_path(model_path_c_str);
     std::string vae_path(vae_path_c_str);
     std::string taesd_path(taesd_path_c_str);
+    std::string control_net_path(control_net_path_c_str);
+    std::string embd_path(embed_dir_c_str);
     std::string lora_model_dir(lora_model_dir_c_str);
 
     sd_ctx->sd = new StableDiffusionGGML(n_threads,
@@ -1125,10 +1172,13 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
 
     if (!sd_ctx->sd->load_from_file(model_path,
                                     vae_path,
+                                    control_net_path,
+                                    embd_path,
                                     taesd_path,
                                     vae_tiling,
                                     (ggml_type)wtype,
-                                    s)) {
+                                    s,
+                                    keep_control_net_cpu)) {
         delete sd_ctx->sd;
         sd_ctx->sd = NULL;
         free(sd_ctx);
@@ -1155,7 +1205,9 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                     enum sample_method_t sample_method,
                     int sample_steps,
                     int64_t seed,
-                    int batch_count) {
+                    int batch_count,
+                    const sd_image_t* control_cond,
+                    float control_strength) {
     LOG_DEBUG("txt2img %dx%d", width, height);
     if (sd_ctx == NULL) {
         return NULL;
@@ -1223,6 +1275,12 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
         sd_ctx->sd->cond_stage_model.free_params_buffer();
     }
 
+    struct ggml_tensor* image_hint = NULL;
+    if (control_cond != NULL) {
+        image_hint = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
+        sd_image_to_tensor(control_cond->data, image_hint);
+    }
+
     std::vector<struct ggml_tensor*> final_latents;  // collect latents to decode
     int C = 4;
     int W = width / 8;
@@ -1239,7 +1297,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
 
         std::vector<float> sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps);
 
-        struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx, x_t, NULL, c, c_vector, uc, uc_vector, cfg_scale, sample_method, sigmas);
+        struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx, x_t, NULL, c, c_vector, uc, uc_vector, image_hint, cfg_scale, sample_method, sigmas, control_strength);
         // struct ggml_tensor* x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
         // print_ggml_tensor(x_0);
         int64_t sampling_end = ggml_time_ms();
@@ -1392,8 +1450,8 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
     ggml_tensor_set_f32_randn(noise, sd_ctx->sd->rng);
 
     LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]);
-    struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx, init_latent, noise, c, c_vector, uc, uc_vector,
-                                                 cfg_scale, sample_method, sigma_sched);
+    struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx, init_latent, noise, c, c_vector, uc,
+                                                 uc_vector, NULL, cfg_scale, sample_method, sigma_sched, 1.0f);
     // struct ggml_tensor *x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
     // print_ggml_tensor(x_0);
     int64_t t3 = ggml_time_ms();

stable-diffusion.h

@@ -65,12 +65,13 @@ enum sd_type_t {
     SD_TYPE_Q8_0 = 8,
     SD_TYPE_Q8_1 = 9,
     // k-quantizations
-    SD_TYPE_Q2_K = 10,
-    SD_TYPE_Q3_K = 11,
-    SD_TYPE_Q4_K = 12,
-    SD_TYPE_Q5_K = 13,
-    SD_TYPE_Q6_K = 14,
-    SD_TYPE_Q8_K = 15,
+    SD_TYPE_Q2_K    = 10,
+    SD_TYPE_Q3_K    = 11,
+    SD_TYPE_Q4_K    = 12,
+    SD_TYPE_Q5_K    = 13,
+    SD_TYPE_Q6_K    = 14,
+    SD_TYPE_Q8_K    = 15,
+    SD_TYPE_IQ2_XXS = 16,
     SD_TYPE_I8,
     SD_TYPE_I16,
     SD_TYPE_I32,
@@ -104,14 +105,17 @@ typedef struct sd_ctx_t sd_ctx_t;
 SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
                             const char* vae_path,
                             const char* taesd_path,
+                            const char* control_net_path_c_str,
                             const char* lora_model_dir,
+                            const char* embed_dir_c_str,
                             bool vae_decode_only,
                             bool vae_tiling,
                             bool free_params_immediately,
                             int n_threads,
                             enum sd_type_t wtype,
                             enum rng_type_t rng_type,
-                            enum schedule_t s);
+                            enum schedule_t s,
+                            bool keep_control_net_cpu);
 
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
 
@@ -125,7 +129,9 @@ SD_API sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                            enum sample_method_t sample_method,
                            int sample_steps,
                            int64_t seed,
-                           int batch_count);
+                           int batch_count,
+                           const sd_image_t* control_cond,
+                           float control_strength);
 
 SD_API sd_image_t* img2img(sd_ctx_t* sd_ctx,
                            sd_image_t init_image,
@@ -148,7 +154,9 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
                                         enum sd_type_t wtype);
 SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
-SD_API sd_image_t upscale(upscaler_ctx_t*, sd_image_t input_image, uint32_t upscale_factor);
+SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor);
+
+SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type);
 
 #ifdef __cplusplus
 }

tae.hpp

@@ -278,9 +278,6 @@ struct TinyDecoder {
     ggml_tensor* conv_final_w;  // [output_channels, channels, 3, 3]
     ggml_tensor* conv_final_b;  // [output_channels]
 
-    ggml_tensor* in_scale_1d3;  // [1]
-    ggml_tensor* in_scale_3;    // [1]
-
     TinyDecoder() {
         for (int i = 0; i < num_blocks; i++) {
             input_blocks[i].in_channels  = channels;
@@ -351,16 +348,6 @@ struct TinyDecoder {
         }
 
         final_block.init_params(ctx);
-
-        // initialize constants scales
-        in_scale_1d3 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        in_scale_3   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, in_scale_1d3);
-        float scale_1d3 = 1.0f / 3.0f;
-        ggml_backend_tensor_set(in_scale_1d3, &scale_1d3, 0, sizeof(scale_1d3));
-        ggml_allocr_alloc(alloc, in_scale_3);
-        float scale_3 = 3.0f;
-        ggml_backend_tensor_set(in_scale_3, &scale_3, 0, sizeof(scale_3));
     }
 
     void map_by_name(std::map<std::string, ggml_tensor*>& tensors, std::string prefix) {
@@ -391,9 +378,9 @@ struct TinyDecoder {
 
     ggml_tensor* forward(ggml_context* ctx, ggml_tensor* z) {
         // torch.tanh(x / 3) * 3
-        auto h = ggml_scale(ctx, z, in_scale_1d3);
+        auto h = ggml_scale(ctx, z, 1.0f / 3.0f);
         h      = ggml_tanh_inplace(ctx, h);
-        h      = ggml_scale(ctx, h, in_scale_3);
+        h      = ggml_scale(ctx, h, 3.0f);
 
         // conv(4, 64)
         h = ggml_nn_conv_2d(ctx, h, conv_input_w, conv_input_b, 1, 1, 1, 1);

unet.hpp

@@ -3,476 +3,12 @@
 
 #include "common.hpp"
 #include "ggml_extend.hpp"
+#include "model.h"
 
 /*==================================================== UnetModel =====================================================*/
 
 #define UNET_GRAPH_SIZE 10240
 
-struct ResBlock {
-    // network hparams
-    int channels;      // model_channels * (1, 1, 1, 2, 2, 4, 4, 4)
-    int emb_channels;  // time_embed_dim
-    int out_channels;  // mult * model_channels
-
-    // network params
-    // in_layers
-    struct ggml_tensor* in_layer_0_w;  // [channels, ]
-    struct ggml_tensor* in_layer_0_b;  // [channels, ]
-    // in_layer_1 is nn.SILU()
-    struct ggml_tensor* in_layer_2_w;  // [out_channels, channels, 3, 3]
-    struct ggml_tensor* in_layer_2_b;  // [out_channels, ]
-
-    // emb_layers
-    // emb_layer_0 is nn.SILU()
-    struct ggml_tensor* emb_layer_1_w;  // [out_channels, emb_channels]
-    struct ggml_tensor* emb_layer_1_b;  // [out_channels, ]
-
-    // out_layers
-    struct ggml_tensor* out_layer_0_w;  // [out_channels, ]
-    struct ggml_tensor* out_layer_0_b;  // [out_channels, ]
-    // out_layer_1 is nn.SILU()
-    // out_layer_2 is nn.Dropout(), p = 0 for inference
-    struct ggml_tensor* out_layer_3_w;  // [out_channels, out_channels, 3, 3]
-    struct ggml_tensor* out_layer_3_b;  // [out_channels, ]
-
-    // skip connection, only if out_channels != channels
-    struct ggml_tensor* skip_w;  // [out_channels, channels, 1, 1]
-    struct ggml_tensor* skip_b;  // [out_channels, ]
-
-    size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
-        mem_size += 2 * channels * ggml_type_sizef(GGML_TYPE_F32);                         // in_layer_0_w/b
-        mem_size += out_channels * channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);      // in_layer_2_w
-        mem_size += 5 * out_channels * ggml_type_sizef(GGML_TYPE_F32);                     // in_layer_2_b/emb_layer_1_b/out_layer_0_w/out_layer_0_b/out_layer_3_b
-        mem_size += out_channels * emb_channels * ggml_type_sizef(wtype);                  // emb_layer_1_w
-        mem_size += out_channels * out_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // out_layer_3_w
-
-        if (out_channels != channels) {
-            mem_size += out_channels * channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // skip_w
-            mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32);                     // skip_b
-        }
-        return static_cast<size_t>(mem_size);
-    }
-
-    void init_params(struct ggml_context* ctx, ggml_type wtype) {
-        in_layer_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
-        in_layer_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
-        in_layer_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels);
-        in_layer_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-
-        emb_layer_1_w = ggml_new_tensor_2d(ctx, wtype, emb_channels, out_channels);
-        emb_layer_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-
-        out_layer_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-        out_layer_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-        out_layer_3_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels);
-        out_layer_3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-
-        if (out_channels != channels) {
-            skip_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, channels, out_channels);
-            skip_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-        }
-    }
-
-    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
-        tensors[prefix + "in_layers.0.weight"] = in_layer_0_w;
-        tensors[prefix + "in_layers.0.bias"]   = in_layer_0_b;
-        tensors[prefix + "in_layers.2.weight"] = in_layer_2_w;
-        tensors[prefix + "in_layers.2.bias"]   = in_layer_2_b;
-
-        tensors[prefix + "emb_layers.1.weight"] = emb_layer_1_w;
-        tensors[prefix + "emb_layers.1.bias"]   = emb_layer_1_b;
-
-        tensors[prefix + "out_layers.0.weight"] = out_layer_0_w;
-        tensors[prefix + "out_layers.0.bias"]   = out_layer_0_b;
-        tensors[prefix + "out_layers.3.weight"] = out_layer_3_w;
-        tensors[prefix + "out_layers.3.bias"]   = out_layer_3_b;
-
-        if (out_channels != channels) {
-            tensors[prefix + "skip_connection.weight"] = skip_w;
-            tensors[prefix + "skip_connection.bias"]   = skip_b;
-        }
-    }
-
-    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* emb) {
-        // x: [N, channels, h, w]
-        // emb: [N, emb_channels]
-
-        // in_layers
-        auto h = ggml_nn_group_norm(ctx, x, in_layer_0_w, in_layer_0_b);
-        h      = ggml_silu_inplace(ctx, h);
-        h      = ggml_nn_conv_2d(ctx, h, in_layer_2_w, in_layer_2_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
-
-        // emb_layers
-        auto emb_out = ggml_silu(ctx, emb);
-        emb_out      = ggml_nn_linear(ctx, emb_out, emb_layer_1_w, emb_layer_1_b);           // [N, out_channels]
-        emb_out      = ggml_reshape_4d(ctx, emb_out, 1, 1, emb_out->ne[0], emb_out->ne[1]);  // [N, out_channels, 1, 1]
-
-        // out_layers
-        h = ggml_add(ctx, h, emb_out);
-        h = ggml_nn_group_norm(ctx, h, out_layer_0_w, out_layer_0_b);
-        h = ggml_silu_inplace(ctx, h);
-
-        // dropout, skip for inference
-
-        h = ggml_nn_conv_2d(ctx, h, out_layer_3_w, out_layer_3_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
-
-        // skip connection
-        if (out_channels != channels) {
-            x = ggml_nn_conv_2d(ctx, x, skip_w, skip_b);  // [N, out_channels, h, w]
-        }
-
-        h = ggml_add(ctx, h, x);
-        return h;  // [N, out_channels, h, w]
-    }
-};
-
-struct SpatialTransformer {
-    int in_channels;        // mult * model_channels
-    int n_head;             // num_heads
-    int d_head;             // in_channels // n_heads
-    int depth       = 1;    // 1
-    int context_dim = 768;  // hidden_size, 1024 for VERSION_2_x
-
-    // group norm
-    struct ggml_tensor* norm_w;  // [in_channels,]
-    struct ggml_tensor* norm_b;  // [in_channels,]
-
-    // proj_in
-    struct ggml_tensor* proj_in_w;  // [in_channels, in_channels, 1, 1]
-    struct ggml_tensor* proj_in_b;  // [in_channels,]
-
-    // transformer
-    struct Transformer {
-        // layer norm 1
-        struct ggml_tensor* norm1_w;  // [in_channels, ]
-        struct ggml_tensor* norm1_b;  // [in_channels, ]
-
-        // attn1
-        struct ggml_tensor* attn1_q_w;  // [in_channels, in_channels]
-        struct ggml_tensor* attn1_k_w;  // [in_channels, in_channels]
-        struct ggml_tensor* attn1_v_w;  // [in_channels, in_channels]
-
-        struct ggml_tensor* attn1_out_w;  // [in_channels, in_channels]
-        struct ggml_tensor* attn1_out_b;  // [in_channels, ]
-
-        // layer norm 2
-        struct ggml_tensor* norm2_w;  // [in_channels, ]
-        struct ggml_tensor* norm2_b;  // [in_channels, ]
-
-        // attn2
-        struct ggml_tensor* attn2_q_w;  // [in_channels, in_channels]
-        struct ggml_tensor* attn2_k_w;  // [in_channels, context_dim]
-        struct ggml_tensor* attn2_v_w;  // [in_channels, context_dim]
-
-        struct ggml_tensor* attn2_out_w;  // [in_channels, in_channels]
-        struct ggml_tensor* attn2_out_b;  // [in_channels, ]
-
-        // layer norm 3
-        struct ggml_tensor* norm3_w;  // [in_channels, ]
-        struct ggml_tensor* norm3_b;  // [in_channels, ]
-
-        // ff
-        struct ggml_tensor* ff_0_proj_w;  // [in_channels * 4 * 2, in_channels]
-        struct ggml_tensor* ff_0_proj_b;  // [in_channels * 4 * 2]
-
-        struct ggml_tensor* ff_2_w;  // [in_channels, in_channels * 4]
-        struct ggml_tensor* ff_2_b;  // [in_channels,]
-    };
-
-    std::vector<Transformer> transformers;
-
-    struct ggml_tensor* attn_scale;
-
-    // proj_out
-    struct ggml_tensor* proj_out_w;  // [in_channels, in_channels, 1, 1]
-    struct ggml_tensor* proj_out_b;  // [in_channels,]
-
-    SpatialTransformer(int depth = 1)
-        : depth(depth) {
-        transformers.resize(depth);
-    }
-
-    int get_num_tensors() {
-        return depth * 20 + 7;
-    }
-
-    size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
-        mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                        // norm_w/norm_b
-        mem_size += 2 * in_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // proj_in_w/proj_out_w
-        mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                        // proj_in_b/proj_out_b
-        mem_size += 1 * ggml_type_sizef(GGML_TYPE_F32);                                      // attn_scale
-
-        // transformer
-        for (auto& transformer : transformers) {
-            mem_size += 6 * in_channels * ggml_type_sizef(GGML_TYPE_F32);            // norm1-3_w/b
-            mem_size += 6 * in_channels * in_channels * ggml_type_sizef(wtype);      // attn1_q/k/v/out_w attn2_q/out_w
-            mem_size += 2 * in_channels * context_dim * ggml_type_sizef(wtype);      // attn2_k/v_w
-            mem_size += in_channels * 4 * 2 * in_channels * ggml_type_sizef(wtype);  // ff_0_proj_w
-            mem_size += in_channels * 4 * 2 * ggml_type_sizef(GGML_TYPE_F32);        // ff_0_proj_b
-            mem_size += in_channels * 4 * in_channels * ggml_type_sizef(wtype);      // ff_2_w
-            mem_size += in_channels * ggml_type_sizef(GGML_TYPE_F32);                // ff_2_b
-        }
-        return static_cast<size_t>(mem_size);
-    }
-
-    void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
-        norm_w    = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-        norm_b    = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-        proj_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
-        proj_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-        proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
-        proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-        attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, attn_scale);
-        float scale = 1.0f / sqrt((float)d_head);
-        ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));
-
-        // transformer
-        for (auto& transformer : transformers) {
-            transformer.norm1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-            transformer.norm1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-            transformer.attn1_q_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-            transformer.attn1_k_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-            transformer.attn1_v_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-
-            transformer.attn1_out_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-            transformer.attn1_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-            transformer.norm2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-            transformer.norm2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-            transformer.attn2_q_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-            transformer.attn2_k_w = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
-            transformer.attn2_v_w = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
-
-            transformer.attn2_out_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-            transformer.attn2_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-            transformer.norm3_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-            transformer.norm3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-            transformer.ff_0_proj_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels * 4 * 2);
-            transformer.ff_0_proj_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels * 4 * 2);
-
-            transformer.ff_2_w = ggml_new_tensor_2d(ctx, wtype, in_channels * 4, in_channels);
-            transformer.ff_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-        }
-    }
-
-    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
-        tensors[prefix + "norm.weight"]    = norm_w;
-        tensors[prefix + "norm.bias"]      = norm_b;
-        tensors[prefix + "proj_in.weight"] = proj_in_w;
-        tensors[prefix + "proj_in.bias"]   = proj_in_b;
-
-        // transformer
-        for (int i = 0; i < transformers.size(); i++) {
-            auto& transformer                                 = transformers[i];
-            std::string transformer_prefix                    = prefix + "transformer_blocks." + std::to_string(i) + ".";
-            tensors[transformer_prefix + "attn1.to_q.weight"] = transformer.attn1_q_w;
-            tensors[transformer_prefix + "attn1.to_k.weight"] = transformer.attn1_k_w;
-            tensors[transformer_prefix + "attn1.to_v.weight"] = transformer.attn1_v_w;
-
-            tensors[transformer_prefix + "attn1.to_out.0.weight"] = transformer.attn1_out_w;
-            tensors[transformer_prefix + "attn1.to_out.0.bias"]   = transformer.attn1_out_b;
-
-            tensors[transformer_prefix + "ff.net.0.proj.weight"] = transformer.ff_0_proj_w;
-            tensors[transformer_prefix + "ff.net.0.proj.bias"]   = transformer.ff_0_proj_b;
-            tensors[transformer_prefix + "ff.net.2.weight"]      = transformer.ff_2_w;
-            tensors[transformer_prefix + "ff.net.2.bias"]        = transformer.ff_2_b;
-
-            tensors[transformer_prefix + "attn2.to_q.weight"] = transformer.attn2_q_w;
-            tensors[transformer_prefix + "attn2.to_k.weight"] = transformer.attn2_k_w;
-            tensors[transformer_prefix + "attn2.to_v.weight"] = transformer.attn2_v_w;
-
-            tensors[transformer_prefix + "attn2.to_out.0.weight"] = transformer.attn2_out_w;
-            tensors[transformer_prefix + "attn2.to_out.0.bias"]   = transformer.attn2_out_b;
-
-            tensors[transformer_prefix + "norm1.weight"] = transformer.norm1_w;
-            tensors[transformer_prefix + "norm1.bias"]   = transformer.norm1_b;
-            tensors[transformer_prefix + "norm2.weight"] = transformer.norm2_w;
-            tensors[transformer_prefix + "norm2.bias"]   = transformer.norm2_b;
-            tensors[transformer_prefix + "norm3.weight"] = transformer.norm3_w;
-            tensors[transformer_prefix + "norm3.bias"]   = transformer.norm3_b;
-        }
-
-        tensors[prefix + "proj_out.weight"] = proj_out_w;
-        tensors[prefix + "proj_out.bias"]   = proj_out_b;
-    }
-
-    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* context) {
-        // x: [N, in_channels, h, w]
-        // context: [N, max_position, hidden_size(aka context_dim)]
-        auto x_in = x;
-        x         = ggml_nn_group_norm(ctx, x, norm_w, norm_b);
-        // proj_in
-        x = ggml_nn_conv_2d(ctx, x, proj_in_w, proj_in_b);  // [N, in_channels, h, w]
-
-        // transformer
-        const int64_t n            = x->ne[3];
-        const int64_t c            = x->ne[2];
-        const int64_t h            = x->ne[1];
-        const int64_t w            = x->ne[0];
-        const int64_t max_position = context->ne[1];
-        x                          = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3));  // [N, h, w, in_channels]
-
-        for (auto& transformer : transformers) {
-            auto r = x;
-            // layer norm 1
-            x = ggml_reshape_2d(ctx, x, c, w * h * n);
-            x = ggml_nn_layer_norm(ctx, x, transformer.norm1_w, transformer.norm1_b);
-
-            // self-attention
-            {
-                x                     = ggml_reshape_2d(ctx, x, c, h * w * n);        // [N * h * w, in_channels]
-                struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn1_q_w, x);  // [N * h * w, in_channels]
-#if !defined(SD_USE_FLASH_ATTENTION) || defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)
-                q = ggml_scale_inplace(ctx, q, attn_scale);
-#endif
-                q = ggml_reshape_4d(ctx, q, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
-                q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));    // [N, n_head, h * w, d_head]
-                q = ggml_reshape_3d(ctx, q, d_head, h * w, n_head * n);  // [N * n_head, h * w, d_head]
-
-                struct ggml_tensor* k = ggml_mul_mat(ctx, transformer.attn1_k_w, x);         // [N * h * w, in_channels]
-                k                     = ggml_reshape_4d(ctx, k, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
-                k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));    // [N, n_head, h * w, d_head]
-                k                     = ggml_reshape_3d(ctx, k, d_head, h * w, n_head * n);  // [N * n_head, h * w, d_head]
-
-                struct ggml_tensor* v = ggml_mul_mat(ctx, transformer.attn1_v_w, x);         // [N * h * w, in_channels]
-                v                     = ggml_reshape_4d(ctx, v, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
-                v                     = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));    // [N, n_head, d_head, h * w]
-                v                     = ggml_reshape_3d(ctx, v, h * w, d_head, n_head * n);  // [N * n_head, d_head, h * w]
-
-#if defined(SD_USE_FLASH_ATTENTION) && !defined(SD_USE_CUBLAS) && !defined(SD_USE_METAL)
-                struct ggml_tensor* kqv = ggml_flash_attn(ctx, q, k, v, false);  // [N * n_head, h * w, d_head]
-#else
-                struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, h * w, h * w]
-                // kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);
-                kq = ggml_soft_max_inplace(ctx, kq);
-
-                struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, h * w, d_head]
-#endif
-                kqv = ggml_reshape_4d(ctx, kqv, d_head, h * w, n_head, n);
-                kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));  // [N, h * w, n_head, d_head]
-
-                // x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n));
-                x = ggml_reshape_2d(ctx, kqv, d_head * n_head, h * w * n);
-
-                x = ggml_nn_linear(ctx, x, transformer.attn1_out_w, transformer.attn1_out_b);
-
-                x = ggml_reshape_4d(ctx, x, c, w, h, n);
-            }
-
-            x = ggml_add(ctx, x, r);
-            r = x;
-
-            // layer norm 2
-            x = ggml_nn_layer_norm(ctx, x, transformer.norm2_w, transformer.norm2_b);
-
-            // cross-attention
-            {
-                x                     = ggml_reshape_2d(ctx, x, c, h * w * n);                                           // [N * h * w, in_channels]
-                context               = ggml_reshape_2d(ctx, context, context->ne[0], context->ne[1] * context->ne[2]);  // [N * max_position, hidden_size]
-                struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn2_q_w, x);                                     // [N * h * w, in_channels]
-#if !defined(SD_USE_FLASH_ATTENTION) || defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)
-                q = ggml_scale_inplace(ctx, q, attn_scale);
-#endif
-                q = ggml_reshape_4d(ctx, q, d_head, n_head, h * w, n);   // [N, h * w, n_head, d_head]
-                q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));    // [N, n_head, h * w, d_head]
-                q = ggml_reshape_3d(ctx, q, d_head, h * w, n_head * n);  // [N * n_head, h * w, d_head]
-
-                struct ggml_tensor* k = ggml_mul_mat(ctx, transformer.attn2_k_w, context);          // [N * max_position, in_channels]
-                k                     = ggml_reshape_4d(ctx, k, d_head, n_head, max_position, n);   // [N, max_position, n_head, d_head]
-                k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));           // [N, n_head, max_position, d_head]
-                k                     = ggml_reshape_3d(ctx, k, d_head, max_position, n_head * n);  // [N * n_head, max_position, d_head]
-
-                struct ggml_tensor* v = ggml_mul_mat(ctx, transformer.attn2_v_w, context);          // [N * max_position, in_channels]
-                v                     = ggml_reshape_4d(ctx, v, d_head, n_head, max_position, n);   // [N, max_position, n_head, d_head]
-                v                     = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));           // [N, n_head, d_head, max_position]
-                v                     = ggml_reshape_3d(ctx, v, max_position, d_head, n_head * n);  // [N * n_head, d_head, max_position]
-#if defined(SD_USE_FLASH_ATTENTION) && !defined(SD_USE_CUBLAS) && !defined(SD_USE_METAL)
-                struct ggml_tensor* kqv = ggml_flash_attn(ctx, q, k, v, false);  // [N * n_head, h * w, d_head]
-#else
-                struct ggml_tensor* kq  = ggml_mul_mat(ctx, k, q);   // [N * n_head, h * w, max_position]
-                // kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);
-                kq = ggml_soft_max_inplace(ctx, kq);
-
-                struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, h * w, d_head]
-#endif
-                kqv = ggml_reshape_4d(ctx, kqv, d_head, h * w, n_head, n);
-                kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));
-
-                // x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n)); // [N * h * w, in_channels]
-                x = ggml_reshape_2d(ctx, kqv, d_head * n_head, h * w * n);  // [N * h * w, in_channels]
-
-                x = ggml_nn_linear(ctx, x, transformer.attn2_out_w, transformer.attn2_out_b);
-
-                x = ggml_reshape_4d(ctx, x, c, w, h, n);
-            }
-
-            x = ggml_add(ctx, x, r);
-            r = x;
-
-            // layer norm 3
-            x = ggml_reshape_2d(ctx, x, c, h * w * n);  // [N * h * w, in_channels]
-            x = ggml_nn_layer_norm(ctx, x, transformer.norm3_w, transformer.norm3_b);
-
-            // ff
-            {
-                // GEGLU
-                auto x_w    = ggml_view_2d(ctx,
-                                        transformer.ff_0_proj_w,
-                                        transformer.ff_0_proj_w->ne[0],
-                                        transformer.ff_0_proj_w->ne[1] / 2,
-                                        transformer.ff_0_proj_w->nb[1],
-                                        0);  // [in_channels * 4, in_channels]
-                auto x_b    = ggml_view_1d(ctx,
-                                        transformer.ff_0_proj_b,
-                                        transformer.ff_0_proj_b->ne[0] / 2,
-                                        0);  // [in_channels * 4, in_channels]
-                auto gate_w = ggml_view_2d(ctx,
-                                           transformer.ff_0_proj_w,
-                                           transformer.ff_0_proj_w->ne[0],
-                                           transformer.ff_0_proj_w->ne[1] / 2,
-                                           transformer.ff_0_proj_w->nb[1],
-                                           transformer.ff_0_proj_w->nb[1] * transformer.ff_0_proj_w->ne[1] / 2);  // [in_channels * 4, ]
-                auto gate_b = ggml_view_1d(ctx,
-                                           transformer.ff_0_proj_b,
-                                           transformer.ff_0_proj_b->ne[0] / 2,
-                                           transformer.ff_0_proj_b->nb[0] * transformer.ff_0_proj_b->ne[0] / 2);  // [in_channels * 4, ]
-                x           = ggml_reshape_2d(ctx, x, c, w * h * n);
-                auto x_in   = x;
-                x           = ggml_nn_linear(ctx, x_in, x_w, x_b);        // [N * h * w, in_channels * 4]
-                auto gate   = ggml_nn_linear(ctx, x_in, gate_w, gate_b);  // [N * h * w, in_channels * 4]
-
-                gate = ggml_gelu_inplace(ctx, gate);
-
-                x = ggml_mul(ctx, x, gate);  // [N * h * w, in_channels * 4]
-                // fc
-                x = ggml_nn_linear(ctx, x, transformer.ff_2_w, transformer.ff_2_b);  // [N * h * w, in_channels]
-            }
-
-            x = ggml_reshape_4d(ctx, x, c, w, h, n);  // [N, h, w, in_channels]
-
-            // residual
-            x = ggml_add(ctx, x, r);
-        }
-
-        x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3));  // [N, in_channels, h, w]
-
-        // proj_out
-        x = ggml_nn_conv_2d(ctx, x, proj_out_w, proj_out_b);  // [N, in_channels, h, w]
-
-        x = ggml_add(ctx, x, x_in);
-        return x;
-    }
-};
-
 // ldm.modules.diffusionmodules.openaimodel.UNetModel
 struct UNetModel : public GGMLModule {
     SDVersion version = VERSION_1_x;
@@ -644,21 +180,21 @@ struct UNetModel : public GGMLModule {
     }
 
     size_t calculate_mem_size() {
-        double mem_size = 0;
-        mem_size += time_embed_dim * model_channels * ggml_type_sizef(wtype);  // time_embed_0_w
-        mem_size += time_embed_dim * ggml_type_sizef(GGML_TYPE_F32);           // time_embed_0_b
-        mem_size += time_embed_dim * time_embed_dim * ggml_type_sizef(wtype);  // time_embed_2_w
-        mem_size += time_embed_dim * ggml_type_sizef(GGML_TYPE_F32);           // time_embed_2_b
+        size_t mem_size = 0;
+        mem_size += ggml_row_size(wtype, time_embed_dim * model_channels);  // time_embed_0_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);           // time_embed_0_b
+        mem_size += ggml_row_size(wtype, time_embed_dim * time_embed_dim);  // time_embed_2_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);           // time_embed_2_b
 
         if (version == VERSION_XL) {
-            mem_size += time_embed_dim * adm_in_channels * ggml_type_sizef(wtype);  // label_embed_0_w
-            mem_size += time_embed_dim * ggml_type_sizef(GGML_TYPE_F32);            // label_embed_0_b
-            mem_size += time_embed_dim * time_embed_dim * ggml_type_sizef(wtype);   // label_embed_2_w
-            mem_size += time_embed_dim * ggml_type_sizef(GGML_TYPE_F32);            // label_embed_2_b
+            mem_size += ggml_row_size(wtype, time_embed_dim * adm_in_channels);  // label_embed_0_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);            // label_embed_0_b
+            mem_size += ggml_row_size(wtype, time_embed_dim * time_embed_dim);   // label_embed_2_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);            // label_embed_2_b
         }
 
-        mem_size += model_channels * in_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // input_block_0_w
-        mem_size += model_channels * ggml_type_sizef(GGML_TYPE_F32);                        // input_block_0_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, model_channels * in_channels * 3 * 3);  // input_block_0_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, model_channels);                        // input_block_0_b
 
         // input_blocks
         int ds           = 1;
@@ -699,11 +235,11 @@ struct UNetModel : public GGMLModule {
         }
 
         // out
-        mem_size += 2 * model_channels * ggml_type_sizef(GGML_TYPE_F32);                     // out_0_w/b
-        mem_size += out_channels * model_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // out_2_w
-        mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32);                           // out_2_b
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, model_channels);                     // out_0_w/b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * model_channels * 3 * 3);  // out_2_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                           // out_2_b
 
-        return static_cast<size_t>(mem_size);
+        return mem_size;
     }
 
     size_t get_num_tensors() {
@@ -900,6 +436,8 @@ struct UNetModel : public GGMLModule {
                                 struct ggml_tensor* x,
                                 struct ggml_tensor* timesteps,
                                 struct ggml_tensor* context,
+                                std::vector<struct ggml_tensor*> control,
+                                float control_net_strength,
                                 struct ggml_tensor* t_emb = NULL,
                                 struct ggml_tensor* y     = NULL) {
         // x: [N, in_channels, h, w]
@@ -957,12 +495,24 @@ struct UNetModel : public GGMLModule {
         h = middle_block_1.forward(ctx0, h, context);  // [N, 4*model_channels, h/8, w/8]
         h = middle_block_2.forward(ctx0, h, emb);      // [N, 4*model_channels, h/8, w/8]
 
+        if (control.size() > 0) {
+            auto cs = ggml_scale_inplace(ctx0, control[control.size() - 1], control_net_strength);
+            h       = ggml_add(ctx0, h, cs);  // middle control
+        }
+
+        int control_offset = control.size() - 2;
         // output_blocks
         for (int i = (int)len_mults - 1; i >= 0; i--) {
             for (int j = 0; j < num_res_blocks + 1; j++) {
                 auto h_skip = hs.back();
                 hs.pop_back();
 
+                if (control.size() > 0) {
+                    auto cs = ggml_scale_inplace(ctx0, control[control_offset], control_net_strength);
+                    h_skip  = ggml_add(ctx0, h_skip, cs);  // control net condition
+                    control_offset--;
+                }
+
                 h = ggml_concat(ctx0, h, h_skip);
                 h = output_res_blocks[i][j].forward(ctx0, h, emb);
 
@@ -991,8 +541,10 @@ struct UNetModel : public GGMLModule {
     struct ggml_cgraph* build_graph(struct ggml_tensor* x,
                                     struct ggml_tensor* timesteps,
                                     struct ggml_tensor* context,
-                                    struct ggml_tensor* t_emb = NULL,
-                                    struct ggml_tensor* y     = NULL) {
+                                    std::vector<struct ggml_tensor*> control,
+                                    struct ggml_tensor* t_emb  = NULL,
+                                    struct ggml_tensor* y      = NULL,
+                                    float control_net_strength = 1.0) {
         // since we are using ggml-alloc, this buffer only needs enough space to hold the ggml_tensor and ggml_cgraph structs, but not the tensor data
         static size_t buf_size = ggml_tensor_overhead() * UNET_GRAPH_SIZE + ggml_graph_overhead();
         static std::vector<uint8_t> buf(buf_size);
@@ -1014,6 +566,7 @@ struct UNetModel : public GGMLModule {
         struct ggml_tensor* context_t   = NULL;
         struct ggml_tensor* t_emb_t     = NULL;
         struct ggml_tensor* y_t         = NULL;
+        std::vector<struct ggml_tensor*> control_t;
 
         // it's performing a compute, check if backend isn't cpu
         if (!ggml_backend_is_cpu(backend)) {
@@ -1057,7 +610,22 @@ struct UNetModel : public GGMLModule {
             y_t         = y;
         }
 
-        struct ggml_tensor* out = forward(ctx0, x_t, timesteps_t, context_t, t_emb_t, y_t);
+        // offload all controls tensors to gpu
+        if (control.size() > 0 && !ggml_backend_is_cpu(backend) && control[0]->backend != GGML_BACKEND_GPU) {
+            for (int i = 0; i < control.size(); i++) {
+                ggml_tensor* cntl_t = ggml_dup_tensor(ctx0, control[i]);
+                control_t.push_back(cntl_t);
+                ggml_allocr_alloc(compute_allocr, cntl_t);
+                if (!ggml_allocr_is_measure(compute_allocr)) {
+                    ggml_backend_tensor_copy(control[i], control_t[i]);
+                    ggml_backend_synchronize(backend);
+                }
+            }
+        } else {
+            control_t = control;
+        }
+
+        struct ggml_tensor* out = forward(ctx0, x_t, timesteps_t, context_t, control_t, control_net_strength, t_emb_t, y_t);
 
         ggml_build_forward_expand(gf, out);
         ggml_free(ctx0);
@@ -1067,10 +635,12 @@ struct UNetModel : public GGMLModule {
 
     void alloc_compute_buffer(struct ggml_tensor* x,
                               struct ggml_tensor* context,
-                              struct ggml_tensor* t_emb = NULL,
-                              struct ggml_tensor* y     = NULL) {
+                              std::vector<struct ggml_tensor*> control,
+                              struct ggml_tensor* t_emb  = NULL,
+                              struct ggml_tensor* y      = NULL,
+                              float control_net_strength = 1.0) {
         auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(x, NULL, context, t_emb, y);
+            return build_graph(x, NULL, context, control, t_emb, y, control_net_strength);
         };
         GGMLModule::alloc_compute_buffer(get_graph);
     }
@@ -1080,10 +650,12 @@ struct UNetModel : public GGMLModule {
                  struct ggml_tensor* x,
                  struct ggml_tensor* timesteps,
                  struct ggml_tensor* context,
+                 std::vector<struct ggml_tensor*> control,
+                 float control_net_strength,
                  struct ggml_tensor* t_emb = NULL,
                  struct ggml_tensor* y     = NULL) {
         auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(x, timesteps, context, t_emb, y);
+            return build_graph(x, timesteps, context, control, t_emb, y, control_net_strength);
         };
 
         GGMLModule::compute(get_graph, n_threads, work_latent);

util.cpp

@@ -1,6 +1,6 @@
 #include "util.h"
-
 #include <stdarg.h>
+#include <algorithm>
 #include <codecvt>
 #include <fstream>
 #include <locale>
@@ -72,6 +72,20 @@ bool is_directory(const std::string& path) {
     return (attributes != INVALID_FILE_ATTRIBUTES && (attributes & FILE_ATTRIBUTE_DIRECTORY));
 }
 
+std::string get_full_path(const std::string& dir, const std::string& filename) {
+    std::string full_path = dir + "\\" + filename;
+
+    WIN32_FIND_DATA find_file_data;
+    HANDLE hFind = FindFirstFile(full_path.c_str(), &find_file_data);
+
+    if (hFind != INVALID_HANDLE_VALUE) {
+        FindClose(hFind);
+        return full_path;
+    } else {
+        return "";
+    }
+}
+
 #else  // Unix
 #include <dirent.h>
 #include <sys/stat.h>
@@ -86,6 +100,25 @@ bool is_directory(const std::string& path) {
     return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
+std::string get_full_path(const std::string& dir, const std::string& filename) {
+    DIR* dp = opendir(dir.c_str());
+
+    if (dp != nullptr) {
+        struct dirent* entry;
+
+        while ((entry = readdir(dp)) != nullptr) {
+            if (strcasecmp(entry->d_name, filename.c_str()) == 0) {
+                closedir(dp);
+                return dir + "/" + entry->d_name;
+            }
+        }
+
+        closedir(dp);
+    }
+
+    return "";
+}
+
 #endif
 
 // get_num_physical_cores is copy from
@@ -192,6 +225,24 @@ void pretty_progress(int step, int steps, float time) {
     }
 }
 
+std::string ltrim(const std::string& s) {
+    auto it = std::find_if(s.begin(), s.end(), [](int ch) {
+        return !std::isspace(ch);
+    });
+    return std::string(it, s.end());
+}
+
+std::string rtrim(const std::string& s) {
+    auto it = std::find_if(s.rbegin(), s.rend(), [](int ch) {
+        return !std::isspace(ch);
+    });
+    return std::string(s.begin(), it.base());
+}
+
+std::string trim(const std::string& s) {
+    return rtrim(ltrim(s));
+}
+
 static sd_log_cb_t sd_log_cb = NULL;
 void* sd_log_cb_data         = NULL;
 

util.h

@@ -15,6 +15,7 @@ void replace_all_chars(std::string& str, char target, char replacement);
 
 bool file_exists(const std::string& filename);
 bool is_directory(const std::string& path);
+std::string get_full_path(const std::string& dir, const std::string& filename);
 
 std::u32string utf8_to_utf32(const std::string& utf8_str);
 std::string utf32_to_utf8(const std::u32string& utf32_str);
@@ -28,6 +29,8 @@ void pretty_progress(int step, int steps, float time);
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
 
+std::string trim(const std::string& s);
+
 #define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)

vae.hpp

@@ -32,14 +32,14 @@ struct ResnetBlock {
 
     size_t calculate_mem_size(ggml_type wtype) {
         double mem_size = 0;
-        mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                      // norm1_w/b
-        mem_size += out_channels * in_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);   // conv1_w
-        mem_size += 4 * out_channels * ggml_type_sizef(GGML_TYPE_F32);                     // conv1_b/norm2_w/norm2_b/conv2_b
-        mem_size += out_channels * out_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // conv2_w
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, in_channels);                      // norm1_w/b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * in_channels * 3 * 3);   // conv1_w
+        mem_size += 4 * ggml_row_size(GGML_TYPE_F32, out_channels);                     // conv1_b/norm2_w/norm2_b/conv2_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * out_channels * 3 * 3);  // conv2_w
 
         if (out_channels != in_channels) {
-            mem_size += out_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // nin_shortcut_w
-            mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32);                        // nin_shortcut_b
+            mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * in_channels * 1 * 1);  // nin_shortcut_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                        // nin_shortcut_b
         }
         return static_cast<size_t>(mem_size);
     }
@@ -118,12 +118,10 @@ struct AttnBlock {
     struct ggml_tensor* proj_out_w;  // [in_channels, in_channels, 1, 1]
     struct ggml_tensor* proj_out_b;  // [in_channels,]
 
-    struct ggml_tensor* attn_scale;
-
     size_t calculate_mem_size(ggml_type wtype) {
         double mem_size = 0;
-        mem_size += 6 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                        // norm_w/norm_b/q_b/k_v/v_b/proj_out_b
-        mem_size += 4 * in_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // q_w/k_w/v_w/proj_out_w                                            // object overhead
+        mem_size += 6 * ggml_row_size(GGML_TYPE_F32, in_channels);                        // norm_w/norm_b/q_b/k_v/v_b/proj_out_b
+        mem_size += 4 * ggml_row_size(GGML_TYPE_F16, in_channels * in_channels * 1 * 1);  // q_w/k_w/v_w/proj_out_w                                            // object overhead
         return static_cast<size_t>(mem_size);
     }
 
@@ -140,11 +138,6 @@ struct AttnBlock {
 
         proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
         proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-        attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, attn_scale);
-        float scale = 1.0f / sqrt((float)in_channels);
-        ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));
     }
 
     void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
@@ -181,7 +174,7 @@ struct AttnBlock {
         k = ggml_reshape_3d(ctx, k, c, h * w, n);              // [N, h * w, in_channels]
 
         auto w_ = ggml_mul_mat(ctx, k, q);  // [N, h * w, h * w]
-        w_      = ggml_scale_inplace(ctx, w_, attn_scale);
+        w_      = ggml_scale_inplace(ctx, w_, 1.0f / sqrt((float)in_channels));
         w_      = ggml_soft_max_inplace(ctx, w_);
 
         v  = ggml_reshape_3d(ctx, v, h * w, c, n);               // [N, in_channels, h * w]
@@ -276,17 +269,17 @@ struct Encoder {
     }
 
     size_t calculate_mem_size(ggml_type wtype) {
-        double mem_size = 0;
+        size_t mem_size = 0;
         int len_mults   = sizeof(ch_mult) / sizeof(int);
         int block_in    = ch * ch_mult[len_mults - 1];
 
-        mem_size += ch * in_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // conv_in_w
-        mem_size += ch * ggml_type_sizef(GGML_TYPE_F32);                        // conv_in_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, ch * in_channels * 3 * 3);  // conv_in_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, ch);                        // conv_in_b
 
-        mem_size += 2 * block_in * ggml_type_sizef(GGML_TYPE_F32);  // norm_out_w/b
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, block_in);  // norm_out_w/b
 
-        mem_size += z_channels * 2 * block_in * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // conv_out_w
-        mem_size += z_channels * 2 * ggml_type_sizef(GGML_TYPE_F32);                     // conv_out_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, z_channels * 2 * block_in * 3 * 3);  // conv_out_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, z_channels * 2);                     // conv_out_b
 
         mem_size += mid.block_1.calculate_mem_size(wtype);
         mem_size += mid.attn_1.calculate_mem_size(wtype);
@@ -301,7 +294,7 @@ struct Encoder {
             }
         }
 
-        return static_cast<size_t>(mem_size);
+        return mem_size;
     }
 
     void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
@@ -443,13 +436,13 @@ struct Decoder {
         int len_mults   = sizeof(ch_mult) / sizeof(int);
         int block_in    = ch * ch_mult[len_mults - 1];
 
-        mem_size += block_in * z_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // conv_in_w
-        mem_size += block_in * ggml_type_sizef(GGML_TYPE_F32);                       // conv_in_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, block_in * z_channels * 3 * 3);  // conv_in_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, block_in);                       // conv_in_b
 
-        mem_size += 2 * (ch * ch_mult[0]) * ggml_type_sizef(GGML_TYPE_F32);  // norm_out_w/b
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, (ch * ch_mult[0]));  // norm_out_w/b
 
-        mem_size += (ch * ch_mult[0]) * out_ch * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16);  // conv_out_w
-        mem_size += out_ch * ggml_type_sizef(GGML_TYPE_F32);                              // conv_out_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, (ch * ch_mult[0]) * out_ch * 3 * 3);  // conv_out_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, out_ch);                              // conv_out_b
 
         mem_size += mid.block_1.calculate_mem_size(wtype);
         mem_size += mid.attn_1.calculate_mem_size(wtype);
@@ -613,19 +606,19 @@ struct AutoEncoderKL : public GGMLModule {
     }
 
     size_t calculate_mem_size() {
-        double mem_size = 0;
+        size_t mem_size = 0;
 
         if (!decode_only) {
-            mem_size += 2 * embed_dim * 2 * dd_config.z_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // quant_conv_w
-            mem_size += 2 * embed_dim * ggml_type_sizef(GGML_TYPE_F32);                                     // quant_conv_b
+            mem_size += ggml_row_size(GGML_TYPE_F16, 2 * embed_dim * 2 * dd_config.z_channels * 1 * 1);  // quant_conv_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, 2 * embed_dim);                                     // quant_conv_b
             mem_size += encoder.calculate_mem_size(wtype);
         }
 
-        mem_size += dd_config.z_channels * embed_dim * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // post_quant_conv_w
-        mem_size += dd_config.z_channels * ggml_type_sizef(GGML_TYPE_F32);                      // post_quant_conv_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, dd_config.z_channels * embed_dim * 1 * 1);  // post_quant_conv_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, dd_config.z_channels);                      // post_quant_conv_b
 
         mem_size += decoder.calculate_mem_size(wtype);
-        return static_cast<size_t>(mem_size);
+        return mem_size;
     }
 
     size_t get_num_tensors() {