feat: Control Net support + Textual Inversion (embeddings) (#131)

* add controlnet to pipeline

* add cli params

* control strength cli param

* cli param keep controlnet in cpu

* add Textual Inversion

* add canny preprocessor

* refactor: change ggml_type_sizef to ggml_row_size

* process hint once time

* ignore the embedding name case

---------

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

.clang-format

@@ -7,7 +7,7 @@ IndentCaseLabels: false
 ColumnLimit: 0
 AccessModifierOffset: -4
 NamespaceIndentation: All
-FixNamespaceComments: false 
+FixNamespaceComments: false
 AlignAfterOpenBracket: true
 AlignConsecutiveAssignments: true
 IndentCaseLabels: true

.gitignore

@@ -8,6 +8,6 @@ test/
 *.bin
 *.exe
 *.gguf
-*.log
-output.png
-models/
+output*.png
+models*
+*.log

.gitmodules

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

CMakeLists.txt

@@ -25,6 +25,8 @@ 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(BUILD_SHARED_LIBS             "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
@@ -35,21 +37,51 @@ if(SD_CUBLAS)
     add_definitions(-DSD_USE_CUBLAS)
 endif()
 
+if(SD_METAL)
+	message("Use Metal as backend stable-diffusion")
+    set(GGML_METAL ON)
+    add_definitions(-DSD_USE_METAL)
+endif()
+
+if (SD_HIPBLAS)
+    message("Use HIPBLAS as backend stable-diffusion")
+    set(GGML_HIPBLAS ON)
+    add_definitions(-DSD_USE_CUBLAS)
+    if(SD_FAST_SOFTMAX)
+        set(GGML_CUDA_FAST_SOFTMAX ON)
+    endif()
+endif ()
+
 if(SD_FLASH_ATTN)
     message("Use Flash Attention for memory optimization")
     add_definitions(-DSD_USE_FLASH_ATTENTION)
 endif()
 
+set(SD_LIB stable-diffusion)
+
+add_library(${SD_LIB} stable-diffusion.h stable-diffusion.cpp model.h model.cpp util.h util.cpp upscaler.cpp
+             ggml_extend.hpp clip.hpp common.hpp unet.hpp tae.hpp esrgan.hpp lora.hpp denoiser.hpp rng.hpp rng_philox.hpp)
+
+if(BUILD_SHARED_LIBS)
+    message("Build shared library")
+    add_definitions(-DSD_BUILD_SHARED_LIB)
+    target_compile_definitions(${SD_LIB} PRIVATE -DSD_BUILD_DLL)
+    set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+else()
+    message("Build static library")
+endif()
+
 
 set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
+
+# see https://github.com/ggerganov/ggml/pull/682
+add_definitions(-DGGML_MAX_NAME=128)
+
 # deps
 add_subdirectory(ggml)
 
 add_subdirectory(thirdparty)
 
-set(SD_LIB stable-diffusion)
-
-add_library(${SD_LIB} stable-diffusion.h stable-diffusion.cpp model.h model.cpp util.h util.cpp)
 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

@@ -9,22 +9,29 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 ## Features
 
 - Plain C/C++ implementation based on [ggml](https://github.com/ggerganov/ggml), working in the same way as [llama.cpp](https://github.com/ggerganov/llama.cpp)
-- Super lightweight and without external dependencies.
-- SD1.x and SD2.x support
+- Super lightweight and without external dependencies
+- SD1.x, SD2.x and SDXL support
+    - !!!The VAE in SDXL encounters NaN issues under FP16, but unfortunately, the ggml_conv_2d only operates under FP16. Hence, a parameter is needed to specify the VAE that has fixed the FP16 NaN issue. You can find it here: [SDXL VAE FP16 Fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors).
+
+- [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo) and [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo) support
 - 16-bit, 32-bit float support
 - 4-bit, 5-bit and 8-bit integer quantization support
 - Accelerated memory-efficient CPU inference
     - Only requires ~2.3GB when using txt2img with fp16 precision to generate a 512x512 image, enabling Flash Attention just requires ~1.8GB.
 - AVX, AVX2 and AVX512 support for x86 architectures
-- Full CUDA backend for GPU acceleration, for now just for float16 and float32 models. There are some issues with quantized models and CUDA; it will be fixed in the future.
-- Can load ckpt, safetensors and diffusers models/checkpoints. Standalone VAEs models.
+- Full CUDA and Metal backend for GPU acceleration.
+- Can load ckpt, safetensors and diffusers models/checkpoints. Standalone VAEs models
     - No need to convert to `.ggml` or `.gguf` anymore!
-- Flash Attention for memory usage optimization (only cpu for now).
+- Flash Attention for memory usage optimization (only cpu for now)
 - Original `txt2img` and `img2img` mode
 - Negative prompt
 - [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) style tokenizer (not all the features, only token weighting for now)
 - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
 - Latent Consistency Models support (LCM/LCM-LoRA)
+- 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`
@@ -48,8 +55,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - [ ] Make inference faster
     - The current implementation of ggml_conv_2d is slow and has high memory usage
 - [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
-- [ ] Implement BPE Tokenizer
-- [ ] Add [TAESD](https://github.com/madebyollin/taesd) for faster VAE decoding
+- [ ] Implement Inpainting support
 - [ ] k-quants support
 
 ## Usage
@@ -80,7 +86,7 @@ git submodule update
     ```shell
     curl -L -O https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt
     # curl -L -O https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
-    # curl -L -O https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-nonema-pruned.safetensors
+    # curl -L -O https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-nonema-pruned.safetensors
     ```
 
 ### Build
@@ -110,7 +116,27 @@ cmake .. -DSD_CUBLAS=ON
 cmake --build . --config Release
 ```
 
-### Using Flash Attention
+##### 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.
+
+```
+cmake .. -DSD_METAL=ON
+cmake --build . --config Release
+```
+
+##### 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.
 
@@ -126,19 +152,26 @@ usage: ./bin/sd [arguments]
 
 arguments:
   -h, --help                         show this help message and exit
-  -M, --mode [txt2img or img2img]    generation mode (default: txt2img)
+  -M, --mode [MODEL]                 run mode (txt2img or img2img or convert, default: txt2img)
   -t, --threads N                    number of threads to use during computation (default: -1).
                                      If threads <= 0, then threads will be set to the number of CPU physical cores
   -m, --model [MODEL]                path to model
   --vae [VAE]                        path to vae
+  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
+  --control-net [CONTROL_PATH]       path to control net model
+  --embd-dir [EMBEDDING_PATH]        path to embeddings.
+  --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.
   --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
-                                     If not specified, the default is the type of the weight file.  --lora-model-dir [DIR]             lora model directory  
+                                     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)
@@ -149,6 +182,10 @@ arguments:
   -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
   -b, --batch-count COUNT            number of images to generate.
   --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)
+  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)
+                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
+  --vae-tiling                       process vae in tiles to reduce memory usage
+  --control-net-cpu                  keep controlnet in cpu (for low vram)
   -v, --verbose                      print extra info
 ```
 
@@ -162,11 +199,22 @@ 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
 ./bin/sd -m ../models/sd-v1-4.ckpt -p "a lovely cat"
 # ./bin/sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat"
+# ./bin/sd -m ../models/sd_xl_base_1.0.safetensors --vae ../models/sdxl_vae-fp16-fix.safetensors -H 1024 -W 1024 -p "a lovely cat" -v
 ```
 
 Using formats of different precisions will yield results of varying quality.
@@ -218,6 +266,33 @@ Here's a simple example:
 | ----  |----    |
 | ![](./assets/without_lcm.png) |![](./assets/with_lcm.png)  |
 
+#### Using TAESD to faster decoding
+
+You can use TAESD to accelerate the decoding of latent images by following these steps:
+
+- Download the model [weights](https://huggingface.co/madebyollin/taesd/blob/main/diffusion_pytorch_model.safetensors).
+
+Or curl
+
+```bash
+curl -L -O https://huggingface.co/madebyollin/taesd/blob/main/diffusion_pytorch_model.safetensors
+```
+
+- Specify the model path using the `--taesd PATH` parameter. example:
+
+```bash
+sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --taesd ../models/diffusion_pytorch_model.safetensors
+```
+
+#### 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.
+
+- Specify the model path using the `--upscale-model PATH` parameter. example:
+
+```bash
+sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --upscale-model ../models/RealESRGAN_x4plus_anime_6B.pth
+```
 
 ### Docker
 
@@ -256,3 +331,4 @@ Thank you to all the people who have already contributed to stable-diffusion.cpp
 - [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
 - [k-diffusion](https://github.com/crowsonkb/k-diffusion)
 - [latent-consistency-model](https://github.com/luosiallen/latent-consistency-model)
+- [generative-models](https://github.com/Stability-AI/generative-models/)

common.hpp

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

common/zip.h

@@ -1,509 +0,0 @@
-/*
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
- * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
- * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
- * IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- */
-
-#pragma once
-#ifndef ZIP_H
-#define ZIP_H
-
-#include <stdint.h>
-#include <string.h>
-#include <sys/types.h>
-
-#ifndef ZIP_SHARED
-#define ZIP_EXPORT
-#else
-#ifdef _WIN32
-#ifdef ZIP_BUILD_SHARED
-#define ZIP_EXPORT __declspec(dllexport)
-#else
-#define ZIP_EXPORT __declspec(dllimport)
-#endif
-#else
-#define ZIP_EXPORT __attribute__((visibility("default")))
-#endif
-#endif
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-#if !defined(_POSIX_C_SOURCE) && defined(_MSC_VER)
-// 64-bit Windows is the only mainstream platform
-// where sizeof(long) != sizeof(void*)
-#ifdef _WIN64
-typedef long long ssize_t; /* byte count or error */
-#else
-typedef long ssize_t; /* byte count or error */
-#endif
-#endif
-
-/**
- * @mainpage
- *
- * Documentation for @ref zip.
- */
-
-/**
- * @addtogroup zip
- * @{
- */
-
-/**
- * Default zip compression level.
- */
-#define ZIP_DEFAULT_COMPRESSION_LEVEL 6
-
-/**
- * Error codes
- */
-#define ZIP_ENOINIT -1      // not initialized
-#define ZIP_EINVENTNAME -2  // invalid entry name
-#define ZIP_ENOENT -3       // entry not found
-#define ZIP_EINVMODE -4     // invalid zip mode
-#define ZIP_EINVLVL -5      // invalid compression level
-#define ZIP_ENOSUP64 -6     // no zip 64 support
-#define ZIP_EMEMSET -7      // memset error
-#define ZIP_EWRTENT -8      // cannot write data to entry
-#define ZIP_ETDEFLINIT -9   // cannot initialize tdefl compressor
-#define ZIP_EINVIDX -10     // invalid index
-#define ZIP_ENOHDR -11      // header not found
-#define ZIP_ETDEFLBUF -12   // cannot flush tdefl buffer
-#define ZIP_ECRTHDR -13     // cannot create entry header
-#define ZIP_EWRTHDR -14     // cannot write entry header
-#define ZIP_EWRTDIR -15     // cannot write to central dir
-#define ZIP_EOPNFILE -16    // cannot open file
-#define ZIP_EINVENTTYPE -17 // invalid entry type
-#define ZIP_EMEMNOALLOC -18 // extracting data using no memory allocation
-#define ZIP_ENOFILE -19     // file not found
-#define ZIP_ENOPERM -20     // no permission
-#define ZIP_EOOMEM -21      // out of memory
-#define ZIP_EINVZIPNAME -22 // invalid zip archive name
-#define ZIP_EMKDIR -23      // make dir error
-#define ZIP_ESYMLINK -24    // symlink error
-#define ZIP_ECLSZIP -25     // close archive error
-#define ZIP_ECAPSIZE -26    // capacity size too small
-#define ZIP_EFSEEK -27      // fseek error
-#define ZIP_EFREAD -28      // fread error
-#define ZIP_EFWRITE -29     // fwrite error
-#define ZIP_ERINIT -30      // cannot initialize reader
-#define ZIP_EWINIT -31      // cannot initialize writer
-#define ZIP_EWRINIT -32     // cannot initialize writer from reader
-
-/**
- * Looks up the error message string corresponding to an error number.
- * @param errnum error number
- * @return error message string corresponding to errnum or NULL if error is not
- * found.
- */
-extern ZIP_EXPORT const char *zip_strerror(int errnum);
-
-/**
- * @struct zip_t
- *
- * This data structure is used throughout the library to represent zip archive -
- * forward declaration.
- */
-struct zip_t;
-
-/**
- * Opens zip archive with compression level using the given mode.
- *
- * @param zipname zip archive file name.
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *zip_open(const char *zipname, int level,
-                                         char mode);
-
-/**
- * Opens zip archive with compression level using the given mode.
- * The function additionally returns @param errnum -
- *
- * @param zipname zip archive file name.
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- * @param errnum 0 on success, negative number (< 0) on error.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *
-zip_openwitherror(const char *zipname, int level, char mode, int *errnum);
-
-/**
- * Closes the zip archive, releases resources - always finalize.
- *
- * @param zip zip archive handler.
- */
-extern ZIP_EXPORT void zip_close(struct zip_t *zip);
-
-/**
- * Determines if the archive has a zip64 end of central directory headers.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - 1 (true), 0 (false), negative number (< 0) on
- *         error.
- */
-extern ZIP_EXPORT int zip_is64(struct zip_t *zip);
-
-/**
- * Opens an entry by name in the zip archive.
- *
- * For zip archive opened in 'w' or 'a' mode the function will append
- * a new entry. In readonly mode the function tries to locate the entry
- * in global dictionary.
- *
- * @param zip zip archive handler.
- * @param entryname an entry name in local dictionary.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_open(struct zip_t *zip, const char *entryname);
-
-/**
- * Opens an entry by name in the zip archive.
- *
- * For zip archive opened in 'w' or 'a' mode the function will append
- * a new entry. In readonly mode the function tries to locate the entry
- * in global dictionary (case sensitive).
- *
- * @param zip zip archive handler.
- * @param entryname an entry name in local dictionary (case sensitive).
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_opencasesensitive(struct zip_t *zip,
-                                                  const char *entryname);
-
-/**
- * Opens a new entry by index in the zip archive.
- *
- * This function is only valid if zip archive was opened in 'r' (readonly) mode.
- *
- * @param zip zip archive handler.
- * @param index index in local dictionary.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_openbyindex(struct zip_t *zip, size_t index);
-
-/**
- * Closes a zip entry, flushes buffer and releases resources.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_close(struct zip_t *zip);
-
-/**
- * Returns a local name of the current zip entry.
- *
- * The main difference between user's entry name and local entry name
- * is optional relative path.
- * Following .ZIP File Format Specification - the path stored MUST not contain
- * a drive or device letter, or a leading slash.
- * All slashes MUST be forward slashes '/' as opposed to backwards slashes '\'
- * for compatibility with Amiga and UNIX file systems etc.
- *
- * @param zip: zip archive handler.
- *
- * @return the pointer to the current zip entry name, or NULL on error.
- */
-extern ZIP_EXPORT const char *zip_entry_name(struct zip_t *zip);
-
-/**
- * Returns an index of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the index on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entry_index(struct zip_t *zip);
-
-/**
- * Determines if the current zip entry is a directory entry.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - 1 (true), 0 (false), negative number (< 0) on
- *         error.
- */
-extern ZIP_EXPORT int zip_entry_isdir(struct zip_t *zip);
-
-/**
- * Returns the uncompressed size of the current zip entry.
- * Alias for zip_entry_uncomp_size (for backward compatibility).
- *
- * @param zip zip archive handler.
- *
- * @return the uncompressed size in bytes.
- */
-extern ZIP_EXPORT unsigned long long zip_entry_size(struct zip_t *zip);
-
-/**
- * Returns the uncompressed size of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the uncompressed size in bytes.
- */
-extern ZIP_EXPORT unsigned long long zip_entry_uncomp_size(struct zip_t *zip);
-
-/**
- * Returns the compressed size of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the compressed size in bytes.
- */
-extern ZIP_EXPORT unsigned long long zip_entry_comp_size(struct zip_t *zip);
-
-/**
- * Returns CRC-32 checksum of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the CRC-32 checksum.
- */
-extern ZIP_EXPORT unsigned int zip_entry_crc32(struct zip_t *zip);
-
-/**
- * Compresses an input buffer for the current zip entry.
- *
- * @param zip zip archive handler.
- * @param buf input buffer.
- * @param bufsize input buffer size (in bytes).
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_write(struct zip_t *zip, const void *buf,
-                                      size_t bufsize);
-
-/**
- * Compresses a file for the current zip entry.
- *
- * @param zip zip archive handler.
- * @param filename input file.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_fwrite(struct zip_t *zip, const char *filename);
-
-/**
- * Extracts the current zip entry into output buffer.
- *
- * The function allocates sufficient memory for a output buffer.
- *
- * @param zip zip archive handler.
- * @param buf output buffer.
- * @param bufsize output buffer size (in bytes).
- *
- * @note remember to release memory allocated for a output buffer.
- *       for large entries, please take a look at zip_entry_extract function.
- *
- * @return the return code - the number of bytes actually read on success.
- *         Otherwise a negative number (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entry_read(struct zip_t *zip, void **buf,
-                                         size_t *bufsize);
-
-/**
- * Extracts the current zip entry into a memory buffer using no memory
- * allocation.
- *
- * @param zip zip archive handler.
- * @param buf preallocated output buffer.
- * @param bufsize output buffer size (in bytes).
- *
- * @note ensure supplied output buffer is large enough.
- *       zip_entry_size function (returns uncompressed size for the current
- *       entry) can be handy to estimate how big buffer is needed.
- *       For large entries, please take a look at zip_entry_extract function.
- *
- * @return the return code - the number of bytes actually read on success.
- *         Otherwise a negative number (< 0) on error (e.g. bufsize is not large
- * enough).
- */
-extern ZIP_EXPORT ssize_t zip_entry_noallocread(struct zip_t *zip, void *buf,
-                                                size_t bufsize);
-
-/**
- * Extracts the current zip entry into output file.
- *
- * @param zip zip archive handler.
- * @param filename output file.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_fread(struct zip_t *zip, const char *filename);
-
-/**
- * Extracts the current zip entry using a callback function (on_extract).
- *
- * @param zip zip archive handler.
- * @param on_extract callback function.
- * @param arg opaque pointer (optional argument, which you can pass to the
- *        on_extract callback)
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int
-zip_entry_extract(struct zip_t *zip,
-                  size_t (*on_extract)(void *arg, uint64_t offset,
-                                       const void *data, size_t size),
-                  void *arg);
-
-/**
- * Returns the number of all entries (files and directories) in the zip archive.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - the number of entries on success, negative number
- *         (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entries_total(struct zip_t *zip);
-
-/**
- * Deletes zip archive entries.
- *
- * @param zip zip archive handler.
- * @param entries array of zip archive entries to be deleted.
- * @param len the number of entries to be deleted.
- * @return the number of deleted entries, or negative number (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entries_delete(struct zip_t *zip,
-                                             char *const entries[], size_t len);
-
-/**
- * Extracts a zip archive stream into directory.
- *
- * If on_extract is not NULL, the callback will be called after
- * successfully extracted each zip entry.
- * Returning a negative value from the callback will cause abort and return an
- * error. The last argument (void *arg) is optional, which you can use to pass
- * data to the on_extract callback.
- *
- * @param stream zip archive stream.
- * @param size stream size.
- * @param dir output directory.
- * @param on_extract on extract callback.
- * @param arg opaque pointer.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int
-zip_stream_extract(const char *stream, size_t size, const char *dir,
-                   int (*on_extract)(const char *filename, void *arg),
-                   void *arg);
-
-/**
- * Opens zip archive stream into memory.
- *
- * @param stream zip archive stream.
- * @param size stream size.
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *zip_stream_open(const char *stream, size_t size,
-                                                int level, char mode);
-
-/**
- * Opens zip archive stream into memory.
- * The function additionally returns @param errnum -
- *
- * @param stream zip archive stream.
- * @param size stream size.*
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- * @param errnum 0 on success, negative number (< 0) on error.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *zip_stream_openwitherror(const char *stream,
-                                                         size_t size, int level,
-                                                         char mode,
-                                                         int *errnum);
-
-/**
- * Copy zip archive stream output buffer.
- *
- * @param zip zip archive handler.
- * @param buf output buffer. User should free buf.
- * @param bufsize output buffer size (in bytes).
- *
- * @return copy size
- */
-extern ZIP_EXPORT ssize_t zip_stream_copy(struct zip_t *zip, void **buf,
-                                          size_t *bufsize);
-
-/**
- * Close zip archive releases resources.
- *
- * @param zip zip archive handler.
- *
- * @return
- */
-extern ZIP_EXPORT void zip_stream_close(struct zip_t *zip);
-
-/**
- * Creates a new archive and puts files into a single zip archive.
- *
- * @param zipname zip archive file.
- * @param filenames input files.
- * @param len: number of input files.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_create(const char *zipname, const char *filenames[],
-                                 size_t len);
-
-/**
- * Extracts a zip archive file into directory.
- *
- * If on_extract_entry is not NULL, the callback will be called after
- * successfully extracted each zip entry.
- * Returning a negative value from the callback will cause abort and return an
- * error. The last argument (void *arg) is optional, which you can use to pass
- * data to the on_extract_entry callback.
- *
- * @param zipname zip archive file.
- * @param dir output directory.
- * @param on_extract_entry on extract callback.
- * @param arg opaque pointer.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_extract(const char *zipname, const char *dir,
-                                  int (*on_extract_entry)(const char *filename,
-                                                          void *arg),
-                                  void *arg);
-/** @} */
-#ifdef __cplusplus
-}
-#endif
-
-#endif

control.hpp

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

denoiser.hpp

@@ -0,0 +1,125 @@
+#ifndef __DENOISER_HPP__
+#define __DENOISER_HPP__
+
+#include "ggml_extend.hpp"
+
+/*================================================= CompVisDenoiser ==================================================*/
+
+// Ref: https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/external.py
+
+#define TIMESTEPS 1000
+
+struct SigmaSchedule {
+    float alphas_cumprod[TIMESTEPS];
+    float sigmas[TIMESTEPS];
+    float log_sigmas[TIMESTEPS];
+
+    virtual std::vector<float> get_sigmas(uint32_t n) = 0;
+
+    float sigma_to_t(float sigma) {
+        float log_sigma = std::log(sigma);
+        std::vector<float> dists;
+        dists.reserve(TIMESTEPS);
+        for (float log_sigma_val : log_sigmas) {
+            dists.push_back(log_sigma - log_sigma_val);
+        }
+
+        int low_idx = 0;
+        for (size_t i = 0; i < TIMESTEPS; i++) {
+            if (dists[i] >= 0) {
+                low_idx++;
+            }
+        }
+        low_idx      = std::min(std::max(low_idx - 1, 0), TIMESTEPS - 2);
+        int high_idx = low_idx + 1;
+
+        float low  = log_sigmas[low_idx];
+        float high = log_sigmas[high_idx];
+        float w    = (low - log_sigma) / (low - high);
+        w          = std::max(0.f, std::min(1.f, w));
+        float t    = (1.0f - w) * low_idx + w * high_idx;
+
+        return t;
+    }
+
+    float t_to_sigma(float t) {
+        int low_idx     = static_cast<int>(std::floor(t));
+        int high_idx    = static_cast<int>(std::ceil(t));
+        float w         = t - static_cast<float>(low_idx);
+        float log_sigma = (1.0f - w) * log_sigmas[low_idx] + w * log_sigmas[high_idx];
+        return std::exp(log_sigma);
+    }
+};
+
+struct DiscreteSchedule : SigmaSchedule {
+    std::vector<float> get_sigmas(uint32_t n) {
+        std::vector<float> result;
+
+        int t_max = TIMESTEPS - 1;
+
+        if (n == 0) {
+            return result;
+        } else if (n == 1) {
+            result.push_back(t_to_sigma((float)t_max));
+            result.push_back(0);
+            return result;
+        }
+
+        float step = static_cast<float>(t_max) / static_cast<float>(n - 1);
+        for (uint32_t i = 0; i < n; ++i) {
+            float t = t_max - step * i;
+            result.push_back(t_to_sigma(t));
+        }
+        result.push_back(0);
+        return result;
+    }
+};
+
+struct KarrasSchedule : SigmaSchedule {
+    std::vector<float> get_sigmas(uint32_t n) {
+        // These *COULD* be function arguments here,
+        // but does anybody ever bother to touch them?
+        float sigma_min = 0.1f;
+        float sigma_max = 10.f;
+        float rho       = 7.f;
+
+        std::vector<float> result(n + 1);
+
+        float min_inv_rho = pow(sigma_min, (1.f / rho));
+        float max_inv_rho = pow(sigma_max, (1.f / rho));
+        for (uint32_t i = 0; i < n; i++) {
+            // Eq. (5) from Karras et al 2022
+            result[i] = pow(max_inv_rho + (float)i / ((float)n - 1.f) * (min_inv_rho - max_inv_rho), rho);
+        }
+        result[n] = 0.;
+        return result;
+    }
+};
+
+struct Denoiser {
+    std::shared_ptr<SigmaSchedule> schedule              = std::make_shared<DiscreteSchedule>();
+    virtual std::vector<float> get_scalings(float sigma) = 0;
+};
+
+struct CompVisDenoiser : public Denoiser {
+    float sigma_data = 1.0f;
+
+    std::vector<float> get_scalings(float sigma) {
+        float c_out = -sigma;
+        float c_in  = 1.0f / std::sqrt(sigma * sigma + sigma_data * sigma_data);
+        return {c_out, c_in};
+    }
+};
+
+struct CompVisVDenoiser : public Denoiser {
+    float sigma_data = 1.0f;
+
+    std::vector<float> get_scalings(float sigma) {
+        float c_skip = sigma_data * sigma_data / (sigma * sigma + sigma_data * sigma_data);
+        float c_out  = -sigma * sigma_data / std::sqrt(sigma * sigma + sigma_data * sigma_data);
+        float c_in   = 1.0f / std::sqrt(sigma * sigma + sigma_data * sigma_data);
+        return {c_skip, c_out, c_in};
+    }
+};
+
+#endif  // __DENOISER_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

@@ -0,0 +1,418 @@
+#ifndef __ESRGAN_HPP__
+#define __ESRGAN_HPP__
+
+#include "ggml_extend.hpp"
+#include "model.h"
+
+/*
+    ===================================    ESRGAN  ===================================
+    References:
+    https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py
+    https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py
+
+*/
+
+struct ResidualDenseBlock {
+    int num_features;
+    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;
+    }
+
+    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;
+    }
+
+    int get_num_tensors() {
+        int num_tensors = 10;
+        return num_tensors;
+    }
+
+    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);
+    }
+
+    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;
+
+        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, float out_scale, ggml_tensor* x /* feat */) {
+        // x1 = self.lrelu(self.conv1(x))
+        ggml_tensor* x1 = ggml_nn_conv_2d(ctx, x, conv1_w, conv1_b, 1, 1, 1, 1);
+        x1              = ggml_leaky_relu(ctx, x1, 0.2f, true);
+
+        // 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);
+        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);
+        }
+    }
+
+    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();
+        }
+        return num_tensors;
+    }
+
+    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;
+    }
+
+    void init_params(ggml_context* ctx) {
+        for (int i = 0; i < num_residual_blocks; i++) {
+            rd_blocks[i].init_params(ctx);
+        }
+    }
+
+    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) + ".");
+        }
+    }
+
+    ggml_tensor* forward(ggml_context* ctx, float out_scale, ggml_tensor* x) {
+        ggml_tensor* out = x;
+        for (int i = 0; i < num_residual_blocks; i++) {
+            // out = self.rdb...(x)
+            out = rd_blocks[i].forward(ctx, out_scale, out);
+        }
+        // return out * 0.2 + x
+        out = ggml_add(ctx, ggml_scale(ctx, out, out_scale), x);
+        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
+
+    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);
+        }
+    }
+
+    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;
+    }
+
+    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;
+    }
+
+    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);
+    }
+
+    bool load_from_file(const std::string& file_path, ggml_backend_t backend) {
+        LOG_INFO("loading esrgan from '%s'", file_path.c_str());
+
+        if (!alloc_params_buffer(backend)) {
+            return false;
+        }
+
+        std::map<std::string, ggml_tensor*> esrgan_tensors;
+
+        // prepare memory for the weights
+        {
+            init_params();
+            map_by_name(esrgan_tensors);
+        }
+
+        ModelLoader model_loader;
+        if (!model_loader.init_from_file(file_path)) {
+            LOG_ERROR("init esrgan model loader from file failed: '%s'", file_path.c_str());
+            return false;
+        }
+
+        bool success = model_loader.load_tensors(esrgan_tensors, backend);
+
+        if (!success) {
+            LOG_ERROR("load esrgan tensors from model loader failed");
+            return false;
+        }
+
+        LOG_INFO("esrgan model loaded");
+        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, float out_scale, ggml_tensor* x /* feat */) {
+        // feat = self.conv_first(feat)
+        auto h = ggml_nn_conv_2d(ctx0, x, conv_first_w, conv_first_b, 1, 1, 1, 1);
+
+        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;
+        float out_scale = 0.2f;
+
+        // 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, out_scale, x);
+
+        ggml_build_forward_expand(gf, out);
+        ggml_free(ctx0);
+
+        return gf;
+    }
+
+    void alloc_compute_buffer(struct ggml_tensor* x) {
+        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);
+    }
+};
+
+#endif  // __ESRGAN_HPP__

