fix: avoid double free and fix sdxl lora naming conversion

* Fixed a double free issue when running multiple backends on the CPU, eg: CLIP and the primary backend, as this would result in the *_backend pointers both pointing to the same thing resulting in a segfault when calling the StableDiffusionGGML destructor. * Improve logging to allow for a color switch on the command line interface. Changed the base log_printf function to not bake the log level directly into the log buffer as that information is already passed the logging function via the level parameter and it's easier to add in there than strip it out. * Added a fix for certain SDXL LoRAs that don't seem to follow the expected naming convention, converts over the tensor name during the LoRA model loading. Added some logging of useful LoRA loading information. Had to increase the base size of the GGML graph as the existing size results in an insufficient graph memory error when using SDXL LoRAs. * small fixes --------- Co-authored-by: leejet <leejet714@gmail.com>
feat: add TencentARC PhotoMaker support (#179 )
2026-07-23 19:30:54 -05:00 · 2024-03-20 22:00:22 +08:00 · 2024-03-12 23:15:17 +08:00 · 2024-03-10 17:23:11 +08:00 · 2024-03-03 13:27:07 +08:00 · 2024-03-03 13:24:59 +08:00
36 changed files with 4487 additions and 437 deletions
--- a/.github/workflows/build.yml
+++ b/.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' }}
--- a/CMakeLists.txt
+++ b/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,25 @@ 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")
+    message(${SD_LIB_SOURCES})
+    set(BUILD_SHARED_LIBS OFF)
+    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")
+    set(BUILD_SHARED_LIBS OFF)
+    add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
 endif()


--- a/README.md
+++ b/README.md
@@ -14,6 +14,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
    - !!!The VAE in SDXL encounters NaN issues under FP16, but unfortunately, the ggml_conv_2d only operates under FP16. Hence, a parameter is needed to specify the VAE that has fixed the FP16 NaN issue. You can find it here: [SDXL VAE FP16 Fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors).

 - [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo) and [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo) support
+- [PhotoMaker](https://github.com/TencentARC/PhotoMaker) support.
 - 16-bit, 32-bit float support
 - 4-bit, 5-bit and 8-bit integer quantization support
 - Accelerated memory-efficient CPU inference
@@ -60,6 +61,9 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in

 ## Usage

+For most users, you can download the built executable program from the latest [release](https://github.com/leejet/stable-diffusion.cpp/releases/latest).
+If the built product does not meet your requirements, you can choose to build it manually.
+
 ### Get the Code

 ```
@@ -148,7 +152,7 @@ cmake --build . --config Release
 ### Run

 ```
-usage: ./build/bin/sd [arguments]
+usage: ./bin/sd [arguments]

 arguments:
  -h, --help                         show this help message and exit
@@ -160,6 +164,9 @@ arguments:
  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
  --control-net [CONTROL_PATH]       path to control net model
  --embd-dir [EMBEDDING_PATH]        path to embeddings.
+  --stacked-id-embd-dir [DIR]        path to PHOTOMAKER stacked id embeddings.
+  --input-id-images-dir [DIR]        path to PHOTOMAKER input id images dir.
+  --normalize-input                  normalize PHOTOMAKER input id images
  --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.
  --upscale-repeats                  Run the ESRGAN upscaler this many times (default 1)
  --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
@@ -172,6 +179,7 @@ arguments:
  -n, --negative-prompt PROMPT       the negative prompt (default: "")
  --cfg-scale SCALE                  unconditional guidance scale: (default: 7.0)
  --strength STRENGTH                strength for noising/unnoising (default: 0.75)
+  --style-ratio STYLE-RATIO          strength for keeping input identity (default: 20%)
  --control-strength STRENGTH        strength to apply Control Net (default: 0.9)
                                     1.0 corresponds to full destruction of information in init image
  -H, --height H                     image height, in pixel space (default: 512)
@@ -296,6 +304,39 @@ You can use ESRGAN to upscale the generated images. At the moment, only the [Rea
 sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --upscale-model ../models/RealESRGAN_x4plus_anime_6B.pth
 ```

+#### Using PhotoMaker to personalize image generation
+
+You can use [PhotoMaker](https://github.com/TencentARC/PhotoMaker) to personalize generated images with your own ID. 
+
+**NOTE**, currently PhotoMaker **ONLY** works with **SDXL** (any SDXL model files will work).
+
+Download PhotoMaker model file (in safetensor format) [here](https://huggingface.co/bssrdf/PhotoMaker). The official release of the model file (in .bin format) does not work with ```stablediffusion.cpp```.
+
+- Specify the PhotoMaker model path using the `--stacked-id-embd-dir PATH` parameter. 
+- Specify the input images path using the `--input-id-images-dir PATH` parameter. 
+  - input images **must** have the same width and height for preprocessing (to be improved)
+
+In prompt, make sure you have a class word followed by the trigger word ```"img"``` (hard-coded for now). The class word could be one of ```"man, woman, girl, boy"```. If input ID images contain asian faces, add ```Asian``` before the class
+word.
+
+Another PhotoMaker specific parameter:
+
+- ```--style-ratio  (0-100)%```: default is 20 and 10-20 typically gets good results. Lower ratio means more faithfully following input ID (not necessarily better quality).
+
+Other parameters recommended for running Photomaker:
+
+- ```--cfg-scale 5.0```
+- ```-H 1024```
+- ```-W 1024```
+
+If on low memory GPUs (<= 8GB), recommend running with ```--vae-on-cpu``` option to get artifact free images.
+
+Example:
+
+```bash
+bin/sd -m ../models/sdxlUnstableDiffusers_v11.safetensors  --vae ../models/sdxl_vae.safetensors --stacked-id-embd-dir ../models/photomaker-v1.safetensors --input-id-images-dir ../assets/examples/scarletthead_woman -p "a girl img, retro futurism, retro game art style but extremely beautiful, intricate details, masterpiece, best quality, space-themed, cosmic, celestial, stars, galaxies, nebulas, planets, science fiction, highly detailed" -n "realistic, photo-realistic, worst quality, greyscale, bad anatomy, bad hands, error, text" --cfg-scale 5.0  --sampling-method euler -H 1024 -W 1024 --style-ratio 10 --vae-on-cpu -o output.png
+```
+
 ### Docker

 #### Building using Docker
@@ -319,6 +360,13 @@ docker run -v /path/to/models:/models -v /path/to/output/:/output sd [args...]
 |  **Memory** (txt2img - 512 x 512) | ~2.8G | ~2.3G | ~2.1G | ~2.0G | ~2.0G | ~2.0G | ~2.0G |
 |  **Memory** (txt2img - 512 x 512) *with Flash Attention* | ~2.4G | ~1.9G | ~1.6G | ~1.5G | ~1.5G | ~1.5G | ~1.5G |

+## Bindings
+
+These projects wrap `stable-diffusion.cpp` for easier use in other languages/frameworks.
+
+* Golang: [seasonjs/stable-diffusion](https://github.com/seasonjs/stable-diffusion)
+* C#: [DarthAffe/StableDiffusion.NET](https://github.com/DarthAffe/StableDiffusion.NET)
+
 ## Contributors

 Thank you to all the people who have already contributed to stable-diffusion.cpp!
@@ -335,3 +383,4 @@ Thank you to all the people who have already contributed to stable-diffusion.cpp
 - [k-diffusion](https://github.com/crowsonkb/k-diffusion)
 - [latent-consistency-model](https://github.com/luosiallen/latent-consistency-model)
 - [generative-models](https://github.com/Stability-AI/generative-models/)
+- [PhotoMaker](https://github.com/TencentARC/PhotoMaker)
--- a/assets/photomaker_examples/lenna_woman/lenna.jpg
+++ b/assets/photomaker_examples/lenna_woman/lenna.jpg
--- a/assets/photomaker_examples/newton_man/newton_0.jpg
+++ b/assets/photomaker_examples/newton_man/newton_0.jpg
--- a/assets/photomaker_examples/newton_man/newton_1.jpg
+++ b/assets/photomaker_examples/newton_man/newton_1.jpg
--- a/assets/photomaker_examples/newton_man/newton_2.png
+++ b/assets/photomaker_examples/newton_man/newton_2.png
--- a/assets/photomaker_examples/newton_man/newton_3.jpg
+++ b/assets/photomaker_examples/newton_man/newton_3.jpg
--- a/assets/photomaker_examples/scarletthead_woman/scarlett_0.jpg
+++ b/assets/photomaker_examples/scarletthead_woman/scarlett_0.jpg
--- a/assets/photomaker_examples/scarletthead_woman/scarlett_1.jpg
+++ b/assets/photomaker_examples/scarletthead_woman/scarlett_1.jpg
--- a/assets/photomaker_examples/scarletthead_woman/scarlett_2.jpg
+++ b/assets/photomaker_examples/scarletthead_woman/scarlett_2.jpg
--- a/assets/photomaker_examples/scarletthead_woman/scarlett_3.jpg
+++ b/assets/photomaker_examples/scarletthead_woman/scarlett_3.jpg
--- a/assets/photomaker_examples/yangmi_woman/yangmi_1.jpg
+++ b/assets/photomaker_examples/yangmi_woman/yangmi_1.jpg
--- a/assets/photomaker_examples/yangmi_woman/yangmi_2.jpeg
+++ b/assets/photomaker_examples/yangmi_woman/yangmi_2.jpeg
--- a/assets/photomaker_examples/yangmi_woman/yangmi_3.jpg
+++ b/assets/photomaker_examples/yangmi_woman/yangmi_3.jpg
--- a/assets/photomaker_examples/yangmi_woman/yangmi_4.jpg
+++ b/assets/photomaker_examples/yangmi_woman/yangmi_4.jpg
--- a/assets/photomaker_examples/yangmi_woman/yangmi_5.jpg
+++ b/assets/photomaker_examples/yangmi_woman/yangmi_5.jpg
--- a/assets/photomaker_examples/yangmi_woman/yangmi_6.jpg
+++ b/assets/photomaker_examples/yangmi_woman/yangmi_6.jpg
--- a/clip.hpp
+++ b/clip.hpp
@@ -75,9 +75,13 @@ class CLIPTokenizer {
 private:
    SDVersion version = VERSION_1_x;
    std::map<int, std::u32string> byte_encoder;
+    std::map<std::u32string, int> byte_decoder;
    std::map<std::u32string, int> encoder;
+    std::map<int, std::u32string> decoder;
    std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
    std::regex pat;
+    int encoder_len;
+    int bpe_len;

    static std::string strip(const std::string& str) {
        std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
@@ -118,7 +122,11 @@ public:

    void load_from_merges(const std::string& merges_utf8_str) {
        auto byte_unicode_pairs = bytes_to_unicode();
-        byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+        // printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size());
+        byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+        for (auto& pair : byte_unicode_pairs) {
+            byte_decoder[pair.second] = pair.first;
+        }
        // for (auto & pair: byte_unicode_pairs) {
        //     std::cout << pair.first << ": " << pair.second << std::endl;
        // }
@@ -138,6 +146,8 @@ public:
            size_t space_pos = merge.find(' ');
            merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
            // LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
+            // printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
+            //                     utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
        }
        std::vector<std::u32string> vocab;
        for (const auto& pair : byte_unicode_pairs) {
@@ -154,15 +164,36 @@ public:
        LOG_DEBUG("vocab size: %llu", vocab.size());
        int i = 0;
        for (const auto& token : vocab) {
-            encoder[token] = i++;
+            encoder[token] = i;
+            decoder[i]     = token;
+            i++;
+        }
+        encoder_len = i;
+
+        auto it = encoder.find(utf8_to_utf32("img</w>"));
+        if (it != encoder.end()) {
+            LOG_DEBUG(" trigger word img already in vocab");
+        } else {
+            LOG_DEBUG(" trigger word img not in vocab yet");
        }

        int rank = 0;
        for (const auto& merge : merge_pairs) {
            bpe_ranks[merge] = rank++;
        }
+        bpe_len = rank;
    };

+    void add_token(const std::string& text) {
+        std::u32string token = utf8_to_utf32(text);
+        auto it              = encoder.find(token);
+        if (it != encoder.end()) {
+            encoder[token]       = encoder_len;
+            decoder[encoder_len] = token;
+            encoder_len++;
+        }
+    }
+
    std::u32string bpe(const std::u32string& token) {
        std::vector<std::u32string> word;

@@ -243,6 +274,7 @@ public:
                              size_t max_length = 0,
                              bool padding      = false) {
        std::vector<int32_t> tokens = encode(text, on_new_token_cb);
+
        tokens.insert(tokens.begin(), BOS_TOKEN_ID);
        if (max_length > 0) {
            if (tokens.size() > max_length - 1) {
@@ -259,9 +291,34 @@ public:
                }
            }
        }
+
        return tokens;
    }

+    std::string decode(const std::vector<int>& tokens) {
+        std::string text = "";
+        for (int t : tokens) {
+            if (t == 49406 || t == 49407)
+                continue;
+            std::u32string ts = decoder[t];
+            // printf("%d, %s \n", t,  utf32_to_utf8(ts).c_str());
+            std::string s = utf32_to_utf8(ts);
+            if (s.length() >= 4 && ends_with(s, "</w>")) {
+                text += " " + s.replace(s.length() - 4, s.length() - 1, "");
+            } else {
+                text += " " + s;
+            }
+        }
+        // std::vector<unsigned char> bytes;
+        // for (auto c : text){
+        //     bytes.push_back(byte_decoder[c]);
+        // }
+
+        // std::string s((char *)bytes.data());
+        // std::string s = "";
+        return trim(text);
+    }
+
    std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb) {
        std::string original_text = text;
        std::vector<int32_t> bpe_tokens;
@@ -308,7 +365,8 @@ public:
            ss << "\"" << token << "\", ";
        }
        ss << "]";
-        LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
+        // LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
+        // printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str());
        return bpe_tokens;
    }
 };
@@ -469,7 +527,8 @@ public:
        : d_model(d_model),
          n_head(n_head),
          intermediate_size(intermediate_size) {
-        blocks["self_attn"]   = std::shared_ptr<GGMLBlock>(new MultiheadAttention(d_model, n_head));
+        blocks["self_attn"] = std::shared_ptr<GGMLBlock>(new MultiheadAttention(d_model, n_head, true));
+
        blocks["layer_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_model));
        blocks["layer_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_model));

@@ -508,7 +567,7 @@ public:
    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, int clip_skip = -1, bool mask = true) {
        // x: [N, n_token, d_model]
        int layer_idx = n_layer - 1;
-        LOG_DEBUG("clip_skip %d", clip_skip);
+        // LOG_DEBUG("clip_skip %d", clip_skip);
        if (clip_skip > 0) {
            layer_idx = n_layer - clip_skip;
        }
@@ -520,7 +579,7 @@ public:
            }
            std::string name = "layers." + std::to_string(i);
            auto layer       = std::dynamic_pointer_cast<CLIPLayer>(blocks[name]);
-            x                = layer->forward(ctx, x);  // [N, n_token, d_model]
+            x                = layer->forward(ctx, x, mask);  // [N, n_token, d_model]
            // LOG_DEBUG("layer %d", i);
        }
        return x;
@@ -558,11 +617,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;
    }
@@ -717,11 +779,6 @@ public:
 };

 class CLIPVisionModel : public GGMLBlock {
-protected:
-    void init_params(struct ggml_context* ctx, ggml_type wtype) {
-        params["visual_projection"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, projection_dim, hidden_size);
-    }
-
 public:
    // network hparams
    int32_t num_channels      = 3;
@@ -732,16 +789,14 @@ public:
    int32_t intermediate_size = 4096;
    int32_t n_head            = 16;
    int32_t n_layer           = 24;
-    int32_t projection_dim    = 768;

 public:
-    CLIPVisionModel(CLIPVersion version = OPEN_CLIP_VIT_H_14) {
+    CLIPVisionModel(CLIPVersion version = OPENAI_CLIP_VIT_L_14) {
        if (version == OPEN_CLIP_VIT_H_14) {
            hidden_size       = 1280;
            intermediate_size = 5120;
            n_head            = 16;
            n_layer           = 32;
-            projection_dim    = 1024;
        } else if (version == OPEN_CLIP_VIT_BIGG_14) {
            hidden_size       = 1664;
            intermediate_size = 8192;
@@ -755,9 +810,8 @@ public:
        blocks["post_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
    }

-    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* pixel_values) {
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* pixel_values, bool return_pooled = true) {
        // pixel_values: [N, num_channels, image_size, image_size]
-        // return: // [N, projection_dim]
        auto embeddings     = std::dynamic_pointer_cast<CLIPVisionEmbeddings>(blocks["embeddings"]);
        auto pre_layernorm  = std::dynamic_pointer_cast<LayerNorm>(blocks["pre_layernorm"]);
        auto encoder        = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
@@ -765,26 +819,60 @@ public:

        auto x = embeddings->forward(ctx, pixel_values);  // [N, num_positions, embed_dim]
        x      = pre_layernorm->forward(ctx, x);
-        x      = encoder->forward(ctx, x, -1, true);
+        x      = encoder->forward(ctx, x, -1, false);
        x      = post_layernorm->forward(ctx, x);  // [N, n_token, hidden_size]

-        GGML_ASSERT(x->ne[2] == 1);
-        int64_t max_token_idx  = 0;
-        ggml_tensor* pooled    = ggml_view_1d(ctx, x, x->ne[0], x->nb[1] * max_token_idx);  // assert N == 1
-        auto visual_projection = params["visual_projection"];
-        pooled                 = ggml_mul_mat(ctx, ggml_cont(ctx, ggml_transpose(ctx, visual_projection)), pooled);
-        return pooled;  // [N, projection_dim]
+        GGML_ASSERT(x->ne[3] == 1);
+        if (return_pooled) {
+            ggml_tensor* pooled = ggml_cont(ctx, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0));
+            return pooled;  // [N, hidden_size]
+        } else {
+            return x;  // [N, n_token, hidden_size]
+        }
+    }
+};
+
+class CLIPProjection : public UnaryBlock {
+protected:
+    int64_t in_features;
+    int64_t out_features;
+    bool transpose_weight;
+
+    void init_params(struct ggml_context* ctx, ggml_type wtype) {
+        if (transpose_weight) {
+            LOG_ERROR("transpose_weight");
+            params["weight"] = ggml_new_tensor_2d(ctx, wtype, out_features, in_features);
+        } else {
+            params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        }
+    }
+
+public:
+    CLIPProjection(int64_t in_features,
+                   int64_t out_features,
+                   bool transpose_weight = false)
+        : in_features(in_features),
+          out_features(out_features),
+          transpose_weight(transpose_weight) {}
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        struct ggml_tensor* w = params["weight"];
+        if (transpose_weight) {
+            w = ggml_cont(ctx, ggml_transpose(ctx, w));
+        }
+        return ggml_nn_linear(ctx, x, w, NULL);
    }
 };

 class CLIPVisionModelProjection : public GGMLBlock {
 public:
    int32_t hidden_size    = 1024;
-    int32_t projection_dim = 1024;
+    int32_t projection_dim = 768;
    int32_t image_size     = 224;

 public:
-    CLIPVisionModelProjection(CLIPVersion version = OPEN_CLIP_VIT_H_14) {
+    CLIPVisionModelProjection(CLIPVersion version   = OPENAI_CLIP_VIT_L_14,
+                              bool transpose_proj_w = false) {
        if (version == OPEN_CLIP_VIT_H_14) {
            hidden_size    = 1280;
            projection_dim = 1024;
@@ -792,17 +880,17 @@ public:
            hidden_size = 1664;
        }

-        blocks["visual_model"]      = std::shared_ptr<GGMLBlock>(new CLIPVisionModel(version));
-        blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, projection_dim, false));
+        blocks["vision_model"]      = std::shared_ptr<GGMLBlock>(new CLIPVisionModel(version));
+        blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new CLIPProjection(hidden_size, projection_dim, transpose_proj_w));
    }

    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* pixel_values) {
        // pixel_values: [N, num_channels, image_size, image_size]
-        // return: [N, num_positions, projection_dim]
-        auto visual_model      = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["visual_model"]);
-        auto visual_projection = std::dynamic_pointer_cast<Linear>(blocks["visual_projection"]);
+        // return: [N, projection_dim]
+        auto vision_model      = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
+        auto visual_projection = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);

-        auto x = visual_model->forward(ctx, pixel_values);  // [N, embed_dim]
+        auto x = vision_model->forward(ctx, pixel_values);  // [N, hidden_size]
        x      = visual_projection->forward(ctx, x);        // [N, projection_dim]

        return x;  // [N, projection_dim]
@@ -888,8 +976,12 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
            LOG_ERROR("embedding '%s' failed", embd_name.c_str());
            return false;
        }
+        if (std::find(readed_embeddings.begin(), readed_embeddings.end(), embd_name) != readed_embeddings.end()) {
+            LOG_DEBUG("embedding already read in: %s", embd_name.c_str());
+            return true;
+        }
        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 +1016,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 +1056,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 +1097,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);
    }
@@ -1039,6 +1114,205 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
        return tokenize(text, text_model.n_token, padding);
    }

+    std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
+    tokenize_with_trigger_token(std::string text,
+                                int num_input_imgs,
+                                int32_t image_token,
+                                bool padding = false) {
+        return tokenize_with_trigger_token(text, num_input_imgs, image_token,
+                                           text_model.n_token, padding);
+    }
+
+    std::vector<int> convert_token_to_id(std::string text) {
+        auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
+            size_t word_end       = str.find(",");
+            std::string embd_name = word_end == std::string::npos ? str : str.substr(0, word_end);
+            embd_name             = trim(embd_name);
+            std::string embd_path = get_full_path(embd_dir, embd_name + ".pt");
+            if (embd_path.size() == 0) {
+                embd_path = get_full_path(embd_dir, embd_name + ".ckpt");
+            }
+            if (embd_path.size() == 0) {
+                embd_path = get_full_path(embd_dir, embd_name + ".safetensors");
+            }
+            if (embd_path.size() > 0) {
+                if (load_embedding(embd_name, embd_path, bpe_tokens)) {
+                    if (word_end != std::string::npos) {
+                        str = str.substr(word_end);
+                    } else {
+                        str = "";
+                    }
+                    return true;
+                }
+            }
+            return false;
+        };
+        std::vector<int> curr_tokens = tokenizer.encode(text, on_new_token_cb);
+        return curr_tokens;
+    }
+
+    std::string decode(const std::vector<int>& tokens) {
+        return tokenizer.decode(tokens);
+    }
+
+    void pad_tokens(std::vector<int>& tokens,
+                    std::vector<float>& weights,
+                    size_t max_length = 0,
+                    bool padding      = false) {
+        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);
+            }
+        }
+    }
+
+    std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
+    tokenize_with_trigger_token(std::string text,
+                                int num_input_imgs,
+                                int32_t image_token,
+                                size_t max_length = 0,
+                                bool padding      = false) {
+        auto parsed_attention = parse_prompt_attention(text);
+
+        {
+            std::stringstream ss;
+            ss << "[";
+            for (const auto& item : parsed_attention) {
+                ss << "['" << item.first << "', " << item.second << "], ";
+            }
+            ss << "]";
+            LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
+        }
+
+        auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
+            size_t word_end       = str.find(",");
+            std::string embd_name = word_end == std::string::npos ? str : str.substr(0, word_end);
+            embd_name             = trim(embd_name);
+            std::string embd_path = get_full_path(embd_dir, embd_name + ".pt");
+            if (embd_path.size() == 0) {
+                embd_path = get_full_path(embd_dir, embd_name + ".ckpt");
+            }
+            if (embd_path.size() == 0) {
+                embd_path = get_full_path(embd_dir, embd_name + ".safetensors");
+            }
+            if (embd_path.size() > 0) {
+                if (load_embedding(embd_name, embd_path, bpe_tokens)) {
+                    if (word_end != std::string::npos) {
+                        str = str.substr(word_end);
+                    } else {
+                        str = "";
+                    }
+                    return true;
+                }
+            }
+            return false;
+        };
+
+        std::vector<int> tokens;
+        std::vector<float> weights;
+        std::vector<bool> class_token_mask;
+        int32_t class_idx = -1, tokens_acc = 0;
+        for (const auto& item : parsed_attention) {
+            std::vector<int> class_token_index;
+            std::vector<int> clean_input_ids;
+            const std::string& curr_text = item.first;
+            float curr_weight            = item.second;
+            // printf(" %s: %f \n", curr_text.c_str(), curr_weight);
+            std::vector<int> curr_tokens = tokenizer.encode(curr_text, on_new_token_cb);
+            int32_t clean_index          = 0;
+            for (uint32_t i = 0; i < curr_tokens.size(); i++) {
+                int token_id = curr_tokens[i];
+                if (token_id == image_token)
+                    class_token_index.push_back(clean_index - 1);
+                else {
+                    clean_input_ids.push_back(token_id);
+                    clean_index++;
+                }
+            }
+            // GGML_ASSERT(class_token_index.size() == 1); // PhotoMaker currently does not support multiple
+            //     trigger words in a single prompt.
+            if (class_token_index.size() == 1) {
+                // Expand the class word token and corresponding mask
+                int class_token = clean_input_ids[class_token_index[0]];
+                class_idx       = tokens_acc + class_token_index[0];
+                std::vector<int> clean_input_ids_tmp;
+                for (uint32_t i = 0; i < class_token_index[0]; i++)
+                    clean_input_ids_tmp.push_back(clean_input_ids[i]);
+                for (uint32_t i = 0; i < num_input_imgs; i++)
+                    clean_input_ids_tmp.push_back(class_token);
+                for (uint32_t i = class_token_index[0] + 1; i < clean_input_ids.size(); i++)
+                    clean_input_ids_tmp.push_back(clean_input_ids[i]);
+                clean_input_ids.clear();
+                clean_input_ids = clean_input_ids_tmp;
+            }
+            tokens_acc += clean_index;
+            tokens.insert(tokens.end(), clean_input_ids.begin(), clean_input_ids.end());
+            weights.insert(weights.end(), clean_input_ids.size(), curr_weight);
+        }
+        tokens.insert(tokens.begin(), BOS_TOKEN_ID);
+        weights.insert(weights.begin(), 1.0);
+
+        pad_tokens(tokens, weights, max_length, padding);
+
+        for (uint32_t i = 0; i < tokens.size(); i++) {
+            if (class_idx + 1 <= i && i < class_idx + 1 + num_input_imgs)
+                class_token_mask.push_back(true);
+            else
+                class_token_mask.push_back(false);
+        }
+
+        // printf("[");
+        // for (int i = 0; i < tokens.size(); i++) {
+        //     printf("%d, ", class_token_mask[i] ? 1 : 0);
+        // }
+        // printf("]\n");
+
+        // for (int i = 0; i < tokens.size(); i++) {
+        //     std::cout << tokens[i] << ":" << weights[i] << ", ";
+        // }
+        // std::cout << std::endl;
+
+        return std::make_tuple(tokens, weights, class_token_mask);
+    }
+
    std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
                                                             size_t max_length = 0,
                                                             bool padding      = false) {
@@ -1087,28 +1361,8 @@ 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);
-                }
-            }
-        }
+        pad_tokens(tokens, weights, max_length, padding);

        // for (int i = 0; i < tokens.size(); i++) {
        //     std::cout << tokens[i] << ":" << weights[i] << ", ";
@@ -1120,10 +1374,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
 };

 struct FrozenCLIPVisionEmbedder : public GGMLModule {
-    CLIPVisionModel vision_model;
+    CLIPVisionModelProjection vision_model;

    FrozenCLIPVisionEmbedder(ggml_backend_t backend, ggml_type wtype)
-        : GGMLModule(backend, wtype) {
+        : vision_model(OPEN_CLIP_VIT_H_14, true), GGMLModule(backend, wtype) {
        vision_model.init(params_ctx, wtype);
    }

@@ -1140,11 +1394,10 @@ struct FrozenCLIPVisionEmbedder : public GGMLModule {
    }

    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
-        vision_model.get_param_tensors(tensors, prefix + "transformer.visual_model");
+        vision_model.get_param_tensors(tensors, prefix + "transformer");
    }

-    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 +1409,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);
    }
--- a/control.hpp
+++ b/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;
    }

--- a/examples/cli/main.cpp
+++ b/examples/cli/main.cpp
@@ -6,10 +6,11 @@
 #include <string>
 #include <vector>

-#include "preprocessing.hpp"
+// #include "preprocessing.hpp"
 #include "stable-diffusion.h"

 #define STB_IMAGE_IMPLEMENTATION
+#define STB_IMAGE_STATIC
 #include "stb_image.h"

 #define STB_IMAGE_WRITE_IMPLEMENTATION
@@ -65,6 +66,8 @@ struct SDParams {
    std::string esrgan_path;
    std::string controlnet_path;
    std::string embeddings_path;
+    std::string stacked_id_embeddings_path;
+    std::string input_id_images_path;
    sd_type_t wtype = SD_TYPE_COUNT;
    std::string lora_model_dir;
    std::string output_path = "output.png";
@@ -73,12 +76,13 @@ struct SDParams {

    std::string prompt;
    std::string negative_prompt;
-    float min_cfg   = 1.0f;
-    float cfg_scale = 7.0f;
-    int clip_skip   = -1;  // <= 0 represents unspecified
-    int width       = 512;
-    int height      = 512;
-    int batch_count = 1;
+    float min_cfg     = 1.0f;
+    float cfg_scale   = 7.0f;
+    float style_ratio = 20.f;
+    int clip_skip     = -1;  // <= 0 represents unspecified
+    int width         = 512;
+    int height        = 512;
+    int batch_count   = 1;

    int video_frames         = 6;
    int motion_bucket_id     = 127;
@@ -95,7 +99,11 @@ struct SDParams {
    bool verbose                  = false;
    bool vae_tiling               = false;
    bool control_net_cpu          = false;
+    bool normalize_input          = false;
+    bool clip_on_cpu              = false;
+    bool vae_on_cpu               = false;
    bool canny_preprocess         = false;
+    bool color                    = false;
    int upscale_repeats           = 1;
 };

@@ -110,10 +118,16 @@ void print_params(SDParams params) {
    printf("    esrgan_path:       %s\n", params.esrgan_path.c_str());
    printf("    controlnet_path:   %s\n", params.controlnet_path.c_str());
    printf("    embeddings_path:   %s\n", params.embeddings_path.c_str());
+    printf("    stacked_id_embeddings_path:   %s\n", params.stacked_id_embeddings_path.c_str());
+    printf("    input_id_images_path:   %s\n", params.input_id_images_path.c_str());
+    printf("    style ratio:       %.2f\n", params.style_ratio);
+    printf("    normzalize input image :  %s\n", params.normalize_input ? "true" : "false");
    printf("    output_path:       %s\n", params.output_path.c_str());
    printf("    init_img:          %s\n", params.input_path.c_str());
    printf("    control_image:     %s\n", params.control_image_path.c_str());
+    printf("    clip on cpu:       %s\n", params.clip_on_cpu ? "true" : "false");
    printf("    controlnet cpu:    %s\n", params.control_net_cpu ? "true" : "false");
+    printf("    vae decoder on cpu:%s\n", params.vae_on_cpu ? "true" : "false");
    printf("    strength(control): %.2f\n", params.control_strength);
    printf("    prompt:            %s\n", params.prompt.c_str());
    printf("    negative_prompt:   %s\n", params.negative_prompt.c_str());
@@ -146,6 +160,9 @@ void print_usage(int argc, const char* argv[]) {
    printf("  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n");
    printf("  --control-net [CONTROL_PATH]       path to control net model\n");
    printf("  --embd-dir [EMBEDDING_PATH]        path to embeddings.\n");
+    printf("  --stacked-id-embd-dir [DIR]        path to PHOTOMAKER stacked id embeddings.\n");
+    printf("  --input-id-images-dir [DIR]        path to PHOTOMAKER input id images dir.\n");
+    printf("  --normalize-input                  normalize PHOTOMAKER input id images\n");
    printf("  --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n");
    printf("  --upscale-repeats                  Run the ESRGAN upscaler this many times (default 1)\n");
    printf("  --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n");
@@ -158,6 +175,7 @@ void print_usage(int argc, const char* argv[]) {
    printf("  -n, --negative-prompt PROMPT       the negative prompt (default: \"\")\n");
    printf("  --cfg-scale SCALE                  unconditional guidance scale: (default: 7.0)\n");
    printf("  --strength STRENGTH                strength for noising/unnoising (default: 0.75)\n");
+    printf("  --style-ratio STYLE-RATIO          strength for keeping input identity (default: 20%%)\n");
    printf("  --control-strength STRENGTH        strength to apply Control Net (default: 0.9)\n");
    printf("                                     1.0 corresponds to full destruction of information in init image\n");
    printf("  -H, --height H                     image height, in pixel space (default: 512)\n");
@@ -244,6 +262,18 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                break;
            }
            params.embeddings_path = argv[i];
+        } else if (arg == "--stacked-id-embd-dir") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.stacked_id_embeddings_path = argv[i];
+        } else if (arg == "--input-id-images-dir") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.input_id_images_path = argv[i];
        } else if (arg == "--type") {
            if (++i >= argc) {
                invalid_arg = true;
@@ -327,6 +357,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                break;
            }
            params.strength = std::stof(argv[i]);
+        } else if (arg == "--style-ratio") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.style_ratio = std::stof(argv[i]);
        } else if (arg == "--control-strength") {
            if (++i >= argc) {
                invalid_arg = true;
@@ -361,6 +397,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
            params.vae_tiling = true;
        } else if (arg == "--control-net-cpu") {
            params.control_net_cpu = true;
+        } else if (arg == "--normalize-input") {
+            params.normalize_input = true;
+        } else if (arg == "--clip-on-cpu") {
+            params.clip_on_cpu = true;  // will slow down get_learned_condiotion but necessary for low MEM GPUs
+        } else if (arg == "--vae-on-cpu") {
+            params.vae_on_cpu = true;  // will slow down latent decoding but necessary for low MEM GPUs
        } else if (arg == "--canny") {
            params.canny_preprocess = true;
        } else if (arg == "-b" || arg == "--batch-count") {
@@ -428,6 +470,8 @@ void parse_args(int argc, const char** argv, SDParams& params) {
            exit(0);
        } else if (arg == "-v" || arg == "--verbose") {
            params.verbose = true;
+        } else if (arg == "--color") {
+            params.color = true;
        } else {
            fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
            print_usage(argc, argv);
@@ -499,6 +543,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) {
@@ -519,18 +575,47 @@ std::string get_image_params(SDParams params, int64_t seed) {
    return parameter_string;
 }

+/* Enables Printing the log level tag in color using ANSI escape codes */
 void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
    SDParams* params = (SDParams*)data;
-    if (!params->verbose && level <= SD_LOG_DEBUG) {
+    int tag_color;
+    const char* level_str;
+    FILE* out_stream = (level == SD_LOG_ERROR) ? stderr : stdout;
+
+    if (!log || (!params->verbose && level <= SD_LOG_DEBUG)) {
        return;
    }
-    if (level <= SD_LOG_INFO) {
-        fputs(log, stdout);
-        fflush(stdout);
-    } else {
-        fputs(log, stderr);
-        fflush(stderr);
+
+    switch (level) {
+        case SD_LOG_DEBUG:
+            tag_color = 37;
+            level_str = "DEBUG";
+            break;
+        case SD_LOG_INFO:
+            tag_color = 34;
+            level_str = "INFO";
+            break;
+        case SD_LOG_WARN:
+            tag_color = 35;
+            level_str = "WARN";
+            break;
+        case SD_LOG_ERROR:
+            tag_color = 31;
+            level_str = "ERROR";
+            break;
+        default: /* Potential future-proofing */
+            tag_color = 33;
+            level_str = "?????";
+            break;
    }
+
+    if (params->color == true) {
+        fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str);
+    } else {
+        fprintf(out_stream, "[%-5s] ", level_str);
+    }
+    fputs(log, out_stream);
+    fflush(out_stream);
 }

 int main(int argc, const char* argv[]) {
@@ -601,6 +686,7 @@ int main(int argc, const char* argv[]) {
                                  params.controlnet_path.c_str(),
                                  params.lora_model_dir.c_str(),
                                  params.embeddings_path.c_str(),
+                                  params.stacked_id_embeddings_path.c_str(),
                                  vae_decode_only,
                                  params.vae_tiling,
                                  true,
@@ -608,7 +694,9 @@ int main(int argc, const char* argv[]) {
                                  params.wtype,
                                  params.rng_type,
                                  params.schedule,
-                                  params.control_net_cpu);
+                                  params.clip_on_cpu,
+                                  params.control_net_cpu,
+                                  params.vae_on_cpu);

    if (sd_ctx == NULL) {
        printf("new_sd_ctx_t failed\n");
@@ -630,8 +718,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,
@@ -646,7 +740,10 @@ int main(int argc, const char* argv[]) {
                          params.seed,
                          params.batch_count,
                          control_image,
-                          params.control_strength);
+                          params.control_strength,
+                          params.style_ratio,
+                          params.normalize_input,
+                          params.input_id_images_path.c_str());
    } else {
        sd_image_t input_image = {(uint32_t)params.width,
                                  (uint32_t)params.height,
--- a/2
+++ b/2
--- a/ggml_extend.hpp
+++ b/ggml_extend.hpp
@@ -80,8 +80,27 @@ __STATIC_INLINE__ ggml_fp16_t ggml_tensor_get_f16(const ggml_tensor* tensor, int
    return *(ggml_fp16_t*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]);
 }

-__STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
-    printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+static struct ggml_tensor* get_tensor_from_graph(struct ggml_cgraph* gf, const char* name) {
+    struct ggml_tensor* res = NULL;
+    for (int i = 0; i < gf->n_nodes; i++) {
+        // printf("%d, %s \n", i, gf->nodes[i]->name);
+        if (strcmp(ggml_get_name(gf->nodes[i]), name) == 0) {
+            res = gf->nodes[i];
+            break;
+        }
+    }
+    for (int i = 0; i < gf->n_leafs; i++) {
+        // printf("%d, %s \n", i, gf->leafs[i]->name);
+        if (strcmp(ggml_get_name(gf->leafs[i]), name) == 0) {
+            res = gf->leafs[i];
+            break;
+        }
+    }
+    return res;
+}
+
+__STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false, const char* mark = "") {
+    printf("%s (%s): shape(%zu, %zu, %zu, %zu)\n", mark, ggml_type_name(tensor->type), tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
    fflush(stdout);
    if (shape_only) {
        return;
@@ -217,6 +236,23 @@ __STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input) {
    return image_data;
 }

+__STATIC_INLINE__ uint8_t* sd_tensor_to_mul_image(struct ggml_tensor* input, int idx) {
+    int64_t width    = input->ne[0];
+    int64_t height   = input->ne[1];
+    int64_t channels = input->ne[2];
+    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    uint8_t* image_data = (uint8_t*)malloc(width * height * channels);
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            for (int k = 0; k < channels; k++) {
+                float value                                               = ggml_tensor_get_f32(input, ix, iy, k, idx);
+                *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
 __STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
                                          struct ggml_tensor* output,
                                          bool scale = true) {
@@ -237,6 +273,28 @@ __STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
    }
 }

+__STATIC_INLINE__ void sd_mul_images_to_tensor(const uint8_t* image_data,
+                                               struct ggml_tensor* output,
+                                               int idx,
+                                               float* mean = NULL,
+                                               float* std  = NULL) {
+    int64_t width    = output->ne[0];
+    int64_t height   = output->ne[1];
+    int64_t channels = output->ne[2];
+    GGML_ASSERT(channels == 3 && output->type == GGML_TYPE_F32);
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            for (int k = 0; k < channels; k++) {
+                int value       = *(image_data + iy * width * channels + ix * channels + k);
+                float pixel_val = value / 255.0f;
+                if (mean != NULL && std != NULL)
+                    pixel_val = (pixel_val - mean[k]) / std[k];
+                ggml_tensor_set_f32(output, pixel_val, ix, iy, k, idx);
+            }
+        }
+    }
+}
+
 __STATIC_INLINE__ void sd_image_f32_to_tensor(const float* image_data,
                                              struct ggml_tensor* output,
                                              bool scale = true) {
@@ -247,7 +305,7 @@ __STATIC_INLINE__ void sd_image_f32_to_tensor(const float* image_data,
    for (int iy = 0; iy < height; iy++) {
        for (int ix = 0; ix < width; ix++) {
            for (int k = 0; k < channels; k++) {
-                float value = *(image_data + iy * width * channels + ix * channels + k);
+                int value = *(image_data + iy * width * channels + ix * channels + k);
                if (scale) {
                    value /= 255.f;
                }
@@ -606,6 +664,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 +724,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 +735,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;
@@ -683,8 +759,13 @@ __STATIC_INLINE__ struct ggml_tensor* new_timestep_embedding(struct ggml_context
 //     virtual struct ggml_cgraph* get_ggml_cgraph() = 0;
 // };

+/*
 #define MAX_PARAMS_TENSOR_NUM 10240
 #define MAX_GRAPH_SIZE 10240
+*/
+/* SDXL with LoRA requires more space */
+#define MAX_PARAMS_TENSOR_NUM 15360
+#define MAX_GRAPH_SIZE 15360

 struct GGMLModule {
 protected:
@@ -693,9 +774,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 +816,40 @@ 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;
+        size_t compute_buffer_size = ggml_gallocr_get_buffer_size(compute_allocr, 0);
+        LOG_DEBUG("%s compute buffer size: %.2f MB(%s)",
+                  get_desc().c_str(),
+                  compute_buffer_size / 1024.0 / 1024.0,
+                  ggml_backend_is_cpu(backend) ? "RAM" : "VRAM");
+        return true;
+    }

-        // recreate the allocator with the required memory
-        ggml_allocr_free(compute_allocr);
+    void cpy_data_to_backend_tensor() {
+        for (auto& kv : backend_tensor_data_map) {
+            auto tensor = kv.first;
+            auto data   = kv.second;

-        LOG_DEBUG("%s compute buffer size: %.2f MB", get_desc().c_str(), compute_buffer_size / 1024.0 / 1024.0);
+            ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+        }

-        compute_buffer = ggml_backend_alloc_buffer(backend, compute_buffer_size);
-        compute_allocr = ggml_allocr_new_from_buffer(compute_buffer);
+        backend_tensor_data_map.clear();
    }

 public:
@@ -775,31 +874,19 @@ 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();
-
-        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);
-            }
+        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;
        }
-        ggml_allocr_free(alloc);
+        size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer);
+        LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)",
+                  get_desc().c_str(),
+                  params_buffer_size / (1024.0 * 1024.0),
+                  ggml_backend_is_cpu(backend) ? "RAM" : "VRAM",
+                  num_tensors);
        return true;
    }

@@ -812,13 +899,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 +915,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_TYPE_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,12 +931,11 @@ 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);
        }
@@ -862,13 +945,11 @@ public:
            ggml_backend_metal_set_n_cb(backend, n_threads);
        }
 #endif
-
        ggml_backend_graph_compute(backend, gf);

 #ifdef GGML_PERF
        ggml_graph_print(gf);
 #endif
-
        if (output != NULL) {
            auto result = gf->nodes[gf->n_nodes - 1];
            if (*output == NULL && output_ctx != NULL) {
@@ -960,13 +1041,11 @@ public:
        }
        for (auto& pair : blocks) {
            auto& block = pair.second;
-
            block->get_param_tensors(tensors, prefix + pair.first);
        }

        for (auto& pair : params) {
-            struct ggml_tensor* param = pair.second;
-
+            struct ggml_tensor* param    = pair.second;
            tensors[prefix + pair.first] = pair.second;
        }
    }
@@ -1214,9 +1293,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]
@@ -1226,13 +1305,12 @@ public:
        struct ggml_tensor* kqv = ggml_nn_attention(ctx, q, k, v, mask);  // [N * n_head, n_token, d_head]

        kqv = ggml_reshape_4d(ctx, kqv, d_head, n_token, n_head, N);
-        kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));  // [N, n_token, n_head, d_head]
+        kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3));      // [N, n_token, n_head, d_head]
+        x   = ggml_reshape_3d(ctx, kqv, d_head * n_head, n_token, N);  // [N, n_token, d_head * n_head]

-        x = ggml_reshape_2d(ctx, kqv, d_head * n_head, n_token * N);  // [N * n_token, d_head * n_head]
-
-        x = out_proj->forward(ctx, x);
+        x = out_proj->forward(ctx, x);  // [N, n_token, embed_dim]
        return x;
    }
 };

-#endif  // __GGML_EXTEND__HPP__
+#endif  // __GGML_EXTEND__HPP__
--- a/lora.hpp
+++ b/lora.hpp
@@ -14,9 +14,10 @@ struct LoraModel : public GGMLModule {

    LoraModel(ggml_backend_t backend,
              ggml_type wtype,
-              const std::string file_path = "")
+              const std::string& file_path = "",
+              const std::string& prefix    = "")
        : file_path(file_path), GGMLModule(backend, wtype) {
-        if (!model_loader.init_from_file(file_path)) {
+        if (!model_loader.init_from_file(file_path, prefix)) {
            load_failed = true;
        }
    }
@@ -33,37 +34,48 @@ struct LoraModel : public GGMLModule {
        return model_loader.get_params_mem_size(NULL);
    }

-    bool load_from_file() {
+    bool load_from_file(bool filter_tensor = false) {
        LOG_INFO("loading LoRA from '%s'", file_path.c_str());

        if (load_failed) {
            LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
            return false;
        }
-        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);
+            if (filter_tensor && !contains(name, "lora")) {
+                // LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
+                return true;
+            }

-            *dst_tensor = 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;
+            }

-            lora_tensors[name] = real;
            return true;
        };

+        model_loader.load_tensors(on_new_tensor_cb, backend);
+        alloc_params_buffer();
+
+        dry_run = false;
        model_loader.load_tensors(on_new_tensor_cb, backend);

        LOG_DEBUG("finished loaded lora");
-        ggml_allocr_free(alloc);
        return true;
    }

-    struct ggml_cgraph* build_graph(std::map<std::string, struct ggml_tensor*> model_tensors) {
+    struct ggml_cgraph* build_lora_graph(std::map<std::string, struct ggml_tensor*> model_tensors) {
        struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, LORA_GRAPH_SIZE, false);

        std::set<std::string> applied_lora_tensors;
@@ -77,6 +89,10 @@ struct LoraModel : public GGMLModule {
            }
            k_tensor = k_tensor.substr(0, k_pos);
            replace_all_chars(k_tensor, '.', '_');
+            // LOG_DEBUG("k_tensor %s", k_tensor.c_str());
+            if (k_tensor == "model_diffusion_model_output_blocks_2_2_conv") {  // fix for SDXL
+                k_tensor = "model_diffusion_model_output_blocks_2_1_conv";
+            }
            std::string lora_up_name   = "lora." + k_tensor + ".lora_up.weight";
            std::string lora_down_name = "lora." + k_tensor + ".lora_down.weight";
            std::string alpha_name     = "lora." + k_tensor + ".alpha";
@@ -139,21 +155,37 @@ struct LoraModel : public GGMLModule {
            ggml_build_forward_expand(gf, final_weight);
        }

+        size_t total_lora_tensors_count   = 0;
+        size_t applied_lora_tensors_count = 0;
+
        for (auto& kv : lora_tensors) {
+            total_lora_tensors_count++;
            if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
                LOG_WARN("unused lora tensor %s", kv.first.c_str());
+            } else {
+                applied_lora_tensors_count++;
            }
        }
+        /* Don't worry if this message shows up twice in the logs per LoRA,
+         * this function is called once to calculate the required buffer size
+         * and then again to actually generate a graph to be used */
+        if (applied_lora_tensors_count != total_lora_tensors_count) {
+            LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied",
+                     applied_lora_tensors_count, total_lora_tensors_count);
+        } else {
+            LOG_DEBUG("(%lu / %lu) LoRA tensors applied successfully",
+                      applied_lora_tensors_count, total_lora_tensors_count);
+        }

        return gf;
    }

    void apply(std::map<std::string, struct ggml_tensor*> model_tensors, int n_threads) {
        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(model_tensors);
+            return build_lora_graph(model_tensors);
        };
        GGMLModule::compute(get_graph, n_threads, true);
    }
 };

-#endif  // __LORA_HPP__
+#endif  // __LORA_HPP__
--- a/model.cpp
+++ b/model.cpp
@@ -108,14 +108,14 @@ std::unordered_map<std::string, std::string> open_clip_to_hf_clip_model = {
    {"model.positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"},
    {"model.token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"},
    {"model.text_projection", "transformer.text_model.text_projection"},
-    {"model.visual.class_embedding", "transformer.visual_model.embeddings.class_embedding"},
-    {"model.visual.conv1.weight", "transformer.visual_model.embeddings.patch_embedding.weight"},
-    {"model.visual.ln_post.bias", "transformer.visual_model.post_layernorm.bias"},
-    {"model.visual.ln_post.weight", "transformer.visual_model.post_layernorm.weight"},
-    {"model.visual.ln_pre.bias", "transformer.visual_model.pre_layernorm.bias"},
-    {"model.visual.ln_pre.weight", "transformer.visual_model.pre_layernorm.weight"},
-    {"model.visual.positional_embedding", "transformer.visual_model.embeddings.position_embedding.weight"},
-    {"model.visual.proj", "transformer.visual_model.visual_projection"},
+    {"model.visual.class_embedding", "transformer.vision_model.embeddings.class_embedding"},
+    {"model.visual.conv1.weight", "transformer.vision_model.embeddings.patch_embedding.weight"},
+    {"model.visual.ln_post.bias", "transformer.vision_model.post_layernorm.bias"},
+    {"model.visual.ln_post.weight", "transformer.vision_model.post_layernorm.weight"},
+    {"model.visual.ln_pre.bias", "transformer.vision_model.pre_layernorm.bias"},
+    {"model.visual.ln_pre.weight", "transformer.vision_model.pre_layernorm.weight"},
+    {"model.visual.positional_embedding", "transformer.vision_model.embeddings.position_embedding.weight"},
+    {"model.visual.proj", "transformer.visual_projection.weight"},
 };

 std::unordered_map<std::string, std::string> open_clip_to_hk_clip_resblock = {
@@ -157,6 +157,10 @@ std::string convert_open_clip_to_hf_clip(const std::string& name) {
    } else if (starts_with(new_name, "cond_stage_model.")) {
        prefix   = "cond_stage_model.";
        new_name = new_name.substr(strlen("cond_stage_model."));
+    } else if (ends_with(new_name, "vision_model.visual_projection.weight")) {
+        prefix   = new_name.substr(0, new_name.size() - strlen("vision_model.visual_projection.weight"));
+        new_name = prefix + "visual_projection.weight";
+        return new_name;
    } else {
        return new_name;
    }
@@ -186,7 +190,7 @@ std::string convert_open_clip_to_hf_clip(const std::string& name) {
    replace_suffix();

    open_clip_resblock_prefix = "model.visual.transformer.resblocks.";
-    hf_clip_resblock_prefix   = "transformer.visual_model.encoder.layers.";
+    hf_clip_resblock_prefix   = "transformer.vision_model.encoder.layers.";

    replace_suffix();

@@ -200,6 +204,23 @@ std::string convert_vae_decoder_name(const std::string& name) {
    return name;
 }

+/* If not a SDXL LoRA the unet" prefix will have already been replaced by this
+ * point and "te2" and "te1" don't seem to appear in non-SDXL only "te_" */
+std::string convert_sdxl_lora_name(std::string tensor_name) {
+    const std::pair<std::string, std::string> sdxl_lora_name_lookup[] = {
+        {"unet", "model_diffusion_model"},
+        {"te2", "cond_stage_model_1_transformer"},
+        {"te1", "cond_stage_model_transformer"},
+    };
+    for (auto& pair_i : sdxl_lora_name_lookup) {
+        if (tensor_name.compare(0, pair_i.first.length(), pair_i.first) == 0) {
+            tensor_name = std::regex_replace(tensor_name, std::regex(pair_i.first), pair_i.second);
+            break;
+        }
+    }
+    return tensor_name;
+}
+
 std::unordered_map<std::string, std::unordered_map<std::string, std::string>> suffix_conversion_underline = {
    {
        "attentions",
@@ -248,7 +269,7 @@ std::unordered_map<std::string, std::unordered_map<std::string, std::string>> su
    },
 };

-std::string convert_diffusers_name_to_compvis(const std::string& key, char seq) {
+std::string convert_diffusers_name_to_compvis(std::string key, char seq) {
    std::vector<std::string> m;

    auto match = [](std::vector<std::string>& match_list, const std::regex& regex, const std::string& key) {
@@ -282,6 +303,11 @@ std::string convert_diffusers_name_to_compvis(const std::string& key, char seq)
        return inner_key;
    };

+    // convert attn to out
+    if (ends_with(key, "to_out")) {
+        key += format("%c0", seq);
+    }
+
    // unet
    if (match(m, std::regex(format("unet%cconv_in(.*)", seq)), key)) {
        return format("model%cdiffusion_model%cinput_blocks%c0%c0", seq, seq, seq, seq) + m[0];
@@ -391,8 +417,8 @@ std::string convert_diffusers_name_to_compvis(const std::string& key, char seq)
 }

 std::string convert_tensor_name(const std::string& name) {
-    std::string new_name;
-    if (starts_with(name, "cond_stage_model.") || starts_with(name, "conditioner.embedders.")) {
+    std::string new_name = name;
+    if (starts_with(name, "cond_stage_model.") || starts_with(name, "conditioner.embedders.") || ends_with(name, ".vision_model.visual_projection.weight")) {
        new_name = convert_open_clip_to_hf_clip(name);
    } else if (starts_with(name, "first_stage_model.decoder")) {
        new_name = convert_vae_decoder_name(name);
@@ -406,8 +432,12 @@ std::string convert_tensor_name(const std::string& name) {
        if (pos != std::string::npos) {
            std::string name_without_network_parts = name.substr(5, pos - 5);
            std::string network_part               = name.substr(pos + 1);
+
            // LOG_DEBUG("%s %s", name_without_network_parts.c_str(), network_part.c_str());
            std::string new_key = convert_diffusers_name_to_compvis(name_without_network_parts, '_');
+            /* For dealing with the new SDXL LoRA tensor naming convention */
+            new_key = convert_sdxl_lora_name(new_key);
+
            if (new_key.empty()) {
                new_name = name;
            } else {
@@ -416,6 +446,26 @@ std::string convert_tensor_name(const std::string& name) {
        } else {
            new_name = name;
        }
+    } else if (contains(name, "lora_up") || contains(name, "lora_down") || contains(name, "lora.up") || contains(name, "lora.down")) {
+        size_t pos = new_name.find(".processor");
+        if (pos != std::string::npos) {
+            new_name.replace(pos, strlen(".processor"), "");
+        }
+        pos = new_name.find_last_of('_');
+        if (pos != std::string::npos) {
+            std::string name_without_network_parts = new_name.substr(0, pos);
+            std::string network_part               = new_name.substr(pos + 1);
+            // LOG_DEBUG("%s %s", name_without_network_parts.c_str(), network_part.c_str());
+            std::string new_key = convert_diffusers_name_to_compvis(name_without_network_parts, '.');
+            replace_all_chars(new_key, '.', '_');
+            if (starts_with(network_part, "lora.")) {
+                network_part = "lora_" + network_part.substr(5);
+            }
+            if (new_key.size() > 0) {
+                new_name = "lora." + new_key + "." + network_part;
+            }
+            // LOG_DEBUG("new name: %s", new_name.c_str());
+        }
    } else if (starts_with(name, "unet") || starts_with(name, "vae") || starts_with(name, "te")) {  // for diffuser
        size_t pos = name.find_last_of('.');
        if (pos != std::string::npos) {
@@ -498,7 +548,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 +563,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 +594,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);
        }
    }
 }
@@ -820,7 +880,6 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
        }

        TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
-
        tensor_storage.reverse_ne();

        size_t tensor_data_size = end - begin;
@@ -1159,7 +1218,9 @@ bool ModelLoader::parse_data_pkl(uint8_t* buffer,
                    if (reader.phase == PickleTensorReader::READ_DIMENS) {
                        reader.tensor_storage.reverse_ne();
                        reader.tensor_storage.file_index = file_index;
-                        reader.tensor_storage.name       = prefix + reader.tensor_storage.name;
+                        // if(strcmp(prefix.c_str(), "scarlett") == 0)
+                        // printf(" got tensor %s \n ", reader.tensor_storage.name.c_str());
+                        reader.tensor_storage.name = prefix + reader.tensor_storage.name;
                        tensor_storages.push_back(reader.tensor_storage);
                        // LOG_DEBUG("%s", reader.tensor_storage.name.c_str());
                        // reset
@@ -1262,7 +1323,8 @@ std::string ModelLoader::load_merges() {
    return merges_utf8_str;
 }

-void remove_duplicates(std::vector<TensorStorage>& vec) {
+std::vector<TensorStorage> remove_duplicates(const std::vector<TensorStorage>& vec) {
+    std::vector<TensorStorage> res;
    std::unordered_map<std::string, size_t> name_to_index_map;

    for (size_t i = 0; i < vec.size(); ++i) {
@@ -1270,13 +1332,16 @@ void remove_duplicates(std::vector<TensorStorage>& vec) {
        auto it                         = name_to_index_map.find(current_name);

        if (it != name_to_index_map.end()) {
-            vec[it->second] = vec[i];
+            res[it->second] = vec[i];
        } else {
            name_to_index_map[current_name] = i;
+            res.push_back(vec[i]);
        }
    }

-    vec.resize(name_to_index_map.size());
+    // vec.resize(name_to_index_map.size());
+
+    return res;
 }

 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend_t backend) {
@@ -1290,7 +1355,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend

        preprocess_tensor(tensor_storage, processed_tensor_storages);
    }
-    remove_duplicates(processed_tensor_storages);
+    std::vector<TensorStorage> dedup = remove_duplicates(processed_tensor_storages);
+    processed_tensor_storages        = dedup;
+
    bool success = true;
    for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
        std::string file_path = file_paths_[file_index];
@@ -1352,7 +1419,6 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
            if (tensor_storage.file_index != file_index) {
                continue;
            }
-
            ggml_tensor* dst_tensor = NULL;

            success = on_new_tensor_cb(tensor_storage, &dst_tensor);
@@ -1387,7 +1453,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 +1472,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));
                }
            }
@@ -1596,4 +1662,4 @@ bool convert(const char* input_path, const char* vae_path, const char* output_pa
    }
    bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type);
    return success;
-}
+}
--- a/model.h
+++ b/model.h
@@ -7,6 +7,7 @@
 #include <set>
 #include <sstream>
 #include <string>
+#include <tuple>
 #include <vector>

 #include "ggml/ggml-backend.h"
--- a/pmid.hpp
+++ b/pmid.hpp
@@ -0,0 +1,305 @@
+#ifndef __PMI_HPP__
+#define __PMI_HPP__
+
+#include "ggml_extend.hpp"
+
+#include "clip.hpp"
+#include "lora.hpp"
+
+struct FuseBlock : public GGMLBlock {
+    // network hparams
+    int in_dim;
+    int out_dim;
+    int hidden_dim;
+    bool use_residue;
+
+public:
+    FuseBlock(int i_d, int o_d, int h_d, bool use_residue = true)
+        : in_dim(i_d), out_dim(o_d), hidden_dim(h_d), use_residue(use_residue) {
+        blocks["fc1"]       = std::shared_ptr<GGMLBlock>(new Linear(in_dim, hidden_dim, true));
+        blocks["fc2"]       = std::shared_ptr<GGMLBlock>(new Linear(hidden_dim, out_dim, true));
+        blocks["layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(in_dim));
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [N, channels, h, w]
+
+        auto fc1        = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
+        auto fc2        = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
+        auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
+
+        struct ggml_tensor* r = x;
+        // x = ggml_nn_layer_norm(ctx, x, ln_w, ln_b);
+        x = layer_norm->forward(ctx, x);
+        // x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x),  fc1_b);
+        x = fc1->forward(ctx, x);
+        x = ggml_gelu_inplace(ctx, x);
+        x = fc2->forward(ctx, x);
+        // x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x),  fc2_b);
+        if (use_residue)
+            x = ggml_add(ctx, x, r);
+        return x;
+    }
+};
+
+struct FuseModule : public GGMLBlock {
+    // network hparams
+    int embed_dim;
+
+public:
+    FuseModule(int imb_d)
+        : embed_dim(imb_d) {
+        blocks["mlp1"]       = std::shared_ptr<GGMLBlock>(new FuseBlock(imb_d * 2, imb_d, imb_d, false));
+        blocks["mlp2"]       = std::shared_ptr<GGMLBlock>(new FuseBlock(imb_d, imb_d, imb_d, true));
+        blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(embed_dim));
+    }
+
+    struct ggml_tensor* fuse_fn(struct ggml_context* ctx,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* id_embeds) {
+        auto mlp1       = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp1"]);
+        auto mlp2       = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp2"]);
+        auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
+
+        auto prompt_embeds0 = ggml_cont(ctx, ggml_permute(ctx, prompt_embeds, 2, 0, 1, 3));
+        auto id_embeds0     = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 2, 0, 1, 3));
+        // concat is along dim 2
+        auto stacked_id_embeds = ggml_concat(ctx, prompt_embeds0, id_embeds0);
+        stacked_id_embeds      = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 1, 2, 0, 3));
+
+        // stacked_id_embeds = mlp1.forward(ctx, stacked_id_embeds);
+        // stacked_id_embeds = ggml_add(ctx, stacked_id_embeds, prompt_embeds);
+        // stacked_id_embeds = mlp2.forward(ctx, stacked_id_embeds);
+        // stacked_id_embeds = ggml_nn_layer_norm(ctx, stacked_id_embeds, ln_w, ln_b);
+
+        stacked_id_embeds = mlp1->forward(ctx, stacked_id_embeds);
+        stacked_id_embeds = ggml_add(ctx, stacked_id_embeds, prompt_embeds);
+        stacked_id_embeds = mlp2->forward(ctx, stacked_id_embeds);
+        stacked_id_embeds = layer_norm->forward(ctx, stacked_id_embeds);
+
+        return stacked_id_embeds;
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* id_embeds,
+                                struct ggml_tensor* class_tokens_mask,
+                                struct ggml_tensor* class_tokens_mask_pos,
+                                struct ggml_tensor* left,
+                                struct ggml_tensor* right) {
+        // x: [N, channels, h, w]
+
+        struct ggml_tensor* valid_id_embeds = id_embeds;
+        // # slice out the image token embeddings
+        // print_ggml_tensor(class_tokens_mask_pos, false);
+        ggml_set_name(class_tokens_mask_pos, "class_tokens_mask_pos");
+        ggml_set_name(prompt_embeds, "prompt_embeds");
+        // print_ggml_tensor(valid_id_embeds, true, "valid_id_embeds");
+        // print_ggml_tensor(class_tokens_mask_pos, true, "class_tokens_mask_pos");
+        struct ggml_tensor* image_token_embeds = ggml_get_rows(ctx, prompt_embeds, class_tokens_mask_pos);
+        ggml_set_name(image_token_embeds, "image_token_embeds");
+        struct ggml_tensor* stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds);
+
+        stacked_id_embeds = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 0, 2, 1, 3));
+        if (left && right) {
+            stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds);
+            stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right);
+        } else if (left) {
+            stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds);
+        } else if (right) {
+            stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right);
+        }
+        stacked_id_embeds                         = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 0, 2, 1, 3));
+        class_tokens_mask                         = ggml_cont(ctx, ggml_transpose(ctx, class_tokens_mask));
+        class_tokens_mask                         = ggml_repeat(ctx, class_tokens_mask, prompt_embeds);
+        prompt_embeds                             = ggml_mul(ctx, prompt_embeds, class_tokens_mask);
+        struct ggml_tensor* updated_prompt_embeds = ggml_add(ctx, prompt_embeds, stacked_id_embeds);
+        ggml_set_name(updated_prompt_embeds, "updated_prompt_embeds");
+        return updated_prompt_embeds;
+    }
+};
+
+struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
+    PhotoMakerIDEncoderBlock()
+        : CLIPVisionModelProjection(OPENAI_CLIP_VIT_L_14) {
+        blocks["visual_projection_2"] = std::shared_ptr<GGMLBlock>(new Linear(1024, 1280, false));
+        blocks["fuse_module"]         = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx,
+                                struct ggml_tensor* id_pixel_values,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* class_tokens_mask,
+                                struct ggml_tensor* class_tokens_mask_pos,
+                                struct ggml_tensor* left,
+                                struct ggml_tensor* right) {
+        // x: [N, channels, h, w]
+        auto vision_model        = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
+        auto visual_projection   = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);
+        auto visual_projection_2 = std::dynamic_pointer_cast<Linear>(blocks["visual_projection_2"]);
+        auto fuse_module         = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
+
+        struct ggml_tensor* shared_id_embeds = vision_model->forward(ctx, id_pixel_values);          // [N, hidden_size]
+        struct ggml_tensor* id_embeds        = visual_projection->forward(ctx, shared_id_embeds);    // [N, proj_dim(768)]
+        struct ggml_tensor* id_embeds_2      = visual_projection_2->forward(ctx, shared_id_embeds);  // [N, 1280]
+
+        id_embeds   = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 2, 0, 1, 3));
+        id_embeds_2 = ggml_cont(ctx, ggml_permute(ctx, id_embeds_2, 2, 0, 1, 3));
+
+        id_embeds = ggml_concat(ctx, id_embeds, id_embeds_2);  // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
+        id_embeds = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 1, 2, 0, 3));
+
+        struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
+                                                                         prompt_embeds,
+                                                                         id_embeds,
+                                                                         class_tokens_mask,
+                                                                         class_tokens_mask_pos,
+                                                                         left, right);
+        return updated_prompt_embeds;
+    }
+};
+
+struct PhotoMakerIDEncoder : public GGMLModule {
+public:
+    SDVersion version = VERSION_XL;
+    PhotoMakerIDEncoderBlock id_encoder;
+    float style_strength;
+
+    std::vector<float> ctm;
+    std::vector<ggml_fp16_t> ctmf16;
+    std::vector<int> ctmpos;
+
+    std::vector<ggml_fp16_t> zeros_left_16;
+    std::vector<float> zeros_left;
+    std::vector<ggml_fp16_t> zeros_right_16;
+    std::vector<float> zeros_right;
+
+public:
+    PhotoMakerIDEncoder(ggml_backend_t backend, ggml_type wtype, SDVersion version = VERSION_XL, float sty = 20.f)
+        : GGMLModule(backend, wtype),
+          version(version),
+          style_strength(sty) {
+        id_encoder.init(params_ctx, wtype);
+    }
+
+    std::string get_desc() {
+        return "pmid";
+    }
+
+    size_t get_params_mem_size() {
+        size_t params_mem_size = id_encoder.get_params_mem_size();
+        return params_mem_size;
+    }
+
+    size_t get_params_num() {
+        size_t params_num = id_encoder.get_params_num();
+        return params_num;
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        id_encoder.get_param_tensors(tensors, prefix);
+    }
+
+    struct ggml_cgraph* build_graph(  // struct ggml_allocr* allocr,
+        struct ggml_tensor* id_pixel_values,
+        struct ggml_tensor* prompt_embeds,
+        std::vector<bool>& class_tokens_mask) {
+        ctm.clear();
+        ctmf16.clear();
+        ctmpos.clear();
+        zeros_left.clear();
+        zeros_left_16.clear();
+        zeros_right.clear();
+        zeros_right_16.clear();
+
+        ggml_context* ctx0 = compute_ctx;
+
+        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+
+        int64_t hidden_size = prompt_embeds->ne[0];
+        int64_t seq_length  = prompt_embeds->ne[1];
+        ggml_type type      = GGML_TYPE_F32;
+
+        struct ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(ctx0, type, class_tokens_mask.size());
+
+        struct ggml_tensor* id_pixel_values_d = to_backend(id_pixel_values);
+        struct ggml_tensor* prompt_embeds_d   = to_backend(prompt_embeds);
+
+        struct ggml_tensor* left  = NULL;
+        struct ggml_tensor* right = NULL;
+        for (int i = 0; i < class_tokens_mask.size(); i++) {
+            if (class_tokens_mask[i]) {
+                ctm.push_back(0.f);                        // here use 0.f instead of 1.f to make a scale mask
+                ctmf16.push_back(ggml_fp32_to_fp16(0.f));  // here use 0.f instead of 1.f to make a scale mask
+                ctmpos.push_back(i);
+            } else {
+                ctm.push_back(1.f);                        // here use 1.f instead of 0.f to make a scale mask
+                ctmf16.push_back(ggml_fp32_to_fp16(1.f));  // here use 0.f instead of 1.f to make a scale mask
+            }
+        }
+        if (ctmpos[0] > 0) {
+            left = ggml_new_tensor_3d(ctx0, type, hidden_size, 1, ctmpos[0]);
+        }
+        if (ctmpos[ctmpos.size() - 1] < seq_length - 1) {
+            right = ggml_new_tensor_3d(ctx0, type,
+                                       hidden_size, 1, seq_length - ctmpos[ctmpos.size() - 1] - 1);
+        }
+        struct ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ctmpos.size());
+
+        {
+            if (type == GGML_TYPE_F16)
+                set_backend_tensor_data(class_tokens_mask_d, ctmf16.data());
+            else
+                set_backend_tensor_data(class_tokens_mask_d, ctm.data());
+            set_backend_tensor_data(class_tokens_mask_pos, ctmpos.data());
+            if (left) {
+                if (type == GGML_TYPE_F16) {
+                    for (int i = 0; i < ggml_nelements(left); ++i)
+                        zeros_left_16.push_back(ggml_fp32_to_fp16(0.f));
+                    set_backend_tensor_data(left, zeros_left_16.data());
+                } else {
+                    for (int i = 0; i < ggml_nelements(left); ++i)
+                        zeros_left.push_back(0.f);
+                    set_backend_tensor_data(left, zeros_left.data());
+                }
+            }
+            if (right) {
+                if (type == GGML_TYPE_F16) {
+                    for (int i = 0; i < ggml_nelements(right); ++i)
+                        zeros_right_16.push_back(ggml_fp32_to_fp16(0.f));
+                    set_backend_tensor_data(right, zeros_right_16.data());
+                } else {
+                    for (int i = 0; i < ggml_nelements(right); ++i)
+                        zeros_right.push_back(0.f);
+                    set_backend_tensor_data(right, zeros_right.data());
+                }
+            }
+        }
+        struct ggml_tensor* updated_prompt_embeds = id_encoder.forward(ctx0,
+                                                                       id_pixel_values_d,
+                                                                       prompt_embeds_d,
+                                                                       class_tokens_mask_d,
+                                                                       class_tokens_mask_pos,
+                                                                       left, right);
+        ggml_build_forward_expand(gf, updated_prompt_embeds);
+
+        return gf;
+    }
+
+    void compute(const int n_threads,
+                 struct ggml_tensor* id_pixel_values,
+                 struct ggml_tensor* prompt_embeds,
+                 std::vector<bool>& class_tokens_mask,
+                 struct ggml_tensor** updated_prompt_embeds,
+                 ggml_context* output_ctx) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            // return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask);
+            return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask);
+        };
+
+        // GGMLModule::compute(get_graph, n_threads, updated_prompt_embeds);
+        GGMLModule::compute(get_graph, n_threads, true, updated_prompt_embeds, output_ctx);
+    }
+};
+
+#endif  // __PMI_HPP__
--- a/preprocessing.hpp
+++ b/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++) {
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -11,10 +11,19 @@
 #include "denoiser.hpp"
 #include "esrgan.hpp"
 #include "lora.hpp"
+#include "pmid.hpp"
 #include "tae.hpp"
 #include "unet.hpp"
 #include "vae.hpp"

+#define STB_IMAGE_IMPLEMENTATION
+#define STB_IMAGE_STATIC
+#include "stb_image.h"
+
+// #define STB_IMAGE_WRITE_IMPLEMENTATION
+// #define STB_IMAGE_WRITE_STATIC
+// #include "stb_image_write.h"
+
 const char* model_version_to_str[] = {
    "1.x",
    "2.x",
@@ -56,8 +65,11 @@ void calculate_alphas_cumprod(float* alphas_cumprod,

 class StableDiffusionGGML {
 public:
-    ggml_backend_t backend    = NULL;  // general backend
-    ggml_type model_data_type = GGML_TYPE_COUNT;
+    ggml_backend_t backend             = NULL;  // general backend
+    ggml_backend_t clip_backend        = NULL;
+    ggml_backend_t control_net_backend = NULL;
+    ggml_backend_t vae_backend         = NULL;
+    ggml_type model_data_type          = GGML_TYPE_COUNT;

    SDVersion version;
    bool vae_decode_only         = false;
@@ -73,10 +85,13 @@ public:
    std::shared_ptr<AutoEncoderKL> first_stage_model;
    std::shared_ptr<TinyAutoEncoder> tae_first_stage;
    std::shared_ptr<ControlNet> control_net;
+    std::shared_ptr<PhotoMakerIDEncoder> pmid_model;
+    std::shared_ptr<LoraModel> pmid_lora;

    std::string taesd_path;
    bool use_tiny_autoencoder = false;
    bool vae_tiling           = false;
+    bool stacked_id           = false;

    std::map<std::string, struct ggml_tensor*> tensors;

@@ -86,6 +101,8 @@ public:

    std::shared_ptr<Denoiser> denoiser = std::make_shared<CompVisDenoiser>();

+    std::string trigger_word = "img";  // should be user settable
+
    StableDiffusionGGML() = default;

    StableDiffusionGGML(int n_threads,
@@ -105,6 +122,15 @@ public:
    }

    ~StableDiffusionGGML() {
+        if (clip_backend != backend) {
+            ggml_backend_free(clip_backend);
+        }
+        if (control_net_backend != backend) {
+            ggml_backend_free(control_net_backend);
+        }
+        if (vae_backend != backend) {
+            ggml_backend_free(vae_backend);
+        }
        ggml_backend_free(backend);
    }

@@ -112,11 +138,14 @@ public:
                        const std::string& vae_path,
                        const std::string control_net_path,
                        const std::string embeddings_path,
+                        const std::string id_embeddings_path,
                        const std::string& taesd_path,
                        bool vae_tiling_,
                        ggml_type wtype,
                        schedule_t schedule,
-                        bool control_net_cpu) {
+                        bool clip_on_cpu,
+                        bool control_net_cpu,
+                        bool vae_on_cpu) {
        use_tiny_autoencoder = taesd_path.size() > 0;
 #ifdef SD_USE_CUBLAS
        LOG_DEBUG("Using CUDA backend");
@@ -124,7 +153,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

@@ -161,6 +190,7 @@ public:
            LOG_ERROR("get sd version from file failed: '%s'", model_path.c_str());
            return false;
        }
+
        LOG_INFO("Stable Diffusion %s ", model_version_to_str[version]);
        if (wtype == GGML_TYPE_COUNT) {
            model_data_type = model_loader.get_sd_wtype();
@@ -195,7 +225,12 @@ public:
            first_stage_model->alloc_params_buffer();
            first_stage_model->get_param_tensors(tensors, "first_stage_model");
        } else {
-            cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(backend, model_data_type, version);
+            clip_backend = backend;
+            if (clip_on_cpu && !ggml_backend_is_cpu(backend)) {
+                LOG_INFO("CLIP: Using CPU backend");
+                clip_backend = ggml_backend_cpu_init();
+            }
+            cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(clip_backend, model_data_type, version);
            cond_stage_model->alloc_params_buffer();
            cond_stage_model->get_param_tensors(tensors, "cond_stage_model.");

@@ -211,24 +246,59 @@ public:
            }

            if (!use_tiny_autoencoder) {
-                first_stage_model = std::make_shared<AutoEncoderKL>(backend, vae_type, vae_decode_only);
+                if (vae_on_cpu && !ggml_backend_is_cpu(backend)) {
+                    LOG_INFO("VAE Autoencoder: Using CPU backend");
+                    vae_backend = ggml_backend_cpu_init();
+                } else {
+                    vae_backend = backend;
+                }
+                first_stage_model = std::make_shared<AutoEncoderKL>(vae_backend, vae_type, vae_decode_only);
                first_stage_model->alloc_params_buffer();
                first_stage_model->get_param_tensors(tensors, "first_stage_model");
            } else {
                tae_first_stage = std::make_shared<TinyAutoEncoder>(backend, model_data_type, vae_decode_only);
            }
+            // first_stage_model->get_param_tensors(tensors, "first_stage_model.");

            if (control_net_path.size() > 0) {
-                ggml_backend_t cn_backend = NULL;
+                ggml_backend_t controlnet_backend = NULL;
                if (control_net_cpu && !ggml_backend_is_cpu(backend)) {
                    LOG_DEBUG("ControlNet: Using CPU backend");
-                    cn_backend = ggml_backend_cpu_init();
+                    controlnet_backend = ggml_backend_cpu_init();
                } else {
-                    cn_backend = backend;
+                    controlnet_backend = backend;
                }
-                control_net = std::make_shared<ControlNet>(cn_backend, model_data_type, version);
+                control_net = std::make_shared<ControlNet>(controlnet_backend, model_data_type, version);
            }

+            pmid_model = std::make_shared<PhotoMakerIDEncoder>(clip_backend, model_data_type, version);
+            if (id_embeddings_path.size() > 0) {
+                pmid_lora = std::make_shared<LoraModel>(backend, model_data_type, id_embeddings_path, "");
+                if (!pmid_lora->load_from_file(true)) {
+                    LOG_WARN("load photomaker lora tensors from %s failed", id_embeddings_path.c_str());
+                    return false;
+                }
+                LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", id_embeddings_path.c_str());
+                if (!model_loader.init_from_file(id_embeddings_path, "pmid.")) {
+                    LOG_WARN("loading stacked ID embedding from '%s' failed", id_embeddings_path.c_str());
+                } else {
+                    stacked_id = true;
+                }
+            }
+            if (stacked_id) {
+                if (!pmid_model->alloc_params_buffer()) {
+                    LOG_ERROR(" pmid model params buffer allocation failed");
+                    return false;
+                }
+                // LOG_INFO("pmid param memory buffer size = %.2fMB ",
+                //     pmid_model->params_buffer_size / 1024.0 / 1024.0);
+                pmid_model->get_param_tensors(tensors, "pmid");
+            }
+            // if(stacked_id){
+            //    pmid_model.init_params(GGML_TYPE_F32);
+            //    pmid_model.map_by_name(tensors, "pmid.");
+            // }
+
            LOG_DEBUG("loading vocab");
            std::string merges_utf8_str = model_loader.load_merges();
            if (merges_utf8_str.size() == 0) {
@@ -250,6 +320,7 @@ public:

        // load weights
        LOG_DEBUG("loading weights");
+
        int64_t t0 = ggml_time_ms();

        std::set<std::string> ignore_tensors;
@@ -257,6 +328,10 @@ public:
        if (use_tiny_autoencoder) {
            ignore_tensors.insert("first_stage_model.");
        }
+        if (stacked_id) {
+            ignore_tensors.insert("lora.");
+        }
+
        if (vae_decode_only) {
            ignore_tensors.insert("first_stage_model.encoder");
            ignore_tensors.insert("first_stage_model.quant");
@@ -296,14 +371,54 @@ public:
                }
                control_net_params_mem_size = control_net->get_params_mem_size();
            }
+            size_t pmid_params_mem_size = 0;
+            if (stacked_id) {
+                pmid_params_mem_size = pmid_model->get_params_mem_size();
+            }

-            size_t total_params_size = clip_params_mem_size + clip_params_mem_size + clip_params_mem_size + control_net_params_mem_size;
-            LOG_INFO("total params memory size = %.2fMB (clip %.2fMB, unet %.2fMB, vae %.2fMB, controlnet %.2fMB)",
-                     total_params_size / 1024.0 / 1024.0,
-                     clip_params_mem_size / 1024.0 / 1024.0,
-                     unet_params_mem_size / 1024.0 / 1024.0,
-                     vae_params_mem_size / 1024.0 / 1024.0,
-                     control_net_params_mem_size / 1024.0 / 1024.0);
+            size_t total_params_ram_size  = 0;
+            size_t total_params_vram_size = 0;
+            if (ggml_backend_is_cpu(clip_backend)) {
+                total_params_ram_size += clip_params_mem_size + pmid_params_mem_size;
+            } else {
+                total_params_vram_size += clip_params_mem_size + pmid_params_mem_size;
+            }
+
+            if (ggml_backend_is_cpu(backend)) {
+                total_params_ram_size += unet_params_mem_size;
+            } else {
+                total_params_vram_size += unet_params_mem_size;
+            }
+
+            if (ggml_backend_is_cpu(vae_backend)) {
+                total_params_ram_size += vae_params_mem_size;
+            } else {
+                total_params_vram_size += vae_params_mem_size;
+            }
+
+            if (ggml_backend_is_cpu(control_net_backend)) {
+                total_params_ram_size += control_net_params_mem_size;
+            } else {
+                total_params_vram_size += control_net_params_mem_size;
+            }
+
+            size_t total_params_size = total_params_ram_size + total_params_vram_size;
+            LOG_INFO(
+                "total params memory size = %.2fMB (VRAM %.2fMB, RAM %.2fMB): "
+                "clip %.2fMB(%s), unet %.2fMB(%s), vae %.2fMB(%s), controlnet %.2fMB(%s), pmid %.2fMB(%s)",
+                total_params_size / 1024.0 / 1024.0,
+                total_params_vram_size / 1024.0 / 1024.0,
+                total_params_ram_size / 1024.0 / 1024.0,
+                clip_params_mem_size / 1024.0 / 1024.0,
+                ggml_backend_is_cpu(clip_backend) ? "RAM" : "VRAM",
+                unet_params_mem_size / 1024.0 / 1024.0,
+                ggml_backend_is_cpu(backend) ? "RAM" : "VRAM",
+                vae_params_mem_size / 1024.0 / 1024.0,
+                ggml_backend_is_cpu(vae_backend) ? "RAM" : "VRAM",
+                control_net_params_mem_size / 1024.0 / 1024.0,
+                ggml_backend_is_cpu(control_net_backend) ? "RAM" : "VRAM",
+                pmid_params_mem_size / 1024.0 / 1024.0,
+                ggml_backend_is_cpu(clip_backend) ? "RAM" : "VRAM");
        }

        int64_t t1 = ggml_time_ms();
@@ -363,9 +478,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();

@@ -411,9 +527,7 @@ public:

        int64_t t1 = ggml_time_ms();

-        LOG_INFO("lora '%s' applied, taking %.2fs",
-                 lora_name.c_str(),
-                 (t1 - t0) * 1.0f / 1000);
+        LOG_INFO("lora '%s' applied, taking %.2fs", lora_name.c_str(), (t1 - t0) * 1.0f / 1000);
    }

    void apply_loras(const std::unordered_map<std::string, float>& lora_state) {
@@ -436,6 +550,8 @@ public:
            }
        }

+        LOG_INFO("Attempting to apply %lu LoRAs", lora_state.size());
+
        for (auto& kv : lora_state_diff) {
            apply_lora(kv.first, kv.second);
        }
@@ -443,52 +559,154 @@ public:
        curr_lora_state = lora_state;
    }

+    std::string remove_trigger_from_prompt(ggml_context* work_ctx,
+                                           const std::string& prompt) {
+        auto image_tokens = cond_stage_model->convert_token_to_id(trigger_word);
+        GGML_ASSERT(image_tokens.size() == 1);
+        auto tokens_and_weights  = cond_stage_model->tokenize(prompt, false);
+        std::vector<int>& tokens = tokens_and_weights.first;
+        auto it                  = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
+        GGML_ASSERT(it != tokens.end());  // prompt must have trigger word
+        tokens.erase(it);
+        return cond_stage_model->decode(tokens);
+    }
+
+    std::tuple<ggml_tensor*, ggml_tensor*, std::vector<bool>>
+    get_learned_condition_with_trigger(ggml_context* work_ctx,
+                                       const std::string& text,
+                                       int clip_skip,
+                                       int width,
+                                       int height,
+                                       int num_input_imgs,
+                                       bool force_zero_embeddings = false) {
+        auto image_tokens = cond_stage_model->convert_token_to_id(trigger_word);
+        // if(image_tokens.size() == 1){
+        //     printf(" image token id is: %d \n", image_tokens[0]);
+        // }
+        GGML_ASSERT(image_tokens.size() == 1);
+        auto tokens_and_weights     = cond_stage_model->tokenize_with_trigger_token(text,
+                                                                                    num_input_imgs,
+                                                                                    image_tokens[0],
+                                                                                    true);
+        std::vector<int>& tokens    = std::get<0>(tokens_and_weights);
+        std::vector<float>& weights = std::get<1>(tokens_and_weights);
+        std::vector<bool>& clsm     = std::get<2>(tokens_and_weights);
+        // printf("tokens: \n");
+        // for(int i = 0; i < tokens.size(); ++i)
+        //    printf("%d ", tokens[i]);
+        // printf("\n");
+        // printf("clsm: \n");
+        // for(int i = 0; i < clsm.size(); ++i)
+        //    printf("%d ", clsm[i]?1:0);
+        // printf("\n");
+        auto cond = get_learned_condition_common(work_ctx, tokens, weights, clip_skip, width, height, force_zero_embeddings);
+        return std::make_tuple(cond.first, cond.second, clsm);
+    }
+
+    ggml_tensor* id_encoder(ggml_context* work_ctx,
+                            ggml_tensor* init_img,
+                            ggml_tensor* prompts_embeds,
+                            std::vector<bool>& class_tokens_mask) {
+        ggml_tensor* res = NULL;
+        pmid_model->compute(n_threads, init_img, prompts_embeds, class_tokens_mask, &res, work_ctx);
+
+        return res;
+    }
+
    std::pair<ggml_tensor*, ggml_tensor*> get_learned_condition(ggml_context* work_ctx,
                                                                const std::string& text,
                                                                int clip_skip,
                                                                int width,
                                                                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;
+        return get_learned_condition_common(work_ctx, tokens, weights, clip_skip, width, height, force_zero_embeddings);
+    }

-        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);
+    std::pair<ggml_tensor*, ggml_tensor*> get_learned_condition_common(ggml_context* work_ctx,
+                                                                       std::vector<int>& tokens,
+                                                                       std::vector<float>& weights,
+                                                                       int clip_skip,
+                                                                       int width,
+                                                                       int height,
+                                                                       bool force_zero_embeddings = false) {
+        cond_stage_model->set_clip_skip(clip_skip);
+        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;
+
+        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 +744,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,
@@ -625,7 +843,10 @@ public:
                        float min_cfg,
                        float cfg_scale,
                        sample_method_t method,
-                        const std::vector<float>& sigmas) {
+                        const std::vector<float>& sigmas,
+                        int start_merge_step,
+                        ggml_tensor* c_id,
+                        ggml_tensor* c_vec_id) {
        size_t steps = sigmas.size() - 1;
        // x_t = load_tensor_from_file(work_ctx, "./rand0.bin");
        // print_ggml_tensor(x_t);
@@ -675,7 +896,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
@@ -690,17 +912,30 @@ public:
                // GGML_ASSERT(0);
            }

-            // cond
-            diffusion_model->compute(n_threads,
-                                     noised_input,
-                                     timesteps,
-                                     c,
-                                     c_concat,
-                                     c_vector,
-                                     -1,
-                                     controls,
-                                     control_strength,
-                                     &out_cond);
+            if (start_merge_step == -1 || step <= start_merge_step) {
+                // cond
+                diffusion_model->compute(n_threads,
+                                         noised_input,
+                                         timesteps,
+                                         c,
+                                         c_concat,
+                                         c_vector,
+                                         -1,
+                                         controls,
+                                         control_strength,
+                                         &out_cond);
+            } else {
+                diffusion_model->compute(n_threads,
+                                         noised_input,
+                                         timesteps,
+                                         c_id,
+                                         c_concat,
+                                         c_vec_id,
+                                         -1,
+                                         controls,
+                                         control_strength,
+                                         &out_cond);
+            }

            float* negative_data = NULL;
            if (has_unconditioned) {
@@ -1243,6 +1478,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                     const char* control_net_path_c_str,
                     const char* lora_model_dir_c_str,
                     const char* embed_dir_c_str,
+                     const char* id_embed_dir_c_str,
                     bool vae_decode_only,
                     bool vae_tiling,
                     bool free_params_immediately,
@@ -1250,7 +1486,9 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                     enum sd_type_t wtype,
                     enum rng_type_t rng_type,
                     enum schedule_t s,
-                     bool keep_control_net_cpu) {
+                     bool keep_clip_on_cpu,
+                     bool keep_control_net_cpu,
+                     bool keep_vae_on_cpu) {
    sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
    if (sd_ctx == NULL) {
        return NULL;
@@ -1260,6 +1498,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
    std::string taesd_path(taesd_path_c_str);
    std::string control_net_path(control_net_path_c_str);
    std::string embd_path(embed_dir_c_str);
+    std::string id_embd_path(id_embed_dir_c_str);
    std::string lora_model_dir(lora_model_dir_c_str);

    sd_ctx->sd = new StableDiffusionGGML(n_threads,
@@ -1275,11 +1514,14 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                                    vae_path,
                                    control_net_path,
                                    embd_path,
+                                    id_embd_path,
                                    taesd_path,
                                    vae_tiling,
                                    (ggml_type)wtype,
                                    s,
-                                    keep_control_net_cpu)) {
+                                    keep_clip_on_cpu,
+                                    keep_control_net_cpu,
+                                    keep_vae_on_cpu)) {
        delete sd_ctx->sd;
        sd_ctx->sd = NULL;
        free(sd_ctx);
@@ -1308,7 +1550,10 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                    int64_t seed,
                    int batch_count,
                    const sd_image_t* control_cond,
-                    float control_strength) {
+                    float control_strength,
+                    float style_ratio,
+                    bool normalize_input,
+                    const char* input_id_images_path_c_str) {
    LOG_DEBUG("txt2img %dx%d", width, height);
    if (sd_ctx == NULL) {
        return NULL;
@@ -1316,6 +1561,35 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
    // LOG_DEBUG("%s %s %f %d %d %d", prompt_c_str, negative_prompt_c_str, cfg_scale, sample_steps, seed, batch_count);
    std::string prompt(prompt_c_str);
    std::string negative_prompt(negative_prompt_c_str);
+    std::string input_id_images_path(input_id_images_path_c_str);
+
+    // preprocess input id images
+    std::vector<sd_image_t*> input_id_images;
+    if (sd_ctx->sd->pmid_model && input_id_images_path.size() > 0) {
+        std::vector<std::string> img_files = get_files_from_dir(input_id_images_path);
+        for (std::string img_file : img_files) {
+            int c = 0;
+            int width, height;
+            uint8_t* input_image_buffer = stbi_load(img_file.c_str(), &width, &height, &c, 3);
+            if (input_image_buffer == NULL) {
+                LOG_ERROR("PhotoMaker load image from '%s' failed", img_file.c_str());
+                continue;
+            } else {
+                LOG_INFO("PhotoMaker loaded image from '%s'", img_file.c_str());
+            }
+            sd_image_t* input_image = NULL;
+            input_image             = new sd_image_t{(uint32_t)width,
+                                         (uint32_t)height,
+                                         3,
+                                         input_image_buffer};
+            input_image             = preprocess_id_image(input_image);
+            if (input_image == NULL) {
+                LOG_ERROR("preprocess input id image from '%s' failed", img_file.c_str());
+                continue;
+            }
+            input_id_images.push_back(input_image);
+        }
+    }

    // extract and remove lora
    auto result_pair                                = extract_and_remove_lora(prompt);
@@ -1332,8 +1606,22 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
    sd_ctx->sd->apply_loras(lora_f2m);
    int64_t t1 = ggml_time_ms();
    LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
+
+    if (sd_ctx->sd->stacked_id) {
+        t0 = ggml_time_ms();
+        sd_ctx->sd->pmid_lora->apply(sd_ctx->sd->tensors, sd_ctx->sd->n_threads);
+        t1 = ggml_time_ms();
+        LOG_INFO("pmid_lora apply completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
+        if (sd_ctx->sd->free_params_immediately) {
+            sd_ctx->sd->pmid_lora->free_params_buffer();
+        }
+    }
+
    struct ggml_init_params params;
    params.mem_size = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
+    if (sd_ctx->sd->stacked_id) {
+        params.mem_size += static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
+    }
    params.mem_size += width * height * 3 * sizeof(float);
    params.mem_size *= batch_count;
    params.mem_buffer = NULL;
@@ -1354,10 +1642,67 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
        seed = rand();
    }

-    t0                            = ggml_time_ms();
-    auto cond_pair                = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height);
-    ggml_tensor* c                = cond_pair.first;
-    ggml_tensor* c_vector         = cond_pair.second;  // [adm_in_channels, ]
+    std::string prompt_text_only;
+    ggml_tensor* init_img              = NULL;
+    ggml_tensor* prompts_embeds        = NULL;
+    ggml_tensor* pooled_prompts_embeds = NULL;
+    // ggml_tensor* class_tokens_mask = NULL;
+    std::vector<bool> class_tokens_mask;
+    if (sd_ctx->sd->stacked_id) {
+        if (input_id_images.size() > 0) {
+            sd_ctx->sd->pmid_model->style_strength = style_ratio;
+            int32_t w                              = input_id_images[0]->width;
+            int32_t h                              = input_id_images[0]->height;
+            int32_t channels                       = input_id_images[0]->channel;
+            int32_t num_input_images               = (int32_t)input_id_images.size();
+            init_img                               = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, w, h, channels, num_input_images);
+            // TODO: move these to somewhere else and be user settable
+            float mean[] = {0.48145466f, 0.4578275f, 0.40821073f};
+            float std[]  = {0.26862954f, 0.26130258f, 0.27577711f};
+            for (int i = 0; i < num_input_images; i++) {
+                sd_image_t* init_image = input_id_images[i];
+                if (normalize_input)
+                    sd_mul_images_to_tensor(init_image->data, init_img, i, mean, std);
+                else
+                    sd_mul_images_to_tensor(init_image->data, init_img, i, NULL, NULL);
+            }
+            t0                    = ggml_time_ms();
+            auto cond_tup         = sd_ctx->sd->get_learned_condition_with_trigger(work_ctx, prompt,
+                                                                                   clip_skip, width, height, num_input_images);
+            prompts_embeds        = std::get<0>(cond_tup);
+            pooled_prompts_embeds = std::get<1>(cond_tup);  // [adm_in_channels, ]
+            class_tokens_mask     = std::get<2>(cond_tup);  //
+
+            prompts_embeds = sd_ctx->sd->id_encoder(work_ctx, init_img, prompts_embeds, class_tokens_mask);
+            t1             = ggml_time_ms();
+            LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0);
+            if (sd_ctx->sd->free_params_immediately) {
+                sd_ctx->sd->pmid_model->free_params_buffer();
+            }
+            // Encode input prompt without the trigger word for delayed conditioning
+            prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt);
+            // printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str());
+            prompt = prompt_text_only;  //
+            if (sample_steps < 50) {
+                LOG_INFO("sampling steps increases from %d to 50 for PHOTOMAKER", sample_steps);
+                sample_steps = 50;
+            }
+        } else {
+            LOG_WARN("Provided PhotoMaker model file, but NO input ID images");
+            LOG_WARN("Turn off PhotoMaker");
+            sd_ctx->sd->stacked_id = false;
+        }
+    }
+    for (sd_image_t* img : input_id_images) {
+        free(img->data);
+    }
+    input_id_images.clear();
+
+    t0                    = ggml_time_ms();
+    auto cond_pair        = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height);
+    ggml_tensor* c        = cond_pair.first;
+    ggml_tensor* c_vector = cond_pair.second;  // [adm_in_channels, ]
+
    struct ggml_tensor* uc        = NULL;
    struct ggml_tensor* uc_vector = NULL;
    if (cfg_scale != 1.0) {
@@ -1398,6 +1743,14 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,

        std::vector<float> sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps);

+        int start_merge_step = -1;
+        if (sd_ctx->sd->stacked_id) {
+            start_merge_step = int(sd_ctx->sd->pmid_model->style_strength / 100.f * sample_steps);
+            if (start_merge_step > 30)
+                start_merge_step = 30;
+            LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step);
+        }
+
        struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx,
                                                     x_t,
                                                     NULL,
@@ -1412,7 +1765,10 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                                                     cfg_scale,
                                                     cfg_scale,
                                                     sample_method,
-                                                     sigmas);
+                                                     sigmas,
+                                                     start_merge_step,
+                                                     prompts_embeds,
+                                                     pooled_prompts_embeds);
        // struct ggml_tensor* x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
        // print_ggml_tensor(x_0);
        int64_t sampling_end = ggml_time_ms();
@@ -1579,7 +1935,10 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
                                                 cfg_scale,
                                                 cfg_scale,
                                                 sample_method,
-                                                 sigma_sched);
+                                                 sigma_sched,
+                                                 -1,
+                                                 NULL,
+                                                 NULL);
    // struct ggml_tensor *x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
    // print_ggml_tensor(x_0);
    int64_t t3 = ggml_time_ms();
@@ -1715,7 +2074,10 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
                                                 min_cfg,
                                                 cfg_scale,
                                                 sample_method,
-                                                 sigmas);
+                                                 sigmas,
+                                                 -1,
+                                                 NULL,
+                                                 NULL);

    int64_t t2 = ggml_time_ms();
    LOG_INFO("sampling completed, taking %.2fs", (t2 - t1) * 1.0f / 1000);
@@ -1753,4 +2115,4 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
    LOG_INFO("img2vid completed in %.2fs", (t3 - t0) * 1.0f / 1000);

    return result_images;
-}
+}
--- a/stable-diffusion.h
+++ b/stable-diffusion.h
@@ -73,6 +73,12 @@ 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_IQ3_S   = 21,
+    SD_TYPE_IQ2_S   = 22,
+    SD_TYPE_IQ4_XS  = 23,
    SD_TYPE_I8,
    SD_TYPE_I16,
    SD_TYPE_I32,
@@ -89,8 +95,10 @@ enum sd_log_level_t {
 };

 typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
+typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);

 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
+SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
 SD_API int32_t get_num_physical_cores();
 SD_API const char* sd_get_system_info();

@@ -109,6 +117,7 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
                            const char* control_net_path_c_str,
                            const char* lora_model_dir,
                            const char* embed_dir_c_str,
+                            const char* stacked_id_embed_dir_c_str,
                            bool vae_decode_only,
                            bool vae_tiling,
                            bool free_params_immediately,
@@ -116,7 +125,9 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
                            enum sd_type_t wtype,
                            enum rng_type_t rng_type,
                            enum schedule_t s,
-                            bool keep_control_net_cpu);
+                            bool keep_clip_on_cpu,
+                            bool keep_control_net_cpu,
+                            bool keep_vae_on_cpu);

 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);

@@ -132,7 +143,10 @@ SD_API sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                           int64_t seed,
                           int batch_count,
                           const sd_image_t* control_cond,
-                           float control_strength);
+                           float control_strength,
+                           float style_strength,
+                           bool normalize_input,
+                           const char* input_id_images_path);

 SD_API sd_image_t* img2img(sd_ctx_t* sd_ctx,
                           sd_image_t init_image,
@@ -174,8 +188,17 @@ 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

-#endif  // __STABLE_DIFFUSION_H__
+#endif  // __STABLE_DIFFUSION_H__
--- a/thirdparty/stb_image_resize.h
+++ b/thirdparty/stb_image_resize.h
--- a/unet.hpp
+++ b/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);

--- a/upscaler.cpp
+++ b/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

--- a/util.cpp
+++ b/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>
@@ -24,6 +25,9 @@
 #include "ggml/ggml.h"
 #include "stable-diffusion.h"

+#define STB_IMAGE_RESIZE_IMPLEMENTATION
+#include "stb_image_resize.h"
+
 bool ends_with(const std::string& str, const std::string& ending) {
    if (str.length() >= ending.length()) {
        return (str.compare(str.length() - ending.length(), ending.length(), ending) == 0);
@@ -39,6 +43,13 @@ bool starts_with(const std::string& str, const std::string& start) {
    return false;
 }

+bool contains(const std::string& str, const std::string& substr) {
+    if (str.find(substr) != std::string::npos) {
+        return true;
+    }
+    return false;
+}
+
 void replace_all_chars(std::string& str, char target, char replacement) {
    for (size_t i = 0; i < str.length(); ++i) {
        if (str[i] == target) {
@@ -87,6 +98,43 @@ std::string get_full_path(const std::string& dir, const std::string& filename) {
    }
 }

+std::vector<std::string> get_files_from_dir(const std::string& dir) {
+    std::vector<std::string> files;
+
+    WIN32_FIND_DATA findFileData;
+    HANDLE hFind;
+
+    char currentDirectory[MAX_PATH];
+    GetCurrentDirectory(MAX_PATH, currentDirectory);
+
+    char directoryPath[MAX_PATH];  // this is absolute path
+    sprintf(directoryPath, "%s\\%s\\*", currentDirectory, dir.c_str());
+
+    // Find the first file in the directory
+    hFind = FindFirstFile(directoryPath, &findFileData);
+
+    // Check if the directory was found
+    if (hFind == INVALID_HANDLE_VALUE) {
+        printf("Unable to find directory.\n");
+        return files;
+    }
+
+    // Loop through all files in the directory
+    do {
+        // Check if the found file is a regular file (not a directory)
+        if (!(findFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)) {
+            files.push_back(std::string(currentDirectory) + "\\" + dir + "\\" + std::string(findFileData.cFileName));
+        }
+    } while (FindNextFile(hFind, &findFileData) != 0);
+
+    // Close the handle
+    FindClose(hFind);
+
+    sort(files.begin(), files.end());
+
+    return files;
+}
+
 #else  // Unix
 #include <dirent.h>
 #include <sys/stat.h>
@@ -101,6 +149,7 @@ bool is_directory(const std::string& path) {
    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }

+// TODO: add windows version
 std::string get_full_path(const std::string& dir, const std::string& filename) {
    DIR* dp = opendir(dir.c_str());

@@ -120,6 +169,27 @@ std::string get_full_path(const std::string& dir, const std::string& filename) {
    return "";
 }

+std::vector<std::string> get_files_from_dir(const std::string& dir) {
+    std::vector<std::string> files;
+
+    DIR* dp = opendir(dir.c_str());
+
+    if (dp != nullptr) {
+        struct dirent* entry;
+
+        while ((entry = readdir(dp)) != nullptr) {
+            std::string fname = dir + "/" + entry->d_name;
+            if (!is_directory(fname))
+                files.push_back(fname);
+        }
+        closedir(dp);
+    }
+
+    sort(files.begin(), files.end());
+
+    return files;
+}
+
 #endif

 // get_num_physical_cores is copy from
@@ -160,6 +230,9 @@ int32_t get_num_physical_cores() {
    return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
 }

+static sd_progress_cb_t sd_progress_cb = NULL;
+void* sd_progress_cb_data              = NULL;
+
 std::u32string utf8_to_utf32(const std::string& utf8_str) {
    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
    return converter.from_bytes(utf8_str);
@@ -175,7 +248,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);
@@ -203,7 +276,44 @@ std::string path_join(const std::string& p1, const std::string& p2) {
    return p1 + "/" + p2;
 }

+sd_image_t* preprocess_id_image(sd_image_t* img) {
+    int shortest_edge   = 224;
+    int size            = shortest_edge;
+    sd_image_t* resized = NULL;
+    uint32_t w          = img->width;
+    uint32_t h          = img->height;
+    uint32_t c          = img->channel;
+
+    // 1. do resize using stb_resize functions
+
+    unsigned char* buf = (unsigned char*)malloc(sizeof(unsigned char) * 3 * size * size);
+    if (!stbir_resize_uint8(img->data, w, h, 0,
+                            buf, size, size, 0,
+                            c)) {
+        fprintf(stderr, "%s: resize operation failed \n ", __func__);
+        return resized;
+    }
+
+    // 2. do center crop (likely unnecessary due to step 1)
+
+    // 3. do rescale
+
+    // 4. do normalize
+
+    // 3 and 4 will need to be done in float format.
+
+    resized = new sd_image_t{(uint32_t)shortest_edge,
+                             (uint32_t)shortest_edge,
+                             3,
+                             buf};
+    return resized;
+}
+
 void pretty_progress(int step, int steps, float time) {
+    if (sd_progress_cb) {
+        sd_progress_cb(step, steps, time, sd_progress_cb_data);
+        return;
+    }
    if (step == 0) {
        return;
    }
@@ -256,23 +366,13 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
    va_list args;
    va_start(args, format);

-    const char* level_str = "DEBUG";
-    if (level == SD_LOG_INFO) {
-        level_str = "INFO ";
-    } else if (level == SD_LOG_WARN) {
-        level_str = "WARN ";
-    } else if (level == SD_LOG_ERROR) {
-        level_str = "ERROR";
-    }
-
-    static char log_buffer[LOG_BUFFER_SIZE];
-
-    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "[%s] %s:%-4d - ", level_str, sd_basename(file).c_str(), line);
+    static char log_buffer[LOG_BUFFER_SIZE + 1];
+    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", 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);
@@ -285,7 +385,10 @@ void sd_set_log_callback(sd_log_cb_t cb, void* data) {
    sd_log_cb      = cb;
    sd_log_cb_data = data;
 }
-
+void sd_set_progress_callback(sd_progress_cb_t cb, void* data) {
+    sd_progress_cb      = cb;
+    sd_progress_cb_data = data;
+}
 const char* sd_get_system_info() {
    static char buffer[1024];
    std::stringstream ss;
@@ -459,4 +562,4 @@ sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size) {
    }

    return result;
-}
+}
--- a/util.h
+++ b/util.h
@@ -3,11 +3,13 @@

 #include <cstdint>
 #include <string>
+#include <vector>

 #include "stable-diffusion.h"

 bool ends_with(const std::string& str, const std::string& ending);
 bool starts_with(const std::string& str, const std::string& start);
+bool contains(const std::string& str, const std::string& substr);

 std::string format(const char* fmt, ...);

@@ -17,11 +19,15 @@ bool file_exists(const std::string& filename);
 bool is_directory(const std::string& path);
 std::string get_full_path(const std::string& dir, const std::string& filename);

+std::vector<std::string> get_files_from_dir(const std::string& dir);
+
 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);
+sd_image_t* preprocess_id_image(sd_image_t* img);
+
+// std::string sd_basename(const std::string& path);

 typedef struct {
    uint32_t width;
--- a/vae.hpp
+++ b/vae.hpp
@@ -6,7 +6,7 @@

 /*================================================== AutoEncoderKL ===================================================*/

-#define VAE_GRAPH_SIZE 10240
+#define VAE_GRAPH_SIZE 20480

 class ResnetBlock : public UnaryBlock {
 protected:
Author	SHA1	Message	Date
Grauho	48bcce493f	fix: avoid double free and fix sdxl lora naming conversion * Fixed a double free issue when running multiple backends on the CPU, eg: CLIP and the primary backend, as this would result in the _backend pointers both pointing to the same thing resulting in a segfault when calling the StableDiffusionGGML destructor. Improve logging to allow for a color switch on the command line interface. Changed the base log_printf function to not bake the log level directly into the log buffer as that information is already passed the logging function via the level parameter and it's easier to add in there than strip it out. * Added a fix for certain SDXL LoRAs that don't seem to follow the expected naming convention, converts over the tensor name during the LoRA model loading. Added some logging of useful LoRA loading information. Had to increase the base size of the GGML graph as the existing size results in an insufficient graph memory error when using SDXL LoRAs. * small fixes --------- Co-authored-by: leejet <leejet714@gmail.com>	2024-03-20 22:00:22 +08:00
bssrdf	a469688e30	feat: add TencentARC PhotoMaker support (#179 ) * first efforts at implementing photomaker; lots more to do * added PhotoMakerIDEncoder model in SD * fixed soem bugs; now photomaker model weights can be loaded into their tensor buffers * added input id image loading * added preprocessing inpit id images * finished get_num_tensors * fixed a bug in remove_duplicates * add a get_learned_condition_with_trigger function to do photomaker stuff * add a convert_token_to_id function for photomaker to extract trigger word's token id * making progress; need to implement tokenizer decoder * making more progress; finishing vision model forward * debugging vision_model outputs * corrected clip vision model output * continue making progress in id fusion process * finished stacked id embedding; to be tested * remove garbage file * debuging graph compute * more progress; now alloc buffer failed * fixed wtype issue; input images can only be 1 because issue with transformer when batch size > 1 (to be investigated) * added delayed subject conditioning; now photomaker runs and generates images * fixed stat_merge_step * added photomaker lora model (to be tested) * reworked pmid lora * finished applying pmid lora; to be tested * finalized pmid lora * add a few print tensor; tweak in sample again * small tweak; still not getting ID faces * fixed a bug in FuseBlock forward; also remove diag_mask op in for vision transformer; getting better results * disable pmid lora apply for now; 1 input image seems working; > 1 not working * turn pmid lora apply back on * fixed a decode bug * fixed a bug in ggml's conv_2d, and now > 1 input images working * add style_ratio as a cli param; reworked encode with trigger for attention weights * merge commit fixing lora free param buffer error * change default style ratio to 10% * added an option to offload vae decoder to CPU for mem-limited gpus * removing image normalization step seems making ID fidelity much higher * revert default style ratio back ro 20% * added an option for normalizing input ID images; cleaned up debugging code * more clean up * fixed bugs; now failed with cuda error; likely out-of-mem on GPU * free pmid model params when required * photomaker working properly now after merging and adapting to GGMLBlock API * remove tensor renaming; fixing names in the photomaker model file * updated README.md to include instructions and notes for running PhotoMaker * a bit clean up * remove -DGGML_CUDA_FORCE_MMQ; more clean up and README update * add input image requirement in README * bring back freeing pmid lora params buffer; simply pooled output of CLIPvision * remove MultiheadAttention2; customized MultiheadAttention * added a WIN32 get_files_from_dir; turn off Photomakder if receiving no input images * update docs * fix ci error * make stable-diffusion.h a pure c header file This reverts commit `27887b630d`. * fix ci error * format code * reuse get_learned_condition * reuse pad_tokens * reuse CLIPVisionModel * reuse LoraModel * add --clip-on-cpu * fix lora name conversion for SDXL --------- Co-authored-by: bssrdf <bssrdf@gmail.com> Co-authored-by: leejet <leejet714@gmail.com>	2024-03-12 23:15:17 +08:00
leejet	61980171a1	sync: update ggml	2024-03-10 17:23:11 +08:00
Cyberhan123	583cc5bba2	docs: add binding (#189 )	2024-03-03 13:27:07 +08:00
Phu Tran	1ce9470f27	fix: fix building shared library (#188 )	2024-03-03 13:24:59 +08:00
leejet	a65c410463	sync: update ggml	2024-03-02 19:49:41 +08:00
leejet	a17ae7b7d2	sync: update ggml	2024-03-02 19:23:11 +08:00
leejet	e1b37b4ef6	fix: update ggml submodule url	2024-03-02 17:34:08 +08:00
fszontagh	7be65faa7c	feat: add progress callback (#170 )	2024-03-02 17:28:41 +08:00
Phu Tran	d164236b2a	fix: fix metal build issues (#183 )	2024-03-02 17:17:57 +08:00
leejet	ef5c3f7401	feat: add support for prompt longer than 77	2024-03-02 17:13:18 +08:00
Cyberhan123	b7870a0f89	chore: improve ci (#150 ) --------- Co-authored-by: leejet <leejet714@gmail.com>	2024-02-26 22:01:34 +08:00
leejet	4a8190405a	fix: fix the issue with dynamic linking	2024-02-25 21:39:01 +08:00
leejet	730585d515	sync: update ggml (#180 )	2024-02-25 21:11:01 +08:00