feat: added prediction argument (#334)

* add ref latent support for qwen image

* optimize clip_preprocess and fix get_first_stage_encoding

* add qwen2vl vit support

* add qwen image edit support

* fix qwen image edit pipeline

* add mmproj file support

* support dynamic number of Qwen image transformer blocks

* set prompt_template_encode_start_idx every time

* to_add_out precision fix

* to_out.0 precision fix

* update docs

.github/workflows/build.yml

@@ -149,7 +149,7 @@ jobs:
     runs-on: windows-2025
 
     env:
-      VULKAN_VERSION: 1.3.261.1
+      VULKAN_VERSION: 1.4.328.1
 
     strategy:
       matrix:
@@ -199,9 +199,9 @@ jobs:
           version: 1.11.1
       - name: Install Vulkan SDK
         id: get_vulkan
-        if: ${{ matrix.build == 'vulkan' }}
+        if: ${{ matrix.build == 'vulkan' }} https://sdk.lunarg.com/sdk/download/1.4.328.1/windows/vulkansdk-windows-X64-1.4.328.1.exe
         run: |
-          curl.exe -o $env:RUNNER_TEMP/VulkanSDK-Installer.exe -L "https://sdk.lunarg.com/sdk/download/${env:VULKAN_VERSION}/windows/VulkanSDK-${env:VULKAN_VERSION}-Installer.exe"
+          curl.exe -o $env:RUNNER_TEMP/VulkanSDK-Installer.exe -L "https://sdk.lunarg.com/sdk/download/${env:VULKAN_VERSION}/windows/vulkansdk-windows-X64-${env:VULKAN_VERSION}.exe"
           & "$env:RUNNER_TEMP\VulkanSDK-Installer.exe" --accept-licenses --default-answer --confirm-command install
           Add-Content $env:GITHUB_ENV "VULKAN_SDK=C:\VulkanSDK\${env:VULKAN_VERSION}"
           Add-Content $env:GITHUB_PATH "C:\VulkanSDK\${env:VULKAN_VERSION}\bin"
@@ -254,7 +254,7 @@ jobs:
 
       - 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' }}
+        if: ${{ matrix.build == 'cuda12' && (github.event_name == 'push' && github.ref == 'refs/heads/master' || github.event.inputs.create_release == 'true') }}
         run: |
           echo "Cuda install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
           $dst='.\build\bin\cudart\'
@@ -262,7 +262,7 @@ jobs:
           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' }}
+        if: ${{ matrix.build == 'cuda12' && (github.event_name == 'push' && github.ref == 'refs/heads/master' || github.event.inputs.create_release == 'true') }}
         uses: actions/upload-artifact@v4
         with:
           name: sd-cudart-sd-bin-win-cu12-x64.zip
@@ -288,6 +288,11 @@ jobs:
       - windows-latest-cmake
 
     steps:
+      - name: Clone
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
       - name: Download artifacts
         id: download-artifact
         uses: actions/download-artifact@v4
@@ -296,20 +301,27 @@ jobs:
           pattern: sd-*
           merge-multiple: true
 
+      - name: Get commit count
+        id: commit_count
+        run: |
+          echo "count=$(git rev-list --count HEAD)" >> $GITHUB_OUTPUT
+
       - name: Get commit hash
         id: commit
         uses: pr-mpt/actions-commit-hash@v2
 
       - name: Create release
         id: create_release
+        if: ${{ github.event_name == 'workflow_dispatch' || github.ref_name == 'master' }}
         uses: anzz1/action-create-release@v1
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         with:
-          tag_name: ${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}
+          tag_name: ${{ format('{0}-{1}-{2}', env.BRANCH_NAME, steps.commit_count.outputs.count, steps.commit.outputs.short) }}
 
       - name: Upload release
         id: upload_release
+        if: ${{ github.event_name == 'workflow_dispatch' || github.ref_name == 'master' }}
         uses: actions/github-script@v3
         with:
           github-token: ${{secrets.GITHUB_TOKEN}}

CMakeLists.txt

@@ -33,6 +33,7 @@ option(SD_SYCL                       "sd: sycl backend" OFF)
 option(SD_MUSA                       "sd: musa backend" 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_SHARED_GGML_LIB      "sd: build ggml as a separate shared lib" OFF)
 option(SD_USE_SYSTEM_GGML            "sd: use system-installed GGML library" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
@@ -86,18 +87,21 @@ file(GLOB SD_LIB_SOURCES
     "*.hpp"
 )
 
-# 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)
+    if(NOT SD_BUILD_SHARED_GGML_LIB)
+        set(BUILD_SHARED_LIBS OFF)
+    endif()
     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)
+    if(NOT SD_BUILD_SHARED_GGML_LIB)
+        set(BUILD_SHARED_LIBS OFF)
+    endif()
     add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
 endif()
 

Dockerfile

@@ -1,16 +1,21 @@
 ARG UBUNTU_VERSION=22.04
 
-FROM ubuntu:$UBUNTU_VERSION as build
+FROM ubuntu:$UBUNTU_VERSION AS build
 
-RUN apt-get update && apt-get install -y build-essential git cmake
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake
 
 WORKDIR /sd.cpp
 
 COPY . .
 
-RUN mkdir build && cd build && cmake .. && cmake --build . --config Release
+RUN cmake . -B ./build
+RUN cmake --build ./build --config Release --parallel
 
-FROM ubuntu:$UBUNTU_VERSION as runtime
+FROM ubuntu:$UBUNTU_VERSION AS runtime
+
+RUN apt-get update && \
+    apt-get install --yes --no-install-recommends libgomp1 && \
+    apt-get clean
 
 COPY --from=build /sd.cpp/build/bin/sd /sd
 

README.md

@@ -21,8 +21,10 @@ API and command-line option may change frequently.***
     - [SD3/SD3.5](./docs/sd3.md)
     - [Flux-dev/Flux-schnell](./docs/flux.md)
     - [Chroma](./docs/chroma.md)
+    - [Qwen Image](./docs/qwen_image.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
+    - [Qwen Image Edit/Qwen Image Edit 2509](./docs/qwen_image_edit.md)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
   - [PhotoMaker](https://github.com/TencentARC/PhotoMaker) support.
@@ -60,14 +62,6 @@ API and command-line option may change frequently.***
     - Windows
     - Android (via Termux, [Local Diffusion](https://github.com/rmatif/Local-Diffusion))
 
-### TODO
-
-- [ ] More sampling methods
-- [ ] Make inference faster
-    - The current implementation of ggml_conv_2d is slow and has high memory usage
-- [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
-- [ ] Implement Inpainting support
-
 ## Usage
 
 For most users, you can download the built executable program from the latest [release](https://github.com/leejet/stable-diffusion.cpp/releases/latest).
@@ -133,13 +127,14 @@ cmake --build . --config Release
 
 ##### Using HipBLAS
 This provides BLAS acceleration using the ROCm cores of your AMD GPU. Make sure to have the ROCm toolkit installed.
+To build for another GPU architecture than installed in your system, set `$GFX_NAME` manually to the desired architecture (replace first command). This is also necessary if your GPU is not officially supported by ROCm, for example you have to set `$GFX_NAME` manually to `gfx1030` for consumer RDNA2 cards.
 
 Windows User Refer to [docs/hipBLAS_on_Windows.md](docs%2FhipBLAS_on_Windows.md) for a comprehensive guide.
 
 ```
-export GFX_NAME=$(rocminfo | grep -m 1 -E "gfx[^0]{1}" | sed -e 's/ *Name: *//' | awk '{$1=$1; print}' || echo "rocminfo missing")
-echo $GFX_NAME
-cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DGPU_TARGETS=$GFX_NAME -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON
+if command -v rocminfo; then export GFX_NAME=$(rocminfo | awk '/ *Name: +gfx[1-9]/ {print $2; exit}'); else echo "rocminfo missing!"; fi
+if [ -z "${GFX_NAME}" ]; then echo "Error: Couldn't detect GPU!"; else echo "Building for GPU: ${GFX_NAME}"; fi
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DGPU_TARGETS=$GFX_NAME -DAMDGPU_TARGETS=$GFX_NAME -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_POSITION_INDEPENDENT_CODE=ON
 cmake --build . --config Release
 ```
 
@@ -292,7 +287,7 @@ usage: ./bin/sd [arguments]
 
 arguments:
   -h, --help                         show this help message and exit
-  -M, --mode [MODE]                  run mode, one of: [img_gen, vid_gen, convert], default: img_gen
+  -M, --mode [MODE]                  run mode, one of: [img_gen, vid_gen, upscale, convert], default: img_gen
   -t, --threads N                    number of threads to use during computation (default: -1)
                                      If threads <= 0, then threads will be set to the number of CPU physical cores
   --offload-to-cpu                   place the weights in RAM to save VRAM, and automatically load them into VRAM when needed
@@ -303,14 +298,13 @@ arguments:
   --clip_g                           path to the clip-g text encoder
   --clip_vision                      path to the clip-vision encoder
   --t5xxl                            path to the t5xxl text encoder
+  --qwen2vl                          path to the qwen2vl text encoder
+  --qwen2vl_vision                   path to the qwen2vl vit
   --vae [VAE]                        path to vae
   --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
   --control-net [CONTROL_PATH]       path to control net model
   --embd-dir [EMBEDDING_PATH]        path to embeddings
-  --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-model [ESRGAN_PATH]      path to esrgan model. For img_gen mode, 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 (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K)
                                      If not specified, the default is the type of the weight file
@@ -321,6 +315,9 @@ arguments:
   -i, --end-img [IMAGE]              path to the end image, required by flf2v
   --control-image [IMAGE]            path to image condition, control net
   -r, --ref-image [PATH]             reference image for Flux Kontext models (can be used multiple times)
+  --control-video [PATH]             path to control video frames, It must be a directory path.
+                                     The video frames inside should be stored as images in lexicographical (character) order
+                                     For example, if the control video path is `frames`, the directory contain images such as 00.png, 01.png, 鈥?etc.
   --increase-ref-index               automatically increase the indices of references images based on the order they are listed (starting with 1).
   -o, --output OUTPUT                path to write result image to (default: ./output.png)
   -p, --prompt [PROMPT]              the prompt to render
@@ -334,9 +331,10 @@ arguments:
   --skip-layers LAYERS               Layers to skip for SLG steps: (default: [7,8,9])
   --skip-layer-start START           SLG enabling point: (default: 0.01)
   --skip-layer-end END               SLG disabling point: (default: 0.2)
-  --scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)
+  --scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)
   --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}
                                      sampling method (default: "euler" for Flux/SD3/Wan, "euler_a" otherwise)
+  --timestep-shift N                 shift timestep for NitroFusion models, default: 0, recommended N for NitroSD-Realism around 250 and 500 for NitroSD-Vibrant
   --steps  STEPS                     number of sample steps (default: 20)
   --high-noise-cfg-scale SCALE       (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale SCALE   (high noise) image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)
@@ -347,13 +345,12 @@ arguments:
   --high-noise-skip-layers LAYERS    (high noise) Layers to skip for SLG steps: (default: [7,8,9])
   --high-noise-skip-layer-start      (high noise) SLG enabling point: (default: 0.01)
   --high-noise-skip-layer-end END    (high noise) SLG disabling point: (default: 0.2)
-  --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)
+  --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)
   --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}
                                      (high noise) sampling method (default: "euler_a")
   --high-noise-steps  STEPS          (high noise) number of sample steps (default: -1 = auto)
                                      SLG will be enabled at step int([STEPS]*[START]) and disabled at int([STEPS]*[END])
   --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)
@@ -361,9 +358,13 @@ arguments:
   --rng {std_default, cuda}          RNG (default: cuda)
   -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
   -b, --batch-count COUNT            number of images to generate
-  --clip-skip N                      ignore last_dot_pos layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)
+  --prediction {eps, v, edm_v, sd3_flow, flux_flow} Prediction type override
+  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)
                                      <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
   --vae-tiling                       process vae in tiles to reduce memory usage
+  --vae-tile-size [X]x[Y]            tile size for vae tiling (default: 32x32)
+  --vae-relative-tile-size [X]x[Y]   relative tile size for vae tiling, in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)
+  --vae-tile-overlap OVERLAP         tile overlap for vae tiling, in fraction of tile size (default: 0.5)
   --vae-on-cpu                       keep vae in cpu (for low vram)
   --clip-on-cpu                      keep clip in cpu (for low vram)
   --diffusion-fa                     use flash attention in the diffusion model (for low vram)
@@ -384,6 +385,11 @@ arguments:
   --moe-boundary BOUNDARY            timestep boundary for Wan2.2 MoE model. (default: 0.875)
                                      only enabled if `--high-noise-steps` is set to -1
   --flow-shift SHIFT                 shift value for Flow models like SD3.x or WAN (default: auto)
+  --vace-strength                    wan vace strength
+  --photo-maker                      path to PHOTOMAKER model
+  --pm-id-images-dir [DIR]           path to PHOTOMAKER input id images dir
+  --pm-id-embed-path [PATH]          path to PHOTOMAKER v2 id embed
+  --pm-style-strength                strength for keeping PHOTOMAKER input identity (default: 20)
   -v, --verbose                      print extra info
 ```
 
@@ -393,9 +399,9 @@ arguments:
 ./bin/sd -m ../models/sd-v1-4.ckpt -p "a lovely cat"
 # ./bin/sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat"
 # ./bin/sd -m ../models/sd_xl_base_1.0.safetensors --vae ../models/sdxl_vae-fp16-fix.safetensors -H 1024 -W 1024 -p "a lovely cat" -v
-# ./bin/sd -m ../models/sd3_medium_incl_clips_t5xxlfp16.safetensors -H 1024 -W 1024 -p 'a lovely cat holding a sign says \"Stable Diffusion CPP\"' --cfg-scale 4.5 --sampling-method euler -v
-# ./bin/sd --diffusion-model  ../models/flux1-dev-q3_k.gguf --vae ../models/ae.sft --clip_l ../models/clip_l.safetensors --t5xxl ../models/t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'" --cfg-scale 1.0 --sampling-method euler -v
-# ./bin/sd -m  ..\models\sd3.5_large.safetensors --clip_l ..\models\clip_l.safetensors --clip_g ..\models\clip_g.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -H 1024 -W 1024 -p 'a lovely cat holding a sign says \"Stable diffusion 3.5 Large\"' --cfg-scale 4.5 --sampling-method euler -v
+# ./bin/sd -m ../models/sd3_medium_incl_clips_t5xxlfp16.safetensors -H 1024 -W 1024 -p 'a lovely cat holding a sign says \"Stable Diffusion CPP\"' --cfg-scale 4.5 --sampling-method euler -v --clip-on-cpu
+# ./bin/sd --diffusion-model  ../models/flux1-dev-q3_k.gguf --vae ../models/ae.sft --clip_l ../models/clip_l.safetensors --t5xxl ../models/t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'" --cfg-scale 1.0 --sampling-method euler -v --clip-on-cpu
+# ./bin/sd -m  ..\models\sd3.5_large.safetensors --clip_l ..\models\clip_l.safetensors --clip_g ..\models\clip_g.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -H 1024 -W 1024 -p 'a lovely cat holding a sign says \"Stable diffusion 3.5 Large\"' --cfg-scale 4.5 --sampling-method euler -v --clip-on-cpu
 ```
 
 Using formats of different precisions will yield results of varying quality.
@@ -448,6 +454,7 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio
 - [Local Diffusion](https://github.com/rmatif/Local-Diffusion)
 - [sd.cpp-webui](https://github.com/daniandtheweb/sd.cpp-webui)
 - [LocalAI](https://github.com/mudler/LocalAI)
+- [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel)
 
 ## Contributors
 
@@ -462,6 +469,7 @@ Thank you to all the people who have already contributed to stable-diffusion.cpp
 ## References
 
 - [ggml](https://github.com/ggerganov/ggml)
+- [diffusers](https://github.com/huggingface/diffusers)
 - [stable-diffusion](https://github.com/CompVis/stable-diffusion)
 - [sd3-ref](https://github.com/Stability-AI/sd3-ref)
 - [stable-diffusion-stability-ai](https://github.com/Stability-AI/stablediffusion)
@@ -472,4 +480,4 @@ Thank you to all the people who have already contributed to stable-diffusion.cpp
 - [generative-models](https://github.com/Stability-AI/generative-models/)
 - [PhotoMaker](https://github.com/TencentARC/PhotoMaker)
 - [Wan2.1](https://github.com/Wan-Video/Wan2.1)
-- [Wan2.2](https://github.com/Wan-Video/Wan2.2)
+- [Wan2.2](https://github.com/Wan-Video/Wan2.2)

clip.hpp

@@ -6,7 +6,7 @@
 
 /*================================================== CLIPTokenizer ===================================================*/
 
-std::pair<std::unordered_map<std::string, float>, std::string> extract_and_remove_lora(std::string text) {
+__STATIC_INLINE__ std::pair<std::unordered_map<std::string, float>, std::string> extract_and_remove_lora(std::string text) {
     std::regex re("<lora:([^:]+):([^>]+)>");
     std::smatch matches;
     std::unordered_map<std::string, float> filename2multiplier;
@@ -31,7 +31,7 @@ std::pair<std::unordered_map<std::string, float>, std::string> extract_and_remov
     return std::make_pair(filename2multiplier, text);
 }
 
-std::vector<std::pair<int, std::u32string>> bytes_to_unicode() {
+__STATIC_INLINE__ std::vector<std::pair<int, std::u32string>> bytes_to_unicode() {
     std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;
     std::set<int> byte_set;
     for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
@@ -553,12 +553,12 @@ protected:
     void init_params(struct ggml_context* ctx, const String2GGMLType& tensor_types = {}, const std::string prefix = "") {
         enum ggml_type token_wtype = GGML_TYPE_F32;
         if (!force_clip_f32) {
-            auto tensor_type = tensor_types.find(prefix + "token_embedding.weight");
-            if (tensor_type != tensor_types.end())
-                token_wtype = tensor_type->second;
+            token_wtype = get_type(prefix + "token_embedding.weight", tensor_types, GGML_TYPE_F32);
+            if (!support_get_rows(token_wtype)) {
+                token_wtype = GGML_TYPE_F32;
+            }
         }
-        enum ggml_type position_wtype = GGML_TYPE_F32;
-
+        enum ggml_type position_wtype       = GGML_TYPE_F32;
         params["token_embedding.weight"]    = ggml_new_tensor_2d(ctx, token_wtype, embed_dim, vocab_size);
         params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions);
     }

common.hpp

@@ -177,7 +177,7 @@ public:
     }
 };
 
-class GEGLU : public GGMLBlock {
+class GEGLU : public UnaryBlock {
 protected:
     int64_t dim_in;
     int64_t dim_out;
@@ -216,23 +216,57 @@ public:
     }
 };
 
+class GELU : public UnaryBlock {
+public:
+    GELU(int64_t dim_in, int64_t dim_out, bool bias = true) {
+        blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out, bias));
+    }
+
+    struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
+        // x: [ne3, ne2, ne1, dim_in]
+        // return: [ne3, ne2, ne1, dim_out]
+        auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+
+        x = proj->forward(ctx, x);
+        x = ggml_gelu_inplace(ctx, x);
+        return x;
+    }
+};
+
 class FeedForward : public GGMLBlock {
 public:
+    enum class Activation {
+        GEGLU,
+        GELU
+    };
     FeedForward(int64_t dim,
                 int64_t dim_out,
-                int64_t mult = 4) {
+                int64_t mult          = 4,
+                Activation activation = Activation::GEGLU,
+                bool precision_fix    = false) {
         int64_t inner_dim = dim * mult;
+        if (activation == Activation::GELU) {
+            blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GELU(dim, inner_dim));
+        } else {
+            blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GEGLU(dim, inner_dim));
+        }
 
-        blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GEGLU(dim, inner_dim));
         // net_1 is nn.Dropout(), skip for inference
-        blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out));
+        float scale = 1.f;
+        if (precision_fix) {
+            scale = 1.f / 128.f;
+        }
+        // The purpose of the scale here is to prevent NaN issues in certain situations.
+        // For example, when using Vulkan without enabling force_prec_f32,
+        // or when using CUDA but the weights are k-quants.
+        blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, false, scale));
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         // x: [ne3, ne2, ne1, dim]
         // return: [ne3, ne2, ne1, dim_out]
 
-        auto net_0 = std::dynamic_pointer_cast<GEGLU>(blocks["net.0"]);
+        auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
         auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
 
         x = net_0->forward(ctx, x);  // [ne3, ne2, ne1, inner_dim]

conditioner.hpp

@@ -2,6 +2,7 @@
 #define __CONDITIONER_HPP__
 
 #include "clip.hpp"
+#include "qwenvl.hpp"
 #include "t5.hpp"
 
 struct SDCondition {
@@ -14,30 +15,34 @@ struct SDCondition {
         : c_crossattn(c_crossattn), c_vector(c_vector), c_concat(c_concat) {}
 };
 
+struct ConditionerParams {
+    std::string text;
+    int clip_skip                       = -1;
+    int width                           = -1;
+    int height                          = -1;
+    int adm_in_channels                 = -1;
+    bool zero_out_masked                = false;
+    int num_input_imgs                  = 0;   // for photomaker
+    std::vector<sd_image_t*> ref_images = {};  // for qwen image edit
+};
+
 struct Conditioner {
     virtual SDCondition get_learned_condition(ggml_context* work_ctx,
                                               int n_threads,
-                                              const std::string& text,
-                                              int clip_skip,
-                                              int width,
-                                              int height,
-                                              int adm_in_channels  = -1,
-                                              bool zero_out_masked = false)                                             = 0;
-    virtual void alloc_params_buffer()                                                                                  = 0;
-    virtual void free_params_buffer()                                                                                   = 0;
-    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors)                                 = 0;
-    virtual size_t get_params_buffer_size()                                                                             = 0;
+                                              const ConditionerParams& conditioner_params) = 0;
+    virtual void alloc_params_buffer()                                                     = 0;
+    virtual void free_params_buffer()                                                      = 0;
+    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors)    = 0;
+    virtual size_t get_params_buffer_size()                                                = 0;
     virtual std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(ggml_context* work_ctx,
                                                                                           int n_threads,
-                                                                                          const std::string& text,
-                                                                                          int clip_skip,
-                                                                                          int width,
-                                                                                          int height,
-                                                                                          int num_input_imgs,
-                                                                                          int adm_in_channels  = -1,
-                                                                                          bool zero_out_masked = false) = 0;
+                                                                                          const ConditionerParams& conditioner_params) {
+        GGML_ABORT("Not implemented yet!");
+    }
     virtual std::string remove_trigger_from_prompt(ggml_context* work_ctx,
-                                                   const std::string& prompt)                                           = 0;
+                                                   const std::string& prompt) {
+        GGML_ABORT("Not implemented yet!");
+    }
 };
 
 // ldm.modules.encoders.modules.FrozenCLIPEmbedder
@@ -141,7 +146,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             }
             return true;
         };
-        model_loader.load_tensors(on_load);
+        model_loader.load_tensors(on_load, 1);
         readed_embeddings.push_back(embd_name);
         if (embd) {
             int64_t hidden_size = text_model->model.hidden_size;
@@ -550,20 +555,14 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::tuple<SDCondition, std::vector<bool>>
     get_learned_condition_with_trigger(ggml_context* work_ctx,
                                        int n_threads,
-                                       const std::string& text,
-                                       int clip_skip,
-                                       int width,
-                                       int height,
-                                       int num_input_imgs,
-                                       int adm_in_channels  = -1,
-                                       bool zero_out_masked = false) {
+                                       const ConditionerParams& conditioner_params) {
         auto image_tokens = 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     = tokenize_with_trigger_token(text,
-                                                                  num_input_imgs,
+        auto tokens_and_weights     = tokenize_with_trigger_token(conditioner_params.text,
+                                                                  conditioner_params.num_input_imgs,
                                                                   image_tokens[0],
                                                                   true);
         std::vector<int>& tokens    = std::get<0>(tokens_and_weights);
@@ -577,7 +576,15 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         // for(int i = 0; i < clsm.size(); ++i)
         //    printf("%d ", clsm[i]?1:0);
         // printf("\n");
-        auto cond = get_learned_condition_common(work_ctx, n_threads, tokens, weights, clip_skip, width, height, adm_in_channels, zero_out_masked);
+        auto cond = get_learned_condition_common(work_ctx,
+                                                 n_threads,
+                                                 tokens,
+                                                 weights,
+                                                 conditioner_params.clip_skip,
+                                                 conditioner_params.width,
+                                                 conditioner_params.height,
+                                                 conditioner_params.adm_in_channels,
+                                                 conditioner_params.zero_out_masked);
         return std::make_tuple(cond, clsm);
     }
 
@@ -595,16 +602,19 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
 
     SDCondition get_learned_condition(ggml_context* work_ctx,
                                       int n_threads,
-                                      const std::string& text,
-                                      int clip_skip,
-                                      int width,
-                                      int height,
-                                      int adm_in_channels  = -1,
-                                      bool zero_out_masked = false) {
-        auto tokens_and_weights     = tokenize(text, true);
+                                      const ConditionerParams& conditioner_params) {
+        auto tokens_and_weights     = tokenize(conditioner_params.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, n_threads, tokens, weights, clip_skip, width, height, adm_in_channels, zero_out_masked);
+        return get_learned_condition_common(work_ctx,
+                                            n_threads,
+                                            tokens,
+                                            weights,
+                                            conditioner_params.clip_skip,
+                                            conditioner_params.width,
+                                            conditioner_params.height,
+                                            conditioner_params.adm_in_channels,
+                                            conditioner_params.zero_out_masked);
     }
 };
 
@@ -969,31 +979,13 @@ struct SD3CLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(ggml_context* work_ctx,
                                       int n_threads,
-                                      const std::string& text,
-                                      int clip_skip,
-                                      int width,
-                                      int height,
-                                      int adm_in_channels  = -1,
-                                      bool zero_out_masked = false) {
-        auto tokens_and_weights = tokenize(text, 77, true);
-        return get_learned_condition_common(work_ctx, n_threads, tokens_and_weights, clip_skip, zero_out_masked);
-    }
-
-    std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(ggml_context* work_ctx,
-                                                                                  int n_threads,
-                                                                                  const std::string& text,
-                                                                                  int clip_skip,
-                                                                                  int width,
-                                                                                  int height,
-                                                                                  int num_input_imgs,
-                                                                                  int adm_in_channels  = -1,
-                                                                                  bool zero_out_masked = false) {
-        GGML_ASSERT(0 && "Not implemented yet!");
-    }
-
-    std::string remove_trigger_from_prompt(ggml_context* work_ctx,
-                                           const std::string& prompt) {
-        GGML_ASSERT(0 && "Not implemented yet!");
+                                      const ConditionerParams& conditioner_params) {
+        auto tokens_and_weights = tokenize(conditioner_params.text, 77, true);
+        return get_learned_condition_common(work_ctx,
+                                            n_threads,
+                                            tokens_and_weights,
+                                            conditioner_params.clip_skip,
+                                            conditioner_params.zero_out_masked);
     }
 };
 
@@ -1186,31 +1178,13 @@ struct FluxCLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(ggml_context* work_ctx,
                                       int n_threads,
-                                      const std::string& text,
-                                      int clip_skip,
-                                      int width,
-                                      int height,
-                                      int adm_in_channels  = -1,
-                                      bool zero_out_masked = false) {
-        auto tokens_and_weights = tokenize(text, chunk_len, true);
-        return get_learned_condition_common(work_ctx, n_threads, tokens_and_weights, clip_skip, zero_out_masked);
-    }
-
-    std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(ggml_context* work_ctx,
-                                                                                  int n_threads,
-                                                                                  const std::string& text,
-                                                                                  int clip_skip,
-                                                                                  int width,
-                                                                                  int height,
-                                                                                  int num_input_imgs,
-                                                                                  int adm_in_channels  = -1,
-                                                                                  bool zero_out_masked = false) {
-        GGML_ASSERT(0 && "Not implemented yet!");
-    }
-
-    std::string remove_trigger_from_prompt(ggml_context* work_ctx,
-                                           const std::string& prompt) {
-        GGML_ASSERT(0 && "Not implemented yet!");
+                                      const ConditionerParams& conditioner_params) {
+        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
+        return get_learned_condition_common(work_ctx,
+                                            n_threads,
+                                            tokens_and_weights,
+                                            conditioner_params.clip_skip,
+                                            conditioner_params.zero_out_masked);
     }
 };
 
@@ -1389,31 +1363,214 @@ struct T5CLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(ggml_context* work_ctx,
                                       int n_threads,
-                                      const std::string& text,
-                                      int clip_skip,
-                                      int width,
-                                      int height,
-                                      int adm_in_channels  = -1,
-                                      bool zero_out_masked = false) {
-        auto tokens_and_weights = tokenize(text, chunk_len, true);
-        return get_learned_condition_common(work_ctx, n_threads, tokens_and_weights, clip_skip, zero_out_masked);
+                                      const ConditionerParams& conditioner_params) {
+        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
+        return get_learned_condition_common(work_ctx,
+                                            n_threads,
+                                            tokens_and_weights,
+                                            conditioner_params.clip_skip,
+                                            conditioner_params.zero_out_masked);
+    }
+};
+
+struct Qwen2_5_VLCLIPEmbedder : public Conditioner {
+    Qwen::Qwen2Tokenizer tokenizer;
+    std::shared_ptr<Qwen::Qwen2_5_VLRunner> qwenvl;
+
+    Qwen2_5_VLCLIPEmbedder(ggml_backend_t backend,
+                           bool offload_params_to_cpu,
+                           const String2GGMLType& tensor_types = {},
+                           const std::string prefix            = "",
+                           bool enable_vision                  = false) {
+        qwenvl = std::make_shared<Qwen::Qwen2_5_VLRunner>(backend,
+                                                          offload_params_to_cpu,
+                                                          tensor_types,
+                                                          "text_encoders.qwen2vl",
+                                                          enable_vision);
     }
 
-    std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(ggml_context* work_ctx,
-                                                                                  int n_threads,
-                                                                                  const std::string& text,
-                                                                                  int clip_skip,
-                                                                                  int width,
-                                                                                  int height,
-                                                                                  int num_input_imgs,
-                                                                                  int adm_in_channels  = -1,
-                                                                                  bool zero_out_masked = false) {
-        GGML_ASSERT(0 && "Not implemented yet!");
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) {
+        qwenvl->get_param_tensors(tensors, "text_encoders.qwen2vl");
     }
 
-    std::string remove_trigger_from_prompt(ggml_context* work_ctx,
-                                           const std::string& prompt) {
-        GGML_ASSERT(0 && "Not implemented yet!");
+    void alloc_params_buffer() {
+        qwenvl->alloc_params_buffer();
+    }
+
+    void free_params_buffer() {
+        qwenvl->free_params_buffer();
+    }
+
+    size_t get_params_buffer_size() {
+        size_t buffer_size = 0;
+        buffer_size += qwenvl->get_params_buffer_size();
+        return buffer_size;
+    }
+
+    std::tuple<std::vector<int>, std::vector<float>> tokenize(std::string text,
+                                                              size_t max_length           = 0,
+                                                              size_t system_prompt_length = 0,
+                                                              bool padding                = false) {
+        std::vector<std::pair<std::string, float>> parsed_attention;
+        if (system_prompt_length > 0) {
+            parsed_attention.emplace_back(text.substr(0, system_prompt_length), 1.f);
+            auto new_parsed_attention = parse_prompt_attention(text.substr(system_prompt_length, text.size() - system_prompt_length));
+            parsed_attention.insert(parsed_attention.end(),
+                                    new_parsed_attention.begin(),
+                                    new_parsed_attention.end());
+        } else {
+            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());
+        }
+
+        std::vector<int> tokens;
+        std::vector<float> weights;
+        for (const auto& item : parsed_attention) {
+            const std::string& curr_text = item.first;
+            float curr_weight            = item.second;
+            std::vector<int> curr_tokens = tokenizer.tokenize(curr_text, nullptr);
+            tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
+            weights.insert(weights.end(), curr_tokens.size(), curr_weight);
+        }
+
+        tokenizer.pad_tokens(tokens, weights, max_length, padding);
+
+        // for (int i = 0; i < tokens.size(); i++) {
+        //     std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl;
+        // }
+        // std::cout << std::endl;
+
+        return {tokens, weights};
+    }
+
+    SDCondition get_learned_condition(ggml_context* work_ctx,
+                                      int n_threads,
+                                      const ConditionerParams& conditioner_params) {
+        std::string prompt;
+        std::vector<std::pair<int, ggml_tensor*>> image_embeds;
+        size_t system_prompt_length = 0;
+        int prompt_template_encode_start_idx = 34;
+        if (qwenvl->enable_vision && conditioner_params.ref_images.size() > 0) {
+            LOG_INFO("QwenImageEditPlusPipeline");
+            prompt_template_encode_start_idx = 64;
+            int image_embed_idx              = 64 + 6;
+
+            int min_pixels          = 384 * 384;
+            int max_pixels          = 560 * 560;
+            std::string placeholder = "<|image_pad|>";
+            std::string img_prompt;
+
+            for (int i = 0; i < conditioner_params.ref_images.size(); i++) {
+                sd_image_f32_t image = sd_image_t_to_sd_image_f32_t(*conditioner_params.ref_images[i]);
+                double factor        = qwenvl->params.vision.patch_size * qwenvl->params.vision.spatial_merge_size;
+                int height           = image.height;
+                int width            = image.width;
+                int h_bar            = static_cast<int>(std::round(height / factor)) * factor;
+                int w_bar            = static_cast<int>(std::round(width / factor)) * factor;
+
+                if (static_cast<double>(h_bar) * w_bar > max_pixels) {
+                    double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
+                    h_bar       = std::max(static_cast<int>(factor),
+                                           static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
+                    w_bar       = std::max(static_cast<int>(factor),
+                                           static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
+                } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
+                    double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
+                    h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                    w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+                }
+
+                LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, image.height, image.width, h_bar, w_bar);
+
+                sd_image_f32_t resized_image = clip_preprocess(image, w_bar, h_bar);
+                free(image.data);
+                image.data = nullptr;
+
+                ggml_tensor* image_tensor = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, resized_image.width, resized_image.height, 3, 1);
+                sd_image_f32_to_tensor(resized_image, image_tensor, false);
+                free(resized_image.data);
+                resized_image.data = nullptr;
+
+                ggml_tensor* image_embed = nullptr;
+                qwenvl->encode_image(n_threads, image_tensor, &image_embed, work_ctx);
+                image_embeds.emplace_back(image_embed_idx, image_embed);
+                image_embed_idx += 1 + image_embed->ne[1] + 6;
+
+                img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";  // [24669, 220, index, 25, 220, 151652]
+                int64_t num_image_tokens = image_embed->ne[1];
+                img_prompt.reserve(num_image_tokens * placeholder.size());
+                for (int j = 0; j < num_image_tokens; j++) {
+                    img_prompt += placeholder;
+                }
+                img_prompt += "<|vision_end|>";
+            }
+
+            prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
+
+            system_prompt_length = prompt.size();
+
+            prompt += img_prompt;
+            prompt += conditioner_params.text;
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
+        } else {
+            prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n" + conditioner_params.text + "<|im_end|>\n<|im_start|>assistant\n";
+        }
+
+        auto tokens_and_weights = tokenize(prompt, 0, system_prompt_length, false);
+        auto& tokens            = std::get<0>(tokens_and_weights);
+        auto& weights           = std::get<1>(tokens_and_weights);
+
+        int64_t t0                        = ggml_time_ms();
+        struct ggml_tensor* hidden_states = NULL;  // [N, n_token, 3584]
+
+        auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
+
+        qwenvl->compute(n_threads,
+                        input_ids,
+                        image_embeds,
+                        &hidden_states,
+                        work_ctx);
+        {
+            auto tensor         = hidden_states;
+            float original_mean = ggml_tensor_mean(tensor);
+            for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+                for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                    for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                        float value = ggml_tensor_get_f32(tensor, i0, i1, i2);
+                        value *= weights[i1];
+                        ggml_tensor_set_f32(tensor, value, i0, i1, i2);
+                    }
+                }
+            }
+            float new_mean = ggml_tensor_mean(tensor);
+            ggml_tensor_scale(tensor, (original_mean / new_mean));
+        }
+
+        GGML_ASSERT(hidden_states->ne[1] > prompt_template_encode_start_idx);
+
+        ggml_tensor* new_hidden_states = ggml_new_tensor_3d(work_ctx,
+                                                            GGML_TYPE_F32,
+                                                            hidden_states->ne[0],
+                                                            hidden_states->ne[1] - prompt_template_encode_start_idx,
+                                                            hidden_states->ne[2]);
+
+        ggml_tensor_iter(new_hidden_states, [&](ggml_tensor* new_hidden_states, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+            float value = ggml_tensor_get_f32(hidden_states, i0, i1 + prompt_template_encode_start_idx, i2, i3);
+            ggml_tensor_set_f32(new_hidden_states, value, i0, i1, i2, i3);
+        });
+
+        int64_t t1 = ggml_time_ms();
+        LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+        return SDCondition(new_hidden_states, nullptr, nullptr);
     }
 };
 

control.hpp

@@ -445,7 +445,7 @@ struct ControlNet : public GGMLRunner {
         guided_hint_cached = true;
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading control net from '%s'", file_path.c_str());
         alloc_params_buffer();
         std::map<std::string, ggml_tensor*> tensors;
@@ -458,7 +458,7 @@ struct ControlNet : public GGMLRunner {
             return false;
         }
 
-        bool success = model_loader.load_tensors(tensors, ignore_tensors);
+        bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);
 
         if (!success) {
             LOG_ERROR("load control net tensors from model loader failed");

denoiser.hpp

@@ -232,6 +232,25 @@ struct GITSSchedule : SigmaSchedule {
     }
 };
 
+struct SGMUniformSchedule : SigmaSchedule {
+    std::vector<float> get_sigmas(uint32_t n, float sigma_min_in, float sigma_max_in, t_to_sigma_t t_to_sigma_func) override {
+        std::vector<float> result;
+        if (n == 0) {
+            result.push_back(0.0f);
+            return result;
+        }
+        result.reserve(n + 1);
+        int t_max                    = TIMESTEPS - 1;
+        int t_min                    = 0;
+        std::vector<float> timesteps = linear_space(static_cast<float>(t_max), static_cast<float>(t_min), n + 1);
+        for (int i = 0; i < n; i++) {
+            result.push_back(t_to_sigma_func(timesteps[i]));
+        }
+        result.push_back(0.0f);
+        return result;
+    }
+};
+
 struct KarrasSchedule : SigmaSchedule {
     std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) {
         // These *COULD* be function arguments here,
@@ -251,6 +270,35 @@ struct KarrasSchedule : SigmaSchedule {
     }
 };
 
