sync: update ggml (#180)

* Add in upscale repeater logic

---------

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,15 +43,30 @@ 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
+             control.hpp preprocessing.hpp)
+
 if(BUILD_SHARED_LIBS)
     message("Build shared library")
     add_definitions(-DSD_BUILD_SHARED_LIB)
-    add_definitions(-DSD_BUILD_DLL)
+    target_compile_definitions(${SD_LIB} PRIVATE -DSD_BUILD_DLL)
     set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 else()
     message("Build static library")
@@ -62,15 +74,15 @@ 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.
 
@@ -138,26 +148,31 @@ cmake --build . --config Release
 ### Run
 
 ```
-usage: ./bin/sd [arguments]
+usage: ./build/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.
+  --upscale-repeats                  Run the ESRGAN upscaler this many times (default 1)
   --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 +183,11 @@ 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)
+  --canny                            apply canny preprocessor (edge detection)
   -v, --verbose                      print extra info
 ```
 
@@ -183,6 +201,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 +268,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 +286,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.
 
@@ -303,6 +331,7 @@ Thank you to all the people who have already contributed to stable-diffusion.cpp
 - [stable-diffusion](https://github.com/CompVis/stable-diffusion)
 - [stable-diffusion-stability-ai](https://github.com/Stability-AI/stablediffusion)
 - [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
+- [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
 - [k-diffusion](https://github.com/crowsonkb/k-diffusion)
 - [latent-consistency-model](https://github.com/luosiallen/latent-consistency-model)
 - [generative-models](https://github.com/Stability-AI/generative-models/)

common.hpp

@@ -3,82 +3,525 @@
 
 #include "ggml_extend.hpp"
 
-struct DownSample {
-    // hparams
+class DownSampleBlock : public GGMLBlock {
+protected:
     int channels;
     int out_channels;
+    bool vae_downsample;
 
-    // conv2d params
-    struct ggml_tensor* op_w;  // [out_channels, channels, 3, 3]
-    struct ggml_tensor* op_b;  // [out_channels,]
-
-    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);
-    }
-
-    void init_params(struct ggml_context* ctx, ggml_type wtype) {
-        op_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels);
-        op_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) {
+public:
+    DownSampleBlock(int channels,
+                    int out_channels,
+                    bool vae_downsample = false)
+        : channels(channels),
+          out_channels(out_channels),
+          vae_downsample(vae_downsample) {
         if (vae_downsample) {
-            tensors[prefix + "conv.weight"] = op_w;
-            tensors[prefix + "conv.bias"]   = op_b;
+            blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {2, 2}, {0, 0}));
         } else {
-            tensors[prefix + "op.weight"] = op_w;
-            tensors[prefix + "op.bias"]   = op_b;
+            blocks["op"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {2, 2}, {1, 1}));
         }
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
-        struct ggml_tensor* c = NULL;
         if (vae_downsample) {
-            c = ggml_pad(ctx, x, 1, 1, 0, 0);
-            c = ggml_nn_conv_2d(ctx, c, op_w, op_b, 2, 2, 0, 0);
+            auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
+
+            x = ggml_pad(ctx, x, 1, 1, 0, 0);
+            x = conv->forward(ctx, x);
         } else {
-            c = ggml_nn_conv_2d(ctx, x, op_w, op_b, 2, 2, 1, 1);
+            auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["op"]);
+
+            x = conv->forward(ctx, x);
         }
-        return c;  // [N, out_channels, h/2, w/2]
+        return x;  // [N, out_channels, h/2, w/2]
     }
 };
 
-struct UpSample {
-    // hparams
+class UpSampleBlock : public GGMLBlock {
+protected:
     int channels;
     int out_channels;
 
-    // conv2d params
-    struct ggml_tensor* conv_w;  // [out_channels, channels, 3, 3]
-    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);
-    }
-
-    void init_params(struct ggml_context* ctx, ggml_type wtype) {
-        conv_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels);
-        conv_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 + "conv.weight"] = conv_w;
-        tensors[prefix + "conv.bias"]   = conv_b;
+public:
+    UpSampleBlock(int channels,
+                  int out_channels)
+        : channels(channels),
+          out_channels(out_channels) {
+        blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
-        x = ggml_upscale(ctx, x, 2);                              // [N, channels, h*2, w*2]
-        x = ggml_nn_conv_2d(ctx, x, conv_w, conv_b, 1, 1, 1, 1);  // [N, out_channels, h*2, w*2]
+        auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
+
+        x = ggml_upscale(ctx, x, 2);  // [N, channels, h*2, w*2]
+        x = conv->forward(ctx, x);    // [N, out_channels, h*2, w*2]
+        return x;
+    }
+};
+
+class ResBlock : public GGMLBlock {
+protected:
+    // network hparams
+    int64_t channels;      // model_channels * (1, 1, 1, 2, 2, 4, 4, 4)
+    int64_t emb_channels;  // time_embed_dim
+    int64_t out_channels;  // mult * model_channels
+    std::pair<int, int> kernel_size;
+    int dims;
+    bool skip_t_emb;
+    bool exchange_temb_dims;
+
+    std::shared_ptr<GGMLBlock> conv_nd(int dims,
+                                       int64_t in_channels,
+                                       int64_t out_channels,
+                                       std::pair<int, int> kernel_size,
+                                       std::pair<int, int> padding) {
+        GGML_ASSERT(dims == 2 || dims == 3);
+        if (dims == 3) {
+            return std::shared_ptr<GGMLBlock>(new Conv3dnx1x1(in_channels, out_channels, kernel_size.first, 1, padding.first));
+        } else {
+            return std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, out_channels, kernel_size, {1, 1}, padding));
+        }
+    }
+
+public:
+    ResBlock(int64_t channels,
+             int64_t emb_channels,
+             int64_t out_channels,
+             std::pair<int, int> kernel_size = {3, 3},
+             int dims                        = 2,
+             bool exchange_temb_dims         = false,
+             bool skip_t_emb                 = false)
+        : channels(channels),
+          emb_channels(emb_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          dims(dims),
+          skip_t_emb(skip_t_emb),
+          exchange_temb_dims(exchange_temb_dims) {
+        std::pair<int, int> padding = {kernel_size.first / 2, kernel_size.second / 2};
+        blocks["in_layers.0"]       = std::shared_ptr<GGMLBlock>(new GroupNorm32(channels));
+        // in_layer_1 is nn.SILU()
+        blocks["in_layers.2"] = conv_nd(dims, channels, out_channels, kernel_size, padding);
+
+        if (!skip_t_emb) {
+            // emb_layer_0 is nn.SILU()
+            blocks["emb_layers.1"] = std::shared_ptr<GGMLBlock>(new Linear(emb_channels, out_channels));
+        }
+
+        blocks["out_layers.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(out_channels));
+        // out_layer_1 is nn.SILU()
+        // out_layer_2 is nn.Dropout(), skip for inference
+        blocks["out_layers.3"] = conv_nd(dims, out_channels, out_channels, kernel_size, padding);
+
+        if (out_channels != channels) {
+            blocks["skip_connection"] = conv_nd(dims, channels, out_channels, {1, 1}, {0, 0});
+        }
+    }
+
+    virtual struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* emb = NULL) {
+        // For dims==3, we reduce dimension from 5d to 4d by merging h and w, in order not to change ggml
+        // [N, c, t, h, w] => [N, c, t, h * w]
+        // x: [N, channels, h, w] if dims == 2 else [N, channels, t, h, w]
+        // emb: [N, emb_channels] if dims == 2 else [N, t, emb_channels]
+        auto in_layers_0  = std::dynamic_pointer_cast<GroupNorm32>(blocks["in_layers.0"]);
+        auto in_layers_2  = std::dynamic_pointer_cast<UnaryBlock>(blocks["in_layers.2"]);
+        auto out_layers_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["out_layers.0"]);
+        auto out_layers_3 = std::dynamic_pointer_cast<UnaryBlock>(blocks["out_layers.3"]);
+
+        if (emb == NULL) {
+            GGML_ASSERT(skip_t_emb);
+        }
+
+        // in_layers
+        auto h = in_layers_0->forward(ctx, x);
+        h      = ggml_silu_inplace(ctx, h);
+        h      = in_layers_2->forward(ctx, h);  // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
+
+        // emb_layers
+        if (!skip_t_emb) {
+            auto emb_layer_1 = std::dynamic_pointer_cast<Linear>(blocks["emb_layers.1"]);
+
+            auto emb_out = ggml_silu(ctx, emb);
+            emb_out      = emb_layer_1->forward(ctx, emb_out);  // [N, out_channels] if dims == 2 else [N, t, out_channels]
+
+            if (dims == 2) {
+                emb_out = ggml_reshape_4d(ctx, emb_out, 1, 1, emb_out->ne[0], emb_out->ne[1]);  // [N, out_channels, 1, 1]
+            } else {
+                emb_out = ggml_reshape_4d(ctx, emb_out, 1, emb_out->ne[0], emb_out->ne[1], emb_out->ne[2]);  // [N, t, out_channels, 1]
+                if (exchange_temb_dims) {
+                    // emb_out = rearrange(emb_out, "b t c ... -> b c t ...")
+                    emb_out = ggml_cont(ctx, ggml_permute(ctx, emb_out, 0, 2, 1, 3));  // [N, out_channels, t, 1]
+                }
+            }
+
+            h = ggml_add(ctx, h, emb_out);  // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
+        }
+
+        // out_layers
+        h = out_layers_0->forward(ctx, h);
+        h = ggml_silu_inplace(ctx, h);
+        // dropout, skip for inference
+        h = out_layers_3->forward(ctx, h);
+
+        // skip connection
+        if (out_channels != channels) {
+            auto skip_connection = std::dynamic_pointer_cast<UnaryBlock>(blocks["skip_connection"]);
+            x                    = skip_connection->forward(ctx, x);  // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
+        }
+
+        h = ggml_add(ctx, h, x);
+        return h;  // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
+    }
+};
+
+class GEGLU : public GGMLBlock {
+protected:
+    int64_t dim_in;
+    int64_t dim_out;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        params["proj.weight"] = ggml_new_tensor_2d(ctx, wtype, dim_in, dim_out * 2);
+        params["proj.bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim_out * 2);
+    }
+
+public:
+    GEGLU(int64_t dim_in, int64_t dim_out)
+        : dim_in(dim_in), dim_out(dim_out) {}
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [ne3, ne2, ne1, dim_in]
+        // return: [ne3, ne2, ne1, dim_out]
+        struct ggml_tensor* w = params["proj.weight"];
+        struct ggml_tensor* b = params["proj.bias"];
+
+        auto x_w    = ggml_view_2d(ctx, w, w->ne[0], w->ne[1] / 2, w->nb[1], 0);                        // [dim_out, dim_in]
+        auto x_b    = ggml_view_1d(ctx, b, b->ne[0] / 2, 0);                                            // [dim_out, dim_in]
+        auto gate_w = ggml_view_2d(ctx, w, w->ne[0], w->ne[1] / 2, w->nb[1], w->nb[1] * w->ne[1] / 2);  // [dim_out, ]
+        auto gate_b = ggml_view_1d(ctx, b, b->ne[0] / 2, b->nb[0] * b->ne[0] / 2);                      // [dim_out, ]
+
+        auto x_in = x;
+        x         = ggml_nn_linear(ctx, x_in, x_w, x_b);        // [ne3, ne2, ne1, dim_out]
+        auto gate = ggml_nn_linear(ctx, x_in, gate_w, gate_b);  // [ne3, ne2, ne1, dim_out]
+
+        gate = ggml_gelu_inplace(ctx, gate);
+
+        x = ggml_mul(ctx, x, gate);  // [ne3, ne2, ne1, dim_out]
+
+        return x;
+    }
+};
+
+class FeedForward : public GGMLBlock {
+public:
+    FeedForward(int64_t dim,
+                int64_t dim_out,
+                int64_t mult = 4) {
+        int64_t inner_dim = dim * mult;
+
+        blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GEGLU(dim, inner_dim));
+        // net_1 is nn.Dropout(), skip for inference
+        blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out));
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [ne3, ne2, ne1, dim]
+        // return: [ne3, ne2, ne1, dim_out]
+
+        auto net_0 = std::dynamic_pointer_cast<GEGLU>(blocks["net.0"]);
+        auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
+
+        x = net_0->forward(ctx, x);  // [ne3, ne2, ne1, inner_dim]
+        x = net_2->forward(ctx, x);  // [ne3, ne2, ne1, dim_out]
+        return x;
+    }
+};
+
+class CrossAttention : public GGMLBlock {
+protected:
+    int64_t query_dim;
+    int64_t context_dim;
+    int64_t n_head;
+    int64_t d_head;
+
+public:
+    CrossAttention(int64_t query_dim,
+                   int64_t context_dim,
+                   int64_t n_head,
+                   int64_t d_head)
+        : n_head(n_head),
+          d_head(d_head),
+          query_dim(query_dim),
+          context_dim(context_dim) {
+        int64_t inner_dim = d_head * n_head;
+
+        blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, false));
+        blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(context_dim, inner_dim, false));
+        blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(context_dim, inner_dim, false));
+
+        blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, query_dim));
+        // to_out_1 is nn.Dropout(), skip for inference
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* context) {
+        // x: [N, n_token, query_dim]
+        // context: [N, n_context, context_dim]
+        // return: [N, n_token, query_dim]
+        auto to_q     = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
+        auto to_k     = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
+        auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
+        auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+
+        int64_t n         = x->ne[2];
+        int64_t n_token   = x->ne[1];
+        int64_t n_context = context->ne[1];
+        int64_t inner_dim = d_head * n_head;
+
+        auto q = to_q->forward(ctx, x);                                 // [N, n_token, inner_dim]
+        q      = ggml_reshape_4d(ctx, q, d_head, n_head, n_token, n);   // [N, n_token, n_head, d_head]
+        q      = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));      // [N, n_head, n_token, d_head]
+        q      = ggml_reshape_3d(ctx, q, d_head, n_token, n_head * n);  // [N * n_head, n_token, d_head]
+
+        auto k = to_k->forward(ctx, context);                             // [N, n_context, inner_dim]
+        k      = ggml_reshape_4d(ctx, k, d_head, n_head, n_context, n);   // [N, n_context, n_head, d_head]
+        k      = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));        // [N, n_head, n_context, d_head]
+        k      = ggml_reshape_3d(ctx, k, d_head, n_context, n_head * n);  // [N * n_head, n_context, d_head]
+
+        auto v = to_v->forward(ctx, context);                             // [N, n_context, inner_dim]
+        v      = ggml_reshape_4d(ctx, v, d_head, n_head, n_context, n);   // [N, n_context, n_head, d_head]
+        v      = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));        // [N, n_head, d_head, n_context]
+        v      = ggml_reshape_3d(ctx, v, n_context, d_head, n_head * n);  // [N * n_head, d_head, n_context]
+
+        auto kqv = ggml_nn_attention(ctx, q, k, v, false);  // [N * n_head, n_token, d_head]
+        kqv      = ggml_reshape_4d(ctx, kqv, d_head, n_token, n_head, n);
+        kqv      = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));  // [N, n_token, n_head, d_head]
+
+        x = ggml_reshape_3d(ctx, kqv, d_head * n_head, n_token, n);  // [N, n_token, inner_dim]
+
+        x = to_out_0->forward(ctx, x);  // [N, n_token, query_dim]
+        return x;
+    }
+};
+
+class BasicTransformerBlock : public GGMLBlock {
+protected:
+    int64_t n_head;
+    int64_t d_head;
+    bool ff_in;
+
+public:
+    BasicTransformerBlock(int64_t dim,
+                          int64_t n_head,
+                          int64_t d_head,
+                          int64_t context_dim,
+                          bool ff_in = false)
+        : n_head(n_head), d_head(d_head), ff_in(ff_in) {
+        // disable_self_attn is always False
+        // disable_temporal_crossattention is always False
+        // switch_temporal_ca_to_sa is always False
+        // inner_dim is always None or equal to dim
+        // gated_ff is always True
+        blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
+        blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head));
+        blocks["ff"]    = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
+        blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+        blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+        blocks["norm3"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+
+        if (ff_in) {
+            blocks["norm_in"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+            blocks["ff_in"]   = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
+        }
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* context) {
+        // x: [N, n_token, query_dim]
+        // context: [N, n_context, context_dim]
+        // return: [N, n_token, query_dim]
+
+        auto attn1 = std::dynamic_pointer_cast<CrossAttention>(blocks["attn1"]);
+        auto attn2 = std::dynamic_pointer_cast<CrossAttention>(blocks["attn2"]);
+        auto ff    = std::dynamic_pointer_cast<FeedForward>(blocks["ff"]);
+        auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+        auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
+        auto norm3 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm3"]);
+
+        if (ff_in) {
+            auto norm_in = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_in"]);
+            auto ff_in   = std::dynamic_pointer_cast<FeedForward>(blocks["ff_in"]);
+
+            auto x_skip = x;
+            x           = norm_in->forward(ctx, x);
+            x           = ff_in->forward(ctx, x);
+            // self.is_res is always True
+            x = ggml_add(ctx, x, x_skip);
+        }
+
+        auto r = x;
+        x      = norm1->forward(ctx, x);
+        x      = attn1->forward(ctx, x, x);  // self-attention
+        x      = ggml_add(ctx, x, r);
+        r      = x;
+        x      = norm2->forward(ctx, x);
+        x      = attn2->forward(ctx, x, context);  // cross-attention
+        x      = ggml_add(ctx, x, r);
+        r      = x;
+        x      = norm3->forward(ctx, x);
+        x      = ff->forward(ctx, x);
+        x      = ggml_add(ctx, x, r);
+
+        return x;
+    }
+};
+
+class SpatialTransformer : public GGMLBlock {
+protected:
+    int64_t in_channels;  // mult * model_channels
+    int64_t n_head;
+    int64_t d_head;
+    int64_t depth       = 1;    // 1
+    int64_t context_dim = 768;  // hidden_size, 1024 for VERSION_2_x
+
+public:
+    SpatialTransformer(int64_t in_channels,
+                       int64_t n_head,
+                       int64_t d_head,
+                       int64_t depth,
+                       int64_t context_dim)
+        : in_channels(in_channels),
+          n_head(n_head),
+          d_head(d_head),
+          depth(depth),
+          context_dim(context_dim) {
+        // We will convert unet transformer linear to conv2d 1x1 when loading the weights, so use_linear is always False
+        // disable_self_attn is always False
+        int64_t inner_dim = n_head * d_head;  // in_channels
+        blocks["norm"]    = std::shared_ptr<GGMLBlock>(new GroupNorm32(in_channels));
+        blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, inner_dim, {1, 1}));
+
+        for (int i = 0; i < depth; i++) {
+            std::string name = "transformer_blocks." + std::to_string(i);
+            blocks[name]     = std::shared_ptr<GGMLBlock>(new BasicTransformerBlock(inner_dim, n_head, d_head, context_dim));
+        }
+
+        blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(inner_dim, in_channels, {1, 1}));
+    }
+
+    virtual 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(aka n_token), hidden_size(aka context_dim)]
+        auto norm     = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
+        auto proj_in  = std::dynamic_pointer_cast<Conv2d>(blocks["proj_in"]);
+        auto proj_out = std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"]);
+
+        auto x_in         = x;
+        int64_t n         = x->ne[3];
+        int64_t h         = x->ne[1];
+        int64_t w         = x->ne[0];
+        int64_t inner_dim = n_head * d_head;
+
+        x = norm->forward(ctx, x);
+        x = proj_in->forward(ctx, x);  // [N, inner_dim, h, w]
+
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3));  // [N, h, w, inner_dim]
+        x = ggml_reshape_3d(ctx, x, inner_dim, w * h, n);      // [N, h * w, inner_dim]
+
+        for (int i = 0; i < depth; i++) {
+            std::string name       = "transformer_blocks." + std::to_string(i);
+            auto transformer_block = std::dynamic_pointer_cast<BasicTransformerBlock>(blocks[name]);
+
+            x = transformer_block->forward(ctx, x, context);
+        }
+
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));  // [N, inner_dim, h * w]
+        x = ggml_reshape_4d(ctx, x, w, h, inner_dim, n);       // [N, inner_dim, h, w]
+
+        // proj_out
+        x = proj_out->forward(ctx, x);  // [N, in_channels, h, w]
+
+        x = ggml_add(ctx, x, x_in);
+        return x;
+    }
+};
+
+class AlphaBlender : public GGMLBlock {
+protected:
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        params["mix_factor"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+    }
+
+    float get_alpha() {
+        // image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,]
+        // so learned_with_images is same as learned
+        float alpha = ggml_backend_tensor_get_f32(params["mix_factor"]);
+        return sigmoid(alpha);
+    }
+
+public:
+    AlphaBlender() {
+        // merge_strategy is always learned_with_images
+        // for inference, we don't need to set alpha
+        // since mix_factor.shape is [1,], we don't need rearrange using  rearrange_pattern
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx,
+                                struct ggml_tensor* x_spatial,
+                                struct ggml_tensor* x_temporal) {
+        // image_only_indicator is always tensor([0.])
+        float alpha = get_alpha();
+        auto x      = ggml_add(ctx,
+                               ggml_scale(ctx, x_spatial, alpha),
+                               ggml_scale(ctx, x_temporal, 1.0f - alpha));
+        return x;
+    }
+};
+
+class VideoResBlock : public ResBlock {
+public:
+    VideoResBlock(int channels,
+                  int emb_channels,
+                  int out_channels,
+                  std::pair<int, int> kernel_size = {3, 3},
+                  int64_t video_kernel_size       = 3,
+                  int dims                        = 2)  // always 2
+        : ResBlock(channels, emb_channels, out_channels, kernel_size, dims) {
+        blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, emb_channels, out_channels, kernel_size, 3, true));
+        blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* emb,
+                                int num_video_frames) {
+        // x: [N, channels, h, w] aka [b*t, channels, h, w]
+        // emb: [N, emb_channels] aka [b*t, emb_channels]
+        // image_only_indicator is always tensor([0.])
+        auto time_stack = std::dynamic_pointer_cast<ResBlock>(blocks["time_stack"]);
+        auto time_mixer = std::dynamic_pointer_cast<AlphaBlender>(blocks["time_mixer"]);
+
+        x = ResBlock::forward(ctx, x, emb);
+
+        int64_t T = num_video_frames;
+        int64_t B = x->ne[3] / T;
+        int64_t C = x->ne[2];
+        int64_t H = x->ne[1];
+        int64_t W = x->ne[0];
+
+        x          = ggml_reshape_4d(ctx, x, W * H, C, T, B);           // (b t) c h w -> b t c (h w)
+        x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // b t c (h w) -> b c t (h w)
+        auto x_mix = x;
+
+        emb = ggml_reshape_4d(ctx, emb, emb->ne[0], T, B, emb->ne[3]);  // (b t) ... -> b t ...
+
+        x = time_stack->forward(ctx, x, emb);  // b t c (h w)
+
+        x = time_mixer->forward(ctx, x_mix, x);  // b t c (h w)
+
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // b c t (h w) -> b t c (h w)
+        x = ggml_reshape_4d(ctx, x, W, H, C, T * B);           // b t c (h w) -> (b t) c h w
+
         return x;
     }
 };

