fix: fix metal build issues (#183)

---------

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

.github/workflows/build.yml

@@ -30,7 +30,6 @@ jobs:
         with:
           submodules: recursive
 
-
       - name: Dependencies
         id: depends
         run: |
@@ -42,14 +41,37 @@ jobs:
         run: |
           mkdir build
           cd build
-          cmake ..
+          cmake .. -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON
           cmake --build . --config Release
 
-      #- name: Test
-        #id: cmake_test
-        #run: |
-          #cd build
-          #ctest --verbose --timeout 900
+      - name: Get commit hash
+        id: commit
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/main' ) || github.event.inputs.create_release == 'true' }}
+        uses: pr-mpt/actions-commit-hash@v2
+
+      - name: Fetch system info
+        id: system-info
+        run: |
+          echo "CPU_ARCH=`uname -m`" >> "$GITHUB_OUTPUT"
+          echo "OS_NAME=`lsb_release -s -i`" >> "$GITHUB_OUTPUT"
+          echo "OS_VERSION=`lsb_release -s -r`" >> "$GITHUB_OUTPUT"
+          echo "OS_TYPE=`uname -s`" >> "$GITHUB_OUTPUT"
+
+      - name: Pack artifacts
+        id: pack_artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        run: |
+          cp ggml/LICENSE ./build/bin/ggml.txt
+          cp LICENSE ./build/bin/stable-diffusion.cpp.txt
+          zip -j sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip ./build/bin/*
+
+      - name: Upload artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        uses: actions/upload-artifact@v3
+        with:
+          path: |
+            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
+
 
   macOS-latest-cmake:
     runs-on: macos-latest
@@ -63,9 +85,8 @@ jobs:
 
       - name: Dependencies
         id: depends
-        continue-on-error: true
         run: |
-          brew update
+          brew install zip
 
       - name: Build
         id: cmake_build
@@ -73,14 +94,37 @@ jobs:
           sysctl -a
           mkdir build
           cd build
-          cmake ..
+          cmake .. -DGGML_AVX2=ON -DCMAKE_OSX_ARCHITECTURES="arm64;x86_64" -DSD_BUILD_SHARED_LIBS=ON
           cmake --build . --config Release
 
-      #- name: Test
-        #id: cmake_test
-        #run: |
-          #cd build
-          #ctest --verbose --timeout 900
+      - name: Get commit hash
+        id: commit
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/main' ) || github.event.inputs.create_release == 'true' }}
+        uses: pr-mpt/actions-commit-hash@v2
+
+      - name: Fetch system info
+        id: system-info
+        run: |
+          echo "CPU_ARCH=`uname -m`" >> "$GITHUB_OUTPUT"
+          echo "OS_NAME=`sw_vers -productName`" >> "$GITHUB_OUTPUT"
+          echo "OS_VERSION=`sw_vers -productVersion`" >> "$GITHUB_OUTPUT"
+          echo "OS_TYPE=`uname -s`" >> "$GITHUB_OUTPUT"
+
+      - name: Pack artifacts
+        id: pack_artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        run: |
+          cp ggml/LICENSE ./build/bin/ggml.txt
+          cp LICENSE ./build/bin/stable-diffusion.cpp.txt
+          zip -j sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip ./build/bin/*
+
+      - name: Upload artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        uses: actions/upload-artifact@v3
+        with:
+          path: |
+            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
+
 
   windows-latest-cmake:
     runs-on: windows-latest
@@ -89,14 +133,17 @@ jobs:
       matrix:
         include:
           - build: 'noavx'
-            defines: '-DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF'
+            defines: '-DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF -DSD_BUILD_SHARED_LIBS=ON'
           - build: 'avx2'
-            defines: '-DGGML_AVX2=ON'
+            defines: '-DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON'
           - build: 'avx'
-            defines: '-DGGML_AVX2=OFF'
+            defines: '-DGGML_AVX2=OFF -DSD_BUILD_SHARED_LIBS=ON'
           - build: 'avx512'
-            defines: '-DGGML_AVX512=ON'
-
+            defines: '-DGGML_AVX512=ON -DSD_BUILD_SHARED_LIBS=ON'
+          - build: 'cuda12'
+            defines: '-DSD_CUBLAS=ON -DSD_BUILD_SHARED_LIBS=ON'
+          - build: 'rocm5.5'
+            defines: '-G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS="gfx1100;gfx1102;gfx1030" -DSD_BUILD_SHARED_LIBS=ON'
     steps:
       - name: Clone
         id: checkout
@@ -104,6 +151,29 @@ jobs:
         with:
           submodules: recursive
 
+      - name: Install cuda-toolkit
+        id: cuda-toolkit
+        if: ${{ matrix.build == 'cuda12' }}
+        uses: Jimver/cuda-toolkit@v0.2.11
+        with:
+          cuda: '12.2.0'
+          method: 'network'
+          sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "thrust", "visual_studio_integration"]'
+
+      - name: Install rocm-toolkit
+        id: rocm-toolkit
+        if: ${{ matrix.build == 'rocm5.5' }}
+        uses: Cyberhan123/rocm-toolkit@v0.1.0
+        with:
+          rocm: '5.5.0'
+
+      - name: Install Ninja
+        id: install-ninja
+        if: ${{ matrix.build == 'rocm5.5' }}
+        uses: urkle/action-get-ninja@v1
+        with:
+          version: 1.11.1
+
       - name: Build
         id: cmake_build
         run: |
@@ -125,12 +195,6 @@ jobs:
           & $cl /O2 /GS- /kernel avx512f.c /link /nodefaultlib /entry:main
           .\avx512f.exe && echo "AVX512F: YES" && ( echo HAS_AVX512F=1 >> $env:GITHUB_ENV ) || echo "AVX512F: NO"
 
-      #- name: Test
-        #id: cmake_test
-        #run: |
-          #cd build
-          #ctest -C Release --verbose --timeout 900
-
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -140,9 +204,37 @@ jobs:
         id: pack_artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         run: |
-          Copy-Item ggml/LICENSE .\build\bin\Release\ggml.txt
-          Copy-Item LICENSE .\build\bin\Release\stable-diffusion.cpp.txt
-          7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip .\build\bin\Release\*
+          $filePath = ".\build\bin\Release\*"
+          if (Test-Path $filePath) {
+            echo "Exists at path $filePath"
+            Copy-Item ggml/LICENSE .\build\bin\Release\ggml.txt
+            Copy-Item LICENSE .\build\bin\Release\stable-diffusion.cpp.txt
+          } elseif (Test-Path ".\build\bin\stable-diffusion.dll") {
+          $filePath = ".\build\bin\*"
+            echo "Exists at path $filePath"
+            Copy-Item ggml/LICENSE .\build\bin\ggml.txt
+            Copy-Item LICENSE .\build\bin\stable-diffusion.cpp.txt
+          } else {
+            ls .\build\bin
+            throw "Can't find stable-diffusion.dll"
+          }
+          7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip $filePath
+
+      - name: Copy and pack Cuda runtime
+        id: pack_cuda_runtime
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' && matrix.build == 'cuda12' ) || github.event.inputs.create_release == 'true' }}
+        run: |
+          echo "Cuda install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
+          $dst='.\build\bin\cudart\'
+          robocopy "${{steps.cuda-toolkit.outputs.CUDA_PATH}}\bin" $dst cudart64_*.dll cublas64_*.dll cublasLt64_*.dll
+          7z a cudart-sd-bin-win-cu12-x64.zip $dst\*
+
+      - name: Upload Cuda runtime
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' && matrix.build == 'cuda12' ) || github.event.inputs.create_release == 'true' }}
+        uses: actions/upload-artifact@v3
+        with:
+          path: |
+            cudart-sd-bin-win-cu12-x64.zip
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}

CMakeLists.txt

@@ -28,7 +28,8 @@ 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_FAST_SOFTMAX               "sd: x1.5 faster softmax, indeterministic (sometimes, same seed don't generate same image), cuda only" OFF)
+option(SD_BUILD_SHARED_LIBS          "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
 if(SD_CUBLAS)
@@ -59,17 +60,24 @@ endif()
 
 set(SD_LIB stable-diffusion)
 
-add_library(${SD_LIB} stable-diffusion.h stable-diffusion.cpp model.h model.cpp util.h util.cpp upscaler.cpp
-             ggml_extend.hpp clip.hpp common.hpp unet.hpp tae.hpp esrgan.hpp lora.hpp denoiser.hpp rng.hpp rng_philox.hpp
-             control.hpp preprocessing.hpp)
+file(GLOB SD_LIB_SOURCES 
+    "*.h"
+    "*.cpp"
+    "*.hpp"
+)
 
-if(BUILD_SHARED_LIBS)
+# we can get only one share lib
+if(SD_BUILD_SHARED_LIBS)
     message("Build shared library")
+    set(BUILD_SHARED_LIBS OFF)
+    message(${SD_LIB_SOURCES})
+    add_library(${SD_LIB} SHARED ${SD_LIB_SOURCES})
     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")
+    add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
 endif()
 
 

clip.hpp

@@ -558,11 +558,14 @@ public:
         auto token_embed_weight    = params["token_embedding.weight"];
         auto position_embed_weight = params["position_embedding.weight"];
 
-        GGML_ASSERT(input_ids->ne[0] <= position_embed_weight->ne[0]);
+        GGML_ASSERT(input_ids->ne[0] == position_embed_weight->ne[1]);
+        input_ids            = ggml_reshape_3d(ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
+        auto token_embedding = ggml_get_rows(ctx, custom_embed_weight != NULL ? custom_embed_weight : token_embed_weight, input_ids);
+        token_embedding      = ggml_reshape_3d(ctx, token_embedding, token_embedding->ne[0], token_embedding->ne[1], token_embedding->ne[3]);
 
         // token_embedding + position_embedding
         auto x = ggml_add(ctx,
-                          ggml_get_rows(ctx, custom_embed_weight != NULL ? custom_embed_weight : token_embed_weight, input_ids),
+                          token_embedding,
                           position_embed_weight);  // [N, n_token, embed_dim]
         return x;
     }
@@ -700,7 +703,7 @@ public:
         auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
 
         auto x = embeddings->forward(ctx, input_ids, tkn_embeddings);  // [N, n_token, hidden_size]
-        x      = encoder->forward(ctx, x, return_pooled ? -1 : clip_skip, true);
+        x = encoder->forward(ctx, x, return_pooled ? -1 : clip_skip, true);
         if (return_pooled || with_final_ln) {
             x = final_layer_norm->forward(ctx, x);
         }
@@ -889,7 +892,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
             return false;
         }
         struct ggml_init_params params;
-        params.mem_size               = 32 * 1024;  // max for custom embeddings 32 KB
+        params.mem_size               = 10 * 1024 * 1024;  // max for custom embeddings 10 MB
         params.mem_buffer             = NULL;
         params.no_alloc               = false;
         struct ggml_context* embd_ctx = ggml_init(params);
@@ -924,9 +927,21 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
                                 struct ggml_tensor* embeddings,
                                 size_t max_token_idx = 0,
                                 bool return_pooled   = false) {
+        size_t N       = input_ids->ne[1];
+        size_t n_token = input_ids->ne[0];
+        if (input_ids != NULL && input_ids->ne[0] > text_model.n_token) {
+            GGML_ASSERT(input_ids->ne[0] % text_model.n_token == 0);
+            input_ids = ggml_reshape_2d(ctx, input_ids, text_model.n_token, input_ids->ne[0] / text_model.n_token);
+        }
+        if (input_ids2 != NULL && input_ids2->ne[0] > text_model2.n_token) {
+            GGML_ASSERT(input_ids2->ne[0] % text_model2.n_token == 0);
+            input_ids2 = ggml_reshape_2d(ctx, input_ids2, text_model2.n_token, input_ids2->ne[0] / text_model2.n_token);
+        }
+
         if (return_pooled) {
             return text_model2.forward(ctx, input_ids2, NULL, max_token_idx, return_pooled);
         }
+
         auto hidden_states = text_model.forward(ctx, input_ids, embeddings);  // [N, n_token, hidden_size]
         // LOG_DEBUG("hidden_states: %d %d %d %d", hidden_states->ne[0], hidden_states->ne[1], hidden_states->ne[2], hidden_states->ne[3]);
         if (version == VERSION_XL) {
@@ -952,68 +967,37 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
 
             hidden_states = ggml_cont(ctx, ggml_permute(ctx, hidden_states, 1, 2, 0, 3));
         }
+        hidden_states = ggml_reshape_3d(ctx, hidden_states, hidden_states->ne[0], n_token, N);
         // LOG_DEBUG("hidden_states: %d %d %d %d", hidden_states->ne[0], hidden_states->ne[1], hidden_states->ne[2], hidden_states->ne[3]);
         return hidden_states;
     }
 
-    struct ggml_cgraph* build_graph(struct ggml_allocr* allocr, std::vector<int> tokens, bool return_pooled = false) {
+    struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
+                                    struct ggml_tensor* input_ids2 = NULL,
+                                    size_t max_token_idx           = 0,
+                                    bool return_pooled             = false) {
         struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
-        struct ggml_tensor* input_ids = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, tokens.size());
-        ggml_allocr_alloc(allocr, input_ids);
-
-        if (!ggml_allocr_is_measure(allocr)) {
-            ggml_backend_tensor_set(input_ids, tokens.data(), 0, tokens.size() * ggml_element_size(input_ids));
-        }
-
-        struct ggml_tensor* input_ids2 = NULL;
-        size_t max_token_idx           = 0;
-        if (version == VERSION_XL) {
-            input_ids2 = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, tokens.size());
-            ggml_allocr_alloc(allocr, input_ids2);
-
-            auto it = std::find(tokens.begin(), tokens.end(), EOS_TOKEN_ID);
-            if (it != tokens.end()) {
-                std::fill(std::next(it), tokens.end(), 0);
-            }
-
-            max_token_idx = std::min<size_t>(std::distance(tokens.begin(), it), tokens.size() - 1);
-
-            // for (int i = 0; i < tokens.size(); i++) {
-            //     printf("%d ", tokens[i]);
-            // }
-            // printf("\n");
-
-            if (!ggml_allocr_is_measure(allocr)) {
-                ggml_backend_tensor_set(input_ids2, tokens.data(), 0, tokens.size() * ggml_element_size(input_ids2));
-            }
+        input_ids2 = to_backend(input_ids2);
+        if (!return_pooled) {
+            input_ids = to_backend(input_ids);
         }
 
         struct ggml_tensor* embeddings = NULL;
 
         if (num_custom_embeddings > 0 && version != VERSION_XL) {
-            embeddings = ggml_new_tensor_2d(compute_ctx,
-                                            wtype,
-                                            text_model.hidden_size,
-                                            text_model.vocab_size + num_custom_embeddings /* custom placeholder */);
-            ggml_allocr_alloc(allocr, embeddings);
-            if (!ggml_allocr_is_measure(allocr)) {
-                // really bad, there is memory inflexibility (this is for host<->device memory conflicts)
-                auto token_embed_weight = text_model.get_token_embed_weight();
-                void* freeze_data       = malloc(ggml_nbytes(token_embed_weight));
-                ggml_backend_tensor_get_and_sync(backend,
-                                                 token_embed_weight,
-                                                 freeze_data,
-                                                 0,
-                                                 ggml_nbytes(token_embed_weight));
-                ggml_backend_tensor_set(embeddings, freeze_data, 0, ggml_nbytes(token_embed_weight));
-                free(freeze_data);
-                // concatenate custom embeddings
-                ggml_backend_tensor_set(embeddings,
-                                        (const void*)token_embed_custom.data(),
-                                        ggml_nbytes(token_embed_weight),
-                                        num_custom_embeddings * text_model.hidden_size * ggml_type_size(wtype));
-            }
+            auto custom_embeddings = ggml_new_tensor_3d(compute_ctx,
+                                                        wtype,
+                                                        text_model.hidden_size,
+                                                        1,
+                                                        num_custom_embeddings);
+            set_backend_tensor_data(custom_embeddings, token_embed_custom.data());
+
+            auto token_embed_weight = text_model.get_token_embed_weight();
+            token_embed_weight      = ggml_reshape_3d(compute_ctx, token_embed_weight, token_embed_weight->ne[0], 1, token_embed_weight->ne[1]);
+            // concatenate custom embeddings
+            embeddings = ggml_concat(compute_ctx, token_embed_weight, custom_embeddings);
+            embeddings = ggml_reshape_2d(compute_ctx, embeddings, embeddings->ne[0], embeddings->ne[2]);
         }
 
         struct ggml_tensor* hidden_states = forward(compute_ctx, input_ids, input_ids2, embeddings, max_token_idx, return_pooled);