+struct SimpleSchedule : SigmaSchedule {
+    std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
+        std::vector<float> result_sigmas;
+
+        if (n == 0) {
+            return result_sigmas;
+        }
+
+        result_sigmas.reserve(n + 1);
+
+        int model_sigmas_len = TIMESTEPS;
+
+        float step_factor = static_cast<float>(model_sigmas_len) / static_cast<float>(n);
+
+        for (uint32_t i = 0; i < n; ++i) {
+            int offset_from_start_of_py_array = static_cast<int>(static_cast<float>(i) * step_factor);
+            int timestep_index                = model_sigmas_len - 1 - offset_from_start_of_py_array;
+
+            if (timestep_index < 0) {
+                timestep_index = 0;
+            }
+
+            result_sigmas.push_back(t_to_sigma(static_cast<float>(timestep_index)));
+        }
+        result_sigmas.push_back(0.0f);
+        return result_sigmas;
+    }
+};
+
 // Close to Beta Schedule, but increadably simple in code.
 struct SmoothStepSchedule : SigmaSchedule {
     static constexpr float smoothstep(float x) {
@@ -722,7 +770,6 @@ static void sample_k_diffusion(sample_method_t method,
         } break;
         case DPMPP2S_A: {
             struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* d     = ggml_dup_tensor(work_ctx, x);
             struct ggml_tensor* x2    = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
@@ -737,22 +784,15 @@ static void sample_k_diffusion(sample_method_t method,
                 auto sigma_fn    = [](float t) -> float { return exp(-t); };
 
                 if (sigma_down == 0) {
-                    // Euler step
-                    float* vec_d        = (float*)d->data;
+                    // d = (x - denoised) / sigmas[i];
+                    // dt = sigma_down - sigmas[i];
+                    // x += d * dt;
+                    // => x = denoised
                     float* vec_x        = (float*)x->data;
                     float* vec_denoised = (float*)denoised->data;
 
-                    for (int j = 0; j < ggml_nelements(d); j++) {
-                        vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i];
-                    }
-
-                    // TODO: If sigma_down == 0, isn't this wrong?
-                    // But
-                    // https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L525
-                    // has this exactly the same way.
-                    float dt = sigma_down - sigmas[i];
-                    for (int j = 0; j < ggml_nelements(d); j++) {
-                        vec_x[j] = vec_x[j] + vec_d[j] * dt;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_x[j] = vec_denoised[j];
                     }
                 } else {
                     // DPM-Solver++(2S)
@@ -761,7 +801,6 @@ static void sample_k_diffusion(sample_method_t method,
                     float h      = t_next - t;
                     float s      = t + 0.5f * h;
 
-                    float* vec_d        = (float*)d->data;
                     float* vec_x        = (float*)x->data;
                     float* vec_x2       = (float*)x2->data;
                     float* vec_denoised = (float*)denoised->data;

diffusion_model.hpp

@@ -3,26 +3,33 @@
 
 #include "flux.hpp"
 #include "mmdit.hpp"
+#include "qwen_image.hpp"
 #include "unet.hpp"
 #include "wan.hpp"
 
+struct DiffusionParams {
+    struct ggml_tensor* x                     = NULL;
+    struct ggml_tensor* timesteps             = NULL;
+    struct ggml_tensor* context               = NULL;
+    struct ggml_tensor* c_concat              = NULL;
+    struct ggml_tensor* y                     = NULL;
+    struct ggml_tensor* guidance              = NULL;
+    std::vector<ggml_tensor*> ref_latents     = {};
+    bool increase_ref_index                   = false;
+    int num_video_frames                      = -1;
+    std::vector<struct ggml_tensor*> controls = {};
+    float control_strength                    = 0.f;
+    struct ggml_tensor* vace_context          = NULL;
+    float vace_strength                       = 1.f;
+    std::vector<int> skip_layers              = {};
+};
+
 struct DiffusionModel {
     virtual std::string get_desc()                                                      = 0;
     virtual void compute(int n_threads,
-                         struct ggml_tensor* x,
-                         struct ggml_tensor* timesteps,
-                         struct ggml_tensor* context,
-                         struct ggml_tensor* c_concat,
-                         struct ggml_tensor* y,
-                         struct ggml_tensor* guidance,
-                         std::vector<ggml_tensor*> ref_latents     = {},
-                         bool increase_ref_index                   = false,
-                         int num_video_frames                      = -1,
-                         std::vector<struct ggml_tensor*> controls = {},
-                         float control_strength                    = 0.f,
-                         struct ggml_tensor** output               = NULL,
-                         struct ggml_context* output_ctx           = NULL,
-                         std::vector<int> skip_layers              = std::vector<int>())             = 0;
+                         DiffusionParams diffusion_params,
+                         struct ggml_tensor** output     = NULL,
+                         struct ggml_context* output_ctx = NULL)                        = 0;
     virtual void alloc_params_buffer()                                                  = 0;
     virtual void free_params_buffer()                                                   = 0;
     virtual void free_compute_buffer()                                                  = 0;
@@ -71,22 +78,18 @@ struct UNetModel : public DiffusionModel {
     }
 
     void compute(int n_threads,
-                 struct ggml_tensor* x,
-                 struct ggml_tensor* timesteps,
-                 struct ggml_tensor* context,
-                 struct ggml_tensor* c_concat,
-                 struct ggml_tensor* y,
-                 struct ggml_tensor* guidance,
-                 std::vector<ggml_tensor*> ref_latents     = {},
-                 bool increase_ref_index                   = false,
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        (void)skip_layers;  // SLG doesn't work with UNet models
-        return unet.compute(n_threads, x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, output, output_ctx);
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = NULL,
+                 struct ggml_context* output_ctx = NULL) {
+        return unet.compute(n_threads,
+                            diffusion_params.x,
+                            diffusion_params.timesteps,
+                            diffusion_params.context,
+                            diffusion_params.c_concat,
+                            diffusion_params.y,
+                            diffusion_params.num_video_frames,
+                            diffusion_params.controls,
+                            diffusion_params.control_strength, output, output_ctx);
     }
 };
 
@@ -129,21 +132,17 @@ struct MMDiTModel : public DiffusionModel {
     }
 
     void compute(int n_threads,
-                 struct ggml_tensor* x,
-                 struct ggml_tensor* timesteps,
-                 struct ggml_tensor* context,
-                 struct ggml_tensor* c_concat,
-                 struct ggml_tensor* y,
-                 struct ggml_tensor* guidance,
-                 std::vector<ggml_tensor*> ref_latents     = {},
-                 bool increase_ref_index                   = false,
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        return mmdit.compute(n_threads, x, timesteps, context, y, output, output_ctx, skip_layers);
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = NULL,
+                 struct ggml_context* output_ctx = NULL) {
+        return mmdit.compute(n_threads,
+                             diffusion_params.x,
+                             diffusion_params.timesteps,
+                             diffusion_params.context,
+                             diffusion_params.y,
+                             output,
+                             output_ctx,
+                             diffusion_params.skip_layers);
     }
 };
 
@@ -188,21 +187,21 @@ struct FluxModel : public DiffusionModel {
     }
 
     void compute(int n_threads,
-                 struct ggml_tensor* x,
-                 struct ggml_tensor* timesteps,
-                 struct ggml_tensor* context,
-                 struct ggml_tensor* c_concat,
-                 struct ggml_tensor* y,
-                 struct ggml_tensor* guidance,
-                 std::vector<ggml_tensor*> ref_latents     = {},
-                 bool increase_ref_index                   = false,
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        return flux.compute(n_threads, x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, output, output_ctx, skip_layers);
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = NULL,
+                 struct ggml_context* output_ctx = NULL) {
+        return flux.compute(n_threads,
+                            diffusion_params.x,
+                            diffusion_params.timesteps,
+                            diffusion_params.context,
+                            diffusion_params.c_concat,
+                            diffusion_params.y,
+                            diffusion_params.guidance,
+                            diffusion_params.ref_latents,
+                            diffusion_params.increase_ref_index,
+                            output,
+                            output_ctx,
+                            diffusion_params.skip_layers);
     }
 };
 
@@ -248,21 +247,76 @@ struct WanModel : public DiffusionModel {
     }
 
     void compute(int n_threads,
-                 struct ggml_tensor* x,
-                 struct ggml_tensor* timesteps,
-                 struct ggml_tensor* context,
-                 struct ggml_tensor* c_concat,
-                 struct ggml_tensor* y,
-                 struct ggml_tensor* guidance,
-                 std::vector<ggml_tensor*> ref_latents     = {},
-                 bool increase_ref_index                   = false,
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        return wan.compute(n_threads, x, timesteps, context, y, c_concat, NULL, output, output_ctx);
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = NULL,
+                 struct ggml_context* output_ctx = NULL) {
+        return wan.compute(n_threads,
+                           diffusion_params.x,
+                           diffusion_params.timesteps,
+                           diffusion_params.context,
+                           diffusion_params.y,
+                           diffusion_params.c_concat,
+                           NULL,
+                           diffusion_params.vace_context,
+                           diffusion_params.vace_strength,
+                           output,
+                           output_ctx);
+    }
+};
+
+struct QwenImageModel : public DiffusionModel {
+    std::string prefix;
+    Qwen::QwenImageRunner qwen_image;
+
+    QwenImageModel(ggml_backend_t backend,
+                   bool offload_params_to_cpu,
+                   const String2GGMLType& tensor_types = {},
+                   const std::string prefix            = "model.diffusion_model",
+                   SDVersion version                   = VERSION_QWEN_IMAGE,
+                   bool flash_attn                     = false)
+        : prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_types, prefix, version, flash_attn) {
+    }
+
+    std::string get_desc() {
+        return qwen_image.get_desc();
+    }
+
+    void alloc_params_buffer() {
+        qwen_image.alloc_params_buffer();
+    }
+
+    void free_params_buffer() {
+        qwen_image.free_params_buffer();
+    }
+
+    void free_compute_buffer() {
+        qwen_image.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) {
+        qwen_image.get_param_tensors(tensors, prefix);
+    }
+
+    size_t get_params_buffer_size() {
+        return qwen_image.get_params_buffer_size();
+    }
+
+    int64_t get_adm_in_channels() {
+        return 768;
+    }
+
+    void compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = NULL,
+                 struct ggml_context* output_ctx = NULL) {
+        return qwen_image.compute(n_threads,
+                                  diffusion_params.x,
+                                  diffusion_params.timesteps,
+                                  diffusion_params.context,
+                                  diffusion_params.ref_latents,
+                                  true,  // increase_ref_index
+                                  output,
+                                  output_ctx);
     }
 };
 

docs/chroma.md