control.hpp

@@ -0,0 +1,466 @@
+#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
+
+*/
+class ControlNetBlock : public GGMLBlock {
+protected:
+    SDVersion version = VERSION_1_x;
+    // network hparams
+    int in_channels                        = 4;
+    int out_channels                       = 4;
+    int hint_channels                      = 3;
+    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;  // 1024 for VERSION_2_x, 2048 for VERSION_XL
+
+public:
+    int model_channels  = 320;
+    int adm_in_channels = 2816;  // only for VERSION_XL
+
+    ControlNetBlock(SDVersion version = VERSION_1_x)
+        : version(version) {
+        if (version == VERSION_2_x) {
+            context_dim       = 1024;
+            num_head_channels = 64;
+            num_heads         = -1;
+        } else if (version == VERSION_XL) {
+            context_dim           = 2048;
+            attention_resolutions = {4, 2};
+            channel_mult          = {1, 2, 4};
+            transformer_depth     = {1, 2, 10};
+            num_head_channels     = 64;
+            num_heads             = -1;
+        } else if (version == VERSION_SVD) {
+            in_channels       = 8;
+            out_channels      = 4;
+            context_dim       = 1024;
+            adm_in_channels   = 768;
+            num_head_channels = 64;
+            num_heads         = -1;
+        }
+
+        blocks["time_embed.0"] = std::shared_ptr<GGMLBlock>(new Linear(model_channels, time_embed_dim));
+        // time_embed_1 is nn.SiLU()
+        blocks["time_embed.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
+
+        if (version == VERSION_XL || version == VERSION_SVD) {
+            blocks["label_emb.0.0"] = std::shared_ptr<GGMLBlock>(new Linear(adm_in_channels, time_embed_dim));
+            // label_emb_1 is nn.SiLU()
+            blocks["label_emb.0.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
+        }
+
+        // input_blocks
+        blocks["input_blocks.0.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, model_channels, {3, 3}, {1, 1}, {1, 1}));
+
+        std::vector<int> input_block_chans;
+        input_block_chans.push_back(model_channels);
+        int ch              = model_channels;
+        int input_block_idx = 0;
+        int ds              = 1;
+
+        auto get_resblock = [&](int64_t channels, int64_t emb_channels, int64_t out_channels) -> ResBlock* {
+            return new ResBlock(channels, emb_channels, out_channels);
+        };
+
+        auto get_attention_layer = [&](int64_t in_channels,
+                                       int64_t n_head,
+                                       int64_t d_head,
+                                       int64_t depth,
+                                       int64_t context_dim) -> SpatialTransformer* {
+            return new SpatialTransformer(in_channels, n_head, d_head, depth, context_dim);
+        };
+
+        auto make_zero_conv = [&](int64_t channels) {
+            return new Conv2d(channels, channels, {1, 1});
+        };
+
+        blocks["zero_convs.0.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(model_channels));
+
+        blocks["input_hint_block.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(hint_channels, 16, {3, 3}, {1, 1}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(16, 16, {3, 3}, {1, 1}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(16, 32, {3, 3}, {2, 2}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.6"] = std::shared_ptr<GGMLBlock>(new Conv2d(32, 32, {3, 3}, {1, 1}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.8"] = std::shared_ptr<GGMLBlock>(new Conv2d(32, 96, {3, 3}, {2, 2}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.10"] = std::shared_ptr<GGMLBlock>(new Conv2d(96, 96, {3, 3}, {1, 1}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.12"] = std::shared_ptr<GGMLBlock>(new Conv2d(96, 256, {3, 3}, {2, 2}, {1, 1}));
+        // nn.SiLU()
+        blocks["input_hint_block.14"] = std::shared_ptr<GGMLBlock>(new Conv2d(256, model_channels, {3, 3}, {1, 1}, {1, 1}));
+
+        size_t 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_block_idx += 1;
+                std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
+                blocks[name]     = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, 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;
+                    }
+                    std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
+                    blocks[name]     = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
+                                                                                      n_head,
+                                                                                      d_head,
+                                                                                      transformer_depth[i],
+                                                                                      context_dim));
+                }
+                blocks["zero_convs." + std::to_string(input_block_idx) + ".0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
+                input_block_chans.push_back(ch);
+            }
+            if (i != len_mults - 1) {
+                input_block_idx += 1;
+                std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
+                blocks[name]     = std::shared_ptr<GGMLBlock>(new DownSampleBlock(ch, ch));
+
+                blocks["zero_convs." + std::to_string(input_block_idx) + ".0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
+
+                input_block_chans.push_back(ch);
+                ds *= 2;
+            }
+        }
+
+        // middle blocks
+        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;
+        }
+        blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
+        blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
+                                                                                  n_head,
+                                                                                  d_head,
+                                                                                  transformer_depth[transformer_depth.size() - 1],
+                                                                                  context_dim));
+        blocks["middle_block.2"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
+
+        // middle_block_out
+        blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
+    }
+
+    struct ggml_tensor* resblock_forward(std::string name,
+                                         struct ggml_context* ctx,
+                                         struct ggml_tensor* x,
+                                         struct ggml_tensor* emb) {
+        auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
+        return block->forward(ctx, x, emb);
+    }
+
+    struct ggml_tensor* attention_layer_forward(std::string name,
+                                                struct ggml_context* ctx,
+                                                struct ggml_tensor* x,
+                                                struct ggml_tensor* context) {
+        auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
+        return block->forward(ctx, x, context);
+    }
+
+    struct ggml_tensor* input_hint_block_forward(struct ggml_context* ctx,
+                                                 struct ggml_tensor* hint,
+                                                 struct ggml_tensor* emb,
+                                                 struct ggml_tensor* context) {
+        int num_input_blocks = 15;
+        auto h               = hint;
+        for (int i = 0; i < num_input_blocks; i++) {
+            if (i % 2 == 0) {
+                auto block = std::dynamic_pointer_cast<Conv2d>(blocks["input_hint_block." + std::to_string(i)]);
+
+                h = block->forward(ctx, h);
+            } else {
+                h = ggml_silu_inplace(ctx, h);
+            }
+        }
+        return h;
+    }
+
+    std::vector<struct ggml_tensor*> forward(struct ggml_context* ctx,
+                                             struct ggml_tensor* x,
+                                             struct ggml_tensor* hint,
+                                             struct ggml_tensor* guided_hint,
+                                             struct ggml_tensor* timesteps,
+                                             struct ggml_tensor* context,
+                                             struct ggml_tensor* y = NULL) {
+        // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
+        // timesteps: [N,]
+        // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
+        // y: [N, adm_in_channels] or [1, adm_in_channels]
+        if (context != NULL) {
+            if (context->ne[2] != x->ne[3]) {
+                context = ggml_repeat(ctx, context, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, context->ne[0], context->ne[1], x->ne[3]));
+            }
+        }
+
+        if (y != NULL) {
+            if (y->ne[1] != x->ne[3]) {
+                y = ggml_repeat(ctx, y, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, y->ne[0], x->ne[3]));
+            }
+        }
+
+        auto time_embed_0     = std::dynamic_pointer_cast<Linear>(blocks["time_embed.0"]);
+        auto time_embed_2     = std::dynamic_pointer_cast<Linear>(blocks["time_embed.2"]);
+        auto input_blocks_0_0 = std::dynamic_pointer_cast<Conv2d>(blocks["input_blocks.0.0"]);
+        auto zero_convs_0     = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs.0.0"]);
+
+        auto middle_block_out = std::dynamic_pointer_cast<Conv2d>(blocks["middle_block_out.0"]);
+
+        auto t_emb = ggml_nn_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
+
+        auto emb = time_embed_0->forward(ctx, t_emb);
+        emb      = ggml_silu_inplace(ctx, emb);
+        emb      = time_embed_2->forward(ctx, emb);  // [N, time_embed_dim]
+
+        // SDXL/SVD
+        if (y != NULL) {
+            auto label_embed_0 = std::dynamic_pointer_cast<Linear>(blocks["label_emb.0.0"]);
+            auto label_embed_2 = std::dynamic_pointer_cast<Linear>(blocks["label_emb.0.2"]);
+
+            auto label_emb = label_embed_0->forward(ctx, y);
+            label_emb      = ggml_silu_inplace(ctx, label_emb);
+            label_emb      = label_embed_2->forward(ctx, label_emb);  // [N, time_embed_dim]
+
+            emb = ggml_add(ctx, emb, label_emb);  // [N, time_embed_dim]
+        }
+
+        std::vector<struct ggml_tensor*> outs;
+
+        if (guided_hint == NULL) {
+            guided_hint = input_hint_block_forward(ctx, hint, emb, context);
+        }
+        outs.push_back(guided_hint);
+
+        // input_blocks
+
+        // input block 0
+        auto h = input_blocks_0_0->forward(ctx, x);
+        h      = ggml_add(ctx, h, guided_hint);
+        outs.push_back(zero_convs_0->forward(ctx, h));
+
+        // input block 1-11
+        size_t len_mults    = channel_mult.size();
+        int input_block_idx = 0;
+        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++) {
+                input_block_idx += 1;
+                std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
+                h                = resblock_forward(name, ctx, h, emb);  // [N, mult*model_channels, h, w]
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
+                    h                = attention_layer_forward(name, ctx, h, context);  // [N, mult*model_channels, h, w]
+                }
+
+                auto zero_conv = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs." + std::to_string(input_block_idx) + ".0"]);
+
+                outs.push_back(zero_conv->forward(ctx, h));
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                input_block_idx += 1;
+
+                std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
+                auto block       = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
+
+                h = block->forward(ctx, h);  // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
+
+                auto zero_conv = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs." + std::to_string(input_block_idx) + ".0"]);
+
+                outs.push_back(zero_conv->forward(ctx, h));
+            }
+        }
+        // [N, 4*model_channels, h/8, w/8]
+
+        // middle_block
+        h = resblock_forward("middle_block.0", ctx, h, emb);             // [N, 4*model_channels, h/8, w/8]
+        h = attention_layer_forward("middle_block.1", ctx, h, context);  // [N, 4*model_channels, h/8, w/8]
+        h = resblock_forward("middle_block.2", ctx, h, emb);             // [N, 4*model_channels, h/8, w/8]
+
+        // out
+        outs.push_back(middle_block_out->forward(ctx, h));
+        return outs;
+    }
+};
+
+struct ControlNet : public GGMLModule {
+    SDVersion version = VERSION_1_x;
+    ControlNetBlock control_net;
+
+    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
+    struct ggml_tensor* guided_hint = NULL;     // guided_hint cache, for faster inference
+    bool guided_hint_cached         = false;
+
+    ControlNet(ggml_backend_t backend,
+               ggml_type wtype,
+               SDVersion version = VERSION_1_x)
+        : GGMLModule(backend, wtype), control_net(version) {
+        control_net.init(params_ctx, wtype);
+    }
+
+    ~ControlNet() {
+        free_control_ctx();
+    }
+
+    void alloc_control_ctx(std::vector<struct ggml_tensor*> outs) {
+        struct ggml_init_params params;
+        params.mem_size   = static_cast<size_t>(outs.size() * ggml_tensor_overhead()) + 1024 * 1024;
+        params.mem_buffer = NULL;
+        params.no_alloc   = true;
+        control_ctx       = ggml_init(params);
+
+        controls.resize(outs.size() - 1);
+
+        size_t control_buffer_size = 0;
+
+        guided_hint = ggml_dup_tensor(control_ctx, outs[0]);
+        control_buffer_size += ggml_nbytes(guided_hint);
+
+        for (int i = 0; i < outs.size() - 1; i++) {
+            controls[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
+            control_buffer_size += ggml_nbytes(controls[i]);
+        }
+
+        control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, backend);
+
+        LOG_DEBUG("control buffer size %.2fMB", control_buffer_size * 1.f / 1024.f / 1024.f);
+    }
+
+    void free_control_ctx() {
+        if (control_buffer != NULL) {
+            ggml_backend_buffer_free(control_buffer);
+            control_buffer = NULL;
+        }
+        if (control_ctx != NULL) {
+            ggml_free(control_ctx);
+            control_ctx = NULL;
+        }
+        guided_hint        = NULL;
+        guided_hint_cached = false;
+        controls.clear();
+    }
+
+    std::string get_desc() {
+        return "control_net";
+    }
+
+    size_t get_params_mem_size() {
+        return control_net.get_params_mem_size();
+    }
+
+    size_t get_params_num() {
+        return control_net.get_params_num();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        control_net.get_param_tensors(tensors, prefix);
+    }
+
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                    struct ggml_tensor* hint,
+                                    struct ggml_tensor* timesteps,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* y = NULL) {
+        struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, CONTROL_NET_GRAPH_SIZE, false);
+
+        x         = to_backend(x);
+        if (guided_hint_cached) {
+            hint = NULL;
+        } else {
+            hint      = to_backend(hint);
+        }
+        context   = to_backend(context);
+        y         = to_backend(y);
+        timesteps = to_backend(timesteps);
+
+        auto outs = control_net.forward(compute_ctx,
+                                        x,
+                                        hint,
+                                        guided_hint_cached ? guided_hint : NULL,
+                                        timesteps,
+                                        context,
+                                        y);
+
+        if (control_ctx == NULL) {
+            alloc_control_ctx(outs);
+        }
+
+        ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint));
+        for (int i = 0; i < outs.size() - 1; i++) {
+            ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], controls[i]));
+        }
+
+        return gf;
+    }
+
+    void compute(int n_threads,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* hint,
+                 struct ggml_tensor* timesteps,
+                 struct ggml_tensor* context,
+                 struct ggml_tensor* y,
+                 struct ggml_tensor** output     = NULL,
+                 struct ggml_context* output_ctx = NULL) {
+        // x: [N, in_channels, h, w]
+        // timesteps: [N, ]
+        // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
+        // y: [N, adm_in_channels] or [1, adm_in_channels]
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(x, hint, timesteps, context, y);
+        };
+
+        GGMLModule::compute(get_graph, n_threads, false, output, output_ctx);
+        guided_hint_cached = true;
+    }
+
+    bool load_from_file(const std::string& file_path) {
+        LOG_INFO("loading control net from '%s'", file_path.c_str());
+        alloc_params_buffer();
+        std::map<std::string, ggml_tensor*> tensors;
+        control_net.get_param_tensors(tensors);
+        std::set<std::string> ignore_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;
+        }
+
+        bool success = model_loader.load_tensors(tensors, backend, ignore_tensors);
+
+        if (!success) {
+            LOG_ERROR("load control net tensors from model loader failed");
+            return false;
+        }
+
+        LOG_INFO("control net model loaded");
+        return success;
+    }
+};
+
+#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