@@ -1024,12 +1008,14 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
     }
 
     void compute(const int n_threads,
-                 std::vector<int> tokens,
+                 struct ggml_tensor* input_ids,
+                 struct ggml_tensor* input_ids2,
+                 size_t max_token_idx,
                  bool return_pooled,
                  ggml_tensor** output,
                  ggml_context* output_ctx = NULL) {
         auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(compute_allocr, tokens, return_pooled);
+            return build_graph(input_ids, input_ids2, max_token_idx, return_pooled);
         };
         GGMLModule::compute(get_graph, n_threads, true, output, output_ctx);
     }
@@ -1087,26 +1073,48 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
-        tokens.insert(tokens.begin(), BOS_TOKEN_ID);
-        weights.insert(weights.begin(), 1.0);
 
-        if (max_length > 0) {
-            if (tokens.size() > max_length - 1) {
-                tokens.resize(max_length - 1);
-                weights.resize(max_length - 1);
-                tokens.push_back(EOS_TOKEN_ID);
-                weights.push_back(1.0);
-            } else {
-                tokens.push_back(EOS_TOKEN_ID);
-                weights.push_back(1.0);
-                if (padding) {
-                    int pad_token_id = PAD_TOKEN_ID;
-                    if (version == VERSION_2_x) {
-                        pad_token_id = 0;
-                    }
-                    tokens.insert(tokens.end(), max_length - tokens.size(), pad_token_id);
-                    weights.insert(weights.end(), max_length - weights.size(), 1.0);
+        if (max_length > 0 && padding) {
+            size_t n = std::ceil(tokens.size() * 1.0 / (max_length - 2));
+            if (n == 0) {
+                n = 1;
+            }
+            size_t length = max_length * n;
+            LOG_DEBUG("token length: %llu", length);
+            std::vector<int> new_tokens;
+            std::vector<float> new_weights;
+            new_tokens.push_back(BOS_TOKEN_ID);
+            new_weights.push_back(1.0);
+            int token_idx = 0;
+            for (int i = 1; i < length; i++) {
+                if (token_idx >= tokens.size()) {
+                    break;
                 }
+                if (i % max_length == 0) {
+                    new_tokens.push_back(BOS_TOKEN_ID);
+                    new_weights.push_back(1.0);
+                } else if (i % max_length == max_length - 1) {
+                    new_tokens.push_back(EOS_TOKEN_ID);
+                    new_weights.push_back(1.0);
+                } else {
+                    new_tokens.push_back(tokens[token_idx]);
+                    new_weights.push_back(weights[token_idx]);
+                    token_idx++;
+                }
+            }
+
+            new_tokens.push_back(EOS_TOKEN_ID);
+            new_weights.push_back(1.0);
+            tokens  = new_tokens;
+            weights = new_weights;
+
+            if (padding) {
+                int pad_token_id = PAD_TOKEN_ID;
+                if (version == VERSION_2_x) {
+                    pad_token_id = 0;
+                }
+                tokens.insert(tokens.end(), length - tokens.size(), pad_token_id);
+                weights.insert(weights.end(), length - weights.size(), 1.0);
             }
         }
 
@@ -1143,8 +1151,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLModule {
         vision_model.get_param_tensors(tensors, prefix + "transformer.visual_model");
     }
 
-    struct ggml_cgraph* build_graph(struct ggml_allocr* allocr,
-                                    struct ggml_tensor* pixel_values) {
+    struct ggml_cgraph* build_graph(struct ggml_tensor* pixel_values) {
         struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
         pixel_values = to_backend(pixel_values);
@@ -1156,19 +1163,12 @@ struct FrozenCLIPVisionEmbedder : public GGMLModule {
         return gf;
     }
 
-    void alloc_compute_buffer(ggml_context* work_ctx, ggml_tensor* pixel_values) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(compute_allocr, pixel_values);
-        };
-        GGMLModule::alloc_compute_buffer(get_graph);
-    }
-
     void compute(const int n_threads,
                  ggml_tensor* pixel_values,
                  ggml_tensor** output,
                  ggml_context* output_ctx) {
         auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(compute_allocr, pixel_values);
+            return build_graph(pixel_values);
         };
         GGMLModule::compute(get_graph, n_threads, true, output, output_ctx);
     }

control.hpp

@@ -166,7 +166,6 @@ public:
 
     struct ggml_tensor* resblock_forward(std::string name,
                                          struct ggml_context* ctx,
-                                         struct ggml_allocr* allocr,
                                          struct ggml_tensor* x,
                                          struct ggml_tensor* emb) {
         auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
@@ -175,7 +174,6 @@ public:
 
     struct ggml_tensor* attention_layer_forward(std::string name,
                                                 struct ggml_context* ctx,
-                                                struct ggml_allocr* allocr,
                                                 struct ggml_tensor* x,
                                                 struct ggml_tensor* context) {
         auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
@@ -201,11 +199,10 @@ public:
     }
 
     std::vector<struct ggml_tensor*> forward(struct ggml_context* ctx,
-                                             struct ggml_allocr* allocr,
                                              struct ggml_tensor* x,
                                              struct ggml_tensor* hint,
                                              struct ggml_tensor* guided_hint,
-                                             std::vector<float> timesteps,
+                                             struct ggml_tensor* timesteps,
                                              struct ggml_tensor* context,
                                              struct ggml_tensor* y = NULL) {
         // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
@@ -231,7 +228,7 @@ public:
 
         auto middle_block_out = std::dynamic_pointer_cast<Conv2d>(blocks["middle_block_out.0"]);
 
-        auto t_emb = new_timestep_embedding(ctx, allocr, timesteps, model_channels);  // [N, model_channels]
+        auto t_emb = ggml_nn_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
 
         auto emb = time_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx, emb);
@@ -272,10 +269,10 @@ public:
             for (int j = 0; j < num_res_blocks; j++) {
                 input_block_idx += 1;
                 std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
-                h                = resblock_forward(name, ctx, allocr, h, emb);  // [N, mult*model_channels, h, w]
+                h                = resblock_forward(name, ctx, h, emb);  // [N, mult*model_channels, h, w]
                 if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
                     std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
-                    h                = attention_layer_forward(name, ctx, allocr, h, context);  // [N, mult*model_channels, h, w]
+                    h                = attention_layer_forward(name, ctx, h, context);  // [N, mult*model_channels, h, w]
                 }
 
                 auto zero_conv = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs." + std::to_string(input_block_idx) + ".0"]);
@@ -299,9 +296,9 @@ public:
         // [N, 4*model_channels, h/8, w/8]
 
         // middle_block
-        h = resblock_forward("middle_block.0", ctx, allocr, h, emb);             // [N, 4*model_channels, h/8, w/8]
-        h = attention_layer_forward("middle_block.1", ctx, allocr, h, context);  // [N, 4*model_channels, h/8, w/8]
-        h = resblock_forward("middle_block.2", ctx, allocr, h, emb);             // [N, 4*model_channels, h/8, w/8]
+        h = resblock_forward("middle_block.0", ctx, h, emb);             // [N, 4*model_channels, h/8, w/8]
+        h = attention_layer_forward("middle_block.1", ctx, h, context);  // [N, 4*model_channels, h/8, w/8]
+        h = resblock_forward("middle_block.2", ctx, h, emb);             // [N, 4*model_channels, h/8, w/8]
 
         // out
         outs.push_back(middle_block_out->forward(ctx, h));
@@ -386,18 +383,22 @@ struct ControlNet : public GGMLModule {
 
     struct ggml_cgraph* build_graph(struct ggml_tensor* x,
                                     struct ggml_tensor* hint,
-                                    std::vector<float> timesteps,
+                                    struct ggml_tensor* timesteps,
                                     struct ggml_tensor* context,
                                     struct ggml_tensor* y = NULL) {
         struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, CONTROL_NET_GRAPH_SIZE, false);
 
-        x       = to_backend(x);
-        hint    = to_backend(hint);
-        context = to_backend(context);
-        y       = to_backend(y);
+        x = to_backend(x);
+        if (guided_hint_cached) {
+            hint = NULL;
+        } else {
+            hint = to_backend(hint);
+        }
+        context   = to_backend(context);
+        y         = to_backend(y);
+        timesteps = to_backend(timesteps);
 
         auto outs = control_net.forward(compute_ctx,
-                                        compute_allocr,
                                         x,
                                         hint,
                                         guided_hint_cached ? guided_hint : NULL,
@@ -420,7 +421,7 @@ struct ControlNet : public GGMLModule {
     void compute(int n_threads,
                  struct ggml_tensor* x,
                  struct ggml_tensor* hint,
-                 std::vector<float> timesteps,
+                 struct ggml_tensor* timesteps,
                  struct ggml_tensor* context,
                  struct ggml_tensor* y,
                  struct ggml_tensor** output     = NULL,
@@ -434,7 +435,6 @@ struct ControlNet : public GGMLModule {
         };
 
         GGMLModule::compute(get_graph, n_threads, false, output, output_ctx);
-
         guided_hint_cached = true;
     }
 

examples/cli/main.cpp

@@ -6,7 +6,7 @@
 #include <string>
 #include <vector>
 
-#include "preprocessing.hpp"
+// #include "preprocessing.hpp"
 #include "stable-diffusion.h"
 
 #define STB_IMAGE_IMPLEMENTATION
@@ -499,6 +499,18 @@ void parse_args(int argc, const char** argv, SDParams& params) {
     }
 }
 
+static std::string sd_basename(const std::string& path) {
+    size_t pos = path.find_last_of('/');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    pos = path.find_last_of('\\');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    return path;
+}
+
 std::string get_image_params(SDParams params, int64_t seed) {
     std::string parameter_string = params.prompt + "\n";
     if (params.negative_prompt.size() != 0) {
@@ -630,8 +642,14 @@ int main(int argc, const char* argv[]) {
                                            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);
+                control_image->data = preprocess_canny(control_image->data,
+                                                       control_image->width,
+                                                       control_image->height,
+                                                       0.08f,
+                                                       0.08f,
+                                                       0.8f,
+                                                       1.0f,
+                                                       false);
             }
         }
         results = txt2img(sd_ctx,

ggml_extend.hpp

@@ -606,6 +606,20 @@ __STATIC_INLINE__ float ggml_backend_tensor_get_f32(ggml_tensor* tensor) {
     return value;
 }
 
+__STATIC_INLINE__ struct ggml_tensor* vector_to_ggml_tensor(struct ggml_context* ctx,
+                                                            const std::vector<float>& vec) {
+    struct ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+__STATIC_INLINE__ struct ggml_tensor* vector_to_ggml_tensor_i32(struct ggml_context* ctx,
+                                                                const std::vector<int>& vec) {
+    struct ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
 __STATIC_INLINE__ std::vector<float> arange(float start, float end, float step = 1.f) {
     std::vector<float> result;
 
@@ -652,7 +666,6 @@ __STATIC_INLINE__ void set_timestep_embedding(std::vector<float> timesteps,
 }
 
 __STATIC_INLINE__ struct ggml_tensor* new_timestep_embedding(struct ggml_context* ctx,
-                                                             struct ggml_allocr* allocr,
                                                              std::vector<float> timesteps,
                                                              int dim,
                                                              int max_period = 10000) {
@@ -664,17 +677,22 @@ __STATIC_INLINE__ struct ggml_tensor* new_timestep_embedding(struct ggml_context
         acutual_dim = dim + 1;
     }
     struct ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, acutual_dim, timesteps.size());
-    if (allocr != NULL) {
-        ggml_allocr_alloc(allocr, embedding);
-        if (!ggml_allocr_is_measure(allocr)) {
-            ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding));
-        }
-    } else {
+    if (embedding->data != NULL) {
         memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+    } else {
+        ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding));
     }
     return embedding;
 }
 
+__STATIC_INLINE__ struct ggml_tensor* ggml_nn_timestep_embedding(
+    struct ggml_context* ctx,
+    struct ggml_tensor* timesteps,
+    int dim,
+    int max_period = 10000) {
+    return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
+}
+
 // struct GGMLComputeGraph {
 //     virtual void init(struct ggml_context* ctx, ggml_type wtype) = 0;
 //     virtual std::string get_desc() = 0;
@@ -693,9 +711,10 @@ protected:
     struct ggml_context* params_ctx     = NULL;
     ggml_backend_buffer_t params_buffer = NULL;
 
-    struct ggml_context* compute_ctx     = NULL;
-    ggml_backend_buffer_t compute_buffer = NULL;  // for compute
-    struct ggml_allocr* compute_allocr   = NULL;
+    struct ggml_context* compute_ctx    = NULL;
+    struct ggml_gallocr* compute_allocr = NULL;
+
+    std::map<struct ggml_tensor*, const void*> backend_tensor_data_map;
 
     ggml_type wtype        = GGML_TYPE_F32;
     ggml_backend_t backend = NULL;
@@ -734,23 +753,37 @@ protected:
         }
     }
 
-    void alloc_compute_buffer(get_graph_cb_t get_graph) {
-        // alignment required by the backend
-        compute_allocr = ggml_allocr_new_measure_from_backend(backend);
-
+    bool alloc_compute_buffer(get_graph_cb_t get_graph) {
+        if (compute_allocr != NULL) {
+            return true;
+        }
         reset_compute_ctx();
         struct ggml_cgraph* gf = get_graph();
+        backend_tensor_data_map.clear();
+        compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
+
+        if (!ggml_gallocr_reserve(compute_allocr, gf)) {
+            // failed to allocate the compute buffer
+            LOG_ERROR("%s: failed to allocate the compute buffer\n", get_desc().c_str());
+            free_compute_buffer();
+            return false;
+        }
 
         // compute the required memory
-        size_t compute_buffer_size = ggml_allocr_alloc_graph(compute_allocr, gf) + 1024 * 1024;
-
-        // recreate the allocator with the required memory
-        ggml_allocr_free(compute_allocr);
-
+        size_t compute_buffer_size = ggml_gallocr_get_buffer_size(compute_allocr, 0);
         LOG_DEBUG("%s compute buffer size: %.2f MB", get_desc().c_str(), compute_buffer_size / 1024.0 / 1024.0);
+        return true;
+    }
 
-        compute_buffer = ggml_backend_alloc_buffer(backend, compute_buffer_size);
-        compute_allocr = ggml_allocr_new_from_buffer(compute_buffer);
+    void cpy_data_to_backend_tensor() {
+        for (auto& kv : backend_tensor_data_map) {
+            auto tensor = kv.first;
+            auto data   = kv.second;
+
+            ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+        }
+
+        backend_tensor_data_map.clear();
     }
 
 public:
@@ -775,31 +808,16 @@ public:
         alloc_compute_ctx();
     }
 
-    void reset_compute_allocr(get_graph_cb_t get_graph) {
-        if (compute_allocr != NULL) {
-            ggml_allocr_reset(compute_allocr);
-        } else {
-            alloc_compute_buffer(get_graph);
-        }
-    }
-
     bool alloc_params_buffer() {
-        size_t params_buffer_size = 10 * 1024 * 1024;  // 10 MB, for padding
-        params_buffer_size += get_params_mem_size();
         size_t num_tensors = get_params_num();
-
+        params_buffer      = ggml_backend_alloc_ctx_tensors(params_ctx, backend);
+        if (params_buffer == NULL) {
+            LOG_ERROR("%s alloc params backend buffer failed", get_desc().c_str());
+            return false;
+        }
+        size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer);
         LOG_DEBUG("%s params backend buffer size = % 6.2f MB (%i tensors)",
                   get_desc().c_str(), params_buffer_size / (1024.0 * 1024.0), num_tensors);
-        params_buffer = ggml_backend_alloc_buffer(backend, params_buffer_size);
-
-        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
-        // alloc all tensors linked to params_ctx
-        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);
         return true;
     }
 
@@ -812,13 +830,14 @@ public:
 
     void free_compute_buffer() {
         if (compute_allocr != NULL) {
-            ggml_allocr_free(compute_allocr);
+            ggml_gallocr_free(compute_allocr);
             compute_allocr = NULL;
         }
-        if (compute_buffer != NULL) {
-            ggml_backend_buffer_free(compute_buffer);
-            compute_buffer = NULL;
-        }
+    }
+
+    // do copy after alloc graph
+    void set_backend_tensor_data(struct ggml_tensor* tensor, const void* data) {
+        backend_tensor_data_map[tensor] = data;
     }
 
     struct ggml_tensor* to_backend(struct ggml_tensor* tensor) {
@@ -827,15 +846,11 @@ public:
             return NULL;
         }
         // it's performing a compute, check if backend isn't cpu
-        if (!ggml_backend_is_cpu(backend)) {
+        if (!ggml_backend_is_cpu(backend) && tensor->backend == GGML_BACKEND_CPU) {
             // pass input tensors to gpu memory
             auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
-            ggml_allocr_alloc(compute_allocr, backend_tensor);
 
-            // pass data to device backend
-            if (!ggml_allocr_is_measure(compute_allocr)) {
-                ggml_backend_tensor_set(backend_tensor, tensor->data, 0, ggml_nbytes(tensor));
-            }
+            set_backend_tensor_data(backend_tensor, tensor->data);
             return backend_tensor;
         } else {
             return tensor;
@@ -847,11 +862,13 @@ public:
                  bool free_compute_buffer_immediately = true,
                  struct ggml_tensor** output          = NULL,
                  struct ggml_context* output_ctx      = NULL) {
-        reset_compute_allocr(get_graph);
+        alloc_compute_buffer(get_graph);
         reset_compute_ctx();
         struct ggml_cgraph* gf = get_graph();
 
-        ggml_allocr_alloc_graph(compute_allocr, gf);
+        GGML_ASSERT(ggml_gallocr_alloc_graph(compute_allocr, gf));
+
+        cpy_data_to_backend_tensor();
 
         if (ggml_backend_is_cpu(backend)) {
             ggml_backend_cpu_set_n_threads(backend, n_threads);
@@ -1214,9 +1231,9 @@ public:
         q                     = ggml_reshape_3d(ctx, q, d_head, n_token, n_head * N);  // [N * n_head, n_token, d_head]
 
         struct ggml_tensor* k = k_proj->forward(ctx, x);
-        k                     = ggml_reshape_4d(ctx, k, d_head, n_head, n_token, N);  // [N, n_token, n_head, d_head]
-        k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));     // [N, n_head, n_token, d_head]
-        k                     = ggml_reshape_3d(ctx, k, d_head, n_token, n_head);     // [N * n_head, n_token, d_head]
+        k                     = ggml_reshape_4d(ctx, k, d_head, n_head, n_token, N);   // [N, n_token, n_head, d_head]
+        k                     = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));      // [N, n_head, n_token, d_head]
+        k                     = ggml_reshape_3d(ctx, k, d_head, n_token, n_head * N);  // [N * n_head, n_token, d_head]
 
         struct ggml_tensor* v = v_proj->forward(ctx, x);
         v                     = ggml_reshape_4d(ctx, v, d_head, n_head, n_token, N);   // [N, n_token, n_head, d_head]