@@ -24,7 +24,7 @@ You can download the preconverted gguf weights from [silveroxides/Chroma-GGUF](h
 For example:
 
 ```
- .\bin\Release\sd.exe --diffusion-model  ..\models\chroma-unlocked-v40-q8_0.gguf --vae ..\models\ae.sft --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'chroma.cpp'" --cfg-scale 4.0 --sampling-method euler -v --chroma-disable-dit-mask
+ .\bin\Release\sd.exe --diffusion-model  ..\models\chroma-unlocked-v40-q8_0.gguf --vae ..\models\ae.sft --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'chroma.cpp'" --cfg-scale 4.0 --sampling-method euler -v --chroma-disable-dit-mask --clip-on-cpu
 ```
 
 ![](../assets/flux/chroma_v40.png)

docs/flux.md

@@ -28,7 +28,7 @@ Using fp16 will lead to overflow, but ggml's support for bf16 is not yet fully d
 For example:
 
 ```
- .\bin\Release\sd.exe --diffusion-model  ..\models\flux1-dev-q8_0.gguf --vae ..\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'" --cfg-scale 1.0 --sampling-method euler -v
+ .\bin\Release\sd.exe --diffusion-model  ..\models\flux1-dev-q8_0.gguf --vae ..\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'" --cfg-scale 1.0 --sampling-method euler -v --clip-on-cpu
 ```
 
 Using formats of different precisions will yield results of varying quality.
@@ -44,7 +44,7 @@ Using formats of different precisions will yield results of varying quality.
 
 
 ```
- .\bin\Release\sd.exe --diffusion-model  ..\models\flux1-schnell-q8_0.gguf --vae ..\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'" --cfg-scale 1.0 --sampling-method euler -v --steps 4
+ .\bin\Release\sd.exe --diffusion-model  ..\models\flux1-schnell-q8_0.gguf --vae ..\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'" --cfg-scale 1.0 --sampling-method euler -v --steps 4 --clip-on-cpu
 ```
 
 | q8_0  |
@@ -60,7 +60,7 @@ Since many flux LoRA training libraries have used various LoRA naming formats, i
 - LoRA model from https://huggingface.co/XLabs-AI/flux-lora-collection/tree/main (using comfy converted version!!!)
 
 ```
-.\bin\Release\sd.exe --diffusion-model  ..\models\flux1-dev-q8_0.gguf --vae ...\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'<lora:realism_lora_comfy_converted:1>" --cfg-scale 1.0 --sampling-method euler -v --lora-model-dir ../models
+.\bin\Release\sd.exe --diffusion-model  ..\models\flux1-dev-q8_0.gguf --vae ...\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'flux.cpp'<lora:realism_lora_comfy_converted:1>" --cfg-scale 1.0 --sampling-method euler -v --lora-model-dir ../models --clip-on-cpu
 ```
 
 ![output](../assets/flux/flux1-dev-q8_0%20with%20lora.png)

docs/kontext.md

@@ -27,7 +27,7 @@ You can download the preconverted gguf weights from [FLUX.1-Kontext-dev-GGUF](ht
 For example:
 
 ```
- .\bin\Release\sd.exe -r .\flux1-dev-q8_0.png --diffusion-model  ..\models\flux1-kontext-dev-q8_0.gguf --vae ..\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors -p "change 'flux.cpp' to 'kontext.cpp'" --cfg-scale 1.0 --sampling-method euler -v
+ .\bin\Release\sd.exe -r .\flux1-dev-q8_0.png --diffusion-model  ..\models\flux1-kontext-dev-q8_0.gguf --vae ..\models\ae.sft --clip_l ..\models\clip_l.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors -p "change 'flux.cpp' to 'kontext.cpp'" --cfg-scale 1.0 --sampling-method euler -v --clip-on-cpu
 ```
 
 

docs/lora.md

@@ -20,20 +20,30 @@ Here's a simple example:
 
 NOTE: The other backends may have different support.
 
-| Quant / Type | CUDA |
-|--------------|------|
-| F32          | ✔️   |
-| F16          | ✔️   |
-| BF16         | ✔️   |
-| I32          | ✔️   |
-| Q4_0         | ✔️   |
-| Q4_1         | ✔️   |
-| Q5_0         | ✔️   |
-| Q5_1         | ✔️   |
-| Q8_0         | ✔️   |
-| Q2_K         | ❌   |
-| Q3_K         | ❌   |
-| Q4_K         | ❌   |
-| Q5_K         | ❌   |
-| Q6_K         | ❌   |
-| Q8_K         | ❌   |
+| Quant / Type | CUDA | Vulkan |
+|--------------|------|--------|
+| F32          | ✔️   | ✔️   |
+| F16          | ✔️   | ✔️   |
+| BF16         | ✔️   | ✔️   |
+| I32          | ✔️   | ❌   |
+| Q4_0         | ✔️   | ✔️   |
+| Q4_1         | ✔️   | ✔️   |
+| Q5_0         | ✔️   | ✔️   |
+| Q5_1         | ✔️   | ✔️   |
+| Q8_0         | ✔️   | ✔️   |
+| Q2_K         | ❌   | ❌   |
+| Q3_K         | ❌   | ❌   |
+| Q4_K         | ❌   | ❌   |
+| Q5_K         | ❌   | ❌   |
+| Q6_K         | ❌   | ❌   |
+| Q8_K         | ❌   | ❌   |
+| IQ1_S        | ❌   | ✔️   |
+| IQ1_M        | ❌   | ✔️   |
+| IQ2_XXS      | ❌   | ✔️   |
+| IQ2_XS       | ❌   | ✔️   |
+| IQ2_S        | ❌   | ✔️   |
+| IQ3_XXS      | ❌   | ✔️   |
+| IQ3_S        | ❌   | ✔️   |
+| IQ4_XS       | ❌   | ✔️   |
+| IQ4_NL       | ❌   | ✔️   |
+| MXFP4        | ❌   | ✔️   |

docs/photo_maker.md

@@ -6,16 +6,15 @@ You can use [PhotoMaker](https://github.com/TencentARC/PhotoMaker) to personaliz
 
 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)
+- Specify the PhotoMaker model path using the `--photo-maker PATH` parameter.
+- Specify the input images path using the `--pm-id-images-dir PATH` parameter.
 
 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).
+- ```--pm-style-strength  (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:
 
@@ -28,7 +27,7 @@ If on low memory GPUs (<= 8GB), recommend running with ```--vae-on-cpu``` option
 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/photomaker_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
+bin/sd -m ../models/sdxlUnstableDiffusers_v11.safetensors  --vae ../models/sdxl_vae.safetensors --photo-maker ../models/photomaker-v1.safetensors --pm-id-images-dir ../assets/photomaker_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 --pm-style-strength 10 --vae-on-cpu --steps 50
 ```
 
 ## PhotoMaker Version 2

docs/qwen_image.md

@@ -0,0 +1,23 @@
+# How to Use
+
+## Download weights
+
+- Download Qwen Image
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/diffusion_models
+    - gguf: https://huggingface.co/QuantStack/Qwen-Image-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/vae
+- Download qwen_2.5_vl 7b
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/text_encoders
+    - gguf: https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-GGUF/tree/main
+
+## Examples
+
+```
+.\bin\Release\sd.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\qwen-image-Q8_0.gguf --vae ..\..\ComfyUI\models\vae\qwen_image_vae.safetensors  --qwen2vl ..\..\ComfyUI\models\text_encoders\Qwen2.5-VL-7B-Instruct-Q8_0.gguf  -p '一个穿着"QWEN"标志的T恤的中国美女正拿着黑色的马克笔面相镜头微笑。她身后的玻璃板上手写体写着 “一、Qwen-Image的技术路线： 探索视觉生成基础模型的极限，开创理解与生成一体化的未来。二、Qwen-Image的模型特色：1、复杂文字渲染。支持中英渲染、自动布局； 2、精准图像编辑。支持文字编辑、物体增减、风格变换。三、Qwen-Image的未来愿景：赋能专业内容创作、助力生成式AI发展。”' --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu -H 1024 -W 1024 --diffusion-fa --flow-shift 3
+```
+
+<img alt="qwen example" src="../assets/qwen/example.png" />
+
+
+

docs/qwen_image_edit.md

@@ -0,0 +1,35 @@
+# How to Use
+
+## Download weights
+
+- Download Qwen Image
+    - Qwen Image Edit
+        - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image-Edit_ComfyUI/tree/main/split_files/diffusion_models
+        - gguf: https://huggingface.co/QuantStack/Qwen-Image-Edit-GGUF/tree/main
+    - Qwen Image Edit 2509
+        - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image-Edit_ComfyUI/tree/main/split_files/diffusion_models
+        - gguf: https://huggingface.co/QuantStack/Qwen-Image-Edit-2509-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/vae
+- Download qwen_2.5_vl 7b
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/text_encoders
+    - gguf: https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-GGUF/tree/main
+
+## Examples
+
+### Qwen Image Edit
+
+```
+.\bin\Release\sd.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\Qwen_Image_Edit-Q8_0.gguf --vae ..\..\ComfyUI\models\vae\qwen_image_vae.safetensors  --qwen2vl ..\..\ComfyUI\models\text_encoders\qwen_2.5_vl_7b.safetensors --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'edit.cpp'" --seed 1118877715456453
+```
+
+<img alt="qwen_image_edit" src="../assets/qwen/qwen_image_edit.png" />
+
+
+### Qwen Image Edit 2509
+
+```
+.\bin\Release\sd.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\Qwen-Image-Edit-2509-Q4_K_S.gguf --vae ..\..\ComfyUI\models\vae\qwen_image_vae.safetensors  --qwen2vl ..\..\ComfyUI\models\text_encoders\Qwen2.5-VL-7B-Instruct-Q8_0.gguf --qwen2vl_vision ..\..\ComfyUI\models\text_encoders\Qwen2.5-VL-7B-Instruct.mmproj-Q8_0.gguf --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'Qwen Image Edit 2509'"
+```
+
+<img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2509.png" />

docs/sd3.md

@@ -14,7 +14,7 @@
 For example:
 
 ```
-.\bin\Release\sd.exe -m  ..\models\sd3.5_large.safetensors --clip_l ..\models\clip_l.safetensors --clip_g ..\models\clip_g.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -H 1024 -W 1024 -p 'a lovely cat holding a sign says \"Stable diffusion 3.5 Large\"' --cfg-scale 4.5 --sampling-method euler -v
+.\bin\Release\sd.exe -m  ..\models\sd3.5_large.safetensors --clip_l ..\models\clip_l.safetensors --clip_g ..\models\clip_g.safetensors --t5xxl ..\models\t5xxl_fp16.safetensors  -H 1024 -W 1024 -p 'a lovely cat holding a sign says \"Stable diffusion 3.5 Large\"' --cfg-scale 4.5 --sampling-method euler -v --clip-on-cpu
 ```
 
 ![](../assets/sd3.5_large.png)

docs/wan.md

@@ -18,6 +18,12 @@
         - Wan2.1 FLF2V 14B 720P
             - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/tree/main/split_files/diffusion_models
             - gguf: https://huggingface.co/city96/Wan2.1-FLF2V-14B-720P-gguf/tree/main
+        - Wan2.1 VACE 1.3B
+            - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/tree/main/split_files/diffusion_models
+            - gguf: https://huggingface.co/calcuis/wan-1.3b-gguf/tree/main
+        - Wan2.1 VACE 14B
+            - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/tree/main/split_files/diffusion_models
+            - gguf: https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/tree/main
     - Wan2.2
         - Wan2.2 TI2V 5B
             - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models
@@ -137,3 +143,62 @@
 ```
 
 <video src=../assets/wan/Wan2.2_14B_flf2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+### Wan2.1 VACE 1.3B
+
+#### T2V
+
+```
+.\bin\Release\sd.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\wan2.1-vace-1.3b-q8_0.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors --t5xxl ..\..\ComfyUI\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa --video-frames 1 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.1_1.3B_vace_t2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+
+#### R2V
+
+```
+.\bin\Release\sd.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\wan2.1-vace-1.3b-q8_0.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors --t5xxl ..\..\ComfyUI\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa -i ..\assets\cat_with_sd_cpp_42.png --video-frames 33 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.1_1.3B_vace_r2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+
+#### V2V
+
+```
+mkdir post+depth
+ffmpeg -i ..\..\ComfyUI\input\post+depth.mp4 -qscale:v 1 -vf fps=8 post+depth\frame_%04d.jpg
+.\bin\Release\sd.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\wan2.1-vace-1.3b-q8_0.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors --t5xxl ..\..\ComfyUI\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "The girl is dancing in a sea of flowers, slowly moving her hands. There is a close - up shot of her upper body. The character is surrounded by other transparent glass flowers in the style of Nicoletta Ceccoli, creating a beautiful, surreal, and emotionally expressive movie scene with a white. transparent feel and a dreamyl atmosphere." --cfg-scale 6.0 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 480 -H 832 --diffusion-fa -i ..\..\ComfyUI\input\dance_girl.jpg --control-video ./post+depth --video-frames 33 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.1_1.3B_vace_v2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+### Wan2.1 VACE 14B
+
+#### T2V
+
+```
+.\bin\Release\sd.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\Wan2.1_14B_VACE-Q8_0.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors --t5xxl ..\..\ComfyUI\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa --video-frames 33 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.1_14B_vace_t2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+
+#### R2V
+
+```
+.\bin\Release\sd.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\Wan2.1_14B_VACE-Q8_0.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors --t5xxl ..\..\ComfyUI\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa -i ..\assets\cat_with_sd_cpp_42.png --video-frames 33 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.1_14B_vace_r2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+
+
+#### V2V
+
+```
+.\bin\Release\sd.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\Wan2.1_14B_VACE-Q8_0.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors --t5xxl ..\..\ComfyUI\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "The girl is dancing in a sea of flowers, slowly moving her hands. There is a close - up shot of her upper body. The character is surrounded by other transparent glass flowers in the style of Nicoletta Ceccoli, creating a beautiful, surreal, and emotionally expressive movie scene with a white. transparent feel and a dreamyl atmosphere." --cfg-scale 6.0 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 480 -H 832 --diffusion-fa -i ..\..\ComfyUI\input\dance_girl.jpg --control-video ./post+depth --video-frames 33 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.1_14B_vace_v2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>

esrgan.hpp

@@ -83,39 +83,44 @@ public:
 
 class RRDBNet : public GGMLBlock {
 protected:
-    int scale       = 4;  // default RealESRGAN_x4plus_anime_6B
-    int num_block   = 6;  // default RealESRGAN_x4plus_anime_6B
+    int scale       = 4;
+    int num_block   = 23;
     int num_in_ch   = 3;
     int num_out_ch  = 3;
-    int num_feat    = 64;  // default RealESRGAN_x4plus_anime_6B
-    int num_grow_ch = 32;  // default RealESRGAN_x4plus_anime_6B
+    int num_feat    = 64;
+    int num_grow_ch = 32;
 
 public:
-    RRDBNet() {
+    RRDBNet(int scale, int num_block, int num_in_ch, int num_out_ch, int num_feat, int num_grow_ch)
+        : scale(scale), num_block(num_block), num_in_ch(num_in_ch), num_out_ch(num_out_ch), num_feat(num_feat), num_grow_ch(num_grow_ch) {
         blocks["conv_first"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_in_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
         for (int i = 0; i < num_block; i++) {
             std::string name = "body." + std::to_string(i);
             blocks[name]     = std::shared_ptr<GGMLBlock>(new RRDB(num_feat, num_grow_ch));
         }
         blocks["conv_body"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        // upsample
-        blocks["conv_up1"]  = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv_up2"]  = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        if (scale >= 2) {
+            blocks["conv_up1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        }
+        if (scale == 4) {
+            blocks["conv_up2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
+        }
         blocks["conv_hr"]   = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
         blocks["conv_last"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_out_ch, {3, 3}, {1, 1}, {1, 1}));
     }
 
+    int get_scale() { return scale; }
+    int get_num_block() { return num_block; }
+
     struct ggml_tensor* lrelu(struct ggml_context* ctx, struct ggml_tensor* x) {
         return ggml_leaky_relu(ctx, x, 0.2f, true);
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         // x: [n, num_in_ch, h, w]
-        // return: [n, num_out_ch, h*4, w*4]
+        // return: [n, num_out_ch, h*scale, w*scale]
         auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
         auto conv_body  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_body"]);
-        auto conv_up1   = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
-        auto conv_up2   = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
         auto conv_hr    = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
         auto conv_last  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
 
@@ -130,15 +135,22 @@ public:
         body_feat = conv_body->forward(ctx, body_feat);
         feat      = ggml_add(ctx, feat, body_feat);
         // upsample
-        feat     = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
-        feat     = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+        if (scale >= 2) {
+            auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
+            feat          = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+            if (scale == 4) {
+                auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
+                feat          = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+            }
+        }
+        // for all scales
         auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
         return out;
     }
 };
 
 struct ESRGAN : public GGMLRunner {
-    RRDBNet rrdb_net;
+    std::unique_ptr<RRDBNet> rrdb_net;
     int scale     = 4;
     int tile_size = 128;  // avoid cuda OOM for 4gb VRAM
 
@@ -146,12 +158,14 @@ struct ESRGAN : public GGMLRunner {
            bool offload_params_to_cpu,
            const String2GGMLType& tensor_types = {})
         : GGMLRunner(backend, offload_params_to_cpu) {
-        rrdb_net.init(params_ctx, tensor_types, "");
+        // rrdb_net will be created in load_from_file
     }
 
     void enable_conv2d_direct() {
+        if (!rrdb_net)
+            return;
         std::vector<GGMLBlock*> blocks;
-        rrdb_net.get_all_blocks(blocks);
+        rrdb_net->get_all_blocks(blocks);
         for (auto block : blocks) {
             if (block->get_desc() == "Conv2d") {
                 auto conv_block = (Conv2d*)block;
@@ -164,34 +178,188 @@ struct ESRGAN : public GGMLRunner {
         return "esrgan";
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading esrgan from '%s'", file_path.c_str());
 
-        alloc_params_buffer();
-        std::map<std::string, ggml_tensor*> esrgan_tensors;
-        rrdb_net.get_param_tensors(esrgan_tensors);
-
         ModelLoader model_loader;
         if (!model_loader.init_from_file(file_path)) {
             LOG_ERROR("init esrgan model loader from file failed: '%s'", file_path.c_str());
             return false;
         }
 
-        bool success = model_loader.load_tensors(esrgan_tensors);
+        // Get tensor names
+        auto tensor_names = model_loader.get_tensor_names();
+
+        // Detect if it's ESRGAN format
+        bool is_ESRGAN = std::find(tensor_names.begin(), tensor_names.end(), "model.0.weight") != tensor_names.end();
+
+        // Detect parameters from tensor names
+        int detected_num_block = 0;
+        if (is_ESRGAN) {
+            for (const auto& name : tensor_names) {
+                if (name.find("model.1.sub.") == 0) {
+                    size_t first_dot = name.find('.', 12);
+                    if (first_dot != std::string::npos) {
+                        size_t second_dot = name.find('.', first_dot + 1);
+                        if (second_dot != std::string::npos && name.substr(first_dot + 1, 3) == "RDB") {
+                            try {
+                                int idx            = std::stoi(name.substr(12, first_dot - 12));
+                                detected_num_block = std::max(detected_num_block, idx + 1);
+                            } catch (...) {
+                            }
+                        }
+                    }
+                }
+            }
+        } else {
+            // Original format
+            for (const auto& name : tensor_names) {
+                if (name.find("body.") == 0) {
+                    size_t pos = name.find('.', 5);
+                    if (pos != std::string::npos) {
+                        try {
+                            int idx            = std::stoi(name.substr(5, pos - 5));
+                            detected_num_block = std::max(detected_num_block, idx + 1);
+                        } catch (...) {
+                        }
+                    }
+                }
+            }
+        }
+
+        int detected_scale = 4;  // default
+        if (is_ESRGAN) {
+            // For ESRGAN format, detect scale by highest model number
+            int max_model_num = 0;
+            for (const auto& name : tensor_names) {
+                if (name.find("model.") == 0) {
+                    size_t dot_pos = name.find('.', 6);
+                    if (dot_pos != std::string::npos) {
+                        try {
+                            int num       = std::stoi(name.substr(6, dot_pos - 6));
+                            max_model_num = std::max(max_model_num, num);
+                        } catch (...) {
+                        }
+                    }
+                }
+            }
+            if (max_model_num <= 4) {
+                detected_scale = 1;
+            } else if (max_model_num <= 7) {
+                detected_scale = 2;
+            } else {
+                detected_scale = 4;
+            }
+        } else {
+            // Original format
+            bool has_conv_up2 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
+                return name == "conv_up2.weight";
+            });
+            bool has_conv_up1 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
+                return name == "conv_up1.weight";
+            });
+            if (has_conv_up2) {
+                detected_scale = 4;
+            } else if (has_conv_up1) {
+                detected_scale = 2;
+            } else {
+                detected_scale = 1;
+            }
+        }
+
+        int detected_num_in_ch   = 3;
+        int detected_num_out_ch  = 3;
+        int detected_num_feat    = 64;
+        int detected_num_grow_ch = 32;
+
+        // Create RRDBNet with detected parameters
+        rrdb_net = std::make_unique<RRDBNet>(detected_scale, detected_num_block, detected_num_in_ch, detected_num_out_ch, detected_num_feat, detected_num_grow_ch);
+        rrdb_net->init(params_ctx, {}, "");
+
+        alloc_params_buffer();
+        std::map<std::string, ggml_tensor*> esrgan_tensors;
+        rrdb_net->get_param_tensors(esrgan_tensors);
+
+        bool success;
+        if (is_ESRGAN) {
+            // Build name mapping for ESRGAN format
+            std::map<std::string, std::string> expected_to_model;
+            expected_to_model["conv_first.weight"] = "model.0.weight";
+            expected_to_model["conv_first.bias"]   = "model.0.bias";
+
+            for (int i = 0; i < detected_num_block; i++) {
+                for (int j = 1; j <= 3; j++) {
+                    for (int k = 1; k <= 5; k++) {
+                        std::string expected_weight        = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".weight";
+                        std::string model_weight           = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.weight";
+                        expected_to_model[expected_weight] = model_weight;
+
+                        std::string expected_bias        = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".bias";
+                        std::string model_bias           = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.bias";
+                        expected_to_model[expected_bias] = model_bias;
+                    }
+                }
+            }
+
+            if (detected_scale == 1) {
+                expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
+                expected_to_model["conv_body.bias"]   = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
+                expected_to_model["conv_hr.weight"]   = "model.2.weight";
+                expected_to_model["conv_hr.bias"]     = "model.2.bias";
+                expected_to_model["conv_last.weight"] = "model.4.weight";
+                expected_to_model["conv_last.bias"]   = "model.4.bias";
+            } else {
+                expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
+                expected_to_model["conv_body.bias"]   = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
+                if (detected_scale >= 2) {
+                    expected_to_model["conv_up1.weight"] = "model.3.weight";
+                    expected_to_model["conv_up1.bias"]   = "model.3.bias";
+                }
+                if (detected_scale == 4) {
+                    expected_to_model["conv_up2.weight"]  = "model.6.weight";
+                    expected_to_model["conv_up2.bias"]    = "model.6.bias";
+                    expected_to_model["conv_hr.weight"]   = "model.8.weight";
+                    expected_to_model["conv_hr.bias"]     = "model.8.bias";
+                    expected_to_model["conv_last.weight"] = "model.10.weight";
+                    expected_to_model["conv_last.bias"]   = "model.10.bias";
+                } else if (detected_scale == 2) {
+                    expected_to_model["conv_hr.weight"]   = "model.5.weight";
+                    expected_to_model["conv_hr.bias"]     = "model.5.bias";
+                    expected_to_model["conv_last.weight"] = "model.7.weight";
+                    expected_to_model["conv_last.bias"]   = "model.7.bias";
+                }
+            }
+
+            std::map<std::string, ggml_tensor*> model_tensors;
+            for (auto& p : esrgan_tensors) {
+                auto it = expected_to_model.find(p.first);
+                if (it != expected_to_model.end()) {
+                    model_tensors[it->second] = p.second;
+                }
+            }
+
+            success = model_loader.load_tensors(model_tensors, {}, n_threads);
+        } else {
+            success = model_loader.load_tensors(esrgan_tensors, {}, n_threads);
+        }
 
         if (!success) {
             LOG_ERROR("load esrgan tensors from model loader failed");
             return false;
         }
 
-        LOG_INFO("esrgan model loaded");
+        scale = rrdb_net->get_scale();
+        LOG_INFO("esrgan model loaded with scale=%d, num_block=%d", scale, detected_num_block);
         return success;
     }
 
     struct ggml_cgraph* build_graph(struct ggml_tensor* x) {
-        struct ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
-        x                       = to_backend(x);
-        struct ggml_tensor* out = rrdb_net.forward(compute_ctx, x);
+        if (!rrdb_net)
+            return nullptr;
+        constexpr int kGraphNodes = 1 << 16;  // 65k
+        struct ggml_cgraph* gf    = ggml_new_graph_custom(compute_ctx, kGraphNodes, /*grads*/ false);
+        x                         = to_backend(x);
+        struct ggml_tensor* out   = rrdb_net->forward(compute_ctx, x);
         ggml_build_forward_expand(gf, out);
         return gf;
     }

examples/cli/main.cpp

@@ -35,17 +35,21 @@
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
+namespace fs = std::filesystem;
+
 const char* modes_str[] = {
     "img_gen",
     "vid_gen",
     "convert",
+    "upscale",
 };
-#define SD_ALL_MODES_STR "img_gen, vid_gen, convert"
+#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale"
 
 enum SDMode {
     IMG_GEN,
     VID_GEN,
     CONVERT,
+    UPSCALE,
     MODE_COUNT
 };
 
@@ -57,6 +61,8 @@ struct SDParams {
     std::string clip_g_path;
     std::string clip_vision_path;
     std::string t5xxl_path;
+    std::string qwen2vl_path;
+    std::string qwen2vl_vision_path;
     std::string diffusion_model_path;
     std::string high_noise_diffusion_model_path;
     std::string vae_path;
@@ -64,8 +70,6 @@ struct SDParams {
     std::string esrgan_path;
     std::string control_net_path;
     std::string embedding_dir;
-    std::string stacked_id_embed_dir;
-    std::string input_id_images_path;
     sd_type_t wtype = SD_TYPE_COUNT;
     std::string tensor_type_rules;
     std::string lora_model_dir;
@@ -75,15 +79,16 @@ struct SDParams {
     std::string mask_image_path;
     std::string control_image_path;
     std::vector<std::string> ref_image_paths;
+    std::string control_video_path;
     bool increase_ref_index = false;
 
     std::string prompt;
     std::string negative_prompt;
-    float style_ratio = 20.f;
-    int clip_skip     = -1;  // <= 0 represents unspecified
-    int width         = 512;
-    int height        = 512;
-    int batch_count   = 1;
+
+    int clip_skip   = -1;  // <= 0 represents unspecified
+    int width       = 512;
+    int height      = 512;
+    int batch_count = 1;
 
     std::vector<int> skip_layers = {7, 8, 9};
     sd_sample_params_t sample_params;
@@ -91,20 +96,18 @@ struct SDParams {
     std::vector<int> high_noise_skip_layers = {7, 8, 9};
     sd_sample_params_t high_noise_sample_params;
 
-    float moe_boundary = 0.875f;
-
-    int video_frames = 1;
-    int fps          = 16;
+    float moe_boundary  = 0.875f;
+    int video_frames    = 1;
+    int fps             = 16;
+    float vace_strength = 1.f;
 
     float strength             = 0.75f;
     float control_strength     = 0.9f;
     rng_type_t rng_type        = CUDA_RNG;
     int64_t seed               = 42;
     bool verbose               = false;
-    bool vae_tiling            = false;
     bool offload_params_to_cpu = false;
     bool control_net_cpu       = false;
-    bool normalize_input       = false;
     bool clip_on_cpu           = false;
     bool vae_on_cpu            = false;
     bool diffusion_flash_attn  = false;
@@ -114,11 +117,21 @@ struct SDParams {
     bool color                 = false;
     int upscale_repeats        = 1;
 
+    // Photo Maker
+    std::string photo_maker_path;
+    std::string pm_id_images_dir;
+    std::string pm_id_embed_path;
+    float pm_style_strength = 20.f;
+
     bool chroma_use_dit_mask = true;
     bool chroma_use_t5_mask  = false;
     int chroma_t5_mask_pad   = 1;
     float flow_shift         = INFINITY;
 
+    prediction_t prediction = DEFAULT_PRED;
+
+    sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
+
     SDParams() {
         sd_sample_params_init(&sample_params);
         sd_sample_params_init(&high_noise_sample_params);
@@ -138,6 +151,8 @@ void print_params(SDParams params) {
     printf("    clip_g_path:                       %s\n", params.clip_g_path.c_str());
     printf("    clip_vision_path:                  %s\n", params.clip_vision_path.c_str());
     printf("    t5xxl_path:                        %s\n", params.t5xxl_path.c_str());
+    printf("    qwen2vl_path:                      %s\n", params.qwen2vl_path.c_str());
+    printf("    qwen2vl_vision_path:               %s\n", params.qwen2vl_vision_path.c_str());
     printf("    diffusion_model_path:              %s\n", params.diffusion_model_path.c_str());
     printf("    high_noise_diffusion_model_path:   %s\n", params.high_noise_diffusion_model_path.c_str());
     printf("    vae_path:                          %s\n", params.vae_path.c_str());
@@ -145,10 +160,10 @@ void print_params(SDParams params) {
     printf("    esrgan_path:                       %s\n", params.esrgan_path.c_str());
     printf("    control_net_path:                  %s\n", params.control_net_path.c_str());
     printf("    embedding_dir:                     %s\n", params.embedding_dir.c_str());
-    printf("    stacked_id_embed_dir:              %s\n", params.stacked_id_embed_dir.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("    normalize input image:             %s\n", params.normalize_input ? "true" : "false");
+    printf("    photo_maker_path:                  %s\n", params.photo_maker_path.c_str());
+    printf("    pm_id_images_dir:                  %s\n", params.pm_id_images_dir.c_str());
+    printf("    pm_id_embed_path:                  %s\n", params.pm_id_embed_path.c_str());
+    printf("    pm_style_strength:                 %.2f\n", params.pm_style_strength);
     printf("    output_path:                       %s\n", params.output_path.c_str());
     printf("    init_image_path:                   %s\n", params.init_image_path.c_str());
     printf("    end_image_path:                    %s\n", params.end_image_path.c_str());
@@ -158,6 +173,7 @@ void print_params(SDParams params) {
     for (auto& path : params.ref_image_paths) {
         printf("        %s\n", path.c_str());
     };
+    printf("    control_video_path:                %s\n", params.control_video_path.c_str());
     printf("    increase_ref_index:                %s\n", params.increase_ref_index ? "true" : "false");
     printf("    offload_params_to_cpu:             %s\n", params.offload_params_to_cpu ? "true" : "false");
     printf("    clip_on_cpu:                       %s\n", params.clip_on_cpu ? "true" : "false");
@@ -175,17 +191,19 @@ void print_params(SDParams params) {
     printf("    sample_params:                     %s\n", SAFE_STR(sample_params_str));
     printf("    high_noise_sample_params:          %s\n", SAFE_STR(high_noise_sample_params_str));
     printf("    moe_boundary:                      %.3f\n", params.moe_boundary);
+    printf("    prediction:                        %s\n", sd_prediction_name(params.prediction));
     printf("    flow_shift:                        %.2f\n", params.flow_shift);
     printf("    strength(img2img):                 %.2f\n", params.strength);
     printf("    rng:                               %s\n", sd_rng_type_name(params.rng_type));
-    printf("    seed:                              %ld\n", params.seed);
+    printf("    seed:                              %zd\n", params.seed);
     printf("    batch_count:                       %d\n", params.batch_count);
-    printf("    vae_tiling:                        %s\n", params.vae_tiling ? "true" : "false");
+    printf("    vae_tiling:                        %s\n", params.vae_tiling_params.enabled ? "true" : "false");
     printf("    upscale_repeats:                   %d\n", params.upscale_repeats);
     printf("    chroma_use_dit_mask:               %s\n", params.chroma_use_dit_mask ? "true" : "false");
     printf("    chroma_use_t5_mask:                %s\n", params.chroma_use_t5_mask ? "true" : "false");
     printf("    chroma_t5_mask_pad:                %d\n", params.chroma_t5_mask_pad);
     printf("    video_frames:                      %d\n", params.video_frames);
+    printf("    vace_strength:                     %.2f\n", params.vace_strength);
     printf("    fps:                               %d\n", params.fps);
     free(sample_params_str);
     free(high_noise_sample_params_str);
@@ -196,7 +214,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("\n");
     printf("arguments:\n");
     printf("  -h, --help                         show this help message and exit\n");
-    printf("  -M, --mode [MODE]                  run mode, one of: [img_gen, vid_gen, convert], default: img_gen\n");
+    printf("  -M, --mode [MODE]                  run mode, one of: [img_gen, vid_gen, upscale, convert], default: img_gen\n");
     printf("  -t, --threads N                    number of threads to use during computation (default: -1)\n");
     printf("                                     If threads <= 0, then threads will be set to the number of CPU physical cores\n");
     printf("  --offload-to-cpu                   place the weights in RAM to save VRAM, and automatically load them into VRAM when needed\n");
@@ -207,14 +225,13 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --clip_g                           path to the clip-g text encoder\n");
     printf("  --clip_vision                      path to the clip-vision encoder\n");
     printf("  --t5xxl                            path to the t5xxl text encoder\n");
+    printf("  --qwen2vl                          path to the qwen2vl text encoder\n");
+    printf("  --qwen2vl_vision                   path to the qwen2vl vit\n");
     printf("  --vae [VAE]                        path to vae\n");
     printf("  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n");
     printf("  --control-net [CONTROL_PATH]       path to control net model\n");
     printf("  --embd-dir [EMBEDDING_PATH]        path to embeddings\n");
-    printf("  --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-model [ESRGAN_PATH]      path to esrgan model. For img_gen mode, 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 (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K)\n");
     printf("                                     If not specified, the default is the type of the weight file\n");
@@ -225,6 +242,9 @@ void print_usage(int argc, const char* argv[]) {
     printf("  -i, --end-img [IMAGE]              path to the end image, required by flf2v\n");
     printf("  --control-image [IMAGE]            path to image condition, control net\n");
     printf("  -r, --ref-image [PATH]             reference image for Flux Kontext models (can be used multiple times) \n");
+    printf("  --control-video [PATH]             path to control video frames, It must be a directory path.\n");
+    printf("                                     The video frames inside should be stored as images in lexicographical (character) order\n");
+    printf("                                     For example, if the control video path is `frames`, the directory contain images such as 00.png, 01.png, … etc.\n");
     printf("  --increase-ref-index               automatically increase the indices of references images based on the order they are listed (starting with 1).\n");
     printf("  -o, --output OUTPUT                path to write result image to (default: ./output.png)\n");
     printf("  -p, --prompt [PROMPT]              the prompt to render\n");
@@ -238,9 +258,10 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --skip-layers LAYERS               Layers to skip for SLG steps: (default: [7,8,9])\n");
     printf("  --skip-layer-start START           SLG enabling point: (default: 0.01)\n");
     printf("  --skip-layer-end END               SLG disabling point: (default: 0.2)\n");
-    printf("  --scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)\n");
+    printf("  --scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)\n");
     printf("  --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}\n");
     printf("                                     sampling method (default: \"euler\" for Flux/SD3/Wan, \"euler_a\" otherwise)\n");
+    printf("  --timestep-shift N                 shift timestep for NitroFusion models, default: 0, recommended N for NitroSD-Realism around 250 and 500 for NitroSD-Vibrant\n");
     printf("  --steps  STEPS                     number of sample steps (default: 20)\n");
     printf("  --high-noise-cfg-scale SCALE       (high noise) unconditional guidance scale: (default: 7.0)\n");
     printf("  --high-noise-img-cfg-scale SCALE   (high noise) image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)\n");
@@ -251,13 +272,12 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --high-noise-skip-layers LAYERS    (high noise) Layers to skip for SLG steps: (default: [7,8,9])\n");
     printf("  --high-noise-skip-layer-start      (high noise) SLG enabling point: (default: 0.01)\n");
     printf("  --high-noise-skip-layer-end END    (high noise) SLG disabling point: (default: 0.2)\n");
-    printf("  --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)\n");
+    printf("  --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple} Denoiser sigma scheduler (default: discrete)\n");
     printf("  --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}\n");
     printf("                                     (high noise) sampling method (default: \"euler_a\")\n");
     printf("  --high-noise-steps  STEPS          (high noise) number of sample steps (default: -1 = auto)\n");
     printf("                                     SLG will be enabled at step int([STEPS]*[START]) and disabled at int([STEPS]*[END])\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");
@@ -265,9 +285,13 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --rng {std_default, cuda}          RNG (default: cuda)\n");
     printf("  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)\n");
     printf("  -b, --batch-count COUNT            number of images to generate\n");
-    printf("  --clip-skip N                      ignore last_dot_pos layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n");
+    printf("  --prediction {eps, v, edm_v, sd3_flow, flux_flow}        Prediction type override.\n");
+    printf("  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n");
     printf("                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n");
     printf("  --vae-tiling                       process vae in tiles to reduce memory usage\n");
+    printf("  --vae-tile-size [X]x[Y]            tile size for vae tiling (default: 32x32)\n");
+    printf("  --vae-relative-tile-size [X]x[Y]   relative tile size for vae tiling, in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)\n");
+    printf("  --vae-tile-overlap OVERLAP         tile overlap for vae tiling, in fraction of tile size (default: 0.5)\n");
     printf("  --vae-on-cpu                       keep vae in cpu (for low vram)\n");
     printf("  --clip-on-cpu                      keep clip in cpu (for low vram)\n");
     printf("  --diffusion-fa                     use flash attention in the diffusion model (for low vram)\n");
@@ -288,6 +312,11 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --moe-boundary BOUNDARY            timestep boundary for Wan2.2 MoE model. (default: 0.875)\n");
     printf("                                     only enabled if `--high-noise-steps` is set to -1\n");
     printf("  --flow-shift SHIFT                 shift value for Flow models like SD3.x or WAN (default: auto)\n");
+    printf("  --vace-strength                    wan vace strength\n");
+    printf("  --photo-maker                      path to PHOTOMAKER model\n");
+    printf("  --pm-id-images-dir [DIR]           path to PHOTOMAKER input id images dir\n");
+    printf("  --pm-id-embed-path [PATH]          path to PHOTOMAKER v2 id embed\n");
+    printf("  --pm-style-strength                strength for keeping PHOTOMAKER input identity (default: 20)\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -468,24 +497,27 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--clip_g", "", &params.clip_g_path},
         {"", "--clip_vision", "", &params.clip_vision_path},
         {"", "--t5xxl", "", &params.t5xxl_path},
+        {"", "--qwen2vl", "", &params.qwen2vl_path},
+        {"", "--qwen2vl_vision", "", &params.qwen2vl_vision_path},
         {"", "--diffusion-model", "", &params.diffusion_model_path},
         {"", "--high-noise-diffusion-model", "", &params.high_noise_diffusion_model_path},
         {"", "--vae", "", &params.vae_path},
         {"", "--taesd", "", &params.taesd_path},
         {"", "--control-net", "", &params.control_net_path},
         {"", "--embd-dir", "", &params.embedding_dir},
-        {"", "--stacked-id-embd-dir", "", &params.stacked_id_embed_dir},
         {"", "--lora-model-dir", "", &params.lora_model_dir},
         {"-i", "--init-img", "", &params.init_image_path},
         {"", "--end-img", "", &params.end_image_path},
         {"", "--tensor-type-rules", "", &params.tensor_type_rules},
-        {"", "--input-id-images-dir", "", &params.input_id_images_path},
+        {"", "--photo-maker", "", &params.photo_maker_path},
+        {"", "--pm-id-images-dir", "", &params.pm_id_images_dir},
+        {"", "--pm-id-embed-path", "", &params.pm_id_embed_path},
         {"", "--mask", "", &params.mask_image_path},
         {"", "--control-image", "", &params.control_image_path},
+        {"", "--control-video", "", &params.control_video_path},
         {"-o", "--output", "", &params.output_path},
         {"-p", "--prompt", "", &params.prompt},
         {"-n", "--negative-prompt", "", &params.negative_prompt},
-
         {"", "--upscale-model", "", &params.esrgan_path},
     };
 
@@ -501,6 +533,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--chroma-t5-mask-pad", "", &params.chroma_t5_mask_pad},
         {"", "--video-frames", "", &params.video_frames},
         {"", "--fps", "", &params.fps},
+        {"", "--timestep-shift", "", &params.sample_params.shifted_timestep},
     };
 
     options.float_options = {
@@ -519,17 +552,18 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--high-noise-skip-layer-end", "", &params.high_noise_sample_params.guidance.slg.layer_end},
         {"", "--high-noise-eta", "", &params.high_noise_sample_params.eta},
         {"", "--strength", "", &params.strength},
-        {"", "--style-ratio", "", &params.style_ratio},
+        {"", "--pm-style-strength", "", &params.pm_style_strength},
         {"", "--control-strength", "", &params.control_strength},
         {"", "--moe-boundary", "", &params.moe_boundary},
         {"", "--flow-shift", "", &params.flow_shift},
+        {"", "--vace-strength", "", &params.vace_strength},
+        {"", "--vae-tile-overlap", "", &params.vae_tiling_params.target_overlap},
     };
 
     options.bool_options = {
-        {"", "--vae-tiling", "", true, &params.vae_tiling},
+        {"", "--vae-tiling", "", true, &params.vae_tiling_params.enabled},
         {"", "--offload-to-cpu", "", true, &params.offload_params_to_cpu},
         {"", "--control-net-cpu", "", true, &params.control_net_cpu},
-        {"", "--normalize-input", "", true, &params.normalize_input},
         {"", "--clip-on-cpu", "", true, &params.clip_on_cpu},
         {"", "--vae-on-cpu", "", true, &params.vae_on_cpu},
         {"", "--diffusion-fa", "", true, &params.diffusion_flash_attn},
@@ -622,6 +656,20 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         return 1;
     };
 
+    auto on_prediction_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            return -1;
+        }
+        const char* arg   = argv[index];
+        params.prediction = str_to_prediction(arg);
+        if (params.prediction == PREDICTION_COUNT) {
+            fprintf(stderr, "error: invalid prediction type %s\n",
+                    arg);
+            return -1;
+        }
+        return 1;
+    };
+
     auto on_sample_method_arg = [&](int argc, const char** argv, int index) {
         if (++index >= argc) {
             return -1;
@@ -726,12 +774,59 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         return 1;
     };
 
+    auto on_tile_size_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            return -1;
+        }
+        std::string tile_size_str = argv[index];
+        size_t x_pos              = tile_size_str.find('x');
+        try {
+            if (x_pos != std::string::npos) {
+                std::string tile_x_str               = tile_size_str.substr(0, x_pos);
+                std::string tile_y_str               = tile_size_str.substr(x_pos + 1);
+                params.vae_tiling_params.tile_size_x = std::stoi(tile_x_str);
+                params.vae_tiling_params.tile_size_y = std::stoi(tile_y_str);
+            } else {
+                params.vae_tiling_params.tile_size_x = params.vae_tiling_params.tile_size_y = std::stoi(tile_size_str);
+            }
+        } catch (const std::invalid_argument& e) {
+            return -1;
+        } catch (const std::out_of_range& e) {
+            return -1;
+        }
+        return 1;
+    };
+
+    auto on_relative_tile_size_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            return -1;
+        }
+        std::string rel_size_str = argv[index];
+        size_t x_pos             = rel_size_str.find('x');
+        try {
+            if (x_pos != std::string::npos) {
+                std::string rel_x_str               = rel_size_str.substr(0, x_pos);
+                std::string rel_y_str               = rel_size_str.substr(x_pos + 1);
+                params.vae_tiling_params.rel_size_x = std::stof(rel_x_str);
+                params.vae_tiling_params.rel_size_y = std::stof(rel_y_str);
+            } else {
+                params.vae_tiling_params.rel_size_x = params.vae_tiling_params.rel_size_y = std::stof(rel_size_str);
+            }
+        } catch (const std::invalid_argument& e) {
+            return -1;
+        } catch (const std::out_of_range& e) {
+            return -1;
+        }
+        return 1;
+    };
+
     options.manual_options = {
         {"-M", "--mode", "", on_mode_arg},
         {"", "--type", "", on_type_arg},
         {"", "--rng", "", on_rng_arg},
         {"-s", "--seed", "", on_seed_arg},
         {"", "--sampling-method", "", on_sample_method_arg},
+        {"", "--prediction", "", on_prediction_arg},
         {"", "--scheduler", "", on_schedule_arg},
         {"", "--skip-layers", "", on_skip_layers_arg},
         {"", "--high-noise-sampling-method", "", on_high_noise_sample_method_arg},
@@ -739,6 +834,8 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--high-noise-skip-layers", "", on_high_noise_skip_layers_arg},
         {"-r", "--ref-image", "", on_ref_image_arg},
         {"-h", "--help", "", on_help_arg},
+        {"", "--vae-tile-size", "", on_tile_size_arg},
+        {"", "--vae-relative-tile-size", "", on_relative_tile_size_arg},
     };
 
     if (!parse_options(argc, argv, options)) {
@@ -750,13 +847,13 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         params.n_threads = get_num_physical_cores();
     }
 
-    if (params.mode != CONVERT && params.mode != VID_GEN && params.prompt.length() == 0) {
+    if ((params.mode == IMG_GEN || params.mode == VID_GEN) && params.prompt.length() == 0) {
         fprintf(stderr, "error: the following arguments are required: prompt\n");
         print_usage(argc, argv);
         exit(1);
     }
 
-    if (params.model_path.length() == 0 && params.diffusion_model_path.length() == 0) {
+    if (params.mode != UPSCALE && params.model_path.length() == 0 && params.diffusion_model_path.length() == 0) {
         fprintf(stderr, "error: the following arguments are required: model_path/diffusion_model\n");
         print_usage(argc, argv);
         exit(1);
@@ -806,11 +903,27 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         exit(1);
     }
 
+    if (params.sample_params.shifted_timestep < 0 || params.sample_params.shifted_timestep > 1000) {
+        fprintf(stderr, "error: timestep-shift must be between 0 and 1000\n");
+        exit(1);
+    }
+
     if (params.upscale_repeats < 1) {
         fprintf(stderr, "error: upscale multiplier must be at least 1\n");
         exit(1);
     }
 
+    if (params.mode == UPSCALE) {
+        if (params.esrgan_path.length() == 0) {
+            fprintf(stderr, "error: upscale mode needs an upscaler model (--upscale-model)\n");
+            exit(1);
+        }
+        if (params.init_image_path.length() == 0) {
+            fprintf(stderr, "error: upscale mode needs an init image (--init-img)\n");
+            exit(1);
+        }
+    }
+
     if (params.seed < 0) {
         srand((int)time(NULL));
         params.seed = rand();
@@ -821,14 +934,6 @@ void parse_args(int argc, const char** argv, SDParams& params) {
             params.output_path = "output.gguf";
         }
     }
-
-    if (!isfinite(params.sample_params.guidance.img_cfg)) {
-        params.sample_params.guidance.img_cfg = params.sample_params.guidance.txt_cfg;
-    }
-
-    if (!isfinite(params.high_noise_sample_params.guidance.img_cfg)) {
-        params.high_noise_sample_params.guidance.img_cfg = params.high_noise_sample_params.guidance.txt_cfg;
-    }
 }
 
 static std::string sd_basename(const std::string& path) {
@@ -871,7 +976,7 @@ std::string get_image_params(SDParams params, int64_t seed) {
         parameter_string += " " + std::string(sd_schedule_name(params.sample_params.scheduler));
     }
     parameter_string += ", ";
-    for (const auto& te : {params.clip_l_path, params.clip_g_path, params.t5xxl_path}) {
+    for (const auto& te : {params.clip_l_path, params.clip_g_path, params.t5xxl_path, params.qwen2vl_path, params.qwen2vl_vision_path}) {
         if (!te.empty()) {
             parameter_string += "TE: " + sd_basename(te) + ", ";
         }
@@ -1012,14 +1117,67 @@ uint8_t* load_image(const char* image_path, int& width, int& height, int expecte
                      STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
                      STBIR_FILTER_BOX, STBIR_FILTER_BOX,
                      STBIR_COLORSPACE_SRGB, nullptr);
-
-        // Save resized result
+        width  = resized_width;
+        height = resized_height;
         free(image_buffer);
         image_buffer = resized_image_buffer;
     }
     return image_buffer;
 }
 
+bool load_images_from_dir(const std::string dir,
+                          std::vector<sd_image_t>& images,
+                          int expected_width  = 0,
+                          int expected_height = 0,
+                          int max_image_num   = 0,
+                          bool verbose        = false) {
+    if (!fs::exists(dir) || !fs::is_directory(dir)) {
+        fprintf(stderr, "'%s' is not a valid directory\n", dir.c_str());
+        return false;
+    }
+
+    std::vector<fs::directory_entry> entries;
+    for (const auto& entry : fs::directory_iterator(dir)) {
+        if (entry.is_regular_file()) {
+            entries.push_back(entry);
+        }
+    }
+
+    std::sort(entries.begin(), entries.end(),
+              [](const fs::directory_entry& a, const fs::directory_entry& b) {
+                  return a.path().filename().string() < b.path().filename().string();
+              });
+
+    for (const auto& entry : entries) {
+        std::string path = entry.path().string();
+        std::string ext  = entry.path().extension().string();
+        std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
+
+        if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp") {
+            if (verbose) {
+                printf("load image %zu from '%s'\n", images.size(), path.c_str());
+            }
+            int width             = 0;
+            int height            = 0;
+            uint8_t* image_buffer = load_image(path.c_str(), width, height, expected_width, expected_height);
+            if (image_buffer == NULL) {
+                fprintf(stderr, "load image from '%s' failed\n", path.c_str());
+                return false;
+            }
+
+            images.push_back({(uint32_t)width,
+                              (uint32_t)height,
+                              3,
+                              image_buffer});
+
+            if (max_image_num > 0 && images.size() >= max_image_num) {
+                break;
+            }
+        }
+    }
+    return true;
+}
+
 int main(int argc, const char* argv[]) {
     SDParams params;
     parse_args(argc, argv, params);
@@ -1059,17 +1217,29 @@ int main(int argc, const char* argv[]) {
     sd_image_t control_image = {(uint32_t)params.width, (uint32_t)params.height, 3, NULL};
     sd_image_t mask_image    = {(uint32_t)params.width, (uint32_t)params.height, 1, NULL};
     std::vector<sd_image_t> ref_images;
+    std::vector<sd_image_t> pmid_images;
+    std::vector<sd_image_t> control_frames;
 
     auto release_all_resources = [&]() {
         free(init_image.data);
         free(end_image.data);
         free(control_image.data);
         free(mask_image.data);
-        for (auto ref_image : ref_images) {
-            free(ref_image.data);
-            ref_image.data = NULL;
+        for (auto image : ref_images) {
+            free(image.data);
+            image.data = NULL;
         }
         ref_images.clear();
+        for (auto image : pmid_images) {
+            free(image.data);
+            image.data = NULL;
+        }
+        pmid_images.clear();
+        for (auto image : control_frames) {
+            free(image.data);
+            image.data = NULL;
+        }
+        control_frames.clear();
     };
 
     if (params.init_image_path.size() > 0) {
@@ -1118,7 +1288,7 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    if (params.control_net_path.size() > 0 && params.control_image_path.size() > 0) {
+    if (params.control_image_path.size() > 0) {
         int width          = 0;
         int height         = 0;
         control_image.data = load_image(params.control_image_path.c_str(), width, height, params.width, params.height);
@@ -1128,14 +1298,12 @@ int main(int argc, const char* argv[]) {
             return 1;
         }
         if (params.canny_preprocess) {  // apply preprocessor
-            control_image.data = preprocess_canny(control_image.data,
-                                                  control_image.width,
-                                                  control_image.height,
-                                                  0.08f,
-                                                  0.08f,
-                                                  0.8f,
-                                                  1.0f,
-                                                  false);
+            preprocess_canny(control_image,
+                             0.08f,
+                             0.08f,
+                             0.8f,
+                             1.0f,
+                             false);
         }
     }
 
@@ -1157,6 +1325,30 @@ int main(int argc, const char* argv[]) {
         }
     }
 
+    if (!params.control_video_path.empty()) {
+        if (!load_images_from_dir(params.control_video_path,
+                                  control_frames,
+                                  params.width,
+                                  params.height,
+                                  params.video_frames,
+                                  params.verbose)) {
+            release_all_resources();
+            return 1;
+        }
+    }
+
+    if (!params.pm_id_images_dir.empty()) {
+        if (!load_images_from_dir(params.pm_id_images_dir,
+                                  pmid_images,
+                                  0,
+                                  0,
+                                  0,
+                                  params.verbose)) {
+            release_all_resources();
+            return 1;
+        }
+    }
+
     if (params.mode == VID_GEN) {
         vae_decode_only = false;
     }
@@ -1167,6 +1359,8 @@ int main(int argc, const char* argv[]) {
         params.clip_g_path.c_str(),
         params.clip_vision_path.c_str(),
         params.t5xxl_path.c_str(),
+        params.qwen2vl_path.c_str(),
+        params.qwen2vl_vision_path.c_str(),
         params.diffusion_model_path.c_str(),
         params.high_noise_diffusion_model_path.c_str(),
         params.vae_path.c_str(),
@@ -1174,13 +1368,13 @@ int main(int argc, const char* argv[]) {
         params.control_net_path.c_str(),
         params.lora_model_dir.c_str(),
         params.embedding_dir.c_str(),
-        params.stacked_id_embed_dir.c_str(),
+        params.photo_maker_path.c_str(),
         vae_decode_only,
-        params.vae_tiling,
         true,
         params.n_threads,
         params.wtype,
         params.rng_type,
+        params.prediction,
         params.offload_params_to_cpu,
         params.clip_on_cpu,
         params.control_net_cpu,
@@ -1194,69 +1388,92 @@ int main(int argc, const char* argv[]) {
         params.flow_shift,
     };
 
-    sd_ctx_t* sd_ctx = new_sd_ctx(&sd_ctx_params);
+    sd_image_t* results = nullptr;
+    int num_results     = 0;
 
-    if (sd_ctx == NULL) {
-        printf("new_sd_ctx_t failed\n");
-        release_all_resources();
-        return 1;
-    }
+    if (params.mode == UPSCALE) {
+        num_results = 1;
+        results     = (sd_image_t*)calloc(num_results, sizeof(sd_image_t));
+        if (results == NULL) {
+            printf("failed to allocate results array\n");
+            release_all_resources();
+            return 1;
+        }
 
-    if (params.sample_params.sample_method == SAMPLE_METHOD_DEFAULT) {
-        params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
-    }
+        results[0]      = init_image;
+        init_image.data = NULL;
+    } else {
+        sd_ctx_t* sd_ctx = new_sd_ctx(&sd_ctx_params);
 
-    sd_image_t* results;
-    int num_results = 1;
-    if (params.mode == IMG_GEN) {
-        sd_img_gen_params_t img_gen_params = {
-            params.prompt.c_str(),
-            params.negative_prompt.c_str(),
-            params.clip_skip,
-            init_image,
-            ref_images.data(),
-            (int)ref_images.size(),
-            params.increase_ref_index,
-            mask_image,
-            params.width,
-            params.height,
-            params.sample_params,
-            params.strength,
-            params.seed,
-            params.batch_count,
-            control_image,
-            params.control_strength,
-            params.style_ratio,
-            params.normalize_input,
-            params.input_id_images_path.c_str(),
-        };
+        if (sd_ctx == NULL) {
+            printf("new_sd_ctx_t failed\n");
+            release_all_resources();
+            return 1;
+        }
 
-        results     = generate_image(sd_ctx, &img_gen_params);
-        num_results = params.batch_count;
-    } else if (params.mode == VID_GEN) {
-        sd_vid_gen_params_t vid_gen_params = {
-            params.prompt.c_str(),
-            params.negative_prompt.c_str(),
-            params.clip_skip,
-            init_image,
-            end_image,
-            params.width,
-            params.height,
-            params.sample_params,
-            params.high_noise_sample_params,
-            params.moe_boundary,
-            params.strength,
-            params.seed,
-            params.video_frames,
-        };
+        if (params.sample_params.sample_method == SAMPLE_METHOD_DEFAULT) {
+            params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
+        }
 
-        results = generate_video(sd_ctx, &vid_gen_params, &num_results);
-    }
+        if (params.mode == IMG_GEN) {
+            sd_img_gen_params_t img_gen_params = {
+                params.prompt.c_str(),
+                params.negative_prompt.c_str(),
+                params.clip_skip,
+                init_image,
+                ref_images.data(),
+                (int)ref_images.size(),
+                params.increase_ref_index,
+                mask_image,
+                params.width,
+                params.height,
+                params.sample_params,
+                params.strength,
+                params.seed,
+                params.batch_count,
+                control_image,
+                params.control_strength,
+                {
+                    pmid_images.data(),
+                    (int)pmid_images.size(),
+                    params.pm_id_embed_path.c_str(),
+                    params.pm_style_strength,
+                },  // pm_params
+                params.vae_tiling_params,
+            };
+
+            results     = generate_image(sd_ctx, &img_gen_params);
+            num_results = params.batch_count;
+        } else if (params.mode == VID_GEN) {
+            sd_vid_gen_params_t vid_gen_params = {
+                params.prompt.c_str(),
+                params.negative_prompt.c_str(),
+                params.clip_skip,
+                init_image,
+                end_image,
+                control_frames.data(),
+                (int)control_frames.size(),
+                params.width,
+                params.height,
+                params.sample_params,
+                params.high_noise_sample_params,
+                params.moe_boundary,
+                params.strength,
+                params.seed,
+                params.video_frames,
+                params.vace_strength,
+            };
+
+            results = generate_video(sd_ctx, &vid_gen_params, &num_results);
+        }
+
+        if (results == NULL) {
+            printf("generate failed\n");
+            free_sd_ctx(sd_ctx);
+            return 1;
+        }
 
-    if (results == NULL) {
-        printf("generate failed\n");
         free_sd_ctx(sd_ctx);
-        return 1;
     }
 
     int upscale_factor = 4;  // unused for RealESRGAN_x4plus_anime_6B.pth
@@ -1269,7 +1486,7 @@ int main(int argc, const char* argv[]) {
         if (upscaler_ctx == NULL) {
             printf("new_upscaler_ctx failed\n");
         } else {
-            for (int i = 0; i < params.batch_count; i++) {
+            for (int i = 0; i < num_results; i++) {
                 if (results[i].data == NULL) {
                     continue;
                 }
@@ -1290,7 +1507,6 @@ int main(int argc, const char* argv[]) {
 
     // create directory if not exists
     {
-        namespace fs            = std::filesystem;
         const fs::path out_path = params.output_path;
         if (const fs::path out_dir = out_path.parent_path(); !out_dir.empty()) {
             std::error_code ec;
@@ -1356,7 +1572,6 @@ int main(int argc, const char* argv[]) {
         results[i].data = NULL;
     }
     free(results);
-    free_sd_ctx(sd_ctx);
 
     release_all_resources();
 

flux.hpp

@@ -81,56 +81,6 @@ namespace Flux {
         }
     };
 
-    __STATIC_INLINE__ struct ggml_tensor* apply_rope(struct ggml_context* ctx,
-                                                     struct ggml_tensor* x,
-                                                     struct ggml_tensor* pe) {
-        // x: [N, L, n_head, d_head]
-        // pe: [L, d_head/2, 2, 2]
-        int64_t d_head = x->ne[0];
-        int64_t n_head = x->ne[1];
-        int64_t L      = x->ne[2];
-        int64_t N      = x->ne[3];
-        x              = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));       // [N, n_head, L, d_head]
-        x              = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N);  // [N * n_head, L, d_head/2, 2]
-        x              = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2));       // [2, N * n_head, L, d_head/2]
-
-        int64_t offset = x->nb[2] * x->ne[2];
-        auto x_0       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0);  // [N * n_head, L, d_head/2]
-        auto x_1       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1);  // [N * n_head, L, d_head/2]
-        x_0            = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]);                    // [N * n_head, L, d_head/2, 1]
-        x_1            = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]);                    // [N * n_head, L, d_head/2, 1]
-        auto temp_x    = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]);
-        x_0            = ggml_repeat(ctx, x_0, temp_x);  // [N * n_head, L, d_head/2, 2]
-        x_1            = ggml_repeat(ctx, x_1, temp_x);  // [N * n_head, L, d_head/2, 2]
-
-        pe        = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2));  // [2, L, d_head/2, 2]
-        offset    = pe->nb[2] * pe->ne[2];
-        auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0);  // [L, d_head/2, 2]
-        auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1);  // [L, d_head/2, 2]
-
-        auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1));  // [N * n_head, L, d_head/2, 2]
-        x_out      = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N);                         // [N*n_head, L, d_head]
-        return x_out;
-    }
-
-    __STATIC_INLINE__ struct ggml_tensor* attention(struct ggml_context* ctx,
-                                                    ggml_backend_t backend,
-                                                    struct ggml_tensor* q,
-                                                    struct ggml_tensor* k,
-                                                    struct ggml_tensor* v,
-                                                    struct ggml_tensor* pe,
-                                                    struct ggml_tensor* mask,
-                                                    bool flash_attn) {
-        // q,k,v: [N, L, n_head, d_head]
-        // pe: [L, d_head/2, 2, 2]
-        // return: [N, L, n_head*d_head]
-        q = apply_rope(ctx, q, pe);  // [N*n_head, L, d_head]
-        k = apply_rope(ctx, k, pe);  // [N*n_head, L, d_head]
-
-        auto x = ggml_nn_attention_ext(ctx, backend, q, k, v, v->ne[1], mask, false, true, flash_attn);  // [N, L, n_head*d_head]
-        return x;
-    }
-
     struct SelfAttention : public GGMLBlock {
     public:
         int64_t num_heads;
@@ -178,9 +128,9 @@ namespace Flux {
             // x: [N, n_token, dim]
             // pe: [n_token, d_head/2, 2, 2]
             // return [N, n_token, dim]
-            auto qkv = pre_attention(ctx, x);                                                  // q,k,v: [N, n_token, n_head, d_head]
-            x        = attention(ctx, backend, qkv[0], qkv[1], qkv[2], pe, mask, flash_attn);  // [N, n_token, dim]
-            x        = post_attention(ctx, x);                                                 // [N, n_token, dim]
+            auto qkv = pre_attention(ctx, x);                                                        // q,k,v: [N, n_token, n_head, d_head]
+            x        = Rope::attention(ctx, backend, qkv[0], qkv[1], qkv[2], pe, mask, flash_attn);  // [N, n_token, dim]
+            x        = post_attention(ctx, x);                                                       // [N, n_token, dim]
             return x;
         }
     };
@@ -368,8 +318,8 @@ namespace Flux {
             auto k = ggml_concat(ctx, txt_k, img_k, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
             auto v = ggml_concat(ctx, txt_v, img_v, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
 
-            auto attn         = attention(ctx, backend, q, k, v, pe, mask, flash_attn);  // [N, n_txt_token + n_img_token, n_head*d_head]
-            attn              = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3));     // [n_txt_token + n_img_token, N, hidden_size]
+            auto attn         = Rope::attention(ctx, backend, q, k, v, pe, mask, flash_attn);  // [N, n_txt_token + n_img_token, n_head*d_head]
+            attn              = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3));           // [n_txt_token + n_img_token, N, hidden_size]
             auto txt_attn_out = ggml_view_3d(ctx,
                                              attn,
                                              attn->ne[0],
@@ -503,7 +453,7 @@ namespace Flux {
             auto v           = ggml_reshape_4d(ctx, qkv_vec[2], head_dim, num_heads, qkv_vec[2]->ne[1], qkv_vec[2]->ne[2]);  // [N, n_token, n_head, d_head]
             q                = norm->query_norm(ctx, q);
             k                = norm->key_norm(ctx, k);
-            auto attn        = attention(ctx, backend, q, k, v, pe, mask, flash_attn);  // [N, n_token, hidden_size]
+            auto attn        = Rope::attention(ctx, backend, q, k, v, pe, mask, flash_attn);  // [N, n_token, hidden_size]
 
             auto attn_mlp = ggml_concat(ctx, attn, ggml_gelu_inplace(ctx, mlp), 0);  // [N, n_token, hidden_size + mlp_hidden_dim]
             auto output   = linear2->forward(ctx, attn_mlp);                         // [N, n_token, hidden_size]
@@ -615,6 +565,7 @@ namespace Flux {
         bool guidance_embed         = true;
         bool flash_attn             = true;
         bool is_chroma              = false;
+        SDVersion version           = VERSION_FLUX;
     };
 
     struct Flux : public GGMLBlock {
@@ -849,7 +800,8 @@ namespace Flux {
             auto img            = process_img(ctx, x);
             uint64_t img_tokens = img->ne[1];
 
-            if (c_concat != NULL) {
+            if (params.version == VERSION_FLUX_FILL) {
+                GGML_ASSERT(c_concat != NULL);
                 ggml_tensor* masked = ggml_view_4d(ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
                 ggml_tensor* mask   = ggml_view_4d(ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
 
@@ -857,6 +809,27 @@ namespace Flux {
                 mask   = process_img(ctx, mask);
 
                 img = ggml_concat(ctx, img, ggml_concat(ctx, masked, mask, 0), 0);
+            } else if (params.version == VERSION_FLEX_2) {
+                GGML_ASSERT(c_concat != NULL);
+                ggml_tensor* masked  = ggml_view_4d(ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
+                ggml_tensor* mask    = ggml_view_4d(ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
+                ggml_tensor* control = ggml_view_4d(ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * (C + 1));
+
+                masked  = ggml_pad(ctx, masked, pad_w, pad_h, 0, 0);
+                mask    = ggml_pad(ctx, mask, pad_w, pad_h, 0, 0);
+                control = ggml_pad(ctx, control, pad_w, pad_h, 0, 0);
+
+                masked  = patchify(ctx, masked, patch_size);
+                mask    = patchify(ctx, mask, patch_size);
+                control = patchify(ctx, control, patch_size);
+
+                img = ggml_concat(ctx, img, ggml_concat(ctx, ggml_concat(ctx, masked, mask, 0), control, 0), 0);
+            } else if (params.version == VERSION_FLUX_CONTROLS) {
+                GGML_ASSERT(c_concat != NULL);
+
+                ggml_tensor* control = ggml_pad(ctx, c_concat, pad_w, pad_h, 0, 0);
+                control              = patchify(ctx, control, patch_size);
+                img                  = ggml_concat(ctx, img, control, 0);
             }
 
             if (ref_latents.size() > 0) {
@@ -867,6 +840,7 @@ namespace Flux {
             }
 
             auto out = forward_orig(ctx, backend, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers);  // [N, num_tokens, C * patch_size * patch_size]
+
             if (out->ne[1] > img_tokens) {
                 out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3));  // [num_tokens, N, C * patch_size * patch_size]
                 out = ggml_view_3d(ctx, out, out->ne[0], out->ne[1], img_tokens, out->nb[1], out->nb[2], 0);
@@ -896,13 +870,18 @@ namespace Flux {
                    SDVersion version                   = VERSION_FLUX,
                    bool flash_attn                     = false,
                    bool use_mask                       = false)
-            : GGMLRunner(backend, offload_params_to_cpu), use_mask(use_mask) {
+            : GGMLRunner(backend, offload_params_to_cpu), version(version), use_mask(use_mask) {
+            flux_params.version             = version;
             flux_params.flash_attn          = flash_attn;
             flux_params.guidance_embed      = false;
             flux_params.depth               = 0;
             flux_params.depth_single_blocks = 0;
             if (version == VERSION_FLUX_FILL) {
                 flux_params.in_channels = 384;
+            } else if (version == VERSION_FLUX_CONTROLS) {
+                flux_params.in_channels = 128;
+            } else if (version == VERSION_FLEX_2) {
+                flux_params.in_channels = 196;
             }
             for (auto pair : tensor_types) {
                 std::string tensor_name = pair.first;

ggml_extend.hpp

@@ -56,6 +56,10 @@
 #define __STATIC_INLINE__ static inline
 #endif
 
+#ifndef SD_UNUSED
+#define SD_UNUSED(x) (void)(x)
+#endif
+
 __STATIC_INLINE__ void ggml_log_callback_default(ggml_log_level level, const char* text, void*) {
     switch (level) {
         case GGML_LOG_LEVEL_DEBUG:
@@ -185,17 +189,20 @@ __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 struct ggml_tensor* get_tensor_from_graph(struct ggml_cgraph* gf, const char* name) {
-    struct ggml_tensor* res = NULL;
-    for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-        struct ggml_tensor* node = ggml_graph_node(gf, i);
-        // printf("%d, %s \n", i, ggml_get_name(node));
-        if (strcmp(ggml_get_name(node), name) == 0) {
-            res = node;
-            break;
-        }
+__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int iw, int ih, int ic, bool scale = true) {
+    float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic);
+    if (scale) {
+        value /= 255.f;
     }
-    return res;
+    return value;
+}
+
+__STATIC_INLINE__ float sd_image_get_f32(sd_image_f32_t image, int iw, int ih, int ic, bool scale = true) {
+    float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic);
+    if (scale) {
+        value /= 255.f;
+    }
+    return value;
 }
 
 __STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false, const char* mark = "") {
@@ -235,6 +242,52 @@ __STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_
     }
 }
 
+__STATIC_INLINE__ void ggml_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i3 = 0; i3 < n3; i3++) {
+        for (int64_t i2 = 0; i2 < n2; i2++) {
+            for (int64_t i1 = 0; i1 < n1; i1++) {
+                for (int64_t i0 = 0; i0 < n0; i0++) {
+                    fn(tensor, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+}
+
+__STATIC_INLINE__ void ggml_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i = 0; i < ggml_nelements(tensor); i++) {
+        fn(tensor, i);
+    }
+}
+
+__STATIC_INLINE__ void ggml_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap = 0.1f) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    ggml_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float a_value = ggml_tensor_get_f32(a, i0, i1, i2, i3);
+        float b_value = ggml_tensor_get_f32(b, i0, i1, i2, i3);
+        if (abs(a_value - b_value) > gap) {
+            LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value);
+        }
+    });
+}
+
 __STATIC_INLINE__ ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) {
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
@@ -322,12 +375,14 @@ __STATIC_INLINE__ float sigmoid(float x) {
 
 // SPECIAL OPERATIONS WITH TENSORS
 
-__STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input) {
+__STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input, uint8_t* image_data = nullptr) {
     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);
+    if (image_data == nullptr) {
+        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++) {
@@ -366,104 +421,58 @@ __STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input, int idx
     return image_data;
 }
 
-__STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
-                                          struct ggml_tensor* output,
+__STATIC_INLINE__ void sd_image_to_tensor(sd_image_t image,
+                                          ggml_tensor* tensor,
                                           bool scale = true) {
-    int64_t width    = output->ne[0];
-    int64_t height   = output->ne[1];
-    int64_t channels = output->ne[2];
-    GGML_ASSERT(channels == 3 && output->type == GGML_TYPE_F32);
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
-            for (int k = 0; k < channels; k++) {
-                float value = *(image_data + iy * width * channels + ix * channels + k);
-                if (scale) {
-                    value /= 255.f;
-                }
-                ggml_tensor_set_f32(output, value, ix, iy, k);
-            }
-        }
-    }
-}
-
-__STATIC_INLINE__ void sd_mask_to_tensor(const uint8_t* image_data,
-                                         struct ggml_tensor* output,
-                                         bool scale = true) {
-    int64_t width    = output->ne[0];
-    int64_t height   = output->ne[1];
-    int64_t channels = output->ne[2];
-    GGML_ASSERT(channels == 1 && output->type == GGML_TYPE_F32);
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
-            float value = *(image_data + iy * width * channels + ix);
-            if (scale) {
-                value /= 255.f;
-            }
-            ggml_tensor_set_f32(output, value, ix, iy);
-        }
-    }
+    GGML_ASSERT(image.width == tensor->ne[0]);
+    GGML_ASSERT(image.height == tensor->ne[1]);
+    GGML_ASSERT(image.channel == tensor->ne[2]);
+    GGML_ASSERT(1 == tensor->ne[3]);
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+    ggml_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float value = sd_image_get_f32(image, i0, i1, i2, scale);
+        ggml_tensor_set_f32(tensor, value, i0, i1, i2, i3);
+    });
 }
 
 __STATIC_INLINE__ void sd_apply_mask(struct ggml_tensor* image_data,
                                      struct ggml_tensor* mask,
-                                     struct ggml_tensor* output) {
+                                     struct ggml_tensor* output,
+                                     float masked_value = 0.5f) {
     int64_t width    = output->ne[0];
     int64_t height   = output->ne[1];
     int64_t channels = output->ne[2];
+    float rescale_mx = mask->ne[0] / output->ne[0];
+    float rescale_my = mask->ne[1] / output->ne[1];
     GGML_ASSERT(output->type == GGML_TYPE_F32);
     for (int ix = 0; ix < width; ix++) {
         for (int iy = 0; iy < height; iy++) {
-            float m = ggml_tensor_get_f32(mask, ix, iy);
+            int mx  = (int)(ix * rescale_mx);
+            int my  = (int)(iy * rescale_my);
+            float m = ggml_tensor_get_f32(mask, mx, my);
             m       = round(m);  // inpaint models need binary masks
-            ggml_tensor_set_f32(mask, m, ix, iy);
+            ggml_tensor_set_f32(mask, m, mx, my);
             for (int k = 0; k < channels; k++) {
-                float value = (1 - m) * (ggml_tensor_get_f32(image_data, ix, iy, k) - .5) + .5;
+                float value = ggml_tensor_get_f32(image_data, ix, iy, k);
+                value       = (1 - m) * (value - masked_value) + masked_value;
                 ggml_tensor_set_f32(output, value, ix, iy, k);
             }
         }
     }
 }
 
-__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,
+__STATIC_INLINE__ void sd_image_f32_to_tensor(sd_image_f32_t image,
+                                              ggml_tensor* tensor,
                                               bool scale = true) {
-    int64_t width    = output->ne[0];
-    int64_t height   = output->ne[1];
-    int64_t channels = output->ne[2];
-    GGML_ASSERT(channels == 3 && output->type == GGML_TYPE_F32);
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
-            for (int k = 0; k < channels; k++) {
-                int value = *(image_data + iy * width * channels + ix * channels + k);
-                if (scale) {
-                    value /= 255.f;
-                }
-                ggml_tensor_set_f32(output, value, ix, iy, k);
-            }
-        }
-    }
+    GGML_ASSERT(image.width == tensor->ne[0]);
+    GGML_ASSERT(image.height == tensor->ne[1]);
+    GGML_ASSERT(image.channel == tensor->ne[2]);
+    GGML_ASSERT(1 == tensor->ne[3]);
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+    ggml_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float value = sd_image_get_f32(image, i0, i1, i2, scale);
+        ggml_tensor_set_f32(tensor, value, i0, i1, i2, i3);
+    });
 }
 
 __STATIC_INLINE__ void ggml_split_tensor_2d(struct ggml_tensor* input,
@@ -473,12 +482,15 @@ __STATIC_INLINE__ void ggml_split_tensor_2d(struct ggml_tensor* input,
     int64_t width    = output->ne[0];
     int64_t height   = output->ne[1];
     int64_t channels = output->ne[2];
+    int64_t ne3      = output->ne[3];
     GGML_ASSERT(input->type == GGML_TYPE_F32 && output->type == GGML_TYPE_F32);
     for (int iy = 0; iy < height; iy++) {
         for (int ix = 0; ix < width; ix++) {
             for (int k = 0; k < channels; k++) {
-                float value = ggml_tensor_get_f32(input, ix + x, iy + y, k);
-                ggml_tensor_set_f32(output, value, ix, iy, k);
+                for (int l = 0; l < ne3; l++) {
+                    float value = ggml_tensor_get_f32(input, ix + x, iy + y, k, l);
+                    ggml_tensor_set_f32(output, value, ix, iy, k, l);
+                }
             }
         }
     }
@@ -494,36 +506,42 @@ __STATIC_INLINE__ void ggml_merge_tensor_2d(struct ggml_tensor* input,
                                             struct ggml_tensor* output,
                                             int x,
                                             int y,
-                                            int overlap) {
+                                            int overlap_x,
+                                            int overlap_y,
+                                            int x_skip = 0,
+                                            int y_skip = 0) {
     int64_t width    = input->ne[0];
     int64_t height   = input->ne[1];
     int64_t channels = input->ne[2];
+    int64_t ne3      = input->ne[3];
 
     int64_t img_width  = output->ne[0];
     int64_t img_height = output->ne[1];
 
     GGML_ASSERT(input->type == GGML_TYPE_F32 && output->type == GGML_TYPE_F32);
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
+    for (int iy = y_skip; iy < height; iy++) {
+        for (int ix = x_skip; ix < width; ix++) {
             for (int k = 0; k < channels; k++) {
-                float new_value = ggml_tensor_get_f32(input, ix, iy, k);
-                if (overlap > 0) {  // blend colors in overlapped area
-                    float old_value = ggml_tensor_get_f32(output, x + ix, y + iy, k);
+                for (int l = 0; l < ne3; l++) {
+                    float new_value = ggml_tensor_get_f32(input, ix, iy, k, l);
+                    if (overlap_x > 0 || overlap_y > 0) {  // blend colors in overlapped area
+                        float old_value = ggml_tensor_get_f32(output, x + ix, y + iy, k, l);
 
-                    const float x_f_0 = (x > 0) ? ix / float(overlap) : 1;
-                    const float x_f_1 = (x < (img_width - width)) ? (width - ix) / float(overlap) : 1;
-                    const float y_f_0 = (y > 0) ? iy / float(overlap) : 1;
-                    const float y_f_1 = (y < (img_height - height)) ? (height - iy) / float(overlap) : 1;
+                        const float x_f_0 = (overlap_x > 0 && x > 0) ? (ix - x_skip) / float(overlap_x) : 1;
+                        const float x_f_1 = (overlap_x > 0 && x < (img_width - width)) ? (width - ix) / float(overlap_x) : 1;
+                        const float y_f_0 = (overlap_y > 0 && y > 0) ? (iy - y_skip) / float(overlap_y) : 1;
+                        const float y_f_1 = (overlap_y > 0 && y < (img_height - height)) ? (height - iy) / float(overlap_y) : 1;
 
-                    const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f);
-                    const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f);
+                        const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f);
+                        const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f);
 
-                    ggml_tensor_set_f32(
-                        output,
-                        old_value + new_value * ggml_smootherstep_f32(y_f) * ggml_smootherstep_f32(x_f),
-                        x + ix, y + iy, k);
-                } else {
-                    ggml_tensor_set_f32(output, new_value, x + ix, y + iy, k);
+                        ggml_tensor_set_f32(
+                            output,
+                            old_value + new_value * ggml_smootherstep_f32(y_f) * ggml_smootherstep_f32(x_f),
+                            x + ix, y + iy, k, l);
+                    } else {
+                        ggml_tensor_set_f32(output, new_value, x + ix, y + iy, k, l);
+                    }
                 }
             }
         }
@@ -745,22 +763,102 @@ __STATIC_INLINE__ std::vector<struct ggml_tensor*> ggml_chunk(struct ggml_contex
 
 typedef std::function<void(ggml_tensor*, ggml_tensor*, bool)> on_tile_process;
 
+__STATIC_INLINE__ void sd_tiling_calc_tiles(int& num_tiles_dim,
+                                            float& tile_overlap_factor_dim,
+                                            int small_dim,
+                                            int tile_size,
+                                            const float tile_overlap_factor) {
+    int tile_overlap     = (tile_size * tile_overlap_factor);
+    int non_tile_overlap = tile_size - tile_overlap;
+
+    num_tiles_dim     = (small_dim - tile_overlap) / non_tile_overlap;
+    int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim;
+
+    if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) {
+        // if tiles don't fit perfectly using the desired overlap
+        // and there is enough room to squeeze an extra tile without overlap becoming >0.5
+        num_tiles_dim++;
+    }
+
+    tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1));
+    if (num_tiles_dim <= 2) {
+        if (small_dim <= tile_size) {
+            num_tiles_dim           = 1;
+            tile_overlap_factor_dim = 0;
+        } else {
+            num_tiles_dim           = 2;
+            tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size;
+        }
+    }
+}
+
 // Tiling
-__STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const int scale, const int tile_size, const float tile_overlap_factor, on_tile_process on_processing) {
+__STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input,
+                                            ggml_tensor* output,
+                                            const int scale,
+                                            const int p_tile_size_x,
+                                            const int p_tile_size_y,
+                                            const float tile_overlap_factor,
+                                            on_tile_process on_processing) {
     output = ggml_set_f32(output, 0);
 
     int input_width   = (int)input->ne[0];
     int input_height  = (int)input->ne[1];
     int output_width  = (int)output->ne[0];
     int output_height = (int)output->ne[1];
+
+    GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) &&
+                ((output_width / input_width) == (output_height / input_height)));
+    GGML_ASSERT(((input_width / output_width) == scale) ||
+                ((output_width / input_width) == scale));
+
+    int small_width  = output_width;
+    int small_height = output_height;
+
+    bool decode = output_width > input_width;
+    if (decode) {
+        small_width  = input_width;
+        small_height = input_height;
+    }
+
+    int num_tiles_x;
+    float tile_overlap_factor_x;
+    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor);
+
+    int num_tiles_y;
+    float tile_overlap_factor_y;
+    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor);
+
+    LOG_DEBUG("num tiles : %d, %d ", num_tiles_x, num_tiles_y);
+    LOG_DEBUG("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor);
+
     GGML_ASSERT(input_width % 2 == 0 && input_height % 2 == 0 && output_width % 2 == 0 && output_height % 2 == 0);  // should be multiple of 2
 
-    int tile_overlap     = (int32_t)(tile_size * tile_overlap_factor);
-    int non_tile_overlap = tile_size - tile_overlap;
+    int tile_overlap_x     = (int32_t)(p_tile_size_x * tile_overlap_factor_x);
+    int non_tile_overlap_x = p_tile_size_x - tile_overlap_x;
+
+    int tile_overlap_y     = (int32_t)(p_tile_size_y * tile_overlap_factor_y);
+    int non_tile_overlap_y = p_tile_size_y - tile_overlap_y;
+
+    int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width;
+    int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height;
+
+    int input_tile_size_x  = tile_size_x;
+    int input_tile_size_y  = tile_size_y;
+    int output_tile_size_x = tile_size_x;
+    int output_tile_size_y = tile_size_y;
+
+    if (decode) {
+        output_tile_size_x *= scale;
+        output_tile_size_y *= scale;
+    } else {
+        input_tile_size_x *= scale;
+        input_tile_size_y *= scale;
+    }
 
     struct ggml_init_params params = {};
-    params.mem_size += tile_size * tile_size * input->ne[2] * sizeof(float);                       // input chunk
-    params.mem_size += (tile_size * scale) * (tile_size * scale) * output->ne[2] * sizeof(float);  // output chunk
+    params.mem_size += input_tile_size_x * input_tile_size_y * input->ne[2] * input->ne[3] * sizeof(float);      // input chunk
+    params.mem_size += output_tile_size_x * output_tile_size_y * output->ne[2] * output->ne[3] * sizeof(float);  // output chunk
     params.mem_size += 3 * ggml_tensor_overhead();
     params.mem_buffer = NULL;
     params.no_alloc   = false;
@@ -775,29 +873,50 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const
     }
 
     // tiling
-    ggml_tensor* input_tile  = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size, tile_size, input->ne[2], 1);
-    ggml_tensor* output_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size * scale, tile_size * scale, output->ne[2], 1);
-    on_processing(input_tile, NULL, true);
-    int num_tiles = ceil((float)input_width / non_tile_overlap) * ceil((float)input_height / non_tile_overlap);
+    ggml_tensor* input_tile  = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, input_tile_size_x, input_tile_size_y, input->ne[2], input->ne[3]);
+    ggml_tensor* output_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, output_tile_size_x, output_tile_size_y, output->ne[2], output->ne[3]);
+    int num_tiles            = num_tiles_x * num_tiles_y;
     LOG_INFO("processing %i tiles", num_tiles);
-    pretty_progress(1, num_tiles, 0.0f);
+    pretty_progress(0, num_tiles, 0.0f);
     int tile_count = 1;
     bool last_y = false, last_x = false;
     float last_time = 0.0f;
-    for (int y = 0; y < input_height && !last_y; y += non_tile_overlap) {
-        if (y + tile_size >= input_height) {
-            y      = input_height - tile_size;
+    for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) {
+        int dy = 0;
+        if (y + tile_size_y >= small_height) {
+            int _y = y;
+            y      = small_height - tile_size_y;
+            dy     = _y - y;
+            if (decode) {
+                dy *= scale;
+            }
             last_y = true;
         }
-        for (int x = 0; x < input_width && !last_x; x += non_tile_overlap) {
-            if (x + tile_size >= input_width) {
-                x      = input_width - tile_size;
+        for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) {
+            int dx = 0;
+            if (x + tile_size_x >= small_width) {
+                int _x = x;
+                x      = small_width - tile_size_x;
+                dx     = _x - x;
+                if (decode) {
+                    dx *= scale;
+                }
                 last_x = true;
             }
+
+            int x_in  = decode ? x : scale * x;
+            int y_in  = decode ? y : scale * y;
+            int x_out = decode ? x * scale : x;
+            int y_out = decode ? y * scale : y;
+
+            int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x;
+            int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y;
+
             int64_t t1 = ggml_time_ms();
-            ggml_split_tensor_2d(input, input_tile, x, y);
+            ggml_split_tensor_2d(input, input_tile, x_in, y_in);
             on_processing(input_tile, output_tile, false);
-            ggml_merge_tensor_2d(output_tile, output, x * scale, y * scale, tile_overlap * scale);
+            ggml_merge_tensor_2d(output_tile, output, x_out, y_out, overlap_x_out, overlap_y_out, dx, dy);
+
             int64_t t2 = ggml_time_ms();
             last_time  = (t2 - t1) / 1000.0f;
             pretty_progress(tile_count, num_tiles, last_time);
@@ -811,6 +930,15 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const
     ggml_free(tiles_ctx);
 }
 
+__STATIC_INLINE__ void sd_tiling(ggml_tensor* input,
+                                 ggml_tensor* output,
+                                 const int scale,
+                                 const int tile_size,
+                                 const float tile_overlap_factor,
+                                 on_tile_process on_processing) {
+    sd_tiling_non_square(input, output, scale, tile_size, tile_size, tile_overlap_factor, on_processing);
+}
+
 __STATIC_INLINE__ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx,
                                                          struct ggml_tensor* a) {
     const float eps = 1e-6f;  // default eps parameter
@@ -820,8 +948,19 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ct
 __STATIC_INLINE__ struct ggml_tensor* ggml_nn_linear(struct ggml_context* ctx,
                                                      struct ggml_tensor* x,
                                                      struct ggml_tensor* w,
-                                                     struct ggml_tensor* b) {
+                                                     struct ggml_tensor* b,
+                                                     bool force_prec_f32 = false,
+                                                     float scale         = 1.f) {
+    if (scale != 1.f) {
+        x = ggml_scale(ctx, x, scale);
+    }
     x = ggml_mul_mat(ctx, w, x);
+    if (force_prec_f32) {
+        ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+    }
+    if (scale != 1.f) {
+        x = ggml_scale(ctx, x, 1.f / scale);
+    }
     if (b != NULL) {
         x = ggml_add_inplace(ctx, x, b);
     }
@@ -1006,9 +1145,9 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention(struct ggml_context* ctx
     return kqv;
 }
 
-// q: [N, L_q, C] or [N*n_head, L_q, d_head]
-// k: [N, L_k, C] or [N*n_head, L_k, d_head]
-// v: [N, L_k, C] or [N, L_k, n_head, d_head]
+// q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head]
+// k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head]
+// v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head]
 // mask: [N, L_q, L_k]
 // return: [N, L_q, C]
 __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context* ctx,
@@ -1020,33 +1159,38 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
                                                             struct ggml_tensor* mask = NULL,
                                                             bool diag_mask_inf       = false,
                                                             bool skip_reshape        = false,
-                                                            bool flash_attn          = false) {
+                                                            bool flash_attn          = false,  // avoid overflow
+                                                            float kv_scale           = 1.0f) {
     int64_t L_q;
     int64_t L_k;
     int64_t C;
     int64_t N;
     int64_t d_head;
+    int64_t n_kv_head;
     if (!skip_reshape) {
-        L_q    = q->ne[1];
-        L_k    = k->ne[1];
-        C      = q->ne[0];
-        N      = q->ne[2];
-        d_head = C / n_head;
-        q      = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N);      // [N, L_q, n_head, d_head]
-        q      = ggml_nn_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));  // [N, n_head, L_q, d_head]
-        q      = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N);     // [N * n_head, L_q, d_head]
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        C         = q->ne[0];
+        N         = q->ne[2];
+        d_head    = C / n_head;
+        n_kv_head = k->ne[0] / d_head;
 
-        k = ggml_reshape_4d(ctx, k, d_head, n_head, L_k, N);      // [N, L_k, n_head, d_head]
-        k = ggml_nn_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));  // [N, n_head, L_k, d_head]
-        k = ggml_reshape_3d(ctx, k, d_head, L_k, n_head * N);     // [N * n_head, L_k, d_head]
+        q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N);      // [N, L_q, n_head, d_head]
+        q = ggml_nn_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));  // [N, n_head, L_q, d_head]
+        q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N);     // [N * n_head, L_q, d_head]
 
-        v = ggml_reshape_4d(ctx, v, d_head, n_head, L_k, N);  // [N, L_k, n_head, d_head]
+        k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N);   // [N, L_k, n_kv_head, d_head]
+        k = ggml_nn_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));  // [N, n_kv_head, L_k, d_head]
+        k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N);  // [N * n_kv_head, L_k, d_head]
+
+        v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N);  // [N, L_k, n_kv_head, d_head]
     } else {
-        L_q    = q->ne[1];
-        L_k    = k->ne[1];
-        d_head = v->ne[0];
-        N      = v->ne[3];
-        C      = d_head * n_head;
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        d_head    = v->ne[0];
+        N         = v->ne[3];
+        n_kv_head = k->ne[2] / N;
+        C         = d_head * n_head;
     }
 
     float scale = (1.0f / sqrt((float)d_head));
@@ -1058,13 +1202,19 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
         if (kv_pad != 0) {
             k_in = ggml_pad(ctx, k_in, 0, kv_pad, 0, 0);
         }
+        if (kv_scale != 1.0f) {
+            k_in = ggml_scale(ctx, k_in, kv_scale);
+        }
         k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16);
 
         v_in = ggml_nn_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
-        v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_head * N);
+        v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
         if (kv_pad != 0) {
             v_in = ggml_pad(ctx, v_in, 0, kv_pad, 0, 0);
         }
+        if (kv_scale != 1.0f) {
+            v_in = ggml_scale(ctx, v_in, kv_scale);
+        }
         v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
 
         if (mask_in != nullptr) {
@@ -1088,8 +1238,11 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
             mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
         }
 
-        auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale, 0, 0);
+        auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
         ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
+        if (kv_scale != 1.0f) {
+            out = ggml_scale(ctx, out, 1.0f / kv_scale);
+        }
         return out;
     };
 
@@ -1119,8 +1272,8 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
         // if (flash_attn) {
         //     LOG_DEBUG("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
         // }
-        v = ggml_nn_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));  // [N, n_head, d_head, L_k]
-        v = ggml_reshape_3d(ctx, v, L_k, d_head, n_head * N);     // [N * n_head, d_head, L_k]
+        v = ggml_nn_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));  // [N, n_kv_head, d_head, L_k]
+        v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N);  // [N * n_kv_head, d_head, L_k]
 
         auto kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, L_q, L_k]
         kq      = ggml_scale_inplace(ctx, kq, scale);
@@ -1236,15 +1389,13 @@ __STATIC_INLINE__ std::vector<float> arange(float start, float end, float step =
 // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
 __STATIC_INLINE__ std::vector<float> timestep_embedding(std::vector<float> timesteps,
                                                         int dim,
-                                                        int max_period = 10000) {
+                                                        int max_period       = 10000,
+                                                        bool flip_sin_to_cos = true,
+                                                        float scale          = 1.f) {
     // timesteps: [N,]
     // embedding: [N, dim]
-    size_t N        = timesteps.size();
-    int acutual_dim = dim;
-    if (dim % 2 != 0) {
-        acutual_dim = dim + 1;
-    }
-    std::vector<float> embedding(N * acutual_dim, 0.f);
+    size_t N = timesteps.size();
+    std::vector<float> embedding(N * dim, 0.f);
     int half = dim / 2;
     std::vector<float> freqs(half);
     for (int i = 0; i < half; ++i) {
@@ -1252,9 +1403,14 @@ __STATIC_INLINE__ std::vector<float> timestep_embedding(std::vector<float> times
     }
     for (int i = 0; i < N; ++i) {
         for (int j = 0; j < half; ++j) {
-            float arg                             = timesteps[i] * freqs[j];
-            embedding[i * acutual_dim + j]        = std::cos(arg);
-            embedding[i * acutual_dim + j + half] = std::sin(arg);
+            float arg = timesteps[i] * freqs[j] * scale;
+            if (flip_sin_to_cos) {
+                embedding[i * dim + j]        = std::cos(arg);
+                embedding[i * dim + j + half] = std::sin(arg);
+            } else {
+                embedding[i * dim + j]        = std::sin(arg);
+                embedding[i * dim + j + half] = std::cos(arg);
+            }
         }
     }
     return embedding;
@@ -1275,11 +1431,7 @@ __STATIC_INLINE__ struct ggml_tensor* new_timestep_embedding(struct ggml_context
     // timesteps: [N,]
     // embedding: [N, dim]
     std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
-    int acutual_dim                  = dim;
-    if (dim % 2 != 0) {
-        acutual_dim = dim + 1;
-    }
-    struct ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, acutual_dim, timesteps.size());
+    struct ggml_tensor* embedding    = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps.size());
     if (embedding->data != NULL) {
         memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
     } else {
@@ -1523,6 +1675,7 @@ protected:
             ggml_backend_tensor_copy(t, offload_t);
             std::swap(t->buffer, offload_t->buffer);
             std::swap(t->data, offload_t->data);
+            std::swap(t->extra, offload_t->extra);
 
             t         = ggml_get_next_tensor(params_ctx, t);
             offload_t = ggml_get_next_tensor(offload_ctx, offload_t);
@@ -1553,8 +1706,10 @@ protected:
         while (t != NULL && offload_t != NULL) {
             t->buffer         = offload_t->buffer;
             t->data           = offload_t->data;
+            t->extra          = offload_t->extra;
             offload_t->buffer = NULL;
             offload_t->data   = NULL;
+            offload_t->extra  = NULL;
 
             t         = ggml_get_next_tensor(params_ctx, t);
             offload_t = ggml_get_next_tensor(offload_ctx, offload_t);
@@ -1818,6 +1973,8 @@ protected:
     int64_t out_features;
     bool bias;
     bool force_f32;
+    bool force_prec_f32;
+    float scale;
 
     void init_params(struct ggml_context* ctx, const String2GGMLType& tensor_types = {}, const std::string prefix = "") {
         enum ggml_type wtype = get_type(prefix + "weight", tensor_types, GGML_TYPE_F32);
@@ -1834,12 +1991,16 @@ protected:
 public:
     Linear(int64_t in_features,
            int64_t out_features,
-           bool bias      = true,
-           bool force_f32 = false)
+           bool bias           = true,
+           bool force_f32      = false,
+           bool force_prec_f32 = false,
+           float scale         = 1.f)
         : in_features(in_features),
           out_features(out_features),
           bias(bias),
-          force_f32(force_f32) {}
+          force_f32(force_f32),
+          force_prec_f32(force_prec_f32),
+          scale(scale) {}
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         struct ggml_tensor* w = params["weight"];
@@ -1847,17 +2008,28 @@ public:
         if (bias) {
             b = params["bias"];
         }
-        return ggml_nn_linear(ctx, x, w, b);
+        return ggml_nn_linear(ctx, x, w, b, force_prec_f32, scale);
     }
 };
 
+__STATIC_INLINE__ bool support_get_rows(ggml_type wtype) {
+    std::set<ggml_type> allow_types = {GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0};
+    if (allow_types.find(wtype) != allow_types.end()) {
+        return true;
+    }
+    return false;
+}
+
 class Embedding : public UnaryBlock {
 protected:
     int64_t embedding_dim;
     int64_t num_embeddings;
     void init_params(struct ggml_context* ctx, const String2GGMLType& tensor_types, const std::string prefix = "") {
         enum ggml_type wtype = get_type(prefix + "weight", tensor_types, GGML_TYPE_F32);
-        params["weight"]     = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
+        if (!support_get_rows(wtype)) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
     }
 
 public:

lora.hpp

@@ -1,6 +1,7 @@
 #ifndef __LORA_HPP__
 #define __LORA_HPP__
 
+#include <mutex>
 #include "ggml_extend.hpp"
 
 #define LORA_GRAPH_BASE_SIZE 10240
@@ -115,7 +116,7 @@ struct LoraModel : public GGMLRunner {
         return "lora";
     }
 
-    bool load_from_file(bool filter_tensor = false) {
+    bool load_from_file(bool filter_tensor, int n_threads) {
         LOG_INFO("loading LoRA from '%s'", file_path.c_str());
 
         if (load_failed) {
@@ -123,41 +124,53 @@ struct LoraModel : public GGMLRunner {
             return false;
         }
 
+        std::unordered_map<std::string, TensorStorage> tensors_to_create;
+        std::mutex lora_mutex;
         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;
+            if (dry_run) {
+                const std::string& name = tensor_storage.name;
 
-            if (filter_tensor && !contains(name, "lora")) {
-                // LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
-                return true;
-            }
-            // LOG_INFO("lora_tensor %s", name.c_str());
-            for (int i = 0; i < LORA_TYPE_COUNT; i++) {
-                if (name.find(type_fingerprints[i]) != std::string::npos) {
-                    type = (lora_t)i;
-                    break;
+                if (filter_tensor && !contains(name, "lora")) {
+                    return true;
+                }
+
+                {
+                    std::lock_guard<std::mutex> lock(lora_mutex);
+                    for (int i = 0; i < LORA_TYPE_COUNT; i++) {
+                        if (name.find(type_fingerprints[i]) != std::string::npos) {
+                            type = (lora_t)i;
+                            break;
+                        }
+                    }
+                    tensors_to_create[name] = tensor_storage;
+                }
+            } else {
+                const std::string& name = tensor_storage.name;
+                auto iter               = lora_tensors.find(name);
+                if (iter != lora_tensors.end()) {
+                    *dst_tensor = iter->second;
                 }
             }
-
-            if (dry_run) {
-                struct ggml_tensor* real = ggml_new_tensor(params_ctx,
-                                                           tensor_storage.type,
-                                                           tensor_storage.n_dims,
-                                                           tensor_storage.ne);
-                lora_tensors[name]       = real;
-            } else {
-                auto real   = lora_tensors[name];
-                *dst_tensor = real;
-            }
-
             return true;
         };
 
-        model_loader.load_tensors(on_new_tensor_cb);
+        model_loader.load_tensors(on_new_tensor_cb, n_threads);
+
+        for (const auto& pair : tensors_to_create) {
+            const auto& name         = pair.first;
+            const auto& ts           = pair.second;
+            struct ggml_tensor* real = ggml_new_tensor(params_ctx,
+                                                       ts.type,
+                                                       ts.n_dims,
+                                                       ts.ne);
+            lora_tensors[name]       = real;
+        }
+
         alloc_params_buffer();
-        // exit(0);
+
         dry_run = false;
-        model_loader.load_tensors(on_new_tensor_cb);
+        model_loader.load_tensors(on_new_tensor_cb, n_threads);
 
         LOG_DEBUG("lora type: \"%s\"/\"%s\"", lora_downs[type].c_str(), lora_ups[type].c_str());
 

model.cpp

@@ -1,8 +1,14 @@
 #include <stdarg.h>
+#include <algorithm>
+#include <atomic>
+#include <chrono>
 #include <fstream>
+#include <functional>
+#include <mutex>
 #include <regex>
 #include <set>
 #include <string>
+#include <thread>
 #include <unordered_map>
 #include <vector>
 
@@ -11,6 +17,7 @@
 #include "stable-diffusion.h"
 #include "util.h"
 #include "vocab.hpp"
+#include "vocab_qwen.hpp"
 #include "vocab_umt5.hpp"
 
 #include "ggml-alloc.h"
@@ -104,6 +111,8 @@ const char* unused_tensors[] = {
     "embedding_manager",
     "denoiser.sigmas",
     "text_encoders.t5xxl.transformer.encoder.embed_tokens.weight",  // only used during training
+    "text_encoders.qwen2vl.output.weight",
+    "text_encoders.qwen2vl.lm_head.",
 };
 
 bool is_unused_tensor(std::string name) {
@@ -187,6 +196,39 @@ std::unordered_map<std::string, std::string> pmid_v2_name_map = {
      "pmid.qformer_perceiver.token_proj.fc2.weight"},
 };
 
+std::unordered_map<std::string, std::string> qwenvl_name_map{
+    {"token_embd.", "model.embed_tokens."},
+    {"blk.", "model.layers."},
+    {"attn_q.", "self_attn.q_proj."},
+    {"attn_k.", "self_attn.k_proj."},
+    {"attn_v.", "self_attn.v_proj."},
+    {"attn_output.", "self_attn.o_proj."},
+    {"attn_norm.", "input_layernorm."},
+    {"ffn_down.", "mlp.down_proj."},
+    {"ffn_gate.", "mlp.gate_proj."},
+    {"ffn_up.", "mlp.up_proj."},
+    {"ffn_norm.", "post_attention_layernorm."},
+    {"output_norm.", "model.norm."},
+};
+
+std::unordered_map<std::string, std::string> qwenvl_vision_name_map{
+    {"mm.", "merger.mlp."},
+    {"v.post_ln.", "merger.ln_q."},
+    {"v.patch_embd.weight", "patch_embed.proj.0.weight"},
+    {"patch_embed.proj.0.weight.1", "patch_embed.proj.1.weight"},
+    {"v.patch_embd.weight.1", "patch_embed.proj.1.weight"},
+    {"v.blk.", "blocks."},
+    {"attn_q.", "attn.q_proj."},
+    {"attn_k.", "attn.k_proj."},
+    {"attn_v.", "attn.v_proj."},
+    {"attn_out.", "attn.proj."},
+    {"ffn_down.", "mlp.down_proj."},
+    {"ffn_gate.", "mlp.gate_proj."},
+    {"ffn_up.", "mlp.up_proj."},
+    {"ln1.", "norm1."},
+    {"ln2.", "norm2."},
+};
+
 std::string convert_cond_model_name(const std::string& name) {
     std::string new_name = name;
     std::string prefix;
@@ -244,6 +286,22 @@ std::string convert_cond_model_name(const std::string& name) {
         if (pos != std::string::npos) {
             new_name.replace(pos, 11, "layer.0.SelfAttention.relative_attention_bias.");
         }
+    } else if (contains(name, "qwen2vl")) {
+        if (contains(name, "qwen2vl.visual")) {
+            for (auto kv : qwenvl_vision_name_map) {
+                size_t pos = new_name.find(kv.first);
+                if (pos != std::string::npos) {
+                    new_name.replace(pos, kv.first.size(), kv.second);
+                }
+            }
+        } else {
+            for (auto kv : qwenvl_name_map) {
+                size_t pos = new_name.find(kv.first);
+                if (pos != std::string::npos) {
+                    new_name.replace(pos, kv.first.size(), kv.second);
+                }
+            }
+        }
     } else if (name == "text_encoders.t5xxl.transformer.token_embd.weight") {
         new_name = "text_encoders.t5xxl.transformer.shared.weight";
     }
@@ -574,7 +632,11 @@ std::string convert_tensor_name(std::string name) {
     //     name.replace(pos, strlen("lora_B"), "lora_down");
     // }
     std::string new_name = name;
-    if (starts_with(name, "cond_stage_model.") || starts_with(name, "conditioner.embedders.") || starts_with(name, "text_encoders.") || ends_with(name, ".vision_model.visual_projection.weight")) {
+    if (starts_with(name, "cond_stage_model.") ||
+        starts_with(name, "conditioner.embedders.") ||
+        starts_with(name, "text_encoders.") ||
+        ends_with(name, ".vision_model.visual_projection.weight") ||
+        starts_with(name, "qwen2vl")) {
         new_name = convert_cond_model_name(name);
     } else if (starts_with(name, "first_stage_model.decoder")) {
         new_name = convert_vae_decoder_name(name);
@@ -693,6 +755,7 @@ void preprocess_tensor(TensorStorage tensor_storage,
 
     // convert unet transformer linear to conv2d 1x1
     if (starts_with(new_name, "model.diffusion_model.") &&
+        !starts_with(new_name, "model.diffusion_model.proj_out.") &&
         (ends_with(new_name, "proj_in.weight") || ends_with(new_name, "proj_out.weight"))) {
         tensor_storage.unsqueeze();
     }
@@ -1726,6 +1789,9 @@ SDVersion ModelLoader::get_sd_version() {
             if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) {
                 return VERSION_SD3;
             }
+            if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
+                return VERSION_QWEN_IMAGE;
+            }
             if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
                 is_wan = true;
             }
@@ -1797,10 +1863,15 @@ SDVersion ModelLoader::get_sd_version() {
     }
 
     if (is_flux) {
-        is_inpaint = input_block_weight.ne[0] == 384;
-        if (is_inpaint) {
+        if (input_block_weight.ne[0] == 384) {
             return VERSION_FLUX_FILL;
         }
+        if (input_block_weight.ne[0] == 128) {
+            return VERSION_FLUX_CONTROLS;
+        }
+        if (input_block_weight.ne[0] == 196) {
+            return VERSION_FLEX_2;
+        }
         return VERSION_FLUX;
     }
 
@@ -1934,6 +2005,11 @@ std::string ModelLoader::load_merges() {
     return merges_utf8_str;
 }
 
+std::string ModelLoader::load_qwen2_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(qwen2_merges_utf8_c_str), sizeof(qwen2_merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
 std::string ModelLoader::load_t5_tokenizer_json() {
     std::string json_str(reinterpret_cast<const char*>(t5_tokenizer_json_str), sizeof(t5_tokenizer_json_str));
     return json_str;
@@ -1944,292 +2020,357 @@ std::string ModelLoader::load_umt5_tokenizer_json() {
     return json_str;
 }
 
-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;
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p) {
+    int64_t process_time_ms = 0;
+    std::atomic<int64_t> read_time_ms(0);
+    std::atomic<int64_t> memcpy_time_ms(0);
+    std::atomic<int64_t> copy_to_backend_time_ms(0);
+    std::atomic<int64_t> convert_time_ms(0);
 
-    for (size_t i = 0; i < vec.size(); ++i) {
-        const std::string& current_name = vec[i].name;
-        auto it                         = name_to_index_map.find(current_name);
+    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : get_num_physical_cores();
+    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
 
-        if (it != name_to_index_map.end()) {
-            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());
-
-    return res;
-}
-
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {
-    int64_t process_time_ms         = 0;
-    int64_t read_time_ms            = 0;
-    int64_t memcpy_time_ms          = 0;
-    int64_t copy_to_backend_time_ms = 0;
-    int64_t convert_time_ms         = 0;
-
-    int64_t prev_time_ms = 0;
-    int64_t curr_time_ms = 0;
-    int64_t start_time   = ggml_time_ms();
-    prev_time_ms         = start_time;
+    int64_t start_time = ggml_time_ms();
     std::vector<TensorStorage> processed_tensor_storages;
-    for (auto& tensor_storage : tensor_storages) {
-        // LOG_DEBUG("%s", name.c_str());
 
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
+    {
+        struct IndexedStorage {
+            size_t index;
+            TensorStorage ts;
+        };
+
+        std::mutex vec_mutex;
+        std::vector<IndexedStorage> all_results;
+
+        int n_threads = std::min(num_threads_to_use, (int)tensor_storages.size());
+        if (n_threads < 1) {
+            n_threads = 1;
+        }
+        std::vector<std::thread> workers;
+
+        for (int i = 0; i < n_threads; ++i) {
+            workers.emplace_back([&, thread_id = i]() {
+                std::vector<IndexedStorage> local_results;
+                std::vector<TensorStorage> temp_storages;
+
+                for (size_t j = thread_id; j < tensor_storages.size(); j += n_threads) {
+                    const auto& tensor_storage = tensor_storages[j];
+                    if (is_unused_tensor(tensor_storage.name)) {
+                        continue;
+                    }
+
+                    temp_storages.clear();
+                    preprocess_tensor(tensor_storage, temp_storages);
+
+                    for (const auto& ts : temp_storages) {
+                        local_results.push_back({j, ts});
+                    }
+                }
+
+                if (!local_results.empty()) {
+                    std::lock_guard<std::mutex> lock(vec_mutex);
+                    all_results.insert(all_results.end(),
+                                       local_results.begin(), local_results.end());
+                }
+            });
+        }
+        for (auto& w : workers) {
+            w.join();
         }
 
-        preprocess_tensor(tensor_storage, processed_tensor_storages);
-    }
-    std::vector<TensorStorage> dedup = remove_duplicates(processed_tensor_storages);
-    processed_tensor_storages        = dedup;
-    curr_time_ms                     = ggml_time_ms();
-    process_time_ms                  = curr_time_ms - prev_time_ms;
-    prev_time_ms                     = curr_time_ms;
+        std::vector<IndexedStorage> deduplicated;
+        deduplicated.reserve(all_results.size());
+        std::unordered_map<std::string, size_t> name_to_pos;
+        for (auto& entry : all_results) {
+            auto it = name_to_pos.find(entry.ts.name);
+            if (it == name_to_pos.end()) {
+                name_to_pos.emplace(entry.ts.name, deduplicated.size());
+                deduplicated.push_back(entry);
+            } else if (deduplicated[it->second].index < entry.index) {
+                deduplicated[it->second] = entry;
+            }
+        }
+
+        std::sort(deduplicated.begin(), deduplicated.end(), [](const IndexedStorage& a, const IndexedStorage& b) {
+            return a.index < b.index;
+        });
+
+        processed_tensor_storages.reserve(deduplicated.size());
+        for (auto& entry : deduplicated) {
+            processed_tensor_storages.push_back(entry.ts);
+        }
+    }
+
+    process_time_ms = ggml_time_ms() - start_time;
+
+    bool success                          = true;
+    size_t total_tensors_processed        = 0;
+    const size_t total_tensors_to_process = processed_tensor_storages.size();
+    const int64_t t_start                 = ggml_time_ms();
+    int last_n_threads                    = 1;
 
-    bool success = true;
     for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
         std::string file_path = file_paths_[file_index];
         LOG_DEBUG("loading tensors from %s", file_path.c_str());
 
-        std::ifstream file(file_path, std::ios::binary);
-        if (!file.is_open()) {
-            LOG_ERROR("failed to open '%s'", file_path.c_str());
-            return false;
+        std::vector<const TensorStorage*> file_tensors;
+        for (const auto& ts : processed_tensor_storages) {
+            if (ts.file_index == file_index) {
+                file_tensors.push_back(&ts);
+            }
+        }
+        if (file_tensors.empty()) {
+            continue;
         }
 
         bool is_zip = false;
-        for (auto& tensor_storage : tensor_storages) {
-            if (tensor_storage.file_index != file_index) {
-                continue;
-            }
-            if (tensor_storage.index_in_zip >= 0) {
+        for (auto const& ts : file_tensors) {
+            if (ts->index_in_zip >= 0) {
                 is_zip = true;
                 break;
             }
         }
 
-        struct zip_t* zip = NULL;
-        if (is_zip) {
-            zip = zip_open(file_path.c_str(), 0, 'r');
-            if (zip == NULL) {
-                LOG_ERROR("failed to open zip '%s'", file_path.c_str());
-                return false;
-            }
+        int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size());
+        if (n_threads < 1) {
+            n_threads = 1;
+        }
+        last_n_threads = n_threads;
+
+        std::atomic<size_t> tensor_idx(0);
+        std::atomic<bool> failed(false);
+        std::vector<std::thread> workers;
+
+        for (int i = 0; i < n_threads; ++i) {
+            workers.emplace_back([&, file_path, is_zip]() {
+                std::ifstream file;
+                struct zip_t* zip = NULL;
+                if (is_zip) {
+                    zip = zip_open(file_path.c_str(), 0, 'r');
+                    if (zip == NULL) {
+                        LOG_ERROR("failed to open zip '%s'", file_path.c_str());
+                        failed = true;
+                        return;
+                    }
+                } else {
+                    file.open(file_path, std::ios::binary);
+                    if (!file.is_open()) {
+                        LOG_ERROR("failed to open '%s'", file_path.c_str());
+                        failed = true;
+                        return;
+                    }
+                }
+
+                std::vector<uint8_t> read_buffer;
+                std::vector<uint8_t> convert_buffer;
+
+                while (true) {
+                    int64_t t0, t1;
+                    size_t idx = tensor_idx.fetch_add(1);
+                    if (idx >= file_tensors.size() || failed) {
+                        break;
+                    }
+
+                    const TensorStorage& tensor_storage = *file_tensors[idx];
+                    ggml_tensor* dst_tensor             = NULL;
+
+                    t0 = ggml_time_ms();
+
+                    if (!on_new_tensor_cb(tensor_storage, &dst_tensor)) {
+                        LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
+                        failed = true;
+                        break;
+                    }
+
+                    if (dst_tensor == NULL) {
+                        t1 = ggml_time_ms();
+                        read_time_ms.fetch_add(t1 - t0);
+                        continue;
+                    }
+
+                    size_t nbytes_to_read = tensor_storage.nbytes_to_read();
+
+                    auto read_data = [&](char* buf, size_t n) {
+                        if (zip != NULL) {
+                            zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
+                            size_t entry_size = zip_entry_size(zip);
+                            if (entry_size != n) {
+                                int64_t t_memcpy_start;
+                                read_buffer.resize(entry_size);
+                                zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
+                                t_memcpy_start = ggml_time_ms();
+                                memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
+                                memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
+                            } else {
+                                zip_entry_noallocread(zip, (void*)buf, n);
+                            }
+                            zip_entry_close(zip);
+                        } else {
+                            file.seekg(tensor_storage.offset);
+                            file.read(buf, n);
+                            if (!file) {
+                                LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
+                                failed = true;
+                            }
+                        }
+                    };
+
+                    if (dst_tensor->buffer == NULL || ggml_backend_buffer_is_host(dst_tensor->buffer)) {
+                        if (tensor_storage.type == dst_tensor->type) {
+                            GGML_ASSERT(ggml_nbytes(dst_tensor) == tensor_storage.nbytes());
+                            if (tensor_storage.is_f64 || tensor_storage.is_i64) {
+                                read_buffer.resize(tensor_storage.nbytes_to_read());
+                                read_data((char*)read_buffer.data(), nbytes_to_read);
+                            } else {
+                                read_data((char*)dst_tensor->data, nbytes_to_read);
+                            }
+                            t1 = ggml_time_ms();
+                            read_time_ms.fetch_add(t1 - t0);
+
+                            t0 = ggml_time_ms();
+                            if (tensor_storage.is_bf16) {
+                                // inplace op
+                                bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
+                            } else if (tensor_storage.is_f8_e4m3) {
+                                // inplace op
+                                f8_e4m3_to_f16_vec((uint8_t*)dst_tensor->data, (uint16_t*)dst_tensor->data, tensor_storage.nelements());
+                            } else if (tensor_storage.is_f8_e5m2) {
+                                // inplace op
+                                f8_e5m2_to_f16_vec((uint8_t*)dst_tensor->data, (uint16_t*)dst_tensor->data, tensor_storage.nelements());
+                            } else if (tensor_storage.is_f64) {
+                                f64_to_f32_vec((double*)read_buffer.data(), (float*)dst_tensor->data, tensor_storage.nelements());
+                            } else if (tensor_storage.is_i64) {
+                                i64_to_i32_vec((int64_t*)read_buffer.data(), (int32_t*)dst_tensor->data, tensor_storage.nelements());
+                            }
+                            t1 = ggml_time_ms();
+                            convert_time_ms.fetch_add(t1 - t0);
+                        } else {
+                            read_buffer.resize(std::max(tensor_storage.nbytes(), tensor_storage.nbytes_to_read()));
+                            read_data((char*)read_buffer.data(), nbytes_to_read);
+                            t1 = ggml_time_ms();
+                            read_time_ms.fetch_add(t1 - t0);
+
+                            t0 = ggml_time_ms();
+                            if (tensor_storage.is_bf16) {
+                                // inplace op
+                                bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                            } else if (tensor_storage.is_f8_e4m3) {
+                                // inplace op
+                                f8_e4m3_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
+                            } else if (tensor_storage.is_f8_e5m2) {
+                                // inplace op
+                                f8_e5m2_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
+                            } else if (tensor_storage.is_f64) {
+                                // inplace op
+                                f64_to_f32_vec((double*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                            } else if (tensor_storage.is_i64) {
+                                // inplace op
+                                i64_to_i32_vec((int64_t*)read_buffer.data(), (int32_t*)read_buffer.data(), tensor_storage.nelements());
+                            }
+                            convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data, dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
+                            t1 = ggml_time_ms();
+                            convert_time_ms.fetch_add(t1 - t0);
+                        }
+                    } else {
+                        read_buffer.resize(std::max(tensor_storage.nbytes(), tensor_storage.nbytes_to_read()));
+                        read_data((char*)read_buffer.data(), nbytes_to_read);
+                        t1 = ggml_time_ms();
+                        read_time_ms.fetch_add(t1 - t0);
+
+                        t0 = ggml_time_ms();
+                        if (tensor_storage.is_bf16) {
+                            // inplace op
+                            bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                        } else if (tensor_storage.is_f8_e4m3) {
+                            // inplace op
+                            f8_e4m3_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
+                        } else if (tensor_storage.is_f8_e5m2) {
+                            // inplace op
+                            f8_e5m2_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
+                        } else if (tensor_storage.is_f64) {
+                            // inplace op
+                            f64_to_f32_vec((double*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                        } else if (tensor_storage.is_i64) {
+                            // inplace op
+                            i64_to_i32_vec((int64_t*)read_buffer.data(), (int32_t*)read_buffer.data(), tensor_storage.nelements());
+                        }
+
+                        if (tensor_storage.type == dst_tensor->type) {
+                            // copy to device memory
+                            t1 = ggml_time_ms();
+                            convert_time_ms.fetch_add(t1 - t0);
+                            t0 = ggml_time_ms();
+                            ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
+                            t1 = ggml_time_ms();
+                            copy_to_backend_time_ms.fetch_add(t1 - t0);
+                        } else {
+                            // convert first, then copy to device memory
+
+                            convert_buffer.resize(ggml_nbytes(dst_tensor));
+                            convert_tensor((void*)read_buffer.data(), tensor_storage.type, (void*)convert_buffer.data(), dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
+                            t1 = ggml_time_ms();
+                            convert_time_ms.fetch_add(t1 - t0);
+                            t0 = ggml_time_ms();
+                            ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
+                            t1 = ggml_time_ms();
+                            copy_to_backend_time_ms.fetch_add(t1 - t0);
+                        }
+                    }
+                }
+                if (zip != NULL) {
+                    zip_close(zip);
+                }
+            });
         }
 
-        std::vector<uint8_t> read_buffer;
-        std::vector<uint8_t> convert_buffer;
-
-        auto read_data = [&](const TensorStorage& tensor_storage, char* buf, size_t n) {
-            if (zip != NULL) {
-                zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
-                size_t entry_size = zip_entry_size(zip);
-                if (entry_size != n) {
-                    read_buffer.resize(entry_size);
-                    prev_time_ms = ggml_time_ms();
-                    zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
-                    curr_time_ms = ggml_time_ms();
-                    read_time_ms += curr_time_ms - prev_time_ms;
-                    prev_time_ms = curr_time_ms;
-                    memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
-                    curr_time_ms = ggml_time_ms();
-                    memcpy_time_ms += curr_time_ms - prev_time_ms;
-                } else {
-                    prev_time_ms = ggml_time_ms();
-                    zip_entry_noallocread(zip, (void*)buf, n);
-                    curr_time_ms = ggml_time_ms();
-                    read_time_ms += curr_time_ms - prev_time_ms;
-                }
-                zip_entry_close(zip);
-            } else {
-                prev_time_ms = ggml_time_ms();
-                file.seekg(tensor_storage.offset);
-                file.read(buf, n);
-                curr_time_ms = ggml_time_ms();
-                read_time_ms += curr_time_ms - prev_time_ms;
-                if (!file) {
-                    LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                    return false;
-                }
-            }
-            return true;
-        };
-        int tensor_count = 0;
-        int64_t t0       = ggml_time_ms();
-        int64_t t1       = t0;
-        bool partial     = true;
-        int tensor_max   = (int)processed_tensor_storages.size();
-        pretty_progress(0, tensor_max, 0.0f);
-        for (auto& tensor_storage : processed_tensor_storages) {
-            if (tensor_storage.file_index != file_index) {
-                ++tensor_count;
-                continue;
-            }
-            ggml_tensor* dst_tensor = NULL;
-
-            success = on_new_tensor_cb(tensor_storage, &dst_tensor);
-            if (!success) {
-                LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
+        while (true) {
+            size_t current_idx = tensor_idx.load();
+            if (current_idx >= file_tensors.size() || failed) {
                 break;
             }
-
-            if (dst_tensor == NULL) {
-                ++tensor_count;
-                continue;
-            }
-
-            size_t nbytes_to_read = tensor_storage.nbytes_to_read();
-
-            if (dst_tensor->buffer == NULL || ggml_backend_buffer_is_host(dst_tensor->buffer)) {
-                // for the CPU and Metal backend, we can copy directly into the tensor
-                if (tensor_storage.type == dst_tensor->type) {
-                    GGML_ASSERT(ggml_nbytes(dst_tensor) == tensor_storage.nbytes());
-                    if (tensor_storage.is_f64 || tensor_storage.is_i64) {
-                        read_buffer.resize(tensor_storage.nbytes_to_read());
-                        read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-                    } else {
-                        read_data(tensor_storage, (char*)dst_tensor->data, nbytes_to_read);
-                    }
-
-                    prev_time_ms = ggml_time_ms();
-                    if (tensor_storage.is_bf16) {
-                        // inplace op
-                        bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
-                    } else if (tensor_storage.is_f8_e4m3) {
-                        // inplace op
-                        f8_e4m3_to_f16_vec((uint8_t*)dst_tensor->data, (uint16_t*)dst_tensor->data, tensor_storage.nelements());
-                    } else if (tensor_storage.is_f8_e5m2) {
-                        // inplace op
-                        f8_e5m2_to_f16_vec((uint8_t*)dst_tensor->data, (uint16_t*)dst_tensor->data, tensor_storage.nelements());
-                    } else if (tensor_storage.is_f64) {
-                        f64_to_f32_vec((double*)read_buffer.data(), (float*)dst_tensor->data, tensor_storage.nelements());
-                    } else if (tensor_storage.is_i64) {
-                        i64_to_i32_vec((int64_t*)read_buffer.data(), (int32_t*)dst_tensor->data, tensor_storage.nelements());
-                    }
-                    curr_time_ms = ggml_time_ms();
-                    convert_time_ms += curr_time_ms - prev_time_ms;
-                } else {
-                    read_buffer.resize(std::max(tensor_storage.nbytes(), tensor_storage.nbytes_to_read()));
-                    read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-
-                    prev_time_ms = ggml_time_ms();
-                    if (tensor_storage.is_bf16) {
-                        // inplace op
-                        bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                    } else if (tensor_storage.is_f8_e4m3) {
-                        // inplace op
-                        f8_e4m3_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
-                    } else if (tensor_storage.is_f8_e5m2) {
-                        // inplace op
-                        f8_e5m2_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
-                    } else if (tensor_storage.is_f64) {
-                        // inplace op
-                        f64_to_f32_vec((double*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                    } else if (tensor_storage.is_i64) {
-                        // inplace op
-                        i64_to_i32_vec((int64_t*)read_buffer.data(), (int32_t*)read_buffer.data(), tensor_storage.nelements());
-                    }
-
-                    convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
-                                   dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
-                    curr_time_ms = ggml_time_ms();
-                    convert_time_ms += curr_time_ms - prev_time_ms;
-                }
-            } else {
-                read_buffer.resize(std::max(tensor_storage.nbytes(), tensor_storage.nbytes_to_read()));
-                read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-
-                prev_time_ms = ggml_time_ms();
-                if (tensor_storage.is_bf16) {
-                    // inplace op
-                    bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                } else if (tensor_storage.is_f8_e4m3) {
-                    // inplace op
-                    f8_e4m3_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
-                } else if (tensor_storage.is_f8_e5m2) {
-                    // inplace op
-                    f8_e5m2_to_f16_vec((uint8_t*)read_buffer.data(), (uint16_t*)read_buffer.data(), tensor_storage.nelements());
-                } else if (tensor_storage.is_f64) {
-                    // inplace op
-                    f64_to_f32_vec((double*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                } else if (tensor_storage.is_i64) {
-                    // inplace op
-                    i64_to_i32_vec((int64_t*)read_buffer.data(), (int32_t*)read_buffer.data(), tensor_storage.nelements());
-                }
-
-                if (tensor_storage.type == dst_tensor->type) {
-                    // copy to device memory
-                    curr_time_ms = ggml_time_ms();
-                    convert_time_ms += curr_time_ms - prev_time_ms;
-                    prev_time_ms = curr_time_ms;
-                    ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
-                    curr_time_ms = ggml_time_ms();
-                    copy_to_backend_time_ms += curr_time_ms - prev_time_ms;
-                } else {
-                    // convert first, then copy to device memory
-                    convert_buffer.resize(ggml_nbytes(dst_tensor));
-                    convert_tensor((void*)read_buffer.data(), tensor_storage.type,
-                                   (void*)convert_buffer.data(), dst_tensor->type,
-                                   (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0]);
-                    curr_time_ms = ggml_time_ms();
-                    convert_time_ms += curr_time_ms - prev_time_ms;
-                    prev_time_ms = curr_time_ms;
-                    ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
-                    curr_time_ms = ggml_time_ms();
-                    copy_to_backend_time_ms += curr_time_ms - prev_time_ms;
-                }
-            }
-            ++tensor_count;
-            int64_t t2 = ggml_time_ms();
-            if ((t2 - t1) >= 200) {
-                t1 = t2;
-                pretty_progress(tensor_count, tensor_max, (t1 - t0) / (1000.0f * tensor_count));
-                partial = tensor_count != tensor_max;
-            }
+            size_t curr_num = total_tensors_processed + current_idx;
+            pretty_progress(curr_num, total_tensors_to_process, (ggml_time_ms() - t_start) / 1000.0f / (curr_num + 1e-6f));
+            std::this_thread::sleep_for(std::chrono::milliseconds(200));
         }
 
-        if (partial) {
-            if (tensor_count >= 1) {
-                t1 = ggml_time_ms();
-                pretty_progress(tensor_count, tensor_max, (t1 - t0) / (1000.0f * tensor_count));
-            }
-            if (tensor_count < tensor_max) {
-                printf("\n");
-            }
+        for (auto& w : workers) {
+            w.join();
         }
 
-        if (zip != NULL) {
-            zip_close(zip);
-        }
-
-        if (!success) {
+        if (failed) {
+            success = false;
             break;
         }
+        total_tensors_processed += file_tensors.size();
+        pretty_progress(total_tensors_processed, total_tensors_to_process, (ggml_time_ms() - t_start) / 1000.0f / (total_tensors_processed + 1e-6f));
+        if (total_tensors_processed < total_tensors_to_process) {
+            printf("\n");
+        }
     }
+
     int64_t end_time = ggml_time_ms();
     LOG_INFO("loading tensors completed, taking %.2fs (process: %.2fs, read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
              (end_time - start_time) / 1000.f,
              process_time_ms / 1000.f,
-             read_time_ms / 1000.f,
-             memcpy_time_ms / 1000.f,
-             convert_time_ms / 1000.f,
-             copy_to_backend_time_ms / 1000.f);
+             (read_time_ms.load() / (float)last_n_threads) / 1000.f,
+             (memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
+             (convert_time_ms.load() / (float)last_n_threads) / 1000.f,
+             (copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
     return success;
 }
 
 bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
-                               std::set<std::string> ignore_tensors) {
+                               std::set<std::string> ignore_tensors,
+                               int n_threads) {
     std::set<std::string> tensor_names_in_file;
+    std::mutex tensor_names_mutex;
     auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
         const std::string& name = tensor_storage.name;
         // LOG_DEBUG("%s", tensor_storage.to_string().c_str());
-        tensor_names_in_file.insert(name);
+        {
+            std::lock_guard<std::mutex> lock(tensor_names_mutex);
+            tensor_names_in_file.insert(name);
+        }
 
         struct ggml_tensor* real;
         if (tensors.find(name) != tensors.end()) {
@@ -2263,7 +2404,7 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
         return true;
     };
 
-    bool success = load_tensors(on_new_tensor_cb);
+    bool success = load_tensors(on_new_tensor_cb, n_threads);
     if (!success) {
         LOG_ERROR("load tensors from file failed");
         return false;
@@ -2351,6 +2492,8 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
             // Pass, do not convert. For MMDiT
         } else if (contains(name, "time_embed.") || contains(name, "label_emb.")) {
             // Pass, do not convert. For Unet
+        } else if (contains(name, "embedding")) {
+            // Pass, do not convert embedding
         } else {
             return true;
         }
@@ -2370,6 +2513,7 @@ bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type
 
     auto tensor_type_rules = parse_tensor_type_rules(tensor_type_rules_str);
 
+    std::mutex tensor_mutex;
     auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
         const std::string& name = tensor_storage.name;
         ggml_type tensor_type   = tensor_storage.type;
@@ -2387,6 +2531,7 @@ bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type
             tensor_type = dst_type;
         }
 
+        std::lock_guard<std::mutex> lock(tensor_mutex);
         ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
         if (tensor == NULL) {
             LOG_ERROR("ggml_new_tensor failed");

model.h

@@ -31,9 +31,12 @@ enum SDVersion {
     VERSION_SD3,
     VERSION_FLUX,
     VERSION_FLUX_FILL,
+    VERSION_FLUX_CONTROLS,
+    VERSION_FLEX_2,
     VERSION_WAN2,
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
+    VERSION_QWEN_IMAGE,
     VERSION_COUNT,
 };
 
@@ -66,7 +69,7 @@ static inline bool sd_version_is_sd3(SDVersion version) {
 }
 
 static inline bool sd_version_is_flux(SDVersion version) {
-    if (version == VERSION_FLUX || version == VERSION_FLUX_FILL) {
+    if (version == VERSION_FLUX || version == VERSION_FLUX_FILL || version == VERSION_FLUX_CONTROLS || version == VERSION_FLEX_2) {
         return true;
     }
     return false;
@@ -79,15 +82,25 @@ static inline bool sd_version_is_wan(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_qwen_image(SDVersion version) {
+    if (version == VERSION_QWEN_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_inpaint(SDVersion version) {
-    if (version == VERSION_SD1_INPAINT || version == VERSION_SD2_INPAINT || version == VERSION_SDXL_INPAINT || version == VERSION_FLUX_FILL) {
+    if (version == VERSION_SD1_INPAINT || version == VERSION_SD2_INPAINT || version == VERSION_SDXL_INPAINT || version == VERSION_FLUX_FILL || version == VERSION_FLEX_2) {
         return true;
     }
     return false;
 }
 
 static inline bool sd_version_is_dit(SDVersion version) {
-    if (sd_version_is_flux(version) || sd_version_is_sd3(version) || sd_version_is_wan(version)) {
+    if (sd_version_is_flux(version) ||
+        sd_version_is_sd3(version) ||
+        sd_version_is_wan(version) ||
+        sd_version_is_qwen_image(version)) {
         return true;
     }
     return false;
@@ -97,8 +110,12 @@ static inline bool sd_version_is_unet_edit(SDVersion version) {
     return version == VERSION_SD1_PIX2PIX || version == VERSION_SDXL_PIX2PIX;
 }
 
+static inline bool sd_version_is_control(SDVersion version) {
+    return version == VERSION_FLUX_CONTROLS || version == VERSION_FLEX_2;
+}
+
 static bool sd_version_is_inpaint_or_unet_edit(SDVersion version) {
-    return sd_version_is_unet_edit(version) || sd_version_is_inpaint(version);
+    return sd_version_is_unet_edit(version) || sd_version_is_inpaint(version) || sd_version_is_control(version);
 }
 
 enum PMVersion {
@@ -247,9 +264,18 @@ public:
     ggml_type get_diffusion_model_wtype();
     ggml_type get_vae_wtype();
     void set_wtype_override(ggml_type wtype, std::string prefix = "");
-    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb);
+    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0);
     bool load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
-                      std::set<std::string> ignore_tensors = {});
+                      std::set<std::string> ignore_tensors = {},
+                      int n_threads                        = 0);
+
+    std::vector<std::string> get_tensor_names() const {
+        std::vector<std::string> names;
+        for (const auto& ts : tensor_storages) {
+            names.push_back(ts.name);
+        }
+        return names;
+    }
 
     bool save_to_gguf_file(const std::string& file_path, ggml_type type, const std::string& tensor_type_rules);
     bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
@@ -257,6 +283,7 @@ public:
     ~ModelLoader() = default;
 
     static std::string load_merges();
+    static std::string load_qwen2_merges();
     static std::string load_t5_tokenizer_json();
     static std::string load_umt5_tokenizer_json();
 };

pmid.hpp

@@ -42,41 +42,6 @@ public:
     }
 };
 
-/*
-class QFormerPerceiver(nn.Module):
-    def __init__(self, id_embeddings_dim, cross_attention_dim, num_tokens, embedding_dim=1024, use_residual=True, ratio=4):
-        super().__init__()
-
-        self.num_tokens = num_tokens
-        self.cross_attention_dim = cross_attention_dim
-        self.use_residual = use_residual
-        print(cross_attention_dim*num_tokens)
-        self.token_proj = nn.Sequential(
-            nn.Linear(id_embeddings_dim, id_embeddings_dim*ratio),
-            nn.GELU(),
-            nn.Linear(id_embeddings_dim*ratio, cross_attention_dim*num_tokens),
-        )
-        self.token_norm = nn.LayerNorm(cross_attention_dim)
-        self.perceiver_resampler = FacePerceiverResampler(
-            dim=cross_attention_dim,
-            depth=4,
-            dim_head=128,
-            heads=cross_attention_dim // 128,
-            embedding_dim=embedding_dim,
-            output_dim=cross_attention_dim,
-            ff_mult=4,
-        )
-
-    def forward(self, x, last_hidden_state):
-        x = self.token_proj(x)
-        x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
-        x = self.token_norm(x) # cls token
-        out = self.perceiver_resampler(x, last_hidden_state) # retrieve from patch tokens
-        if self.use_residual: # TODO: if use_residual is not true
-            out = x + 1.0 * out
-        return out
-*/
-
 struct PMFeedForward : public GGMLBlock {
     // network hparams
     int dim;
@@ -122,17 +87,8 @@ public:
         int64_t ne[4];
         for (int i = 0; i < 4; ++i)
             ne[i] = x->ne[i];
-        // print_ggml_tensor(x, true, "PerceiverAttention reshape x 0: ");
-        // printf("heads = %d \n", heads);
-        // x = ggml_view_4d(ctx, x, x->ne[0], x->ne[1], heads, x->ne[2]/heads,
-        //                          x->nb[1], x->nb[2], x->nb[3], 0);
         x = ggml_reshape_4d(ctx, x, x->ne[0] / heads, heads, x->ne[1], x->ne[2]);
-        // x = ggml_view_4d(ctx, x, x->ne[0]/heads, heads, x->ne[1], x->ne[2],
-        //                          x->nb[1], x->nb[2], x->nb[3], 0);
-        // x = ggml_cont(ctx, x);
         x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));
-        // print_ggml_tensor(x, true, "PerceiverAttention reshape x 1: ");
-        // x  = ggml_reshape_4d(ctx, x, ne[0], heads, ne[1], ne[2]/heads);
         return x;
     }
 
@@ -269,17 +225,6 @@ public:
             4));
     }
 
-    /*
-    def forward(self, x, last_hidden_state):
-        x = self.token_proj(x)
-        x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
-        x = self.token_norm(x) # cls token
-        out = self.perceiver_resampler(x, last_hidden_state) # retrieve from patch tokens
-        if self.use_residual: # TODO: if use_residual is not true
-            out = x + 1.0 * out
-        return out
-    */
-
     struct ggml_tensor* forward(struct ggml_context* ctx,
                                 struct ggml_tensor* x,
                                 struct ggml_tensor* last_hidden_state) {
@@ -299,113 +244,6 @@ public:
     }
 };
 
-/*
-class FacePerceiverResampler(torch.nn.Module):
-    def __init__(
-        self,
-        *,
-        dim=768,
-        depth=4,
-        dim_head=64,
-        heads=16,
-        embedding_dim=1280,
-        output_dim=768,
-        ff_mult=4,
-    ):
-        super().__init__()
-
-        self.proj_in = torch.nn.Linear(embedding_dim, dim)
-        self.proj_out = torch.nn.Linear(dim, output_dim)
-        self.norm_out = torch.nn.LayerNorm(output_dim)
-        self.layers = torch.nn.ModuleList([])
-        for _ in range(depth):
-            self.layers.append(
-                torch.nn.ModuleList(
-                    [
-                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
-                        FeedForward(dim=dim, mult=ff_mult),
-                    ]
-                )
-            )
-
-    def forward(self, latents, x):
-        x = self.proj_in(x)
-        for attn, ff in self.layers:
-            latents = attn(x, latents) + latents
-            latents = ff(latents) + latents
-        latents = self.proj_out(latents)
-        return self.norm_out(latents)
-*/
-
-/*
-
-def FeedForward(dim, mult=4):
-    inner_dim = int(dim * mult)
-    return nn.Sequential(
-        nn.LayerNorm(dim),
-        nn.Linear(dim, inner_dim, bias=False),
-        nn.GELU(),
-        nn.Linear(inner_dim, dim, bias=False),
-    )
-
-def reshape_tensor(x, heads):
-    bs, length, width = x.shape
-    # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
-    x = x.view(bs, length, heads, -1)
-    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
-    x = x.transpose(1, 2)
-    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
-    x = x.reshape(bs, heads, length, -1)
-    return x
-
-class PerceiverAttention(nn.Module):
-    def __init__(self, *, dim, dim_head=64, heads=8):
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.dim_head = dim_head
-        self.heads = heads
-        inner_dim = dim_head * heads
-
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-
-        self.to_q = nn.Linear(dim, inner_dim, bias=False)
-        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
-        self.to_out = nn.Linear(inner_dim, dim, bias=False)
-
-    def forward(self, x, latents):
-        """
-        Args:
-            x (torch.Tensor): image features
-                shape (b, n1, D)
-            latent (torch.Tensor): latent features
-                shape (b, n2, D)
-        """
-        x = self.norm1(x)
-        latents = self.norm2(latents)
-
-        b, l, _ = latents.shape
-
-        q = self.to_q(latents)
-        kv_input = torch.cat((x, latents), dim=-2)
-        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
-
-        q = reshape_tensor(q, self.heads)
-        k = reshape_tensor(k, self.heads)
-        v = reshape_tensor(v, self.heads)
-
-        # attention
-        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
-        weight = (q * scale) @ (k * scale).transpose(-2, -1)  # More stable with f16 than dividing afterwards
-        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
-        out = weight @ v
-
-        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
-
-        return self.to_out(out)
-
-*/
-
 struct FuseModule : public GGMLBlock {
     // network hparams
     int embed_dim;
@@ -425,31 +263,13 @@ public:
         auto mlp2       = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp2"]);
         auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
 
-        // print_ggml_tensor(id_embeds, true, "Fuseblock id_embeds: ");
-        // print_ggml_tensor(prompt_embeds, true, "Fuseblock prompt_embeds: ");
-
-        // 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));
-        // print_ggml_tensor(id_embeds0, true, "Fuseblock id_embeds0: ");
-        // print_ggml_tensor(prompt_embeds0, true, "Fuseblock prompt_embeds0: ");
-        // concat is along dim 2
-        // auto stacked_id_embeds = ggml_concat(ctx, prompt_embeds0, id_embeds0, 2);
         auto stacked_id_embeds = ggml_concat(ctx, prompt_embeds, id_embeds, 0);
-        // print_ggml_tensor(stacked_id_embeds, true, "Fuseblock stacked_id_embeds 0: ");
-        // stacked_id_embeds      = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 1, 2, 0, 3));
-        // print_ggml_tensor(stacked_id_embeds, true, "Fuseblock stacked_id_embeds 1: ");
-        // 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);
 
-        // print_ggml_tensor(stacked_id_embeds, true, "Fuseblock stacked_id_embeds 1: ");
-
         return stacked_id_embeds;
     }
 
@@ -464,21 +284,14 @@ public:
 
         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");
         valid_id_embeds                       = ggml_reshape_2d(ctx, valid_id_embeds, valid_id_embeds->ne[0],
                                                                 ggml_nelements(valid_id_embeds) / valid_id_embeds->ne[0]);
         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));
-        // print_ggml_tensor(stacked_id_embeds, true, "AA stacked_id_embeds");
-        // print_ggml_tensor(left, true, "AA left");
-        // print_ggml_tensor(right, true, "AA right");
         if (left && right) {
             stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds, 1);
             stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right, 1);
@@ -487,15 +300,12 @@ public:
         } else if (right) {
             stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right, 1);
         }
-        // print_ggml_tensor(stacked_id_embeds, true, "BB stacked_id_embeds");
-        // stacked_id_embeds                         = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 0, 2, 1, 3));
-        // print_ggml_tensor(stacked_id_embeds, true, "CC stacked_id_embeds");
+
         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");
-        // print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds: ");
         return updated_prompt_embeds;
     }
 };
@@ -551,34 +361,11 @@ struct PhotoMakerIDEncoder_CLIPInsightfaceExtendtokenBlock : public CLIPVisionMo
           num_tokens(2) {
         blocks["visual_projection_2"] = std::shared_ptr<GGMLBlock>(new Linear(1024, 1280, false));
         blocks["fuse_module"]         = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
-        /*
-        cross_attention_dim = 2048
-        # projection
-        self.num_tokens = 2
-        self.cross_attention_dim = cross_attention_dim
-        self.qformer_perceiver = QFormerPerceiver(
-                                    id_embeddings_dim,
-                                    cross_attention_dim,
-                                    self.num_tokens,
-                                )*/
-        blocks["qformer_perceiver"] = std::shared_ptr<GGMLBlock>(new QFormerPerceiver(id_embeddings_dim,
-                                                                                      cross_attention_dim,
-                                                                                      num_tokens));
+        blocks["qformer_perceiver"]   = std::shared_ptr<GGMLBlock>(new QFormerPerceiver(id_embeddings_dim,
+                                                                                        cross_attention_dim,
+                                                                                        num_tokens));
     }
 
-    /*
-    def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds):
-        b, num_inputs, c, h, w = id_pixel_values.shape
-        id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w)
-
-        last_hidden_state = self.vision_model(id_pixel_values)[0]
-        id_embeds = id_embeds.view(b * num_inputs, -1)
-
-        id_embeds = self.qformer_perceiver(id_embeds, last_hidden_state)
-        id_embeds = id_embeds.view(b, num_inputs, self.num_tokens, -1)
-        updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask)
-    */
-
     struct ggml_tensor* forward(struct ggml_context* ctx,
                                 ggml_backend_t backend,
                                 struct ggml_tensor* id_pixel_values,
@@ -804,7 +591,7 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
         return "id_embeds";
     }
 
-    bool load_from_file(bool filter_tensor = false) {
+    bool load_from_file(bool filter_tensor, int n_threads) {
         LOG_INFO("loading PhotoMaker ID Embeds from '%s'", file_path.c_str());
 
         if (load_failed) {
@@ -812,7 +599,8 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             return false;
         }
 
-        bool dry_run          = true;
+        bool dry_run = true;
+        std::mutex tensor_mutex;
         auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
             const std::string& name = tensor_storage.name;
 
@@ -821,6 +609,7 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
                 return true;
             }
             if (dry_run) {
+                std::lock_guard<std::mutex> lock(tensor_mutex);
                 struct ggml_tensor* real = ggml_new_tensor(params_ctx,
                                                            tensor_storage.type,
                                                            tensor_storage.n_dims,
@@ -834,11 +623,11 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             return true;
         };
 
-        model_loader->load_tensors(on_new_tensor_cb);
+        model_loader->load_tensors(on_new_tensor_cb, n_threads);
         alloc_params_buffer();
 
         dry_run = false;
-        model_loader->load_tensors(on_new_tensor_cb);
+        model_loader->load_tensors(on_new_tensor_cb, n_threads);
 
         LOG_DEBUG("finished loading PhotoMaker ID Embeds ");
         return true;

preprocessing.hpp

@@ -6,7 +6,7 @@
 
 void convolve(struct ggml_tensor* input, struct ggml_tensor* output, struct ggml_tensor* kernel, int padding) {
     struct ggml_init_params params;
-    params.mem_size                 = 20 * 1024 * 1024;  // 10
+    params.mem_size                 = 80 * input->ne[0] * input->ne[1];  // 20M for 512x512
     params.mem_buffer               = NULL;
     params.no_alloc                 = false;
     struct ggml_context* ctx0       = ggml_init(params);
@@ -162,16 +162,16 @@ void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float lo
     }
 }
 
-uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
+bool preprocess_canny(sd_image_t img, 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);  // 10MB
+    params.mem_size               = static_cast<size_t>(40 * img.width * img.height);  // 10MB for 512x512
     params.mem_buffer             = NULL;
     params.no_alloc               = false;
     struct ggml_context* work_ctx = ggml_init(params);
 
     if (!work_ctx) {
         LOG_ERROR("ggml_init() failed");
-        return NULL;
+        return false;
     }
 
     float kX[9] = {
@@ -192,8 +192,8 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
     struct ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
     memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
     gaussian_kernel(gkernel);
-    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
-    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 1, 1);
+    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
+    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
     struct ggml_tensor* iX         = ggml_dup_tensor(work_ctx, image_gray);
     struct ggml_tensor* iY         = ggml_dup_tensor(work_ctx, image_gray);
     struct ggml_tensor* G          = ggml_dup_tensor(work_ctx, image_gray);
@@ -209,8 +209,8 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
     non_max_supression(image_gray, G, tetha);
     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++) {
+    for (int iy = 0; iy < img.height; iy++) {
+        for (int ix = 0; ix < img.width; ix++) {
             float gray = ggml_tensor_get_f32(image_gray, ix, iy);
             gray       = inverse ? 1.0f - gray : gray;
             ggml_tensor_set_f32(image, gray, ix, iy);
@@ -218,10 +218,9 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
             ggml_tensor_set_f32(image, gray, ix, iy, 2);
         }
     }
-    free(img);
-    uint8_t* output = sd_tensor_to_image(image);
+    sd_tensor_to_image(image, img.data);
     ggml_free(work_ctx);
-    return output;
+    return true;
 }
 
 #endif  // __PREPROCESSING_HPP__

qwen_image.hpp

@@ -0,0 +1,694 @@
+#ifndef __QWEN_IMAGE_HPP__
+#define __QWEN_IMAGE_HPP__
+
+#include "common.hpp"
+#include "flux.hpp"
+#include "ggml_extend.hpp"
+
+namespace Qwen {
+    constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480;
+
+    struct TimestepEmbedding : public GGMLBlock {
+    public:
+        TimestepEmbedding(int64_t in_channels,
+                          int64_t time_embed_dim,
+                          int64_t out_dim       = 0,
+                          int64_t cond_proj_dim = 0,
+                          bool sample_proj_bias = true) {
+            blocks["linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, time_embed_dim, sample_proj_bias));
+            if (cond_proj_dim > 0) {
+                blocks["cond_proj"] = std::shared_ptr<GGMLBlock>(new Linear(cond_proj_dim, in_channels, false));
+            }
+            if (out_dim <= 0) {
+                out_dim = time_embed_dim;
+            }
+            blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, out_dim, sample_proj_bias));
+        }
+
+        struct ggml_tensor* forward(struct ggml_context* ctx,
+                                    struct ggml_tensor* sample,
+                                    struct ggml_tensor* condition = nullptr) {
+            if (condition != nullptr) {
+                auto cond_proj = std::dynamic_pointer_cast<Linear>(blocks["cond_proj"]);
+                sample         = ggml_add(ctx, sample, cond_proj->forward(ctx, condition));
+            }
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
+
+            sample = linear_1->forward(ctx, sample);
+            sample = ggml_silu_inplace(ctx, sample);
+            sample = linear_2->forward(ctx, sample);
+            return sample;
+        }
+    };
+
+    struct QwenTimestepProjEmbeddings : public GGMLBlock {
+    public:
+        QwenTimestepProjEmbeddings(int64_t embedding_dim) {
+            blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
+        }
+
+        struct ggml_tensor* forward(struct ggml_context* ctx,
+                                    struct ggml_tensor* timesteps) {
+            // timesteps: [N,]
+            // return: [N, embedding_dim]
+            auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
+
+            auto timesteps_proj = ggml_nn_timestep_embedding(ctx, timesteps, 256, 10000, 1.f);
+            auto timesteps_emb  = timestep_embedder->forward(ctx, timesteps_proj);
+            return timesteps_emb;
+        }
+    };
+
+    struct QwenImageAttention : public GGMLBlock {
+    protected:
+        int64_t dim_head;
+        bool flash_attn;
+
+    public:
+        QwenImageAttention(int64_t query_dim,
+                           int64_t dim_head,
+                           int64_t num_heads,
+                           int64_t out_dim         = 0,
+                           int64_t out_context_dim = 0,
+                           bool bias               = true,
+                           bool out_bias           = true,
+                           float eps               = 1e-6,
+                           bool flash_attn         = false)
+            : dim_head(dim_head), flash_attn(flash_attn) {
+            int64_t inner_dim = out_dim > 0 ? out_dim : dim_head * num_heads;
+            out_dim           = out_dim > 0 ? out_dim : query_dim;
+            out_context_dim   = out_context_dim > 0 ? out_context_dim : query_dim;
+
+            blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
+            blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
+            blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
+
+            blocks["norm_q"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
+            blocks["norm_k"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
+
+            blocks["add_q_proj"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
+            blocks["add_k_proj"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
+            blocks["add_v_proj"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
+
+            blocks["norm_added_q"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
+            blocks["norm_added_k"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
+
+            float scale = 1.f / 32.f;
+            // The purpose of the scale here is to prevent NaN issues in certain situations.
+            // For example when using CUDA but the weights are k-quants (not all prompts).
+            blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, out_dim, out_bias, false, false, scale));
+            // to_out.1 is nn.Dropout
+
+            blocks["to_add_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, out_context_dim, out_bias, false, false, scale));
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(struct ggml_context* ctx,
+                                                      ggml_backend_t backend,
+                                                      struct ggml_tensor* img,
+                                                      struct ggml_tensor* txt,
+                                                      struct ggml_tensor* pe,
+                                                      struct ggml_tensor* mask = nullptr) {
+            // img: [N, n_img_token, hidden_size]
+            // txt: [N, n_txt_token, hidden_size]
+            // pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
+            // return: ([N, n_img_token, hidden_size], [N, n_txt_token, hidden_size])
+
+            auto norm_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_q"]);
+            auto norm_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_k"]);
+
+            auto to_q     = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
+            auto to_k     = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
+            auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
+            auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+
+            auto norm_added_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_q"]);
+            auto norm_added_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_k"]);
+
+            auto add_q_proj = std::dynamic_pointer_cast<Linear>(blocks["add_q_proj"]);
+            auto add_k_proj = std::dynamic_pointer_cast<Linear>(blocks["add_k_proj"]);
+            auto add_v_proj = std::dynamic_pointer_cast<Linear>(blocks["add_v_proj"]);
+            auto to_add_out = std::dynamic_pointer_cast<Linear>(blocks["to_add_out"]);
+
+            int64_t N           = img->ne[2];
+            int64_t n_img_token = img->ne[1];
+            int64_t n_txt_token = txt->ne[1];
+
+            auto img_q        = to_q->forward(ctx, img);
+            int64_t num_heads = img_q->ne[0] / dim_head;
+            img_q             = ggml_reshape_4d(ctx, img_q, dim_head, num_heads, n_img_token, N);  // [N, n_img_token, n_head, d_head]
+            auto img_k        = to_k->forward(ctx, img);
+            img_k             = ggml_reshape_4d(ctx, img_k, dim_head, num_heads, n_img_token, N);  // [N, n_img_token, n_head, d_head]
+            auto img_v        = to_v->forward(ctx, img);
+            img_v             = ggml_reshape_4d(ctx, img_v, dim_head, num_heads, n_img_token, N);  // [N, n_img_token, n_head, d_head]
+
+            img_q = norm_q->forward(ctx, img_q);
+            img_k = norm_k->forward(ctx, img_k);
+
+            auto txt_q = add_q_proj->forward(ctx, txt);
+            txt_q      = ggml_reshape_4d(ctx, txt_q, dim_head, num_heads, n_txt_token, N);  // [N, n_txt_token, n_head, d_head]
+            auto txt_k = add_k_proj->forward(ctx, txt);
+            txt_k      = ggml_reshape_4d(ctx, txt_k, dim_head, num_heads, n_txt_token, N);  // [N, n_txt_token, n_head, d_head]
+            auto txt_v = add_v_proj->forward(ctx, txt);
+            txt_v      = ggml_reshape_4d(ctx, txt_v, dim_head, num_heads, n_txt_token, N);  // [N, n_txt_token, n_head, d_head]
+
+            txt_q = norm_added_q->forward(ctx, txt_q);
+            txt_k = norm_added_k->forward(ctx, txt_k);
+
+            auto q = ggml_concat(ctx, txt_q, img_q, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
+            auto k = ggml_concat(ctx, txt_k, img_k, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
+            auto v = ggml_concat(ctx, txt_v, img_v, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
+
+            auto attn         = Rope::attention(ctx, backend, q, k, v, pe, mask, flash_attn, (1.0f / 128.f));  // [N, n_txt_token + n_img_token, n_head*d_head]
+            attn              = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3));                           // [n_txt_token + n_img_token, N, hidden_size]
+            auto txt_attn_out = ggml_view_3d(ctx,
+                                             attn,
+                                             attn->ne[0],
+                                             attn->ne[1],
+                                             txt->ne[1],
+                                             attn->nb[1],
+                                             attn->nb[2],
+                                             0);                                              // [n_txt_token, N, hidden_size]
+            txt_attn_out      = ggml_cont(ctx, ggml_permute(ctx, txt_attn_out, 0, 2, 1, 3));  // [N, n_txt_token, hidden_size]
+            auto img_attn_out = ggml_view_3d(ctx,
+                                             attn,
+                                             attn->ne[0],
+                                             attn->ne[1],
+                                             img->ne[1],
+                                             attn->nb[1],
+                                             attn->nb[2],
+                                             attn->nb[2] * txt->ne[1]);                       // [n_img_token, N, hidden_size]
+            img_attn_out      = ggml_cont(ctx, ggml_permute(ctx, img_attn_out, 0, 2, 1, 3));  // [N, n_img_token, hidden_size]
+
+            img_attn_out = to_out_0->forward(ctx, img_attn_out);
+            txt_attn_out = to_add_out->forward(ctx, txt_attn_out);
+
+            return {img_attn_out, txt_attn_out};
+        }
+    };
+
+    class QwenImageTransformerBlock : public GGMLBlock {
+    public:
+        QwenImageTransformerBlock(int64_t dim,
+                                  int64_t num_attention_heads,
+                                  int64_t attention_head_dim,
+                                  float eps       = 1e-6,
+                                  bool flash_attn = false) {
+            // img_mod.0 is nn.SiLU()
+            blocks["img_mod.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, 6 * dim, true));
+
+            blocks["img_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
+            blocks["img_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
+            blocks["img_mlp"]   = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim, 4, FeedForward::Activation::GELU, true));
+
+            // txt_mod.0 is nn.SiLU()
+            blocks["txt_mod.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, 6 * dim, true));
+
+            blocks["txt_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
+            blocks["txt_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
+            blocks["txt_mlp"]   = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim, 4, FeedForward::Activation::GELU));
+
+            blocks["attn"] = std::shared_ptr<GGMLBlock>(new QwenImageAttention(dim,
+                                                                               attention_head_dim,
+                                                                               num_attention_heads,
+                                                                               0,     // out_dim
+                                                                               0,     // out_context-dim
+                                                                               true,  // bias
+                                                                               true,  // out_bias
+                                                                               eps,
+                                                                               flash_attn));
+        }
+
+        virtual std::pair<ggml_tensor*, ggml_tensor*> forward(struct ggml_context* ctx,
+                                                              ggml_backend_t backend,
+                                                              struct ggml_tensor* img,
+                                                              struct ggml_tensor* txt,
+                                                              struct ggml_tensor* t_emb,
+                                                              struct ggml_tensor* pe) {
+            // img: [N, n_img_token, hidden_size]
+            // txt: [N, n_txt_token, hidden_size]
+            // pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
+            // return: ([N, n_img_token, hidden_size], [N, n_txt_token, hidden_size])
+
+            auto img_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["img_mod.1"]);
+            auto img_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"]);
+            auto img_norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"]);
+            auto img_mlp   = std::dynamic_pointer_cast<FeedForward>(blocks["img_mlp"]);
+
+            auto txt_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["txt_mod.1"]);
+            auto txt_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["txt_norm1"]);
+            auto txt_norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["txt_norm2"]);
+            auto txt_mlp   = std::dynamic_pointer_cast<FeedForward>(blocks["txt_mlp"]);
+
+            auto attn = std::dynamic_pointer_cast<QwenImageAttention>(blocks["attn"]);
+
+            auto img_mod_params    = ggml_silu(ctx, t_emb);
+            img_mod_params         = img_mod_1->forward(ctx, img_mod_params);
+            auto img_mod_param_vec = ggml_chunk(ctx, img_mod_params, 6, 0);
+
+            auto txt_mod_params    = ggml_silu(ctx, t_emb);
+            txt_mod_params         = txt_mod_1->forward(ctx, txt_mod_params);
+            auto txt_mod_param_vec = ggml_chunk(ctx, txt_mod_params, 6, 0);
+
+            auto img_normed    = img_norm1->forward(ctx, img);
+            auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1]);
+            auto img_gate1     = img_mod_param_vec[2];
+
+            auto txt_normed    = txt_norm1->forward(ctx, txt);
+            auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]);
+            auto txt_gate1     = txt_mod_param_vec[2];
+
+            auto [img_attn_output, txt_attn_output] = attn->forward(ctx, backend, img_modulated, txt_modulated, pe);
+
+            img = ggml_add(ctx, img, ggml_mul(ctx, img_attn_output, img_gate1));
+            txt = ggml_add(ctx, txt, ggml_mul(ctx, txt_attn_output, txt_gate1));
+
+            auto img_normed2    = img_norm2->forward(ctx, img);
+            auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4]);
+            auto img_gate2      = img_mod_param_vec[5];
+
+            auto txt_normed2    = txt_norm2->forward(ctx, txt);
+            auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]);
+            auto txt_gate2      = txt_mod_param_vec[5];
+
+            auto img_mlp_out = img_mlp->forward(ctx, img_modulated2);
+            auto txt_mlp_out = txt_mlp->forward(ctx, txt_modulated2);
+
+            img = ggml_add(ctx, img, ggml_mul(ctx, img_mlp_out, img_gate2));
+            txt = ggml_add(ctx, txt, ggml_mul(ctx, txt_mlp_out, txt_gate2));
+
+            return {img, txt};
+        }
+    };
+
+    struct AdaLayerNormContinuous : public GGMLBlock {
+    public:
+        AdaLayerNormContinuous(int64_t embedding_dim,
+                               int64_t conditioning_embedding_dim,
+                               bool elementwise_affine = true,
+                               float eps               = 1e-5f,
+                               bool bias               = true) {
+            blocks["norm"]   = std::shared_ptr<GGMLBlock>(new LayerNorm(conditioning_embedding_dim, eps, elementwise_affine, bias));
+            blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(conditioning_embedding_dim, embedding_dim * 2, bias));
+        }
+
+        struct ggml_tensor* forward(struct ggml_context* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* c) {
+            // x: [N, n_token, hidden_size]
+            // c: [N, hidden_size]
+            // return: [N, n_token, patch_size * patch_size * out_channels]
+
+            auto norm   = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+
+            auto emb   = linear->forward(ctx, ggml_silu(ctx, c));
+            auto mods  = ggml_chunk(ctx, emb, 2, 0);
+            auto scale = mods[0];
+            auto shift = mods[1];
+
+            x = norm->forward(ctx, x);
+            x = Flux::modulate(ctx, x, shift, scale);
+
+            return x;
+        }
+    };
+
+    struct QwenImageParams {
+        int64_t patch_size          = 2;
+        int64_t in_channels         = 64;
+        int64_t out_channels        = 16;
+        int64_t num_layers          = 60;
+        int64_t attention_head_dim  = 128;
+        int64_t num_attention_heads = 24;
+        int64_t joint_attention_dim = 3584;
+        float theta                 = 10000;
+        std::vector<int> axes_dim   = {16, 56, 56};
+        int64_t axes_dim_sum        = 128;
+        bool flash_attn             = false;
+    };
+
+    class QwenImageModel : public GGMLBlock {
+    protected:
+        QwenImageParams params;
+
+    public:
+        QwenImageModel() {}
+        QwenImageModel(QwenImageParams params)
+            : params(params) {
+            int64_t inner_dim         = params.num_attention_heads * params.attention_head_dim;
+            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
+            blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(params.joint_attention_dim, 1e-6f));
+            blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(params.in_channels, inner_dim));
+            blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(params.joint_attention_dim, inner_dim));
+
+            // blocks
+            for (int i = 0; i < params.num_layers; i++) {
+                auto block                                        = std::shared_ptr<GGMLBlock>(new QwenImageTransformerBlock(inner_dim,
+                                                                                                                             params.num_attention_heads,
+                                                                                                                             params.attention_head_dim,
+                                                                                                                             1e-6f,
+                                                                                                                             params.flash_attn));
+                blocks["transformer_blocks." + std::to_string(i)] = block;
+            }
+
+            blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new AdaLayerNormContinuous(inner_dim, inner_dim, false, 1e-6f));
+            blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
+        }
+
+        struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
+                                              struct ggml_tensor* x) {
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+
+            int pad_h = (params.patch_size - H % params.patch_size) % params.patch_size;
+            int pad_w = (params.patch_size - W % params.patch_size) % params.patch_size;
+            x         = ggml_pad(ctx, x, pad_w, pad_h, 0, 0);  // [N, C, H + pad_h, W + pad_w]
+            return x;
+        }
+
+        struct ggml_tensor* patchify(struct ggml_context* ctx,
+                                     struct ggml_tensor* x) {
+            // x: [N, C, H, W]
+            // return: [N, h*w, C * patch_size * patch_size]
+            int64_t N = x->ne[3];
+            int64_t C = x->ne[2];
+            int64_t H = x->ne[1];
+            int64_t W = x->ne[0];
+            int64_t p = params.patch_size;
+            int64_t h = H / params.patch_size;
+            int64_t w = W / params.patch_size;
+
+            GGML_ASSERT(h * p == H && w * p == W);
+
+            x = ggml_reshape_4d(ctx, x, p, w, p, h * C * N);       // [N*C*h, p, w, p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*h, w, p, p]
+            x = ggml_reshape_4d(ctx, x, p * p, w * h, C, N);       // [N, C, h*w, p*p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, h*w, C, p*p]
+            x = ggml_reshape_3d(ctx, x, p * p * C, w * h, N);      // [N, h*w, C*p*p]
+            return x;
+        }
+
+        struct ggml_tensor* process_img(struct ggml_context* ctx,
+                                        struct ggml_tensor* x) {
+            x = pad_to_patch_size(ctx, x);
+            x = patchify(ctx, x);
+            return x;
+        }
+
+        struct ggml_tensor* unpatchify(struct ggml_context* ctx,
+                                       struct ggml_tensor* x,
+                                       int64_t h,
+                                       int64_t w) {
+            // x: [N, h*w, C*patch_size*patch_size]
+            // return: [N, C, H, W]
+            int64_t N = x->ne[2];
+            int64_t C = x->ne[0] / params.patch_size / params.patch_size;
+            int64_t H = h * params.patch_size;
+            int64_t W = w * params.patch_size;
+            int64_t p = params.patch_size;
+
+            GGML_ASSERT(C * p * p == x->ne[0]);
+
+            x = ggml_reshape_4d(ctx, x, p * p, C, w * h, N);       // [N, h*w, C, p*p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, C, h*w, p*p]
+            x = ggml_reshape_4d(ctx, x, p, p, w, h * C * N);       // [N*C*h, w, p, p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*h, p, w, p]
+            x = ggml_reshape_4d(ctx, x, W, H, C, N);               // [N, C, h*p, w*p]
+
+            return x;
+        }
+
+        struct ggml_tensor* forward_orig(struct ggml_context* ctx,
+                                         ggml_backend_t backend,
+                                         struct ggml_tensor* x,
+                                         struct ggml_tensor* timestep,
+                                         struct ggml_tensor* context,
+                                         struct ggml_tensor* pe) {
+            auto time_text_embed = std::dynamic_pointer_cast<QwenTimestepProjEmbeddings>(blocks["time_text_embed"]);
+            auto txt_norm        = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
+            auto img_in          = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
+            auto txt_in          = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
+            auto norm_out        = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
+            auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
+
+            auto t_emb = time_text_embed->forward(ctx, timestep);
+            auto img   = img_in->forward(ctx, x);
+            auto txt   = txt_norm->forward(ctx, context);
+            txt        = txt_in->forward(ctx, txt);
+
+            for (int i = 0; i < params.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
+
+                auto result = block->forward(ctx, backend, img, txt, t_emb, pe);
+                img         = result.first;
+                txt         = result.second;
+            }
+
+            img = norm_out->forward(ctx, img, t_emb);
+            img = proj_out->forward(ctx, img);
+
+            return img;
+        }
+
+        struct ggml_tensor* forward(struct ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* timestep,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* pe,
+                                    std::vector<ggml_tensor*> ref_latents = {}) {
+            // Forward pass of DiT.
+            // x: [N, C, H, W]
+            // timestep: [N,]
+            // context: [N, L, D]
+            // pe: [L, d_head/2, 2, 2]
+            // return: [N, C, H, W]
+
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t C = x->ne[2];
+            int64_t N = x->ne[3];
+
+            auto img            = process_img(ctx, x);
+            uint64_t img_tokens = img->ne[1];
+
+            if (ref_latents.size() > 0) {
+                for (ggml_tensor* ref : ref_latents) {
+                    ref = process_img(ctx, ref);
+                    img = ggml_concat(ctx, img, ref, 1);
+                }
+            }
+
+            int64_t h_len = ((H + (params.patch_size / 2)) / params.patch_size);
+            int64_t w_len = ((W + (params.patch_size / 2)) / params.patch_size);
+
+            auto out = forward_orig(ctx, backend, img, timestep, context, pe);  // [N, h_len*w_len, ph*pw*C]
+
+            if (out->ne[1] > img_tokens) {
+                out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3));  // [num_tokens, N, C * patch_size * patch_size]
+                out = ggml_view_3d(ctx, out, out->ne[0], out->ne[1], img_tokens, out->nb[1], out->nb[2], 0);
+                out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3));  // [N, h*w, C * patch_size * patch_size]
+            }
+
+            out = unpatchify(ctx, out, h_len, w_len);  // [N, C, H + pad_h, W + pad_w]
+
+            // slice
+            out = ggml_slice(ctx, out, 1, 0, H);  // [N, C, H, W + pad_w]
+            out = ggml_slice(ctx, out, 0, 0, W);  // [N, C, H, W]
+
+            return out;
+        }
+    };
+
+    struct QwenImageRunner : public GGMLRunner {
+    public:
+        QwenImageParams qwen_image_params;
+        QwenImageModel qwen_image;
+        std::vector<float> pe_vec;
+        SDVersion version;
+
+        QwenImageRunner(ggml_backend_t backend,
+                        bool offload_params_to_cpu,
+                        const String2GGMLType& tensor_types = {},
+                        const std::string prefix            = "",
+                        SDVersion version                   = VERSION_QWEN_IMAGE,
+                        bool flash_attn                     = false)
+            : GGMLRunner(backend, offload_params_to_cpu) {
+            qwen_image_params.flash_attn = flash_attn;
+            qwen_image_params.num_layers = 0;
+            for (auto pair : tensor_types) {
+                std::string tensor_name = pair.first;
+                if (tensor_name.find(prefix) == std::string::npos)
+                    continue;
+                size_t pos = tensor_name.find("transformer_blocks.");
+                if (pos != std::string::npos) {
+                    tensor_name = tensor_name.substr(pos);  // remove prefix
+                    auto items  = split_string(tensor_name, '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > qwen_image_params.num_layers) {
+                            qwen_image_params.num_layers = block_index + 1;
+                        }
+                    }
+                    continue;
+                }
+            }
+            LOG_ERROR("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers);
+            qwen_image                   = QwenImageModel(qwen_image_params);
+            qwen_image.init(params_ctx, tensor_types, prefix);
+        }
+
+        std::string get_desc() {
+            return "qwen_image";
+        }
+
+        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+            qwen_image.get_param_tensors(tensors, prefix);
+        }
+
+        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                        struct ggml_tensor* timesteps,
+                                        struct ggml_tensor* context,
+                                        std::vector<ggml_tensor*> ref_latents = {},
+                                        bool increase_ref_index               = false) {
+            GGML_ASSERT(x->ne[3] == 1);
+            struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, QWEN_IMAGE_GRAPH_SIZE, false);
+
+            x         = to_backend(x);
+            context   = to_backend(context);
+            timesteps = to_backend(timesteps);
+
+            for (int i = 0; i < ref_latents.size(); i++) {
+                ref_latents[i] = to_backend(ref_latents[i]);
+            }
+
+            pe_vec      = Rope::gen_qwen_image_pe(x->ne[1],
+                                                  x->ne[0],
+                                                  qwen_image_params.patch_size,
+                                                  x->ne[3],
+                                                  context->ne[1],
+                                                  ref_latents,
+                                                  increase_ref_index,
+                                                  qwen_image_params.theta,
+                                                  qwen_image_params.axes_dim);
+            int pos_len = pe_vec.size() / qwen_image_params.axes_dim_sum / 2;
+            // LOG_DEBUG("pos_len %d", pos_len);
+            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, qwen_image_params.axes_dim_sum / 2, pos_len);
+            // pe->data = pe_vec.data();
+            // print_ggml_tensor(pe, true, "pe");
+            // pe->data = NULL;
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            struct ggml_tensor* out = qwen_image.forward(compute_ctx,
+                                                         runtime_backend,
+                                                         x,
+                                                         timesteps,
+                                                         context,
+                                                         pe,
+                                                         ref_latents);
+
+            ggml_build_forward_expand(gf, out);
+
+            return gf;
+        }
+
+        void compute(int n_threads,
+                     struct ggml_tensor* x,
+                     struct ggml_tensor* timesteps,
+                     struct ggml_tensor* context,
+                     std::vector<ggml_tensor*> ref_latents = {},
+                     bool increase_ref_index               = false,
+                     struct ggml_tensor** output           = NULL,
+                     struct ggml_context* output_ctx       = NULL) {
+            // x: [N, in_channels, h, w]
+            // timesteps: [N, ]
+            // context: [N, max_position, hidden_size]
+            auto get_graph = [&]() -> struct ggml_cgraph* {
+                return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
+            };
+
+            GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        }
+
+        void test() {
+            struct ggml_init_params params;
+            params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
+            params.mem_buffer = NULL;
+            params.no_alloc   = false;
+
+            struct ggml_context* work_ctx = ggml_init(params);
+            GGML_ASSERT(work_ctx != NULL);
+
+            {
+                // auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 16, 16, 16, 1);
+                // ggml_set_f32(x, 0.01f);
+                auto x = load_tensor_from_file(work_ctx, "./qwen_image_x.bin");
+                print_ggml_tensor(x);
+
+                std::vector<float> timesteps_vec(1, 1000.f);
+                auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
+
+                // auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 3584, 256, 1);
+                // ggml_set_f32(context, 0.01f);
+                auto context = load_tensor_from_file(work_ctx, "./qwen_image_context.bin");
+                print_ggml_tensor(context);
+
+                struct ggml_tensor* out = NULL;
+
+                int t0 = ggml_time_ms();
+                compute(8, x, timesteps, context, {}, false, &out, work_ctx);
+                int t1 = ggml_time_ms();
+
+                print_ggml_tensor(out);
+                LOG_DEBUG("qwen_image test done in %dms", t1 - t0);
+            }
+        }
+
+        static void load_from_file_and_test(const std::string& file_path) {
+            // cuda q8: pass
+            // cuda q8 fa: nan
+            // ggml_backend_t backend    = ggml_backend_cuda_init(0);
+            ggml_backend_t backend    = ggml_backend_cpu_init();
+            ggml_type model_data_type = GGML_TYPE_Q8_0;
+
+            ModelLoader model_loader;
+            if (!model_loader.init_from_file(file_path, "model.diffusion_model.")) {
+                LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
+                return;
+            }
+
+            auto tensor_types = model_loader.tensor_storages_types;
+            for (auto& item : tensor_types) {
+                // LOG_DEBUG("%s %u", item.first.c_str(), item.second);
+                if (ends_with(item.first, "weight")) {
+                    item.second = model_data_type;
+                }
+            }
+
+            std::shared_ptr<QwenImageRunner> qwen_image = std::shared_ptr<QwenImageRunner>(new QwenImageRunner(backend,
+                                                                                                               false,
+                                                                                                               tensor_types,
+                                                                                                               "model.diffusion_model",
+                                                                                                               VERSION_QWEN_IMAGE,
+                                                                                                               true));
+
+            qwen_image->alloc_params_buffer();
+            std::map<std::string, ggml_tensor*> tensors;
+            qwen_image->get_param_tensors(tensors, "model.diffusion_model");
+
+            bool success = model_loader.load_tensors(tensors);
+
+            if (!success) {
+                LOG_ERROR("load tensors from model loader failed");
+                return;
+            }
+
+            LOG_INFO("qwen_image model loaded");
+            qwen_image->test();
+        }
+    };
+
+}  // namespace name
+
+#endif  // __QWEN_IMAGE_HPP__

rope.hpp

@@ -4,9 +4,9 @@
 #include <vector>
 #include "ggml_extend.hpp"
 
-struct Rope {
+namespace Rope {
     template <class T>
-    static std::vector<T> linspace(T start, T end, int num) {
+    __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
         if (num == 1) {
             result[0] = start;
@@ -19,7 +19,7 @@ struct Rope {
         return result;
     }
 
-    static std::vector<std::vector<float>> transpose(const std::vector<std::vector<float>>& mat) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> transpose(const std::vector<std::vector<float>>& mat) {
         int rows = mat.size();
         int cols = mat[0].size();
         std::vector<std::vector<float>> transposed(cols, std::vector<float>(rows));
@@ -31,7 +31,7 @@ struct Rope {
         return transposed;
     }
 
-    static std::vector<float> flatten(const std::vector<std::vector<float>>& vec) {
+    __STATIC_INLINE__ std::vector<float> flatten(const std::vector<std::vector<float>>& vec) {
         std::vector<float> flat_vec;
         for (const auto& sub_vec : vec) {
             flat_vec.insert(flat_vec.end(), sub_vec.begin(), sub_vec.end());
@@ -39,7 +39,7 @@ struct Rope {
         return flat_vec;
     }
 
-    static std::vector<std::vector<float>> rope(const std::vector<float>& pos, int dim, int theta) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos, int dim, int theta) {
         assert(dim % 2 == 0);
         int half_dim = dim / 2;
 
@@ -72,11 +72,11 @@ struct Rope {
     }
 
     // Generate IDs for image patches and text
-    static std::vector<std::vector<float>> gen_txt_ids(int bs, int context_len) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_txt_ids(int bs, int context_len) {
         return std::vector<std::vector<float>>(bs * context_len, std::vector<float>(3, 0.0));
     }
 
-    static std::vector<std::vector<float>> gen_img_ids(int h, int w, int patch_size, int bs, int index = 0, int h_offset = 0, int w_offset = 0) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_img_ids(int h, int w, int patch_size, int bs, int index = 0, int h_offset = 0, int w_offset = 0) {
         int h_len = (h + (patch_size / 2)) / patch_size;
         int w_len = (w + (patch_size / 2)) / patch_size;
 
@@ -102,9 +102,9 @@ struct Rope {
         return img_ids_repeated;
     }
 
-    static std::vector<std::vector<float>> concat_ids(const std::vector<std::vector<float>>& a,
-                                                      const std::vector<std::vector<float>>& b,
-                                                      int bs) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> concat_ids(const std::vector<std::vector<float>>& a,
+                                                                 const std::vector<std::vector<float>>& b,
+                                                                 int bs) {
         size_t a_len = a.size() / bs;
         size_t b_len = b.size() / bs;
         std::vector<std::vector<float>> ids(a.size() + b.size(), std::vector<float>(3));
@@ -119,10 +119,10 @@ struct Rope {
         return ids;
     }
 
-    static std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
-                                       int bs,
-                                       int theta,
-                                       const std::vector<int>& axes_dim) {
+    __STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
+                                                  int bs,
+                                                  int theta,
+                                                  const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> trans_ids = transpose(ids);
         size_t pos_len                            = ids.size() / bs;
         int num_axes                              = axes_dim.size();
@@ -151,17 +151,11 @@ struct Rope {
         return flatten(emb);
     }
 
-    static std::vector<std::vector<float>> gen_flux_ids(int h,
-                                                        int w,
-                                                        int patch_size,
-                                                        int bs,
-                                                        int context_len,
-                                                        std::vector<ggml_tensor*> ref_latents,
-                                                        bool increase_ref_index) {
-        auto txt_ids = gen_txt_ids(bs, context_len);
-        auto img_ids = gen_img_ids(h, w, patch_size, bs);
-
-        auto ids               = concat_ids(txt_ids, img_ids, bs);
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
+                                                                   int bs,
+                                                                   const std::vector<ggml_tensor*>& ref_latents,
+                                                                   bool increase_ref_index) {
+        std::vector<std::vector<float>> ids;
         uint64_t curr_h_offset = 0;
         uint64_t curr_w_offset = 0;
         int index              = 1;
@@ -189,30 +183,88 @@ struct Rope {
         return ids;
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_ids(int h,
+                                                                   int w,
+                                                                   int patch_size,
+                                                                   int bs,
+                                                                   int context_len,
+                                                                   const std::vector<ggml_tensor*>& ref_latents,
+                                                                   bool increase_ref_index) {
+        auto txt_ids = gen_txt_ids(bs, context_len);
+        auto img_ids = gen_img_ids(h, w, patch_size, bs);
+
+        auto ids = concat_ids(txt_ids, img_ids, bs);
+        if (ref_latents.size() > 0) {
+            auto refs_ids = gen_refs_ids(patch_size, bs, ref_latents, increase_ref_index);
+            ids           = concat_ids(ids, refs_ids, bs);
+        }
+        return ids;
+    }
+
     // Generate flux positional embeddings
-    static std::vector<float> gen_flux_pe(int h,
-                                          int w,
-                                          int patch_size,
-                                          int bs,
-                                          int context_len,
-                                          std::vector<ggml_tensor*> ref_latents,
-                                          bool increase_ref_index,
-                                          int theta,
-                                          const std::vector<int>& axes_dim) {
+    __STATIC_INLINE__ std::vector<float> gen_flux_pe(int h,
+                                                     int w,
+                                                     int patch_size,
+                                                     int bs,
+                                                     int context_len,
+                                                     const std::vector<ggml_tensor*>& ref_latents,
+                                                     bool increase_ref_index,
+                                                     int theta,
+                                                     const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_flux_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
         return embed_nd(ids, bs, theta, axes_dim);
     }
 
-    static std::vector<std::vector<float>> gen_vid_ids(int t,
-                                                       int h,
-                                                       int w,
-                                                       int pt,
-                                                       int ph,
-                                                       int pw,
-                                                       int bs,
-                                                       int t_offset = 0,
-                                                       int h_offset = 0,
-                                                       int w_offset = 0) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
+                                                                         int w,
+                                                                         int patch_size,
+                                                                         int bs,
+                                                                         int context_len,
+                                                                         const std::vector<ggml_tensor*>& ref_latents,
+                                                                         bool increase_ref_index) {
+        int h_len        = (h + (patch_size / 2)) / patch_size;
+        int w_len        = (w + (patch_size / 2)) / patch_size;
+        int txt_id_start = std::max(h_len, w_len);
+        auto txt_ids     = linspace<float>(txt_id_start, context_len + txt_id_start, context_len);
+        std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < txt_ids.size(); ++j) {
+                txt_ids_repeated[i * txt_ids.size() + j] = {txt_ids[j], txt_ids[j], txt_ids[j]};
+            }
+        }
+        auto img_ids = gen_img_ids(h, w, patch_size, bs);
+        auto ids     = concat_ids(txt_ids_repeated, img_ids, bs);
+        if (ref_latents.size() > 0) {
+            auto refs_ids = gen_refs_ids(patch_size, bs, ref_latents, increase_ref_index);
+            ids           = concat_ids(ids, refs_ids, bs);
+        }
+        return ids;
+    }
+
+    // Generate qwen_image positional embeddings
+    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
+                                                           int w,
+                                                           int patch_size,
+                                                           int bs,
+                                                           int context_len,
+                                                           const std::vector<ggml_tensor*>& ref_latents,
+                                                           bool increase_ref_index,
+                                                           int theta,
+                                                           const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
+        return embed_nd(ids, bs, theta, axes_dim);
+    }
+
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
+                                                                  int h,
+                                                                  int w,
+                                                                  int pt,
+                                                                  int ph,
+                                                                  int pw,
+                                                                  int bs,
+                                                                  int t_offset = 0,
+                                                                  int h_offset = 0,
+                                                                  int w_offset = 0) {
         int t_len = (t + (pt / 2)) / pt;
         int h_len = (h + (ph / 2)) / ph;
         int w_len = (w + (pw / 2)) / pw;
@@ -244,18 +296,115 @@ struct Rope {
     }
 
     // Generate wan positional embeddings
-    static std::vector<float> gen_wan_pe(int t,
-                                         int h,
-                                         int w,
-                                         int pt,
-                                         int ph,
-                                         int pw,
-                                         int bs,
-                                         int theta,
-                                         const std::vector<int>& axes_dim) {
+    __STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
+                                                    int h,
+                                                    int w,
+                                                    int pt,
+                                                    int ph,
+                                                    int pw,
+                                                    int bs,
+                                                    int theta,
+                                                    const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs);
         return embed_nd(ids, bs, theta, axes_dim);
     }
-};  // struct Rope
+
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
+                                                                      int grid_w,
+                                                                      int merge_size,
+                                                                      const std::vector<int>& window_index) {
+        std::vector<std::vector<float>> ids(grid_h * grid_w, std::vector<float>(2, 0.0));
+        int index = 0;
+        for (int ih = 0; ih < grid_h; ih += merge_size) {
+            for (int iw = 0; iw < grid_w; iw += merge_size) {
+                for (int iy = 0; iy < merge_size; iy++) {
+                    for (int ix = 0; ix < merge_size; ix++) {
+                        int inverse_index = window_index[index / (merge_size * merge_size)];
+                        int i             = inverse_index * (merge_size * merge_size) + index % (merge_size * merge_size);
+
+                        GGML_ASSERT(i < grid_h * grid_w);
+
+                        ids[i][0] = ih + iy;
+                        ids[i][1] = iw + ix;
+                        index++;
+                    }
+                }
+            }
+        }
+        return ids;
+    }
+
+    // Generate qwen2vl positional embeddings
+    __STATIC_INLINE__ std::vector<float> gen_qwen2vl_pe(int grid_h,
+                                                        int grid_w,
+                                                        int merge_size,
+                                                        const std::vector<int>& window_index,
+                                                        int theta,
+                                                        const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_qwen2vl_ids(grid_h, grid_w, merge_size, window_index);
+        return embed_nd(ids, 1, theta, axes_dim);
+    }
+
+    __STATIC_INLINE__ struct ggml_tensor* apply_rope(struct ggml_context* ctx,
+                                                     struct ggml_tensor* x,
+                                                     struct ggml_tensor* pe,
+                                                     bool rope_interleaved = true) {
+        // x: [N, L, n_head, d_head]
+        // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
+        int64_t d_head = x->ne[0];
+        int64_t n_head = x->ne[1];
+        int64_t L      = x->ne[2];
+        int64_t N      = x->ne[3];
+        x              = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, n_head, L, d_head]
+        if (rope_interleaved) {
+            x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N);  // [N * n_head, L, d_head/2, 2]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2));       // [2, N * n_head, L, d_head/2]
+        } else {
+            x = ggml_reshape_4d(ctx, x, d_head / 2, 2, L, n_head * N);   // [N * n_head, L, 2, d_head/2]
+            x = ggml_cont(ctx, ggml_torch_permute(ctx, x, 0, 2, 3, 1));  // [2, N * n_head, L, d_head/2]
+        }
+
+        int64_t offset = x->nb[2] * x->ne[2];
+        auto x_0       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0);  // [N * n_head, L, d_head/2]
+        auto x_1       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1);  // [N * n_head, L, d_head/2]
+        x_0            = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]);                    // [N * n_head, L, d_head/2, 1]
+        x_1            = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]);                    // [N * n_head, L, d_head/2, 1]
+        auto temp_x    = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]);
+        x_0            = ggml_repeat(ctx, x_0, temp_x);  // [N * n_head, L, d_head/2, 2]
+        x_1            = ggml_repeat(ctx, x_1, temp_x);  // [N * n_head, L, d_head/2, 2]
+
+        pe        = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2));  // [2, L, d_head/2, 2]
+        offset    = pe->nb[2] * pe->ne[2];
+        auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0);  // [L, d_head/2, 2]
+        auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1);  // [L, d_head/2, 2]
+
+        auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1));  // [N * n_head, L, d_head/2, 2]
+        if (!rope_interleaved) {
+            x_out = ggml_cont(ctx, ggml_permute(ctx, x_out, 1, 0, 2, 3));  // [N * n_head, L, x, d_head/2]
+        }
+        x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N);  // [N*n_head, L, d_head]
+        return x_out;
+    }
+
+    __STATIC_INLINE__ struct ggml_tensor* attention(struct ggml_context* ctx,
+                                                    ggml_backend_t backend,
+                                                    struct ggml_tensor* q,
+                                                    struct ggml_tensor* k,
+                                                    struct ggml_tensor* v,
+                                                    struct ggml_tensor* pe,
+                                                    struct ggml_tensor* mask,
+                                                    bool flash_attn,
+                                                    float kv_scale        = 1.0f,
+                                                    bool rope_interleaved = true) {
+        // q,k,v: [N, L, n_head, d_head]
+        // pe: [L, d_head/2, 2, 2]
+        // return: [N, L, n_head*d_head]
+        q = apply_rope(ctx, q, pe, rope_interleaved);  // [N*n_head, L, d_head]
+        k = apply_rope(ctx, k, pe, rope_interleaved);  // [N*n_head, L, d_head]
+
+        auto x = ggml_nn_attention_ext(ctx, backend, q, k, v, v->ne[1], mask, false, true, flash_attn, kv_scale);  // [N, L, n_head*d_head]
+        return x;
+    }
+};  // namespace Rope
 
 #endif  // __ROPE_HPP__

stable-diffusion.h

@@ -58,10 +58,22 @@ enum scheduler_t {
     EXPONENTIAL,
     AYS,
     GITS,
+    SGM_UNIFORM,
+    SIMPLE,
     SMOOTHSTEP,
     SCHEDULE_COUNT
 };
 
+enum prediction_t {
+    DEFAULT_PRED,
+    EPS_PRED,
+    V_PRED,
+    EDM_V_PRED,
+    SD3_FLOW_PRED,
+    FLUX_FLOW_PRED,
+    PREDICTION_COUNT
+};
+
 // same as enum ggml_type
 enum sd_type_t {
     SD_TYPE_F32  = 0,
@@ -114,12 +126,23 @@ enum sd_log_level_t {
     SD_LOG_ERROR
 };
 
+typedef struct {
+    bool enabled;
+    int tile_size_x;
+    int tile_size_y;
+    float target_overlap;
+    float rel_size_x;
+    float rel_size_y;
+} sd_tiling_params_t;
+
 typedef struct {
     const char* model_path;
     const char* clip_l_path;
     const char* clip_g_path;
     const char* clip_vision_path;
     const char* t5xxl_path;
+    const char* qwen2vl_path;
+    const char* qwen2vl_vision_path;
     const char* diffusion_model_path;
     const char* high_noise_diffusion_model_path;
     const char* vae_path;
@@ -127,13 +150,13 @@ typedef struct {
     const char* control_net_path;
     const char* lora_model_dir;
     const char* embedding_dir;
-    const char* stacked_id_embed_dir;
+    const char* photo_maker_path;
     bool vae_decode_only;
-    bool vae_tiling;
     bool free_params_immediately;
     int n_threads;
     enum sd_type_t wtype;
     enum rng_type_t rng_type;
+    enum prediction_t prediction;
     bool offload_params_to_cpu;
     bool keep_clip_on_cpu;
     bool keep_control_net_on_cpu;
@@ -175,8 +198,16 @@ typedef struct {
     enum sample_method_t sample_method;
     int sample_steps;
     float eta;
+    int shifted_timestep;
 } sd_sample_params_t;
 
+typedef struct {
+    sd_image_t* id_images;
+    int id_images_count;
+    const char* id_embed_path;
+    float style_strength;
+} sd_pm_params_t;  // photo maker
+
 typedef struct {
     const char* prompt;
     const char* negative_prompt;
@@ -194,9 +225,8 @@ typedef struct {
     int batch_count;
     sd_image_t control_image;
     float control_strength;
-    float style_strength;
-    bool normalize_input;
-    const char* input_id_images_path;
+    sd_pm_params_t pm_params;
+    sd_tiling_params_t vae_tiling_params;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -205,6 +235,8 @@ typedef struct {
     int clip_skip;
     sd_image_t init_image;
     sd_image_t end_image;
+    sd_image_t* control_frames;
+    int control_frames_size;
     int width;
     int height;
     sd_sample_params_t sample_params;
@@ -213,6 +245,7 @@ typedef struct {
     float strength;
     int64_t seed;
     int video_frames;
+    float vace_strength;
 } sd_vid_gen_params_t;
 
 typedef struct sd_ctx_t sd_ctx_t;
@@ -233,6 +266,8 @@ SD_API const char* sd_sample_method_name(enum sample_method_t sample_method);
 SD_API enum sample_method_t str_to_sample_method(const char* str);
 SD_API const char* sd_schedule_name(enum scheduler_t scheduler);
 SD_API enum scheduler_t str_to_schedule(const char* str);
+SD_API const char* sd_prediction_name(enum prediction_t prediction);
+SD_API enum prediction_t str_to_prediction(const char* str);
 
 SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params);
 SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
@@ -263,20 +298,20 @@ SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,
                           sd_image_t input_image,
                           uint32_t upscale_factor);
 
+SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
+
 SD_API bool convert(const char* input_path,
                     const char* vae_path,
                     const char* output_path,
                     enum sd_type_t output_type,
                     const char* tensor_type_rules);
 
-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);
+SD_API bool preprocess_canny(sd_image_t image,
+                             float high_threshold,
+                             float low_threshold,
+                             float weak,
+                             float strong,
+                             bool inverse);
 
 #ifdef __cplusplus
 }

t5.hpp

@@ -504,7 +504,9 @@ public:
     T5DenseGatedActDense(int64_t model_dim, int64_t ff_dim) {
         blocks["wi_0"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, ff_dim, false));
         blocks["wi_1"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, ff_dim, false));
-        blocks["wo"]   = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false));
+        float scale    = 1.f / 32.f;
+        // The purpose of the scale here is to prevent NaN issues on some backends(CUDA, ...).
+        blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false, false, false, scale));
     }
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {

tae.hpp

@@ -222,7 +222,7 @@ struct TinyAutoEncoder : public GGMLRunner {
         return "taesd";
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading taesd from '%s', decode_only = %s", file_path.c_str(), decode_only ? "true" : "false");
         alloc_params_buffer();
         std::map<std::string, ggml_tensor*> taesd_tensors;
@@ -238,7 +238,7 @@ struct TinyAutoEncoder : public GGMLRunner {
             return false;
         }
 
-        bool success = model_loader.load_tensors(taesd_tensors, ignore_tensors);
+        bool success = model_loader.load_tensors(taesd_tensors, ignore_tensors, n_threads);
 
         if (!success) {
             LOG_ERROR("load tae tensors from model loader failed");

tokenize_util.cpp

@@ -0,0 +1,985 @@
+#include <algorithm>
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include "tokenize_util.h"
+
+bool is_number(char32_t ch) {
+    return (ch >= U'0' && ch <= U'9');
+}
+
+bool is_letter(char32_t ch) {
+    static const struct { char32_t start, end; } ranges[] = {
+        {0x41, 0x5A},
+        {0x61, 0x7A},
+        {0xAA, 0xAA},
+        {0xB5, 0xB5},
+        {0xBA, 0xBA},
+        {0xC0, 0xD6},
+        {0xD8, 0xF6},
+        {0xF8, 0x2C1},
+        {0x2C6, 0x2D1},
+        {0x2E0, 0x2E4},
+        {0x2EC, 0x2EC},
+        {0x2EE, 0x2EE},
+        {0x370, 0x374},
+        {0x376, 0x377},
+        {0x37A, 0x37D},
+        {0x37F, 0x37F},
+        {0x386, 0x386},
+        {0x388, 0x38A},
+        {0x38C, 0x38C},
+        {0x38E, 0x3A1},
+        {0x3A3, 0x3F5},
+        {0x3F7, 0x481},
+        {0x48A, 0x52F},
+        {0x531, 0x556},
+        {0x559, 0x559},
+        {0x560, 0x588},
+        {0x5D0, 0x5EA},
+        {0x5EF, 0x5F2},
+        {0x620, 0x64A},
+        {0x66E, 0x66F},
+        {0x671, 0x6D3},
+        {0x6D5, 0x6D5},
+        {0x6E5, 0x6E6},
+        {0x6EE, 0x6EF},
+        {0x6FA, 0x6FC},
+        {0x6FF, 0x6FF},
+        {0x710, 0x710},
+        {0x712, 0x72F},
+        {0x74D, 0x7A5},
+        {0x7B1, 0x7B1},
+        {0x7CA, 0x7EA},
+        {0x7F4, 0x7F5},
+        {0x7FA, 0x7FA},
+        {0x800, 0x815},
+        {0x81A, 0x81A},
+        {0x824, 0x824},
+        {0x828, 0x828},
+        {0x840, 0x858},
+        {0x860, 0x86A},
+        {0x870, 0x887},
+        {0x889, 0x88F},
+        {0x8A0, 0x8C9},
+        {0x904, 0x939},
+        {0x93D, 0x93D},
+        {0x950, 0x950},
+        {0x958, 0x961},
+        {0x971, 0x980},
+        {0x985, 0x98C},
+        {0x98F, 0x990},
+        {0x993, 0x9A8},
+        {0x9AA, 0x9B0},
+        {0x9B2, 0x9B2},
+        {0x9B6, 0x9B9},
+        {0x9BD, 0x9BD},
+        {0x9CE, 0x9CE},
+        {0x9DC, 0x9DD},
+        {0x9DF, 0x9E1},
+        {0x9F0, 0x9F1},
+        {0x9FC, 0x9FC},
+        {0xA05, 0xA0A},
+        {0xA0F, 0xA10},
+        {0xA13, 0xA28},
+        {0xA2A, 0xA30},
+        {0xA32, 0xA33},
+        {0xA35, 0xA36},
+        {0xA38, 0xA39},
+        {0xA59, 0xA5C},
+        {0xA5E, 0xA5E},
+        {0xA72, 0xA74},
+        {0xA85, 0xA8D},
+        {0xA8F, 0xA91},
+        {0xA93, 0xAA8},
+        {0xAAA, 0xAB0},
+        {0xAB2, 0xAB3},
+        {0xAB5, 0xAB9},
+        {0xABD, 0xABD},
+        {0xAD0, 0xAD0},
+        {0xAE0, 0xAE1},
+        {0xAF9, 0xAF9},
+        {0xB05, 0xB0C},
+        {0xB0F, 0xB10},
+        {0xB13, 0xB28},
+        {0xB2A, 0xB30},
+        {0xB32, 0xB33},
+        {0xB35, 0xB39},
+        {0xB3D, 0xB3D},
+        {0xB5C, 0xB5D},
+        {0xB5F, 0xB61},
+        {0xB71, 0xB71},
+        {0xB83, 0xB83},
+        {0xB85, 0xB8A},
+        {0xB8E, 0xB90},
+        {0xB92, 0xB95},
+        {0xB99, 0xB9A},
+        {0xB9C, 0xB9C},
+        {0xB9E, 0xB9F},
+        {0xBA3, 0xBA4},
+        {0xBA8, 0xBAA},
+        {0xBAE, 0xBB9},
+        {0xBD0, 0xBD0},
+        {0xC05, 0xC0C},
+        {0xC0E, 0xC10},
+        {0xC12, 0xC28},
+        {0xC2A, 0xC39},
+        {0xC3D, 0xC3D},
+        {0xC58, 0xC5A},
+        {0xC5C, 0xC5D},
+        {0xC60, 0xC61},
+        {0xC80, 0xC80},
+        {0xC85, 0xC8C},
+        {0xC8E, 0xC90},
+        {0xC92, 0xCA8},
+        {0xCAA, 0xCB3},
+        {0xCB5, 0xCB9},
+        {0xCBD, 0xCBD},
+        {0xCDC, 0xCDE},
+        {0xCE0, 0xCE1},
+        {0xCF1, 0xCF2},
+        {0xD04, 0xD0C},
+        {0xD0E, 0xD10},
+        {0xD12, 0xD3A},
+        {0xD3D, 0xD3D},
+        {0xD4E, 0xD4E},
+        {0xD54, 0xD56},
+        {0xD5F, 0xD61},
+        {0xD7A, 0xD7F},
+        {0xD85, 0xD96},
+        {0xD9A, 0xDB1},
+        {0xDB3, 0xDBB},
+        {0xDBD, 0xDBD},
+        {0xDC0, 0xDC6},
+        {0xE01, 0xE30},
+        {0xE32, 0xE33},
+        {0xE40, 0xE46},
+        {0xE81, 0xE82},
+        {0xE84, 0xE84},
+        {0xE86, 0xE8A},
+        {0xE8C, 0xEA3},
+        {0xEA5, 0xEA5},
+        {0xEA7, 0xEB0},
+        {0xEB2, 0xEB3},
+        {0xEBD, 0xEBD},
+        {0xEC0, 0xEC4},
+        {0xEC6, 0xEC6},
+        {0xEDC, 0xEDF},
+        {0xF00, 0xF00},
+        {0xF40, 0xF47},
+        {0xF49, 0xF6C},
+        {0xF88, 0xF8C},
+        {0x1000, 0x102A},
+        {0x103F, 0x103F},
+        {0x1050, 0x1055},
+        {0x105A, 0x105D},
+        {0x1061, 0x1061},
+        {0x1065, 0x1066},
+        {0x106E, 0x1070},
+        {0x1075, 0x1081},
+        {0x108E, 0x108E},
+        {0x10A0, 0x10C5},
+        {0x10C7, 0x10C7},
+        {0x10CD, 0x10CD},
+        {0x10D0, 0x10FA},
+        {0x10FC, 0x1248},
+        {0x124A, 0x124D},
+        {0x1250, 0x1256},
+        {0x1258, 0x1258},
+        {0x125A, 0x125D},
+        {0x1260, 0x1288},
+        {0x128A, 0x128D},
+        {0x1290, 0x12B0},
+        {0x12B2, 0x12B5},
+        {0x12B8, 0x12BE},
+        {0x12C0, 0x12C0},
+        {0x12C2, 0x12C5},
+        {0x12C8, 0x12D6},
+        {0x12D8, 0x1310},
+        {0x1312, 0x1315},
+        {0x1318, 0x135A},
+        {0x1380, 0x138F},
+        {0x13A0, 0x13F5},
+        {0x13F8, 0x13FD},
+        {0x1401, 0x166C},
+        {0x166F, 0x167F},
+        {0x1681, 0x169A},
+        {0x16A0, 0x16EA},
+        {0x16F1, 0x16F8},
+        {0x1700, 0x1711},
+        {0x171F, 0x1731},
+        {0x1740, 0x1751},
+        {0x1760, 0x176C},
+        {0x176E, 0x1770},
+        {0x1780, 0x17B3},
+        {0x17D7, 0x17D7},
+        {0x17DC, 0x17DC},
+        {0x1820, 0x1878},
+        {0x1880, 0x1884},
+        {0x1887, 0x18A8},
+        {0x18AA, 0x18AA},
+        {0x18B0, 0x18F5},
+        {0x1900, 0x191E},
+        {0x1950, 0x196D},
+        {0x1970, 0x1974},
+        {0x1980, 0x19AB},
+        {0x19B0, 0x19C9},
+        {0x1A00, 0x1A16},
+        {0x1A20, 0x1A54},
+        {0x1AA7, 0x1AA7},
+        {0x1B05, 0x1B33},
+        {0x1B45, 0x1B4C},
+        {0x1B83, 0x1BA0},
+        {0x1BAE, 0x1BAF},
+        {0x1BBA, 0x1BE5},
+        {0x1C00, 0x1C23},
+        {0x1C4D, 0x1C4F},
+        {0x1C5A, 0x1C7D},
+        {0x1C80, 0x1C8A},
+        {0x1C90, 0x1CBA},
+        {0x1CBD, 0x1CBF},
+        {0x1CE9, 0x1CEC},
+        {0x1CEE, 0x1CF3},
+        {0x1CF5, 0x1CF6},
+        {0x1CFA, 0x1CFA},
+        {0x1D00, 0x1DBF},
+        {0x1E00, 0x1F15},
+        {0x1F18, 0x1F1D},
+        {0x1F20, 0x1F45},
+        {0x1F48, 0x1F4D},
+        {0x1F50, 0x1F57},
+        {0x1F59, 0x1F59},
+        {0x1F5B, 0x1F5B},
+        {0x1F5D, 0x1F5D},
+        {0x1F5F, 0x1F7D},
+        {0x1F80, 0x1FB4},
+        {0x1FB6, 0x1FBC},
+        {0x1FBE, 0x1FBE},
+        {0x1FC2, 0x1FC4},
+        {0x1FC6, 0x1FCC},
+        {0x1FD0, 0x1FD3},
+        {0x1FD6, 0x1FDB},
+        {0x1FE0, 0x1FEC},
+        {0x1FF2, 0x1FF4},
+        {0x1FF6, 0x1FFC},
+        {0x2071, 0x2071},
+        {0x207F, 0x207F},
+        {0x2090, 0x209C},
+        {0x2102, 0x2102},
+        {0x2107, 0x2107},
+        {0x210A, 0x2113},
+        {0x2115, 0x2115},
+        {0x2119, 0x211D},
+        {0x2124, 0x2124},
+        {0x2126, 0x2126},
+        {0x2128, 0x2128},
+        {0x212A, 0x212D},
+        {0x212F, 0x2139},
+        {0x213C, 0x213F},
+        {0x2145, 0x2149},
+        {0x214E, 0x214E},
+        {0x2183, 0x2184},
+        {0x2C00, 0x2CE4},
+        {0x2CEB, 0x2CEE},
+        {0x2CF2, 0x2CF3},
+        {0x2D00, 0x2D25},
+        {0x2D27, 0x2D27},
+        {0x2D2D, 0x2D2D},
+        {0x2D30, 0x2D67},
+        {0x2D6F, 0x2D6F},
+        {0x2D80, 0x2D96},
+        {0x2DA0, 0x2DA6},
+        {0x2DA8, 0x2DAE},
+        {0x2DB0, 0x2DB6},
+        {0x2DB8, 0x2DBE},
+        {0x2DC0, 0x2DC6},
+        {0x2DC8, 0x2DCE},
+        {0x2DD0, 0x2DD6},
+        {0x2DD8, 0x2DDE},
+        {0x2E2F, 0x2E2F},
+        {0x3005, 0x3006},
+        {0x3031, 0x3035},
+        {0x303B, 0x303C},
+        {0x3041, 0x3096},
+        {0x309D, 0x309F},
+        {0x30A1, 0x30FA},
+        {0x30FC, 0x30FF},
+        {0x3105, 0x312F},
+        {0x3131, 0x318E},
+        {0x31A0, 0x31BF},
+        {0x31F0, 0x31FF},
+        {0x3400, 0x4DBF},
+        {0x4E00, 0xA48C},
+        {0xA4D0, 0xA4FD},
+        {0xA500, 0xA60C},
+        {0xA610, 0xA61F},
+        {0xA62A, 0xA62B},
+        {0xA640, 0xA66E},
+        {0xA67F, 0xA69D},
+        {0xA6A0, 0xA6E5},
+        {0xA717, 0xA71F},
+        {0xA722, 0xA788},
+        {0xA78B, 0xA7DC},
+        {0xA7F1, 0xA801},
+        {0xA803, 0xA805},
+        {0xA807, 0xA80A},
+        {0xA80C, 0xA822},
+        {0xA840, 0xA873},
+        {0xA882, 0xA8B3},
+        {0xA8F2, 0xA8F7},
+        {0xA8FB, 0xA8FB},
+        {0xA8FD, 0xA8FE},
+        {0xA90A, 0xA925},
+        {0xA930, 0xA946},
+        {0xA960, 0xA97C},
+        {0xA984, 0xA9B2},
+        {0xA9CF, 0xA9CF},
+        {0xA9E0, 0xA9E4},
+        {0xA9E6, 0xA9EF},
+        {0xA9FA, 0xA9FE},
+        {0xAA00, 0xAA28},
+        {0xAA40, 0xAA42},
+        {0xAA44, 0xAA4B},
+        {0xAA60, 0xAA76},
+        {0xAA7A, 0xAA7A},
+        {0xAA7E, 0xAAAF},
+        {0xAAB1, 0xAAB1},
+        {0xAAB5, 0xAAB6},
+        {0xAAB9, 0xAABD},
+        {0xAAC0, 0xAAC0},
+        {0xAAC2, 0xAAC2},
+        {0xAADB, 0xAADD},
+        {0xAAE0, 0xAAEA},
+        {0xAAF2, 0xAAF4},
+        {0xAB01, 0xAB06},
+        {0xAB09, 0xAB0E},
+        {0xAB11, 0xAB16},
+        {0xAB20, 0xAB26},
+        {0xAB28, 0xAB2E},
+        {0xAB30, 0xAB5A},
+        {0xAB5C, 0xAB69},
+        {0xAB70, 0xABE2},
+        {0xAC00, 0xD7A3},
+        {0xD7B0, 0xD7C6},
+        {0xD7CB, 0xD7FB},
+        {0xF900, 0xFA6D},
+        {0xFA70, 0xFAD9},
+        {0xFB00, 0xFB06},
+        {0xFB13, 0xFB17},
+        {0xFB1D, 0xFB1D},
+        {0xFB1F, 0xFB28},
+        {0xFB2A, 0xFB36},
+        {0xFB38, 0xFB3C},
+        {0xFB3E, 0xFB3E},
+        {0xFB40, 0xFB41},
+        {0xFB43, 0xFB44},
+        {0xFB46, 0xFBB1},
+        {0xFBD3, 0xFD3D},
+        {0xFD50, 0xFD8F},
+        {0xFD92, 0xFDC7},
+        {0xFDF0, 0xFDFB},
+        {0xFE70, 0xFE74},
+        {0xFE76, 0xFEFC},
+        {0xFF21, 0xFF3A},
+        {0xFF41, 0xFF5A},
+        {0xFF66, 0xFFBE},
+        {0xFFC2, 0xFFC7},
+        {0xFFCA, 0xFFCF},
+        {0xFFD2, 0xFFD7},
+        {0xFFDA, 0xFFDC},
+        {0x10000, 0x1000B},
+        {0x1000D, 0x10026},
+        {0x10028, 0x1003A},
+        {0x1003C, 0x1003D},
+        {0x1003F, 0x1004D},
+        {0x10050, 0x1005D},
+        {0x10080, 0x100FA},
+        {0x10280, 0x1029C},
+        {0x102A0, 0x102D0},
+        {0x10300, 0x1031F},
+        {0x1032D, 0x10340},
+        {0x10342, 0x10349},
+        {0x10350, 0x10375},
+        {0x10380, 0x1039D},
+        {0x103A0, 0x103C3},
+        {0x103C8, 0x103CF},
+        {0x10400, 0x1049D},
+        {0x104B0, 0x104D3},
+        {0x104D8, 0x104FB},
+        {0x10500, 0x10527},
+        {0x10530, 0x10563},
+        {0x10570, 0x1057A},
+        {0x1057C, 0x1058A},
+        {0x1058C, 0x10592},
+        {0x10594, 0x10595},
+        {0x10597, 0x105A1},
+        {0x105A3, 0x105B1},
+        {0x105B3, 0x105B9},
+        {0x105BB, 0x105BC},
+        {0x105C0, 0x105F3},
+        {0x10600, 0x10736},
+        {0x10740, 0x10755},
+        {0x10760, 0x10767},
+        {0x10780, 0x10785},
+        {0x10787, 0x107B0},
+        {0x107B2, 0x107BA},
+        {0x10800, 0x10805},
+        {0x10808, 0x10808},
+        {0x1080A, 0x10835},
+        {0x10837, 0x10838},
+        {0x1083C, 0x1083C},
+        {0x1083F, 0x10855},
+        {0x10860, 0x10876},
+        {0x10880, 0x1089E},
+        {0x108E0, 0x108F2},
+        {0x108F4, 0x108F5},
+        {0x10900, 0x10915},
+        {0x10920, 0x10939},
+        {0x10940, 0x10959},
+        {0x10980, 0x109B7},
+        {0x109BE, 0x109BF},
+        {0x10A00, 0x10A00},
+        {0x10A10, 0x10A13},
+        {0x10A15, 0x10A17},
+        {0x10A19, 0x10A35},
+        {0x10A60, 0x10A7C},
+        {0x10A80, 0x10A9C},
+        {0x10AC0, 0x10AC7},
+        {0x10AC9, 0x10AE4},
+        {0x10B00, 0x10B35},
+        {0x10B40, 0x10B55},
+        {0x10B60, 0x10B72},
+        {0x10B80, 0x10B91},
+        {0x10C00, 0x10C48},
+        {0x10C80, 0x10CB2},
+        {0x10CC0, 0x10CF2},
+        {0x10D00, 0x10D23},
+        {0x10D4A, 0x10D65},
+        {0x10D6F, 0x10D85},
+        {0x10E80, 0x10EA9},
+        {0x10EB0, 0x10EB1},
+        {0x10EC2, 0x10EC7},
+        {0x10F00, 0x10F1C},
+        {0x10F27, 0x10F27},
+        {0x10F30, 0x10F45},
+        {0x10F70, 0x10F81},
+        {0x10FB0, 0x10FC4},
+        {0x10FE0, 0x10FF6},
+        {0x11003, 0x11037},
+        {0x11071, 0x11072},
+        {0x11075, 0x11075},
+        {0x11083, 0x110AF},
+        {0x110D0, 0x110E8},
+        {0x11103, 0x11126},
+        {0x11144, 0x11144},
+        {0x11147, 0x11147},
+        {0x11150, 0x11172},
+        {0x11176, 0x11176},
+        {0x11183, 0x111B2},
+        {0x111C1, 0x111C4},
+        {0x111DA, 0x111DA},
+        {0x111DC, 0x111DC},
+        {0x11200, 0x11211},
+        {0x11213, 0x1122B},
+        {0x1123F, 0x11240},
+        {0x11280, 0x11286},
+        {0x11288, 0x11288},
+        {0x1128A, 0x1128D},
+        {0x1128F, 0x1129D},
+        {0x1129F, 0x112A8},
+        {0x112B0, 0x112DE},
+        {0x11305, 0x1130C},
+        {0x1130F, 0x11310},
+        {0x11313, 0x11328},
+        {0x1132A, 0x11330},
+        {0x11332, 0x11333},
+        {0x11335, 0x11339},
+        {0x1133D, 0x1133D},
+        {0x11350, 0x11350},
+        {0x1135D, 0x11361},
+        {0x11380, 0x11389},
+        {0x1138B, 0x1138B},
+        {0x1138E, 0x1138E},
+        {0x11390, 0x113B5},
+        {0x113B7, 0x113B7},
+        {0x113D1, 0x113D1},
+        {0x113D3, 0x113D3},
+        {0x11400, 0x11434},
+        {0x11447, 0x1144A},
+        {0x1145F, 0x11461},
+        {0x11480, 0x114AF},
+        {0x114C4, 0x114C5},
+        {0x114C7, 0x114C7},
+        {0x11580, 0x115AE},
+        {0x115D8, 0x115DB},
+        {0x11600, 0x1162F},
+        {0x11644, 0x11644},
+        {0x11680, 0x116AA},
+        {0x116B8, 0x116B8},
+        {0x11700, 0x1171A},
+        {0x11740, 0x11746},
+        {0x11800, 0x1182B},
+        {0x118A0, 0x118DF},
+        {0x118FF, 0x11906},
+        {0x11909, 0x11909},
+        {0x1190C, 0x11913},
+        {0x11915, 0x11916},
+        {0x11918, 0x1192F},
+        {0x1193F, 0x1193F},
+        {0x11941, 0x11941},
+        {0x119A0, 0x119A7},
+        {0x119AA, 0x119D0},
+        {0x119E1, 0x119E1},
+        {0x119E3, 0x119E3},
+        {0x11A00, 0x11A00},
+        {0x11A0B, 0x11A32},
+        {0x11A3A, 0x11A3A},
+        {0x11A50, 0x11A50},
+        {0x11A5C, 0x11A89},
+        {0x11A9D, 0x11A9D},
+        {0x11AB0, 0x11AF8},
+        {0x11BC0, 0x11BE0},
+        {0x11C00, 0x11C08},
+        {0x11C0A, 0x11C2E},
+        {0x11C40, 0x11C40},
+        {0x11C72, 0x11C8F},
+        {0x11D00, 0x11D06},
+        {0x11D08, 0x11D09},
+        {0x11D0B, 0x11D30},
+        {0x11D46, 0x11D46},
+        {0x11D60, 0x11D65},
+        {0x11D67, 0x11D68},
+        {0x11D6A, 0x11D89},
+        {0x11D98, 0x11D98},
+        {0x11DB0, 0x11DDB},
+        {0x11EE0, 0x11EF2},
+        {0x11F02, 0x11F02},
+        {0x11F04, 0x11F10},
+        {0x11F12, 0x11F33},
+        {0x11FB0, 0x11FB0},
+        {0x12000, 0x12399},
+        {0x12480, 0x12543},
+        {0x12F90, 0x12FF0},
+        {0x13000, 0x1342F},
+        {0x13441, 0x13446},
+        {0x13460, 0x143FA},
+        {0x14400, 0x14646},
+        {0x16100, 0x1611D},
+        {0x16800, 0x16A38},
+        {0x16A40, 0x16A5E},
+        {0x16A70, 0x16ABE},
+        {0x16AD0, 0x16AED},
+        {0x16B00, 0x16B2F},
+        {0x16B40, 0x16B43},
+        {0x16B63, 0x16B77},
+        {0x16B7D, 0x16B8F},
+        {0x16D40, 0x16D6C},
+        {0x16E40, 0x16E7F},
+        {0x16EA0, 0x16EB8},
+        {0x16EBB, 0x16ED3},
+        {0x16F00, 0x16F4A},
+        {0x16F50, 0x16F50},
+        {0x16F93, 0x16F9F},
+        {0x16FE0, 0x16FE1},
+        {0x16FE3, 0x16FE3},
+        {0x16FF2, 0x16FF3},
+        {0x17000, 0x18CD5},
+        {0x18CFF, 0x18D1E},
+        {0x18D80, 0x18DF2},
+        {0x1AFF0, 0x1AFF3},
+        {0x1AFF5, 0x1AFFB},
+        {0x1AFFD, 0x1AFFE},
+        {0x1B000, 0x1B122},
+        {0x1B132, 0x1B132},
+        {0x1B150, 0x1B152},
+        {0x1B155, 0x1B155},
+        {0x1B164, 0x1B167},
+        {0x1B170, 0x1B2FB},
+        {0x1BC00, 0x1BC6A},
+        {0x1BC70, 0x1BC7C},
+        {0x1BC80, 0x1BC88},
+        {0x1BC90, 0x1BC99},
+        {0x1D400, 0x1D454},
+        {0x1D456, 0x1D49C},
+        {0x1D49E, 0x1D49F},
+        {0x1D4A2, 0x1D4A2},
+        {0x1D4A5, 0x1D4A6},
+        {0x1D4A9, 0x1D4AC},
+        {0x1D4AE, 0x1D4B9},
+        {0x1D4BB, 0x1D4BB},
+        {0x1D4BD, 0x1D4C3},
+        {0x1D4C5, 0x1D505},
+        {0x1D507, 0x1D50A},
+        {0x1D50D, 0x1D514},
+        {0x1D516, 0x1D51C},
+        {0x1D51E, 0x1D539},
+        {0x1D53B, 0x1D53E},
+        {0x1D540, 0x1D544},
+        {0x1D546, 0x1D546},
+        {0x1D54A, 0x1D550},
+        {0x1D552, 0x1D6A5},
+        {0x1D6A8, 0x1D6C0},
+        {0x1D6C2, 0x1D6DA},
+        {0x1D6DC, 0x1D6FA},
+        {0x1D6FC, 0x1D714},
+        {0x1D716, 0x1D734},
+        {0x1D736, 0x1D74E},
+        {0x1D750, 0x1D76E},
+        {0x1D770, 0x1D788},
+        {0x1D78A, 0x1D7A8},
+        {0x1D7AA, 0x1D7C2},
+        {0x1D7C4, 0x1D7CB},
+        {0x1DF00, 0x1DF1E},
+        {0x1DF25, 0x1DF2A},
+        {0x1E030, 0x1E06D},
+        {0x1E100, 0x1E12C},
+        {0x1E137, 0x1E13D},
+        {0x1E14E, 0x1E14E},
+        {0x1E290, 0x1E2AD},
+        {0x1E2C0, 0x1E2EB},
+        {0x1E4D0, 0x1E4EB},
+        {0x1E5D0, 0x1E5ED},
+        {0x1E5F0, 0x1E5F0},
+        {0x1E6C0, 0x1E6DE},
+        {0x1E6E0, 0x1E6E2},
+        {0x1E6E4, 0x1E6E5},
+        {0x1E6E7, 0x1E6ED},
+        {0x1E6F0, 0x1E6F4},
+        {0x1E6FE, 0x1E6FF},
+        {0x1E7E0, 0x1E7E6},
+        {0x1E7E8, 0x1E7EB},
+        {0x1E7ED, 0x1E7EE},
+        {0x1E7F0, 0x1E7FE},
+        {0x1E800, 0x1E8C4},
+        {0x1E900, 0x1E943},
+        {0x1E94B, 0x1E94B},
+        {0x1EE00, 0x1EE03},
+        {0x1EE05, 0x1EE1F},
+        {0x1EE21, 0x1EE22},
+        {0x1EE24, 0x1EE24},
+        {0x1EE27, 0x1EE27},
+        {0x1EE29, 0x1EE32},
+        {0x1EE34, 0x1EE37},
+        {0x1EE39, 0x1EE39},
+        {0x1EE3B, 0x1EE3B},
+        {0x1EE42, 0x1EE42},
+        {0x1EE47, 0x1EE47},
+        {0x1EE49, 0x1EE49},
+        {0x1EE4B, 0x1EE4B},
+        {0x1EE4D, 0x1EE4F},
+        {0x1EE51, 0x1EE52},
+        {0x1EE54, 0x1EE54},
+        {0x1EE57, 0x1EE57},
+        {0x1EE59, 0x1EE59},
+        {0x1EE5B, 0x1EE5B},
+        {0x1EE5D, 0x1EE5D},
+        {0x1EE5F, 0x1EE5F},
+        {0x1EE61, 0x1EE62},
+        {0x1EE64, 0x1EE64},
+        {0x1EE67, 0x1EE6A},
+        {0x1EE6C, 0x1EE72},
+        {0x1EE74, 0x1EE77},
+        {0x1EE79, 0x1EE7C},
+        {0x1EE7E, 0x1EE7E},
+        {0x1EE80, 0x1EE89},
+        {0x1EE8B, 0x1EE9B},
+        {0x1EEA1, 0x1EEA3},
+        {0x1EEA5, 0x1EEA9},
+        {0x1EEAB, 0x1EEBB},
+        {0x20000, 0x2A6DF},
+        {0x2A700, 0x2B81D},
+        {0x2B820, 0x2CEAD},
+        {0x2CEB0, 0x2EBE0},
+        {0x2EBF0, 0x2EE5D},
+        {0x2F800, 0x2FA1D},
+        {0x30000, 0x3134A},
+        {0x31350, 0x33479},
+    };
+
+    for (const auto& r : ranges) {
+        if (ch >= r.start && ch <= r.end)
+            return true;
+    }
+    return false;
+}
+
+bool is_space(char32_t cp) {
+    switch (cp) {
+        case 0x0009:  // TAB \t
+        case 0x000A:  // LF \n
+        case 0x000B:  // VT
+        case 0x000C:  // FF
+        case 0x000D:  // CR \r
+        case 0x0020:  // Space
+        case 0x00A0:  // No-Break Space
+        case 0x1680:  // Ogham Space Mark
+        case 0x2000:  // En Quad
+        case 0x2001:  // Em Quad
+        case 0x2002:  // En Space
+        case 0x2003:  // Em Space
+        case 0x2004:  // Three-Per-Em Space
+        case 0x2005:  // Four-Per-Em Space
+        case 0x2006:  // Six-Per-Em Space
+        case 0x2007:  // Figure Space
+        case 0x2008:  // Punctuation Space
+        case 0x2009:  // Thin Space
+        case 0x200A:  // Hair Space
+        case 0x202F:  // Narrow No-Break Space
+        case 0x205F:  // Medium Mathematical Space
+        case 0x3000:  // Ideographic Space
+            return true;
+        default:
+            return false;
+    }
+}
+
+std::string str_to_lower(const std::string& input) {
+    std::string result = input;
+    std::transform(result.begin(), result.end(), result.begin(),
+                   [](unsigned char c) { return std::tolower(c); });
+    return result;
+}
+
+// UTF-8 -> Unicode code points
+std::vector<char32_t> utf8_to_codepoints(const std::string& str) {
+    std::vector<char32_t> codepoints;
+    size_t i = 0;
+    while (i < str.size()) {
+        unsigned char c    = str[i];
+        char32_t cp        = 0;
+        size_t extra_bytes = 0;
+
+        if ((c & 0x80) == 0)
+            cp = c;
+        else if ((c & 0xE0) == 0xC0) {
+            cp          = c & 0x1F;
+            extra_bytes = 1;
+        } else if ((c & 0xF0) == 0xE0) {
+            cp          = c & 0x0F;
+            extra_bytes = 2;
+        } else if ((c & 0xF8) == 0xF0) {
+            cp          = c & 0x07;
+            extra_bytes = 3;
+        } else {
+            ++i;
+            continue;
+        }  // Invalid UTF-8
+
+        if (i + extra_bytes >= str.size())
+            break;
+
+        for (size_t j = 1; j <= extra_bytes; ++j)
+            cp = (cp << 6) | (str[i + j] & 0x3F);
+
+        codepoints.push_back(cp);
+        i += 1 + extra_bytes;
+    }
+    return codepoints;
+}
+
+// Unicode code point -> UTF-8
+std::string codepoint_to_utf8(char32_t cp) {
+    std::string out;
+    if (cp <= 0x7F)
+        out.push_back(static_cast<char>(cp));
+    else if (cp <= 0x7FF) {
+        out.push_back(static_cast<char>(0xC0 | (cp >> 6)));
+        out.push_back(static_cast<char>(0x80 | (cp & 0x3F)));
+    } else if (cp <= 0xFFFF) {
+        out.push_back(static_cast<char>(0xE0 | (cp >> 12)));
+        out.push_back(static_cast<char>(0x80 | ((cp >> 6) & 0x3F)));
+        out.push_back(static_cast<char>(0x80 | (cp & 0x3F)));
+    } else {
+        out.push_back(static_cast<char>(0xF0 | (cp >> 18)));
+        out.push_back(static_cast<char>(0x80 | ((cp >> 12) & 0x3F)));
+        out.push_back(static_cast<char>(0x80 | ((cp >> 6) & 0x3F)));
+        out.push_back(static_cast<char>(0x80 | (cp & 0x3F)));
+    }
+    return out;
+}
+
+bool starts_with(const std::vector<char32_t>& text,
+                 const std::vector<char32_t>& prefix,
+                 std::size_t index) {
+    if (index > text.size()) {
+        return false;
+    }
+    if (prefix.size() > text.size() - index) {
+        return false;
+    }
+    return std::equal(prefix.begin(), prefix.end(), text.begin() + index);
+}
+
+std::vector<std::string> token_split(const std::string& text) {
+    std::vector<std::string> tokens;
+    auto cps = utf8_to_codepoints(text);
+    size_t i = 0;
+
+    while (i < cps.size()) {
+        char32_t cp = cps[i];
+
+        // `(?i:'s|'t|'re|'ve|'m|'ll|'d)`
+        if (cp == U'\'' && i + 1 < cps.size()) {
+            std::string next = str_to_lower(codepoint_to_utf8(cps[i + 1]));
+            if (next == "s" || next == "t" || next == "m") {
+                tokens.push_back("'" + next);
+                i += 2;
+                continue;
+            }
+            if (i + 2 < cps.size()) {
+                next += str_to_lower(codepoint_to_utf8(cps[i + 2]));
+                if (next == "re" || next == "ve" || next == "ll" || next == "d") {
+                    tokens.push_back("'" + next);
+                    i += 3;
+                    continue;
+                }
+            }
+        }
+
+        // `\p{N}`
+        if (is_number(cp)) {
+            tokens.push_back(codepoint_to_utf8(cp));
+            ++i;
+            continue;
+        }
+
+        // `[^\r\n\p{L}\p{N}]?\p{L}+`
+        {
+            // `[^\r\n\p{L}\p{N}]\p{L}+`
+            if (!is_letter(cp) && cp != U'\r' && cp != U'\n' && i + 1 < cps.size() && is_letter(cps[i + 1])) {
+                std::string token = codepoint_to_utf8(cp);
+                ++i;
+
+                while (i < cps.size() && is_letter(cps[i])) {
+                    token += codepoint_to_utf8(cps[i]);
+                    ++i;
+                }
+                tokens.push_back(token);
+                continue;
+            }
+
+            // `\p{L}+`
+            if (is_letter(cp)) {
+                std::string token = codepoint_to_utf8(cp);
+                ++i;
+                while (i < cps.size() && is_letter(cps[i])) {
+                    token += codepoint_to_utf8(cps[i]);
+                    ++i;
+                }
+                tokens.push_back(token);
+                continue;
+            }
+        }
+
+        // ` ?[^\s\p{L}\p{N}]+[\r\n]*`
+        {
+            // ` [^\s\p{L}\p{N}]+[\r\n]*`
+            if (cp == U' ' && i + 1 < cps.size() && !isspace(cps[i + 1]) && !is_letter(cps[i + 1]) && !is_number(cps[i + 1])) {
+                std::string token = codepoint_to_utf8(cp);
+                token += codepoint_to_utf8(cps[i + 1]);
+                i += 2;
+
+                while (i < cps.size() && !is_letter(cps[i]) && !is_number(cps[i]) && !isspace(cps[i])) {
+                    token += codepoint_to_utf8(cps[i]);
+                    ++i;
+                }
+
+                while (i < cps.size() && (cps[i] == U'\r' || cps[i] == U'\n')) {
+                    token += codepoint_to_utf8(cps[i]);
+                    ++i;
+                }
+
+                tokens.push_back(token);
+                continue;
+            }
+
+            // `[^\s\p{L}\p{N}]+[\r\n]*`
+            std::string token;
+            if (!is_letter(cps[i]) && !is_number(cps[i]) && !isspace(cps[i])) {
+                std::string token = codepoint_to_utf8(cp);
+                ++i;
+
+                while (i < cps.size() && !is_letter(cps[i]) && !is_number(cps[i]) && !isspace(cps[i])) {
+                    token += codepoint_to_utf8(cps[i]);
+                    ++i;
+                }
+
+                while (i < cps.size() && (cps[i] == U'\r' || cps[i] == U'\n')) {
+                    token += codepoint_to_utf8(cps[i]);
+                    ++i;
+                }
+
+                tokens.push_back(token);
+                continue;
+            }
+        }
+
+        // `\s*[\r\n]+|\s+(?!\S)|\s+`
+        if (is_space(cp)) {
+            std::string token = codepoint_to_utf8(cp);
+            ++i;
+
+            while (i < cps.size() && is_space(cps[i])) {
+                token += codepoint_to_utf8(cps[i]);
+                ++i;
+                if (cps[i] == U'\r' || cps[i] == U'\n') {
+                    break;
+                }
+            }
+
+            tokens.push_back(token);
+            continue;
+        }
+
+        // skip
+        ++i;
+    }
+
+    return tokens;
+}
+
+std::vector<std::string> split_with_special_tokens(
+    const std::string& text,
+    const std::vector<std::string>& special_tokens) {
+    std::vector<std::string> result;
+    size_t pos      = 0;
+    size_t text_len = text.size();
+
+    while (pos < text_len) {
+        size_t next_pos = text_len;
+        std::string matched_token;
+
+        for (const auto& token : special_tokens) {
+            size_t token_pos = text.find(token, pos);
+            if (token_pos != std::string::npos && token_pos < next_pos) {
+                next_pos      = token_pos;
+                matched_token = token;
+            }
+        }
+
+        if (next_pos > pos) {
+            result.push_back(text.substr(pos, next_pos - pos));
+        }
+
+        if (!matched_token.empty()) {
+            result.push_back(matched_token);
+            pos = next_pos + matched_token.size();
+        } else {
+            break;
+        }
+    }
+
+    return result;
+}
+
+// int main() {
+//     std::string text = "I'm testing C++ token_split function. 你好，世界! 123";
+//     auto tokens = token_split(text);
+
+//     for (const auto& t : tokens) {
+//         std::cout << "[" << t << "] ";
+//     }
+//     std::cout << "\n";
+//     return 0;
+// }