examples/cli/main.cpp

@@ -1,9 +1,13 @@
 #include <stdio.h>
-#include <ctime>
+#include <string.h>
+#include <time.h>
+#include <iostream>
 #include <random>
-#include "ggml/ggml.h"
+#include <string>
+#include <vector>
+
 #include "stable-diffusion.h"
-#include "util.h"
+#include "preprocessing.hpp"
 
 #define STB_IMAGE_IMPLEMENTATION
 #include "stb_image.h"
@@ -12,11 +16,6 @@
 #define STB_IMAGE_WRITE_STATIC
 #include "stb_image_write.h"
 
-#include <cstring>
-#include <iostream>
-#include <string>
-#include <vector>
-
 const char* rng_type_to_str[] = {
     "std_default",
     "cuda",
@@ -44,11 +43,13 @@ const char* schedule_str[] = {
 const char* modes_str[] = {
     "txt2img",
     "img2img",
+    "convert",
 };
 
 enum SDMode {
     TXT2IMG,
     IMG2IMG,
+    CONVERT,
     MODE_COUNT
 };
 
@@ -58,25 +59,35 @@ struct SDParams {
 
     std::string model_path;
     std::string vae_path;
-    ggml_type wtype = GGML_TYPE_COUNT;
+    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 cfg_scale = 7.0f;
+    int clip_skip   = -1;  // <= 0 represents unspecified
     int width       = 512;
     int height      = 512;
     int batch_count = 1;
 
-    SampleMethod sample_method = EULER_A;
-    Schedule schedule          = DEFAULT;
-    int sample_steps           = 20;
-    float strength             = 0.75f;
-    RNGType rng_type           = CUDA_RNG;
-    int64_t seed               = 42;
-    bool verbose               = false;
+    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;
 };
 
 void print_params(SDParams params) {
@@ -84,13 +95,21 @@ void print_params(SDParams params) {
     printf("    n_threads:         %d\n", params.n_threads);
     printf("    mode:              %s\n", modes_str[params.mode]);
     printf("    model_path:        %s\n", params.model_path.c_str());
-    printf("    wtype:             %s\n", params.wtype < GGML_TYPE_COUNT ? ggml_type_name(params.wtype) : "unspecified");
+    printf("    wtype:             %s\n", params.wtype < SD_TYPE_COUNT ? sd_type_name(params.wtype) : "unspecified");
     printf("    vae_path:          %s\n", params.vae_path.c_str());
+    printf("    taesd_path:        %s\n", params.taesd_path.c_str());
+    printf("    esrgan_path:       %s\n", params.esrgan_path.c_str());
+    printf("    controlnet_path:   %s\n", params.controlnet_path.c_str());
+    printf("    embeddings_path:   %s\n", params.embeddings_path.c_str());
     printf("    output_path:       %s\n", params.output_path.c_str());
     printf("    init_img:          %s\n", params.input_path.c_str());
+    printf("    control_image:     %s\n", params.control_image_path.c_str());
+    printf("    controlnet cpu:    %s\n", params.control_net_cpu ? "true" : "false");
+    printf("    strength(control): %.2f\n", params.control_strength);
     printf("    prompt:            %s\n", params.prompt.c_str());
     printf("    negative_prompt:   %s\n", params.negative_prompt.c_str());
     printf("    cfg_scale:         %.2f\n", params.cfg_scale);
+    printf("    clip_skip:         %d\n", params.clip_skip);
     printf("    width:             %d\n", params.width);
     printf("    height:            %d\n", params.height);
     printf("    sample_method:     %s\n", sample_method_str[params.sample_method]);
@@ -100,6 +119,7 @@ void print_params(SDParams params) {
     printf("    rng:               %s\n", rng_type_to_str[params.rng_type]);
     printf("    seed:              %ld\n", params.seed);
     printf("    batch_count:       %d\n", params.batch_count);
+    printf("    vae_tiling:        %s\n", params.vae_tiling ? "true" : "false");
 }
 
 void print_usage(int argc, const char* argv[]) {
@@ -107,20 +127,26 @@ void print_usage(int argc, const char* argv[]) {
     printf("\n");
     printf("arguments:\n");
     printf("  -h, --help                         show this help message and exit\n");
-    printf("  -M, --mode [txt2img or img2img]    generation mode (default: txt2img)\n");
+    printf("  -M, --mode [MODEL]                 run mode (txt2img or img2img or convert, default: txt2img)\n");
     printf("  -t, --threads N                    number of threads to use during computation (default: -1).\n");
     printf("                                     If threads <= 0, then threads will be set to the number of CPU physical cores\n");
     printf("  -m, --model [MODEL]                path to model\n");
     printf("  --vae [VAE]                        path to vae\n");
+    printf("  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n");
+    printf("  --control-net [CONTROL_PATH]       path to control net model\n");
+    printf("  --embd-dir [EMBEDDING_PATH]        path to embeddings.\n");
+    printf("  --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n");
     printf("  --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n");
-    printf("                                     If not specified, the default is the type of the weight file.");
+    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");
@@ -131,6 +157,11 @@ void print_usage(int argc, const char* argv[]) {
     printf("  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)\n");
     printf("  -b, --batch-count COUNT            number of images to generate.\n");
     printf("  --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)\n");
+    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");
 }
 
@@ -176,26 +207,50 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.vae_path = argv[i];
-        } else if (arg == "--type") {
+        } else if (arg == "--taesd") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                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;
                 break;
             }
             std::string type = argv[i];
             if (type == "f32") {
-                params.wtype = GGML_TYPE_F32;
+                params.wtype = SD_TYPE_F32;
             } else if (type == "f16") {
-                params.wtype = GGML_TYPE_F16;
+                params.wtype = SD_TYPE_F16;
             } else if (type == "q4_0") {
-                params.wtype = GGML_TYPE_Q4_0;
+                params.wtype = SD_TYPE_Q4_0;
             } else if (type == "q4_1") {
-                params.wtype = GGML_TYPE_Q4_1;
+                params.wtype = SD_TYPE_Q4_1;
             } else if (type == "q5_0") {
-                params.wtype = GGML_TYPE_Q5_0;
+                params.wtype = SD_TYPE_Q5_0;
             } else if (type == "q5_1") {
-                params.wtype = GGML_TYPE_Q5_1;
+                params.wtype = SD_TYPE_Q5_1;
             } else if (type == "q8_0") {
-                params.wtype = GGML_TYPE_Q8_0;
+                params.wtype = SD_TYPE_Q8_0;
             } else {
                 fprintf(stderr, "error: invalid weight format %s, must be one of [f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0]\n",
                         type.c_str());
@@ -213,6 +268,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.input_path = argv[i];
+        } else if (arg == "--control-image") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.control_image_path = argv[i];
         } else if (arg == "-o" || arg == "--output") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -243,6 +304,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.strength = std::stof(argv[i]);
+        } else if (arg == "--control-strength") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.control_strength = std::stof(argv[i]);
         } else if (arg == "-H" || arg == "--height") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -261,6 +328,18 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.sample_steps = std::stoi(argv[i]);
+        } else if (arg == "--clip-skip") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            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;
@@ -297,7 +376,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 invalid_arg = true;
                 break;
             }
-            params.schedule = (Schedule)schedule_found;
+            params.schedule = (schedule_t)schedule_found;
         } else if (arg == "-s" || arg == "--seed") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -320,7 +399,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 invalid_arg = true;
                 break;
             }
-            params.sample_method = (SampleMethod)sample_method_found;
+            params.sample_method = (sample_method_t)sample_method_found;
         } else if (arg == "-h" || arg == "--help") {
             print_usage(argc, argv);
             exit(0);
@@ -341,7 +420,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         params.n_threads = get_num_physical_cores();
     }
 