@@ -12,279 +12,160 @@
 
 */
 
-struct ResidualDenseBlock {
-    int num_features;
+class ResidualDenseBlock : public GGMLBlock {
+protected:
+    int num_feat;
     int num_grow_ch;
-    ggml_tensor* conv1_w;  // [num_grow_ch, num_features, 3, 3]
-    ggml_tensor* conv1_b;  // [num_grow_ch]
 
-    ggml_tensor* conv2_w;  // [num_grow_ch, num_features + num_grow_ch, 3, 3]
-    ggml_tensor* conv2_b;  // [num_grow_ch]
-
-    ggml_tensor* conv3_w;  // [num_grow_ch, num_features + 2 * num_grow_ch, 3, 3]
-    ggml_tensor* conv3_b;  // [num_grow_ch]
-
-    ggml_tensor* conv4_w;  // [num_grow_ch, num_features + 3 * num_grow_ch, 3, 3]
-    ggml_tensor* conv4_b;  // [num_grow_ch]
-
-    ggml_tensor* conv5_w;  // [num_features, num_features + 4 * num_grow_ch, 3, 3]
-    ggml_tensor* conv5_b;  // [num_features]
-
-    ResidualDenseBlock() {}
-
-    ResidualDenseBlock(int num_feat, int n_grow_ch) {
-        num_features = num_feat;
-        num_grow_ch  = n_grow_ch;
+public:
+    ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32)
+        : num_feat(num_feat), num_grow_ch(num_grow_ch) {
+        blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv3"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv4"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    size_t calculate_mem_size() {
-        size_t mem_size = num_features * num_grow_ch * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv1_w
-        mem_size += num_grow_ch * ggml_type_size(GGML_TYPE_F32);                               // conv1_b
-
-        mem_size += (num_features + num_grow_ch) * num_grow_ch * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv2_w
-        mem_size += num_grow_ch * ggml_type_size(GGML_TYPE_F32);                                         // conv2_b
-
-        mem_size += (num_features + 2 * num_grow_ch) * num_grow_ch * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv3_w
-        mem_size += num_grow_ch * ggml_type_size(GGML_TYPE_F32);                                             // conv3_w
-
-        mem_size += (num_features + 3 * num_grow_ch) * num_grow_ch * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv4_w
-        mem_size += num_grow_ch * ggml_type_size(GGML_TYPE_F32);                                             // conv4_w
-
-        mem_size += (num_features + 4 * num_grow_ch) * num_features * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv5_w
-        mem_size += num_features * ggml_type_size(GGML_TYPE_F32);                                             // conv5_w
-
-        return mem_size;
+    struct ggml_tensor* lrelu(struct ggml_context* ctx, struct ggml_tensor* x) {
+        return ggml_leaky_relu(ctx, x, 0.2f, true);
     }
 
-    int get_num_tensors() {
-        int num_tensors = 10;
-        return num_tensors;
-    }
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [n, num_feat, h, w]
+        // return: [n, num_feat, h, w]
 
-    void init_params(ggml_context* ctx) {
-        conv1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, num_features, num_grow_ch);
-        conv1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_grow_ch);
-        conv2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, num_features + num_grow_ch, num_grow_ch);
-        conv2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_grow_ch);
-        conv3_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, num_features + 2 * num_grow_ch, num_grow_ch);
-        conv3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_grow_ch);
-        conv4_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, num_features + 3 * num_grow_ch, num_grow_ch);
-        conv4_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_grow_ch);
-        conv5_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, num_features + 4 * num_grow_ch, num_features);
-        conv5_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_features);
-    }
+        auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+        auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
+        auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
+        auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
+        auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
 
-    void map_by_name(std::map<std::string, ggml_tensor*>& tensors, std::string prefix) {
-        tensors[prefix + "conv1.weight"] = conv1_w;
-        tensors[prefix + "conv1.bias"]   = conv1_b;
+        auto x1    = lrelu(ctx, conv1->forward(ctx, x));
+        auto x_cat = ggml_concat(ctx, x, x1);
+        auto x2    = lrelu(ctx, conv2->forward(ctx, x_cat));
+        x_cat      = ggml_concat(ctx, x_cat, x2);
+        auto x3    = lrelu(ctx, conv3->forward(ctx, x_cat));
+        x_cat      = ggml_concat(ctx, x_cat, x3);
+        auto x4    = lrelu(ctx, conv4->forward(ctx, x_cat));
+        x_cat      = ggml_concat(ctx, x_cat, x4);
+        auto x5    = conv5->forward(ctx, x_cat);
 
-        tensors[prefix + "conv2.weight"] = conv2_w;
-        tensors[prefix + "conv2.bias"]   = conv2_b;
-
-        tensors[prefix + "conv3.weight"] = conv3_w;
-        tensors[prefix + "conv3.bias"]   = conv3_b;
-
-        tensors[prefix + "conv4.weight"] = conv4_w;
-        tensors[prefix + "conv4.bias"]   = conv4_b;
-
-        tensors[prefix + "conv5.weight"] = conv5_w;
-        tensors[prefix + "conv5.bias"]   = conv5_b;
-    }
-
-    ggml_tensor* forward(ggml_context* ctx, ggml_tensor* 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);
-
-        // x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
-        ggml_tensor* x_cat = ggml_concat(ctx, x, x1);
-        ggml_tensor* x2    = ggml_nn_conv_2d(ctx, x_cat, conv2_w, conv2_b, 1, 1, 1, 1);
-        x2                 = ggml_leaky_relu(ctx, x2, 0.2f, true);
-
-        // x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
-        x_cat           = ggml_concat(ctx, x_cat, x2);
-        ggml_tensor* x3 = ggml_nn_conv_2d(ctx, x_cat, conv3_w, conv3_b, 1, 1, 1, 1);
-        x3              = ggml_leaky_relu(ctx, x3, 0.2f, true);
-
-        // x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
-        x_cat           = ggml_concat(ctx, x_cat, x3);
-        ggml_tensor* x4 = ggml_nn_conv_2d(ctx, x_cat, conv4_w, conv4_b, 1, 1, 1, 1);
-        x4              = ggml_leaky_relu(ctx, x4, 0.2f, true);
-
-        // self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
-        x_cat           = ggml_concat(ctx, x_cat, x4);
-        ggml_tensor* x5 = ggml_nn_conv_2d(ctx, x_cat, conv5_w, conv5_b, 1, 1, 1, 1);
-
-        // return x5 * 0.2 + x
-        x5 = ggml_add(ctx, ggml_scale(ctx, x5, out_scale), x);
+        x5 = ggml_add(ctx, ggml_scale(ctx, x5, 0.2f), x);
         return x5;
     }
 };
 
-struct EsrganBlock {
-    ResidualDenseBlock rd_blocks[3];
-    int num_residual_blocks = 3;
-
-    EsrganBlock() {}
-
-    EsrganBlock(int num_feat, int num_grow_ch) {
-        for (int i = 0; i < num_residual_blocks; i++) {
-            rd_blocks[i] = ResidualDenseBlock(num_feat, num_grow_ch);
-        }
+class RRDB : public GGMLBlock {
+public:
+    RRDB(int num_feat, int num_grow_ch = 32) {
+        blocks["rdb1"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
+        blocks["rdb2"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
+        blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
     }
 
-    int get_num_tensors() {
-        int num_tensors = 0;
-        for (int i = 0; i < num_residual_blocks; i++) {
-            num_tensors += rd_blocks[i].get_num_tensors();
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [n, num_feat, h, w]
+        // return: [n, num_feat, h, w]
+
+        auto rdb1 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb1"]);
+        auto rdb2 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb2"]);
+        auto rdb3 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb3"]);
+
+        auto out = rdb1->forward(ctx, x);
+        out      = rdb2->forward(ctx, out);
+        out      = rdb3->forward(ctx, out);
+
+        out = ggml_add(ctx, ggml_scale(ctx, out, 0.2f), x);
+        return out;
+    }
+};
+
+class RRDBNet : public GGMLBlock {
+protected:
+    int scale       = 4;  // default RealESRGAN_x4plus_anime_6B
+    int num_block   = 6;  // default RealESRGAN_x4plus_anime_6B
+    int num_in_ch   = 3;
+    int num_out_ch  = 3;
+    int num_feat    = 64;  // default RealESRGAN_x4plus_anime_6B
+    int num_grow_ch = 32;  // default RealESRGAN_x4plus_anime_6B
+
+public:
+    RRDBNet() {
+        blocks["conv_first"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_in_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        for (int i = 0; i < num_block; i++) {
+            std::string name = "body." + std::to_string(i);
+            blocks[name]     = std::shared_ptr<GGMLBlock>(new RRDB(num_feat, num_grow_ch));
         }
-        return num_tensors;
+        blocks["conv_body"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        // upsample
+        blocks["conv_up1"]  = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv_up2"]  = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv_hr"]   = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv_last"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_out_ch, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    size_t calculate_mem_size() {
-        size_t mem_size = 0;
-        for (int i = 0; i < num_residual_blocks; i++) {
-            mem_size += rd_blocks[i].calculate_mem_size();
-        }
-        return mem_size;
+    struct ggml_tensor* lrelu(struct ggml_context* ctx, struct ggml_tensor* x) {
+        return ggml_leaky_relu(ctx, x, 0.2f, true);
     }
 
-    void init_params(ggml_context* ctx) {
-        for (int i = 0; i < num_residual_blocks; i++) {
-            rd_blocks[i].init_params(ctx);
-        }
-    }
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [n, num_in_ch, h, w]
+        // return: [n, num_out_ch, h*4, w*4]
+        auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
+        auto conv_body  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_body"]);
+        auto conv_up1   = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
+        auto conv_up2   = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
+        auto conv_hr    = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
+        auto conv_last  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
 
-    void map_by_name(std::map<std::string, ggml_tensor*>& tensors, std::string prefix) {
-        for (int i = 0; i < num_residual_blocks; i++) {
-            rd_blocks[i].map_by_name(tensors, prefix + "rdb" + std::to_string(i + 1) + ".");
-        }
-    }
+        auto feat      = conv_first->forward(ctx, x);
+        auto body_feat = feat;
+        for (int i = 0; i < num_block; i++) {
+            std::string name = "body." + std::to_string(i);
+            auto block       = std::dynamic_pointer_cast<RRDB>(blocks[name]);
 
-    ggml_tensor* forward(ggml_context* ctx, ggml_tensor* out_scale, ggml_tensor* x) {
-        ggml_tensor* out = x;
-        for (int i = 0; i < num_residual_blocks; i++) {
-            // out = self.rdb...(x)
-            out = rd_blocks[i].forward(ctx, out_scale, out);
+            body_feat = block->forward(ctx, body_feat);
         }
-        // return out * 0.2 + x
-        out = ggml_add(ctx, ggml_scale(ctx, out, out_scale), x);
+        body_feat = conv_body->forward(ctx, body_feat);
+        feat      = ggml_add(ctx, feat, body_feat);
+        // upsample
+        feat     = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx, feat, 2)));
+        feat     = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx, feat, 2)));
+        auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
         return out;
     }
 };
 
 struct ESRGAN : public GGMLModule {
-    int scale        = 4;  // default RealESRGAN_x4plus_anime_6B
-    int num_blocks   = 6;  // default RealESRGAN_x4plus_anime_6B
-    int in_channels  = 3;
-    int out_channels = 3;
-    int num_features = 64;   // default RealESRGAN_x4plus_anime_6B
-    int num_grow_ch  = 32;   // default RealESRGAN_x4plus_anime_6B
-    int tile_size    = 128;  // avoid cuda OOM for 4gb VRAM
+    RRDBNet rrdb_net;
+    int scale     = 4;
+    int tile_size = 128;  // avoid cuda OOM for 4gb VRAM
 
-    ggml_tensor* conv_first_w;  // [num_features, in_channels, 3, 3]
-    ggml_tensor* conv_first_b;  // [num_features]
-
-    EsrganBlock body_blocks[6];
-    ggml_tensor* conv_body_w;  // [num_features, num_features, 3, 3]
-    ggml_tensor* conv_body_b;  // [num_features]
-
-    // upsample
-    ggml_tensor* conv_up1_w;  // [num_features, num_features, 3, 3]
-    ggml_tensor* conv_up1_b;  // [num_features]
-    ggml_tensor* conv_up2_w;  // [num_features, num_features, 3, 3]
-    ggml_tensor* conv_up2_b;  // [num_features]
-
-    ggml_tensor* conv_hr_w;    // [num_features, num_features, 3, 3]
-    ggml_tensor* conv_hr_b;    // [num_features]
-    ggml_tensor* conv_last_w;  // [out_channels, num_features, 3, 3]
-    ggml_tensor* conv_last_b;  // [out_channels]
-
-    bool decode_only = false;
-
-    ESRGAN() {
-        name = "esrgan";
-        for (int i = 0; i < num_blocks; i++) {
-            body_blocks[i] = EsrganBlock(num_features, num_grow_ch);
-        }
+    ESRGAN(ggml_backend_t backend,
+           ggml_type wtype)
+        : GGMLModule(backend, wtype) {
+        rrdb_net.init(params_ctx, wtype);
     }
 
-    size_t calculate_mem_size() {
-        size_t mem_size = num_features * in_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_first_w
-        mem_size += num_features * ggml_type_size(GGML_TYPE_F32);                              // conv_first_b
-
-        for (int i = 0; i < num_blocks; i++) {
-            mem_size += body_blocks[i].calculate_mem_size();
-        }
-
-        mem_size += num_features * num_features * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_body_w
-        mem_size += num_features * ggml_type_size(GGML_TYPE_F32);                         // conv_body_w
-
-        // upsample
-        mem_size += num_features * num_features * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_up1_w
-        mem_size += num_features * ggml_type_size(GGML_TYPE_F32);                         // conv_up1_b
-
-        mem_size += num_features * num_features * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_up2_w
-        mem_size += num_features * ggml_type_size(GGML_TYPE_F32);                         // conv_up2_b
-
-        mem_size += num_features * num_features * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_hr_w
-        mem_size += num_features * ggml_type_size(GGML_TYPE_F32);                         // conv_hr_b
-
-        mem_size += out_channels * num_features * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_last_w
-        mem_size += out_channels * ggml_type_size(GGML_TYPE_F32);                         // conv_last_b
-        return mem_size;
+    std::string get_desc() {
+        return "esrgan";
     }
 
-    size_t get_num_tensors() {
-        size_t num_tensors = 12;
-        for (int i = 0; i < num_blocks; i++) {
-            num_tensors += body_blocks[i].get_num_tensors();
-        }
-        return num_tensors;
+    size_t get_params_mem_size() {
+        return rrdb_net.get_params_mem_size();
     }
 
-    void init_params() {
-        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
-        conv_first_w       = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, in_channels, num_features);
-        conv_first_b       = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, num_features);
-        conv_body_w        = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, num_features, num_features);
-        conv_body_b        = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, num_features);
-        conv_up1_w         = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, num_features, num_features);
-        conv_up1_b         = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, num_features);
-        conv_up2_w         = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, num_features, num_features);
-        conv_up2_b         = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, num_features);
-        conv_hr_w          = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, num_features, num_features);
-        conv_hr_b          = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, num_features);
-        conv_last_w        = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, num_features, out_channels);
-        conv_last_b        = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, out_channels);
-
-        for (int i = 0; i < num_blocks; i++) {
-            body_blocks[i].init_params(params_ctx);
-        }
-
-        // 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);
+    size_t get_params_num() {
+        return rrdb_net.get_params_num();
     }
 