tokenize_util.h

@@ -0,0 +1,10 @@
+#ifndef __TOKENIZE_UTIL__
+#define __TOKENIZE_UTIL__
+
+#include <string>
+#include <vector>
+
+std::vector<std::string> token_split(const std::string& text);
+std::vector<std::string> split_with_special_tokens(const std::string& text, const std::vector<std::string>& special_tokens);
+
+#endif  // __TOKENIZE_UTIL__

upscaler.cpp

@@ -18,7 +18,8 @@ struct UpscalerGGML {
     }
 
     bool load_from_file(const std::string& esrgan_path,
-                        bool offload_params_to_cpu) {
+                        bool offload_params_to_cpu,
+                        int n_threads) {
         ggml_log_set(ggml_log_callback_default, nullptr);
 #ifdef SD_USE_CUDA
         LOG_DEBUG("Using CUDA backend");
@@ -54,7 +55,7 @@ struct UpscalerGGML {
         if (direct) {
             esrgan_upscaler->enable_conv2d_direct();
         }
-        if (!esrgan_upscaler->load_from_file(esrgan_path)) {
+        if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) {
             return false;
         }
         return true;
@@ -69,7 +70,7 @@ struct UpscalerGGML {
                  input_image.width, input_image.height, output_width, output_height);
 
         struct ggml_init_params params;
-        params.mem_size = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
+        params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
         params.mem_buffer = NULL;
         params.no_alloc   = false;
 
@@ -81,7 +82,7 @@ struct UpscalerGGML {
         }
         // LOG_DEBUG("upscale work buffer size: %.2f MB", params.mem_size / 1024.f / 1024.f);
         ggml_tensor* input_image_tensor = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, input_image.width, input_image.height, 3, 1);
-        sd_image_to_tensor(input_image.data, input_image_tensor);
+        sd_image_to_tensor(input_image, input_image_tensor);
 
         ggml_tensor* upscaled = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, output_width, output_height, 3, 1);
         auto on_tiling        = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