-    if (params.prompt.length() == 0) {
+    if (params.mode != CONVERT && params.prompt.length() == 0) {
         fprintf(stderr, "error: the following arguments are required: prompt\n");
         print_usage(argc, argv);
         exit(1);
@@ -389,6 +468,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         srand((int)time(NULL));
         params.seed = rand();
     }
+
+    if (params.mode == CONVERT) {
+        if (params.output_path == "output.png") {
+            params.output_path = "output.gguf";
+        }
+    }
 }
 
 std::string get_image_params(SDParams params, int64_t seed) {
@@ -400,7 +485,7 @@ std::string get_image_params(SDParams params, int64_t seed) {
     parameter_string += "CFG scale: " + std::to_string(params.cfg_scale) + ", ";
     parameter_string += "Seed: " + std::to_string(seed) + ", ";
     parameter_string += "Size: " + std::to_string(params.width) + "x" + std::to_string(params.height) + ", ";
-    parameter_string += "Model: " + basename(params.model_path) + ", ";
+    parameter_string += "Model: " + sd_basename(params.model_path) + ", ";
     parameter_string += "RNG: " + std::string(rng_type_to_str[params.rng_type]) + ", ";
     parameter_string += "Sampler: " + std::string(sample_method_str[params.sample_method]);
     if (params.schedule == KARRAS) {
@@ -411,14 +496,47 @@ std::string get_image_params(SDParams params, int64_t seed) {
     return parameter_string;
 }
 
+void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
+    SDParams* params = (SDParams*)data;
+    if (!params->verbose && level <= SD_LOG_DEBUG) {
+        return;
+    }
+    if (level <= SD_LOG_INFO) {
+        fputs(log, stdout);
+        fflush(stdout);
+    } else {
+        fputs(log, stderr);
+        fflush(stderr);
+    }
+}
+
 int main(int argc, const char* argv[]) {
     SDParams params;
     parse_args(argc, argv, params);
 
+    sd_set_log_callback(sd_log_cb, (void*)&params);
+
     if (params.verbose) {
         print_params(params);
-        printf("%s", sd_get_system_info().c_str());
-        set_sd_log_level(SDLogLevel::DEBUG);
+        printf("%s", sd_get_system_info());
+    }
+
+    if (params.mode == CONVERT) {
+        bool success = convert(params.model_path.c_str(), params.vae_path.c_str(), params.output_path.c_str(), params.wtype);
+        if (!success) {
+            fprintf(stderr,
+                    "convert '%s'/'%s' to '%s' failed\n",
+                    params.model_path.c_str(),
+                    params.vae_path.c_str(),
+                    params.output_path.c_str());
+            return 1;
+        } else {
+            printf("convert '%s'/'%s' to '%s' success\n",
+                   params.model_path.c_str(),
+                   params.vae_path.c_str(),
+                   params.output_path.c_str());
+            return 0;
+        }
     }
 
     bool vae_decode_only        = true;
@@ -449,47 +567,124 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    StableDiffusion sd(params.n_threads, vae_decode_only, true, params.lora_model_dir, params.rng_type);
-    if (!sd.load_from_file(params.model_path, params.vae_path, params.wtype, params.schedule)) {
+    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.control_net_cpu);
+
+    if (sd_ctx == NULL) {
+        printf("new_sd_ctx_t failed\n");
         return 1;
     }
 
-    std::vector<uint8_t*> results;
+    sd_image_t* results;
     if (params.mode == TXT2IMG) {
-        results = sd.txt2img(params.prompt,
-                             params.negative_prompt,
-                             params.cfg_scale,
-                             params.width,
-                             params.height,
-                             params.sample_method,
-                             params.sample_steps,
-                             params.seed,
-                             params.batch_count);
+        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(),
+                          params.clip_skip,
+                          params.cfg_scale,
+                          params.width,
+                          params.height,
+                          params.sample_method,
+                          params.sample_steps,
+                          params.seed,
+                          params.batch_count,
+                          control_image,
+                          params.control_strength);
     } else {
-        results = sd.img2img(input_image_buffer,
-                             params.prompt,
-                             params.negative_prompt,
-                             params.cfg_scale,
-                             params.width,
-                             params.height,
-                             params.sample_method,
-                             params.sample_steps,
-                             params.strength,
-                             params.seed);
+        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 (results.size() == 0 || results.size() != params.batch_count) {
-        LOG_ERROR("generate failed");
+    if (results == NULL) {
+        printf("generate failed\n");
+        free_sd_ctx(sd_ctx);
         return 1;
     }
 
+    int upscale_factor = 4;  // unused for RealESRGAN_x4plus_anime_6B.pth
+    if (params.esrgan_path.size() > 0) {
+        upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(params.esrgan_path.c_str(),
+                                                        params.n_threads,
+                                                        params.wtype);
+
+        if (upscaler_ctx == NULL) {
+            printf("new_upscaler_ctx failed\n");
+        } else {
+            for (int i = 0; i < params.batch_count; i++) {
+                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;
+                }
+                free(results[i].data);
+                results[i] = upscaled_image;
+            }
+        }
+    }
+
     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.batch_count; 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(), params.width, params.height, 3, results[i], 0, get_image_params(params, params.seed + i).c_str());
-        LOG_INFO("save result image to '%s'", final_image_path.c_str());
+        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;
 }

format-code.sh

@@ -1,2 +1,2 @@
-clang-format -style=file -i *.cpp *.h
+clang-format -style=file -i *.cpp *.h *.hpp
 clang-format -style=file -i examples/cli/*.cpp

ggml_extend.hpp

@@ -0,0 +1,646 @@
+#ifndef __GGML_EXTEND_HPP__
+#define __GGML_EXTEND_HPP__
+
+#include <assert.h>
+#include <inttypes.h>
+#include <stdarg.h>
+#include <algorithm>
+#include <cstring>
+#include <fstream>
+#include <functional>
+#include <iostream>
+#include <iterator>
+#include <map>
+#include <random>
+#include <regex>
+#include <set>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "ggml/ggml-alloc.h"
+#include "ggml/ggml-backend.h"
+#include "ggml/ggml.h"
+
+#ifdef SD_USE_CUBLAS
+#include "ggml-cuda.h"
+#endif
+
+#ifdef SD_USE_METAL
+#include "ggml-metal.h"
+#endif
+
+#include "rng.hpp"
+#include "util.h"
+
+#define EPS 1e-05f
+
+#ifndef __STATIC_INLINE__
+#define __STATIC_INLINE__ static inline
+#endif
+
+__STATIC_INLINE__ void ggml_log_callback_default(ggml_log_level level, const char* text, void* user_data) {
+    (void)level;
+    (void)user_data;
+    fputs(text, stderr);
+    fflush(stderr);
+}
+
+__STATIC_INLINE__ void ggml_tensor_set_f32_randn(struct ggml_tensor* tensor, std::shared_ptr<RNG> rng) {
+    uint32_t n                        = (uint32_t)ggml_nelements(tensor);
+    std::vector<float> random_numbers = rng->randn(n);
+    for (uint32_t i = 0; i < n; i++) {
+        ggml_set_f32_1d(tensor, i, random_numbers[i]);
+    }
+}
+
+// set tensor[i, j, k, l]
+// set tensor[l]
+// set tensor[k, l]
+// set tensor[j, k, l]
+__STATIC_INLINE__ void ggml_tensor_set_f32(struct ggml_tensor* tensor, float value, int l, int k = 0, int j = 0, int i = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    *(float*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]) = value;
+}
+
+__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;
+    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]);
+}
+
+__STATIC_INLINE__ ggml_fp16_t ggml_tensor_get_f16(const ggml_tensor* tensor, int l, int k = 0, int j = 0, int i = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t));
+    return *(ggml_fp16_t*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
+    printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    int range = 3;
+    for (int i = 0; i < tensor->ne[3]; i++) {
+        if (i >= range && i + range < tensor->ne[3]) {
+            continue;
+        }
+        for (int j = 0; j < tensor->ne[2]; j++) {
+            if (j >= range && j + range < tensor->ne[2]) {
+                continue;
+            }
+            for (int k = 0; k < tensor->ne[1]; k++) {
+                if (k >= range && k + range < tensor->ne[1]) {
+                    continue;
+                }
+                for (int l = 0; l < tensor->ne[0]; l++) {
+                    if (l >= range && l + range < tensor->ne[0]) {
+                        continue;
+                    }
+                    if (tensor->type == GGML_TYPE_F32) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i, j, k, l, ggml_tensor_get_f32(tensor, l, k, j, i));
+                    } else if (tensor->type == GGML_TYPE_F16) {
+                        printf("  [%d, %d, %d, %d] = %i\n", i, j, k, l, ggml_tensor_get_f16(tensor, l, k, j, i));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+__STATIC_INLINE__ ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        LOG_ERROR("failed to open '%s'", file_path.c_str());
+        return NULL;
+    }
+    int32_t n_dims;
+    int32_t length;
+    int32_t ttype;
+
+    file.read(reinterpret_cast<char*>(&n_dims), sizeof(n_dims));
+    file.read(reinterpret_cast<char*>(&length), sizeof(length));
+    file.read(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+
+    if (file.eof()) {
+        LOG_ERROR("incomplete file '%s'", file_path.c_str());
+        return NULL;
+    }
+
+    int32_t nelements = 1;
+    int32_t ne[4]     = {1, 1, 1, 1};
+    for (int i = 0; i < n_dims; ++i) {
+        file.read(reinterpret_cast<char*>(&ne[i]), sizeof(ne[i]));
+        nelements *= ne[i];
+    }
+    std::string name(length, 0);
+    file.read(&name[0], length);
+    ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]);
+    const size_t bpe    = ggml_type_size(ggml_type(ttype));
+    file.read(reinterpret_cast<char*>(tensor->data), ggml_nbytes(tensor));
+    return tensor;
+}
+
+// __STATIC_INLINE__ void save_tensor_to_file(const std::string& file_name, ggml_tensor* tensor, const std::string & name) {
+//     std::string file_name_ = file_name + ".tensor";
+//     std::string name_ = name;
+//     std::ofstream file("./" + file_name_, std::ios::binary);
+//     file.write(reinterpret_cast<char*>(&tensor->n_dims), sizeof(tensor->n_dims));
+//     int len = (int)name_.size();
+//     file.write(reinterpret_cast<char*>(&len), sizeof(len));
+//     int ttype = (int)tensor->type;
+//     file.write(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+//     for (int i = 0; i < tensor->n_dims; ++i) {
+//         int ne_ = (int) tensor->ne[i];
+//         file.write(reinterpret_cast<char*>(&ne_), sizeof(ne_));
+//     }
+//     file.write(&name_[0], len);
+//     char* data = nullptr;
+//     file.write((char*)tensor->data, ggml_nbytes(tensor));
+//     file.close();
+// }
+
+__STATIC_INLINE__ void copy_ggml_tensor(struct ggml_tensor* dst, struct ggml_tensor* src) {
+    if (dst->type == src->type) {
+        dst->nb[0] = src->nb[0];
+        dst->nb[1] = src->nb[1];
+        dst->nb[2] = src->nb[2];
+        dst->nb[3] = src->nb[3];
+
+        memcpy(((char*)dst->data), ((char*)src->data), ggml_nbytes(dst));
+        return;
+    }
+    struct ggml_init_params params;
+    params.mem_size          = 10 * 1024 * 1024;  // for padding
+    params.mem_buffer        = NULL;
+    params.no_alloc          = false;
+    struct ggml_context* ctx = ggml_init(params);
+    if (!ctx) {
+        LOG_ERROR("ggml_init() failed");
+        return;
+    }
+    ggml_tensor* final = ggml_cpy_inplace(ctx, src, dst);
+
+    struct ggml_cgraph* graph = ggml_new_graph(ctx);
+    ggml_build_forward_expand(graph, final);
+    ggml_graph_compute_with_ctx(ctx, graph, 1);
+    ggml_free(ctx);
+}
+
+// SPECIAL OPERATIONS WITH TENSORS
+
+__STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input) {
+    int64_t width    = input->ne[0];
+    int64_t height   = input->ne[1];
+    int64_t channels = input->ne[2];
+    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    uint8_t* image_data = (uint8_t*)malloc(width * height * channels);
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            for (int k = 0; k < channels; k++) {
+                float value                                               = ggml_tensor_get_f32(input, ix, iy, k);
+                *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+__STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
+                                          struct ggml_tensor* output) {
+    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++) {
+                int value = *(image_data + iy * width * channels + ix * channels + k);
+                ggml_tensor_set_f32(output, value / 255.0f, ix, iy, k);
+            }
+        }
+    }
+}
+
+__STATIC_INLINE__ void ggml_split_tensor_2d(struct ggml_tensor* input,
+                                            struct ggml_tensor* output,
+                                            int x,
+                                            int y) {
+    int64_t width    = output->ne[0];
+    int64_t height   = output->ne[1];
+    int64_t channels = output->ne[2];
+    GGML_ASSERT(input->type == GGML_TYPE_F32 && 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 = ggml_tensor_get_f32(input, ix + x, iy + y, k);
+                ggml_tensor_set_f32(output, value, ix, iy, k);
+            }
+        }
+    }
+}
+
+__STATIC_INLINE__ void ggml_merge_tensor_2d(struct ggml_tensor* input,
+                                            struct ggml_tensor* output,
+                                            int x,
+                                            int y,
+                                            int overlap) {
+    int64_t width    = input->ne[0];
+    int64_t height   = input->ne[1];
+    int64_t channels = input->ne[2];
+    GGML_ASSERT(input->type == GGML_TYPE_F32 && 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 new_value = ggml_tensor_get_f32(input, ix, iy, k);
+                if (overlap > 0) {  // blend colors in overlapped area
+                    float old_value = ggml_tensor_get_f32(output, x + ix, y + iy, k);
+                    if (x > 0 && ix < overlap) {  // in overlapped horizontal
+                        ggml_tensor_set_f32(output, old_value + (new_value - old_value) * (ix / (1.0f * overlap)), x + ix, y + iy, k);
+                        continue;
+                    }
+                    if (y > 0 && iy < overlap) {  // in overlapped vertical
+                        ggml_tensor_set_f32(output, old_value + (new_value - old_value) * (iy / (1.0f * overlap)), x + ix, y + iy, k);
+                        continue;
+                    }
+                }
+                ggml_tensor_set_f32(output, new_value, x + ix, y + iy, k);
+            }
+        }
+    }
+}
+
+__STATIC_INLINE__ float ggml_tensor_mean(struct ggml_tensor* src) {
+    float mean        = 0.0f;
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        mean += data[i] / nelements * 1.0f;
+    }
+    return mean;
+}
+
+// a = a+b
+__STATIC_INLINE__ void ggml_tensor_add(struct ggml_tensor* a, struct ggml_tensor* b) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    int64_t nelements = ggml_nelements(a);
+    float* vec_a      = (float*)a->data;
+    float* vec_b      = (float*)b->data;
+    for (int i = 0; i < nelements; i++) {
+        vec_a[i] = vec_a[i] + vec_b[i];
+    }
+}
+
+__STATIC_INLINE__ void ggml_tensor_scale(struct ggml_tensor* src, float scale) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        data[i] = data[i] * scale;
+    }
+}
+
+__STATIC_INLINE__ void ggml_tensor_clamp(struct ggml_tensor* src, float min, float max) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val < min ? min : (val > max ? max : val);
+    }
+}
+
+// convert values from [0, 1] to [-1, 1]
+__STATIC_INLINE__ void ggml_tensor_scale_input(struct ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val * 2.0f - 1.0f;
+    }
+}
+
+// convert values from [-1, 1] to [0, 1]
+__STATIC_INLINE__ void ggml_tensor_scale_output(struct ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = (val + 1.0f) * 0.5f;
+    }
+}
+
+typedef std::function<void(ggml_tensor*, ggml_tensor*, bool)> on_tile_process;
+
+// Tiling
+__STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const int scale, const int tile_size, const float tile_overlap_factor, on_tile_process on_processing) {
+    int input_width   = (int)input->ne[0];
+    int input_height  = (int)input->ne[1];
+    int output_width  = (int)output->ne[0];
+    int output_height = (int)output->ne[1];
+    GGML_ASSERT(input_width % 2 == 0 && input_height % 2 == 0 && output_width % 2 == 0 && output_height % 2 == 0);  // should be multiple of 2
+
+    int tile_overlap     = (int32_t)(tile_size * tile_overlap_factor);
+    int non_tile_overlap = tile_size - tile_overlap;
+
+    struct ggml_init_params params = {};
+    params.mem_size += tile_size * tile_size * input->ne[2] * sizeof(float);                       // input chunk
+    params.mem_size += (tile_size * scale) * (tile_size * scale) * output->ne[2] * sizeof(float);  // output chunk
+    params.mem_size += 3 * ggml_tensor_overhead();
+    params.mem_buffer = NULL;
+    params.no_alloc   = false;
+
+    LOG_DEBUG("tile work buffer size: %.2f MB", params.mem_size / 1024.f / 1024.f);
+
+    // draft context
+    struct ggml_context* tiles_ctx = ggml_init(params);
+    if (!tiles_ctx) {
+        LOG_ERROR("ggml_init() failed");
+        return;
+    }
+
+    // tiling
+    ggml_tensor* input_tile  = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size, tile_size, input->ne[2], 1);
+    ggml_tensor* output_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size * scale, tile_size * scale, output->ne[2], 1);
+    on_processing(input_tile, NULL, true);
+    int num_tiles = (input_width * input_height) / (non_tile_overlap * non_tile_overlap);
+    LOG_INFO("processing %i tiles", num_tiles);
+    pretty_progress(1, num_tiles, 0.0f);
+    int tile_count = 1;
+    bool last_y = false, last_x = false;
+    float last_time = 0.0f;
+    for (int y = 0; y < input_height && !last_y; y += non_tile_overlap) {
+        if (y + tile_size >= input_height) {
+            y      = input_height - tile_size;
+            last_y = true;
+        }
+        for (int x = 0; x < input_width && !last_x; x += non_tile_overlap) {
+            if (x + tile_size >= input_width) {
+                x      = input_width - tile_size;
+                last_x = true;
+            }
+            int64_t t1 = ggml_time_ms();
+            ggml_split_tensor_2d(input, input_tile, x, y);
+            on_processing(input_tile, output_tile, false);
+            ggml_merge_tensor_2d(output_tile, output, x * scale, y * scale, tile_overlap * scale);
+            int64_t t2 = ggml_time_ms();
+            last_time  = (t2 - t1) / 1000.0f;
+            pretty_progress(tile_count, num_tiles, last_time);
+            tile_count++;
+        }
+        last_x = false;
+    }
+    if (tile_count < num_tiles) {
+        pretty_progress(num_tiles, num_tiles, last_time);
+    }
+}
+
+__STATIC_INLINE__ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx,
+                                                         struct ggml_tensor* a) {
+    return ggml_group_norm(ctx, a, 32);
+}
+
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_linear(struct ggml_context* ctx,
+                                                     struct ggml_tensor* x,
+                                                     struct ggml_tensor* w,
+                                                     struct ggml_tensor* b) {
+    x = ggml_mul_mat(ctx, w, x);
+    x = ggml_add(ctx, x, b);
+    return x;
+}
+
+// w: [OC，IC, KH, KW]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N, OC, OH, OW]
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_conv_2d(struct ggml_context* ctx,
+                                                      struct ggml_tensor* x,
+                                                      struct ggml_tensor* w,
+                                                      struct ggml_tensor* b,
+                                                      int s0 = 1,
+                                                      int s1 = 1,
+                                                      int p0 = 0,
+                                                      int p1 = 0,
+                                                      int d0 = 1,
+                                                      int d1 = 1) {
+    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);
+        x = ggml_add(ctx, x, b);
+    }
+    return x;
+}
+
+__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);
+    return x;
+}
+
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_group_norm(struct ggml_context* ctx,
+                                                         struct ggml_tensor* x,
+                                                         struct ggml_tensor* w,
+                                                         struct ggml_tensor* b,
+                                                         int num_groups = 32) {
+    if (ggml_n_dims(x) >= 3) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+    }
+
+    x = ggml_group_norm(ctx, x, num_groups);
+    x = ggml_mul(ctx, x, w);
+    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
+    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
+}
+
+__STATIC_INLINE__ float ggml_backend_tensor_get_f32(ggml_tensor* tensor) {
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16);
+    float value;
+    if (tensor->type == GGML_TYPE_F32) {
+        ggml_backend_tensor_get(tensor, &value, 0, sizeof(value));
+    } else {  // GGML_TYPE_F16
+        ggml_fp16_t f16_value;
+        ggml_backend_tensor_get(tensor, &f16_value, 0, sizeof(f16_value));
+        value = ggml_fp16_to_fp32(f16_value);
+    }
+    return value;
+}
+
+// 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) {
+    // timesteps: [N,]
+    // embedding: [dim, N]
+    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 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);
+        }
+        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;
+}
+
+struct GGMLModule {
+    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;
+
+    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 = 4 * 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);
+
+        struct ggml_init_params params;
+        params.mem_size   = static_cast<size_t>(num_tensors * ggml_tensor_overhead()) + 1 * 1024 * 1024;
+        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;
+    }
+
+    void free_params_buffer() {
+        if (params_ctx != NULL) {
+            ggml_free(params_ctx);
+            params_ctx = NULL;
+        }
+
+        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);
+        }
+
+        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);
+
+        struct ggml_cgraph* gf = get_graph();
+
+        ggml_allocr_alloc_graph(compute_allocr, gf);
+
+        if (ggml_backend_is_cpu(backend)) {
+            ggml_backend_cpu_set_n_threads(backend, n_threads);
+        }
+
+#ifdef SD_USE_METAL
+        if (ggml_backend_is_metal(backend)) {
+            ggml_backend_metal_set_n_cb(backend, n_threads);
+        }
+#endif
+
+        ggml_backend_graph_compute(backend, gf);
+
+#ifdef GGML_PERF
+        ggml_graph_print(gf);
+#endif
+
+        if (output != NULL) {
+            ggml_backend_tensor_get_and_sync(backend, gf->nodes[gf->n_nodes - 1], output->data, 0, ggml_nbytes(output));
+        }
+    }
+
+    void free_compute_buffer() {
+        ggml_allocr_free(compute_allocr);
+        ggml_backend_buffer_free(compute_buffer);
+        compute_allocr      = NULL;
+        compute_buffer_size = 0;
+    }
+};
+
+#endif  // __GGML_EXTEND__HPP__

lora.hpp

@@ -0,0 +1,178 @@
+#ifndef __LORA_HPP__
+#define __LORA_HPP__
+
+#include "ggml_extend.hpp"
+
+#define LORA_GRAPH_SIZE 10240
+
+struct LoraModel : public GGMLModule {
+    float multiplier = 1.0f;
+    std::map<std::string, struct ggml_tensor*> lora_tensors;
+    std::string file_path;
+    ModelLoader model_loader;
+    bool load_failed = false;
+
+    LoraModel(const std::string file_path = "")
+        : file_path(file_path) {
+        name = "lora";
+        if (!model_loader.init_from_file(file_path)) {
+            load_failed = true;
+        }
+    }
+
+    size_t get_num_tensors() {
+        return LORA_GRAPH_SIZE;
+    }
+
+    size_t calculate_mem_size() {
+        return model_loader.cal_mem_size(NULL);
+    }
+
+    bool load_from_file(ggml_backend_t backend) {
+        if (!alloc_params_buffer(backend)) {
+            return false;
+        }
+        LOG_INFO("loading LoRA from '%s'", file_path.c_str());
+
+        if (load_failed) {
+            LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
+            return false;
+        }
+
+        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
+
+        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;
+            return true;
+        };
+
+        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);
+
+        std::set<std::string> applied_lora_tensors;
+        for (auto it : model_tensors) {
+            std::string k_tensor       = it.first;
+            struct ggml_tensor* weight = model_tensors[it.first];
+
+            size_t k_pos = k_tensor.find(".weight");
+            if (k_pos == std::string::npos) {
+                continue;
+            }
+            k_tensor = k_tensor.substr(0, k_pos);
+            replace_all_chars(k_tensor, '.', '_');
+            std::string lora_up_name   = "lora." + k_tensor + ".lora_up.weight";
+            std::string lora_down_name = "lora." + k_tensor + ".lora_down.weight";
+            std::string alpha_name     = "lora." + k_tensor + ".alpha";
+            std::string scale_name     = "lora." + k_tensor + ".scale";
+
+            ggml_tensor* lora_up   = NULL;
+            ggml_tensor* lora_down = NULL;
+
+            if (lora_tensors.find(lora_up_name) != lora_tensors.end()) {
+                lora_up = lora_tensors[lora_up_name];
+            }
+
+            if (lora_tensors.find(lora_down_name) != lora_tensors.end()) {
+                lora_down = lora_tensors[lora_down_name];
+            }
+
+            if (lora_up == NULL || lora_down == NULL) {
+                continue;
+            }
+
+            applied_lora_tensors.insert(lora_up_name);
+            applied_lora_tensors.insert(lora_down_name);
+            applied_lora_tensors.insert(alpha_name);
+            applied_lora_tensors.insert(scale_name);
+
+            // calc_cale
+            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]);
+            } else if (lora_tensors.find(alpha_name) != lora_tensors.end()) {
+                float alpha = ggml_backend_tensor_get_f32(lora_tensors[alpha_name]);
+                scale_value = alpha / dim;
+            }
+            scale_value *= multiplier;
+
+            // flat lora tensors to multiply it
+            int64_t lora_up_rows   = lora_up->ne[ggml_n_dims(lora_up) - 1];
+            lora_up                = ggml_reshape_2d(ctx0, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
+            int64_t lora_down_rows = lora_down->ne[ggml_n_dims(lora_down) - 1];
+            lora_down              = ggml_reshape_2d(ctx0, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
+
+            // ggml_mul_mat requires tensor b transposed
+            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);
+            GGML_ASSERT(ggml_nelements(updown) == ggml_nelements(weight));
+            updown = ggml_scale_inplace(ctx0, updown, scale_value);
+            ggml_tensor* final_weight;
+            // if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) {
+            //     final_weight = ggml_new_tensor(ctx0, GGML_TYPE_F32, weight->n_dims, weight->ne);
+            //     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);
+            // } else {
+            //     final_weight = ggml_add_inplace(ctx0, weight, updown);
+            // }
+            final_weight = ggml_add_inplace(ctx0, weight, updown);  // apply directly
+            ggml_build_forward_expand(gf, final_weight);
+        }
+
+        for (auto& kv : lora_tensors) {
+            if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
+                LOG_WARN("unused lora tensor %s", kv.first.c_str());
+            }
+        }
+
+        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);
+    }
+};
+
+#endif  // __LORA_HPP__

model.cpp

@@ -1,6 +1,7 @@
 #include <stdarg.h>
 #include <fstream>
 #include <regex>
+#include <set>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -14,6 +15,14 @@
 #include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
 
+#include "stable-diffusion.h"
+
+#ifdef SD_USE_METAL
+#include "ggml-metal.h"
+#endif
+
+#define ST_HEADER_SIZE_LEN 8
+
 uint64_t read_u64(uint8_t* buffer) {
     // little endian
     uint64_t value = 0;
@@ -52,7 +61,6 @@ std::string self_attn_names[] = {
     "self_attn.q_proj.weight",
     "self_attn.k_proj.weight",
     "self_attn.v_proj.weight",
-
     "self_attn.q_proj.bias",
     "self_attn.k_proj.bias",
     "self_attn.v_proj.bias",
@@ -73,13 +81,16 @@ const char* unused_tensors[] = {
     "cond_stage_model.transformer.text_model.embeddings.position_ids",
     "cond_stage_model.model.logit_scale",
     "cond_stage_model.model.text_projection",
+    "conditioner.embedders.0.transformer.text_model.embeddings.position_ids",
+    "conditioner.embedders.0.model.logit_scale",
+    "conditioner.embedders.1.model.logit_scale",
     "model.diffusion_model.time_embedding.cond_proj.weight",
     "unet.time_embedding.cond_proj.weight",
     "model_ema.decay",
     "model_ema.num_updates",
     "model_ema.diffusion_model",
-    "control_model",
     "embedding_manager",
+    "denoiser.sigmas",
 };
 
 bool is_unused_tensor(std::string name) {
@@ -92,11 +103,11 @@ bool is_unused_tensor(std::string name) {
 }
 
 std::unordered_map<std::string, std::string> open_clip_to_hf_clip_model = {
-    {"cond_stage_model.model.ln_final.bias", "cond_stage_model.transformer.text_model.final_layer_norm.bias"},
-    {"cond_stage_model.model.ln_final.weight", "cond_stage_model.transformer.text_model.final_layer_norm.weight"},
-    {"cond_stage_model.model.positional_embedding", "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight"},
-    {"cond_stage_model.model.token_embedding.weight", "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight"},
-
+    {"model.ln_final.bias", "transformer.text_model.final_layer_norm.bias"},
+    {"model.ln_final.weight", "transformer.text_model.final_layer_norm.weight"},
+    {"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"},
 };
 
 std::unordered_map<std::string, std::string> open_clip_to_hk_clip_resblock = {
@@ -124,16 +135,29 @@ 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 open_clip_resblock_prefix = "cond_stage_model.model.transformer.resblocks.";
-    std::string hf_clip_resblock_prefix   = "cond_stage_model.transformer.text_model.encoder.layers.";
+    std::string new_name = name;
+    std::string prefix;
+    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.")) {
+        prefix   = "cond_stage_model.1.";
+        new_name = new_name.substr(strlen("conditioner.embedders.0."));
+    } else if (starts_with(new_name, "cond_stage_model.")) {
+        prefix   = "cond_stage_model.";
+        new_name = new_name.substr(strlen("cond_stage_model."));
+    } 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(name) != open_clip_to_hf_clip_model.end()) {
-        new_name = open_clip_to_hf_clip_model[name];
+    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 (name.find(open_clip_resblock_prefix) == 0) {
-        std::string remain = name.substr(open_clip_resblock_prefix.length());
+    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);
 
@@ -145,7 +169,7 @@ std::string convert_open_clip_to_hf_clip(const std::string& name) {
         }
     }
 
-    return new_name;
+    return prefix + new_name;
 }
 
 std::string convert_vae_decoder_name(const std::string& name) {
@@ -347,10 +371,15 @@ std::string convert_diffusers_name_to_compvis(const std::string& key, char seq)
 
 std::string convert_tensor_name(const std::string& name) {
     std::string new_name;
-    if (starts_with(name, "cond_stage_model.model")) {
+    if (starts_with(name, "cond_stage_model.") || starts_with(name, "conditioner.embedders.")) {
         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) {
@@ -408,7 +437,7 @@ void preprocess_tensor(TensorStorage tensor_storage,
 
     tensor_storage.name = new_name;
 
-    if (starts_with(new_name, "cond_stage_model.transformer.text_model.encoder.layers.") &&
+    if (new_name.find("transformer.text_model.encoder.layers.") != 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);
@@ -420,7 +449,7 @@ void preprocess_tensor(TensorStorage tensor_storage,
 
         processed_tensor_storages.insert(processed_tensor_storages.end(), chunks.begin(), chunks.end());
 
-    } else if (starts_with(new_name, "cond_stage_model.transformer.text_model.encoder.layers.") &&
+    } else if (new_name.find("transformer.text_model.encoder.layers.") != 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);
@@ -533,233 +562,102 @@ std::map<char, int> unicode_to_byte() {
     return byte_decoder;
 }
 
+bool is_zip_file(const std::string& file_path) {
+    struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    if (zip == NULL) {
+        return false;
+    }
+    zip_close(zip);
+    return true;
+}
+
+bool is_gguf_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    char magic[4];
+
+    file.read(magic, sizeof(magic));
+    if (!file) {
+        return false;
+    }
+    for (uint32_t i = 0; i < sizeof(magic); i++) {
+        if (magic[i] != GGUF_MAGIC[i]) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool is_safetensors_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    // get file size
+    file.seekg(0, file.end);
+    size_t file_size_ = file.tellg();
+    file.seekg(0, file.beg);
+
+    // read header size
+    if (file_size_ <= ST_HEADER_SIZE_LEN) {
+        return false;
+    }
+
+    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
+    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
+    if (!file) {
+        return false;
+    }
+
+    size_t header_size_ = read_u64(header_size_buf);
+    if (header_size_ >= file_size_ || header_size_ <= 2) {
+        return false;
+    }
+
+    // read header
+    std::vector<char> header_buf;
+    header_buf.resize(header_size_ + 1);
+    header_buf[header_size_] = '\0';
+    file.read(header_buf.data(), header_size_);
+    if (!file) {
+        return false;
+    }
+    nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
+    if (header_.is_discarded()) {
+        return false;
+    }
+    return true;
+}
+
 bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    file_paths_.push_back(file_path);
-    return true;
-}
-
-bool ModelLoader::init_from_files(const std::vector<std::string>& file_paths) {
-    for (auto& file_path : file_paths) {
-        if (!init_from_file(file_path)) {
-            return false;
-        }
+    if (is_directory(file_path)) {
+        LOG_INFO("load %s using diffusers format", file_path.c_str());
+        return init_from_diffusers_file(file_path, prefix);
+    } else if (is_gguf_file(file_path)) {
+        LOG_INFO("load %s using gguf format", file_path.c_str());
+        return init_from_gguf_file(file_path, prefix);
+    } else if (is_safetensors_file(file_path)) {
+        LOG_INFO("load %s using safetensors format", file_path.c_str());
+        return init_from_safetensors_file(file_path, prefix);
+    } else if (is_zip_file(file_path)) {
+        LOG_INFO("load %s using checkpoint format", file_path.c_str());
+        return init_from_ckpt_file(file_path, prefix);
+    } else {
+        LOG_WARN("unknown format %s", file_path.c_str());
+        return false;
     }
-    return true;
-}
-
-SDVersion ModelLoader::get_sd_version() {
-    TensorStorage token_embedding_weight;
-    for (auto& tensor_storage : tensor_storages) {
-        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" ||
-            tensor_storage.name == "te.text_model.embeddings.token_embedding.weight") {
-            token_embedding_weight = tensor_storage;
-            break;
-        }
-    }
-    if (token_embedding_weight.ne[0] == 768) {
-        return VERSION_1_x;
-    } else if (token_embedding_weight.ne[0] == 1024) {
-        return VERSION_2_x;
-    }
-    return VERSION_COUNT;
-}
-
-ggml_type ModelLoader::get_sd_wtype() {
-    for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
-        if (tensor_storage.name.find(".weight") != std::string::npos &&
-            tensor_storage.name.find("time_embed") != std::string::npos) {
-            return tensor_storage.type;
-        }
-    }
-    return GGML_TYPE_COUNT;
-}
-
-bool ModelLoader::load_vocab(on_new_token_cb_t on_new_token_cb) {
-    char* vocab_buffer          = reinterpret_cast<char*>(vocab_json);
-    nlohmann::json vocab        = nlohmann::json::parse(vocab_buffer);
-    std::map<char, int> decoder = unicode_to_byte();
-    for (auto& it : vocab.items()) {
-        int token_id          = it.value();
-        std::string token_str = it.key();
-        std::string token     = "";
-        for (char c : token_str) {
-            token += decoder[c];
-        }
-        on_new_token_cb(token, token_id);
-    }
-    return true;
-}
-
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {
-    bool success = true;
-    for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
-        std::string file_path = file_paths_[file_index];
-        LOG_DEBUG("loading tensors from %s", file_path.c_str());
-
-        std::ifstream file(file_path, std::ios::binary);
-        if (!file.is_open()) {
-            LOG_ERROR("failed to open '%s'", file_path.c_str());
-            return false;
-        }
-
-        bool is_zip = false;
-        for (auto& tensor_storage : tensor_storages) {
-            if (tensor_storage.index_in_zip >= 0) {
-                is_zip = true;
-                break;
-            }
-        }
-
-        struct zip_t* zip = NULL;
-        if (is_zip) {
-            zip = zip_open(file_path.c_str(), 0, 'r');
-            if (zip == NULL) {
-                LOG_ERROR("failed to open zip '%s'", file_path.c_str());
-                return false;
-            }
-        }
-
-        std::vector<uint8_t> read_buffer;
-        std::vector<uint8_t> convert_buffer;
-
-        auto read_data = [&](const TensorStorage& tensor_storage, char* buf, size_t n) {
-            if (zip != NULL) {
-                zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
-                size_t entry_size = zip_entry_size(zip);
-                if (entry_size != n) {
-                    read_buffer.resize(entry_size);
-                    zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
-                    memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
-                } else {
-                    zip_entry_noallocread(zip, (void*)buf, n);
-                }
-                zip_entry_close(zip);
-            } else {
-                file.seekg(tensor_storage.offset);
-                file.read(buf, n);
-                if (!file) {
-                    LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                    return false;
-                }
-            }
-            return true;
-        };
-
-        std::vector<TensorStorage> processed_tensor_storages;
-        for (auto& tensor_storage : 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", name.c_str());
-
-            ggml_tensor* dst_tensor = NULL;
-
-            success = on_new_tensor_cb(tensor_storage, &dst_tensor);
-            if (!success) {
-                LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
-                break;
-            }
-
-            if (dst_tensor == NULL) {
-                continue;
-            }
-
-            ggml_backend_t backend = ggml_get_backend(dst_tensor);
-
-            size_t nbytes_to_read = tensor_storage.nbytes_to_read();
-
-            if (backend == NULL || ggml_backend_is_cpu(backend)) {
-                // 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) == nbytes_to_read);
-                    read_data(tensor_storage, (char*)dst_tensor->data, nbytes_to_read);
-
-                    if (tensor_storage.is_bf16) {
-                        // inplace op
-                        bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
-                    }
-                } else {
-                    read_buffer.resize(tensor_storage.nbytes());
-                    read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-
-                    if (tensor_storage.is_bf16) {
-                        // inplace op
-                        bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                    }
-
-                    convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
-                                   dst_tensor->type, (int)tensor_storage.nelements());
-                }
-            } else {
-                read_buffer.resize(tensor_storage.nbytes());
-                read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-
-                if (tensor_storage.is_bf16) {
-                    // inplace op
-                    bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                }
-
-                if (tensor_storage.type == dst_tensor->type) {
-                    // copy to device memory
-                    ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
-                } else {
-                    // convert first, then copy to device memory
-                    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());
-                    ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
-                }
-            }
-        }
-
-        if (zip != NULL) {
-            zip_close(zip);
-        }
-
-        if (!success) {
-            break;
-        }
-    }
-    return success;
-}
-
-int64_t ModelLoader::cal_mem_size() {
-    int64_t mem_size = 0;
-    for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
-        mem_size += tensor_storage.nbytes();
-        mem_size += GGML_MEM_ALIGN * 2;  // for lora alphas
-    }
-
-    return mem_size + 10 * 1024 * 1024;
 }
 
 /*================================================= GGUFModelLoader ==================================================*/
 
-bool GGUFModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    LOG_INFO("loading model from '%s'", file_path.c_str());
-    ModelLoader::init_from_file(file_path, prefix);
+bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+    file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
 
     gguf_context* ctx_gguf_ = NULL;
@@ -781,7 +679,7 @@ bool GGUFModelLoader::init_from_file(const std::string& file_path, const std::st
 
         // 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());
 
@@ -796,8 +694,6 @@ bool GGUFModelLoader::init_from_file(const std::string& file_path, const std::st
 
 /*================================================= SafeTensorsModelLoader ==================================================*/
 
-#define ST_HEADER_SIZE_LEN 8
-
 ggml_type str_to_ggml_type(const std::string& dtype) {
     ggml_type ttype = GGML_TYPE_COUNT;
     if (dtype == "F16") {
@@ -811,8 +707,9 @@ ggml_type str_to_ggml_type(const std::string& dtype) {
 }
 
 // https://huggingface.co/docs/safetensors/index
-bool SafeTensorsModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    ModelLoader::init_from_file(file_path, prefix);
+bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+    file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
@@ -913,21 +810,18 @@ bool SafeTensorsModelLoader::init_from_file(const std::string& file_path, const
 
 /*================================================= DiffusersModelLoader ==================================================*/
 
-bool DiffusersModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    if (!is_directory(file_path)) {
-        return SafeTensorsModelLoader::init_from_file(file_path, prefix);
-    }
+bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const std::string& prefix) {
     std::string unet_path = path_join(file_path, "unet/diffusion_pytorch_model.safetensors");
     std::string vae_path  = path_join(file_path, "vae/diffusion_pytorch_model.safetensors");
     std::string clip_path = path_join(file_path, "text_encoder/model.safetensors");
 
-    if (!SafeTensorsModelLoader::init_from_file(unet_path, "unet.")) {
+    if (!init_from_safetensors_file(unet_path, "unet.")) {
         return false;
     }
-    if (!SafeTensorsModelLoader::init_from_file(vae_path, "vae.")) {
+    if (!init_from_safetensors_file(vae_path, "vae.")) {
         return false;
     }
-    if (!SafeTensorsModelLoader::init_from_file(clip_path, "te.")) {
+    if (!init_from_safetensors_file(clip_path, "te.")) {
         return false;
     }
     return true;
@@ -1127,12 +1021,12 @@ int find_char(uint8_t* buffer, int len, char c) {
 
 #define MAX_STRING_BUFFER 512
 
-bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
-                                     size_t buffer_size,
-                                     zip_t* zip,
-                                     std::string dir,
-                                     size_t file_index,
-                                     const std::string& prefix) {
+bool ModelLoader::parse_data_pkl(uint8_t* buffer,
+                                 size_t buffer_size,
+                                 zip_t* zip,
+                                 std::string dir,
+                                 size_t file_index,
+                                 const std::string& prefix) {
     uint8_t* buffer_end = buffer + buffer_size;
     if (buffer[0] == 0x80) {  // proto
         if (buffer[1] != 2) {
@@ -1235,6 +1129,7 @@ bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
                         reader.tensor_storage.file_index = file_index;
                         reader.tensor_storage.name       = prefix + reader.tensor_storage.name;
                         tensor_storages.push_back(reader.tensor_storage);
+                        // LOG_DEBUG("%s", reader.tensor_storage.name.c_str());
                         // reset
                         reader = PickleTensorReader();
                     }
@@ -1250,8 +1145,9 @@ bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
     return true;
 }
 
-bool CkptModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    ModelLoader::init_from_file(file_path, prefix);
+bool ModelLoader::init_from_ckpt_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+    file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
 
     struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
@@ -1271,7 +1167,7 @@ bool CkptModelLoader::init_from_file(const std::string& file_path, const std::st
                 size_t pkl_size;
                 zip_entry_read(zip, &pkl_data, &pkl_size);
 
-                LOG_DEBUG("%lld", pkl_size);
+                // LOG_DEBUG("%lld", pkl_size);
 
                 parse_data_pkl((uint8_t*)pkl_data, pkl_size, zip, dir, file_index, prefix);
 
@@ -1284,29 +1180,382 @@ bool CkptModelLoader::init_from_file(const std::string& file_path, const std::st
     return true;
 }
 
-/*================================================= init_model_loader_from_file ==================================================*/
+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 == "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" ||
+            tensor_storage.name == "te.text_model.embeddings.token_embedding.weight" ||
+            tensor_storage.name == "conditioner.embedders.0.model.token_embedding.weight" ||
+            tensor_storage.name == "conditioner.embedders.0.transformer.text_model.embeddings.token_embedding.weight") {
+            token_embedding_weight = tensor_storage;
+            // break;
+        }
+    }
+    if (token_embedding_weight.ne[0] == 768) {
+        return VERSION_1_x;
+    } else if (token_embedding_weight.ne[0] == 1024) {
+        return VERSION_2_x;
+    }
+    return VERSION_COUNT;
+}
 
-ModelLoader* init_model_loader_from_file(const std::string& file_path) {
-    ModelLoader* model_loader = NULL;
-    if (is_directory(file_path)) {
-        LOG_DEBUG("load %s using diffusers format", file_path.c_str());
-        model_loader = new DiffusersModelLoader();
-    } else if (ends_with(file_path, ".gguf")) {
-        LOG_DEBUG("load %s using gguf format", file_path.c_str());
-        model_loader = new GGUFModelLoader();
-    } else if (ends_with(file_path, ".safetensors")) {
-        LOG_DEBUG("load %s using safetensors format", file_path.c_str());
-        model_loader = new SafeTensorsModelLoader();
-    } else if (ends_with(file_path, ".ckpt")) {
-        LOG_DEBUG("load %s using checkpoint format", file_path.c_str());
-        model_loader = new CkptModelLoader();
-    } else {
-        LOG_DEBUG("unknown format %s", file_path.c_str());
-        return NULL;
+ggml_type ModelLoader::get_sd_wtype() {
+    for (auto& tensor_storage : tensor_storages) {
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+
+        if (tensor_storage.name.find(".weight") != std::string::npos &&
+            tensor_storage.name.find("time_embed") != std::string::npos) {
+            return tensor_storage.type;
+        }
     }
-    if (!model_loader->init_from_file(file_path)) {
-        delete model_loader;
-        model_loader = NULL;
+    return GGML_TYPE_COUNT;
+}
+
+std::string ModelLoader::load_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(merges_utf8_c_str), sizeof(merges_utf8_c_str));
+    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;
+        }
     }
-    return model_loader;
+
+    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];
+        LOG_DEBUG("loading tensors from %s", file_path.c_str());
+
+        std::ifstream file(file_path, std::ios::binary);
+        if (!file.is_open()) {
+            LOG_ERROR("failed to open '%s'", file_path.c_str());
+            return false;
+        }
+
+        bool is_zip = false;
+        for (auto& tensor_storage : tensor_storages) {
+            if (tensor_storage.file_index != file_index) {
+                continue;
+            }
+            if (tensor_storage.index_in_zip >= 0) {
+                is_zip = true;
+                break;
+            }
+        }
+
+        struct zip_t* zip = NULL;
+        if (is_zip) {
+            zip = zip_open(file_path.c_str(), 0, 'r');
+            if (zip == NULL) {
+                LOG_ERROR("failed to open zip '%s'", file_path.c_str());
+                return false;
+            }
+        }
+
+        std::vector<uint8_t> read_buffer;
+        std::vector<uint8_t> convert_buffer;
+
+        auto read_data = [&](const TensorStorage& tensor_storage, char* buf, size_t n) {
+            if (zip != NULL) {
+                zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
+                size_t entry_size = zip_entry_size(zip);
+                if (entry_size != n) {
+                    read_buffer.resize(entry_size);
+                    zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
+                    memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
+                } else {
+                    zip_entry_noallocread(zip, (void*)buf, n);
+                }
+                zip_entry_close(zip);
+            } else {
+                file.seekg(tensor_storage.offset);
+                file.read(buf, n);
+                if (!file) {
+                    LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
+                    return false;
+                }
+            }
+            return true;
+        };
+
+        for (auto& tensor_storage : processed_tensor_storages) {
+            if (tensor_storage.file_index != file_index) {
+                continue;
+            }
+            // LOG_DEBUG("%s", tensor_storage.name.c_str());
+
+            ggml_tensor* dst_tensor = NULL;
+
+            success = on_new_tensor_cb(tensor_storage, &dst_tensor);
+            if (!success) {
+                LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
+                break;
+            }
+
+            if (dst_tensor == NULL) {
+                continue;
+            }
+
+            size_t nbytes_to_read = tensor_storage.nbytes_to_read();
+
+            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());
+                    read_data(tensor_storage, (char*)dst_tensor->data, nbytes_to_read);
+
+                    if (tensor_storage.is_bf16) {
+                        // inplace op
+                        bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
+                    }
+                } else {
+                    read_buffer.resize(tensor_storage.nbytes());
+                    read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
+
+                    if (tensor_storage.is_bf16) {
+                        // inplace op
+                        bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                    }
+
+                    convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
+                                   dst_tensor->type, (int)tensor_storage.nelements());
+                }
+            } else {
+                read_buffer.resize(tensor_storage.nbytes());
+                read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
+
+                if (tensor_storage.is_bf16) {
+                    // inplace op
+                    bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                }
+
+                if (tensor_storage.type == dst_tensor->type) {
+                    // copy to device memory
+                    ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
+                } else {
+                    // convert first, then copy to device memory
+                    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());
+                    ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
+                }
+            }
+        }
+
+        if (zip != NULL) {
+            zip_close(zip);
+        }
+
+        if (!success) {
+            break;
+        }
+    }
+    return success;
+}
+
+bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
+                               ggml_backend_t backend,
+                               std::set<std::string> ignore_tensors) {
+    std::set<std::string> tensor_names_in_file;
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+        tensor_names_in_file.insert(name);
+
+        struct ggml_tensor* real;
+        if (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());
+            }
+            return true;
+        }
+
+        if (
+            real->ne[0] != tensor_storage.ne[0] ||
+            real->ne[1] != tensor_storage.ne[1] ||
+            real->ne[2] != tensor_storage.ne[2] ||
+            real->ne[3] != tensor_storage.ne[3]) {
+            LOG_ERROR(
+                "tensor '%s' has wrong shape in model file: "
+                "got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
+                name.c_str(),
+                (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
+                (int)real->ne[0], (int)real->ne[1], (int)real->ne[2], (int)real->ne[3]);
+            return false;
+        }
+
+        *dst_tensor = real;
+
+        return true;
+    };
+
+    bool success = load_tensors(on_new_tensor_cb, backend);
+    if (!success) {
+        LOG_ERROR("load tensors from file failed");
+        return false;
+    }
+
+    bool some_tensor_not_init = false;
+
+    for (auto pair : tensors) {
+        if (pair.first.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos) {
+            continue;
+        }
+        if (pair.first.find("alphas_cumprod") != std::string::npos) {
+            continue;
+        }
+
+        if (pair.first.find("alphas_cumprod") != std::string::npos) {
+            continue;
+        }
+
+        if (tensor_names_in_file.find(pair.first) == tensor_names_in_file.end()) {
+            LOG_ERROR("tensor '%s' not in model file", pair.first.c_str());
+            some_tensor_not_init = true;
+        }
+    }
+
+    if (some_tensor_not_init) {
+        return false;
+    }
+    return true;
+}
+
+bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type) {
+    auto backend    = ggml_backend_cpu_init();
+    size_t mem_size = 1 * 1024 * 1024;  // for padding
+    mem_size += tensor_storages.size() * ggml_tensor_overhead();
+    mem_size += cal_mem_size(backend, type);
+    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
+    ggml_context* ggml_ctx = ggml_init({mem_size, NULL, false});
+
+    gguf_context* gguf_ctx = gguf_init_empty();
+
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+
+        ggml_type tensor_type = tensor_storage.type;
+        if (type != GGML_TYPE_COUNT) {
+            if (ggml_is_quantized(type) && tensor_storage.ne[0] % 32 != 0) {
+                tensor_type = GGML_TYPE_F16;
+            } else {
+                tensor_type = type;
+            }
+        }
+
+        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
+        if (tensor == NULL) {
+            LOG_ERROR("ggml_new_tensor failed");
+            return false;
+        }
+        ggml_set_name(tensor, name.c_str());
+
+        // LOG_DEBUG("%s %d %s %d[%d %d %d %d] %d[%d %d %d %d]", name.c_str(),
+        // ggml_nbytes(tensor), ggml_type_name(tensor_type),
+        // tensor_storage.n_dims,
+        // tensor_storage.ne[0], tensor_storage.ne[1], tensor_storage.ne[2], tensor_storage.ne[3],
+        // tensor->n_dims, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+
+        *dst_tensor = tensor;
+
+        gguf_add_tensor(gguf_ctx, tensor);
+
+        return true;
+    };
+
+    bool success = load_tensors(on_new_tensor_cb, backend);
+    ggml_backend_free(backend);
+    LOG_INFO("load tensors done");
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    if (success) {
+        gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
+    }
+    ggml_free(ggml_ctx);
+    gguf_free(gguf_ctx);
+    return success;
+}
+
+int64_t ModelLoader::cal_mem_size(ggml_backend_t backend, ggml_type type) {
+    size_t alignment = 128;
+    if (backend != NULL) {
+        alignment = ggml_backend_get_alignment(backend);
+    }
+    int64_t mem_size = 0;
+    std::vector<TensorStorage> processed_tensor_storages;
+    for (auto& tensor_storage : tensor_storages) {
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+        preprocess_tensor(tensor_storage, processed_tensor_storages);
+    }
+
+    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,9 +4,11 @@
 #include <functional>
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
 
+#include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
 #include "json.hpp"
 #include "zip.h"
@@ -91,36 +93,12 @@ struct TensorStorage {
 };
 
 typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
-typedef std::function<void(const std::string&, int32_t)> on_new_token_cb_t;
 
 class ModelLoader {
 protected:
     std::vector<std::string> file_paths_;
     std::vector<TensorStorage> tensor_storages;
 
-public:
-    virtual bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-    virtual bool init_from_files(const std::vector<std::string>& file_paths);
-    virtual SDVersion get_sd_version();
-    virtual ggml_type get_sd_wtype();
-    virtual bool load_vocab(on_new_token_cb_t on_new_token_cb);
-    virtual bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb);
-    virtual int64_t cal_mem_size();
-    virtual ~ModelLoader() = default;
-};
-
-class GGUFModelLoader : public ModelLoader {
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-};
-
-class SafeTensorsModelLoader : public ModelLoader {
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-};
-
-class CkptModelLoader : public ModelLoader {
-private:
     bool parse_data_pkl(uint8_t* buffer,
                         size_t buffer_size,
                         zip_t* zip,
@@ -128,15 +106,22 @@ private:
                         size_t file_index,
                         const std::string& prefix);
 
+    bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_ckpt_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
+
 public:
     bool init_from_file(const std::string& file_path, const std::string& prefix = "");
+    SDVersion get_sd_version();
+    ggml_type get_sd_wtype();
+    std::string load_merges();
+    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend_t backend);
+    bool load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
+                      ggml_backend_t backend,
+                      std::set<std::string> ignore_tensors = {});
+    bool save_to_gguf_file(const std::string& file_path, ggml_type type);
+    int64_t cal_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
+    ~ModelLoader() = default;
 };