-    bool load_from_file(const std::string& file_path, ggml_backend_t backend) {
+    bool load_from_file(const std::string& file_path) {
         LOG_INFO("loading esrgan from '%s'", file_path.c_str());
 
-        if (!alloc_params_buffer(backend)) {
-            return false;
-        }
-
+        alloc_params_buffer();
         std::map<std::string, ggml_tensor*> esrgan_tensors;
-
-        // prepare memory for the weights
-        {
-            init_params();
-            map_by_name(esrgan_tensors);
-        }
+        rrdb_net.get_param_tensors(esrgan_tensors);
 
         ModelLoader model_loader;
         if (!model_loader.init_from_file(file_path)) {
@@ -303,120 +184,22 @@ struct ESRGAN : public GGMLModule {
         return success;
     }
 
-    void map_by_name(std::map<std::string, ggml_tensor*>& tensors) {
-        tensors["conv_first.weight"] = conv_first_w;
-        tensors["conv_first.bias"]   = conv_first_b;
-
-        for (int i = 0; i < num_blocks; i++) {
-            body_blocks[i].map_by_name(tensors, "body." + std::to_string(i) + ".");
-        }
-
-        tensors["conv_body.weight"] = conv_body_w;
-        tensors["conv_body.bias"]   = conv_body_b;
-
-        tensors["conv_up1.weight"] = conv_up1_w;
-        tensors["conv_up1.bias"]   = conv_up1_b;
-        tensors["conv_up2.weight"] = conv_up2_w;
-        tensors["conv_up2.bias"]   = conv_up2_b;
-        tensors["conv_hr.weight"]  = conv_hr_w;
-        tensors["conv_hr.bias"]    = conv_hr_b;
-
-        tensors["conv_last.weight"] = conv_last_w;
-        tensors["conv_last.bias"]   = conv_last_b;
-    }
-
-    ggml_tensor* forward(ggml_context* ctx0, ggml_tensor* 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);
-
-        auto body_h = h;
-        // self.body(feat)
-        for (int i = 0; i < num_blocks; i++) {
-            body_h = body_blocks[i].forward(ctx0, out_scale, body_h);
-        }
-
-        // body_feat = self.conv_body(self.body(feat))
-        body_h = ggml_nn_conv_2d(ctx0, body_h, conv_body_w, conv_body_b, 1, 1, 1, 1);
-
-        // feat = feat + body_feat
-        h = ggml_add(ctx0, h, body_h);
-
-        // upsample
-        // feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
-        h = ggml_upscale(ctx0, h, 2);
-        h = ggml_nn_conv_2d(ctx0, h, conv_up1_w, conv_up1_b, 1, 1, 1, 1);
-        h = ggml_leaky_relu(ctx0, h, 0.2f, true);
-
-        // feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
-        h = ggml_upscale(ctx0, h, 2);
-        h = ggml_nn_conv_2d(ctx0, h, conv_up2_w, conv_up2_b, 1, 1, 1, 1);
-        h = ggml_leaky_relu(ctx0, h, 0.2f, true);
-
-        // out = self.conv_last(self.lrelu(self.conv_hr(feat)))
-        h = ggml_nn_conv_2d(ctx0, h, conv_hr_w, conv_hr_b, 1, 1, 1, 1);
-        h = ggml_leaky_relu(ctx0, h, 0.2f, true);
-
-        h = ggml_nn_conv_2d(ctx0, h, conv_last_w, conv_last_b, 1, 1, 1, 1);
-        return h;
-    }
-
     struct ggml_cgraph* build_graph(struct ggml_tensor* x) {
-        // 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);
-
-        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));
-        }
-
-        // it's performing a compute, check if backend isn't cpu
-        if (!ggml_backend_is_cpu(backend)) {
-            // pass input tensors to gpu memory
-            x_ = ggml_dup_tensor(ctx0, x);
-            ggml_allocr_alloc(compute_allocr, x_);
-
-            // pass data to device backend
-            if (!ggml_allocr_is_measure(compute_allocr)) {
-                ggml_backend_tensor_set(x_, x->data, 0, ggml_nbytes(x));
-            }
-        } else {
-            x_ = x;
-        }
-
-        struct ggml_tensor* out = forward(ctx0, os, x);
-
+        struct ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        x                       = to_backend(x);
+        struct ggml_tensor* out = rrdb_net.forward(compute_ctx, x);
         ggml_build_forward_expand(gf, out);
-        ggml_free(ctx0);
-
         return gf;
     }
 
-    void alloc_compute_buffer(struct ggml_tensor* x) {
+    void compute(const int n_threads,
+                 struct ggml_tensor* x,
+                 ggml_tensor** output,
+                 ggml_context* output_ctx = NULL) {
         auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(x);
         };
-        GGMLModule::alloc_compute_buffer(get_graph);
-    }
-
-    void compute(struct ggml_tensor* work_result, const int n_threads, struct ggml_tensor* x) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(x);
-        };
-        GGMLModule::compute(get_graph, n_threads, work_result);
+        GGMLModule::compute(get_graph, n_threads, false, output, output_ctx);
     }
 };
 

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,15 @@ const char* schedule_str[] = {
 const char* modes_str[] = {
     "txt2img",
     "img2img",
+    "img2vid",
+    "convert",
 };
 
 enum SDMode {
     TXT2IMG,
     IMG2IMG,
+    IMG2VID,
+    CONVERT,
     MODE_COUNT
 };
 
@@ -58,41 +63,42 @@ 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;
+    float min_cfg   = 1.0f;
     float cfg_scale = 7.0f;
     int clip_skip   = -1;  // <= 0 represents unspecified
     int width       = 512;
     int height      = 512;
     int batch_count = 1;
 
+    int video_frames         = 6;
+    int motion_bucket_id     = 127;
+    int fps                  = 6;
+    float augmentation_level = 0.f;
+
     sample_method_t sample_method = EULER_A;
     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;
+    int upscale_repeats           = 1;
 };
 
-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,10 +108,16 @@ 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("    min_cfg:           %.2f\n", params.min_cfg);
     printf("    cfg_scale:         %.2f\n", params.cfg_scale);
     printf("    clip_skip:         %d\n", params.clip_skip);
     printf("    width:             %d\n", params.width);
@@ -118,6 +130,7 @@ void print_params(SDParams params) {
     printf("    seed:              %ld\n", params.seed);
     printf("    batch_count:       %d\n", params.batch_count);
     printf("    vae_tiling:        %s\n", params.vae_tiling ? "true" : "false");
+    printf("    upscale_repeats:   %d\n", params.upscale_repeats);
 }
 
 void print_usage(int argc, const char* argv[]) {
@@ -125,22 +138,27 @@ 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("  --upscale-repeats                  Run the ESRGAN upscaler this many times (default 1)\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 +172,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");
 }
 
@@ -182,7 +202,8 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 }
             }
             if (mode_found == -1) {
-                fprintf(stderr, "error: invalid mode %s, must be one of [txt2img, img2img]\n",
+                fprintf(stderr,
+                        "error: invalid mode %s, must be one of [txt2img, img2img, img2vid, convert]\n",
                         mode_selected);
                 exit(1);
             }
@@ -205,12 +226,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 +281,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;
@@ -260,6 +299,16 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.prompt = argv[i];
+        } else if (arg == "--upscale-repeats") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.upscale_repeats = std::stoi(argv[i]);
+            if (params.upscale_repeats < 1) {
+                fprintf(stderr, "error: upscale multiplier must be at least 1\n");
+                exit(1);
+            }
         } else if (arg == "-n" || arg == "--negative-prompt") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -278,6 +327,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 +359,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 +443,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.mode != IMG2VID && params.prompt.length() == 0) {
         fprintf(stderr, "error: the following arguments are required: prompt\n");
         print_usage(argc, argv);
         exit(1);
@@ -396,7 +455,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         exit(1);
     }
 
-    if (params.mode == IMG2IMG && params.input_path.length() == 0) {
+    if ((params.mode == IMG2IMG || params.mode == IMG2VID) && params.input_path.length() == 0) {
         fprintf(stderr, "error: when using the img2img mode, the following arguments are required: init-img\n");
         print_usage(argc, argv);
         exit(1);
@@ -432,6 +491,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) {
@@ -479,9 +544,32 @@ 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;
+        }
+    }
+
+    if (params.mode == IMG2VID) {
+        fprintf(stderr, "SVD support is broken, do not use it!!!\n");
+        return 1;
+    }
+
     bool vae_decode_only        = true;
     uint8_t* input_image_buffer = NULL;
-    if (params.mode == IMG2IMG) {
+    if (params.mode == IMG2IMG || params.mode == IMG2VID) {
         vae_decode_only = false;
 
         int c              = 0;
@@ -510,14 +598,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 +617,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,26 +644,66 @@ 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,
                                   3,
                                   input_image_buffer};
 
-        results = img2img(sd_ctx,
-                          input_image,
-                          params.prompt.c_str(),
-                          params.negative_prompt.c_str(),
-                          params.clip_skip,
-                          params.cfg_scale,
-                          params.width,
-                          params.height,
-                          params.sample_method,
-                          params.sample_steps,
-                          params.strength,
-                          params.seed,
-                          params.batch_count);
+        if (params.mode == IMG2VID) {
+            results = img2vid(sd_ctx,
+                              input_image,
+                              params.width,
+                              params.height,
+                              params.video_frames,
+                              params.motion_bucket_id,
+                              params.fps,
+                              params.augmentation_level,
+                              params.min_cfg,
+                              params.cfg_scale,
+                              params.sample_method,
+                              params.sample_steps,
+                              params.strength,
+                              params.seed);
+            if (results == NULL) {
+                printf("generate failed\n");
+                free_sd_ctx(sd_ctx);
+                return 1;
+            }
+            size_t last            = params.output_path.find_last_of(".");
+            std::string dummy_name = last != std::string::npos ? params.output_path.substr(0, last) : params.output_path;
+            for (int i = 0; i < params.video_frames; i++) {
+                if (results[i].data == NULL) {
+                    continue;
+                }
+                std::string final_image_path = i > 0 ? dummy_name + "_" + std::to_string(i + 1) + ".png" : dummy_name + ".png";
+                stbi_write_png(final_image_path.c_str(), results[i].width, results[i].height, results[i].channel,
+                               results[i].data, 0, get_image_params(params, params.seed + i).c_str());
+                printf("save result image to '%s'\n", final_image_path.c_str());
+                free(results[i].data);
+                results[i].data = NULL;
+            }
+            free(results);
+            free_sd_ctx(sd_ctx);
+            return 0;
+        } else {
+            results = img2img(sd_ctx,
+                              input_image,
+                              params.prompt.c_str(),
+                              params.negative_prompt.c_str(),
+                              params.clip_skip,
+                              params.cfg_scale,
+                              params.width,
+                              params.height,
+                              params.sample_method,
+                              params.sample_steps,
+                              params.strength,
+                              params.seed,
+                              params.batch_count);
+        }
     }
 
     if (results == NULL) {
@@ -565,7 +713,7 @@ int main(int argc, const char* argv[]) {
     }
 
     int upscale_factor = 4;  // unused for RealESRGAN_x4plus_anime_6B.pth
-    if (params.esrgan_path.size() > 0) {
+    if (params.esrgan_path.size() > 0 && params.upscale_repeats > 0) {
         upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(params.esrgan_path.c_str(),
                                                         params.n_threads,
                                                         params.wtype);
@@ -577,13 +725,17 @@ int main(int argc, const char* argv[]) {
                 if (results[i].data == NULL) {
                     continue;
                 }
-                sd_image_t upscaled_image = upscale(upscaler_ctx, results[i], upscale_factor);
-                if (upscaled_image.data == NULL) {
-                    printf("upscale failed\n");
-                    continue;
+                sd_image_t current_image = results[i];
+                for (int u = 0; u < params.upscale_repeats; ++u) {
+                    sd_image_t upscaled_image = upscale(upscaler_ctx, current_image, upscale_factor);
+                    if (upscaled_image.data == NULL) {
+                        printf("upscale failed\n");
+                        break;
+                    }
+                    free(current_image.data);
+                    current_image = upscaled_image;
                 }
-                free(results[i].data);
-                results[i] = upscaled_image;
+                results[i] = current_image;  // Set the final upscaled image as the result
             }
         }
     }

ggml_extend.hpp

@@ -11,6 +11,7 @@
 #include <iostream>
 #include <iterator>
 #include <map>
+#include <memory>
 #include <random>
 #include <regex>
 #include <set>
@@ -65,9 +66,11 @@ __STATIC_INLINE__ void ggml_tensor_set_f32(struct ggml_tensor* tensor, float val
 }
 
 __STATIC_INLINE__ float ggml_tensor_get_f32(const ggml_tensor* tensor, int l, int k = 0, int j = 0, int i = 0) {
-    // float value;
-    // ggml_backend_tensor_get(tensor, &value, i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0], sizeof(float));
-    // return value;
+    if (tensor->buffer != NULL) {
+        float value;
+        ggml_backend_tensor_get(tensor, &value, i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0], sizeof(float));
+        return value;
+    }
     GGML_ASSERT(tensor->nb[0] == sizeof(float));
     return *(float*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]);
 }
@@ -183,7 +186,7 @@ __STATIC_INLINE__ void copy_ggml_tensor(struct ggml_tensor* dst, struct ggml_ten
         LOG_ERROR("ggml_init() failed");
         return;
     }
-    ggml_tensor* final = ggml_cpy_inplace(ctx, src, dst);
+    ggml_tensor* final = ggml_cpy(ctx, src, dst);
 
     struct ggml_cgraph* graph = ggml_new_graph(ctx);
     ggml_build_forward_expand(graph, final);
@@ -191,6 +194,10 @@ __STATIC_INLINE__ void copy_ggml_tensor(struct ggml_tensor* dst, struct ggml_ten
     ggml_free(ctx);
 }
 
+__STATIC_INLINE__ float sigmoid(float x) {
+    return 1 / (1.0f + expf(-x));
+}
+
 // SPECIAL OPERATIONS WITH TENSORS
 
 __STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input) {
@@ -211,7 +218,8 @@ __STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input) {
 }
 
 __STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
-                                          struct ggml_tensor* output) {
+                                          struct ggml_tensor* output,
+                                          bool scale = true) {
     int64_t width    = output->ne[0];
     int64_t height   = output->ne[1];
     int64_t channels = output->ne[2];
@@ -220,7 +228,30 @@ __STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
         for (int ix = 0; ix < width; ix++) {
             for (int k = 0; k < channels; k++) {
                 float value = *(image_data + iy * width * channels + ix * channels + k);
-                ggml_tensor_set_f32(output, value / 255.0f, ix, iy, k);
+                if (scale) {
+                    value /= 255.f;
+                }
+                ggml_tensor_set_f32(output, value, ix, iy, k);
+            }
+        }
+    }
+}
+
+__STATIC_INLINE__ void sd_image_f32_to_tensor(const float* image_data,
+                                              struct ggml_tensor* output,
+                                              bool scale = true) {
+    int64_t width    = output->ne[0];
+    int64_t height   = output->ne[1];
+    int64_t channels = output->ne[2];
+    GGML_ASSERT(channels == 3 && output->type == GGML_TYPE_F32);
+    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);
+                if (scale) {
+                    value /= 255.f;
+                }
+                ggml_tensor_set_f32(output, value, ix, iy, k);
             }
         }
     }
@@ -407,7 +438,9 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_linear(struct ggml_context* ctx,
                                                      struct ggml_tensor* w,
                                                      struct ggml_tensor* b) {
     x = ggml_mul_mat(ctx, w, x);
-    x = ggml_add(ctx, x, b);
+    if (b != NULL) {
+        x = ggml_add(ctx, x, b);
+    }
     return x;
 }
 
@@ -428,19 +461,103 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_conv_2d(struct ggml_context* ctx,
     x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
     if (b != NULL) {
         b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        // b = ggml_repeat(ctx, b, x);
         x = ggml_add(ctx, x, b);
     }
     return x;
 }
 
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N, OC, OH, OW]
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_conv_3d_nx1x1_bak(struct ggml_context* ctx,
+                                                                struct ggml_tensor* x,
+                                                                struct ggml_tensor* w,
+                                                                struct ggml_tensor* b,
+                                                                int s2 = 1,
+                                                                int p2 = 1,
+                                                                int d2 = 1) {
+    GGML_ASSERT(w->ne[0] == 1);
+    // timesteps = x.shape[0]
+    // x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
+    // x = conv3d(x)
+    // return rearrange(x, "b c t h w -> (b t) c h w")
+    int64_t T = x->ne[3];
+    int64_t B = x->ne[3] / T;
+    int64_t C = x->ne[2];
+    int64_t H = x->ne[1];
+    int64_t W = x->ne[0];
+
+    x = ggml_reshape_4d(ctx, x, W * H, C, T, B);           // (b t) c h w -> b t c (h w)
+    x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // b t c (h w) -> b c t (h w)
+    x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2);      // [B, OC, T, OH * OW]
+    if (b != NULL) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add(ctx, x, b);
+    }
+    x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // b c t (h w) -> b t c (h w)
+    x = ggml_reshape_4d(ctx, x, W, H, C, T * B);           // b t c (h w) -> (b t) c h w
+    return x;                                              // [B*T, OC, OH, OW]
+}
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, ID, IH*IW]
+// b: [OC,]
+// result: [N, OC, OD, OH*OW]
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_conv_3d_nx1x1(struct ggml_context* ctx,
+                                                            struct ggml_tensor* x,
+                                                            struct ggml_tensor* w,
+                                                            struct ggml_tensor* b,
+                                                            int s2 = 1,
+                                                            int p2 = 1,
+                                                            int d2 = 1) {
+    x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2);  // [N, OC, T, OH * OW]
+    if (b != NULL) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add(ctx, x, b);
+    }
+    return x;  // [N, OC, T, OH * OW]
+}
+
+// q: [N * n_head, n_token, d_head]
+// k: [N * n_head, n_k, d_head]
+// v: [N * n_head, d_head, n_k]
+// return: [N * n_head, n_token, d_head]
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention(struct ggml_context* ctx,
+                                                        struct ggml_tensor* q,
+                                                        struct ggml_tensor* k,
+                                                        struct ggml_tensor* v,
+                                                        bool mask = false) {
+#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, n_token, d_head]
+#else
+    float d_head = (float)q->ne[0];
+
+    struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, n_token, n_k]
+    kq                     = ggml_scale_inplace(ctx, kq, 1.0f / sqrt(d_head));
+    if (mask) {
+        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, n_token, d_head]
+#endif
+    return kqv;
+}
+
 __STATIC_INLINE__ struct ggml_tensor* ggml_nn_layer_norm(struct ggml_context* ctx,
                                                          struct ggml_tensor* x,
                                                          struct ggml_tensor* w,
                                                          struct ggml_tensor* b,
                                                          float eps = EPS) {
     x = ggml_norm(ctx, x, eps);
-    x = ggml_mul(ctx, x, w);
-    x = ggml_add(ctx, x, b);
+    if (w != NULL) {
+        x = ggml_mul(ctx, x, w);
+        if (b != NULL) {
+            x = ggml_add(ctx, x, b);
+        }
+    }
     return x;
 }
 
@@ -449,21 +566,28 @@ __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 != NULL && b != NULL) {
         w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
         b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
     }
 
     x = ggml_group_norm(ctx, x, num_groups);
-    x = ggml_mul(ctx, x, w);
-    x = ggml_add(ctx, x, b);
+    if (w != NULL && b != NULL) {
+        x = ggml_mul(ctx, x, w);
+        // b = ggml_repeat(ctx, b, x);
+        x = ggml_add(ctx, x, b);
+    }
     return x;
 }
 
 __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
@@ -482,133 +606,269 @@ __STATIC_INLINE__ float ggml_backend_tensor_get_f32(ggml_tensor* tensor) {
     return value;
 }
 
