feat: add hipBlas support (#94)

.gitmodules

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

CMakeLists.txt

@@ -25,6 +25,7 @@ endif()
 #option(SD_BUILD_TESTS                "sd: build tests"    ${SD_STANDALONE})
 option(SD_BUILD_EXAMPLES             "sd: build examples" ${SD_STANDALONE})
 option(SD_CUBLAS                     "sd: cuda backend" OFF)
+option(SD_HIPBLAS                    "sd: rocm backend" OFF)
 option(SD_METAL                      "sd: metal backend" OFF)
 option(SD_FLASH_ATTN                 "sd: use flash attention for x4 less memory usage" OFF)
 option(SD_FAST_SOFTMAX               "sd: x1.5 faster softmax, indeterministic (sometimes, same seed don't generate same image), cuda only" OFF)
@@ -46,6 +47,15 @@ if(SD_METAL)
     add_definitions(-DSD_USE_METAL)
 endif()
 
+if (SD_HIPBLAS)
+    message("Use HIPBLAS as backend stable-diffusion")
+    set(GGML_HIPBLAS ON)
+    add_definitions(-DSD_USE_CUBLAS)
+    if(SD_FAST_SOFTMAX)
+        set(GGML_CUDA_FAST_SOFTMAX ON)
+    endif()
+endif ()
+
 if(SD_FLASH_ATTN)
     message("Use Flash Attention for memory optimization")
     add_definitions(-DSD_USE_FLASH_ATTENTION)
@@ -67,6 +77,10 @@ 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)
 

README.md

@@ -117,6 +117,17 @@ cmake .. -DSD_CUBLAS=ON
 cmake --build . --config Release
 ```
 
+##### Using HipBLAS
+This provides BLAS acceleration using the ROCm cores of your AMD GPU. Make sure to have the ROCm toolkit installed.
+
+Windows User Refer to [docs/hipBLAS_on_Windows.md](docs%2FhipBLAS_on_Windows.md) for a comprehensive guide.
+
+```
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
+cmake --build . --config Release
+```
+
+
 ##### Using Metal
 
 Using Metal makes the computation run on the GPU. Currently, there are some issues with Metal when performing operations on very large matrices, making it highly inefficient at the moment. Performance improvements are expected in the near future.
@@ -126,7 +137,7 @@ cmake .. -DSD_METAL=ON
 cmake --build . --config Release
 ```
 
-### Using Flash Attention
+##### Using Flash Attention
 
 Enabling flash attention reduces memory usage by at least 400 MB. At the moment, it is not supported when CUBLAS is enabled because the kernel implementation is missing.
 
@@ -142,7 +153,7 @@ 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
@@ -168,7 +179,8 @@ arguments:
   -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
   -b, --batch-count COUNT            number of images to generate.
   --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)
-  --clip-skip N                      number of layers to skip of clip model (default: 0)
+  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)
+                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
   --vae-tiling                       process vae in tiles to reduce memory usage
   -v, --verbose                      print extra info
 ```
@@ -183,6 +195,16 @@ You can specify the model weight type using the `--type` parameter. The weights
 - `q5_0` or `q5_1` for 5-bit integer quantization
 - `q4_0` or `q4_1` for 4-bit integer quantization
 
+#### Convert to GGUF
+
+You can also convert weights in the formats `ckpt/safetensors/diffusers` to gguf and perform quantization in advance, avoiding the need for quantization every time you load them.
+
+For example:
+
+```sh
+./bin/sd -M convert -m ../models/v1-5-pruned-emaonly.safetensors -o  ../models/v1-5-pruned-emaonly.q8_0.gguf -v --type q8_0
+```
+
 #### txt2img example
 
 ```sh
@@ -240,7 +262,7 @@ Here's a simple example:
 | ----  |----    |
 | ![](./assets/without_lcm.png) |![](./assets/with_lcm.png)  |
 
-## Using TAESD to faster decoding
+#### Using TAESD to faster decoding
 
 You can use TAESD to accelerate the decoding of latent images by following these steps:
 