-
-class DiffusersModelLoader : public SafeTensorsModelLoader {
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-};
-
-ModelLoader* init_model_loader_from_file(const std::string& file_path);
-
 #endif  // __MODEL_H__

preprocessing.hpp

@@ -0,0 +1,229 @@
+#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__

rng_philox.hpp

@@ -4,7 +4,7 @@
 #include <cmath>
 #include <vector>
 
-#include "rng.h"
+#include "rng.hpp"
 
 // RNG imitiating torch cuda randn on CPU.
 // Port from: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/modules/rng_philox.py

stable-diffusion.h

@@ -1,16 +1,39 @@
 #ifndef __STABLE_DIFFUSION_H__
 #define __STABLE_DIFFUSION_H__
 
-#include <string>
-#include <memory>
-#include <vector>
+#if defined(_WIN32) || defined(__CYGWIN__)
+#ifndef SD_BUILD_SHARED_LIB
+#define SD_API
+#else
+#ifdef SD_BUILD_DLL
+#define SD_API __declspec(dllexport)
+#else
+#define SD_API __declspec(dllimport)
+#endif
+#endif
+#else
+#if __GNUC__ >= 4
+#define SD_API __attribute__((visibility("default")))
+#else
+#define SD_API
+#endif
+#endif
 
-enum RNGType {
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <string.h>
+
+enum rng_type_t {
     STD_DEFAULT_RNG,
     CUDA_RNG
 };
 
-enum SampleMethod {
+enum sample_method_t {
     EULER_A,
     EULER,
     HEUN,
@@ -22,53 +45,121 @@ enum SampleMethod {
     N_SAMPLE_METHODS
 };
 
-enum Schedule {
+enum schedule_t {
     DEFAULT,
     DISCRETE,
     KARRAS,
     N_SCHEDULES
 };
 
-class StableDiffusionGGML;
-
-class StableDiffusion {
-private:
-    std::shared_ptr<StableDiffusionGGML> sd;
-
-public:
-    StableDiffusion(int n_threads                = -1,
-                    bool vae_decode_only         = false,
-                    bool free_params_immediately = false,
-                    std::string lora_model_dir   = "",
-                    RNGType rng_type             = STD_DEFAULT_RNG);
-    bool load_from_file(const std::string& model_path,
-                        const std::string& vae_path,
-                        ggml_type wtype,
-                        Schedule d = DEFAULT);
-    std::vector<uint8_t*> txt2img(
-        std::string prompt,
-        std::string negative_prompt,
-        float cfg_scale,
-        int width,
-        int height,
-        SampleMethod sample_method,
-        int sample_steps,
-        int64_t seed,
-        int batch_count);
-
-    std::vector<uint8_t*> img2img(
-        const uint8_t* init_img_data,
-        std::string prompt,
-        std::string negative_prompt,
-        float cfg_scale,
-        int width,
-        int height,
-        SampleMethod sample_method,
-        int sample_steps,
-        float strength,
-        int64_t seed);
+// same as enum ggml_type
+enum sd_type_t {
+    SD_TYPE_F32  = 0,
+    SD_TYPE_F16  = 1,
+    SD_TYPE_Q4_0 = 2,
+    SD_TYPE_Q4_1 = 3,
+    // SD_TYPE_Q4_2 = 4, support has been removed
+    // SD_TYPE_Q4_3 (5) support has been removed
+    SD_TYPE_Q5_0 = 6,
+    SD_TYPE_Q5_1 = 7,
+    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_IQ2_XXS = 16,
+    SD_TYPE_I8,
+    SD_TYPE_I16,
+    SD_TYPE_I32,
+    SD_TYPE_COUNT,
 };
 
-std::string sd_get_system_info();
+SD_API const char* sd_type_name(enum sd_type_t type);
+
+enum sd_log_level_t {
+    SD_LOG_DEBUG,
+    SD_LOG_INFO,
+    SD_LOG_WARN,
+    SD_LOG_ERROR
+};
+
+typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
+
+SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
+SD_API int32_t get_num_physical_cores();
+SD_API const char* sd_get_system_info();
+
+typedef struct {
+    uint32_t width;
+    uint32_t height;
+    uint32_t channel;
+    uint8_t* data;
+} sd_image_t;
+
+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,
+                            bool keep_control_net_cpu);
+
+SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
+
+SD_API sd_image_t* txt2img(sd_ctx_t* sd_ctx,
+                           const char* prompt,
+                           const char* negative_prompt,
+                           int clip_skip,
+                           float cfg_scale,
+                           int width,
+                           int height,
+                           enum sample_method_t sample_method,
+                           int sample_steps,
+                           int64_t seed,
+                           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,
+                           const char* prompt,
+                           const char* negative_prompt,
+                           int clip_skip,
+                           float cfg_scale,
+                           int width,
+                           int height,
+                           enum sample_method_t sample_method,
+                           int sample_steps,
+                           float strength,
+                           int64_t seed,
+                           int batch_count);
+
+typedef struct upscaler_ctx_t upscaler_ctx_t;
+
+SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
+                                        int n_threads,
+                                        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* 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
+}
+#endif
 
 #endif  // __STABLE_DIFFUSION_H__

tae.hpp

@@ -0,0 +1,581 @@
+#ifndef __TAE_HPP__
+#define __TAE_HPP__
+
+#include "ggml_extend.hpp"
+
+#include "model.h"
+
+/*
+    ===================================    TinyAutoEncoder  ===================================
+    References:
+    https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoder_tiny.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]
+
+    // 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);
+        }
+    }
+
+    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;
+
+        tensors[prefix + "conv.2.weight"] = conv_1_w;
+        tensors[prefix + "conv.2.bias"]   = conv_1_b;
+
+        tensors[prefix + "conv.4.weight"] = conv_2_w;
+        tensors[prefix + "conv.4.bias"]   = conv_2_b;
+
+        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);
+        }
+
+        h = ggml_add(ctx, h, x);
+        h = ggml_relu_inplace(ctx, h);
+        return h;
+    }
+};
+
+struct TinyEncoder {
+    int in_channels = 3;
+    int z_channels  = 4;
+    int channels    = 64;
+    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]
+
+    TinyEncoder() {
+        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;
+        }
+
+        initial_block.in_channels  = channels;
+        initial_block.out_channels = channels;
+    }
+
+    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
+
+        mem_size += initial_block.calculate_mem_size();
+
+        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) + ".");
+        }
+
+        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;
+    }
+};
+
+struct TinyDecoder {
+    int z_channels      = 4;
+    int channels        = 64;
+    int output_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]
+
+    // 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]
+
+    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;
+        }
+
+        final_block.in_channels  = channels;
+        final_block.out_channels = channels;
+    }
+
+    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
+
+        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);
+    }
+
+    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, 1.0f / 3.0f);
+        h      = ggml_tanh_inplace(ctx, h);
+        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);
+
+        // 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);
+        }
+
+        // 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;
+    }
+};
+
+struct TinyAutoEncoder : public GGMLModule {
+    TinyEncoder encoder;
+    TinyDecoder decoder;
+    bool decode_only = false;
+
+    TinyAutoEncoder(bool decoder_only_ = true)
+        : decode_only(decoder_only_) {
+        name = "tae";
+    }
+
+    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;
+    }
+
+    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;
+    }
+
+    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);
+    }
+
+    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) {
+        LOG_INFO("loading taesd from '%s'", file_path.c_str());
+
+        if (!alloc_params_buffer(backend)) {
+            return false;
+        }
+
+        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;
+        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);
+        }
+
+        ModelLoader model_loader;
+        if (!model_loader.init_from_file(file_path)) {
+            LOG_ERROR("init taesd model loader from file failed: '%s'", file_path.c_str());
+            return false;
+        }
+
+        bool success = model_loader.load_tensors(tensors_need_to_load, backend, ignore_tensors);
+
+        if (!success) {
+            LOG_ERROR("load tae tensors from model loader failed");
+            return false;
+        }
+
+        LOG_INFO("taesd model loaded");
+        return success;
+    }
+
+    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_);
+
+        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) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(z, decode_graph);
+        };
+        GGMLModule::compute(get_graph, n_threads, work_result);
+    }
+};
+
+#endif  // __TAE_HPP__

unet.hpp

@@ -0,0 +1,665 @@
+#ifndef __UNET_HPP__
+#define __UNET_HPP__
+
+#include "common.hpp"
+#include "ggml_extend.hpp"
+#include "model.h"
+
+/*==================================================== UnetModel =====================================================*/
+
+#define UNET_GRAPH_SIZE 10240
+
+// ldm.modules.diffusionmodules.openaimodel.UNetModel
+struct UNetModel : public GGMLModule {
+    SDVersion version = VERSION_1_x;
+    // network hparams
+    int in_channels                        = 4;
+    int model_channels                     = 320;
+    int out_channels                       = 4;
+    int num_res_blocks                     = 2;
+    std::vector<int> attention_resolutions = {4, 2, 1};
+    std::vector<int> channel_mult          = {1, 2, 4, 4};
+    std::vector<int> transformer_depth     = {1, 1, 1, 1};
+    int time_embed_dim                     = 1280;  // model_channels*4
+    int num_heads                          = 8;
+    int num_head_channels                  = -1;    // channels // num_heads
+    int context_dim                        = 768;   // 1024 for VERSION_2_x, 2048 for VERSION_XL
+    int adm_in_channels                    = 2816;  // only for VERSION_XL
+
+    // network params
+    struct ggml_tensor* time_embed_0_w;  // [time_embed_dim, model_channels]
+    struct ggml_tensor* time_embed_0_b;  // [time_embed_dim, ]
+    // time_embed_1 is nn.SILU()
+    struct ggml_tensor* time_embed_2_w;  // [time_embed_dim, time_embed_dim]
+    struct ggml_tensor* time_embed_2_b;  // [time_embed_dim, ]
+
+    struct ggml_tensor* label_embed_0_w;  // [time_embed_dim, adm_in_channels]
+    struct ggml_tensor* label_embed_0_b;  // [time_embed_dim, ]
+    // label_embed_1 is nn.SILU()
+    struct ggml_tensor* label_embed_2_w;  // [time_embed_dim, time_embed_dim]
+    struct ggml_tensor* label_embed_2_b;  // [time_embed_dim, ]
+
+    struct ggml_tensor* input_block_0_w;  // [model_channels, in_channels, 3, 3]
+    struct ggml_tensor* input_block_0_b;  // [model_channels, ]
+
+    // input_blocks
+    ResBlock input_res_blocks[4][2];
+    SpatialTransformer input_transformers[3][2];
+    DownSample input_down_samples[3];
+
+    // middle_block
+    ResBlock middle_block_0;
+    SpatialTransformer middle_block_1;
+    ResBlock middle_block_2;
+
+    // output_blocks
+    ResBlock output_res_blocks[4][3];
+    SpatialTransformer output_transformers[3][3];
+    UpSample output_up_samples[3];
+
+    // out
+    // group norm 32
+    struct ggml_tensor* out_0_w;  // [model_channels, ]
+    struct ggml_tensor* out_0_b;  // [model_channels, ]
+    // out 1 is nn.SILU()
+    struct ggml_tensor* out_2_w;  // [out_channels, model_channels, 3, 3]
+    struct ggml_tensor* out_2_b;  // [out_channels, ]
+
+    UNetModel(SDVersion version = VERSION_1_x)
+        : version(version) {
+        name = "unet";
+        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;
+        }
+        // set up hparams of blocks
+
+        // input_blocks
+        std::vector<int> input_block_chans;
+        input_block_chans.push_back(model_channels);
+        int ch = model_channels;
+        int ds = 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_res_blocks[i][j].channels     = ch;
+                input_res_blocks[i][j].emb_channels = time_embed_dim;
+                input_res_blocks[i][j].out_channels = mult * model_channels;
+
+                ch = mult * model_channels;
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    int n_head = num_heads;
+                    int d_head = ch / num_heads;
+                    if (num_head_channels != -1) {
+                        d_head = num_head_channels;
+                        n_head = ch / d_head;
+                    }
+                    input_transformers[i][j]             = SpatialTransformer(transformer_depth[i]);
+                    input_transformers[i][j].in_channels = ch;
+                    input_transformers[i][j].n_head      = n_head;
+                    input_transformers[i][j].d_head      = d_head;
+                    input_transformers[i][j].context_dim = context_dim;
+                }
+                input_block_chans.push_back(ch);
+            }
+            if (i != len_mults - 1) {
+                input_down_samples[i].channels     = ch;
+                input_down_samples[i].out_channels = ch;
+                input_block_chans.push_back(ch);
+
+                ds *= 2;
+            }
+        }
+
+        // middle blocks
+        middle_block_0.channels     = ch;
+        middle_block_0.emb_channels = time_embed_dim;
+        middle_block_0.out_channels = ch;
+
+        int n_head = num_heads;
+        int d_head = ch / num_heads;
+        if (num_head_channels != -1) {
+            d_head = num_head_channels;
+            n_head = ch / d_head;
+        }
+        middle_block_1             = SpatialTransformer(transformer_depth[transformer_depth.size() - 1]);
+        middle_block_1.in_channels = ch;
+        middle_block_1.n_head      = n_head;
+        middle_block_1.d_head      = d_head;
+        middle_block_1.context_dim = context_dim;
+
+        middle_block_2.channels     = ch;
+        middle_block_2.emb_channels = time_embed_dim;
+        middle_block_2.out_channels = ch;
+
+        // output blocks
+        for (int i = (int)len_mults - 1; i >= 0; i--) {
+            int mult = channel_mult[i];
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                int ich = input_block_chans.back();
+                input_block_chans.pop_back();
+
+                output_res_blocks[i][j].channels     = ch + ich;
+                output_res_blocks[i][j].emb_channels = time_embed_dim;
+                output_res_blocks[i][j].out_channels = mult * model_channels;
+
+                ch = mult * model_channels;
+
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    int n_head = num_heads;
+                    int d_head = ch / num_heads;
+                    if (num_head_channels != -1) {
+                        d_head = num_head_channels;
+                        n_head = ch / d_head;
+                    }
+                    output_transformers[i][j]             = SpatialTransformer(transformer_depth[i]);
+                    output_transformers[i][j].in_channels = ch;
+                    output_transformers[i][j].n_head      = n_head;
+                    output_transformers[i][j].d_head      = d_head;
+                    output_transformers[i][j].context_dim = context_dim;
+                }
+
+                if (i > 0 && j == num_res_blocks) {
+                    output_up_samples[i - 1].channels     = ch;
+                    output_up_samples[i - 1].out_channels = ch;
+
+                    ds /= 2;
+                }
+            }
+        }
+    }
+
+    size_t calculate_mem_size() {
+        size_t mem_size = 0;
+        mem_size += ggml_row_size(wtype, time_embed_dim * model_channels);  // time_embed_0_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);           // time_embed_0_b
+        mem_size += ggml_row_size(wtype, time_embed_dim * time_embed_dim);  // time_embed_2_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);           // time_embed_2_b
+
+        if (version == VERSION_XL) {
+            mem_size += ggml_row_size(wtype, time_embed_dim * adm_in_channels);  // label_embed_0_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);            // label_embed_0_b
+            mem_size += ggml_row_size(wtype, time_embed_dim * time_embed_dim);   // label_embed_2_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, time_embed_dim);            // label_embed_2_b
+        }
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, model_channels * in_channels * 3 * 3);  // input_block_0_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, model_channels);                        // input_block_0_b
+
+        // input_blocks
+        int ds           = 1;
+        size_t len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                mem_size += input_res_blocks[i][j].calculate_mem_size(wtype);
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    mem_size += input_transformers[i][j].calculate_mem_size(wtype);
+                }
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                mem_size += input_down_samples[i].calculate_mem_size(wtype);
+            }
+        }
+
+        // middle_block
+        mem_size += middle_block_0.calculate_mem_size(wtype);
+        mem_size += middle_block_1.calculate_mem_size(wtype);
+        mem_size += middle_block_2.calculate_mem_size(wtype);
+
+        // output_blocks
+        for (int i = (int)len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                mem_size += output_res_blocks[i][j].calculate_mem_size(wtype);
+
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    mem_size += output_transformers[i][j].calculate_mem_size(wtype);
+                }
+
+                if (i > 0 && j == num_res_blocks) {
+                    mem_size += output_up_samples[i - 1].calculate_mem_size(wtype);
+
+                    ds /= 2;
+                }
+            }
+        }
+
+        // out
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, model_channels);                     // out_0_w/b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * model_channels * 3 * 3);  // out_2_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                           // out_2_b
+
+        return mem_size;
+    }
+
+    size_t get_num_tensors() {
+        // in
+        int num_tensors = 6;
+        if (version == VERSION_XL) {
+            num_tensors += 4;
+        }
+
+        // input blocks
+        int ds           = 1;
+        size_t len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                num_tensors += 12;
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    num_tensors += input_transformers[i][j].get_num_tensors();
+                }
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                num_tensors += 2;
+            }
+        }
+
+        // middle blocks
+        num_tensors += 13 * 2;
+        num_tensors += middle_block_1.get_num_tensors();
+
+        // output blocks
+        for (int i = (int)len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                num_tensors += 12;
+
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    num_tensors += output_transformers[i][j].get_num_tensors();
+                }
+
+                if (i > 0 && j == num_res_blocks) {
+                    num_tensors += 2;
+
+                    ds /= 2;
+                }
+            }
+        }
+
+        // out
+        num_tensors += 4;
+        return num_tensors;
+    }
+
+    void init_params() {
+        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
+        time_embed_0_w     = ggml_new_tensor_2d(params_ctx, wtype, model_channels, time_embed_dim);
+        time_embed_0_b     = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, time_embed_dim);
+        time_embed_2_w     = ggml_new_tensor_2d(params_ctx, wtype, time_embed_dim, time_embed_dim);
+        time_embed_2_b     = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, time_embed_dim);
+
+        // SDXL
+        if (version == VERSION_XL) {
+            label_embed_0_w = ggml_new_tensor_2d(params_ctx, wtype, adm_in_channels, time_embed_dim);
+            label_embed_0_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, time_embed_dim);
+            label_embed_2_w = ggml_new_tensor_2d(params_ctx, wtype, time_embed_dim, time_embed_dim);
+            label_embed_2_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, time_embed_dim);
+        }
+
+        // input_blocks
+        input_block_0_w = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, in_channels, model_channels);
+        input_block_0_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, model_channels);
+
+        int ds           = 1;
+        size_t len_mults = channel_mult.size();
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                input_res_blocks[i][j].init_params(params_ctx, wtype);
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    input_transformers[i][j].init_params(params_ctx, alloc, wtype);
+                }
+            }
+            if (i != len_mults - 1) {
+                input_down_samples[i].init_params(params_ctx, wtype);
+                ds *= 2;
+            }
+        }
+
+        // middle_blocks
+        middle_block_0.init_params(params_ctx, wtype);
+        middle_block_1.init_params(params_ctx, alloc, wtype);
+        middle_block_2.init_params(params_ctx, wtype);
+
+        // output_blocks
+        for (int i = (int)len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                output_res_blocks[i][j].init_params(params_ctx, wtype);
+
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    output_transformers[i][j].init_params(params_ctx, alloc, wtype);
+                }
+
+                if (i > 0 && j == num_res_blocks) {
+                    output_up_samples[i - 1].init_params(params_ctx, wtype);
+
+                    ds /= 2;
+                }
+            }
+        }
+
+        // out
+        out_0_w = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, model_channels);
+        out_0_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, model_channels);
+
+        out_2_w = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 3, 3, model_channels, out_channels);
+        out_2_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, out_channels);
+
+        // 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);
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "time_embed.0.weight"] = time_embed_0_w;
+        tensors[prefix + "time_embed.0.bias"]   = time_embed_0_b;
+        tensors[prefix + "time_embed.2.weight"] = time_embed_2_w;
+        tensors[prefix + "time_embed.2.bias"]   = time_embed_2_b;
+
+        if (version == VERSION_XL) {
+            tensors[prefix + "label_emb.0.0.weight"] = label_embed_0_w;
+            tensors[prefix + "label_emb.0.0.bias"]   = label_embed_0_b;
+            tensors[prefix + "label_emb.0.2.weight"] = label_embed_2_w;
+            tensors[prefix + "label_emb.0.2.bias"]   = label_embed_2_b;
+        }
+
+        // input_blocks
+        tensors[prefix + "input_blocks.0.0.weight"] = input_block_0_w;
+        tensors[prefix + "input_blocks.0.0.bias"]   = input_block_0_b;
+
+        size_t len_mults    = channel_mult.size();
+        int input_block_idx = 0;
+        int ds              = 1;
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                input_block_idx += 1;
+                input_res_blocks[i][j].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".0.");
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    input_transformers[i][j].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".1.");
+                }
+            }
+            if (i != len_mults - 1) {
+                input_block_idx += 1;
+                input_down_samples[i].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".0.");
+                ds *= 2;
+            }
+        }
+
+        // middle_blocks
+        middle_block_0.map_by_name(tensors, prefix + "middle_block.0.");
+        middle_block_1.map_by_name(tensors, prefix + "middle_block.1.");
+        middle_block_2.map_by_name(tensors, prefix + "middle_block.2.");
+
+        // output_blocks
+        int output_block_idx = 0;
+        for (int i = (int)len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                output_res_blocks[i][j].map_by_name(tensors, prefix + "output_blocks." + std::to_string(output_block_idx) + ".0.");
+
+                int up_sample_idx = 1;
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    output_transformers[i][j].map_by_name(tensors, prefix + "output_blocks." + std::to_string(output_block_idx) + ".1.");
+                    up_sample_idx++;
+                }
+
+                if (i > 0 && j == num_res_blocks) {
+                    output_up_samples[i - 1].map_by_name(tensors, prefix + "output_blocks." + std::to_string(output_block_idx) + "." + std::to_string(up_sample_idx) + ".");
+
+                    ds /= 2;
+                }
+                output_block_idx += 1;
+            }
+        }
+
+        // out
+        tensors[prefix + "out.0.weight"] = out_0_w;
+        tensors[prefix + "out.0.bias"]   = out_0_b;
+        tensors[prefix + "out.2.weight"] = out_2_w;
+        tensors[prefix + "out.2.bias"]   = out_2_b;
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx0,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* timesteps,
+                                struct ggml_tensor* context,
+                                std::vector<struct ggml_tensor*> control,
+                                float control_net_strength,
+                                struct ggml_tensor* t_emb = NULL,
+                                struct ggml_tensor* y     = NULL) {
+        // x: [N, in_channels, h, w]
+        // timesteps: [N, ]
+        // t_emb: [N, model_channels]
+        // context: [N, max_position, hidden_size]([N, 77, 768])
+        // y: [adm_in_channels]
+        if (t_emb == NULL && timesteps != NULL) {
+            t_emb = new_timestep_embedding(ctx0, compute_allocr, timesteps, model_channels);  // [N, model_channels]
+        }
+
+        // time_embed = nn.Sequential
+        auto emb = ggml_nn_linear(ctx0, t_emb, time_embed_0_w, time_embed_0_b);
+        emb      = ggml_silu_inplace(ctx0, emb);
+        emb      = ggml_nn_linear(ctx0, emb, time_embed_2_w, time_embed_2_b);  // [N, time_embed_dim]
+
+        // SDXL
+        if (y != NULL) {
+            auto label_emb = ggml_nn_linear(ctx0, y, label_embed_0_w, label_embed_0_b);
+            label_emb      = ggml_silu_inplace(ctx0, label_emb);
+            label_emb      = ggml_nn_linear(ctx0, label_emb, label_embed_2_w, label_embed_2_b);
+            emb            = ggml_add(params_ctx, emb, label_emb);  // [N, time_embed_dim]
+        }
+
+        // input_blocks
+        std::vector<struct ggml_tensor*> hs;
+
+        // input block 0
+        struct ggml_tensor* h = ggml_nn_conv_2d(ctx0, x, input_block_0_w, input_block_0_b, 1, 1, 1, 1);  // [N, model_channels, h, w]
+
+        ggml_set_name(h, "bench-start");
+        hs.push_back(h);
+        // input block 1-11
+        size_t len_mults = channel_mult.size();
+        int ds           = 1;
+        for (int i = 0; i < len_mults; i++) {
+            int mult = channel_mult[i];
+            for (int j = 0; j < num_res_blocks; j++) {
+                h = input_res_blocks[i][j].forward(ctx0, h, emb);  // [N, mult*model_channels, h, w]
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    h = input_transformers[i][j].forward(ctx0, h, context);  // [N, mult*model_channels, h, w]
+                }
+                hs.push_back(h);
+            }
+            if (i != len_mults - 1) {
+                ds *= 2;
+                h = input_down_samples[i].forward(ctx0, h);  // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
+                hs.push_back(h);
+            }
+        }
+        // [N, 4*model_channels, h/8, w/8]
+
+        // middle_block
+        h = middle_block_0.forward(ctx0, h, emb);      // [N, 4*model_channels, h/8, w/8]
+        h = middle_block_1.forward(ctx0, h, context);  // [N, 4*model_channels, h/8, w/8]
+        h = middle_block_2.forward(ctx0, h, emb);      // [N, 4*model_channels, h/8, w/8]
+
+        if(control.size() > 0) {
+            auto cs = ggml_scale_inplace(ctx0, control[control.size() - 1], control_net_strength);
+            h = ggml_add(ctx0, h, cs); // middle control
+        }
+
+        int control_offset = control.size() - 2;
+        // output_blocks
+        for (int i = (int)len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                auto h_skip = hs.back();
+                hs.pop_back();
+
+                if(control.size() > 0) {
+                    auto cs = ggml_scale_inplace(ctx0,  control[control_offset], control_net_strength);
+                    h_skip = ggml_add(ctx0, h_skip, cs); // control net condition
+                    control_offset--;
+                }
+
+                h = ggml_concat(ctx0, h, h_skip);
+                h = output_res_blocks[i][j].forward(ctx0, h, emb);
+
+                if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
+                    h = output_transformers[i][j].forward(ctx0, h, context);
+                }
+
+                if (i > 0 && j == num_res_blocks) {
+                    h = output_up_samples[i - 1].forward(ctx0, h);
+
+                    ds /= 2;
+                }
+            }
+        }
+
+        // out
+        h = ggml_nn_group_norm(ctx0, h, out_0_w, out_0_b);
+        h = ggml_silu_inplace(ctx0, h);
+
+        // conv2d
+        h = ggml_nn_conv_2d(ctx0, h, out_2_w, out_2_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
+        ggml_set_name(h, "bench-end");
+        return h;
+    }
+
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                    struct ggml_tensor* timesteps,
+                                    struct ggml_tensor* context,
+                                    std::vector<struct ggml_tensor*> control,
+                                    struct ggml_tensor* t_emb = NULL,
+                                    struct ggml_tensor* y     = NULL,
+                                    float control_net_strength = 1.0) {
+        // since we are using ggml-alloc, this buffer only needs enough space to hold the ggml_tensor and ggml_cgraph structs, but not the tensor data
+        static size_t buf_size = ggml_tensor_overhead() * UNET_GRAPH_SIZE + ggml_graph_overhead();
+        static std::vector<uint8_t> buf(buf_size);
+
+        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, UNET_GRAPH_SIZE, false);
+
+        // temporal tensors for transfer tensors from cpu to gpu if needed
+        struct ggml_tensor* x_t         = NULL;
+        struct ggml_tensor* timesteps_t = NULL;
+        struct ggml_tensor* context_t   = NULL;
+        struct ggml_tensor* t_emb_t     = NULL;
+        struct ggml_tensor* y_t         = NULL;
+        std::vector<struct ggml_tensor*> control_t;
+
+        // it's performing a compute, check if backend isn't cpu
+        if (!ggml_backend_is_cpu(backend)) {
+            // pass input tensors to gpu memory
+            x_t       = ggml_dup_tensor(ctx0, x);
+            context_t = ggml_dup_tensor(ctx0, context);
+            ggml_allocr_alloc(compute_allocr, x_t);
+            if (timesteps != NULL) {
+                timesteps_t = ggml_dup_tensor(ctx0, timesteps);
+                ggml_allocr_alloc(compute_allocr, timesteps_t);
+            }
+            ggml_allocr_alloc(compute_allocr, context_t);
+            if (t_emb != NULL) {
+                t_emb_t = ggml_dup_tensor(ctx0, t_emb);
+                ggml_allocr_alloc(compute_allocr, t_emb_t);
+            }
+            if (y != NULL) {
+                y_t = ggml_dup_tensor(ctx0, y);
+                ggml_allocr_alloc(compute_allocr, y_t);
+            }
+            // pass data to device backend
+            if (!ggml_allocr_is_measure(compute_allocr)) {
+                ggml_backend_tensor_set(x_t, x->data, 0, ggml_nbytes(x));
+                ggml_backend_tensor_set(context_t, context->data, 0, ggml_nbytes(context));
+                if (timesteps_t != NULL) {
+                    ggml_backend_tensor_set(timesteps_t, timesteps->data, 0, ggml_nbytes(timesteps));
+                }
+                if (t_emb_t != NULL) {
+                    ggml_backend_tensor_set(t_emb_t, t_emb->data, 0, ggml_nbytes(t_emb));
+                }
+                if (y != NULL) {
+                    ggml_backend_tensor_set(y_t, y->data, 0, ggml_nbytes(y));
+                }
+            }
+        } else {
+            // if it's cpu backend just pass the same tensors
+            x_t         = x;
+            timesteps_t = timesteps;
+            context_t   = context;
+            t_emb_t     = t_emb;
+            y_t         = y;
+        }
+
+        // offload all controls tensors to gpu
+        if(control.size() > 0 && !ggml_backend_is_cpu(backend) && control[0]->backend != GGML_BACKEND_GPU) {
+            for(int i = 0; i < control.size(); i++) {
+                ggml_tensor* cntl_t = ggml_dup_tensor(ctx0, control[i]);
+                control_t.push_back(cntl_t);
+                ggml_allocr_alloc(compute_allocr, cntl_t);
+                if(!ggml_allocr_is_measure(compute_allocr)) {
+                    ggml_backend_tensor_copy(control[i], control_t[i]);
+                    ggml_backend_synchronize(backend);
+                }
+            }
+        } else {
+            control_t = control;
+        }
+
+        struct ggml_tensor* out = forward(ctx0, x_t, timesteps_t, context_t, control_t, control_net_strength, t_emb_t, y_t);
+
+        ggml_build_forward_expand(gf, out);
+        ggml_free(ctx0);
+
+        return gf;
+    }
+
+    void alloc_compute_buffer(struct ggml_tensor* x,
+                              struct ggml_tensor* context,
+                              std::vector<struct ggml_tensor*> control,
+                              struct ggml_tensor* t_emb = NULL,
+                              struct ggml_tensor* y     = NULL,
+                              float control_net_strength = 1.0) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(x, NULL, context, control, t_emb, y, control_net_strength);
+        };
+        GGMLModule::alloc_compute_buffer(get_graph);
+    }
+
+    void compute(struct ggml_tensor* work_latent,
+                 int n_threads,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* timesteps,
+                 struct ggml_tensor* context,
+                 std::vector<struct ggml_tensor*> control,
+                 float control_net_strength,
+                 struct ggml_tensor* t_emb = NULL,
+                 struct ggml_tensor* y     = NULL) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(x, timesteps, context, control, t_emb, y, control_net_strength);
+        };
+
+        GGMLModule::compute(get_graph, n_threads, work_latent);
+    }
+};
+
+#endif  // __UNET_HPP__