+__STATIC_INLINE__ struct ggml_tensor* vector_to_ggml_tensor(struct ggml_context* ctx,
+                                                            const std::vector<float>& vec) {
+    struct ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+__STATIC_INLINE__ struct ggml_tensor* vector_to_ggml_tensor_i32(struct ggml_context* ctx,
+                                                                const std::vector<int>& vec) {
+    struct ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+__STATIC_INLINE__ std::vector<float> arange(float start, float end, float step = 1.f) {
+    std::vector<float> result;
+
+    for (float value = start; value < end; value += step) {
+        result.push_back(value);
+    }
+
+    return result;
+}
+
 // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
-__STATIC_INLINE__ void set_timestep_embedding(struct ggml_tensor* timesteps, struct ggml_tensor* embedding, int dim, int max_period = 10000) {
+__STATIC_INLINE__ std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                                        int dim,
+                                                        int max_period = 10000) {
     // timesteps: [N,]
-    // embedding: [dim, N]
+    // embedding: [N, dim]
+    size_t N        = timesteps.size();
+    int acutual_dim = dim;
+    if (dim % 2 != 0) {
+        acutual_dim = dim + 1;
+    }
+    std::vector<float> embedding(N * acutual_dim, 0.f);
     int half = dim / 2;
     std::vector<float> freqs(half);
     for (int i = 0; i < half; ++i) {
         freqs[i] = (float)std::exp(-std::log(max_period) * i / half);
     }
-    for (int i = 0; i < timesteps->ne[0]; ++i) {
+    for (int i = 0; i < N; ++i) {
         for (int j = 0; j < half; ++j) {
-            float arg = ggml_get_f32_1d(timesteps, i) * freqs[j];
-            ggml_tensor_set_f32(embedding, std::cos(arg), j, i);
-            ggml_tensor_set_f32(embedding, std::sin(arg), j + half, i);
+            float arg                             = timesteps[i] * freqs[j];
+            embedding[i * acutual_dim + j]        = std::cos(arg);
+            embedding[i * acutual_dim + j + half] = std::sin(arg);
         }
-        if (dim % 2 != 0) {
-            *(float*)((char*)embedding->data + i * embedding->nb[1] + dim * embedding->nb[0]) = 0;
-        }
-    }
-}
-
-__STATIC_INLINE__ struct ggml_tensor* new_timestep_embedding(struct ggml_context* ctx,
-                                                             struct ggml_allocr* allocr,
-                                                             struct ggml_tensor* timesteps,
-                                                             int dim,
-                                                             int max_period = 10000) {
-    // timesteps: [N,]
-    // embedding: [dim, N]
-    int acutual_dim = dim;
-    if (dim % 2 != 0) {
-        acutual_dim = dim + 1;
-    }
-    struct ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, acutual_dim, timesteps->ne[0]);
-    if (allocr != NULL) {
-        ggml_allocr_alloc(allocr, embedding);
-    }
-    if (allocr != NULL && !ggml_allocr_is_measure(allocr)) {
-        set_timestep_embedding(timesteps, embedding, dim, max_period);
     }
     return embedding;
 }
 
+__STATIC_INLINE__ void set_timestep_embedding(std::vector<float> timesteps,
+                                              struct ggml_tensor* embedding,
+                                              int dim,
+                                              int max_period = 10000) {
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+}
+
+__STATIC_INLINE__ struct ggml_tensor* new_timestep_embedding(struct ggml_context* ctx,
+                                                             std::vector<float> timesteps,
+                                                             int dim,
+                                                             int max_period = 10000) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    int acutual_dim                  = dim;
+    if (dim % 2 != 0) {
+        acutual_dim = dim + 1;
+    }
+    struct ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, acutual_dim, timesteps.size());
+    if (embedding->data != NULL) {
+        memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+    } else {
+        ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding));
+    }
+    return embedding;
+}
+
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_timestep_embedding(
+    struct ggml_context* ctx,
+    struct ggml_tensor* timesteps,
+    int dim,
+    int max_period = 10000) {
+    return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
+}
+
+// struct GGMLComputeGraph {
+//     virtual void init(struct ggml_context* ctx, ggml_type wtype) = 0;
+//     virtual std::string get_desc() = 0;
+//     virtual size_t get_params_mem_size()   = 0;
+//     virtual size_t get_params_num() = 0;
+//     virtual struct ggml_cgraph* get_ggml_cgraph() = 0;
+// };
+
+#define MAX_PARAMS_TENSOR_NUM 10240
+#define MAX_GRAPH_SIZE 10240
+
 struct GGMLModule {
+protected:
     typedef std::function<struct ggml_cgraph*()> get_graph_cb_t;
 
-    std::string name                     = "ggml module";
-    struct ggml_context* params_ctx      = NULL;
-    size_t params_buffer_size            = 0;
-    size_t compute_buffer_size           = 0;
-    ggml_backend_buffer_t params_buffer  = NULL;
-    ggml_backend_buffer_t compute_buffer = NULL;  // for compute
-    struct ggml_allocr* compute_allocr   = NULL;
+    struct ggml_context* params_ctx     = NULL;
+    ggml_backend_buffer_t params_buffer = NULL;
+
+    struct ggml_context* compute_ctx    = NULL;
+    struct ggml_gallocr* compute_allocr = NULL;
+
+    std::map<struct ggml_tensor*, const void*> backend_tensor_data_map;
 
     ggml_type wtype        = GGML_TYPE_F32;
     ggml_backend_t backend = NULL;
 
-    virtual size_t calculate_mem_size() = 0;
-    virtual size_t get_num_tensors()    = 0;
-
-    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 += calculate_mem_size();
-        size_t num_tensors = get_num_tensors();
-
-        LOG_DEBUG("%s params backend buffer size = % 6.2f MB (%i tensors)",
-                  name.c_str(), params_buffer_size / (1024.0 * 1024.0), num_tensors);
-
+    void alloc_params_ctx() {
         struct ggml_init_params params;
-        params.mem_size   = static_cast<size_t>(num_tensors * ggml_tensor_overhead()) + 1 * 1024 * 1024;
+        params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
         params.mem_buffer = NULL;
         params.no_alloc   = true;
-        // LOG_DEBUG("mem_size %u ", params.mem_size);
 
         params_ctx = ggml_init(params);
-        if (!params_ctx) {
-            LOG_ERROR("ggml_init() failed");
-            return false;
-        }
-
-        params_buffer = ggml_backend_alloc_buffer(backend, params_buffer_size);
-        return true;
+        GGML_ASSERT(params_ctx != NULL);
     }
 
-    void free_params_buffer() {
+    void free_params_ctx() {
         if (params_ctx != NULL) {
             ggml_free(params_ctx);
             params_ctx = NULL;
         }
+    }
 
+    void alloc_compute_ctx() {
+        struct ggml_init_params params;
+        params.mem_size   = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
+        params.mem_buffer = NULL;
+        params.no_alloc   = true;
+
+        compute_ctx = ggml_init(params);
+        GGML_ASSERT(compute_ctx != NULL);
+    }
+
+    void free_compute_ctx() {
+        if (compute_ctx != NULL) {
+            ggml_free(compute_ctx);
+            compute_ctx = NULL;
+        }
+    }
+
+    bool alloc_compute_buffer(get_graph_cb_t get_graph) {
+        if (compute_allocr != NULL) {
+            return true;
+        }
+        reset_compute_ctx();
+        struct ggml_cgraph* gf = get_graph();
+        backend_tensor_data_map.clear();
+        compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
+
+        if (!ggml_gallocr_reserve(compute_allocr, gf)) {
+            // failed to allocate the compute buffer
+            LOG_ERROR("%s: failed to allocate the compute buffer\n", get_desc().c_str());
+            free_compute_buffer();
+            return false;
+        }
+
+        // compute the required memory
+        size_t compute_buffer_size = ggml_gallocr_get_buffer_size(compute_allocr, 0);
+        LOG_DEBUG("%s compute buffer size: %.2f MB", get_desc().c_str(), compute_buffer_size / 1024.0 / 1024.0);
+        return true;
+    }
+
+    void cpy_data_to_backend_tensor() {
+        for (auto& kv : backend_tensor_data_map) {
+            auto tensor = kv.first;
+            auto data   = kv.second;
+
+            ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+        }
+
+        backend_tensor_data_map.clear();
+    }
+
+public:
+    virtual size_t get_params_mem_size() = 0;
+    virtual size_t get_params_num()      = 0;
+    virtual std::string get_desc()       = 0;
+
+    GGMLModule(ggml_backend_t backend, ggml_type wtype = GGML_TYPE_F32)
+        : backend(backend), wtype(wtype) {
+        alloc_params_ctx();
+    }
+
+    virtual ~GGMLModule() {
+        free_params_buffer();
+        free_compute_buffer();
+        free_params_ctx();
+        free_compute_ctx();
+    }
+
+    void reset_compute_ctx() {
+        free_compute_ctx();
+        alloc_compute_ctx();
+    }
+
+    bool alloc_params_buffer() {
+        size_t num_tensors = get_params_num();
+        params_buffer      = ggml_backend_alloc_ctx_tensors(params_ctx, backend);
+        if (params_buffer == NULL) {
+            LOG_ERROR("%s alloc params backend buffer failed", get_desc().c_str());
+            return false;
+        }
+        size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer);
+        LOG_DEBUG("%s params backend buffer size = % 6.2f MB (%i tensors)",
+                  get_desc().c_str(), params_buffer_size / (1024.0 * 1024.0), num_tensors);
+        return true;
+    }
+
+    void free_params_buffer() {
         if (params_buffer != NULL) {
             ggml_backend_buffer_free(params_buffer);
             params_buffer = NULL;
         }
     }
 
-    ~GGMLModule() {
-        free_params_buffer();
-    }
-
-    void alloc_compute_buffer(get_graph_cb_t get_graph) {
-        if (compute_buffer_size == 0) {
-            // alignment required by the backend
-            compute_allocr = ggml_allocr_new_measure_from_backend(backend);
-
-            struct ggml_cgraph* gf = get_graph();
-
-            // compute the required memory
-            compute_buffer_size = ggml_allocr_alloc_graph(compute_allocr, gf) + 1024 * 1024;
-
-            // recreate the allocator with the required memory
-            ggml_allocr_free(compute_allocr);
-
-            LOG_DEBUG("%s compute buffer size: %.2f MB", name.c_str(), compute_buffer_size / 1024.0 / 1024.0);
+    void free_compute_buffer() {
+        if (compute_allocr != NULL) {
+            ggml_gallocr_free(compute_allocr);
+            compute_allocr = NULL;
         }
-
-        compute_buffer = ggml_backend_alloc_buffer(backend, compute_buffer_size);
-        compute_allocr = ggml_allocr_new_from_buffer(compute_buffer);
     }
 
-    void compute(get_graph_cb_t get_graph, int n_threads, struct ggml_tensor* output = NULL) {
-        ggml_allocr_reset(compute_allocr);
+    // do copy after alloc graph
+    void set_backend_tensor_data(struct ggml_tensor* tensor, const void* data) {
+        backend_tensor_data_map[tensor] = data;
+    }
 
+    struct ggml_tensor* to_backend(struct ggml_tensor* tensor) {
+        GGML_ASSERT(compute_ctx != NULL);
+        if (tensor == NULL) {
+            return NULL;
+        }
+        // it's performing a compute, check if backend isn't cpu
+        if (!ggml_backend_is_cpu(backend) && tensor->backend == GGML_BACKEND_CPU) {
+            // pass input tensors to gpu memory
+            auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
+
+            set_backend_tensor_data(backend_tensor, tensor->data);
+            return backend_tensor;
+        } else {
+            return tensor;
+        }
+    }
+
+    void compute(get_graph_cb_t get_graph,
+                 int n_threads,
+                 bool free_compute_buffer_immediately = true,
+                 struct ggml_tensor** output          = NULL,
+                 struct ggml_context* output_ctx      = NULL) {
+        alloc_compute_buffer(get_graph);
+        reset_compute_ctx();
         struct ggml_cgraph* gf = get_graph();
 
-        ggml_allocr_alloc_graph(compute_allocr, gf);
+        GGML_ASSERT(ggml_gallocr_alloc_graph(compute_allocr, gf));
+
+        cpy_data_to_backend_tensor();
 
         if (ggml_backend_is_cpu(backend)) {
             ggml_backend_cpu_set_n_threads(backend, n_threads);
@@ -627,15 +887,368 @@ struct GGMLModule {
 #endif
 
         if (output != NULL) {
-            ggml_backend_tensor_get_and_sync(backend, gf->nodes[gf->n_nodes - 1], output->data, 0, ggml_nbytes(output));
+            auto result = gf->nodes[gf->n_nodes - 1];
+            if (*output == NULL && output_ctx != NULL) {
+                *output = ggml_dup_tensor(output_ctx, result);
+            }
+            if (*output != NULL) {
+                ggml_backend_tensor_get_and_sync(backend, result, (*output)->data, 0, ggml_nbytes(*output));
+            }
+        }
+
+        if (free_compute_buffer_immediately) {
+            free_compute_buffer();
+        }
+    }
+};
+
+class GGMLBlock {
+private:
+    static char temp_buffer[1024 * 1024 * 10];
+    ggml_context* get_temp_ctx() {
+        struct ggml_init_params params;
+        params.mem_size   = sizeof(temp_buffer);
+        params.mem_buffer = temp_buffer;
+        params.no_alloc   = true;
+
+        ggml_context* temp_ctx = ggml_init(params);
+        GGML_ASSERT(temp_ctx != NULL);
+        return temp_ctx;
+    }
+
+protected:
+    typedef std::unordered_map<std::string, struct ggml_tensor*> ParameterMap;
+    typedef std::unordered_map<std::string, std::shared_ptr<GGMLBlock>> GGMLBlockMap;
+    GGMLBlockMap blocks;
+    ParameterMap params;
+
+    void init_blocks(struct ggml_context* ctx, ggml_type wtype) {
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            block->init(ctx, wtype);
         }
     }
 
-    void free_compute_buffer() {
-        ggml_allocr_free(compute_allocr);
-        ggml_backend_buffer_free(compute_buffer);
-        compute_allocr      = NULL;
-        compute_buffer_size = 0;
+    virtual void init_params(struct ggml_context* ctx, ggml_type wtype) {}
+
+public:
+    void init(struct ggml_context* ctx, ggml_type wtype) {
+        init_blocks(ctx, wtype);
+        init_params(ctx, wtype);
+    }
+
+    std::tuple<size_t, size_t> get_params_info(ggml_type wtype) {
+        ggml_context* temp_ctx = get_temp_ctx();
+        init(temp_ctx, wtype);
+        size_t num_tensors = get_params_num();
+        size_t mem_size    = get_params_mem_size();
+        return {num_tensors, mem_size};
+    }
+
+    size_t get_params_num() {
+        size_t num_tensors = params.size();
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            num_tensors += block->get_params_num();
+        }
+        return num_tensors;
+    };
+
+    size_t get_params_mem_size() {
+        size_t mem_size = 0;
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            mem_size += block->get_params_mem_size();
+        }
+
+        for (auto& pair : params) {
+            mem_size += ggml_nbytes(pair.second);
+        }
+
+        return mem_size;
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            block->get_param_tensors(tensors, prefix + pair.first);
+        }
+
+        for (auto& pair : params) {
+            struct ggml_tensor* param = pair.second;
+
+            tensors[prefix + pair.first] = pair.second;
+        }
+    }
+};
+
+class UnaryBlock : public GGMLBlock {
+public:
+    virtual struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) = 0;
+};
+
+class Linear : public UnaryBlock {
+protected:
+    int64_t in_features;
+    int64_t out_features;
+    bool bias;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
+        }
+    }
+
+public:
+    Linear(int64_t in_features,
+           int64_t out_features,
+           bool bias = true)
+        : in_features(in_features),
+          out_features(out_features),
+          bias(bias) {}
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        struct ggml_tensor* w = params["weight"];
+        struct ggml_tensor* b = NULL;
+        if (bias) {
+            b = params["bias"];
+        }
+        return ggml_nn_linear(ctx, x, w, b);
+    }
+};
+
+class Conv2d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kernel_size.second, kernel_size.first, in_channels, out_channels);
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+public:
+    Conv2d(int64_t in_channels,
+           int64_t out_channels,
+           std::pair<int, int> kernel_size,
+           std::pair<int, int> stride   = {1, 1},
+           std::pair<int, int> padding  = {0, 0},
+           std::pair<int, int> dilation = {1, 1},
+           bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        struct ggml_tensor* w = params["weight"];
+        struct ggml_tensor* b = NULL;
+        if (bias) {
+            b = params["bias"];
+        }
+        return ggml_nn_conv_2d(ctx, x, w, b, stride.second, stride.first, padding.second, padding.first, dilation.second, dilation.first);
+    }
+};
+
+class Conv3dnx1x1 : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    int64_t kernel_size;
+    int64_t stride;
+    int64_t padding;
+    int64_t dilation;
+    bool bias;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, kernel_size, in_channels, out_channels);  // 5d => 4d
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+public:
+    Conv3dnx1x1(int64_t in_channels,
+                int64_t out_channels,
+                int64_t kernel_size,
+                int64_t stride   = 1,
+                int64_t padding  = 0,
+                int64_t dilation = 1,
+                bool bias        = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    // x: [N, IC, ID, IH*IW]
+    // result: [N, OC, OD, OH*OW]
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        struct ggml_tensor* w = params["weight"];
+        struct ggml_tensor* b = NULL;
+        if (bias) {
+            b = params["bias"];
+        }
+        return ggml_nn_conv_3d_nx1x1(ctx, x, w, b, stride, padding, dilation);
+    }
+};
+
+class LayerNorm : public UnaryBlock {
+protected:
+    int64_t normalized_shape;
+    float eps;
+    bool elementwise_affine;
+    bool bias;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        if (elementwise_affine) {
+            params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, normalized_shape);
+            if (bias) {
+                params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, normalized_shape);
+            }
+        }
+    }
+
+public:
+    LayerNorm(int64_t normalized_shape,
+              float eps               = 1e-05f,
+              bool elementwise_affine = true,
+              bool bias               = true)
+        : normalized_shape(normalized_shape),
+          eps(eps),
+          elementwise_affine(elementwise_affine),
+          bias(bias) {}
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        struct ggml_tensor* w = NULL;
+        struct ggml_tensor* b = NULL;
+
+        if (elementwise_affine) {
+            w = params["weight"];
+            if (bias) {
+                b = params["bias"];
+            }
+        }
+        return ggml_nn_layer_norm(ctx, x, w, b, eps);
+    }
+};
+
+class GroupNorm : public GGMLBlock {
+protected:
+    int64_t num_groups;
+    int64_t num_channels;
+    float eps;
+    bool affine;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        if (affine) {
+            params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
+            params["bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
+        }
+    }
+
+public:
+    GroupNorm(int64_t num_groups,
+              int64_t num_channels,
+              float eps   = 1e-05f,
+              bool affine = true)
+        : num_groups(num_groups),
+          num_channels(num_channels),
+          eps(eps),
+          affine(affine) {}
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        struct ggml_tensor* w = NULL;
+        struct ggml_tensor* b = NULL;
+        if (affine) {
+            w = params["weight"];
+            b = params["bias"];
+        }
+        return ggml_nn_group_norm(ctx, x, w, b, num_groups);
+    }
+};
+
+class GroupNorm32 : public GroupNorm {
+public:
+    GroupNorm32(int64_t num_channels)
+        : GroupNorm(32, num_channels, 1e-06f) {}
+};
+
+class MultiheadAttention : public GGMLBlock {
+protected:
+    int64_t embed_dim;
+    int64_t n_head;
+    bool bias;
+    bool mask;
+
+public:
+    MultiheadAttention(int64_t embed_dim,
+                       int64_t n_head,
+                       bool bias = true)
+        : embed_dim(embed_dim),
+          n_head(n_head),
+          bias(bias) {
+        blocks["q_proj"]   = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, bias));
+        blocks["k_proj"]   = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, bias));
+        blocks["v_proj"]   = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, bias));
+        blocks["out_proj"] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, bias));
+    }
+
+    // x: [N, n_token, embed_dim]
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, bool mask = false) {
+        auto q_proj   = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
+        auto k_proj   = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
+        auto v_proj   = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
+        auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
+
+        int64_t N       = x->ne[2];
+        int64_t n_token = x->ne[1];
+        int64_t d_head  = embed_dim / n_head;
+
+        struct ggml_tensor* q = q_proj->forward(ctx, x);
+        q                     = ggml_reshape_4d(ctx, q, d_head, n_head, n_token, N);   // [N, n_token, n_head, d_head]
+        q                     = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));      // [N, n_head, n_token, d_head]
+        q                     = ggml_reshape_3d(ctx, q, d_head, n_token, n_head * N);  // [N * n_head, n_token, d_head]
+
+        struct ggml_tensor* k = k_proj->forward(ctx, x);
+        k                     = ggml_reshape_4d(ctx, k, d_head, n_head, n_token, N);  // [N, n_token, n_head, d_head]
+        k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));     // [N, n_head, n_token, d_head]
+        k                     = ggml_reshape_3d(ctx, k, d_head, n_token, n_head);     // [N * n_head, n_token, d_head]
+
+        struct ggml_tensor* v = v_proj->forward(ctx, x);
+        v                     = ggml_reshape_4d(ctx, v, d_head, n_head, n_token, N);   // [N, n_token, n_head, d_head]
+        v                     = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));      // [N, n_head, d_head, n_token]
+        v                     = ggml_reshape_3d(ctx, v, n_token, d_head, n_head * N);  // [N * n_head, d_head, n_token]
+
+        struct ggml_tensor* kqv = ggml_nn_attention(ctx, q, k, v, mask);  // [N * n_head, n_token, d_head]
+
+        kqv = ggml_reshape_4d(ctx, kqv, d_head, n_token, n_head, N);
+        kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));  // [N, n_token, n_head, d_head]
+
+        x = ggml_reshape_2d(ctx, kqv, d_head * n_head, n_token * N);  // [N * n_token, d_head * n_head]
+
+        x = out_proj->forward(ctx, x);
+        return x;
     }
 };
 