@@ -124,7 +125,7 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
         return NULL;
     }
 
-    if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, offload_params_to_cpu)) {
+    if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, offload_params_to_cpu, n_threads)) {
         delete upscaler_ctx->upscaler;
         upscaler_ctx->upscaler = NULL;
         free(upscaler_ctx);
@@ -137,6 +138,13 @@ sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_
     return upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
 }
 
+int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {
+    if (upscaler_ctx == NULL || upscaler_ctx->upscaler == NULL || upscaler_ctx->upscaler->esrgan_upscaler == NULL) {
+        return 1;
+    }
+    return upscaler_ctx->upscaler->esrgan_upscaler->scale;
+}
+
 void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) {
     if (upscaler_ctx->upscaler != NULL) {
         delete upscaler_ctx->upscaler;

util.cpp

@@ -84,6 +84,7 @@ int round_up_to(int value, int base) {
 }
 
 #ifdef _WIN32  // code for windows
+#define NOMINMAX
 #include <windows.h>
 
 bool file_exists(const std::string& filename) {
@@ -110,56 +111,6 @@ 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);
-    bool isAbsolutePath = false;
-    // Check if the directory was found
-    if (hFind == INVALID_HANDLE_VALUE) {
-        printf("Unable to find directory. Try with original path \n");
-
-        char directoryPathAbsolute[MAX_PATH];
-        sprintf(directoryPathAbsolute, "%s*", dir.c_str());
-
-        hFind          = FindFirstFile(directoryPathAbsolute, &findFileData);
-        isAbsolutePath = true;
-        if (hFind == INVALID_HANDLE_VALUE) {
-            printf("Absolute path was also wrong.\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)) {
-            if (isAbsolutePath) {
-                files.push_back(dir + "\\" + std::string(findFileData.cFileName));
-            } else {
-                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>
@@ -194,27 +145,6 @@ 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
@@ -318,39 +248,6 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
     return result;
 }
 
-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);
@@ -402,7 +299,7 @@ std::string trim(const std::string& s) {
 static sd_log_cb_t sd_log_cb = NULL;
 void* sd_log_cb_data         = NULL;
 
-#define LOG_BUFFER_SIZE 1024
+#define LOG_BUFFER_SIZE 4096
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
     va_list args;
@@ -491,10 +388,10 @@ sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int
             float original_x = (float)x * image.width / target_width;
             float original_y = (float)y * image.height / target_height;
 
-            int x1 = (int)original_x;
-            int y1 = (int)original_y;
-            int x2 = x1 + 1;
-            int y2 = y1 + 1;
+            uint32_t x1 = (uint32_t)original_x;
+            uint32_t y1 = (uint32_t)original_y;
+            uint32_t x2 = std::min(x1 + 1, image.width - 1);
+            uint32_t y2 = std::min(y1 + 1, image.height - 1);
 
             for (int k = 0; k < image.channel; k++) {
                 float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
@@ -531,23 +428,26 @@ float means[3] = {0.48145466, 0.4578275, 0.40821073};
 float stds[3]  = {0.26862954, 0.26130258, 0.27577711};
 
 // Function to clip and preprocess sd_image_f32_t
-sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size) {
-    float scale = (float)size / fmin(image.width, image.height);
+sd_image_f32_t clip_preprocess(sd_image_f32_t image, int target_width, int target_height) {
+    float width_scale  = (float)target_width / image.width;
+    float height_scale = (float)target_height / image.height;
+
+    float scale = std::fmax(width_scale, height_scale);
 
     // Interpolation
-    int new_width       = (int)(scale * image.width);
-    int new_height      = (int)(scale * image.height);
-    float* resized_data = (float*)malloc(new_width * new_height * image.channel * sizeof(float));
+    int resized_width   = (int)(scale * image.width);
+    int resized_height  = (int)(scale * image.height);
+    float* resized_data = (float*)malloc(resized_width * resized_height * image.channel * sizeof(float));
 
-    for (int y = 0; y < new_height; y++) {
-        for (int x = 0; x < new_width; x++) {
-            float original_x = (float)x * image.width / new_width;
-            float original_y = (float)y * image.height / new_height;
+    for (int y = 0; y < resized_height; y++) {
+        for (int x = 0; x < resized_width; x++) {
+            float original_x = (float)x * image.width / resized_width;
+            float original_y = (float)y * image.height / resized_height;
 
-            int x1 = (int)original_x;
-            int y1 = (int)original_y;
-            int x2 = x1 + 1;
-            int y2 = y1 + 1;
+            uint32_t x1 = (uint32_t)original_x;
+            uint32_t y1 = (uint32_t)original_y;
+            uint32_t x2 = std::min(x1 + 1, image.width - 1);
+            uint32_t y2 = std::min(y1 + 1, image.height - 1);
 
             for (int k = 0; k < image.channel; k++) {
                 float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
@@ -560,26 +460,28 @@ sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size) {
 
                 float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
 
-                *(resized_data + y * new_width * image.channel + x * image.channel + k) = value;
+                *(resized_data + y * resized_width * image.channel + x * image.channel + k) = value;
             }
         }
     }
 
     // Clip and preprocess
-    int h = (new_height - size) / 2;
-    int w = (new_width - size) / 2;
+    int h_offset = std::max((int)(resized_height - target_height) / 2, 0);
+    int w_offset = std::max((int)(resized_width - target_width) / 2, 0);
 
     sd_image_f32_t result;
-    result.width   = size;
-    result.height  = size;
+    result.width   = target_width;
+    result.height  = target_height;
     result.channel = image.channel;
-    result.data    = (float*)malloc(size * size * image.channel * sizeof(float));
+    result.data    = (float*)malloc(target_height * target_width * image.channel * sizeof(float));
 
     for (int k = 0; k < image.channel; k++) {
-        for (int i = 0; i < size; i++) {
-            for (int j = 0; j < size; j++) {
-                *(result.data + i * size * image.channel + j * image.channel + k) =
-                    fmin(fmax(*(resized_data + (i + h) * new_width * image.channel + (j + w) * image.channel + k), 0.0f), 255.0f) / 255.0f;
+        for (int i = 0; i < result.height; i++) {
+            for (int j = 0; j < result.width; j++) {
+                int src_y = std::min(i + h_offset, resized_height - 1);
+                int src_x = std::min(j + w_offset, resized_width - 1);
+                *(result.data + i * result.width * image.channel + j * image.channel + k) =
+                    fmin(fmax(*(resized_data + src_y * resized_width * image.channel + src_x * image.channel + k), 0.0f), 255.0f) / 255.0f;
             }
         }
     }
@@ -589,10 +491,10 @@ sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size) {
 
     // Normalize
     for (int k = 0; k < image.channel; k++) {
-        for (int i = 0; i < size; i++) {
-            for (int j = 0; j < size; j++) {
+        for (int i = 0; i < result.height; i++) {
+            for (int j = 0; j < result.width; j++) {
                 // *(result.data + i * size * image.channel + j * image.channel + k) = 0.5f;
-                int offset  = i * size * image.channel + j * image.channel + k;
+                int offset  = i * result.width * image.channel + j * image.channel + k;
                 float value = *(result.data + offset);
                 value       = (value - means[k]) / stds[k];
                 // value = 0.5f;

util.h

@@ -24,14 +24,9 @@ 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);
-
-sd_image_t* preprocess_id_image(sd_image_t* img);
-
 // std::string sd_basename(const std::string& path);
 
 typedef struct {
@@ -47,7 +42,7 @@ sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image);
 
 sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height);
 
-sd_image_f32_t clip_preprocess(sd_image_f32_t image, int size);
+sd_image_f32_t clip_preprocess(sd_image_f32_t image, int target_width, int target_height);
 
 std::string path_join(const std::string& p1, const std::string& p2);
 std::vector<std::string> split_string(const std::string& str, char delimiter);

vae.hpp

@@ -583,6 +583,7 @@ struct AutoEncoderKL : public VAE {
                  bool decode_graph,
                  struct ggml_tensor** output,
                  struct ggml_context* output_ctx = NULL) {
+        GGML_ASSERT(!decode_only || decode_graph);
         auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(z, decode_graph);
         };

wan.hpp

@@ -1219,7 +1219,7 @@ namespace WAN {
 
         void test() {
             struct ggml_init_params params;
-            params.mem_size = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
+            params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
             params.mem_buffer = NULL;
             params.no_alloc   = false;
 
@@ -1333,7 +1333,7 @@ namespace WAN {
             k = ggml_reshape_4d(ctx, k, head_dim, num_heads, n_token, N);  // [N, n_token, n_head, d_head]
             v = ggml_reshape_4d(ctx, v, head_dim, num_heads, n_token, N);  // [N, n_token, n_head, d_head]
 
-            x = Flux::attention(ctx, backend, q, k, v, pe, mask, flash_attn);  // [N, n_token, dim]
+            x = Rope::attention(ctx, backend, q, k, v, pe, mask, flash_attn);  // [N, n_token, dim]
 
             x = o_proj->forward(ctx, x);  // [N, n_token, dim]
             return x;
@@ -1532,13 +1532,13 @@ namespace WAN {
             blocks["ffn.2"] = std::shared_ptr<GGMLBlock>(new Linear(ffn_dim, dim));
         }
 
-        struct ggml_tensor* forward(struct ggml_context* ctx,
-                                    ggml_backend_t backend,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* e,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* context,
-                                    int64_t context_img_len = 257) {
+        virtual struct ggml_tensor* forward(struct ggml_context* ctx,
+                                            ggml_backend_t backend,
+                                            struct ggml_tensor* x,
+                                            struct ggml_tensor* e,
+                                            struct ggml_tensor* pe,
+                                            struct ggml_tensor* context,
+                                            int64_t context_img_len = 257) {
             // x: [N, n_token, dim]
             // e: [N, 6, dim] or [N, T, 6, dim]
             // context: [N, context_img_len + context_txt_len, dim]
@@ -1584,6 +1584,59 @@ namespace WAN {
         }
     };
 
+    class VaceWanAttentionBlock : public WanAttentionBlock {
+    protected:
+        int block_id;
+        void init_params(struct ggml_context* ctx, const String2GGMLType& tensor_types = {}, const std::string prefix = "") {
+            enum ggml_type wtype = get_type(prefix + "weight", tensor_types, GGML_TYPE_F32);
+            params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
+        }
+
+    public:
+        VaceWanAttentionBlock(bool t2v_cross_attn,
+                              int64_t dim,
+                              int64_t ffn_dim,
+                              int64_t num_heads,
+                              bool qk_norm         = true,
+                              bool cross_attn_norm = false,
+                              float eps            = 1e-6,
+                              int block_id         = 0,
+                              bool flash_attn      = false)
+            : WanAttentionBlock(t2v_cross_attn, dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps, flash_attn), block_id(block_id) {
+            if (block_id == 0) {
+                blocks["before_proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim));
+            }
+            blocks["after_proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim));
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(struct ggml_context* ctx,
+                                                      ggml_backend_t backend,
+                                                      struct ggml_tensor* c,
+                                                      struct ggml_tensor* x,
+                                                      struct ggml_tensor* e,
+                                                      struct ggml_tensor* pe,
+                                                      struct ggml_tensor* context,
+                                                      int64_t context_img_len = 257) {
+            // x: [N, n_token, dim]
+            // e: [N, 6, dim] or [N, T, 6, dim]
+            // context: [N, context_img_len + context_txt_len, dim]
+            // return [N, n_token, dim]
+            if (block_id == 0) {
+                auto before_proj = std::dynamic_pointer_cast<Linear>(blocks["before_proj"]);
+
+                c = before_proj->forward(ctx, c);
+                c = ggml_add(ctx, c, x);
+            }
+
+            auto after_proj = std::dynamic_pointer_cast<Linear>(blocks["after_proj"]);
+
+            c           = WanAttentionBlock::forward(ctx, backend, c, e, pe, context, context_img_len);
+            auto c_skip = after_proj->forward(ctx, c);
+
+            return {c_skip, c};
+        }
+    };
+
     class Head : public GGMLBlock {
     protected:
         int dim;
@@ -1680,22 +1733,25 @@ namespace WAN {
     };
 
     struct WanParams {
-        std::string model_type               = "t2v";
-        std::tuple<int, int, int> patch_size = {1, 2, 2};
-        int64_t text_len                     = 512;
-        int64_t in_dim                       = 16;
-        int64_t dim                          = 2048;
-        int64_t ffn_dim                      = 8192;
-        int64_t freq_dim                     = 256;
-        int64_t text_dim                     = 4096;
-        int64_t out_dim                      = 16;
-        int64_t num_heads                    = 16;
-        int64_t num_layers                   = 32;
-        bool qk_norm                         = true;
-        bool cross_attn_norm                 = true;
-        float eps                            = 1e-6;
-        int64_t flf_pos_embed_token_number   = 0;
-        int theta                            = 10000;
+        std::string model_type                 = "t2v";
+        std::tuple<int, int, int> patch_size   = {1, 2, 2};
+        int64_t text_len                       = 512;
+        int64_t in_dim                         = 16;
+        int64_t dim                            = 2048;
+        int64_t ffn_dim                        = 8192;
+        int64_t freq_dim                       = 256;
+        int64_t text_dim                       = 4096;
+        int64_t out_dim                        = 16;
+        int64_t num_heads                      = 16;
+        int64_t num_layers                     = 32;
+        int64_t vace_layers                    = 0;
+        int64_t vace_in_dim                    = 96;
+        std::map<int, int> vace_layers_mapping = {};
+        bool qk_norm                           = true;
+        bool cross_attn_norm                   = true;
+        float eps                              = 1e-6;
+        int64_t flf_pos_embed_token_number     = 0;
+        int theta                              = 10000;
         // wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24
         std::vector<int> axes_dim = {44, 42, 42};
         int64_t axes_dim_sum      = 128;
@@ -1746,13 +1802,38 @@ namespace WAN {
             if (params.model_type == "i2v") {
                 blocks["img_emb"] = std::shared_ptr<GGMLBlock>(new MLPProj(1280, params.dim, params.flf_pos_embed_token_number));
             }
+
+            // vace
+            if (params.vace_layers > 0) {
+                for (int i = 0; i < params.vace_layers; i++) {
+                    auto block                                 = std::shared_ptr<GGMLBlock>(new VaceWanAttentionBlock(params.model_type == "t2v",
+                                                                                                                      params.dim,
+                                                                                                                      params.ffn_dim,
+                                                                                                                      params.num_heads,
+                                                                                                                      params.qk_norm,
+                                                                                                                      params.cross_attn_norm,
+                                                                                                                      params.eps,
+                                                                                                                      i,
+                                                                                                                      params.flash_attn));
+                    blocks["vace_blocks." + std::to_string(i)] = block;
+                }
+
+                int step = params.num_layers / params.vace_layers;
+                int n    = 0;
+                for (int i = 0; i < params.num_layers; i += step) {
+                    this->params.vace_layers_mapping[i] = n;
+                    n++;
+                }
+
+                blocks["vace_patch_embedding"] = std::shared_ptr<GGMLBlock>(new Conv3d(params.vace_in_dim, params.dim, params.patch_size, params.patch_size));
+            }
         }
 
         struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
                                               struct ggml_tensor* x) {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
-            int64_t T = x->ne[1];
+            int64_t T = x->ne[2];
 
             int pad_t = (std::get<0>(params.patch_size) - T % std::get<0>(params.patch_size)) % std::get<0>(params.patch_size);
             int pad_h = (std::get<1>(params.patch_size) - H % std::get<1>(params.patch_size)) % std::get<1>(params.patch_size);
@@ -1795,9 +1876,12 @@ namespace WAN {
                                          struct ggml_tensor* timestep,
                                          struct ggml_tensor* context,
                                          struct ggml_tensor* pe,
-                                         struct ggml_tensor* clip_fea = NULL,
-                                         int64_t N                    = 1) {
+                                         struct ggml_tensor* clip_fea     = NULL,
+                                         struct ggml_tensor* vace_context = NULL,
+                                         float vace_strength              = 1.f,
+                                         int64_t N                        = 1) {
             // x: [N*C, T, H, W], C => in_dim
+            // vace_context: [N*vace_in_dim, T, H, W]
             // timestep: [N,] or [T]
             // context: [N, L, text_dim]
             // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
@@ -1845,10 +1929,35 @@ namespace WAN {
                 context_img_len = clip_fea->ne[1];  // 257
             }
 
+            // vace_patch_embedding
+            ggml_tensor* c = NULL;
+            if (params.vace_layers > 0) {
+                auto vace_patch_embedding = std::dynamic_pointer_cast<Conv3d>(blocks["vace_patch_embedding"]);
+
+                c = vace_patch_embedding->forward(ctx, vace_context);                          // [N*dim, t_len, h_len, w_len]
+                c = ggml_reshape_3d(ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
+                c = ggml_nn_cont(ctx, ggml_torch_permute(ctx, c, 1, 0, 2, 3));                 // [N, t_len*h_len*w_len, dim]
+            }
+
+            auto x_orig = x;
+
             for (int i = 0; i < params.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<WanAttentionBlock>(blocks["blocks." + std::to_string(i)]);
 
                 x = block->forward(ctx, backend, x, e0, pe, context, context_img_len);
+
+                auto iter = params.vace_layers_mapping.find(i);
+                if (iter != params.vace_layers_mapping.end()) {
+                    int n = iter->second;
+
+                    auto vace_block = std::dynamic_pointer_cast<VaceWanAttentionBlock>(blocks["vace_blocks." + std::to_string(n)]);
+
+                    auto result = vace_block->forward(ctx, backend, c, x_orig, e0, pe, context, context_img_len);
+                    auto c_skip = result.first;
+                    c           = result.second;
+                    c_skip      = ggml_scale(ctx, c_skip, vace_strength);
+                    x           = ggml_add(ctx, x, c_skip);
+                }
             }
 
             x = head->forward(ctx, x, e);  // [N, t_len*h_len*w_len, pt*ph*pw*out_dim]
@@ -1864,6 +1973,8 @@ namespace WAN {
                                     struct ggml_tensor* pe,
                                     struct ggml_tensor* clip_fea        = NULL,
                                     struct ggml_tensor* time_dim_concat = NULL,
+                                    struct ggml_tensor* vace_context    = NULL,
+                                    float vace_strength                 = 1.f,
                                     int64_t N                           = 1) {
             // Forward pass of DiT.
             // x: [N*C, T, H, W]
@@ -1892,7 +2003,7 @@ namespace WAN {
                 t_len           = ((x->ne[2] + (std::get<0>(params.patch_size) / 2)) / std::get<0>(params.patch_size));
             }
 
-            auto out = forward_orig(ctx, backend, x, timestep, context, pe, clip_fea, N);  // [N, t_len*h_len*w_len, pt*ph*pw*C]
+            auto out = forward_orig(ctx, backend, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N);  // [N, t_len*h_len*w_len, pt*ph*pw*C]
 
             out = unpatchify(ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
 
@@ -1927,7 +2038,19 @@ namespace WAN {
                 std::string tensor_name = pair.first;
                 if (tensor_name.find(prefix) == std::string::npos)
                     continue;
-                size_t pos = tensor_name.find("blocks.");
+                size_t pos = tensor_name.find("vace_blocks.");
+                if (pos != std::string::npos) {
+                    tensor_name = tensor_name.substr(pos);  // remove prefix
+                    auto items  = split_string(tensor_name, '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > wan_params.vace_layers) {
+                            wan_params.vace_layers = block_index + 1;
+                        }
+                    }
+                    continue;
+                }
+                pos = tensor_name.find("blocks.");
                 if (pos != std::string::npos) {
                     tensor_name = tensor_name.substr(pos);  // remove prefix
                     auto items  = split_string(tensor_name, '.');
@@ -1937,6 +2060,7 @@ namespace WAN {
                             wan_params.num_layers = block_index + 1;
                         }
                     }
+                    continue;
                 }
                 if (tensor_name.find("img_emb") != std::string::npos) {
                     wan_params.model_type = "i2v";
@@ -1958,7 +2082,11 @@ namespace WAN {
                     wan_params.out_dim   = 48;
                     wan_params.text_len  = 512;
                 } else {
-                    desc                 = "Wan2.1-T2V-1.3B";
+                    if (wan_params.vace_layers > 0) {
+                        desc = "Wan2.1-VACE-1.3B";
+                    } else {
+                        desc = "Wan2.1-T2V-1.3B";
+                    }
                     wan_params.dim       = 1536;
                     wan_params.eps       = 1e-06;
                     wan_params.ffn_dim   = 8960;
@@ -1974,7 +2102,11 @@ namespace WAN {
                         desc              = "Wan2.2-I2V-14B";
                         wan_params.in_dim = 36;
                     } else {
-                        desc              = "Wan2.x-T2V-14B";
+                        if (wan_params.vace_layers > 0) {
+                            desc = "Wan2.x-VACE-14B";
+                        } else {
+                            desc = "Wan2.x-T2V-14B";
+                        }
                         wan_params.in_dim = 16;
                     }
                 } else {
@@ -2015,7 +2147,9 @@ namespace WAN {
                                         struct ggml_tensor* context,
                                         struct ggml_tensor* clip_fea        = NULL,
                                         struct ggml_tensor* c_concat        = NULL,
-                                        struct ggml_tensor* time_dim_concat = NULL) {
+                                        struct ggml_tensor* time_dim_concat = NULL,
+                                        struct ggml_tensor* vace_context    = NULL,
+                                        float vace_strength                 = 1.f) {
             struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, WAN_GRAPH_SIZE, false);
 
             x               = to_backend(x);
@@ -2024,6 +2158,7 @@ namespace WAN {
             clip_fea        = to_backend(clip_fea);
             c_concat        = to_backend(c_concat);
             time_dim_concat = to_backend(time_dim_concat);
+            vace_context    = to_backend(vace_context);
 
             pe_vec      = Rope::gen_wan_pe(x->ne[2],
                                            x->ne[1],
@@ -2053,7 +2188,9 @@ namespace WAN {
                                                   context,
                                                   pe,
                                                   clip_fea,
-                                                  time_dim_concat);
+                                                  time_dim_concat,
+                                                  vace_context,
+                                                  vace_strength);
 
             ggml_build_forward_expand(gf, out);
 
@@ -2067,10 +2204,12 @@ namespace WAN {
                      struct ggml_tensor* clip_fea        = NULL,
                      struct ggml_tensor* c_concat        = NULL,
                      struct ggml_tensor* time_dim_concat = NULL,
+                     struct ggml_tensor* vace_context    = NULL,
+                     float vace_strength                 = 1.f,
                      struct ggml_tensor** output         = NULL,
                      struct ggml_context* output_ctx     = NULL) {
             auto get_graph = [&]() -> struct ggml_cgraph* {
-                return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat);
+                return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
             };
 
             GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
@@ -2108,7 +2247,7 @@ namespace WAN {
                 struct ggml_tensor* out = NULL;
 
                 int t0 = ggml_time_ms();
-                compute(8, x, timesteps, context, NULL, NULL, NULL, &out, work_ctx);
+                compute(8, x, timesteps, context, NULL, NULL, NULL, NULL, 1.f, &out, work_ctx);
                 int t1 = ggml_time_ms();
 
                 print_ggml_tensor(out);