@@ -1228,7 +1245,7 @@ public:
         kqv = ggml_reshape_4d(ctx, kqv, d_head, n_token, n_head, N);
         kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));  // [N, n_token, n_head, d_head]
 
-        x = ggml_reshape_2d(ctx, kqv, d_head * n_head, n_token * N);  // [N * n_token, d_head * n_head]
+        x = ggml_reshape_3d(ctx, kqv, d_head * n_head, n_token, N);  // [N * n_token, d_head * n_head]
 
         x = out_proj->forward(ctx, x);
         return x;

lora.hpp

@@ -40,26 +40,31 @@ struct LoraModel : public GGMLModule {
             LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
             return false;
         }
-        alloc_params_buffer();
-
-        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
 
+        bool dry_run          = true;
         auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
             const std::string& name = tensor_storage.name;
 
-            struct ggml_tensor* real = ggml_new_tensor(params_ctx, tensor_storage.type, tensor_storage.n_dims, tensor_storage.ne);
-            ggml_allocr_alloc(alloc, real);
-
-            *dst_tensor = real;
-
-            lora_tensors[name] = real;
+            if (dry_run) {
+                struct ggml_tensor* real = ggml_new_tensor(params_ctx,
+                                                           tensor_storage.type,
+                                                           tensor_storage.n_dims,
+                                                           tensor_storage.ne);
+                lora_tensors[name]       = real;
+            } else {
+                auto real   = lora_tensors[name];
+                *dst_tensor = real;
+            }
             return true;
         };
 