lora.hpp

@@ -12,26 +12,28 @@ struct LoraModel : public GGMLModule {
     ModelLoader model_loader;
     bool load_failed = false;
 
-    LoraModel(const std::string file_path = "")
-        : file_path(file_path) {
-        name = "lora";
+    LoraModel(ggml_backend_t backend,
+              ggml_type wtype,
+              const std::string file_path = "")
+        : file_path(file_path), GGMLModule(backend, wtype) {
         if (!model_loader.init_from_file(file_path)) {
             load_failed = true;
         }
     }
 
-    size_t get_num_tensors() {
+    std::string get_desc() {
+        return "lora";
+    }
+
+    size_t get_params_num() {
         return LORA_GRAPH_SIZE;
     }
 
-    size_t calculate_mem_size() {
-        return model_loader.cal_mem_size(NULL);
+    size_t get_params_mem_size() {
+        return model_loader.get_params_mem_size(NULL);
     }
 
-    bool load_from_file(ggml_backend_t backend) {
-        if (!alloc_params_buffer(backend)) {
-            return false;
-        }
+    bool load_from_file() {
         LOG_INFO("loading LoRA from '%s'", file_path.c_str());
 
         if (load_failed) {
@@ -39,42 +41,35 @@ struct LoraModel : public GGMLModule {
             return false;
         }
 
-        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
-
+        bool dry_run          = true;
         auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
             const std::string& name = tensor_storage.name;
 
-            struct ggml_tensor* real = ggml_new_tensor(params_ctx, tensor_storage.type, tensor_storage.n_dims, tensor_storage.ne);
-            ggml_allocr_alloc(alloc, real);
-
-            *dst_tensor = real;
-
-            lora_tensors[name] = real;
+            if (dry_run) {
+                struct ggml_tensor* real = ggml_new_tensor(params_ctx,
+                                                           tensor_storage.type,
+                                                           tensor_storage.n_dims,
+                                                           tensor_storage.ne);
+                lora_tensors[name]       = real;
+            } else {
+                auto real   = lora_tensors[name];
+                *dst_tensor = real;
+            }
             return true;
         };
 
+        model_loader.load_tensors(on_new_tensor_cb, backend);
+        alloc_params_buffer();
+
+        dry_run = false;
         model_loader.load_tensors(on_new_tensor_cb, backend);
 
         LOG_DEBUG("finished loaded lora");
-        ggml_allocr_free(alloc);
         return true;
     }
 
     struct ggml_cgraph* build_graph(std::map<std::string, struct ggml_tensor*> model_tensors) {
-        // make a graph to compute all lora, expected lora and models tensors are in the same backend
-        // 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() * LORA_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, LORA_GRAPH_SIZE, false);
+        struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, LORA_GRAPH_SIZE, false);
 
         std::set<std::string> applied_lora_tensors;
         for (auto it : model_tensors) {
@@ -113,7 +108,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,36 +118,29 @@ 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];
-            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];
-            lora_down              = ggml_reshape_2d(ctx0, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
+            int64_t lora_up_rows   = lora_up->ne[ggml_n_dims(lora_up) - 1];
+            lora_up                = ggml_reshape_2d(compute_ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
+            int64_t lora_down_rows = lora_down->ne[ggml_n_dims(lora_down) - 1];
+            lora_down              = ggml_reshape_2d(compute_ctx, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
 
             // ggml_mul_mat requires tensor b transposed
-            lora_down                  = ggml_cont(ctx0, ggml_transpose(ctx0, lora_down));
-            struct ggml_tensor* updown = ggml_mul_mat(ctx0, lora_up, lora_down);
-            updown                     = ggml_cont(ctx0, ggml_transpose(ctx0, updown));
-            updown                     = ggml_reshape(ctx0, updown, weight);
+            lora_down                  = ggml_cont(compute_ctx, ggml_transpose(compute_ctx, lora_down));
+            struct ggml_tensor* updown = ggml_mul_mat(compute_ctx, lora_up, lora_down);
+            updown                     = ggml_cont(compute_ctx, ggml_transpose(compute_ctx, updown));
+            updown                     = ggml_reshape(compute_ctx, updown, weight);
             GGML_ASSERT(ggml_nelements(updown) == ggml_nelements(weight));
-            updown = ggml_scale_inplace(ctx0, updown, lora_scale);
+            updown = ggml_scale_inplace(compute_ctx, 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);
-            //     final_weight = ggml_cpy_inplace(ctx0, weight, final_weight);
-            //     final_weight = ggml_add_inplace(ctx0, final_weight, updown);
-            //     final_weight = ggml_cpy_inplace(ctx0, final_weight, weight);
+            //     final_weight = ggml_new_tensor(compute_ctx, GGML_TYPE_F32, weight->n_dims, weight->ne);
+            //     final_weight = ggml_cpy_inplace(compute_ctx, weight, final_weight);
+            //     final_weight = ggml_add_inplace(compute_ctx, final_weight, updown);
+            //     final_weight = ggml_cpy_inplace(compute_ctx, final_weight, weight);
             // } else {
-            //     final_weight = ggml_add_inplace(ctx0, weight, updown);
+            //     final_weight = ggml_add_inplace(compute_ctx, weight, updown);
             // }
-            final_weight = ggml_add_inplace(ctx0, weight, updown);  // apply directly
+            final_weight = ggml_add_inplace(compute_ctx, weight, updown);  // apply directly
             ggml_build_forward_expand(gf, final_weight);
         }
 
@@ -165,20 +153,11 @@ struct LoraModel : public GGMLModule {
         return gf;
     }
 
-    void alloc_compute_buffer(std::map<std::string, struct ggml_tensor*> model_tensors) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(model_tensors);
-        };
-        GGMLModule::alloc_compute_buffer(get_graph);
-    }
-
     void apply(std::map<std::string, struct ggml_tensor*> model_tensors, int n_threads) {
-        alloc_compute_buffer(model_tensors);
-
         auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(model_tensors);
         };
-        GGMLModule::compute(get_graph, n_threads);
+        GGMLModule::compute(get_graph, n_threads, true);
     }
 };
 

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",
 };
@@ -107,6 +108,14 @@ std::unordered_map<std::string, std::string> open_clip_to_hf_clip_model = {
     {"model.positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"},
     {"model.token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"},
     {"model.text_projection", "transformer.text_model.text_projection"},
+    {"model.visual.class_embedding", "transformer.visual_model.embeddings.class_embedding"},
+    {"model.visual.conv1.weight", "transformer.visual_model.embeddings.patch_embedding.weight"},
+    {"model.visual.ln_post.bias", "transformer.visual_model.post_layernorm.bias"},
+    {"model.visual.ln_post.weight", "transformer.visual_model.post_layernorm.weight"},
+    {"model.visual.ln_pre.bias", "transformer.visual_model.pre_layernorm.bias"},
+    {"model.visual.ln_pre.weight", "transformer.visual_model.pre_layernorm.weight"},
+    {"model.visual.positional_embedding", "transformer.visual_model.embeddings.position_embedding.weight"},
+    {"model.visual.proj", "transformer.visual_model.visual_projection"},
 };
 
 std::unordered_map<std::string, std::string> open_clip_to_hk_clip_resblock = {
@@ -136,7 +145,10 @@ std::unordered_map<std::string, std::string> vae_decoder_name_map = {
 std::string convert_open_clip_to_hf_clip(const std::string& name) {
     std::string new_name = name;
     std::string prefix;
-    if (starts_with(new_name, "conditioner.embedders.0.")) {
+    if (starts_with(new_name, "conditioner.embedders.0.open_clip.")) {
+        prefix   = "cond_stage_model.";
+        new_name = new_name.substr(strlen("conditioner.embedders.0.open_clip."));
+    } else if (starts_with(new_name, "conditioner.embedders.0.")) {
         prefix   = "cond_stage_model.";
         new_name = new_name.substr(strlen("conditioner.embedders.0."));
     } else if (starts_with(new_name, "conditioner.embedders.1.")) {
@@ -148,25 +160,35 @@ std::string convert_open_clip_to_hf_clip(const std::string& name) {
     } else {
         return new_name;
     }
-    std::string open_clip_resblock_prefix = "model.transformer.resblocks.";
-    std::string hf_clip_resblock_prefix   = "transformer.text_model.encoder.layers.";
 
     if (open_clip_to_hf_clip_model.find(new_name) != open_clip_to_hf_clip_model.end()) {
         new_name = open_clip_to_hf_clip_model[new_name];
     }
 
-    if (new_name.find(open_clip_resblock_prefix) == 0) {
-        std::string remain = new_name.substr(open_clip_resblock_prefix.length());
-        std::string idx    = remain.substr(0, remain.find("."));
-        std::string suffix = remain.substr(idx.length() + 1);
+    std::string open_clip_resblock_prefix = "model.transformer.resblocks.";
+    std::string hf_clip_resblock_prefix   = "transformer.text_model.encoder.layers.";
 
-        if (suffix == "attn.in_proj_weight" || suffix == "attn.in_proj_bias") {
-            new_name = hf_clip_resblock_prefix + idx + "." + suffix;
-        } else if (open_clip_to_hk_clip_resblock.find(suffix) != open_clip_to_hk_clip_resblock.end()) {
-            std::string new_suffix = open_clip_to_hk_clip_resblock[suffix];
-            new_name               = hf_clip_resblock_prefix + idx + "." + new_suffix;
+    auto replace_suffix = [&]() {
+        if (new_name.find(open_clip_resblock_prefix) == 0) {
+            std::string remain = new_name.substr(open_clip_resblock_prefix.length());
+            std::string idx    = remain.substr(0, remain.find("."));
+            std::string suffix = remain.substr(idx.length() + 1);
+
+            if (suffix == "attn.in_proj_weight" || suffix == "attn.in_proj_bias") {
+                new_name = hf_clip_resblock_prefix + idx + "." + suffix;
+            } else if (open_clip_to_hk_clip_resblock.find(suffix) != open_clip_to_hk_clip_resblock.end()) {
+                std::string new_suffix = open_clip_to_hk_clip_resblock[suffix];
+                new_name               = hf_clip_resblock_prefix + idx + "." + new_suffix;
+            }
         }
-    }
+    };
+
+    replace_suffix();
+
+    open_clip_resblock_prefix = "model.visual.transformer.resblocks.";
+    hf_clip_resblock_prefix   = "transformer.visual_model.encoder.layers.";
+
+    replace_suffix();
 
     return prefix + new_name;
 }
@@ -374,6 +396,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) {
@@ -431,7 +458,7 @@ void preprocess_tensor(TensorStorage tensor_storage,
 
     tensor_storage.name = new_name;
 
-    if (new_name.find("transformer.text_model.encoder.layers.") != std::string::npos &&
+    if (new_name.find("cond_stage_model") != std::string::npos &&
         ends_with(new_name, "attn.in_proj_weight")) {
         size_t prefix_size = new_name.find("attn.in_proj_weight");
         std::string prefix = new_name.substr(0, prefix_size);
@@ -443,7 +470,7 @@ void preprocess_tensor(TensorStorage tensor_storage,
 
         processed_tensor_storages.insert(processed_tensor_storages.end(), chunks.begin(), chunks.end());
 
-    } else if (new_name.find("transformer.text_model.encoder.layers.") != std::string::npos &&
+    } else if (new_name.find("cond_stage_model") != std::string::npos &&
                ends_with(new_name, "attn.in_proj_bias")) {
         size_t prefix_size = new_name.find("attn.in_proj_bias");
         std::string prefix = new_name.substr(0, prefix_size);
@@ -471,7 +498,13 @@ void bf16_to_f32_vec(uint16_t* src, float* dst, int64_t n) {
     }
 }
 
-void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type, int n) {
+void convert_tensor(void* src,
+                    ggml_type src_type,
+                    void* dst,
+                    ggml_type dst_type,
+                    int nrows,
+                    int n_per_row) {
+    int n = nrows * n_per_row;
     if (src_type == dst_type) {
         size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
         memcpy(((char*)dst), ((char*)src), nbytes);
@@ -480,7 +513,9 @@ void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type
             ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
         } else {
             int64_t hist[16];
-            ggml_quantize_chunk(dst_type, (float*)src, dst, 0, n, hist);
+            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            const float* im = imatrix.data();
+            ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, hist, im);
         }
     } else if (dst_type == GGML_TYPE_F32) {
         if (src_type == GGML_TYPE_F16) {
@@ -509,7 +544,9 @@ void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type
             ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
         } else {
             int64_t hist[16];
-            ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, n, hist);
+            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            const float* im = imatrix.data();
+            ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, hist, im);
         }
     }
 }
@@ -609,7 +646,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 +710,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());
 
@@ -772,17 +809,26 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
             return false;
         }
 
-        if (shape.size() > 4) {
+        if (shape.size() > SD_MAX_DIMS) {
             LOG_ERROR("invalid tensor '%s'", name.c_str());
             return false;
         }
 
-        int n_dims    = (int)shape.size();
-        int64_t ne[4] = {1, 1, 1, 1};
+        int n_dims              = (int)shape.size();
+        int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
         for (int i = 0; i < n_dims; i++) {
             ne[i] = shape[i].get<int64_t>();
         }
 
+        if (n_dims == 5) {
+            if (ne[3] == 1 && ne[4] == 1) {
+                n_dims = 4;
+            } else {
+                LOG_ERROR("invalid tensor '%s'", name.c_str());
+                return false;
+            }
+        }
+
         TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
 
         tensor_storage.reverse_ne();
@@ -797,6 +843,8 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
         }
 
         tensor_storages.push_back(tensor_storage);
+
+        // LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
     }
 
     return true;
@@ -940,7 +988,7 @@ struct PickleTensorReader {
                 phase = READ_NAME;
             }
         } else if (phase == READ_DIMENS) {
-            if (tensor_storage.n_dims + 1 > 4) {  // too many dimens
+            if (tensor_storage.n_dims + 1 > SD_MAX_DIMS) {  // too many dimens
                 phase                 = READ_NAME;
                 tensor_storage.n_dims = 0;
             }
@@ -1175,12 +1223,18 @@ bool ModelLoader::init_from_ckpt_file(const std::string& file_path, const std::s
 }
 
 SDVersion ModelLoader::get_sd_version() {
-    // return VERSION_1_x;
     TensorStorage token_embedding_weight;
     for (auto& tensor_storage : tensor_storages) {
         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.find("model.diffusion_model.input_blocks.8.0.time_mixer.mix_factor") != std::string::npos) {
+            return VERSION_SVD;
+        }
+
         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 +1272,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,24 +1358,11 @@ 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;
 
             success = on_new_tensor_cb(tensor_storage, &dst_tensor);
@@ -1308,11 +1377,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());
@@ -1332,7 +1397,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
                     }
 
                     convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
-                                   dst_tensor->type, (int)tensor_storage.nelements());
+                                   dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
                 }
             } else {
                 read_buffer.resize(tensor_storage.nbytes());
@@ -1351,7 +1416,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
                     convert_buffer.resize(ggml_nbytes(dst_tensor));
                     convert_tensor((void*)read_buffer.data(), tensor_storage.type,
                                    (void*)convert_buffer.data(), dst_tensor->type,
-                                   (int)tensor_storage.nelements());
+                                   (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
                     ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
                 }
             }
@@ -1374,15 +1439,19 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
     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;