upscaler.cpp

@@ -0,0 +1,126 @@
+#include "esrgan.hpp"
+#include "ggml_extend.hpp"
+#include "model.h"
+#include "stable-diffusion.h"
+
+struct UpscalerGGML {
+    ggml_backend_t backend    = NULL;  // general backend
+    ggml_type model_data_type = GGML_TYPE_F16;
+    ESRGAN esrgan_upscaler;
+    std::string esrgan_path;
+    int n_threads;
+
+    UpscalerGGML(int n_threads)
+        : n_threads(n_threads) {
+    }
+
+    bool load_from_file(const std::string& esrgan_path) {
+#ifdef SD_USE_CUBLAS
+        LOG_DEBUG("Using CUDA backend");
+        backend = ggml_backend_cuda_init(0);
+#endif
+#ifdef SD_USE_METAL
+        LOG_DEBUG("Using Metal backend");
+        ggml_metal_log_set_callback(ggml_log_callback_default, nullptr);
+        backend = ggml_backend_metal_init();
+#endif
+
+        if (!backend) {
+            LOG_DEBUG("Using CPU backend");
+            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)) {
+            return false;
+        }
+        return true;
+    }
+
+    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;
+        LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
+                 input_image.width, input_image.height, output_width, output_height);
+
+        struct ggml_init_params params;
+        params.mem_size = output_width * output_height * 3 * sizeof(float) * 2;
+        params.mem_size += 2 * ggml_tensor_overhead();
+        params.mem_buffer = NULL;
+        params.no_alloc   = false;
+
+        // draft context
+        struct ggml_context* upscale_ctx = ggml_init(params);
+        if (!upscale_ctx) {
+            LOG_ERROR("ggml_init() failed");
+            return upscaled_image;
+        }
+        LOG_DEBUG("upscale work buffer size: %.2f MB", params.mem_size / 1024.f / 1024.f);
+        ggml_tensor* input_image_tensor = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, input_image.width, input_image.height, 3, 1);
+        sd_image_to_tensor(input_image.data, input_image_tensor);
+
+        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);
+            }
+        };
+        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();
+        ggml_tensor_clamp(upscaled, 0.f, 1.f);
+        uint8_t* upscaled_data = sd_tensor_to_image(upscaled);
+        ggml_free(upscale_ctx);
+        int64_t t3 = ggml_time_ms();
+        LOG_INFO("input_image_tensor upscaled, taking %.2fs", (t3 - t0) / 1000.0f);
+        upscaled_image = {
+            (uint32_t)output_width,
+            (uint32_t)output_height,
+            3,
+            upscaled_data,
+        };
+        return upscaled_image;
+    }
+};
+
+struct upscaler_ctx_t {
+    UpscalerGGML* upscaler = NULL;
+};
+
+upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
+                                 int n_threads,
+                                 enum sd_type_t wtype) {
+    upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t));
+    if (upscaler_ctx == NULL) {
+        return NULL;
+    }
+    std::string esrgan_path(esrgan_path_c_str);
+
+    upscaler_ctx->upscaler = new UpscalerGGML(n_threads);
+    if (upscaler_ctx->upscaler == NULL) {
+        return NULL;
+    }
+
+    if (!upscaler_ctx->upscaler->load_from_file(esrgan_path)) {
+        delete upscaler_ctx->upscaler;
+        upscaler_ctx->upscaler = NULL;
+        free(upscaler_ctx);
+        return NULL;
+    }
+    return upscaler_ctx;
+}
+
+sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor) {
+    return upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
+}
+
+void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) {
+    if (upscaler_ctx->upscaler != NULL) {
+        delete upscaler_ctx->upscaler;
+        upscaler_ctx->upscaler = NULL;
+    }
+    free(upscaler_ctx);
+}