@@ -258,7 +280,7 @@ curl -L -O https://huggingface.co/madebyollin/taesd/blob/main/diffusion_pytorch_
 sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --taesd ../models/diffusion_pytorch_model.safetensors
 ```
 
-## Using ESRGAN to upscale results
+#### Using ESRGAN to upscale results
 
 You can use ESRGAN to upscale the generated images. At the moment, only the [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) model is supported. Support for more models of this architecture will be added soon.
 

clip.hpp

@@ -443,8 +443,6 @@ struct ResidualAttentionBlock {
     struct ggml_tensor* ln2_w;  // [hidden_size, ]
     struct ggml_tensor* ln2_b;  // [hidden_size, ]
 
-    struct ggml_tensor* attn_scale;  // [hidden_size, ]
-
     size_t calculate_mem_size(ggml_type wtype) {
         double mem_size = 0;
         mem_size += 4 * hidden_size * hidden_size * ggml_type_sizef(wtype);        // q_w/k_w/v_w/out_w
@@ -452,7 +450,6 @@ struct ResidualAttentionBlock {
         mem_size += 2 * hidden_size * intermediate_size * ggml_type_sizef(wtype);  // fc1_w/fc2_w
         mem_size += intermediate_size * ggml_type_sizef(GGML_TYPE_F32);            // fc1_b
         mem_size += hidden_size * ggml_type_sizef(GGML_TYPE_F32);                  // fc2_b
-        mem_size += ggml_type_sizef(GGML_TYPE_F32);                                // attn_scale
         return static_cast<size_t>(mem_size);
     }
 
@@ -479,10 +476,6 @@ struct ResidualAttentionBlock {
         ln2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);
         ln2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);
 
-        attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, attn_scale);
-        float scale = 1.0f / sqrt((float)d_model);
-        ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));
     }
 
     void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
@@ -521,7 +514,7 @@ struct ResidualAttentionBlock {
         // self-attention
         {
             struct ggml_tensor* q = ggml_nn_linear(ctx, x, q_w, q_b);
-            q                     = ggml_scale_inplace(ctx, q, attn_scale);
+            q                     = ggml_scale_inplace(ctx, q, 1.0f / sqrt((float)d_model));
             q                     = ggml_reshape_4d(ctx, q, d_model, n_head, n_token, N);   // [N, n_token, n_head, d_model]
             q                     = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));       // [N, n_head, n_token, d_model]
             q                     = ggml_reshape_3d(ctx, q, d_model, n_token, n_head * N);  // [N * n_head, n_token, d_model]

docs/hipBLAS_on_Windows.md

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

esrgan.hpp

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

examples/cli/main.cpp

@@ -42,11 +42,13 @@ const char* schedule_str[] = {
 const char* modes_str[] = {
     "txt2img",
     "img2img",
+    "convert",
 };
 
 enum SDMode {
     TXT2IMG,
     IMG2IMG,
+    CONVERT,
     MODE_COUNT
 };
 
@@ -125,7 +127,7 @@ 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");
@@ -384,7 +386,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         params.n_threads = get_num_physical_cores();
     }
 
-    if (params.prompt.length() == 0) {
+    if (params.mode != CONVERT && params.prompt.length() == 0) {
         fprintf(stderr, "error: the following arguments are required: prompt\n");
         print_usage(argc, argv);
         exit(1);
@@ -432,6 +434,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         srand((int)time(NULL));
         params.seed = rand();
     }
+
+    if (params.mode == CONVERT) {
+        if (params.output_path == "output.png") {
+            params.output_path = "output.gguf";
+        }
+    }
 }
 
 std::string get_image_params(SDParams params, int64_t seed) {
@@ -479,6 +487,24 @@ int main(int argc, const char* argv[]) {
         printf("%s", sd_get_system_info());
     }
 
+    if (params.mode == CONVERT) {
+        bool success = convert(params.model_path.c_str(), params.vae_path.c_str(), params.output_path.c_str(), params.wtype);
+        if (!success) {
+            fprintf(stderr,
+                    "convert '%s'/'%s' to '%s' failed\n",
+                    params.model_path.c_str(),
+                    params.vae_path.c_str(),
+                    params.output_path.c_str());
+            return 1;
+        } else {
+            printf("convert '%s'/'%s' to '%s' success\n",
+                   params.model_path.c_str(),
+                   params.vae_path.c_str(),
+                   params.output_path.c_str());
+            return 0;
+        }
+    }
+
     bool vae_decode_only        = true;
     uint8_t* input_image_buffer = NULL;
     if (params.mode == IMG2IMG) {

ggml_extend.hpp

@@ -449,7 +449,7 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_group_norm(struct ggml_context* ct
                                                          struct ggml_tensor* w,
                                                          struct ggml_tensor* b,
                                                          int num_groups = 32) {
-    if (x->n_dims == 4) {
+    if (ggml_n_dims(x) >= 3) {
         w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
         b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
     }

lora.hpp

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

model.cpp

@@ -15,6 +15,8 @@
 #include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
 
+#include "stable-diffusion.h"
+
 #ifdef SD_USE_METAL
 #include "ggml-metal.h"
 #endif
@@ -609,7 +611,7 @@ bool is_safetensors_file(const std::string& file_path) {
     }
 
     size_t header_size_ = read_u64(header_size_buf);
-    if (header_size_ >= file_size_) {
+    if (header_size_ >= file_size_ || header_size_ <= 2) {
         return false;
     }
 
@@ -673,7 +675,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
 
         // LOG_DEBUG("%s", name.c_str());
 
-        TensorStorage tensor_storage(prefix + name, dummy->type, dummy->ne, dummy->n_dims, file_index, offset);
+        TensorStorage tensor_storage(prefix + name, dummy->type, dummy->ne, ggml_n_dims(dummy), file_index, offset);
 
         GGML_ASSERT(ggml_nbytes(dummy) == tensor_storage.nbytes());
 
@@ -1181,6 +1183,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("conditioner.embedders.1") != std::string::npos) {
             return VERSION_XL;
         }
+        if (tensor_storage.name.find("cond_stage_model.1") != std::string::npos) {
+            return VERSION_XL;
+        }
         if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
             tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
             tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
@@ -1218,7 +1223,35 @@ std::string ModelLoader::load_merges() {
     return merges_utf8_str;
 }
 
+void remove_duplicates(std::vector<TensorStorage>& vec) {
+    std::unordered_map<std::string, size_t> name_to_index_map;
+
+    for (size_t i = 0; i < vec.size(); ++i) {
+        const std::string& current_name = vec[i].name;
+        auto it                         = name_to_index_map.find(current_name);
+
+        if (it != name_to_index_map.end()) {
+            vec[it->second] = vec[i];
+        } else {
+            name_to_index_map[current_name] = i;
+        }
+    }
+
+    vec.resize(name_to_index_map.size());
+}
+
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend_t backend) {
+    std::vector<TensorStorage> processed_tensor_storages;
+    for (auto& tensor_storage : tensor_storages) {
+        // LOG_DEBUG("%s", name.c_str());
+
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+
+        preprocess_tensor(tensor_storage, processed_tensor_storages);
+    }
+    remove_duplicates(processed_tensor_storages);
     bool success = true;
     for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
         std::string file_path = file_paths_[file_index];
@@ -1276,22 +1309,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
             return true;
         };
 
-        std::vector<TensorStorage> processed_tensor_storages;
-        for (auto& tensor_storage : tensor_storages) {
+        for (auto& tensor_storage : processed_tensor_storages) {
             if (tensor_storage.file_index != file_index) {
                 continue;
             }
-
-            // LOG_DEBUG("%s", name.c_str());
-
-            if (is_unused_tensor(tensor_storage.name)) {
-                continue;
-            }
-
-            preprocess_tensor(tensor_storage, processed_tensor_storages);
-        }
-
-        for (auto& tensor_storage : processed_tensor_storages) {
             // LOG_DEBUG("%s", tensor_storage.name.c_str());
 
             ggml_tensor* dst_tensor = NULL;
@@ -1417,6 +1438,9 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
         if (pair.first.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos) {
             continue;
         }
+        if (pair.first.find("alphas_cumprod") != std::string::npos) {
+            continue;
+        }
 
         if (pair.first.find("alphas_cumprod") != std::string::npos) {
             continue;
@@ -1434,7 +1458,61 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
     return true;
 }
 
-int64_t ModelLoader::cal_mem_size(ggml_backend_t backend) {
+bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type) {
+    auto backend    = ggml_backend_cpu_init();
+    size_t mem_size = 1 * 1024 * 1024;  // for padding
+    mem_size += tensor_storages.size() * ggml_tensor_overhead();
+    mem_size += cal_mem_size(backend, type);
+    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
+    ggml_context* ggml_ctx = ggml_init({mem_size, NULL, false});
+
+    gguf_context* gguf_ctx = gguf_init_empty();
+
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+
+        ggml_type tensor_type = tensor_storage.type;
+        if (type != GGML_TYPE_COUNT) {
+            if (ggml_is_quantized(type) && tensor_storage.ne[0] % 32 != 0) {
+                tensor_type = GGML_TYPE_F16;
+            } else {
+                tensor_type = type;
+            }
+        }
+
+        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
+        if (tensor == NULL) {
+            LOG_ERROR("ggml_new_tensor failed");
+            return false;
+        }
+        ggml_set_name(tensor, name.c_str());
+
+        // LOG_DEBUG("%s %d %s %d[%d %d %d %d] %d[%d %d %d %d]", name.c_str(),
+        // ggml_nbytes(tensor), ggml_type_name(tensor_type),
+        // tensor_storage.n_dims,
+        // tensor_storage.ne[0], tensor_storage.ne[1], tensor_storage.ne[2], tensor_storage.ne[3],
+        // tensor->n_dims, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+
+        *dst_tensor = tensor;
+
+        gguf_add_tensor(gguf_ctx, tensor);
+
+        return true;
+    };
+
+    bool success = load_tensors(on_new_tensor_cb, backend);
+    ggml_backend_free(backend);
+    LOG_INFO("load tensors done");
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    if (success) {
+        gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
+    }
+    ggml_free(ggml_ctx);
+    gguf_free(gguf_ctx);
+    return success;
+}
+
+int64_t ModelLoader::cal_mem_size(ggml_backend_t backend, ggml_type type) {
     size_t alignment = 128;
     if (backend != NULL) {
         alignment = ggml_backend_get_alignment(backend);
@@ -1449,8 +1527,35 @@ int64_t ModelLoader::cal_mem_size(ggml_backend_t backend) {
     }
 
     for (auto& tensor_storage : processed_tensor_storages) {
+        ggml_type tensor_type = tensor_storage.type;
+        if (type != GGML_TYPE_COUNT) {
+            if (ggml_is_quantized(type) && tensor_storage.ne[0] % 32 != 0) {
+                tensor_type = GGML_TYPE_F16;
+            } else {
+                tensor_type = type;
+            }
+        }
+        tensor_storage.type = tensor_type;
         mem_size += tensor_storage.nbytes() + alignment;
     }
 
     return mem_size;
 }
+
+bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type) {
+    ModelLoader model_loader;
+
+    if (!model_loader.init_from_file(input_path)) {
+        LOG_ERROR("init model loader from file failed: '%s'", input_path);
+        return false;
+    }
+
+    if (vae_path != NULL && strlen(vae_path) > 0) {
+        if (!model_loader.init_from_file(vae_path, "vae.")) {
+            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
+            return false;
+        }
+    }
+    bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type);
+    return success;
+}

model.h

@@ -4,9 +4,9 @@
 #include <functional>
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
-#include <set>
 
 #include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
@@ -121,7 +121,8 @@ public:
     bool load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
                       ggml_backend_t backend,
                       std::set<std::string> ignore_tensors = {});
-    int64_t cal_mem_size(ggml_backend_t backend);
+    bool save_to_gguf_file(const std::string& file_path, ggml_type type);
+    int64_t cal_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
     ~ModelLoader() = default;
 };
 #endif  // __MODEL_H__

stable-diffusion.h

@@ -71,6 +71,7 @@ enum sd_type_t {
     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,
@@ -148,7 +149,9 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
                                         enum sd_type_t wtype);
 SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
-SD_API sd_image_t upscale(upscaler_ctx_t*, sd_image_t input_image, uint32_t upscale_factor);
+SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor);
+
+SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type);
 
 #ifdef __cplusplus
 }

tae.hpp

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

unet.hpp

@@ -182,8 +182,6 @@ struct SpatialTransformer {
 
     std::vector<Transformer> transformers;
 
-    struct ggml_tensor* attn_scale;
-
     // proj_out
     struct ggml_tensor* proj_out_w;  // [in_channels, in_channels, 1, 1]
     struct ggml_tensor* proj_out_b;  // [in_channels,]
@@ -202,7 +200,6 @@ struct SpatialTransformer {
         mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                        // norm_w/norm_b
         mem_size += 2 * in_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16);  // proj_in_w/proj_out_w
         mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                        // proj_in_b/proj_out_b
-        mem_size += 1 * ggml_type_sizef(GGML_TYPE_F32);                                      // attn_scale
 
         // transformer
         for (auto& transformer : transformers) {
@@ -226,11 +223,6 @@ struct SpatialTransformer {
         proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
         proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
 
-        attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, attn_scale);
-        float scale = 1.0f / sqrt((float)d_head);
-        ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));
-
         // transformer
         for (auto& transformer : transformers) {
             transformer.norm1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
@@ -332,7 +324,7 @@ struct SpatialTransformer {
                 x                     = ggml_reshape_2d(ctx, x, c, h * w * n);        // [N * h * w, in_channels]
                 struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn1_q_w, x);  // [N * h * w, in_channels]
 #if !defined(SD_USE_FLASH_ATTENTION) || defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)
-                q = ggml_scale_inplace(ctx, q, attn_scale);
+                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]
@@ -380,7 +372,7 @@ struct SpatialTransformer {
                 context               = ggml_reshape_2d(ctx, context, context->ne[0], context->ne[1] * context->ne[2]);  // [N * max_position, hidden_size]
                 struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn2_q_w, x);                                     // [N * h * w, in_channels]
 #if !defined(SD_USE_FLASH_ATTENTION) || defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)
-                q = ggml_scale_inplace(ctx, q, attn_scale);
+                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]

vae.hpp

@@ -118,8 +118,6 @@ struct AttnBlock {
     struct ggml_tensor* proj_out_w;  // [in_channels, in_channels, 1, 1]
     struct ggml_tensor* proj_out_b;  // [in_channels,]
 
-    struct ggml_tensor* attn_scale;
-
     size_t calculate_mem_size(ggml_type wtype) {
         double mem_size = 0;
         mem_size += 6 * in_channels * ggml_type_sizef(GGML_TYPE_F32);                        // norm_w/norm_b/q_b/k_v/v_b/proj_out_b
@@ -140,11 +138,6 @@ struct AttnBlock {
 
         proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels);
         proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-
-        attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-        ggml_allocr_alloc(alloc, attn_scale);
-        float scale = 1.0f / sqrt((float)in_channels);
-        ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));
     }
 
     void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
@@ -181,7 +174,7 @@ struct AttnBlock {
         k = ggml_reshape_3d(ctx, k, c, h * w, n);              // [N, h * w, in_channels]
 
         auto w_ = ggml_mul_mat(ctx, k, q);  // [N, h * w, h * w]
-        w_      = ggml_scale_inplace(ctx, w_, attn_scale);
+        w_      = ggml_scale_inplace(ctx, w_, 1.0f / sqrt((float)in_channels));
         w_      = ggml_soft_max_inplace(ctx, w_);
 
         v  = ggml_reshape_3d(ctx, v, h * w, c, n);               // [N, in_channels, h * w]