+        // LOG_DEBUG("%s", tensor_storage.to_string().c_str());
         tensor_names_in_file.insert(name);
 
         struct ggml_tensor* real;
         if (tensors.find(name) != tensors.end()) {
             real = tensors[name];
         } else {
-            if (ignore_tensors.find(name) == ignore_tensors.end()) {
-                LOG_WARN("unknown tensor '%s' in model file", name.c_str());
+            for (auto& ignore_tensor : ignore_tensors) {
+                if (starts_with(name, ignore_tensor)) {
+                    return true;
+                }
             }
+            LOG_INFO("unknown tensor '%s' in model file", tensor_storage.to_string().c_str());
             return true;
         }
 
@@ -1417,6 +1486,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 +1506,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 += get_params_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::get_params_mem_size(ggml_backend_t backend, ggml_type type) {
     size_t alignment = 128;
     if (backend != NULL) {
         alignment = ggml_backend_get_alignment(backend);
@@ -1449,8 +1575,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,8 @@
 #include <functional>
 #include <map>
 #include <memory>
+#include <set>
+#include <sstream>
 #include <string>
 #include <vector>
 
@@ -12,19 +14,22 @@
 #include "json.hpp"
 #include "zip.h"
 
+#define SD_MAX_DIMS 5
+
 enum SDVersion {
     VERSION_1_x,
     VERSION_2_x,
     VERSION_XL,
+    VERSION_SVD,
     VERSION_COUNT,
 };
 
 struct TensorStorage {
     std::string name;
-    ggml_type type = GGML_TYPE_F32;
-    bool is_bf16   = false;
-    int64_t ne[4]  = {1, 1, 1, 1};
-    int n_dims     = 0;
+    ggml_type type          = GGML_TYPE_F32;
+    bool is_bf16            = false;
+    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+    int n_dims              = 0;
 
     size_t file_index = 0;
     int index_in_zip  = -1;  // >= means stored in a zip file
@@ -40,7 +45,11 @@ struct TensorStorage {
     }
 
     int64_t nelements() const {
-        return ne[0] * ne[1] * ne[2] * ne[3];
+        int64_t n = 1;
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            n *= ne[i];
+        }
+        return n;
     }
 
     int64_t nbytes() const {
@@ -68,6 +77,7 @@ struct TensorStorage {
     std::vector<TensorStorage> chunk(size_t n) {
         std::vector<TensorStorage> chunks;
         size_t chunk_size = nbytes_to_read() / n;
+        // printf("%d/%d\n", chunk_size, nbytes_to_read());
         reverse_ne();
         for (int i = 0; i < n; i++) {
             TensorStorage chunk_i = *this;
@@ -81,7 +91,7 @@ struct TensorStorage {
     }
 
     void reverse_ne() {
-        int64_t new_ne[4] = {1, 1, 1, 1};
+        int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
         for (int i = 0; i < n_dims; i++) {
             new_ne[i] = ne[n_dims - 1 - i];
         }
@@ -89,10 +99,27 @@ struct TensorStorage {
             ne[i] = new_ne[i];
         }
     }
+
+    std::string to_string() const {
+        std::stringstream ss;
+        const char* type_name = ggml_type_name(type);
+        if (is_bf16) {
+            type_name = "bf16";
+        }
+        ss << name << " | " << type_name << " | ";
+        ss << n_dims << " [";
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            ss << ne[i];
+            if (i != SD_MAX_DIMS - 1) {
+                ss << ", ";
+            }
+        }
+        ss << "]";
+        return ss.str();
+    }
 };
 
 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 +147,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 get_params_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.h

@@ -65,12 +65,17 @@ 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_IQ2_XS  = 17,
+    SD_TYPE_IQ3_XXS = 18,
+    SD_TYPE_IQ1_S   = 19,
+    SD_TYPE_IQ4_NL  = 20,
     SD_TYPE_I8,
     SD_TYPE_I16,
     SD_TYPE_I32,
@@ -104,14 +109,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 +133,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,
@@ -141,6 +151,21 @@ SD_API sd_image_t* img2img(sd_ctx_t* sd_ctx,
                            int64_t seed,
                            int batch_count);
 
+SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
+                           sd_image_t init_image,
+                           int width,
+                           int height,
+                           int video_frames,
+                           int motion_bucket_id,
+                           int fps,
+                           float augmentation_level,
+                           float min_cfg,
+                           float cfg_scale,
+                           enum sample_method_t sample_method,
+                           int sample_steps,
+                           float strength,
+                           int64_t seed);
+
 typedef struct upscaler_ctx_t upscaler_ctx_t;
 
 SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
@@ -148,7 +173,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

@@ -8,88 +8,45 @@
 /*
     ===================================    TinyAutoEncoder  ===================================
     References:
-    https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoder_tiny.py
+    https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py
     https://github.com/madebyollin/taesd/blob/main/taesd.py
 
 */
-struct TAEBlock {
-    int in_channels;
-    int out_channels;
 
-    // conv
-    ggml_tensor* conv_0_w;  // [in_channels, out_channels, 3, 3]
-    ggml_tensor* conv_0_b;  // [in_channels]
-    ggml_tensor* conv_1_w;  // [out_channels, out_channels, 3, 3]
-    ggml_tensor* conv_1_b;  // [out_channels]
-    ggml_tensor* conv_2_w;  // [out_channels, out_channels, 3, 3]
-    ggml_tensor* conv_2_b;  // [out_channels]
+class TAEBlock : public UnaryBlock {
+protected:
+    int n_in;
+    int n_out;
 
-    // skip
-    ggml_tensor* conv_skip_w;  // [in_channels, out_channels, 1, 1]
-
-    size_t calculate_mem_size() {
-        size_t mem_size = in_channels * out_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_0_w
-        mem_size += in_channels * ggml_type_size(GGML_TYPE_F32);                               // conv_0_b
-        mem_size += out_channels * out_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);       // conv_1_w
-        mem_size += out_channels * ggml_type_size(GGML_TYPE_F32);                              // conv_1_b
-        mem_size += out_channels * out_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);       // conv_1_w
-        mem_size += out_channels * ggml_type_size(GGML_TYPE_F32);                              // conv_1_b
-        mem_size += out_channels * out_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);       // conv_2_w
-        mem_size += out_channels * ggml_type_size(GGML_TYPE_F32);                              // conv_2_b
-
-        if (in_channels != out_channels) {
-            mem_size += in_channels * out_channels * ggml_type_size(GGML_TYPE_F16);  // conv_skip_w
-        }
-        return mem_size;
-    }
-
-    int get_num_tensors() {
-        return 6 + (in_channels != out_channels ? 1 : 0);
-    }
-
-    void init_params(ggml_context* ctx) {
-        conv_0_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, in_channels);
-        conv_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-        conv_1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels);
-        conv_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-
-        conv_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels);
-        conv_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
-
-        if (in_channels != out_channels) {
-            conv_skip_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, out_channels, in_channels);
+public:
+    TAEBlock(int n_in, int n_out)
+        : n_in(n_in), n_out(n_out) {
+        blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
+        if (n_in != n_out) {
+            blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {1, 1}, {1, 1}, {1, 1}, {1, 1}, false));
         }
     }
 
-    void map_by_name(std::map<std::string, ggml_tensor*>& tensors, std::string prefix) {
-        tensors[prefix + "conv.0.weight"] = conv_0_w;
-        tensors[prefix + "conv.0.bias"]   = conv_0_b;
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [n, n_in, h, w]
+        // return: [n, n_out, h, w]
 
-        tensors[prefix + "conv.2.weight"] = conv_1_w;
-        tensors[prefix + "conv.2.bias"]   = conv_1_b;
+        auto conv_0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
+        auto conv_2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
+        auto conv_4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
 
-        tensors[prefix + "conv.4.weight"] = conv_2_w;
-        tensors[prefix + "conv.4.bias"]   = conv_2_b;
+        auto h = conv_0->forward(ctx, x);
+        h      = ggml_relu_inplace(ctx, h);
+        h      = conv_2->forward(ctx, h);
+        h      = ggml_relu_inplace(ctx, h);
+        h      = conv_4->forward(ctx, h);
 
-        if (in_channels != out_channels) {
-            tensors[prefix + "skip.weight"] = conv_skip_w;
-        }
-    }
-
-    ggml_tensor* forward(ggml_context* ctx, ggml_tensor* x) {
-        // conv(n_in, n_out)
-        ggml_tensor* h;
-        h = ggml_nn_conv_2d(ctx, x, conv_0_w, conv_0_b, 1, 1, 1, 1);
-        h = ggml_relu_inplace(ctx, h);
-        h = ggml_nn_conv_2d(ctx, h, conv_1_w, conv_1_b, 1, 1, 1, 1);
-        h = ggml_relu_inplace(ctx, h);
-        h = ggml_nn_conv_2d(ctx, h, conv_2_w, conv_2_b, 1, 1, 1, 1);
-
-        // skip connection
-        if (in_channels != out_channels) {
-            // skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
-            x = ggml_nn_conv_2d(ctx, x, conv_skip_w, NULL, 1, 1, 1, 1);
+        if (n_in != n_out) {
+            auto skip = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
+            LOG_DEBUG("skip");
+            x = skip->forward(ctx, x);
         }
 
         h = ggml_add(ctx, h, x);
@@ -98,425 +55,167 @@ struct TAEBlock {
     }
 };
 
-struct TinyEncoder {
+class TinyEncoder : public UnaryBlock {
     int in_channels = 3;
-    int z_channels  = 4;
     int channels    = 64;
+    int z_channels  = 4;
     int num_blocks  = 3;
 
-    // input
-    ggml_tensor* conv_input_w;  // [channels, in_channels, 3, 3]
-    ggml_tensor* conv_input_b;  // [channels]
-    TAEBlock initial_block;
-
-    ggml_tensor* conv_1_w;  // [channels, channels, 3, 3]
-    TAEBlock input_blocks[3];
-
-    // middle
-    ggml_tensor* conv_2_w;  // [channels, channels, 3, 3]
-    TAEBlock middle_blocks[3];
-
-    // output
-    ggml_tensor* conv_3_w;  // [channels, channels, 3, 3]
-    TAEBlock output_blocks[3];
-
-    // final
-    ggml_tensor* conv_final_w;  // [z_channels, channels, 3, 3]
-    ggml_tensor* conv_final_b;  // [z_channels]
-
+public:
     TinyEncoder() {
+        int index                       = 0;
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, channels, {3, 3}, {1, 1}, {1, 1}));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
         for (int i = 0; i < num_blocks; i++) {
-            input_blocks[i].in_channels  = channels;
-            input_blocks[i].out_channels = channels;
-
-            middle_blocks[i].in_channels  = channels;
-            middle_blocks[i].out_channels = channels;
-
-            output_blocks[i].in_channels  = channels;
-            output_blocks[i].out_channels = channels;
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
         }
 
-        initial_block.in_channels  = channels;
-        initial_block.out_channels = channels;
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    size_t calculate_mem_size() {
-        size_t mem_size = channels * in_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_input_w
-        mem_size += channels * ggml_type_size(GGML_TYPE_F32);                              // conv_input_b
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [n, in_channels, h, w]
+        // return: [n, z_channels, h/8, w/8]
 
-        mem_size += initial_block.calculate_mem_size();
+        for (int i = 0; i < num_blocks * 3 + 6; i++) {
+            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
 
-        mem_size += channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_1_w
-        mem_size += channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_2_w
-        mem_size += channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_3_w
-
-        for (int i = 0; i < num_blocks; i++) {
-            mem_size += input_blocks[i].calculate_mem_size();
-            mem_size += middle_blocks[i].calculate_mem_size();
-            mem_size += output_blocks[i].calculate_mem_size();
-        }
-        mem_size += z_channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_input_w
-        mem_size += z_channels * ggml_type_size(GGML_TYPE_F32);                     // conv_input_b
-        return mem_size;
-    }
-
-    int get_num_tensors() {
-        int num_tensors = 7;
-        for (int i = 0; i < num_blocks; i++) {
-            num_tensors += input_blocks[i].get_num_tensors();
-            num_tensors += middle_blocks[i].get_num_tensors();
-            num_tensors += output_blocks[i].get_num_tensors();
-        }
-        num_tensors += initial_block.get_num_tensors();
-        return num_tensors;
-    }
-
-    void init_params(ggml_context* ctx) {
-        conv_input_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, in_channels, channels);
-        conv_input_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
-
-        initial_block.init_params(ctx);
-
-        conv_1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, channels);
-        conv_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, channels);
-        conv_3_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, channels);
-
-        conv_final_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, z_channels);
-        conv_final_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_channels);
-
-        for (int i = 0; i < num_blocks; i++) {
-            input_blocks[i].init_params(ctx);
-            middle_blocks[i].init_params(ctx);
-            output_blocks[i].init_params(ctx);
-        }
-    }
-
-    void map_by_name(std::map<std::string, ggml_tensor*>& tensors, std::string prefix) {
-        tensors[prefix + "0.weight"] = conv_input_w;
-        tensors[prefix + "0.bias"]   = conv_input_b;
-
-        initial_block.map_by_name(tensors, prefix + "1.");
-
-        tensors[prefix + "2.weight"] = conv_1_w;
-        for (int i = 0; i < num_blocks; i++) {
-            input_blocks[i].map_by_name(tensors, prefix + std::to_string(i + 3) + ".");
+            x = block->forward(ctx, x);
         }
 
-        tensors[prefix + "6.weight"] = conv_2_w;
-        for (int i = 0; i < num_blocks; i++) {
-            middle_blocks[i].map_by_name(tensors, prefix + std::to_string(i + 7) + ".");
-        }
-
-        tensors[prefix + "10.weight"] = conv_3_w;
-        for (int i = 0; i < num_blocks; i++) {
-            output_blocks[i].map_by_name(tensors, prefix + std::to_string(i + 11) + ".");
-        }
-
-        tensors[prefix + "14.weight"] = conv_final_w;
-        tensors[prefix + "14.bias"]   = conv_final_b;
-    }
-
-    ggml_tensor* forward(ggml_context* ctx, ggml_tensor* x) {
-        // conv(3, 64)
-        auto z = ggml_nn_conv_2d(ctx, x, conv_input_w, conv_input_b, 1, 1, 1, 1);
-
-        // Block(64, 64)
-        z = initial_block.forward(ctx, z);
-
-        // conv(64, 64, stride=2, bias=False)
-        z = ggml_nn_conv_2d(ctx, z, conv_1_w, NULL, 2, 2, 1, 1);
-
-        // Block(64, 64), Block(64, 64), Block(64, 64)
-        for (int i = 0; i < num_blocks; i++) {
-            z = input_blocks[i].forward(ctx, z);
-        }
-
-        // conv(64, 64, stride=2, bias=False)
-        z = ggml_nn_conv_2d(ctx, z, conv_2_w, NULL, 2, 2, 1, 1);
-
-        // Block(64, 64), Block(64, 64), Block(64, 64)
-        for (int i = 0; i < num_blocks; i++) {
-            z = middle_blocks[i].forward(ctx, z);
-        }
-
-        // conv(64, 64, stride=2, bias=False)
-        z = ggml_nn_conv_2d(ctx, z, conv_3_w, NULL, 2, 2, 1, 1);
-
-        // Block(64, 64), Block(64, 64), Block(64, 64)
-        for (int i = 0; i < num_blocks; i++) {
-            z = output_blocks[i].forward(ctx, z);
-        }
-
-        // conv(64, 4)
-        z = ggml_nn_conv_2d(ctx, z, conv_final_w, conv_final_b, 1, 1, 1, 1);
-        return z;
+        return x;
     }
 };
 
-struct TinyDecoder {
-    int z_channels      = 4;
-    int channels        = 64;
-    int output_channels = 3;
-    int num_blocks      = 3;
+class TinyDecoder : public UnaryBlock {
+    int z_channels   = 4;
+    int channels     = 64;
+    int out_channels = 3;
+    int num_blocks   = 3;
 
-    // input
-    ggml_tensor* conv_input_w;  // [channels, z_channels, 3, 3]
-    ggml_tensor* conv_input_b;  // [channels]
-    TAEBlock input_blocks[3];
-    ggml_tensor* conv_1_w;  // [channels, channels, 3, 3]
+public:
+    TinyDecoder(int index = 0) {
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels, {3, 3}, {1, 1}, {1, 1}));
+        index++;  // nn.ReLU()
 
-    // middle
-    TAEBlock middle_blocks[3];
-    ggml_tensor* conv_2_w;  // [channels, channels, 3, 3]
-
-    // output
-    TAEBlock output_blocks[3];
-    ggml_tensor* conv_3_w;  // [channels, channels, 3, 3]
-
-    // final
-    TAEBlock final_block;
-    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;
-            input_blocks[i].out_channels = channels;
-
-            middle_blocks[i].in_channels  = channels;
-            middle_blocks[i].out_channels = channels;
-
-            output_blocks[i].in_channels  = channels;
-            output_blocks[i].out_channels = channels;
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
         }
+        index++;  // nn.Upsample()
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
 