+        model_loader.load_tensors(on_new_tensor_cb, backend);
+        alloc_params_buffer();
+
+        dry_run = false;
         model_loader.load_tensors(on_new_tensor_cb, backend);
 
         LOG_DEBUG("finished loaded lora");
-        ggml_allocr_free(alloc);
         return true;
     }
 

model.cpp

@@ -498,7 +498,13 @@ void bf16_to_f32_vec(uint16_t* src, float* dst, int64_t n) {
     }
 }
 
-void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type, int n) {
+void convert_tensor(void* src,
+                    ggml_type src_type,
+                    void* dst,
+                    ggml_type dst_type,
+                    int nrows,
+                    int n_per_row) {
+    int n = nrows * n_per_row;
     if (src_type == dst_type) {
         size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
         memcpy(((char*)dst), ((char*)src), nbytes);
@@ -507,7 +513,9 @@ void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type
             ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
         } else {
             int64_t hist[16];
-            ggml_quantize_chunk(dst_type, (float*)src, dst, 0, n, hist);
+            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            const float* im = imatrix.data();
+            ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, hist, im);
         }
     } else if (dst_type == GGML_TYPE_F32) {
         if (src_type == GGML_TYPE_F16) {
@@ -536,7 +544,9 @@ void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type
             ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
         } else {
             int64_t hist[16];
-            ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, n, hist);
+            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            const float* im = imatrix.data();
+            ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, hist, im);
         }
     }
 }