util.cpp

@@ -1,7 +1,11 @@
 #include "util.h"
-
+#include <algorithm>
 #include <stdarg.h>
+#include <codecvt>
 #include <fstream>
+#include <locale>
+#include <sstream>
+#include <string>
 #include <thread>
 #include <unordered_set>
 #include <vector>
@@ -16,6 +20,9 @@
 #include <unistd.h>
 #endif
 
+#include "ggml/ggml.h"
+#include "stable-diffusion.h"
+
 bool ends_with(const std::string& str, const std::string& ending) {
     if (str.length() >= ending.length()) {
         return (str.compare(str.length() - ending.length(), ending.length(), ending) == 0);
@@ -65,6 +72,20 @@ bool is_directory(const std::string& path) {
     return (attributes != INVALID_FILE_ATTRIBUTES && (attributes & FILE_ATTRIBUTE_DIRECTORY));
 }
 
+std::string get_full_path(const std::string& dir, const std::string& filename) {
+    std::string full_path = dir + "\\" + filename;
+
+    WIN32_FIND_DATA find_file_data;
+    HANDLE hFind = FindFirstFile(full_path.c_str(), &find_file_data);
+
+    if (hFind != INVALID_HANDLE_VALUE) {
+        FindClose(hFind);
+        return full_path;
+    } else {
+        return "";
+    }
+}
+
 #else  // Unix
 #include <dirent.h>
 #include <sys/stat.h>
@@ -79,6 +100,25 @@ bool is_directory(const std::string& path) {
     return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
+std::string get_full_path(const std::string& dir, const std::string& filename) {
+    DIR* dp = opendir(dir.c_str());
+
+    if (dp != nullptr) {
+        struct dirent* entry;
+
+        while ((entry = readdir(dp)) != nullptr) {
+            if (strcasecmp(entry->d_name, filename.c_str()) == 0) {
+                closedir(dp);
+                return dir + "/" + entry->d_name;
+            }
+        }
+
+        closedir(dp);
+    }
+
+    return "";
+}
+
 #endif
 
 // get_num_physical_cores is copy from
@@ -119,7 +159,22 @@ int32_t get_num_physical_cores() {
     return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
 }
 
-std::string basename(const std::string& path) {
+std::u32string utf8_to_utf32(const std::string& utf8_str) {
+    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
+    return converter.from_bytes(utf8_str);
+}
+
+std::string utf32_to_utf8(const std::u32string& utf32_str) {
+    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
+    return converter.to_bytes(utf32_str);
+}
+
+std::u32string unicode_value_to_utf32(int unicode_value) {
+    std::u32string utf32_string = {static_cast<char32_t>(unicode_value)};
+    return utf32_string;
+}
+
+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);
@@ -147,40 +202,108 @@ std::string path_join(const std::string& p1, const std::string& p2) {
     return p1 + "/" + p2;
 }
 
-static SDLogLevel log_level = SDLogLevel::INFO;
-
-void set_sd_log_level(SDLogLevel level) {
-    log_level = level;
+void pretty_progress(int step, int steps, float time) {
+    std::string progress = "  |";
+    int max_progress     = 50;
+    int32_t current      = (int32_t)(step * 1.f * max_progress / steps);
+    for (int i = 0; i < 50; i++) {
+        if (i > current) {
+            progress += " ";
+        } else if (i == current && i != max_progress - 1) {
+            progress += ">";
+        } else {
+            progress += "=";
+        }
+    }
+    progress += "|";
+    printf(time > 1.0f ? "\r%s %i/%i - %.2fs/it" : "\r%s %i/%i - %.2fit/s",
+           progress.c_str(), step, steps,
+           time > 1.0f || time == 0 ? time : (1.0f / time));
+    fflush(stdout);  // for linux
+    if (step == steps) {
+        printf("\n");
+    }
 }
 
-void log_printf(SDLogLevel level, const char* file, int line, const char* format, ...) {
-    if (level < log_level) {
-        return;
-    }
+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;
+
+#define LOG_BUFFER_SIZE 1024
+
+void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
     va_list args;
     va_start(args, format);
 
-    if (level == SDLogLevel::DEBUG) {
-        printf("[DEBUG] %s:%-4d - ", basename(file).c_str(), line);
-        vprintf(format, args);
-        printf("\n");
-        fflush(stdout);
-    } else if (level == SDLogLevel::INFO) {
-        printf("[INFO]  %s:%-4d - ", basename(file).c_str(), line);
-        vprintf(format, args);
-        printf("\n");
-        fflush(stdout);
-    } else if (level == SDLogLevel::WARN) {
-        fprintf(stdout, "[WARN]  %s:%-4d - ", basename(file).c_str(), line);
-        vfprintf(stdout, format, args);
-        fprintf(stdout, "\n");
-        fflush(stdout);
-    } else {
-        fprintf(stderr, "[ERROR] %s:%-4d - ", basename(file).c_str(), line);
-        vfprintf(stderr, format, args);
-        fprintf(stderr, "\n");
-        fflush(stderr);
+    const char* level_str = "DEBUG";
+    if (level == SD_LOG_INFO) {
+        level_str = "INFO ";
+    } else if (level == SD_LOG_WARN) {
+        level_str = "WARN ";
+    } else if (level == SD_LOG_ERROR) {
+        level_str = "ERROR";
+    }
+
+    static char log_buffer[LOG_BUFFER_SIZE];
+
+    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "[%s] %s:%-4d - ", level_str, sd_basename(file).c_str(), line);
+
+    if (written >= 0 && written < LOG_BUFFER_SIZE) {
+        vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
+        strncat(log_buffer, "\n", LOG_BUFFER_SIZE - strlen(log_buffer) - 1);
+    }
+
+    if (sd_log_cb) {
+        sd_log_cb(level, log_buffer, sd_log_cb_data);
     }
 
     va_end(args);
 }
+
+void sd_set_log_callback(sd_log_cb_t cb, void* data) {
+    sd_log_cb      = cb;
+    sd_log_cb_data = data;
+}
+
+const char* sd_get_system_info() {
+    static char buffer[1024];
+    std::stringstream ss;
+    ss << "System Info: \n";
+    ss << "    BLAS = " << ggml_cpu_has_blas() << std::endl;
+    ss << "    SSE3 = " << ggml_cpu_has_sse3() << std::endl;
+    ss << "    AVX = " << ggml_cpu_has_avx() << std::endl;
+    ss << "    AVX2 = " << ggml_cpu_has_avx2() << std::endl;
+    ss << "    AVX512 = " << ggml_cpu_has_avx512() << std::endl;
+    ss << "    AVX512_VBMI = " << ggml_cpu_has_avx512_vbmi() << std::endl;
+    ss << "    AVX512_VNNI = " << ggml_cpu_has_avx512_vnni() << std::endl;
+    ss << "    FMA = " << ggml_cpu_has_fma() << std::endl;
+    ss << "    NEON = " << ggml_cpu_has_neon() << std::endl;
+    ss << "    ARM_FMA = " << ggml_cpu_has_arm_fma() << std::endl;
+    ss << "    F16C = " << ggml_cpu_has_f16c() << std::endl;
+    ss << "    FP16_VA = " << ggml_cpu_has_fp16_va() << std::endl;
+    ss << "    WASM_SIMD = " << ggml_cpu_has_wasm_simd() << std::endl;
+    ss << "    VSX = " << ggml_cpu_has_vsx() << std::endl;
+    snprintf(buffer, sizeof(buffer), "%s", ss.str().c_str());
+    return buffer;
+}
+
+const char* sd_type_name(enum sd_type_t type) {
+    return ggml_type_name((ggml_type)type);
+}

util.h

@@ -1,8 +1,11 @@
 #ifndef __UTIL_H__
 #define __UTIL_H__
 
+#include <cstdint>
 #include <string>
 
+#include "stable-diffusion.h"
+
 bool ends_with(const std::string& str, const std::string& ending);
 bool starts_with(const std::string& str, const std::string& start);
 
@@ -12,26 +15,24 @@ 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::string basename(const std::string& path);
+std::u32string utf8_to_utf32(const std::string& utf8_str);
+std::string utf32_to_utf8(const std::u32string& utf32_str);
+std::u32string unicode_value_to_utf32(int unicode_value);
+
+std::string sd_basename(const std::string& path);
 
 std::string path_join(const std::string& p1, const std::string& p2);
 
-int32_t get_num_physical_cores();
+void pretty_progress(int step, int steps, float time);
 
-enum SDLogLevel {
-    DEBUG,
-    INFO,
-    WARN,
-    ERROR
-};
+void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
 
-void set_sd_log_level(SDLogLevel level);
+std::string trim(const std::string& s);
 
-void log_printf(SDLogLevel level, const char* file, int line, const char* format, ...);
-
-#define LOG_DEBUG(format, ...) log_printf(SDLogLevel::DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#define LOG_INFO(format, ...) log_printf(SDLogLevel::INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#define LOG_WARN(format, ...) log_printf(SDLogLevel::WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#define LOG_ERROR(format, ...) log_printf(SDLogLevel::ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#endif  // __UTIL_H__
+#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__)
+#define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#endif  // __UTIL_H__

vae.hpp

@@ -0,0 +1,740 @@
+#ifndef __VAE_HPP__
+#define __VAE_HPP__
+
+#include "common.hpp"
+#include "ggml_extend.hpp"
+
+/*================================================== AutoEncoderKL ===================================================*/
+
+#define VAE_GRAPH_SIZE 10240
+
+struct ResnetBlock {
+    // network hparams
+    int in_channels;
+    int out_channels;
+
+    // network params
+    struct ggml_tensor* norm1_w;  // [in_channels, ]
+    struct ggml_tensor* norm1_b;  // [in_channels, ]
+
+    struct ggml_tensor* conv1_w;  // [out_channels, in_channels, 3, 3]
+    struct ggml_tensor* conv1_b;  // [out_channels, ]
+
+    struct ggml_tensor* norm2_w;  // [out_channels, ]
+    struct ggml_tensor* norm2_b;  // [out_channels, ]
+
+    struct ggml_tensor* conv2_w;  // [out_channels, out_channels, 3, 3]
+    struct ggml_tensor* conv2_b;  // [out_channels, ]
+
+    // nin_shortcut, only if out_channels != in_channels
+    struct ggml_tensor* nin_shortcut_w;  // [out_channels, in_channels, 1, 1]
+    struct ggml_tensor* nin_shortcut_b;  // [out_channels, ]
+
+    size_t calculate_mem_size(ggml_type wtype) {
+        double mem_size = 0;
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, in_channels);                      // norm1_w/b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * in_channels * 3 * 3);   // conv1_w
+        mem_size += 4 * ggml_row_size(GGML_TYPE_F32, out_channels);                     // conv1_b/norm2_w/norm2_b/conv2_b
+        mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * out_channels * 3 * 3);  // conv2_w
+
+        if (out_channels != in_channels) {
+            mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * in_channels * 1 * 1);  // nin_shortcut_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, out_channels);                        // nin_shortcut_b
+        }
+        return static_cast<size_t>(mem_size);
+    }
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        norm1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        norm1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        conv1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, in_channels, out_channels);
+        conv1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+
+        norm2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        norm2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        conv2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels);
+        conv2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+
+        if (out_channels != in_channels) {
+            nin_shortcut_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, out_channels);
+            nin_shortcut_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 + "norm1.weight"] = norm1_w;
+        tensors[prefix + "norm1.bias"]   = norm1_b;
+        tensors[prefix + "conv1.weight"] = conv1_w;
+        tensors[prefix + "conv1.bias"]   = conv1_b;
+
+        tensors[prefix + "norm2.weight"] = norm2_w;
+        tensors[prefix + "norm2.bias"]   = norm2_b;
+        tensors[prefix + "conv2.weight"] = conv2_w;
+        tensors[prefix + "conv2.bias"]   = conv2_b;
+
+        if (out_channels != in_channels) {
+            tensors[prefix + "nin_shortcut.weight"] = nin_shortcut_w;
+            tensors[prefix + "nin_shortcut.bias"]   = nin_shortcut_b;
+        }
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* z) {
+        // z: [N, in_channels, h, w]
+
+        auto h = ggml_nn_group_norm(ctx, z, norm1_w, norm1_b);
+        h      = ggml_silu_inplace(ctx, h);
+        h      = ggml_nn_conv_2d(ctx, h, conv1_w, conv1_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
+        h      = ggml_nn_group_norm(ctx, h, norm2_w, norm2_b);
+        h      = ggml_silu_inplace(ctx, h);
+        // dropout, skip for inference
+        h = ggml_nn_conv_2d(ctx, h, conv2_w, conv2_b, 1, 1, 1, 1);  // [N, out_channels, h, w]
+
+        // skip connection
+        if (out_channels != in_channels) {
+            z = ggml_nn_conv_2d(ctx, z, nin_shortcut_w, nin_shortcut_b);  // [N, out_channels, h, w]
+        }
+
+        h = ggml_add(ctx, h, z);
+        return h;  // [N, out_channels, h, w]
+    }
+};
+
+struct AttnBlock {
+    int in_channels;  // mult * model_channels
+
+    // group norm
+    struct ggml_tensor* norm_w;  // [in_channels,]
+    struct ggml_tensor* norm_b;  // [in_channels,]
+
+    // q/k/v
+    struct ggml_tensor* q_w;  // [in_channels, in_channels, 1, 1]
+    struct ggml_tensor* q_b;  // [in_channels,]
+    struct ggml_tensor* k_w;  // [in_channels, in_channels, 1, 1]
+    struct ggml_tensor* k_b;  // [in_channels,]
+    struct ggml_tensor* v_w;  // [in_channels, in_channels, 1, 1]
+    struct ggml_tensor* v_b;  // [in_channels,]
+
+    // proj_out
+    struct ggml_tensor* proj_out_w;  // [in_channels, in_channels, 1, 1]
+    struct ggml_tensor* proj_out_b;  // [in_channels,]
+
+    size_t calculate_mem_size(ggml_type wtype) {
+        double mem_size = 0;
+        mem_size += 6 * ggml_row_size(GGML_TYPE_F32, in_channels);                        // norm_w/norm_b/q_b/k_v/v_b/proj_out_b
+        mem_size += 4 * ggml_row_size(GGML_TYPE_F16,  in_channels * in_channels * 1 * 1);  // q_w/k_w/v_w/proj_out_w                                            // object overhead
+        return static_cast<size_t>(mem_size);
+    }
+
+    void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
+        norm_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        norm_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+        q_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
+        q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        k_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
+        k_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+        v_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
+        v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+
+        proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
+        proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "norm.weight"]     = norm_w;
+        tensors[prefix + "norm.bias"]       = norm_b;
+        tensors[prefix + "q.weight"]        = q_w;
+        tensors[prefix + "q.bias"]          = q_b;
+        tensors[prefix + "k.weight"]        = k_w;
+        tensors[prefix + "k.bias"]          = k_b;
+        tensors[prefix + "v.weight"]        = v_w;
+        tensors[prefix + "v.bias"]          = v_b;
+        tensors[prefix + "proj_out.weight"] = proj_out_w;
+        tensors[prefix + "proj_out.bias"]   = proj_out_b;
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [N, in_channels, h, w]
+
+        auto h_ = ggml_nn_group_norm(ctx, x, norm_w, norm_b);
+
+        const int64_t n = h_->ne[3];
+        const int64_t c = h_->ne[2];
+        const int64_t h = h_->ne[1];
+        const int64_t w = h_->ne[0];
+
+        auto q = ggml_nn_conv_2d(ctx, h_, q_w, q_b);  // [N, in_channels, h, w]
+        auto k = ggml_nn_conv_2d(ctx, h_, k_w, k_b);  // [N, in_channels, h, w]
+        auto v = ggml_nn_conv_2d(ctx, h_, v_w, v_b);  // [N, in_channels, h, w]
+
+        q = ggml_cont(ctx, ggml_permute(ctx, q, 1, 2, 0, 3));  // [N, h, w, in_channels]
+        q = ggml_reshape_3d(ctx, q, c, h * w, n);              // [N, h * w, in_channels]
+
+        k = ggml_cont(ctx, ggml_permute(ctx, k, 1, 2, 0, 3));  // [N, h, w, in_channels]
+        k = ggml_reshape_3d(ctx, k, c, h * w, n);              // [N, h * w, in_channels]
+
+        auto w_ = ggml_mul_mat(ctx, k, q);  // [N, h * w, h * w]
+        w_      = ggml_scale_inplace(ctx, w_, 1.0f / sqrt((float)in_channels));
+        w_      = ggml_soft_max_inplace(ctx, w_);
+
+        v  = ggml_reshape_3d(ctx, v, h * w, c, n);               // [N, in_channels, h * w]
+        h_ = ggml_mul_mat(ctx, v, w_);                           // [N, h * w, in_channels]
+        h_ = ggml_cont(ctx, ggml_permute(ctx, h_, 1, 0, 2, 3));  // [N, in_channels, h * w]
+        h_ = ggml_reshape_4d(ctx, h_, w, h, c, n);               // [N, in_channels, h, w]
+
+        // proj_out
+        h_ = ggml_nn_conv_2d(ctx, h_, proj_out_w, proj_out_b);  // [N, in_channels, h, w]
+
+        h_ = ggml_add(ctx, h_, x);
+        return h_;
+    }
+};
+
+// ldm.modules.diffusionmodules.model.Encoder
+struct Encoder {
+    int embed_dim      = 4;
+    int ch             = 128;
+    int z_channels     = 4;
+    int in_channels    = 3;
+    int num_res_blocks = 2;
+    int ch_mult[4]     = {1, 2, 4, 4};
+
+    struct ggml_tensor* conv_in_w;  // [ch, in_channels, 3, 3]
+    struct ggml_tensor* conv_in_b;  // [ch, ]
+
+    ResnetBlock down_blocks[4][2];
+    DownSample down_samples[3];
+
+    struct
+    {
+        ResnetBlock block_1;
+        AttnBlock attn_1;
+        ResnetBlock block_2;
+    } mid;
+
+    // block_in = ch * ch_mult[len_mults - 1]
+    struct ggml_tensor* norm_out_w;  // [block_in, ]
+    struct ggml_tensor* norm_out_b;  // [block_in, ]
+
+    struct ggml_tensor* conv_out_w;  // [embed_dim*2, block_in, 3, 3]
+    struct ggml_tensor* conv_out_b;  // [embed_dim*2, ]
+
+    Encoder() {
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+
+        int block_in = 1;
+        for (int i = 0; i < len_mults; i++) {
+            if (i == 0) {
+                block_in = ch;
+            } else {
+                block_in = ch * ch_mult[i - 1];
+            }
+            int block_out = ch * ch_mult[i];
+            for (int j = 0; j < num_res_blocks; j++) {
+                down_blocks[i][j].in_channels  = block_in;
+                down_blocks[i][j].out_channels = block_out;
+                block_in                       = block_out;
+            }
+            if (i != len_mults - 1) {
+                down_samples[i].channels       = block_in;
+                down_samples[i].out_channels   = block_in;
+                down_samples[i].vae_downsample = true;
+            }
+        }
+
+        mid.block_1.in_channels  = block_in;
+        mid.block_1.out_channels = block_in;
+        mid.attn_1.in_channels   = block_in;
+        mid.block_2.in_channels  = block_in;
+        mid.block_2.out_channels = block_in;
+    }
+
+    size_t get_num_tensors() {
+        int num_tensors = 6;
+
+        // mid
+        num_tensors += 10 * 3;
+
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                num_tensors += 10;
+            }
+
+            if (i != 0) {
+                num_tensors += 2;
+            }
+        }
+        return num_tensors;
+    }
+
+    size_t calculate_mem_size(ggml_type wtype) {
+        size_t mem_size = 0;
+        int len_mults   = sizeof(ch_mult) / sizeof(int);
+        int block_in    = ch * ch_mult[len_mults - 1];
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, ch * in_channels * 3 * 3);  // conv_in_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, ch);                        // conv_in_b
+
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, block_in);  // norm_out_w/b
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, z_channels * 2 * block_in * 3 * 3);  // conv_out_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, z_channels * 2);                     // conv_out_b
+
+        mem_size += mid.block_1.calculate_mem_size(wtype);
+        mem_size += mid.attn_1.calculate_mem_size(wtype);
+        mem_size += mid.block_2.calculate_mem_size(wtype);
+
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                mem_size += down_blocks[i][j].calculate_mem_size(wtype);
+            }
+            if (i != 0) {
+                mem_size += down_samples[i - 1].calculate_mem_size(wtype);
+            }
+        }
+
+        return mem_size;
+    }
+
+    void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        int block_in  = ch * ch_mult[len_mults - 1];
+
+        conv_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, in_channels, ch);
+        conv_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ch);
+
+        norm_out_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, block_in);
+        norm_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, block_in);
+
+        conv_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, block_in, z_channels * 2);
+        conv_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_channels * 2);
+
+        mid.block_1.init_params(ctx, wtype);
+        mid.attn_1.init_params(ctx, alloc, wtype);
+        mid.block_2.init_params(ctx, wtype);
+
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                down_blocks[i][j].init_params(ctx, wtype);
+            }
+            if (i != len_mults - 1) {
+                down_samples[i].init_params(ctx, wtype);
+            }
+        }
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "norm_out.weight"] = norm_out_w;
+        tensors[prefix + "norm_out.bias"]   = norm_out_b;
+        tensors[prefix + "conv_in.weight"]  = conv_in_w;
+        tensors[prefix + "conv_in.bias"]    = conv_in_b;
+        tensors[prefix + "conv_out.weight"] = conv_out_w;
+        tensors[prefix + "conv_out.bias"]   = conv_out_b;
+
+        mid.block_1.map_by_name(tensors, prefix + "mid.block_1.");
+        mid.attn_1.map_by_name(tensors, prefix + "mid.attn_1.");
+        mid.block_2.map_by_name(tensors, prefix + "mid.block_2.");
+
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                down_blocks[i][j].map_by_name(tensors, prefix + "down." + std::to_string(i) + ".block." + std::to_string(j) + ".");
+            }
+            if (i != len_mults - 1) {
+                down_samples[i].map_by_name(tensors, prefix + "down." + std::to_string(i) + ".downsample.");
+            }
+        }
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [N, in_channels, h, w]
+
+        // conv_in
+        auto h = ggml_nn_conv_2d(ctx, x, conv_in_w, conv_in_b, 1, 1, 1, 1);  // [N, ch, h, w]
+        ggml_set_name(h, "b-start");
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        for (int i = 0; i < len_mults; i++) {
+            for (int j = 0; j < num_res_blocks; j++) {
+                h = down_blocks[i][j].forward(ctx, h);
+            }
+            if (i != len_mults - 1) {
+                h = down_samples[i].forward(ctx, h);
+            }
+        }
+
+        h = mid.block_1.forward(ctx, h);
+        h = mid.attn_1.forward(ctx, h);
+        h = mid.block_2.forward(ctx, h);  // [N, block_in, h, w]
+
+        h = ggml_nn_group_norm(ctx, h, norm_out_w, norm_out_b);
+        h = ggml_silu_inplace(ctx, h);
+
+        // conv_out
+        h = ggml_nn_conv_2d(ctx, h, conv_out_w, conv_out_b, 1, 1, 1, 1);  // [N, z_channels*2, h, w]
+
+        return h;
+    }
+};
+
+// ldm.modules.diffusionmodules.model.Decoder
+struct Decoder {
+    int embed_dim      = 4;
+    int ch             = 128;
+    int z_channels     = 4;
+    int out_ch         = 3;
+    int num_res_blocks = 2;
+    int ch_mult[4]     = {1, 2, 4, 4};
+
+    // block_in = ch *  ch_mult[-1], 512
+    struct ggml_tensor* conv_in_w;  // [block_in, z_channels, 3, 3]
+    struct ggml_tensor* conv_in_b;  // [block_in, ]
+
+    struct
+    {
+        ResnetBlock block_1;
+        AttnBlock attn_1;
+        ResnetBlock block_2;
+    } mid;
+
+    ResnetBlock up_blocks[4][3];
+    UpSample up_samples[3];
+
+    struct ggml_tensor* norm_out_w;  // [ch *  ch_mult[0], ]
+    struct ggml_tensor* norm_out_b;  // [ch *  ch_mult[0], ]
+
+    struct ggml_tensor* conv_out_w;  // [out_ch, ch *  ch_mult[0], 3, 3]
+    struct ggml_tensor* conv_out_b;  // [out_ch, ]
+
+    Decoder() {
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        int block_in  = ch * ch_mult[len_mults - 1];
+
+        mid.block_1.in_channels  = block_in;
+        mid.block_1.out_channels = block_in;
+        mid.attn_1.in_channels   = block_in;
+        mid.block_2.in_channels  = block_in;
+        mid.block_2.out_channels = block_in;
+
+        for (int i = len_mults - 1; i >= 0; i--) {
+            int mult      = ch_mult[i];
+            int block_out = ch * mult;
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                up_blocks[i][j].in_channels  = block_in;
+                up_blocks[i][j].out_channels = block_out;
+                block_in                     = block_out;
+            }
+            if (i != 0) {
+                up_samples[i - 1].channels     = block_in;
+                up_samples[i - 1].out_channels = block_in;
+            }
+        }
+    }
+
+    size_t calculate_mem_size(ggml_type wtype) {
+        double mem_size = 0;
+        int len_mults   = sizeof(ch_mult) / sizeof(int);
+        int block_in    = ch * ch_mult[len_mults - 1];
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, block_in * z_channels * 3 * 3);  // conv_in_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, block_in);                       // conv_in_b
+
+        mem_size += 2 * ggml_row_size(GGML_TYPE_F32, (ch * ch_mult[0]));  // norm_out_w/b
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, (ch * ch_mult[0]) * out_ch * 3 * 3);  // conv_out_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, out_ch);                              // conv_out_b
+
+        mem_size += mid.block_1.calculate_mem_size(wtype);
+        mem_size += mid.attn_1.calculate_mem_size(wtype);
+        mem_size += mid.block_2.calculate_mem_size(wtype);
+
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                mem_size += up_blocks[i][j].calculate_mem_size(wtype);
+            }
+            if (i != 0) {
+                mem_size += up_samples[i - 1].calculate_mem_size(wtype);
+            }
+        }
+
+        return static_cast<size_t>(mem_size);
+    }
+
+    size_t get_num_tensors() {
+        int num_tensors = 8;
+
+        // mid
+        num_tensors += 10 * 3;
+
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                num_tensors += 10;
+            }
+
+            if (i != 0) {
+                num_tensors += 2;
+            }
+        }
+        return num_tensors;
+    }
+
+    void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        int block_in  = ch * ch_mult[len_mults - 1];
+
+        norm_out_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ch * ch_mult[0]);
+        norm_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ch * ch_mult[0]);
+
+        conv_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, z_channels, block_in);
+        conv_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, block_in);
+
+        conv_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, ch * ch_mult[0], out_ch);
+        conv_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_ch);
+
+        mid.block_1.init_params(ctx, wtype);
+        mid.attn_1.init_params(ctx, alloc, wtype);
+        mid.block_2.init_params(ctx, wtype);
+
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                up_blocks[i][j].init_params(ctx, wtype);
+            }
+
+            if (i != 0) {
+                up_samples[i - 1].init_params(ctx, wtype);
+            }
+        }
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "norm_out.weight"] = norm_out_w;
+        tensors[prefix + "norm_out.bias"]   = norm_out_b;
+        tensors[prefix + "conv_in.weight"]  = conv_in_w;
+        tensors[prefix + "conv_in.bias"]    = conv_in_b;
+        tensors[prefix + "conv_out.weight"] = conv_out_w;
+        tensors[prefix + "conv_out.bias"]   = conv_out_b;
+
+        mid.block_1.map_by_name(tensors, prefix + "mid.block_1.");
+        mid.attn_1.map_by_name(tensors, prefix + "mid.attn_1.");
+        mid.block_2.map_by_name(tensors, prefix + "mid.block_2.");
+
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                up_blocks[i][j].map_by_name(tensors, prefix + "up." + std::to_string(i) + ".block." + std::to_string(j) + ".");
+            }
+            if (i != 0) {
+                up_samples[i - 1].map_by_name(tensors, prefix + "up." + std::to_string(i) + ".upsample.");
+            }
+        }
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* z) {
+        // z: [N, z_channels, h, w]
+        // conv_in
+        auto h = ggml_nn_conv_2d(ctx, z, conv_in_w, conv_in_b, 1, 1, 1, 1);  // [N, block_in, h, w]
+
+        h = mid.block_1.forward(ctx, h);
+        h = mid.attn_1.forward(ctx, h);
+        h = mid.block_2.forward(ctx, h);  // [N, block_in, h, w]
+
+        int len_mults = sizeof(ch_mult) / sizeof(int);
+        for (int i = len_mults - 1; i >= 0; i--) {
+            for (int j = 0; j < num_res_blocks + 1; j++) {
+                h = up_blocks[i][j].forward(ctx, h);
+            }
+            if (i != 0) {
+                h = up_samples[i - 1].forward(ctx, h);
+            }
+        }
+
+        // group norm 32
+        h = ggml_nn_group_norm(ctx, h, norm_out_w, norm_out_b);
+        h = ggml_silu_inplace(ctx, h);
+
+        // conv_out
+        h = ggml_nn_conv_2d(ctx, h, conv_out_w, conv_out_b, 1, 1, 1, 1);  // [N, out_ch, h, w]
+        return h;
+    }
+};
+
+// ldm.models.autoencoder.AutoencoderKL
+struct AutoEncoderKL : public GGMLModule {
+    bool decode_only = true;
+    int embed_dim    = 4;
+    struct {
+        int z_channels     = 4;
+        int resolution     = 256;
+        int in_channels    = 3;
+        int out_ch         = 3;
+        int ch             = 128;
+        int ch_mult[4]     = {1, 2, 4, 4};
+        int num_res_blocks = 2;
+    } dd_config;
+
+    struct ggml_tensor* quant_conv_w;  // [2*embed_dim, 2*z_channels, 1, 1]
+    struct ggml_tensor* quant_conv_b;  // [2*embed_dim, ]
+
+    struct ggml_tensor* post_quant_conv_w;  // [z_channels, embed_dim, 1, 1]
+    struct ggml_tensor* post_quant_conv_b;  // [z_channels, ]
+
+    Encoder encoder;
+    Decoder decoder;
+
+    AutoEncoderKL(bool decode_only = false)
+        : decode_only(decode_only) {
+        name = "vae";
+        assert(sizeof(dd_config.ch_mult) == sizeof(encoder.ch_mult));
+        assert(sizeof(dd_config.ch_mult) == sizeof(decoder.ch_mult));
+
+        encoder.embed_dim      = embed_dim;
+        decoder.embed_dim      = embed_dim;
+        encoder.ch             = dd_config.ch;
+        decoder.ch             = dd_config.ch;
+        encoder.z_channels     = dd_config.z_channels;
+        decoder.z_channels     = dd_config.z_channels;
+        encoder.in_channels    = dd_config.in_channels;
+        decoder.out_ch         = dd_config.out_ch;
+        encoder.num_res_blocks = dd_config.num_res_blocks;
+
+        int len_mults = sizeof(dd_config.ch_mult) / sizeof(int);
+        for (int i = 0; i < len_mults; i++) {
+            encoder.ch_mult[i] = dd_config.ch_mult[i];
+            decoder.ch_mult[i] = dd_config.ch_mult[i];
+        }
+    }
+
+    size_t calculate_mem_size() {
+        size_t mem_size = 0;
+
+        if (!decode_only) {
+            mem_size += ggml_row_size(GGML_TYPE_F16, 2 * embed_dim * 2 * dd_config.z_channels * 1 * 1);  // quant_conv_w
+            mem_size += ggml_row_size(GGML_TYPE_F32, 2 * embed_dim);                                     // quant_conv_b
+            mem_size += encoder.calculate_mem_size(wtype);
+        }
+
+        mem_size += ggml_row_size(GGML_TYPE_F16, dd_config.z_channels * embed_dim * 1 * 1);  // post_quant_conv_w
+        mem_size += ggml_row_size(GGML_TYPE_F32, dd_config.z_channels);                      // post_quant_conv_b
+
+        mem_size += decoder.calculate_mem_size(wtype);
+        return mem_size;
+    }
+
+    size_t get_num_tensors() {
+        size_t num_tensors = decoder.get_num_tensors();
+        if (!decode_only) {
+            num_tensors += 2;
+            num_tensors += encoder.get_num_tensors();
+        }
+        return num_tensors;
+    }
+
+    void init_params() {
+        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
+
+        if (!decode_only) {
+            quant_conv_w = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 1, 1, 2 * dd_config.z_channels, 2 * embed_dim);
+            quant_conv_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, 2 * embed_dim);
+            encoder.init_params(params_ctx, alloc, wtype);
+        }
+
+        post_quant_conv_w = ggml_new_tensor_4d(params_ctx, GGML_TYPE_F16, 1, 1, embed_dim, dd_config.z_channels);
+        post_quant_conv_b = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, dd_config.z_channels);
+        decoder.init_params(params_ctx, alloc, wtype);
+
+        // 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);
+    }
+
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        tensors[prefix + "quant_conv.weight"] = quant_conv_w;
+        tensors[prefix + "quant_conv.bias"]   = quant_conv_b;
+        encoder.map_by_name(tensors, prefix + "encoder.");
+
+        tensors[prefix + "post_quant_conv.weight"] = post_quant_conv_w;
+        tensors[prefix + "post_quant_conv.bias"]   = post_quant_conv_b;
+        decoder.map_by_name(tensors, prefix + "decoder.");
+    }
+
+    struct ggml_tensor* decode(struct ggml_context* ctx0, struct ggml_tensor* z) {
+        // z: [N, z_channels, h, w]
+        // post_quant_conv
+        auto h = ggml_nn_conv_2d(ctx0, z, post_quant_conv_w, post_quant_conv_b);  // [N, z_channels, h, w]
+        ggml_set_name(h, "bench-start");
+        h = decoder.forward(ctx0, h);
+        ggml_set_name(h, "bench-end");
+        return h;
+    }
+
+    struct ggml_tensor* encode(struct ggml_context* ctx0, struct ggml_tensor* x) {
+        // x: [N, in_channels, h, w]
+        auto h = encoder.forward(ctx0, x);  // [N, 2*z_channels, h/8, w/8]
+        // quant_conv
+        h = ggml_nn_conv_2d(ctx0, h, quant_conv_w, quant_conv_b);  // [N, 2*embed_dim, h/8, w/8]
+        ggml_set_name(h, "b-end");
+        return h;
+    }
+
+    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() * VAE_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 ? decode(ctx0, z_) : encode(ctx0, 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, const int n_threads, struct ggml_tensor* z, bool decode_graph) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(z, decode_graph);
+        };
+
+        GGMLModule::compute(get_graph, n_threads, work_result);
+    }
+};
+
+#endif