-        final_block.in_channels  = channels;
-        final_block.out_channels = channels;
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+        index++;  // nn.Upsample()
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+        index++;  // nn.Upsample()
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    size_t calculate_mem_size() {
-        size_t mem_size = channels * z_channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_input_w
-        mem_size += channels * ggml_type_size(GGML_TYPE_F32);                             // conv_input_b
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* z) {
+        // z: [n, z_channels, h, w]
+        // return: [n, out_channels, h*8, w*8]
 
-        for (int i = 0; i < num_blocks; i++) {
-            mem_size += input_blocks[i].calculate_mem_size();
-        }
-        mem_size += channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_1_w
-
-        for (int i = 0; i < num_blocks; i++) {
-            mem_size += middle_blocks[i].calculate_mem_size();
-        }
-        mem_size += channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_2_w
-
-        for (int i = 0; i < num_blocks; i++) {
-            mem_size += output_blocks[i].calculate_mem_size();
-        }
-        mem_size += channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_3_w
-
-        mem_size += final_block.calculate_mem_size();
-        mem_size += output_channels * channels * 3 * 3 * ggml_type_size(GGML_TYPE_F16);  // conv_input_w
-        mem_size += output_channels * ggml_type_size(GGML_TYPE_F32);                     // conv_input_b
-        return mem_size;
-    }
-
-    int get_num_tensors() {
-        int num_tensors = 9;
-        for (int i = 0; i < num_blocks; i++) {
-            num_tensors += input_blocks[i].get_num_tensors();
-            num_tensors += middle_blocks[i].get_num_tensors();
-            num_tensors += output_blocks[i].get_num_tensors();
-        }
-        num_tensors += final_block.get_num_tensors();
-        return num_tensors;
-    }
-
-    void init_params(ggml_allocr* alloc, ggml_context* ctx) {
-        conv_input_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, z_channels, channels);
-        conv_input_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
-
-        conv_1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, channels);
-        conv_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, channels);
-        conv_3_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, channels);
-
-        conv_final_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, output_channels);
-        conv_final_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_channels);
-
-        for (int i = 0; i < num_blocks; i++) {
-            input_blocks[i].init_params(ctx);
-            middle_blocks[i].init_params(ctx);
-            output_blocks[i].init_params(ctx);
-        }
-
-        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) {
-        tensors[prefix + "0.weight"] = conv_input_w;
-        tensors[prefix + "0.bias"]   = conv_input_b;
-
-        for (int i = 0; i < num_blocks; i++) {
-            input_blocks[i].map_by_name(tensors, prefix + std::to_string(i + 2) + ".");
-        }
-
-        tensors[prefix + "6.weight"] = conv_1_w;
-        for (int i = 0; i < num_blocks; i++) {
-            middle_blocks[i].map_by_name(tensors, prefix + std::to_string(i + 7) + ".");
-        }
-
-        tensors[prefix + "11.weight"] = conv_2_w;
-        for (int i = 0; i < num_blocks; i++) {
-            output_blocks[i].map_by_name(tensors, prefix + std::to_string(i + 12) + ".");
-        }
-
-        tensors[prefix + "16.weight"] = conv_3_w;
-
-        final_block.map_by_name(tensors, prefix + "17.");
-
-        tensors[prefix + "18.weight"] = conv_final_w;
-        tensors[prefix + "18.bias"]   = conv_final_b;
-    }
-
-    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);
+        for (int i = 0; i < num_blocks * 3 + 10; i++) {
+            if (blocks.find(std::to_string(i)) == blocks.end()) {
+                if (i == 1) {
+                    h = ggml_relu_inplace(ctx, h);
+                } else {
+                    h = ggml_upscale(ctx, h, 2);
+                }
+                continue;
+            }
+            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
 
-        // nn.ReLU()
-        h = ggml_relu_inplace(ctx, h);
-
-        // Block(64, 64), Block(64, 64), Block(64, 64)
-        for (int i = 0; i < num_blocks; i++) {
-            h = input_blocks[i].forward(ctx, h);
+            h = block->forward(ctx, h);
         }
 
-        // nn.Upsample(scale_factor=2)
-        h = ggml_upscale(ctx, h, 2);
-
-        // conv(64, 64, bias=False)
-        h = ggml_nn_conv_2d(ctx, h, conv_1_w, NULL, 1, 1, 1, 1);
-
-        // Block(64, 64), Block(64, 64), Block(64, 64)
-        for (int i = 0; i < num_blocks; i++) {
-            h = middle_blocks[i].forward(ctx, h);
-        }
-
-        // nn.Upsample(scale_factor=2)
-        h = ggml_upscale(ctx, h, 2);
-
-        // conv(64, 64, bias=False)
-        h = ggml_nn_conv_2d(ctx, h, conv_2_w, NULL, 1, 1, 1, 1);
-
-        // Block(64, 64), Block(64, 64), Block(64, 64)
-        for (int i = 0; i < num_blocks; i++) {
-            h = output_blocks[i].forward(ctx, h);
-        }
-
-        // nn.Upsample(scale_factor=2)
-        h = ggml_upscale(ctx, h, 2);
-
-        // conv(64, 64, bias=False)
-        h = ggml_nn_conv_2d(ctx, h, conv_3_w, NULL, 1, 1, 1, 1);
-
-        // Block(64, 64)
-        h = final_block.forward(ctx, h);
-
-        // conv(64, 3)
-        h = ggml_nn_conv_2d(ctx, h, conv_final_w, conv_final_b, 1, 1, 1, 1);
         return h;
     }
 };
 
+class TAESD : public GGMLBlock {
+protected:
+    bool decode_only;
+
+public:
+    TAESD(bool decode_only = true)
+        : decode_only(decode_only) {
+        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder());
+
+        if (!decode_only) {
+            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder());
+        }
+    }
+
+    struct ggml_tensor* decode(struct ggml_context* ctx, struct ggml_tensor* z) {
+        auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
+        return decoder->forward(ctx, z);
+    }
+
+    struct ggml_tensor* encode(struct ggml_context* ctx, struct ggml_tensor* x) {
+        auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
+        return encoder->forward(ctx, x);
+    }
+};
+
 struct TinyAutoEncoder : public GGMLModule {
-    TinyEncoder encoder;
-    TinyDecoder decoder;
+    TAESD taesd;
     bool decode_only = false;
 
-    TinyAutoEncoder(bool decoder_only_ = true)
-        : decode_only(decoder_only_) {
-        name = "tae";
+    TinyAutoEncoder(ggml_backend_t backend,
+                    ggml_type wtype,
+                    bool decoder_only = true)
+        : decode_only(decoder_only),
+          taesd(decode_only),
+          GGMLModule(backend, wtype) {
+        taesd.init(params_ctx, wtype);
     }
 
-    size_t calculate_mem_size() {
-        size_t mem_size = decoder.calculate_mem_size();
-        if (!decode_only) {
-            mem_size += encoder.calculate_mem_size();
-        }
-        mem_size += 1024;  // padding
-        return mem_size;
+    std::string get_desc() {
+        return "taesd";
     }
 
-    size_t get_num_tensors() {
-        size_t num_tensors = decoder.get_num_tensors();
-        if (!decode_only) {
-            num_tensors += encoder.get_num_tensors();
-        }
-        return num_tensors;
+    size_t get_params_mem_size() {
+        return taesd.get_params_mem_size();
     }
 
-    void init_params() {
-        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
-        decoder.init_params(alloc, params_ctx);
-        if (!decode_only) {
-            encoder.init_params(params_ctx);
-        }
-
-        // 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);
+    size_t get_params_num() {
+        return taesd.get_params_num();
     }
 
-    void map_by_name(std::map<std::string, ggml_tensor*>& tensors) {
-        decoder.map_by_name(tensors, "decoder.layers.");
-        encoder.map_by_name(tensors, "encoder.layers.");
-    }
-
-    bool load_from_file(const std::string& file_path, ggml_backend_t backend) {
+    bool load_from_file(const std::string& file_path) {
         LOG_INFO("loading taesd from '%s'", file_path.c_str());
-
-        if (!alloc_params_buffer(backend)) {
-            return false;
-        }
-
+        alloc_params_buffer();
         std::map<std::string, ggml_tensor*> taesd_tensors;
-
-        // prepare memory for the weights
-        {
-            init_params();
-            map_by_name(taesd_tensors);
-        }
-
-        std::map<std::string, struct ggml_tensor*> tensors_need_to_load;
+        taesd.get_param_tensors(taesd_tensors);
         std::set<std::string> ignore_tensors;
-        for (auto& pair : taesd_tensors) {
-            const std::string& name = pair.first;
-
-            if (decode_only && starts_with(name, "encoder")) {
-                ignore_tensors.insert(name);
-                continue;
-            }
-
-            tensors_need_to_load.insert(pair);
+        if (decode_only) {
+            ignore_tensors.insert("encoder.");
         }
 
         ModelLoader model_loader;
@@ -525,7 +224,7 @@ struct TinyAutoEncoder : public GGMLModule {
             return false;
         }
 
-        bool success = model_loader.load_tensors(tensors_need_to_load, backend, ignore_tensors);
+        bool success = model_loader.load_tensors(taesd_tensors, backend, ignore_tensors);
 
         if (!success) {
             LOG_ERROR("load tae tensors from model loader failed");
@@ -537,57 +236,23 @@ struct TinyAutoEncoder : public GGMLModule {
     }
 
     struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
-        // 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()
-        };
-        // LOG_DEBUG("mem_size %u ", params.mem_size);
-
-        struct ggml_context* ctx0 = ggml_init(params);
-
-        struct ggml_cgraph* gf = ggml_new_graph(ctx0);
-
-        struct ggml_tensor* z_ = NULL;
-
-        // it's performing a compute, check if backend isn't cpu
-        if (!ggml_backend_is_cpu(backend)) {
-            // pass input tensors to gpu memory
-            z_ = ggml_dup_tensor(ctx0, z);
-            ggml_allocr_alloc(compute_allocr, z_);
-
-            // pass data to device backend
-            if (!ggml_allocr_is_measure(compute_allocr)) {
-                ggml_backend_tensor_set(z_, z->data, 0, ggml_nbytes(z));
-            }
-        } else {
-            z_ = z;
-        }
-
-        struct ggml_tensor* out = decode_graph ? decoder.forward(ctx0, z_) : encoder.forward(ctx0, z_);
-
+        struct ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        z                       = to_backend(z);
+        struct ggml_tensor* out = decode_graph ? taesd.decode(compute_ctx, z) : taesd.encode(compute_ctx, z);
         ggml_build_forward_expand(gf, out);
-        ggml_free(ctx0);
-
         return gf;
     }
 
-    void alloc_compute_buffer(struct ggml_tensor* x, bool decode) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(x, decode);
-        };
-        GGMLModule::alloc_compute_buffer(get_graph);
-    }
-
-    void compute(struct ggml_tensor* work_result, int n_threads, struct ggml_tensor* z, bool decode_graph) {
+    void compute(const int n_threads,
+                 struct ggml_tensor* z,
+                 bool decode_graph,
+                 struct ggml_tensor** output,
+                 struct ggml_context* output_ctx = NULL) {
         auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
-        GGMLModule::compute(get_graph, n_threads, work_result);
+
+        GGMLModule::compute(get_graph, n_threads, false, output, output_ctx);
     }
 };
 

upscaler.cpp

@@ -6,7 +6,7 @@
 struct UpscalerGGML {
     ggml_backend_t backend    = NULL;  // general backend
     ggml_type model_data_type = GGML_TYPE_F16;
-    ESRGAN esrgan_upscaler;
+    std::shared_ptr<ESRGAN> esrgan_upscaler;
     std::string esrgan_path;
     int n_threads;
 
@@ -30,7 +30,8 @@ struct UpscalerGGML {
             backend = ggml_backend_cpu_init();
         }
         LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
-        if (!esrgan_upscaler.load_from_file(esrgan_path, backend)) {
+        esrgan_upscaler = std::make_shared<ESRGAN>(backend, model_data_type);
+        if (!esrgan_upscaler->load_from_file(esrgan_path)) {
             return false;
         }
         return true;
@@ -39,8 +40,8 @@ struct UpscalerGGML {
     sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor) {
         // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
         sd_image_t upscaled_image = {0, 0, 0, NULL};
-        int output_width          = (int)input_image.width * esrgan_upscaler.scale;
-        int output_height         = (int)input_image.height * esrgan_upscaler.scale;
+        int output_width          = (int)input_image.width * esrgan_upscaler->scale;
+        int output_height         = (int)input_image.height * esrgan_upscaler->scale;
         LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
                  input_image.width, input_image.height, output_width, output_height);
 
@@ -62,15 +63,11 @@ struct UpscalerGGML {
 
         ggml_tensor* upscaled = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, output_width, output_height, 3, 1);
         auto on_tiling        = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
-            if (init) {
-                esrgan_upscaler.alloc_compute_buffer(in);
-            } else {
-                esrgan_upscaler.compute(out, n_threads, in);
-            }
+            esrgan_upscaler->compute(n_threads, in, &out);
         };
         int64_t t0 = ggml_time_ms();
-        sd_tiling(input_image_tensor, upscaled, esrgan_upscaler.scale, esrgan_upscaler.tile_size, 0.25f, on_tiling);
-        esrgan_upscaler.free_compute_buffer();
+        sd_tiling(input_image_tensor, upscaled, esrgan_upscaler->scale, esrgan_upscaler->tile_size, 0.25f, on_tiling);
+        esrgan_upscaler->free_compute_buffer();
         ggml_tensor_clamp(upscaled, 0.f, 1.f);
         uint8_t* upscaled_data = sd_tensor_to_image(upscaled);
         ggml_free(upscale_ctx);

util.cpp

@@ -1,6 +1,7 @@
 #include "util.h"
-
 #include <stdarg.h>
+#include <algorithm>
+#include <cmath>
 #include <codecvt>
 #include <fstream>
 #include <locale>
@@ -72,6 +73,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 +101,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
@@ -170,6 +204,9 @@ std::string path_join(const std::string& p1, const std::string& p2) {
 }
 
 void pretty_progress(int step, int steps, float time) {
+    if (step == 0) {
+        return;
+    }
     std::string progress = "  |";
     int max_progress     = 50;
     int32_t current      = (int32_t)(step * 1.f * max_progress / steps);
@@ -192,6 +229,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;
 
@@ -256,3 +311,152 @@ const char* sd_get_system_info() {
 const char* sd_type_name(enum sd_type_t type) {
     return ggml_type_name((ggml_type)type);
 }
+
+sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image) {
+    sd_image_f32_t converted_image;
+    converted_image.width   = image.width;
+    converted_image.height  = image.height;
+    converted_image.channel = image.channel;
+
+    // Allocate memory for float data
+    converted_image.data = (float*)malloc(image.width * image.height * image.channel * sizeof(float));
+
+    for (int i = 0; i < image.width * image.height * image.channel; i++) {
+        // Convert uint8_t to float
+        converted_image.data[i] = (float)image.data[i];
+    }
+
+    return converted_image;
+}
+
+// Function to perform double linear interpolation
+float interpolate(float v1, float v2, float v3, float v4, float x_ratio, float y_ratio) {
+    return v1 * (1 - x_ratio) * (1 - y_ratio) + v2 * x_ratio * (1 - y_ratio) + v3 * (1 - x_ratio) * y_ratio + v4 * x_ratio * y_ratio;
+}
+
+sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height) {
+    sd_image_f32_t resized_image;
+    resized_image.width   = target_width;
+    resized_image.height  = target_height;
+    resized_image.channel = image.channel;
+
+    // Allocate memory for resized float data
+    resized_image.data = (float*)malloc(target_width * target_height * image.channel * sizeof(float));
+
+    for (int y = 0; y < target_height; y++) {
+        for (int x = 0; x < target_width; x++) {
+            float original_x = (float)x * image.width / target_width;
+            float original_y = (float)y * image.height / target_height;
+
+            int x1 = (int)original_x;
+            int y1 = (int)original_y;
+            int x2 = x1 + 1;
+            int y2 = y1 + 1;
+
+            for (int k = 0; k < image.channel; k++) {
+                float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
+                float v2 = *(image.data + y1 * image.width * image.channel + x2 * image.channel + k);
+                float v3 = *(image.data + y2 * image.width * image.channel + x1 * image.channel + k);
+                float v4 = *(image.data + y2 * image.width * image.channel + x2 * image.channel + k);
+
+                float x_ratio = original_x - x1;
+                float y_ratio = original_y - y1;
+
+                float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
+
+                *(resized_image.data + y * target_width * image.channel + x * image.channel + k) = value;
+            }
+        }
+    }
+
+    return resized_image;
+}
+
+void normalize_sd_image_f32_t(sd_image_f32_t image, float means[3], float stds[3]) {
+    for (int y = 0; y < image.height; y++) {
+        for (int x = 0; x < image.width; x++) {
+            for (int k = 0; k < image.channel; k++) {
+                int index         = (y * image.width + x) * image.channel + k;
+                image.data[index] = (image.data[index] - means[k]) / stds[k];
+            }
+        }
+    }
+}
+
+// Constants for means and std
+float means[3] = {0.48145466, 0.4578275, 0.40821073};
+float stds[3]  = {0.26862954, 0.26130258, 0.27577711};
+
+// Function to clip and preprocess sd_image_f32_t
+sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size) {
+    float scale = (float)size / fmin(image.width, image.height);
+
+    // Interpolation
+    int new_width       = (int)(scale * image.width);
+    int new_height      = (int)(scale * image.height);
+    float* resized_data = (float*)malloc(new_width * new_height * image.channel * sizeof(float));
+
+    for (int y = 0; y < new_height; y++) {
+        for (int x = 0; x < new_width; x++) {
+            float original_x = (float)x * image.width / new_width;
+            float original_y = (float)y * image.height / new_height;
+
+            int x1 = (int)original_x;
+            int y1 = (int)original_y;
+            int x2 = x1 + 1;
+            int y2 = y1 + 1;
+
+            for (int k = 0; k < image.channel; k++) {
+                float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
+                float v2 = *(image.data + y1 * image.width * image.channel + x2 * image.channel + k);
+                float v3 = *(image.data + y2 * image.width * image.channel + x1 * image.channel + k);
+                float v4 = *(image.data + y2 * image.width * image.channel + x2 * image.channel + k);
+
+                float x_ratio = original_x - x1;
+                float y_ratio = original_y - y1;
+
+                float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
+
+                *(resized_data + y * new_width * image.channel + x * image.channel + k) = value;
+            }
+        }
+    }
+
+    // Clip and preprocess
+    int h = (new_height - size) / 2;
+    int w = (new_width - size) / 2;
+
+    sd_image_f32_t result;
+    result.width   = size;
+    result.height  = size;
+    result.channel = image.channel;
+    result.data    = (float*)malloc(size * size * image.channel * sizeof(float));
+
+    for (int k = 0; k < image.channel; k++) {
+        for (int i = 0; i < size; i++) {
+            for (int j = 0; j < size; j++) {
+                *(result.data + i * size * image.channel + j * image.channel + k) =
+                    fmin(fmax(*(resized_data + (i + h) * new_width * image.channel + (j + w) * image.channel + k), 0.0f), 255.0f) / 255.0f;
+            }
+        }
+    }
+
+    // Free allocated memory
+    free(resized_data);
+
+    // Normalize
+    for (int k = 0; k < image.channel; k++) {
+        for (int i = 0; i < size; i++) {
+            for (int j = 0; j < size; j++) {
+                // *(result.data + i * size * image.channel + j * image.channel + k) = 0.5f;
+                int offset  = i * size * image.channel + j * image.channel + k;
+                float value = *(result.data + offset);
+                value       = (value - means[k]) / stds[k];
+                // value = 0.5f;
+                *(result.data + offset) = value;
+            }
+        }
+    }
+
+    return result;
+}

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);
@@ -22,12 +23,29 @@ std::u32string unicode_value_to_utf32(int unicode_value);
 
 std::string sd_basename(const std::string& path);
 
+typedef struct {
+    uint32_t width;
+    uint32_t height;
+    uint32_t channel;
+    float* data;
+} sd_image_f32_t;
+
+void normalize_sd_image_f32_t(sd_image_f32_t image, float means[3], float stds[3]);
+
+sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image);
+
+sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height);
+
+sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size);
+
 std::string path_join(const std::string& p1, const std::string& p2);
 
 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__)