@@ -1387,7 +1397,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
                     }
 
                     convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
-                                   dst_tensor->type, (int)tensor_storage.nelements());
+                                   dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
                 }
             } else {
                 read_buffer.resize(tensor_storage.nbytes());
@@ -1406,7 +1416,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
                     convert_buffer.resize(ggml_nbytes(dst_tensor));
                     convert_tensor((void*)read_buffer.data(), tensor_storage.type,
                                    (void*)convert_buffer.data(), dst_tensor->type,
-                                   (int)tensor_storage.nelements());
+                                   (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
                     ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
                 }
             }

preprocessing.hpp

@@ -117,15 +117,15 @@ void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struc
     }
 }
 
-void threshold_hystersis(struct ggml_tensor* img, float highThreshold, float lowThreshold, float weak, float strong) {
+void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float low_threshold, 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;
+    float ht = max * high_threshold;
+    float lt = ht * low_threshold;
     for (int i = 0; i < n_elements; i++) {
         float img_v = imd[i];
         if (img_v >= ht) {  // strong pixel
@@ -162,7 +162,7 @@ void threshold_hystersis(struct ggml_tensor* img, float highThreshold, float low
     }
 }
 
-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) {
+uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
     struct ggml_init_params params;
     params.mem_size               = static_cast<size_t>(10 * 1024 * 1024);  // 10
     params.mem_buffer             = NULL;
@@ -207,7 +207,7 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float highThresho
     normalize_tensor(G);
     prop_arctan2(iX, iY, tetha);
     non_max_supression(image_gray, G, tetha);
-    threshold_hystersis(image_gray, highThreshold, lowThreshold, weak, strong);
+    threshold_hystersis(image_gray, high_threshold, low_threshold, weak, strong);
     // to RGB channels
     for (int iy = 0; iy < height; iy++) {
         for (int ix = 0; ix < width; ix++) {

stable-diffusion.cpp

@@ -124,7 +124,7 @@ public:
 #endif
 #ifdef SD_USE_METAL
         LOG_DEBUG("Using Metal backend");
-        ggml_metal_log_set_callback(ggml_log_callback_default, nullptr);
+        ggml_backend_metal_log_set_callback(ggml_log_callback_default, nullptr);
         backend = ggml_backend_metal_init();
 #endif
 
@@ -363,9 +363,10 @@ public:
         struct ggml_tensor* c = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 1024, 2, 1, 1);
         ggml_set_f32(c, 0.5);
 
-        std::vector<float> timesteps = {999.f};  // [N, ]
-        int64_t t0                   = ggml_time_ms();
-        struct ggml_tensor* out      = ggml_dup_tensor(work_ctx, x_t);
+        struct ggml_tensor* timesteps = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 1);
+        ggml_set_f32(timesteps, 999);
+        int64_t t0              = ggml_time_ms();
+        struct ggml_tensor* out = ggml_dup_tensor(work_ctx, x_t);
         diffusion_model->compute(n_threads, x_t, timesteps, c, NULL, NULL, -1, {}, 0.f, &out);
         diffusion_model->free_compute_buffer();
 
@@ -450,45 +451,83 @@ public:
                                                                 int height,
                                                                 bool force_zero_embeddings = false) {
         cond_stage_model->set_clip_skip(clip_skip);
-        auto tokens_and_weights           = cond_stage_model->tokenize(text, true);
-        std::vector<int>& tokens          = tokens_and_weights.first;
-        std::vector<float>& weights       = tokens_and_weights.second;
-        int64_t t0                        = ggml_time_ms();
-        struct ggml_tensor* hidden_states = NULL;  // [N, n_token, hidden_size]
-        struct ggml_tensor* pooled        = NULL;
-        cond_stage_model->compute(n_threads, tokens, false, &hidden_states, work_ctx);
-        if (version == VERSION_XL) {
-            cond_stage_model->compute(n_threads, tokens, true, &pooled, work_ctx);
-        }
-        // if (pooled != NULL) {
-        //     print_ggml_tensor(hidden_states);
-        //     print_ggml_tensor(pooled);
-        // }
+        auto tokens_and_weights                 = cond_stage_model->tokenize(text, true);
+        std::vector<int>& tokens                = tokens_and_weights.first;
+        std::vector<float>& weights             = tokens_and_weights.second;
+        int64_t t0                              = ggml_time_ms();
+        struct ggml_tensor* hidden_states       = NULL;  // [N, n_token, hidden_size]
+        struct ggml_tensor* chunk_hidden_states = NULL;  // [n_token, hidden_size]
+        struct ggml_tensor* pooled              = NULL;
+        std::vector<float> hidden_states_vec;
 
-        int64_t t1 = ggml_time_ms();
-        LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
-        ggml_tensor* result = ggml_dup_tensor(work_ctx, hidden_states);
-        {
-            float original_mean = ggml_tensor_mean(hidden_states);
-            for (int i2 = 0; i2 < hidden_states->ne[2]; i2++) {
-                for (int i1 = 0; i1 < hidden_states->ne[1]; i1++) {
-                    for (int i0 = 0; i0 < hidden_states->ne[0]; i0++) {
-                        float value = ggml_tensor_get_f32(hidden_states, i0, i1, i2);
-                        value *= weights[i1];
-                        ggml_tensor_set_f32(result, value, i0, i1, i2);
+        size_t chunk_len   = 77;
+        size_t chunk_count = tokens.size() / chunk_len;
+        for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
+            std::vector<int> chunk_tokens(tokens.begin() + chunk_idx * chunk_len,
+                                          tokens.begin() + (chunk_idx + 1) * chunk_len);
+            std::vector<float> chunk_weights(weights.begin() + chunk_idx * chunk_len,
+                                           weights.begin() + (chunk_idx + 1) * chunk_len);
+
+            auto input_ids                 = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
+            struct ggml_tensor* input_ids2 = NULL;
+            size_t max_token_idx           = 0;
+            if (version == VERSION_XL) {
+                auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), EOS_TOKEN_ID);
+                if (it != chunk_tokens.end()) {
+                    std::fill(std::next(it), chunk_tokens.end(), 0);
+                }
+
+                max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
+
+                input_ids2 = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
+
+                // for (int i = 0; i < chunk_tokens.size(); i++) {
+                //     printf("%d ", chunk_tokens[i]);
+                // }
+                // printf("\n");
+            }
+
+            cond_stage_model->compute(n_threads, input_ids, input_ids2, max_token_idx, false, &chunk_hidden_states, work_ctx);
+            if (version == VERSION_XL && chunk_idx == 0) {
+                cond_stage_model->compute(n_threads, input_ids, input_ids2, max_token_idx, true, &pooled, work_ctx);
+            }
+            // if (pooled != NULL) {
+            //     print_ggml_tensor(chunk_hidden_states);
+            //     print_ggml_tensor(pooled);
+            // }
+
+            int64_t t1 = ggml_time_ms();
+            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+            ggml_tensor* result = ggml_dup_tensor(work_ctx, chunk_hidden_states);
+            {
+                float original_mean = ggml_tensor_mean(chunk_hidden_states);
+                for (int i2 = 0; i2 < chunk_hidden_states->ne[2]; i2++) {
+                    for (int i1 = 0; i1 < chunk_hidden_states->ne[1]; i1++) {
+                        for (int i0 = 0; i0 < chunk_hidden_states->ne[0]; i0++) {
+                            float value = ggml_tensor_get_f32(chunk_hidden_states, i0, i1, i2);
+                            value *= chunk_weights[i1];
+                            ggml_tensor_set_f32(result, value, i0, i1, i2);
+                        }
                     }
                 }
+                float new_mean = ggml_tensor_mean(result);
+                ggml_tensor_scale(result, (original_mean / new_mean));
             }
-            float new_mean = ggml_tensor_mean(result);
-            ggml_tensor_scale(result, (original_mean / new_mean));
-        }
-        if (force_zero_embeddings) {
-            float* vec = (float*)result->data;
-            for (int i = 0; i < ggml_nelements(result); i++) {
-                vec[i] = 0;
+            if (force_zero_embeddings) {
+                float* vec = (float*)result->data;
+                for (int i = 0; i < ggml_nelements(result); i++) {
+                    vec[i] = 0;
+                }
             }
+            hidden_states_vec.insert(hidden_states_vec.end(), (float*)result->data, ((float*)result->data) + ggml_nelements(result));
         }
 
+        hidden_states = vector_to_ggml_tensor(work_ctx, hidden_states_vec);
+        hidden_states = ggml_reshape_2d(work_ctx,
+                                        hidden_states,
+                                        chunk_hidden_states->ne[0],
+                                        ggml_nelements(hidden_states) / chunk_hidden_states->ne[0]);
+
         ggml_tensor* vec = NULL;
         if (version == VERSION_XL) {
             int out_dim = 256;
@@ -526,7 +565,7 @@ public:
             GGML_ASSERT(offset == ggml_nbytes(vec));
         }
         // print_ggml_tensor(result);
-        return {result, vec};
+        return {hidden_states, vec};
     }
 
     std::tuple<ggml_tensor*, ggml_tensor*, ggml_tensor*> get_svd_condition(ggml_context* work_ctx,
@@ -675,7 +714,8 @@ public:
             }
 
             float t = denoiser->schedule->sigma_to_t(sigma);
-            std::vector<float> timesteps(x->ne[3], t);  // [N, ]
+            std::vector<float> timesteps_vec(x->ne[3], t);  // [N, ]
+            auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
 
             copy_ggml_tensor(noised_input, input);
             // noised_input = noised_input * c_in

stable-diffusion.h

@@ -73,6 +73,9 @@ enum sd_type_t {
     SD_TYPE_Q8_K    = 15,
     SD_TYPE_IQ2_XXS = 16,
     SD_TYPE_IQ2_XS  = 17,
+    SD_TYPE_IQ3_XXS = 18,
+    SD_TYPE_IQ1_S   = 19,
+    SD_TYPE_IQ4_NL  = 20,
     SD_TYPE_I8,
     SD_TYPE_I16,
     SD_TYPE_I32,
@@ -174,6 +177,15 @@ SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image,
 
 SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type);
 
+SD_API uint8_t* preprocess_canny(uint8_t* img,
+                                 int width,
+                                 int height,
+                                 float high_threshold,
+                                 float low_threshold,
+                                 float weak,
+                                 float strong,
+                                 bool inverse);
+
 #ifdef __cplusplus
 }
 #endif

unet.hpp

@@ -61,7 +61,6 @@ public:
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx,
-                                struct ggml_allocr* allocr,
                                 struct ggml_tensor* x,
                                 struct ggml_tensor* context,
                                 int timesteps) {
@@ -112,9 +111,9 @@ public:
         x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3));  // [N, h, w, inner_dim]
         x = ggml_reshape_3d(ctx, x, inner_dim, w * h, n);      // [N, h * w, inner_dim]
 
-        std::vector<float> num_frames = arange(0, timesteps);
+        auto num_frames = ggml_arange(ctx, 0, timesteps, 1);
         // since b is 1, no need to do repeat
-        auto t_emb = new_timestep_embedding(ctx, allocr, num_frames, in_channels, max_time_embed_period);  // [N, in_channels]
+        auto t_emb = ggml_nn_timestep_embedding(ctx, num_frames, in_channels, max_time_embed_period);  // [N, in_channels]
 
         auto emb = time_pos_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx, emb);
@@ -342,7 +341,6 @@ public:
 
     struct ggml_tensor* resblock_forward(std::string name,
                                          struct ggml_context* ctx,
-                                         struct ggml_allocr* allocr,
                                          struct ggml_tensor* x,
                                          struct ggml_tensor* emb,
                                          int num_video_frames) {
@@ -359,14 +357,13 @@ public:
 
     struct ggml_tensor* attention_layer_forward(std::string name,
                                                 struct ggml_context* ctx,
-                                                struct ggml_allocr* allocr,
                                                 struct ggml_tensor* x,
                                                 struct ggml_tensor* context,
                                                 int timesteps) {
         if (version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
 
-            return block->forward(ctx, allocr, x, context, timesteps);
+            return block->forward(ctx, x, context, timesteps);
         } else {
             auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
 
@@ -375,9 +372,8 @@ public:
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx,
-                                struct ggml_allocr* allocr,
                                 struct ggml_tensor* x,
-                                std::vector<float> timesteps,
+                                struct ggml_tensor* timesteps,
                                 struct ggml_tensor* context,
                                 struct ggml_tensor* c_concat              = NULL,
                                 struct ggml_tensor* y                     = NULL,
@@ -386,7 +382,6 @@ public:
                                 float control_strength                    = 0.f) {
         // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
         // timesteps: [N,]
-        // t_emb: [N, model_channels] timestep_embedding(timesteps, model_channels)
         // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
         // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
@@ -417,7 +412,7 @@ public:
         auto out_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["out.0"]);
         auto out_2 = std::dynamic_pointer_cast<Conv2d>(blocks["out.2"]);
 
-        auto t_emb = new_timestep_embedding(ctx, allocr, timesteps, model_channels);  // [N, model_channels]
+        auto t_emb = ggml_nn_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
 
         auto emb = time_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx, emb);
@@ -452,10 +447,10 @@ public:
             for (int j = 0; j < num_res_blocks; j++) {
                 input_block_idx += 1;
                 std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
-                h                = resblock_forward(name, ctx, allocr, h, emb, num_video_frames);  // [N, mult*model_channels, h, w]
+                h                = resblock_forward(name, ctx, h, emb, num_video_frames);  // [N, mult*model_channels, h, w]
                 if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
                     std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
-                    h                = attention_layer_forward(name, ctx, allocr, h, context, num_video_frames);  // [N, mult*model_channels, h, w]
+                    h                = attention_layer_forward(name, ctx, h, context, num_video_frames);  // [N, mult*model_channels, h, w]
                 }
                 hs.push_back(h);
             }
@@ -473,9 +468,9 @@ public:
         // [N, 4*model_channels, h/8, w/8]
 
         // middle_block
-        h = resblock_forward("middle_block.0", ctx, allocr, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
-        h = attention_layer_forward("middle_block.1", ctx, allocr, h, context, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
-        h = resblock_forward("middle_block.2", ctx, allocr, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
+        h = resblock_forward("middle_block.0", ctx, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
+        h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
+        h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
 
         if (controls.size() > 0) {
             auto cs = ggml_scale_inplace(ctx, controls[controls.size() - 1], control_strength);
@@ -500,13 +495,13 @@ public:
 
                 std::string name = "output_blocks." + std::to_string(output_block_idx) + ".0";
 
-                h = resblock_forward(name, ctx, allocr, h, emb, num_video_frames);
+                h = resblock_forward(name, ctx, h, emb, num_video_frames);
 
                 int up_sample_idx = 1;
                 if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
                     std::string name = "output_blocks." + std::to_string(output_block_idx) + ".1";
 
-                    h = attention_layer_forward(name, ctx, allocr, h, context, num_video_frames);
+                    h = attention_layer_forward(name, ctx, h, context, num_video_frames);
 
                     up_sample_idx++;
                 }
@@ -561,7 +556,7 @@ struct UNetModel : public GGMLModule {
     }
 
     struct ggml_cgraph* build_graph(struct ggml_tensor* x,
-                                    std::vector<float> timesteps,
+                                    struct ggml_tensor* timesteps,
                                     struct ggml_tensor* context,
                                     struct ggml_tensor* c_concat              = NULL,
                                     struct ggml_tensor* y                     = NULL,
@@ -574,16 +569,16 @@ struct UNetModel : public GGMLModule {
             num_video_frames = x->ne[3];
         }
 
-        x       = to_backend(x);
-        context = to_backend(context);
-        y       = to_backend(y);
+        x         = to_backend(x);
+        context   = to_backend(context);
+        y         = to_backend(y);
+        timesteps = to_backend(timesteps);
 
         for (int i = 0; i < controls.size(); i++) {
             controls[i] = to_backend(controls[i]);
         }
 
         struct ggml_tensor* out = unet.forward(compute_ctx,
-                                               compute_allocr,
                                                x,
                                                timesteps,
                                                context,
@@ -600,7 +595,7 @@ struct UNetModel : public GGMLModule {
 
     void compute(int n_threads,
                  struct ggml_tensor* x,
-                 std::vector<float> timesteps,
+                 struct ggml_tensor* timesteps,
                  struct ggml_tensor* context,
                  struct ggml_tensor* c_concat,
                  struct ggml_tensor* y,
@@ -638,7 +633,8 @@ struct UNetModel : public GGMLModule {
             int num_video_frames = 3;
 
             auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 8, num_video_frames);
-            std::vector<float> timesteps(num_video_frames, 999.f);
+            std::vector<float> timesteps_vec(num_video_frames, 999.f);
+            auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
             ggml_set_f32(x, 0.5f);
             // print_ggml_tensor(x);
 

upscaler.cpp

@@ -21,7 +21,7 @@ struct UpscalerGGML {
 #endif
 #ifdef SD_USE_METAL
         LOG_DEBUG("Using Metal backend");
-        ggml_metal_log_set_callback(ggml_log_callback_default, nullptr);
+        ggml_backend_metal_log_set_callback(ggml_log_callback_default, nullptr);
         backend = ggml_backend_metal_init();
 #endif
 

util.cpp

@@ -10,6 +10,7 @@
 #include <thread>
 #include <unordered_set>
 #include <vector>
+#include "preprocessing.hpp"
 
 #if defined(__APPLE__) && defined(__MACH__)
 #include <sys/sysctl.h>
@@ -175,7 +176,7 @@ std::u32string unicode_value_to_utf32(int unicode_value) {
     return utf32_string;
 }
 
-std::string sd_basename(const std::string& path) {
+static std::string sd_basename(const std::string& path) {
     size_t pos = path.find_last_of('/');
     if (pos != std::string::npos) {
         return path.substr(pos + 1);
@@ -265,14 +266,14 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
         level_str = "ERROR";
     }
 
-    static char log_buffer[LOG_BUFFER_SIZE];
+    static char log_buffer[LOG_BUFFER_SIZE + 1];
 
     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);
     }
+    strncat(log_buffer, "\n", LOG_BUFFER_SIZE - strlen(log_buffer));
 
     if (sd_log_cb) {
         sd_log_cb(level, log_buffer, sd_log_cb_data);

util.h

@@ -21,8 +21,6 @@ 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);
-
 typedef struct {
     uint32_t width;
     uint32_t height;