feat: add ernie image support (#1427)

Co-authored-by: professor-moody <keys@nimbus.lan>

.github/workflows/build.yml

@@ -21,6 +21,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "examples/server/frontend",
         "examples/server/frontend/**",
       ]
   pull_request:
@@ -35,6 +36,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "examples/server/frontend",
         "examples/server/frontend/**",
       ]
 
@@ -64,7 +66,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Dependencies
         id: depends
@@ -127,7 +129,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Dependencies
         id: depends
@@ -205,7 +207,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Get commit hash
         id: commit
@@ -239,6 +241,7 @@ jobs:
         id: build-push
         uses: docker/build-push-action@v6
         with:
+          context: .
           platforms: linux/amd64
           push: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
           file: Dockerfile.${{ matrix.variant }}
@@ -264,7 +267,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Dependencies
         id: depends
@@ -345,7 +348,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Install cuda-toolkit
         id: cuda-toolkit
@@ -460,7 +463,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Cache ROCm Installation
         id: cache-rocm
@@ -573,7 +576,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Free disk space
         run: |

.gitmodules

@@ -3,4 +3,10 @@
 	url = https://github.com/ggml-org/ggml.git
 [submodule "examples/server/frontend"]
 	path = examples/server/frontend
-	url = https://github.com/leejet/stable-ui.git
+	url = https://github.com/leejet/sdcpp-webui.git
+[submodule "thirdparty/libwebp"]
+	path = thirdparty/libwebp
+	url = https://github.com/webmproject/libwebp.git
+[submodule "thirdparty/libwebm"]
+	path = thirdparty/libwebm
+	url = https://github.com/webmproject/libwebm.git

CMakeLists.txt

@@ -22,6 +22,26 @@ else()
     set(SD_STANDALONE OFF)
 endif()
 
+set(SD_SUBMODULE_WEBP FALSE)
+if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/libwebp/CMakeLists.txt")
+    set(SD_SUBMODULE_WEBP TRUE)
+endif()
+if(SD_SUBMODULE_WEBP)
+    set(SD_WEBP_DEFAULT ON)
+else()
+    set(SD_WEBP_DEFAULT ${SD_USE_SYSTEM_WEBP})
+endif()
+
+set(SD_SUBMODULE_WEBM FALSE)
+if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/libwebm/CMakeLists.txt")
+    set(SD_SUBMODULE_WEBM TRUE)
+endif()
+if(SD_SUBMODULE_WEBM)
+    set(SD_WEBM_DEFAULT ON)
+else()
+    set(SD_WEBM_DEFAULT ${SD_USE_SYSTEM_WEBM})
+endif()
+
 #
 # Option list
 #
@@ -29,6 +49,10 @@ endif()
 # general
 #option(SD_BUILD_TESTS                "sd: build tests"    ${SD_STANDALONE})
 option(SD_BUILD_EXAMPLES             "sd: build examples" ${SD_STANDALONE})
+option(SD_WEBP                       "sd: enable WebP image I/O support" ${SD_WEBP_DEFAULT})
+option(SD_USE_SYSTEM_WEBP            "sd: link against system libwebp" OFF)
+option(SD_WEBM                       "sd: enable WebM video output support" ${SD_WEBM_DEFAULT})
+option(SD_USE_SYSTEM_WEBM            "sd: link against system libwebm" OFF)
 option(SD_CUDA                       "sd: cuda backend" OFF)
 option(SD_HIPBLAS                    "sd: rocm backend" OFF)
 option(SD_METAL                      "sd: metal backend" OFF)
@@ -77,14 +101,65 @@ if(SD_MUSA)
     add_definitions(-DSD_USE_CUDA)
 endif()
 
+if(SD_WEBP)
+    if(NOT SD_SUBMODULE_WEBP AND NOT SD_USE_SYSTEM_WEBP)
+        message(FATAL_ERROR "WebP support enabled but no source found.
+          Either initialize the submodule:\n  git submodule update --init thirdparty/libwebp\n\n"
+          "Or link against system library:\n  cmake (...) -DSD_USE_SYSTEM_WEBP=ON")
+    endif()
+    if(SD_USE_SYSTEM_WEBP)
+        find_package(WebP REQUIRED)
+        add_library(webp ALIAS WebP::webp)
+        # libwebp CMake target naming is not consistent across versions/distros.
+        # Some export WebP::libwebpmux, others export WebP::webpmux.
+        if(TARGET WebP::libwebpmux)
+            add_library(libwebpmux ALIAS WebP::libwebpmux)
+        elseif(TARGET WebP::webpmux)
+            add_library(libwebpmux ALIAS WebP::webpmux)
+        else()
+            message(FATAL_ERROR
+                "Could not find a compatible webpmux target in system WebP package. "
+                "Expected WebP::libwebpmux or WebP::webpmux."
+            )
+        endif()
+    endif()
+endif()
+
+if(SD_WEBM)
+    if(NOT SD_WEBP)
+        message(FATAL_ERROR "SD_WEBM requires SD_WEBP because WebM output reuses libwebp VP8 encoding.")
+    endif()
+    if(NOT SD_SUBMODULE_WEBM AND NOT SD_USE_SYSTEM_WEBM)
+        message(FATAL_ERROR "WebM support enabled but no source found.
+          Either initialize the submodule:\n  git submodule update --init thirdparty/libwebm\n\n"
+          "Or link against system library:\n  cmake (...) -DSD_USE_SYSTEM_WEBM=ON")
+    endif()
+    if(SD_USE_SYSTEM_WEBM)
+        find_path(WEBM_INCLUDE_DIR
+            NAMES mkvmuxer/mkvmuxer.h mkvparser/mkvparser.h common/webmids.h
+            PATH_SUFFIXES webm
+            REQUIRED)
+        find_library(WEBM_LIBRARY
+            NAMES webm libwebm
+            REQUIRED)
+
+        add_library(webm UNKNOWN IMPORTED)
+        set_target_properties(webm PROPERTIES
+            IMPORTED_LOCATION "${WEBM_LIBRARY}"
+            INTERFACE_INCLUDE_DIRECTORIES "${WEBM_INCLUDE_DIR}")
+    endif()
+endif()
+
 set(SD_LIB stable-diffusion)
 
 file(GLOB SD_LIB_SOURCES
     "src/*.h"
     "src/*.cpp"
     "src/*.hpp"
-    "src/vocab/*.h"
-    "src/vocab/*.cpp"
+    "src/tokenizers/*.h"
+    "src/tokenizers/*.cpp"
+    "src/tokenizers/vocab/*.h"
+    "src/tokenizers/vocab/*.cpp"
 )
 
 find_program(GIT_EXE NAMES git git.exe NO_CMAKE_FIND_ROOT_PATH)
@@ -177,7 +252,7 @@ endif()
 add_subdirectory(thirdparty)
 
 target_link_libraries(${SD_LIB} PUBLIC ggml zip)
-target_include_directories(${SD_LIB} PUBLIC . include)
+target_include_directories(${SD_LIB} PUBLIC . src include)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
 target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17)
 

README.md

@@ -15,6 +15,9 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/04/11** 🚀 stable-diffusion.cpp now uses a brand-new embedded web UI.  
+  👉 Details: [PR #1408](https://github.com/leejet/stable-diffusion.cpp/pull/1408)
+
 * **2026/01/18** 🚀 stable-diffusion.cpp now supports **FLUX.2-klein**  
   👉 Details: [PR #1193](https://github.com/leejet/stable-diffusion.cpp/pull/1193)
 
@@ -54,6 +57,7 @@ API and command-line option may change frequently.***
     - [Z-Image](./docs/z_image.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
+    - [ERNIE-Image](./docs/ernie_image.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
@@ -141,6 +145,7 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 - [🔥Z-Image](./docs/z_image.md)
 - [Ovis-Image](./docs/ovis_image.md)
 - [Anima](./docs/anima.md)
+- [ERNIE-Image](./docs/ernie_image.md)
 - [LoRA](./docs/lora.md)
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)

docs/build.md

@@ -16,6 +16,26 @@ git submodule init
 git submodule update
 ```
 
+## WebP and WebM Support in Examples
+
+The example applications (`examples/cli` and `examples/server`) use `libwebp` to support WebP image I/O, and `examples/cli` can also use `libwebm` for `.webm` video output. Both are enabled by default. WebM output currently reuses `libwebp` to encode each frame as VP8 before muxing with `libwebm`.
+
+If you do not want WebP/WebM support, you can disable them at configure time:
+
+```shell
+mkdir build && cd build
+cmake .. -DSD_WEBP=OFF -DSD_WEBM=OFF
+cmake --build . --config Release
+```
+
+If the submodules are not available, you can also link against system packages instead:
+
+```shell
+mkdir build && cd build
+cmake .. -DSD_USE_SYSTEM_WEBP=ON -DSD_USE_SYSTEM_WEBM=ON
+cmake --build . --config Release
+```
+
 ## Build (CPU only)
 
 If you don't have a GPU or CUDA installed, you can build a CPU-only version.

docs/ernie_image.md

@@ -0,0 +1,35 @@
+# How to Use
+
+You can run ERNIE-Image with stable-diffusion.cpp on GPUs with 4GB of VRAM — or even less.
+
+## Download weights
+
+- Download ERNIE-Image-Turbo
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/diffusion_models
+    - gguf: https://huggingface.co/unsloth/ERNIE-Image-Turbo-GGUF/tree/main
+- Download ERNIE-Image
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/diffusion_models
+    - gguf: https://huggingface.co/unsloth/ERNIE-Image-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/vae
+- Download ministral 3b
+    - safetensors: https://huggingface.co/Comfy-Org/ERNIE-Image/tree/main/text_encoders
+    - gguf: https://huggingface.co/unsloth/Ministral-3-3B-Instruct-2512-GGUF/tree/main
+
+## Examples
+
+### ERNIE-Image-Turbo
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ernie-image-turbo.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\ministral-3-3b.safetensors -p "a lovely cat" --cfg-scale 1.0 --steps 8 -v --offload-to-cpu --diffusion-fa
+```
+
+<img width="256" alt="ERNIE-Image Turbo example" src="../assets/ernie_image/turbo_example.png" />
+
+### ERNIE-Image
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ernie-image-UD-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\ministral-3-3b.safetensors -p "a lovely cat" --cfg-scale 5.0 -v --offload-to-cpu --diffusion-fa
+```
+
+<img width="256" alt="ERNIE-Image example" src="../assets/ernie_image/example.png" />

docs/flux2.md

@@ -8,6 +8,8 @@
     - gguf: https://huggingface.co/city96/FLUX.2-dev-gguf/tree/main
 - Download vae
     - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
+- Download FLUX.2-small-decoder (full_encoder_small_decoder.safetensors) as an alternative VAE option
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-small-decoder/tree/main
 - Download Mistral-Small-3.2-24B-Instruct-2506-GGUF
     - gguf: https://huggingface.co/unsloth/Mistral-Small-3.2-24B-Instruct-2506-GGUF/tree/main
 
@@ -31,6 +33,8 @@
     - gguf: https://huggingface.co/leejet/FLUX.2-klein-base-4B-GGUF/tree/main
 - Download vae
     - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
+- Download FLUX.2-small-decoder (full_encoder_small_decoder.safetensors) as an alternative VAE option
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-small-decoder/tree/main
 - Download Qwen3 4b
     - safetensors: https://huggingface.co/Comfy-Org/flux2-klein-4B/tree/main/split_files/text_encoders
     - gguf: https://huggingface.co/unsloth/Qwen3-4B-GGUF/tree/main

examples/cli/CMakeLists.txt

@@ -1,6 +1,20 @@
 set(TARGET sd-cli)
 
-add_executable(${TARGET} main.cpp)
+add_executable(${TARGET}
+    ../common/common.cpp
+    ../common/log.cpp
+    ../common/media_io.cpp
+    image_metadata.cpp
+    main.cpp
+)
 install(TARGETS ${TARGET} RUNTIME)
-target_link_libraries(${TARGET} PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
-target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17)
+target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT})
+if(SD_WEBP)
+    target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBP)
+    target_link_libraries(${TARGET} PRIVATE webp libwebpmux)
+endif()
+if(SD_WEBM)
+    target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBM)
+    target_link_libraries(${TARGET} PRIVATE webm)
+endif()
+target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17)

examples/cli/README.md

@@ -5,18 +5,23 @@ usage: ./bin/sd-cli  [options]
 
 CLI Options:
   -o, --output <string>       path to write result image to. you can use printf-style %d format specifiers for image sequences (default:
-                              ./output.png) (eg. output_%03d.png)
-  --preview-path <string>     path to write preview image to (default: ./preview.png)
+                              ./output.png) (eg. output_%03d.png). For video generation, single-file outputs support .avi, .webm, and animated .webp
+  --preview-path <string>     path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp
   --preview-interval <int>    interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at
                               every step)
   --output-begin-idx <int>    starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)
+  --image <string>            path to the image to inspect (for metadata mode)
+  --metadata-format <string>  metadata output format, one of [text, json] (default: text)
   --canny                     apply canny preprocessor (edge detection)
   --convert-name              convert tensor name (for convert mode)
   -v, --verbose               print extra info
   --color                     colors the logging tags according to level
   --taesd-preview-only        prevents usage of taesd for decoding the final image. (for use with --preview tae)
   --preview-noisy             enables previewing noisy inputs of the models rather than the denoised outputs
-  -M, --mode                  run mode, one of [img_gen, vid_gen, upscale, convert], default: img_gen
+  --metadata-raw              include raw hex previews for unparsed metadata payloads
+  --metadata-brief            truncate long metadata text values in text output
+  --metadata-all              include structural/container entries such as IHDR, IDAT, and non-metadata JPEG segments
+  -M, --mode                  run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen
   --preview                   preview method. must be one of the following [none, proj, tae, vae] (default is none)
   -h, --help                  show this help message and exit
 
@@ -109,7 +114,7 @@ Generation Options:
                                            medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            eta in DDIM, only for DDIM and TCD (default: 0)
+  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
   --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
@@ -117,7 +122,7 @@ Generation Options:
   --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM and TCD (default: 0)
+  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>
   --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image
@@ -148,3 +153,12 @@ Generation Options:
   --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g., "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
   --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
 ```
+
+Metadata mode inspects PNG/JPEG container metadata without loading any model:
+
+```bash
+./bin/sd-cli -M metadata --image ./output.png
+./bin/sd-cli -M metadata --image ./output.jpg --metadata-format json
+./bin/sd-cli -M metadata --image ./output.png --metadata-raw
+./bin/sd-cli -M metadata --image ./output.png --metadata-all
+```

examples/cli/avi_writer.h

@@ -1,217 +0,0 @@
-#ifndef __AVI_WRITER_H__
-#define __AVI_WRITER_H__
-
-#include <cstdint>
-#include <cstdio>
-#include <cstdlib>
-#include <cstring>
-
-#include "stable-diffusion.h"
-
-#ifndef INCLUDE_STB_IMAGE_WRITE_H
-#include "stb_image_write.h"
-#endif
-
-typedef struct {
-    uint32_t offset;
-    uint32_t size;
-} avi_index_entry;
-
-// Write 32-bit little-endian integer
-void write_u32_le(FILE* f, uint32_t val) {
-    fwrite(&val, 4, 1, f);
-}
-
-// Write 16-bit little-endian integer
-void write_u16_le(FILE* f, uint16_t val) {
-    fwrite(&val, 2, 1, f);
-}
-
-/**
- * Create an MJPG AVI file from an array of sd_image_t images.
- * Images are encoded to JPEG using stb_image_write.
- *
- * @param filename Output AVI file name.
- * @param images Array of input images.
- * @param num_images Number of images in the array.
- * @param fps Frames per second for the video.
- * @param quality JPEG quality (0-100).
- * @return 0 on success, -1 on failure.
- */
-int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality = 90) {
-    if (num_images == 0) {
-        fprintf(stderr, "Error: Image array is empty.\n");
-        return -1;
-    }
-
-    FILE* f = fopen(filename, "wb");
-    if (!f) {
-        perror("Error opening file for writing");
-        return -1;
-    }
-
-    uint32_t width    = images[0].width;
-    uint32_t height   = images[0].height;
-    uint32_t channels = images[0].channel;
-    if (channels != 3 && channels != 4) {
-        fprintf(stderr, "Error: Unsupported channel count: %u\n", channels);
-        fclose(f);
-        return -1;
-    }
-
-    // --- RIFF AVI Header ---
-    fwrite("RIFF", 4, 1, f);
-    long riff_size_pos = ftell(f);
-    write_u32_le(f, 0);  // Placeholder for file size
-    fwrite("AVI ", 4, 1, f);
-
-    // 'hdrl' LIST (header list)
-    fwrite("LIST", 4, 1, f);
-    write_u32_le(f, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
-    fwrite("hdrl", 4, 1, f);
-
-    // 'avih' chunk (AVI main header)
-    fwrite("avih", 4, 1, f);
-    write_u32_le(f, 56);
-    write_u32_le(f, 1000000 / fps);       // Microseconds per frame
-    write_u32_le(f, 0);                   // Max bytes per second
-    write_u32_le(f, 0);                   // Padding granularity
-    write_u32_le(f, 0x110);               // Flags (HASINDEX | ISINTERLEAVED)
-    write_u32_le(f, num_images);          // Total frames
-    write_u32_le(f, 0);                   // Initial frames
-    write_u32_le(f, 1);                   // Number of streams
-    write_u32_le(f, width * height * 3);  // Suggested buffer size
-    write_u32_le(f, width);
-    write_u32_le(f, height);
-    write_u32_le(f, 0);  // Reserved
-    write_u32_le(f, 0);  // Reserved
-    write_u32_le(f, 0);  // Reserved
-    write_u32_le(f, 0);  // Reserved
-
-    // 'strl' LIST (stream list)
-    fwrite("LIST", 4, 1, f);
-    write_u32_le(f, 4 + 8 + 56 + 8 + 40);
-    fwrite("strl", 4, 1, f);
-
-    // 'strh' chunk (stream header)
-    fwrite("strh", 4, 1, f);
-    write_u32_le(f, 56);
-    fwrite("vids", 4, 1, f);              // Stream type: video
-    fwrite("MJPG", 4, 1, f);              // Codec: Motion JPEG
-    write_u32_le(f, 0);                   // Flags
-    write_u16_le(f, 0);                   // Priority
-    write_u16_le(f, 0);                   // Language
-    write_u32_le(f, 0);                   // Initial frames
-    write_u32_le(f, 1);                   // Scale
-    write_u32_le(f, fps);                 // Rate
-    write_u32_le(f, 0);                   // Start
-    write_u32_le(f, num_images);          // Length
-    write_u32_le(f, width * height * 3);  // Suggested buffer size
-    write_u32_le(f, (uint32_t)-1);        // Quality
-    write_u32_le(f, 0);                   // Sample size
-    write_u16_le(f, 0);                   // rcFrame.left
-    write_u16_le(f, 0);                   // rcFrame.top
-    write_u16_le(f, 0);                   // rcFrame.right
-    write_u16_le(f, 0);                   // rcFrame.bottom
-
-    // 'strf' chunk (stream format: BITMAPINFOHEADER)
-    fwrite("strf", 4, 1, f);
-    write_u32_le(f, 40);
-    write_u32_le(f, 40);  // biSize
-    write_u32_le(f, width);
-    write_u32_le(f, height);
-    write_u16_le(f, 1);                   // biPlanes
-    write_u16_le(f, 24);                  // biBitCount
-    fwrite("MJPG", 4, 1, f);              // biCompression (FOURCC)
-    write_u32_le(f, width * height * 3);  // biSizeImage
-    write_u32_le(f, 0);                   // XPelsPerMeter
-    write_u32_le(f, 0);                   // YPelsPerMeter
-    write_u32_le(f, 0);                   // Colors used
-    write_u32_le(f, 0);                   // Colors important
-
-    // 'movi' LIST (video frames)
-    // long movi_list_pos = ftell(f);
-    fwrite("LIST", 4, 1, f);
-    long movi_size_pos = ftell(f);
-    write_u32_le(f, 0);  // Placeholder for movi size
-    fwrite("movi", 4, 1, f);
-
-    avi_index_entry* index = (avi_index_entry*)malloc(sizeof(avi_index_entry) * num_images);
-    if (!index) {
-        fclose(f);
-        return -1;
-    }
-
-    // Encode and write each frame as JPEG
-    struct {
-        uint8_t* buf;
-        size_t size;
-    } jpeg_data;
-
-    for (int i = 0; i < num_images; i++) {
-        jpeg_data.buf  = nullptr;
-        jpeg_data.size = 0;
-
-        // Callback function to collect JPEG data into memory
-        auto write_to_buf = [](void* context, void* data, int size) {
-            auto jd = (decltype(jpeg_data)*)context;
-            jd->buf = (uint8_t*)realloc(jd->buf, jd->size + size);
-            memcpy(jd->buf + jd->size, data, size);
-            jd->size += size;
-        };
-
-        // Encode to JPEG in memory
-        stbi_write_jpg_to_func(
-            write_to_buf,
-            &jpeg_data,
-            images[i].width,
-            images[i].height,
-            channels,
-            images[i].data,
-            quality);
-
-        // Write '00dc' chunk (video frame)
-        fwrite("00dc", 4, 1, f);
-        write_u32_le(f, (uint32_t)jpeg_data.size);
-        index[i].offset = ftell(f) - 8;
-        index[i].size   = (uint32_t)jpeg_data.size;
-        fwrite(jpeg_data.buf, 1, jpeg_data.size, f);
-
-        // Align to even byte size
-        if (jpeg_data.size % 2)
-            fputc(0, f);
-
-        free(jpeg_data.buf);
-    }
-
-    // Finalize 'movi' size
-    long cur_pos   = ftell(f);
-    long movi_size = cur_pos - movi_size_pos - 4;
-    fseek(f, movi_size_pos, SEEK_SET);
-    write_u32_le(f, movi_size);
-    fseek(f, cur_pos, SEEK_SET);
-
-    // Write 'idx1' index
-    fwrite("idx1", 4, 1, f);
-    write_u32_le(f, num_images * 16);
-    for (int i = 0; i < num_images; i++) {
-        fwrite("00dc", 4, 1, f);
-        write_u32_le(f, 0x10);
-        write_u32_le(f, index[i].offset);
-        write_u32_le(f, index[i].size);
-    }
-
-    // Finalize RIFF size
-    cur_pos        = ftell(f);
-    long file_size = cur_pos - riff_size_pos - 4;
-    fseek(f, riff_size_pos, SEEK_SET);
-    write_u32_le(f, file_size);
-    fseek(f, cur_pos, SEEK_SET);
-
-    fclose(f);
-    free(index);
-
-    return 0;
-}
-
-#endif  // __AVI_WRITER_H__

examples/cli/image_metadata.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <iosfwd>
+#include <string>
+
+enum class MetadataOutputFormat {
+    TEXT,
+    JSON,
+};
+
+struct MetadataReadOptions {
+    MetadataOutputFormat output_format = MetadataOutputFormat::TEXT;
+    bool include_raw                   = false;
+    bool brief                         = false;
+    bool include_structural            = false;
+};
+
+bool print_image_metadata(const std::string& image_path,
+                          const MetadataReadOptions& options,
+                          std::ostream& out,
+                          std::string& error);

examples/cli/main.cpp

@@ -15,9 +15,12 @@
 // #include "preprocessing.hpp"
 #include "stable-diffusion.h"
 
-#include "common/common.hpp"
+#include "common/common.h"
+#include "common/media_io.h"
+#include "common/resource_owners.hpp"
+#include "image_metadata.h"
 
-#include "avi_writer.h"
+namespace fs = std::filesystem;
 
 const char* previews_str[] = {
     "none",
@@ -32,6 +35,8 @@ struct SDCliParams {
     SDMode mode             = IMG_GEN;
     std::string output_path = "output.png";
     int output_begin_idx    = -1;
+    std::string image_path;
+    std::string metadata_format = "text";
 
     bool verbose          = false;
     bool canny_preprocess = false;
@@ -44,6 +49,9 @@ struct SDCliParams {
     bool taesd_preview       = false;
     bool preview_noisy       = false;
     bool color               = false;
+    bool metadata_raw        = false;
+    bool metadata_brief      = false;
+    bool metadata_all        = false;
 
     bool normal_exit = false;
 
@@ -53,11 +61,19 @@ struct SDCliParams {
         options.string_options = {
             {"-o",
              "--output",
-             "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png)",
+             "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp",
              &output_path},
+            {"",
+             "--image",
+             "path to the image to inspect (for metadata mode)",
+             &image_path},
+            {"",
+             "--metadata-format",
+             "metadata output format, one of [text, json] (default: text)",
+             &metadata_format},
             {"",
              "--preview-path",
-             "path to write preview image to (default: ./preview.png)",
+             "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
              &preview_path},
         };
 
@@ -97,6 +113,18 @@ struct SDCliParams {
              "--preview-noisy",
              "enables previewing noisy inputs of the models rather than the denoised outputs",
              true, &preview_noisy},
+            {"",
+             "--metadata-raw",
+             "include raw hex previews for unparsed metadata payloads",
+             true, &metadata_raw},
+            {"",
+             "--metadata-brief",
+             "truncate long metadata text values in text output",
+             true, &metadata_brief},
+            {"",
+             "--metadata-all",
+             "include structural/container entries such as IHDR, IDAT, and non-metadata JPEG segments",
+             true, &metadata_all},
 
         };
 
@@ -149,7 +177,7 @@ struct SDCliParams {
         options.manual_options = {
             {"-M",
              "--mode",
-             "run mode, one of [img_gen, vid_gen, upscale, convert], default: img_gen",
+             "run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen",
              on_mode_arg},
             {"",
              "--preview",
@@ -164,12 +192,7 @@ struct SDCliParams {
         return options;
     };
 
-    bool process_and_check() {
-        if (output_path.length() == 0) {
-            LOG_ERROR("error: the following arguments are required: output_path");
-            return false;
-        }
-
+    bool resolve() {
         if (mode == CONVERT) {
             if (output_path == "output.png") {
                 output_path = "output.gguf";
@@ -178,11 +201,43 @@ struct SDCliParams {
         return true;
     }
 
+    bool validate() {
+        if (mode != METADATA) {
+            if (output_path.length() == 0) {
+                LOG_ERROR("error: the following arguments are required: output_path");
+                return false;
+            }
+        } else {
+            if (image_path.empty()) {
+                LOG_ERROR("error: metadata mode needs an image path (--image)");
+                return false;
+            }
+            if (metadata_format != "text" && metadata_format != "json") {
+                LOG_ERROR("error: invalid metadata format %s, must be one of [text, json]",
+                          metadata_format.c_str());
+                return false;
+            }
+        }
+        return true;
+    }
+
+    bool resolve_and_validate() {
+        if (!resolve()) {
+            return false;
+        }
+        if (!validate()) {
+            return false;
+        }
+        return true;
+    }
+
     std::string to_string() const {
         std::ostringstream oss;
         oss << "SDCliParams {\n"
             << "  mode: " << modes_str[mode] << ",\n"
             << "  output_path: \"" << output_path << "\",\n"
+            << "  image_path: \"" << image_path << "\",\n"
+            << "  metadata_format: \"" << metadata_format << "\",\n"
             << "  verbose: " << (verbose ? "true" : "false") << ",\n"
             << "  color: " << (color ? "true" : "false") << ",\n"
             << "  canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
@@ -192,7 +247,10 @@ struct SDCliParams {
             << "  preview_path: \"" << preview_path << "\",\n"
             << "  preview_fps: " << preview_fps << ",\n"
             << "  taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n"
-            << "  preview_noisy: " << (preview_noisy ? "true" : "false") << "\n"
+            << "  preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n"
+            << "  metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n"
+            << "  metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n"
+            << "  metadata_all: " << (metadata_all ? "true" : "false") << "\n"
             << "}";
         return oss.str();
     }
@@ -217,9 +275,13 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
         exit(cli_params.normal_exit ? 0 : 1);
     }
 
-    if (!cli_params.process_and_check() ||
-        !ctx_params.process_and_check(cli_params.mode) ||
-        !gen_params.process_and_check(cli_params.mode, ctx_params.lora_model_dir)) {
+    bool valid = cli_params.resolve_and_validate();
+    if (valid && cli_params.mode != METADATA) {
+        valid = ctx_params.resolve_and_validate(cli_params.mode) &&
+                gen_params.resolve_and_validate(cli_params.mode, ctx_params.lora_model_dir);
+    }
+
+    if (!valid) {
         print_usage(argc, argv, options_vec);
         exit(1);
     }
@@ -231,7 +293,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
 }
 
 bool load_images_from_dir(const std::string dir,
-                          std::vector<sd_image_t>& images,
+                          std::vector<SDImageOwner>& images,
                           int expected_width  = 0,
                           int expected_height = 0,
                           int max_image_num   = 0,
@@ -258,7 +320,7 @@ bool load_images_from_dir(const std::string dir,
         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 (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp" || ext == ".webp") {
             LOG_DEBUG("load image %zu from '%s'", images.size(), path.c_str());
             int width             = 0;
             int height            = 0;
@@ -268,12 +330,12 @@ bool load_images_from_dir(const std::string dir,
                 return false;
             }
 
-            images.push_back({(uint32_t)width,
-                              (uint32_t)height,
-                              3,
-                              image_buffer});
+            images.emplace_back(sd_image_t{(uint32_t)width,
+                                           (uint32_t)height,
+                                           3,
+                                           image_buffer});
 
-            if (max_image_num > 0 && images.size() >= max_image_num) {
+            if (max_image_num > 0 && static_cast<int>(images.size()) >= max_image_num) {
                 break;
             }
         }
@@ -288,9 +350,17 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy,
     // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
     // unused in this app, it will either be always noisy or always denoised here
     if (frame_count == 1) {
-        stbi_write_png(cli_params->preview_path.c_str(), image->width, image->height, image->channel, image->data, 0);
+        if (!write_image_to_file(cli_params->preview_path,
+                                 image->data,
+                                 image->width,
+                                 image->height,
+                                 image->channel)) {
+            LOG_ERROR("save preview image to '%s' failed", cli_params->preview_path.c_str());
+        }
     } else {
-        create_mjpg_avi_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps);
+        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps) != 0) {
+            LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
+        }
     }
 }
 
@@ -340,9 +410,13 @@ bool save_results(const SDCliParams& cli_params,
 
     std::string ext_lower = ext.string();
     std::transform(ext_lower.begin(), ext_lower.end(), ext_lower.begin(), ::tolower);
-    bool is_jpg = (ext_lower == ".jpg" || ext_lower == ".jpeg" || ext_lower == ".jpe");
+    const EncodedImageFormat output_format = encoded_image_format_from_path(out_path.string());
     if (!ext.empty()) {
-        if (is_jpg || ext_lower == ".png") {
+        if (output_format == EncodedImageFormat::JPEG ||
+            output_format == EncodedImageFormat::PNG ||
+            output_format == EncodedImageFormat::WEBP ||
+            ext_lower == ".avi" ||
+            ext_lower == ".webm") {
             base_path.replace_extension();
         }
     }
@@ -360,20 +434,15 @@ bool save_results(const SDCliParams& cli_params,
         std::string params = gen_params.embed_image_metadata
                                  ? get_image_params(ctx_params, gen_params, gen_params.seed + idx)
                                  : "";
-        int ok             = 0;
-        if (is_jpg) {
-            ok = stbi_write_jpg(path.string().c_str(), img.width, img.height, img.channel, img.data, 90, params.size() > 0 ? params.c_str() : nullptr);
-        } else {
-            ok = stbi_write_png(path.string().c_str(), img.width, img.height, img.channel, img.data, 0, params.size() > 0 ? params.c_str() : nullptr);
-        }
+        const bool ok      = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, 90);
         LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
-        return ok != 0;
+        return ok;
     };
 
     int sucessful_reults = 0;
 
     if (std::regex_search(cli_params.output_path, format_specifier_regex)) {
-        if (!is_jpg && ext_lower != ".png")
+        if (output_format == EncodedImageFormat::UNKNOWN)
             ext = ".png";
         fs::path pattern = base_path;
         pattern += ext;
@@ -389,20 +458,20 @@ bool save_results(const SDCliParams& cli_params,
     }
 
     if (cli_params.mode == VID_GEN && num_results > 1) {
-        if (ext_lower != ".avi")
+        if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
             ext = ".avi";
         fs::path video_path = base_path;
         video_path += ext;
-        if (create_mjpg_avi_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
-            LOG_INFO("save result MJPG AVI video to '%s'", video_path.string().c_str());
+        if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
+            LOG_INFO("save result video to '%s'", video_path.string().c_str());
             return true;
         } else {
-            LOG_ERROR("Failed to save result MPG AVI video to '%s'", video_path.string().c_str());
+            LOG_ERROR("Failed to save result video to '%s'", video_path.string().c_str());
             return false;
         }
     }
 
-    if (!is_jpg && ext_lower != ".png")
+    if (output_format == EncodedImageFormat::UNKNOWN)
         ext = ".png";
 
     for (int i = 0; i < num_results; ++i) {
@@ -430,6 +499,27 @@ int main(int argc, const char* argv[]) {
     SDGenerationParams gen_params;
 
     parse_args(argc, argv, cli_params, ctx_params, gen_params);
+    sd_set_log_callback(sd_log_cb, (void*)&cli_params);
+    log_verbose = cli_params.verbose;
+    log_color   = cli_params.color;
+
+    if (cli_params.mode == METADATA) {
+        MetadataReadOptions options;
+        options.output_format      = cli_params.metadata_format == "json"
+                                         ? MetadataOutputFormat::JSON
+                                         : MetadataOutputFormat::TEXT;
+        options.include_raw        = cli_params.metadata_raw;
+        options.brief              = cli_params.metadata_brief;
+        options.include_structural = cli_params.metadata_all;
+
+        std::string error;
+        if (!print_image_metadata(cli_params.image_path, options, std::cout, error)) {
+            LOG_ERROR("%s", error.c_str());
+            return 1;
+        }
+        return 0;
+    }
+
     if (gen_params.video_frames > 4) {
         size_t last_dot_pos   = cli_params.preview_path.find_last_of(".");
         std::string base_path = cli_params.preview_path;
@@ -447,9 +537,6 @@ int main(int argc, const char* argv[]) {
     if (cli_params.preview_method == PREVIEW_PROJ)
         cli_params.preview_fps /= 4;
 
-    sd_set_log_callback(sd_log_cb, (void*)&cli_params);
-    log_verbose = cli_params.verbose;
-    log_color   = cli_params.color;
     sd_set_preview_callback(step_callback,
                             cli_params.preview_method,
                             cli_params.preview_interval,
@@ -485,39 +572,10 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    bool vae_decode_only     = true;
-    sd_image_t init_image    = {0, 0, 3, nullptr};
-    sd_image_t end_image     = {0, 0, 3, nullptr};
-    sd_image_t control_image = {0, 0, 3, nullptr};
-    sd_image_t mask_image    = {0, 0, 1, nullptr};
-    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 image : ref_images) {
-            free(image.data);
-            image.data = nullptr;
-        }
-        ref_images.clear();
-        for (auto image : pmid_images) {
-            free(image.data);
-            image.data = nullptr;
-        }
-        pmid_images.clear();
-        for (auto image : control_frames) {
-            free(image.data);
-            image.data = nullptr;
-        }
-        control_frames.clear();
-    };
+    bool vae_decode_only = true;
 
     auto load_image_and_update_size = [&](const std::string& path,
-                                          sd_image_t& image,
+                                          SDImageOwner& image,
                                           bool resize_image    = true,
                                           int expected_channel = 3) -> bool {
         int expected_width  = 0;
@@ -527,74 +585,73 @@ int main(int argc, const char* argv[]) {
             expected_height = gen_params.height;
         }
 
-        if (!load_sd_image_from_file(&image, path.c_str(), expected_width, expected_height, expected_channel)) {
+        if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) {
             LOG_ERROR("load image from '%s' failed", path.c_str());
-            release_all_resources();
             return false;
         }
 
-        gen_params.set_width_and_height_if_unset(image.width, image.height);
+        gen_params.set_width_and_height_if_unset(image.get().width, image.get().height);
         return true;
     };
 
     if (gen_params.init_image_path.size() > 0) {
         vae_decode_only = false;
-        if (!load_image_and_update_size(gen_params.init_image_path, init_image)) {
+        if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
             return 1;
         }
     }
 
     if (gen_params.end_image_path.size() > 0) {
         vae_decode_only = false;
-        if (!load_image_and_update_size(gen_params.end_image_path, end_image)) {
+        if (!load_image_and_update_size(gen_params.end_image_path, gen_params.end_image)) {
             return 1;
         }
     }
 
     if (gen_params.ref_image_paths.size() > 0) {
         vae_decode_only = false;
+        gen_params.ref_images.clear();
         for (auto& path : gen_params.ref_image_paths) {
-            sd_image_t ref_image = {0, 0, 3, nullptr};
+            SDImageOwner ref_image({0, 0, 3, nullptr});
             if (!load_image_and_update_size(path, ref_image, false)) {
                 return 1;
             }
-            ref_images.push_back(ref_image);
+            gen_params.ref_images.push_back(std::move(ref_image));
         }
     }
 
     if (gen_params.mask_image_path.size() > 0) {
-        if (!load_sd_image_from_file(&mask_image,
+        if (!load_sd_image_from_file(gen_params.mask_image.put(),
                                      gen_params.mask_image_path.c_str(),
                                      gen_params.get_resolved_width(),
                                      gen_params.get_resolved_height(),
                                      1)) {
             LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
-            release_all_resources();
             return 1;
         }
     } else {
-        mask_image.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
-        if (mask_image.data == nullptr) {
+        sd_image_t generated_mask = {0, 0, 1, nullptr};
+        generated_mask.data       = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        if (generated_mask.data == nullptr) {
             LOG_ERROR("malloc mask image failed");
-            release_all_resources();
             return 1;
         }
-        mask_image.width  = gen_params.get_resolved_width();
-        mask_image.height = gen_params.get_resolved_height();
-        memset(mask_image.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        generated_mask.width  = gen_params.get_resolved_width();
+        generated_mask.height = gen_params.get_resolved_height();
+        memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        gen_params.mask_image.reset(generated_mask);
     }
 
     if (gen_params.control_image_path.size() > 0) {
-        if (!load_sd_image_from_file(&control_image,
+        if (!load_sd_image_from_file(gen_params.control_image.put(),
                                      gen_params.control_image_path.c_str(),
                                      gen_params.get_resolved_width(),
                                      gen_params.get_resolved_height())) {
             LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
-            release_all_resources();
             return 1;
         }
         if (cli_params.canny_preprocess) {  // apply preprocessor
-            preprocess_canny(control_image,
+            preprocess_canny(gen_params.control_image.get(),
                              0.08f,
                              0.08f,
                              0.8f,
@@ -604,25 +661,25 @@ int main(int argc, const char* argv[]) {
     }
 
     if (!gen_params.control_video_path.empty()) {
+        gen_params.control_frames.clear();
         if (!load_images_from_dir(gen_params.control_video_path,
-                                  control_frames,
+                                  gen_params.control_frames,
                                   gen_params.get_resolved_width(),
                                   gen_params.get_resolved_height(),
                                   gen_params.video_frames,
                                   cli_params.verbose)) {
-            release_all_resources();
             return 1;
         }
     }
 
     if (!gen_params.pm_id_images_dir.empty()) {
+        gen_params.pm_id_images.clear();
         if (!load_images_from_dir(gen_params.pm_id_images_dir,
-                                  pmid_images,
+                                  gen_params.pm_id_images,
                                   0,
                                   0,
                                   0,
                                   cli_params.verbose)) {
-            release_all_resources();
             return 1;
         }
     }
@@ -633,117 +690,56 @@ int main(int argc, const char* argv[]) {
 
     sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
 
-    sd_image_t* results = nullptr;
-    int num_results     = 0;
+    SDImageVec results;
+    int num_results = 0;
 
     if (cli_params.mode == UPSCALE) {
         num_results = 1;
-        results     = (sd_image_t*)calloc(num_results, sizeof(sd_image_t));
-        if (results == nullptr) {
-            LOG_INFO("failed to allocate results array");
-            release_all_resources();
-            return 1;
-        }
-
-        results[0]      = init_image;
-        init_image.data = nullptr;
+        results.push_back(gen_params.init_image.release());
     } else {
-        sd_ctx_t* sd_ctx = new_sd_ctx(&sd_ctx_params);
+        SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
 
         if (sd_ctx == nullptr) {
             LOG_INFO("new_sd_ctx_t failed");
-            release_all_resources();
             return 1;
         }
 
         if (gen_params.sample_params.sample_method == SAMPLE_METHOD_COUNT) {
-            gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
+            gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get());
         }
 
         if (gen_params.high_noise_sample_params.sample_method == SAMPLE_METHOD_COUNT) {
-            gen_params.high_noise_sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
+            gen_params.high_noise_sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get());
         }
 
         if (gen_params.sample_params.scheduler == SCHEDULER_COUNT) {
-            gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx, gen_params.sample_params.sample_method);
+            gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method);
         }
 
         if (cli_params.mode == IMG_GEN) {
-            sd_img_gen_params_t img_gen_params = {
-                gen_params.lora_vec.data(),
-                static_cast<uint32_t>(gen_params.lora_vec.size()),
-                gen_params.prompt.c_str(),
-                gen_params.negative_prompt.c_str(),
-                gen_params.clip_skip,
-                init_image,
-                ref_images.data(),
-                (int)ref_images.size(),
-                gen_params.auto_resize_ref_image,
-                gen_params.increase_ref_index,
-                mask_image,
-                gen_params.get_resolved_width(),
-                gen_params.get_resolved_height(),
-                gen_params.sample_params,
-                gen_params.strength,
-                gen_params.seed,
-                gen_params.batch_count,
-                control_image,
-                gen_params.control_strength,
-                {
-                    pmid_images.data(),
-                    (int)pmid_images.size(),
-                    gen_params.pm_id_embed_path.c_str(),
-                    gen_params.pm_style_strength,
-                },  // pm_params
-                gen_params.vae_tiling_params,
-                gen_params.cache_params,
-            };
+            sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
 
-            results     = generate_image(sd_ctx, &img_gen_params);
             num_results = gen_params.batch_count;
+            results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
         } else if (cli_params.mode == VID_GEN) {
-            sd_vid_gen_params_t vid_gen_params = {
-                gen_params.lora_vec.data(),
-                static_cast<uint32_t>(gen_params.lora_vec.size()),
-                gen_params.prompt.c_str(),
-                gen_params.negative_prompt.c_str(),
-                gen_params.clip_skip,
-                init_image,
-                end_image,
-                control_frames.data(),
-                (int)control_frames.size(),
-                gen_params.get_resolved_width(),
-                gen_params.get_resolved_height(),
-                gen_params.sample_params,
-                gen_params.high_noise_sample_params,
-                gen_params.moe_boundary,
-                gen_params.strength,
-                gen_params.seed,
-                gen_params.video_frames,
-                gen_params.vace_strength,
-                gen_params.vae_tiling_params,
-                gen_params.cache_params,
-            };
-
-            results = generate_video(sd_ctx, &vid_gen_params, &num_results);
+            sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
+            sd_image_t* generated_video        = generate_video(sd_ctx.get(), &vid_gen_params, &num_results);
+            results.adopt(generated_video, num_results);
         }
 
-        if (results == nullptr) {
+        if (!results) {
             LOG_ERROR("generate failed");
-            free_sd_ctx(sd_ctx);
             return 1;
         }
-
-        free_sd_ctx(sd_ctx);
     }
 
     int upscale_factor = 4;  // unused for RealESRGAN_x4plus_anime_6B.pth
     if (ctx_params.esrgan_path.size() > 0 && gen_params.upscale_repeats > 0) {
-        upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(ctx_params.esrgan_path.c_str(),
-                                                        ctx_params.offload_params_to_cpu,
-                                                        ctx_params.diffusion_conv_direct,
-                                                        ctx_params.n_threads,
-                                                        gen_params.upscale_tile_size);
+        UpscalerCtxPtr upscaler_ctx(new_upscaler_ctx(ctx_params.esrgan_path.c_str(),
+                                                     ctx_params.offload_params_to_cpu,
+                                                     ctx_params.diffusion_conv_direct,
+                                                     ctx_params.n_threads,
+                                                     gen_params.upscale_tile_size));
 
         if (upscaler_ctx == nullptr) {
             LOG_ERROR("new_upscaler_ctx failed");
@@ -752,32 +748,24 @@ int main(int argc, const char* argv[]) {
                 if (results[i].data == nullptr) {
                     continue;
                 }
-                sd_image_t current_image = results[i];
+                SDImageOwner current_image(results[i]);
+                results[i] = {0, 0, 0, nullptr};
                 for (int u = 0; u < gen_params.upscale_repeats; ++u) {
-                    sd_image_t upscaled_image = upscale(upscaler_ctx, current_image, upscale_factor);
-                    if (upscaled_image.data == nullptr) {
+                    SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor));
+                    if (upscaled_image.get().data == nullptr) {
                         LOG_ERROR("upscale failed");
                         break;
                     }
-                    free(current_image.data);
-                    current_image = upscaled_image;
+                    current_image = std::move(upscaled_image);
                 }
-                results[i] = current_image;  // Set the final upscaled image as the result
+                results[i] = current_image.release();  // Set the final upscaled image as the result
             }
         }
     }
 
-    if (!save_results(cli_params, ctx_params, gen_params, results, num_results)) {
+    if (!save_results(cli_params, ctx_params, gen_params, results.data(), num_results)) {
         return 1;
     }
 
-    for (int i = 0; i < num_results; i++) {
-        free(results[i].data);
-        results[i].data = nullptr;
-    }
-    free(results);
-
-    release_all_resources();
-
     return 0;
 }

examples/common/common.h

@@ -0,0 +1,239 @@
+#ifndef __EXAMPLES_COMMON_COMMON_H__
+#define __EXAMPLES_COMMON_COMMON_H__
+
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "log.h"
+#include "resource_owners.hpp"
+#include "stable-diffusion.h"
+
+#define SAFE_STR(s) ((s) ? (s) : "")
+#define BOOL_STR(b) ((b) ? "true" : "false")
+
+extern const char* const modes_str[];
+#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale, metadata"
+
+enum SDMode {
+    IMG_GEN,
+    VID_GEN,
+    CONVERT,
+    UPSCALE,
+    METADATA,
+    MODE_COUNT
+};
+
+struct StringOption {
+    std::string short_name;
+    std::string long_name;
+    std::string desc;
+    std::string* target;
+};
+
+struct IntOption {
+    std::string short_name;
+    std::string long_name;
+    std::string desc;
+    int* target;
+};
+
+struct FloatOption {
+    std::string short_name;
+    std::string long_name;
+    std::string desc;
+    float* target;
+};
+
+struct BoolOption {
+    std::string short_name;
+    std::string long_name;
+    std::string desc;
+    bool keep_true;
+    bool* target;
+};
+
+struct ManualOption {
+    std::string short_name;
+    std::string long_name;
+    std::string desc;
+    std::function<int(int argc, const char** argv, int index)> cb;
+};
+
+struct ArgOptions {
+    std::vector<StringOption> string_options;
+    std::vector<IntOption> int_options;
+    std::vector<FloatOption> float_options;
+    std::vector<BoolOption> bool_options;
+    std::vector<ManualOption> manual_options;
+
+    static std::string wrap_text(const std::string& text, size_t width, size_t indent);
+    void print() const;
+};
+
+bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& options_list);
+bool decode_base64_image(const std::string& encoded_input,
+                         int target_channels,
+                         int expected_width,
+                         int expected_height,
+                         SDImageOwner& out_image);
+
+struct SDContextParams {
+    int n_threads = -1;
+    std::string model_path;
+    std::string clip_l_path;
+    std::string clip_g_path;
+    std::string clip_vision_path;
+    std::string t5xxl_path;
+    std::string llm_path;
+    std::string llm_vision_path;
+    std::string diffusion_model_path;
+    std::string high_noise_diffusion_model_path;
+    std::string vae_path;
+    std::string taesd_path;
+    std::string esrgan_path;
+    std::string control_net_path;
+    std::string embedding_dir;
+    std::string photo_maker_path;
+    sd_type_t wtype = SD_TYPE_COUNT;
+    std::string tensor_type_rules;
+    std::string lora_model_dir = ".";
+
+    std::map<std::string, std::string> embedding_map;
+    std::vector<sd_embedding_t> embedding_vec;
+
+    rng_type_t rng_type         = CUDA_RNG;
+    rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
+    bool offload_params_to_cpu  = false;
+    bool enable_mmap            = false;
+    bool control_net_cpu        = false;
+    bool clip_on_cpu            = false;
+    bool vae_on_cpu             = false;
+    bool flash_attn             = false;
+    bool diffusion_flash_attn   = false;
+    bool diffusion_conv_direct  = false;
+    bool vae_conv_direct        = false;
+
+    bool circular   = false;
+    bool circular_x = false;
+    bool circular_y = false;
+
+    bool chroma_use_dit_mask = true;
+    bool chroma_use_t5_mask  = false;
+    int chroma_t5_mask_pad   = 1;
+
+    bool qwen_image_zero_cond_t = false;
+
+    prediction_t prediction           = PREDICTION_COUNT;
+    lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
+
+    bool force_sdxl_vae_conv_scale = false;
+
+    float flow_shift = INFINITY;
+    ArgOptions get_options();
+    void build_embedding_map();
+    bool resolve(SDMode mode);
+    bool validate(SDMode mode);
+    bool resolve_and_validate(SDMode mode);
+    std::string to_string() const;
+    sd_ctx_params_t to_sd_ctx_params_t(bool vae_decode_only, bool free_params_immediately, bool taesd_preview);
+};
+
+struct SDGenerationParams {
+    // User-facing input fields.
+    std::string prompt;
+    std::string negative_prompt;
+    int clip_skip              = -1;  // <= 0 represents unspecified
+    int width                  = -1;
+    int height                 = -1;
+    int batch_count            = 1;
+    int64_t seed               = 42;
+    float strength             = 0.75f;
+    float control_strength     = 0.9f;
+    bool auto_resize_ref_image = true;
+    bool increase_ref_index    = false;
+    bool embed_image_metadata  = true;
+
+    std::string init_image_path;
+    std::string end_image_path;
+    std::string mask_image_path;
+    std::string control_image_path;
+    std::vector<std::string> ref_image_paths;
+    std::string control_video_path;
+
+    sd_sample_params_t sample_params;
+    sd_sample_params_t high_noise_sample_params;
+    std::vector<int> skip_layers            = {7, 8, 9};
+    std::vector<int> high_noise_skip_layers = {7, 8, 9};
+
+    std::vector<float> custom_sigmas;
+
+    std::string cache_mode;
+    std::string cache_option;
+    std::string scm_mask;
+    bool scm_policy_dynamic = true;
+    sd_cache_params_t cache_params{};
+
+    float moe_boundary                   = 0.875f;
+    int video_frames                     = 1;
+    int fps                              = 16;
+    float vace_strength                  = 1.f;
+    sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
+
+    std::string pm_id_images_dir;
+    std::string pm_id_embed_path;
+    float pm_style_strength = 20.f;
+
+    int upscale_repeats   = 1;
+    int upscale_tile_size = 128;
+
+    std::map<std::string, float> lora_map;
+    std::map<std::string, float> high_noise_lora_map;
+
+    // Derived and normalized fields.
+    std::string prompt_with_lora;  // for metadata record only
+    std::vector<sd_lora_t> lora_vec;
+
+    // Owned execution payload.
+    SDImageOwner init_image;
+    SDImageOwner end_image;
+    std::vector<SDImageOwner> ref_images;
+    SDImageOwner mask_image;
+    SDImageOwner control_image;
+    std::vector<SDImageOwner> pm_id_images;
+    std::vector<SDImageOwner> control_frames;
+
+    // Backing storage for sd_img_gen_params_t view fields.
+    std::vector<sd_image_t> ref_image_views;
+    std::vector<sd_image_t> pm_id_image_views;
+    std::vector<sd_image_t> control_frame_views;
+
+    SDGenerationParams();
+    SDGenerationParams(const SDGenerationParams& other)                = default;
+    SDGenerationParams& operator=(const SDGenerationParams& other)     = default;
+    SDGenerationParams(SDGenerationParams&& other) noexcept            = default;
+    SDGenerationParams& operator=(SDGenerationParams&& other) noexcept = default;
+    ArgOptions get_options();
+    bool from_json_str(const std::string& json_str,
+                       const std::function<std::string(const std::string&)>& lora_path_resolver = {});
+    bool initialize_cache_params();
+    void extract_and_remove_lora(const std::string& lora_model_dir);
+    bool width_and_height_are_set() const;
+    void set_width_and_height_if_unset(int w, int h);
+    int get_resolved_width() const;
+    int get_resolved_height() const;
+    bool resolve(const std::string& lora_model_dir, bool strict = false);
+    bool validate(SDMode mode);
+    bool resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict = false);
+    sd_img_gen_params_t to_sd_img_gen_params_t();
+    sd_vid_gen_params_t to_sd_vid_gen_params_t();
+    std::string to_string() const;
+};
+
+std::string version_string();
+std::string get_image_params(const SDContextParams& ctx_params, const SDGenerationParams& gen_params, int64_t seed);
+
+#endif  // __EXAMPLES_COMMON_COMMON_H__

examples/common/log.cpp

@@ -0,0 +1,115 @@
+#include "log.h"
+
+#include <vector>
+
+bool log_verbose = false;
+bool log_color   = false;
+
+std::string sd_basename(const std::string& path) {
+    size_t pos = path.find_last_of('/');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    pos = path.find_last_of('\\');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    return path;
+}
+
+void print_utf8(FILE* stream, const char* utf8) {
+    if (!utf8) {
+        return;
+    }
+
+#ifdef _WIN32
+    HANDLE h = (stream == stderr)
+                   ? GetStdHandle(STD_ERROR_HANDLE)
+                   : GetStdHandle(STD_OUTPUT_HANDLE);
+
+    DWORD mode;
+    BOOL is_console = GetConsoleMode(h, &mode);
+
+    if (is_console) {
+        int wlen = MultiByteToWideChar(CP_UTF8, 0, utf8, -1, NULL, 0);
+        if (wlen <= 0) {
+            return;
+        }
+
+        std::vector<wchar_t> wbuf(static_cast<size_t>(wlen));
+
+        MultiByteToWideChar(CP_UTF8, 0, utf8, -1, wbuf.data(), wlen);
+
+        DWORD written;
+        WriteConsoleW(h, wbuf.data(), wlen - 1, &written, NULL);
+    } else {
+        DWORD written;
+        WriteFile(h, utf8, (DWORD)strlen(utf8), &written, NULL);
+    }
+#else
+    fputs(utf8, stream);
+#endif
+}
+
+void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool color) {
+    int tag_color;
+    const char* level_str;
+    FILE* out_stream = (level == SD_LOG_ERROR) ? stderr : stdout;
+
+    if (!log || (!verbose && level <= SD_LOG_DEBUG)) {
+        return;
+    }
+
+    switch (level) {
+        case SD_LOG_DEBUG:
+            tag_color = 37;
+            level_str = "DEBUG";
+            break;
+        case SD_LOG_INFO:
+            tag_color = 34;
+            level_str = "INFO";
+            break;
+        case SD_LOG_WARN:
+            tag_color = 35;
+            level_str = "WARN";
+            break;
+        case SD_LOG_ERROR:
+            tag_color = 31;
+            level_str = "ERROR";
+            break;
+        default:
+            tag_color = 33;
+            level_str = "?????";
+            break;
+    }
+
+    if (color) {
+        fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str);
+    } else {
+        fprintf(out_stream, "[%-5s] ", level_str);
+    }
+    print_utf8(out_stream, log);
+    fflush(out_stream);
+}
+
+void example_log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
+    constexpr size_t LOG_BUFFER_SIZE = 4096;
+
+    va_list args;
+    va_start(args, format);
+
+    static char log_buffer[LOG_BUFFER_SIZE + 1];
+    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line);
+
+    if (written >= 0 && written < static_cast<int>(LOG_BUFFER_SIZE)) {
+        vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
+    }
+    size_t len = strlen(log_buffer);
+    if (len == 0 || log_buffer[len - 1] != '\n') {
+        strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len);
+    }
+
+    log_print(level, log_buffer, log_verbose, log_color);
+
+    va_end(args);
+}

examples/common/log.h

@@ -0,0 +1,32 @@
+#ifndef __EXAMPLE_LOG_H__
+#define __EXAMPLE_LOG_H__
+
+#include <cstdarg>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+
+#if defined(_WIN32)
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#include <windows.h>
+#endif  // _WIN32
+
+#include "stable-diffusion.h"
+
+extern bool log_verbose;
+extern bool log_color;
+
+std::string sd_basename(const std::string& path);
+void print_utf8(FILE* stream, const char* utf8);
+void log_print(sd_log_level_t level, const char* log, bool verbose, bool color);
+void example_log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
+
+#define LOG_DEBUG(format, ...) example_log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#define LOG_INFO(format, ...) example_log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#define LOG_WARN(format, ...) example_log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#define LOG_ERROR(format, ...) example_log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
+
+#endif  // __EXAMPLE_LOG_H__

examples/common/media_io.h

@@ -0,0 +1,84 @@
+#ifndef __MEDIA_IO_H__
+#define __MEDIA_IO_H__
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "stable-diffusion.h"
+
+enum class EncodedImageFormat {
+    JPEG,
+    PNG,
+    WEBP,
+    UNKNOWN,
+};
+
+EncodedImageFormat encoded_image_format_from_path(const std::string& path);
+
+std::vector<uint8_t> encode_image_to_vector(EncodedImageFormat format,
+                                            const uint8_t* image,
+                                            int width,
+                                            int height,
+                                            int channels,
+                                            const std::string& parameters = "",
+                                            int quality                   = 90);
+
+bool write_image_to_file(const std::string& path,
+                         const uint8_t* image,
+                         int width,
+                         int height,
+                         int channels,
+                         const std::string& parameters = "",
+                         int quality                   = 90);
+
+uint8_t* load_image_from_file(const char* image_path,
+                              int& width,
+                              int& height,
+                              int expected_width   = 0,
+                              int expected_height  = 0,
+                              int expected_channel = 3);
+
+bool load_sd_image_from_file(sd_image_t* image,
+                             const char* image_path,
+                             int expected_width   = 0,
+                             int expected_height  = 0,
+                             int expected_channel = 3);
+
+uint8_t* load_image_from_memory(const char* image_bytes,
+                                int len,
+                                int& width,
+                                int& height,
+                                int expected_width   = 0,
+                                int expected_height  = 0,
+                                int expected_channel = 3);
+
+int create_mjpg_avi_from_sd_images(const char* filename,
+                                   sd_image_t* images,
+                                   int num_images,
+                                   int fps,
+                                   int quality = 90);
+
+#ifdef SD_USE_WEBP
+int create_animated_webp_from_sd_images(const char* filename,
+                                        sd_image_t* images,
+                                        int num_images,
+                                        int fps,
+                                        int quality = 90);
+#endif
+
+#ifdef SD_USE_WEBM
+int create_webm_from_sd_images(const char* filename,
+                               sd_image_t* images,
+                               int num_images,
+                               int fps,
+                               int quality = 90);
+#endif
+
+int create_video_from_sd_images(const char* filename,
+                                sd_image_t* images,
+                                int num_images,
+                                int fps,
+                                int quality = 90);
+
+#endif  // __MEDIA_IO_H__

examples/common/resource_owners.hpp

@@ -0,0 +1,236 @@
+#ifndef __EXAMPLE_RESOURCE_OWNERS_H__
+#define __EXAMPLE_RESOURCE_OWNERS_H__
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "stable-diffusion.h"
+
+struct FreeDeleter {
+    void operator()(void* ptr) const {
+        free(ptr);
+    }
+};
+
+struct FileCloser {
+    void operator()(FILE* file) const {
+        if (file != nullptr) {
+            fclose(file);
+        }
+    }
+};
+
+struct SDCtxDeleter {
+    void operator()(sd_ctx_t* ctx) const {
+        if (ctx != nullptr) {
+            free_sd_ctx(ctx);
+        }
+    }
+};
+
+struct UpscalerCtxDeleter {
+    void operator()(upscaler_ctx_t* ctx) const {
+        if (ctx != nullptr) {
+            free_upscaler_ctx(ctx);
+        }
+    }
+};
+
+template <typename T>
+using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
+
+using FilePtr        = std::unique_ptr<FILE, FileCloser>;
+using SDCtxPtr       = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
+using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
+
+class SDImageOwner {
+private:
+    static sd_image_t copy_image(const sd_image_t& image) {
+        if (image.data == nullptr) {
+            return {image.width, image.height, image.channel, nullptr};
+        }
+
+        const size_t byte_count = static_cast<size_t>(image.width) * image.height * image.channel;
+        uint8_t* raw_copy       = static_cast<uint8_t*>(malloc(byte_count));
+        if (raw_copy == nullptr) {
+            return {0, 0, 0, nullptr};
+        }
+
+        std::memcpy(raw_copy, image.data, byte_count);
+        return {image.width, image.height, image.channel, raw_copy};
+    }
+
+    sd_image_t image_ = {0, 0, 0, nullptr};
+
+public:
+    SDImageOwner() = default;
+    explicit SDImageOwner(sd_image_t image)
+        : image_(image) {
+    }
+
+    SDImageOwner(const SDImageOwner& other)
+        : image_(copy_image(other.image_)) {
+    }
+
+    SDImageOwner& operator=(const SDImageOwner& other) {
+        if (this != &other) {
+            reset(copy_image(other.image_));
+        }
+        return *this;
+    }
+
+    SDImageOwner(SDImageOwner&& other) noexcept
+        : image_(other.release()) {
+    }
+
+    SDImageOwner& operator=(SDImageOwner&& other) noexcept {
+        if (this != &other) {
+            reset();
+            image_ = other.release();
+        }
+        return *this;
+    }
+
+    ~SDImageOwner() {
+        reset();
+    }
+
+    sd_image_t* put() {
+        if (image_.data != nullptr) {
+            free(image_.data);
+            image_.data = nullptr;
+        }
+        image_.width   = 0;
+        image_.height  = 0;
+        image_.channel = 0;
+        return &image_;
+    }
+
+    sd_image_t& get() {
+        return image_;
+    }
+
+    const sd_image_t& get() const {
+        return image_;
+    }
+
+    sd_image_t release() {
+        sd_image_t image = image_;
+        image_           = {0, 0, 0, nullptr};
+        return image;
+    }
+
+    void reset(sd_image_t image = {0, 0, 0, nullptr}) {
+        if (image_.data != nullptr) {
+            free(image_.data);
+        }
+        image_ = image;
+    }
+};
+
+class SDImageVec {
+private:
+    std::vector<sd_image_t> images_;
+
+public:
+    SDImageVec() = default;
+
+    SDImageVec(const SDImageVec&)            = delete;
+    SDImageVec& operator=(const SDImageVec&) = delete;
+
+    SDImageVec(SDImageVec&& other) noexcept
+        : images_(std::move(other.images_)) {
+    }
+
+    SDImageVec& operator=(SDImageVec&& other) noexcept {
+        if (this != &other) {
+            clear();
+            images_ = std::move(other.images_);
+        }
+        return *this;
+    }
+
+    ~SDImageVec() {
+        clear();
+    }
+
+    void push_back(sd_image_t image) {
+        images_.push_back(image);
+    }
+
+    void push_back(SDImageOwner&& image) {
+        images_.push_back(image.release());
+    }
+
+    void reserve(size_t count) {
+        images_.reserve(count);
+    }
+
+    void adopt(sd_image_t* images, int count) {
+        clear();
+        if (images == nullptr || count <= 0) {
+            free(images);
+            return;
+        }
+
+        images_.reserve(static_cast<size_t>(count));
+        for (int i = 0; i < count; ++i) {
+            images_.push_back(images[i]);
+        }
+        free(images);
+    }
+
+    size_t size() const {
+        return images_.size();
+    }
+
+    bool empty() const {
+        return images_.empty();
+    }
+
+    int count() const {
+        return static_cast<int>(images_.size());
+    }
+
+    explicit operator bool() const {
+        return !images_.empty();
+    }
+
+    sd_image_t* data() {
+        return images_.data();
+    }
+
+    const sd_image_t* data() const {
+        return images_.data();
+    }
+
+    sd_image_t& operator[](size_t index) {
+        return images_[index];
+    }
+
+    const sd_image_t& operator[](size_t index) const {
+        return images_[index];
+    }
+
+    std::vector<sd_image_t>& raw() {
+        return images_;
+    }
+
+    const std::vector<sd_image_t>& raw() const {
+        return images_;
+    }
+
+    void clear() {
+        for (sd_image_t& image : images_) {
+            free(image.data);
+            image.data = nullptr;
+        }
+        images_.clear();
+    }
+};
+
+#endif  // __EXAMPLE_RESOURCE_OWNERS_H__

examples/server/CMakeLists.txt

@@ -50,13 +50,30 @@ if(SD_SERVER_BUILD_FRONTEND AND EXISTS "${FRONTEND_DIR}")
 
         set_source_files_properties("${GENERATED_HTML_HEADER}" PROPERTIES GENERATED TRUE)
     else()
-        message(WARNING "pnpm not found, frontend build disabled")
+        if(EXISTS "${GENERATED_HTML_HEADER}")
+            message(STATUS "pnpm not found; using pre-built frontend header detected at ${GENERATED_HTML_HEADER}")
+            set(HAVE_FRONTEND_BUILD ON)
+            add_custom_target(${TARGET}_frontend)
+        else()
+            message(WARNING "pnpm not found; frontend build disabled.")
+        endif()
     endif()
 else()
     message(STATUS "Frontend disabled or directory not found: ${FRONTEND_DIR}")
 endif()
 
-add_executable(${TARGET} main.cpp)
+add_executable(${TARGET}
+    ../common/common.cpp
+    ../common/log.cpp
+    ../common/media_io.cpp
+    main.cpp
+    runtime.cpp
+    async_jobs.cpp
+    routes_index.cpp
+    routes_openai.cpp
+    routes_sdapi.cpp
+    routes_sdcpp.cpp
+)
 
 if(HAVE_FRONTEND_BUILD)
     add_dependencies(${TARGET} ${TARGET}_frontend)
@@ -70,10 +87,18 @@ endif()
 
 install(TARGETS ${TARGET} RUNTIME)
 target_link_libraries(${TARGET} PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
+if(SD_WEBP)
+    target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBP)
+    target_link_libraries(${TARGET} PRIVATE webp libwebpmux)
+endif()
+if(SD_WEBM)
+    target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBM)
+    target_link_libraries(${TARGET} PRIVATE webm)
+endif()
 
 # due to httplib; it contains a pragma for MSVC, but other things need explicit flags
 if(WIN32 AND NOT MSVC)
     target_link_libraries(${TARGET} PRIVATE ws2_32)
 endif()
 
-target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17)
+target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17)

examples/server/README.md

@@ -1,3 +1,33 @@
+# Example
+
+The following example starts `sd-server` with a standalone diffusion model, VAE, and LLM text encoder:
+
+```
+.\bin\Release\sd-server.exe --diffusion-model  ..\models\diffusion_models\z_image_turbo_bf16.safetensors --vae ..\models\vae\ae.sft  --llm ..\models\text_encoders\qwen_3_4b.safetensors --diffusion-fa --offload-to-cpu -v --cfg-scale 1.0
+```
+
+What this example does:
+
+* `--diffusion-model` selects the standalone diffusion model
+* `--vae` selects the VAE decoder
+* `--llm` selects the text encoder / language model used by this pipeline
+* `--diffusion-fa` enables flash attention in the diffusion model
+* `--offload-to-cpu` reduces VRAM pressure by keeping weights in RAM when possible
+* `-v` enables verbose logging
+* `--cfg-scale 1.0` sets the default CFG scale for generation
+
+After the server starts successfully:
+
+* the web UI is available at `http://127.0.0.1:1234/`
+* the native async API is available under `/sdcpp/v1/...`
+* the compatibility APIs are available under `/v1/...` and `/sdapi/v1/...`
+
+If you want to use a different host or port, pass:
+
+```bash
+--listen-ip <ip> --listen-port <port>
+```
+
 # Frontend
 
 ## Build with Frontend
@@ -8,7 +38,7 @@ The server can optionally build the web frontend and embed it into the binary as
 
 Install the following tools:
 
-* **Node.js** ≥ 22.18
+* **Node.js** ≥ 20
   https://nodejs.org/
 
 * **pnpm** ≥ 10
@@ -54,7 +84,7 @@ and embed the generated frontend into the server binary.
 
 ## Frontend Repository
 
-The web frontend is maintained in a **separate repository**, https://github.com/leejet/stable-ui.
+The web frontend is maintained in a **separate repository**, https://github.com/leejet/sdcpp-webui.
 
 If you want to modify the UI or frontend logic, please submit pull requests to the **frontend repository**.
 
@@ -189,7 +219,7 @@ Default Generation Options:
                                            medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            eta in DDIM, only for DDIM and TCD (default: 0)
+  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
   --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
@@ -197,7 +227,7 @@ Default Generation Options:
   --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM and TCD (default: 0)
+  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>
   --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image

examples/server/api.md

@@ -0,0 +1,794 @@
+# stable-diffusion.cpp Server APIs
+
+This document describes the server-facing APIs exposed by `examples/server`.
+
+The server currently exposes three API families:
+
+- `OpenAI API` under `/v1/...`
+- `Stable Diffusion WebUI API` under `/sdapi/v1/...`
+- `sdcpp API` under `/sdcpp/v1/...`
+
+The `sdcpp API` is the native API surface.
+Its request schema is also the canonical schema for `sd_cpp_extra_args`.
+
+Global LoRA rule:
+
+- Server APIs do not parse LoRA tags embedded inside `prompt`.
+- `<lora:...>` prompt syntax is intentionally unsupported in `OpenAI API`, `sdapi`, and `sdcpp API`.
+- LoRA must be passed through structured API fields when the API supports it.
+
+## Overview
+
+### OpenAI API
+
+Compatibility API shaped like OpenAI image endpoints.
+
+Current generation-related endpoints include:
+
+- `POST /v1/images/generations`
+- `POST /v1/images/edits`
+- `GET /v1/models`
+
+### Stable Diffusion WebUI API
+
+Compatibility API shaped like the AUTOMATIC1111 / WebUI endpoints.
+
+Current generation-related endpoints include:
+
+- `POST /sdapi/v1/txt2img`
+- `POST /sdapi/v1/img2img`
+- `GET /sdapi/v1/loras`
+- `GET /sdapi/v1/samplers`
+- `GET /sdapi/v1/schedulers`
+- `GET /sdapi/v1/sd-models`
+- `GET /sdapi/v1/options`
+
+### sdcpp API
+
+Native async API for `stable-diffusion.cpp`.
+
+Current endpoints include:
+
+- `GET /sdcpp/v1/capabilities`
+- `POST /sdcpp/v1/img_gen`
+- `GET /sdcpp/v1/jobs/{id}`
+- `POST /sdcpp/v1/jobs/{id}/cancel`
+- `POST /sdcpp/v1/vid_gen`
+
+`POST /sdcpp/v1/vid_gen` is currently exposed but returns `501 Not Implemented`.
+
+## `sd_cpp_extra_args`
+
+`sd_cpp_extra_args` is an extension mechanism for the compatibility APIs.
+
+Rules:
+
+- Its JSON schema is the same schema used by the native `sdcpp API`.
+- `OpenAI API` and `sdapi` can embed it inside `prompt`.
+- `sdcpp API` does not need it, because the request body already uses the native schema directly.
+
+Embedding format:
+
+```text
+normal prompt text <sd_cpp_extra_args>{"sample_params":{"sample_steps":28}}</sd_cpp_extra_args>
+```
+
+Behavior:
+
+- The server extracts the JSON block.
+- The JSON block is parsed using the same field rules as the `sdcpp API`.
+- The block is removed from the final prompt before generation.
+
+Intended use:
+
+- extend `OpenAI API` requests with native `stable-diffusion.cpp` controls
+- extend `sdapi` requests with native `stable-diffusion.cpp` controls
+
+Not intended use:
+
+- do not use `sd_cpp_extra_args` with `/sdcpp/v1/*`
+
+## OpenAI API
+
+### Purpose
+
+This family exists for client compatibility.
+
+Use it when you want OpenAI-style request and response shapes.
+
+### Native Extension
+
+`OpenAI API` supports `sd_cpp_extra_args` embedded inside `prompt`.
+
+The embedded JSON follows the `sdcpp API` request schema.
+
+### Supported Fields
+
+#### `POST /v1/images/generations`
+
+Currently supported top-level request fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `prompt` | `string` | Required |
+| `n` | `integer` | Number of images |
+| `size` | `string` | Format `WIDTHxHEIGHT` |
+| `output_format` | `string` | `png`, `jpeg`, or `webp` |
+| `output_compression` | `integer` | Range is clamped to `0..100` |
+
+Native extension fields:
+
+- any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
+
+Response fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `created` | `integer` | Unix timestamp |
+| `output_format` | `string` | Final encoded image format |
+| `data` | `array<object>` | Generated image list |
+| `data[].b64_json` | `string` | Base64-encoded image bytes |
+
+#### `POST /v1/images/edits`
+
+Currently supported multipart form fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `prompt` | `string` | Required |
+| `image[]` | `file[]` | Preferred image upload field |
+| `image` | `file` | Legacy single-image upload field |
+| `mask` | `file` | Optional mask image |
+| `n` | `integer` | Number of images |
+| `size` | `string` | Format `WIDTHxHEIGHT` |
+| `output_format` | `string` | `png` or `jpeg` |
+| `output_compression` | `integer` | Range is clamped to `0..100` |
+
+Native extension fields:
+
+- any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
+
+Response fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `created` | `integer` | Unix timestamp |
+| `output_format` | `string` | Final encoded image format |
+| `data` | `array<object>` | Generated image list |
+| `data[].b64_json` | `string` | Base64-encoded image bytes |
+
+#### `GET /v1/models`
+
+Response fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `data` | `array<object>` | Available local models |
+| `data[].id` | `string` | Currently fixed to `sd-cpp-local` |
+| `data[].object` | `string` | Currently fixed to `model` |
+| `data[].owned_by` | `string` | Currently fixed to `local` |
+
+### Output Options
+
+`OpenAI API` supports response serialization controls such as:
+
+- `output_format`
+- `output_compression`
+
+### Notes
+
+- `OpenAI API` is synchronous from the HTTP client's perspective.
+- Native async job polling is not exposed through this family.
+- Prompt-embedded `<lora:...>` tags are intentionally unsupported.
+
+## Stable Diffusion WebUI API
+
+### Purpose
+
+This family exists for client compatibility with WebUI-style tools.
+
+Use it when you want `txt2img` / `img2img`-style endpoints and response shapes.
+
+### Native Extension
+
+`sdapi` supports `sd_cpp_extra_args` embedded inside `prompt`.
+
+The embedded JSON follows the `sdcpp API` request schema.
+
+This allows `sdapi` clients to use native `stable-diffusion.cpp` controls without changing the outer request format.
+
+### Supported Fields
+
+#### `POST /sdapi/v1/txt2img`
+
+Currently supported request fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `prompt` | `string` | Required |
+| `negative_prompt` | `string` | Optional |
+| `width` | `integer` | Positive image width |
+| `height` | `integer` | Positive image height |
+| `steps` | `integer` | Sampling steps |
+| `cfg_scale` | `number` | Text CFG scale |
+| `seed` | `integer` | `-1` means random |
+| `batch_size` | `integer` | Number of images |
+| `clip_skip` | `integer` | Optional |
+| `sampler_name` | `string` | WebUI sampler name |
+| `scheduler` | `string` | Scheduler name |
+| `lora` | `array<object>` | Structured LoRA list |
+| `extra_images` | `array<string>` | Base64 or data URL images |
+
+Native extension fields:
+
+- any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
+
+Response fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `images` | `array<string>` | Base64-encoded PNG images |
+| `parameters` | `object` | Echo of the parsed outer request body |
+| `info` | `string` | Currently empty string |
+
+#### `POST /sdapi/v1/img2img`
+
+Currently supported request fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| all currently supported `txt2img` fields | same as above | Reused |
+| `init_images` | `array<string>` | Base64 or data URL images |
+| `mask` | `string` | Base64 or data URL image |
+| `inpainting_mask_invert` | `integer` or `boolean` | Treated as invert flag |
+| `denoising_strength` | `number` | Clamped to `0.0..1.0` |
+
+Native extension fields:
+
+- any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
+
+Response fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `images` | `array<string>` | Base64-encoded PNG images |
+| `parameters` | `object` | Echo of the parsed outer request body |
+| `info` | `string` | Currently empty string |
+
+#### Discovery / Compatibility Endpoints
+
+Currently exposed:
+
+- `GET /sdapi/v1/loras`
+- `GET /sdapi/v1/samplers`
+- `GET /sdapi/v1/schedulers`
+- `GET /sdapi/v1/sd-models`
+- `GET /sdapi/v1/options`
+
+Response fields:
+
+`GET /sdapi/v1/loras`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | Display name derived from file stem |
+| `[].path` | `string` | Relative path under the configured LoRA directory |
+
+`GET /sdapi/v1/samplers`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | Sampler name |
+| `[].aliases` | `array<string>` | Currently contains the same single sampler name |
+| `[].options` | `object` | Currently empty object |
+
+`GET /sdapi/v1/schedulers`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].name` | `string` | Scheduler name |
+| `[].label` | `string` | Same value as `name` |
+
+`GET /sdapi/v1/sd-models`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `[].title` | `string` | Model stem |
+| `[].model_name` | `string` | Same value as `title` |
+| `[].filename` | `string` | Model filename |
+| `[].hash` | `string` | Placeholder compatibility value |
+| `[].sha256` | `string` | Placeholder compatibility value |
+| `[].config` | `null` | Currently always null |
+
+`GET /sdapi/v1/options`
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `samples_format` | `string` | Currently fixed to `png` |
+| `sd_model_checkpoint` | `string` | Model stem |
+
+### Notes
+
+- `sdapi` is synchronous from the HTTP client's perspective.
+- Prompt-embedded `<lora:...>` tags are intentionally unsupported.
+
+## sdcpp API
+
+### Purpose
+
+This is the native `stable-diffusion.cpp` API.
+
+Use it when you want:
+
+- async job submission
+- explicit native parameter control
+- frontend-oriented capability discovery
+
+### Job Model
+
+All async generation requests create a job.
+
+Job states:
+
+- `queued`
+- `generating`
+- `completed`
+- `failed`
+- `cancelled`
+
+Common job shape:
+
+```json
+{
+  "id": "job_01HTXYZABC",
+  "kind": "img_gen",
+  "status": "queued",
+  "created": 1775401200,
+  "started": null,
+  "completed": null,
+  "queue_position": 2,
+  "result": null,
+  "error": null
+}
+```
+
+Field types:
+
+| Field | Type |
+| --- | --- |
+| `id` | `string` |
+| `kind` | `string` |
+| `status` | `string` |
+| `created` | `integer` |
+| `started` | `integer \| null` |
+| `completed` | `integer \| null` |
+| `queue_position` | `integer` |
+| `result` | `object \| null` |
+| `error` | `object \| null` |
+
+### Endpoints
+
+#### `GET /sdcpp/v1/capabilities`
+
+Returns frontend-friendly capability metadata.
+
+Typical contents:
+
+| Field | Type |
+| --- | --- |
+| `model` | `object` |
+| `defaults` | `object` |
+| `loras` | `array<object>` |
+| `samplers` | `array<string>` |
+| `schedulers` | `array<string>` |
+| `output_formats` | `array<string>` |
+| `limits` | `object` |
+| `features` | `object` |
+
+Nested fields currently returned:
+
+`model`
+
+| Field | Type |
+| --- | --- |
+| `model.name` | `string` |
+| `model.stem` | `string` |
+| `model.path` | `string` |
+
+`defaults`
+
+| Field | Type |
+| --- | --- |
+| `defaults.prompt` | `string` |
+| `defaults.negative_prompt` | `string` |
+| `defaults.clip_skip` | `integer` |
+| `defaults.width` | `integer` |
+| `defaults.height` | `integer` |
+| `defaults.strength` | `number` |
+| `defaults.seed` | `integer` |
+| `defaults.batch_count` | `integer` |
+| `defaults.auto_resize_ref_image` | `boolean` |
+| `defaults.increase_ref_index` | `boolean` |
+| `defaults.control_strength` | `number` |
+| `defaults.sample_params` | `object` |
+| `defaults.sample_params.scheduler` | `string` |
+| `defaults.sample_params.sample_method` | `string` |
+| `defaults.sample_params.sample_steps` | `integer` |
+| `defaults.sample_params.eta` | `number \| null` |
+| `defaults.sample_params.shifted_timestep` | `integer` |
+| `defaults.sample_params.flow_shift` | `number \| null` |
+| `defaults.sample_params.guidance` | `object` |
+| `defaults.sample_params.guidance.txt_cfg` | `number` |
+| `defaults.sample_params.guidance.img_cfg` | `number \| null` |
+| `defaults.sample_params.guidance.distilled_guidance` | `number` |
+| `defaults.sample_params.guidance.slg` | `object` |
+| `defaults.sample_params.guidance.slg.layers` | `array<integer>` |
+| `defaults.sample_params.guidance.slg.layer_start` | `number` |
+| `defaults.sample_params.guidance.slg.layer_end` | `number` |
+| `defaults.sample_params.guidance.slg.scale` | `number` |
+| `defaults.vae_tiling_params` | `object` |
+| `defaults.vae_tiling_params.enabled` | `boolean` |
+| `defaults.vae_tiling_params.tile_size_x` | `integer` |
+| `defaults.vae_tiling_params.tile_size_y` | `integer` |
+| `defaults.vae_tiling_params.target_overlap` | `number` |
+| `defaults.vae_tiling_params.rel_size_x` | `number` |
+| `defaults.vae_tiling_params.rel_size_y` | `number` |
+| `defaults.cache_mode` | `string` |
+| `defaults.cache_option` | `string` |
+| `defaults.scm_mask` | `string` |
+| `defaults.scm_policy_dynamic` | `boolean` |
+| `defaults.output_format` | `string` |
+| `defaults.output_compression` | `integer` |
+
+`loras`
+
+| Field | Type |
+| --- | --- |
+| `loras[].name` | `string` |
+| `loras[].path` | `string` |
+
+`limits`
+
+| Field | Type |
+| --- | --- |
+| `limits.min_width` | `integer` |
+| `limits.max_width` | `integer` |
+| `limits.min_height` | `integer` |
+| `limits.max_height` | `integer` |
+| `limits.max_batch_count` | `integer` |
+| `limits.max_queue_size` | `integer` |
+
+`features`
+
+| Field | Type |
+| --- | --- |
+| `features.init_image` | `boolean` |
+| `features.mask_image` | `boolean` |
+| `features.control_image` | `boolean` |
+| `features.ref_images` | `boolean` |
+| `features.lora` | `boolean` |
+| `features.vae_tiling` | `boolean` |
+| `features.cache` | `boolean` |
+| `features.cancel_queued` | `boolean` |
+| `features.cancel_generating` | `boolean` |
+
+#### `POST /sdcpp/v1/img_gen`
+
+Submits an async image generation job.
+
+Successful submission returns `202 Accepted`.
+
+Example response:
+
+```json
+{
+  "id": "job_01HTXYZABC",
+  "kind": "img_gen",
+  "status": "queued",
+  "created": 1775401200,
+  "poll_url": "/sdcpp/v1/jobs/job_01HTXYZABC"
+}
+```
+
+Response fields:
+
+| Field | Type |
+| --- | --- |
+| `id` | `string` |
+| `kind` | `string` |
+| `status` | `string` |
+| `created` | `integer` |
+| `poll_url` | `string` |
+
+#### `GET /sdcpp/v1/jobs/{id}`
+
+Returns current job status.
+
+Typical status codes:
+
+- `200 OK`
+- `404 Not Found`
+- `410 Gone`
+
+#### `POST /sdcpp/v1/jobs/{id}/cancel`
+
+Attempts to cancel an accepted job.
+
+Typical status codes:
+
+- `200 OK`
+- `404 Not Found`
+- `409 Conflict`
+- `410 Gone`
+
+### Canonical Request Schema
+
+The `sdcpp API` request body is the canonical native schema.
+
+Example:
+
+```json
+{
+  "prompt": "a cat sitting on a chair",
+  "negative_prompt": "",
+  "clip_skip": -1,
+  "width": 1024,
+  "height": 1024,
+  "strength": 0.75,
+  "seed": -1,
+  "batch_count": 1,
+  "auto_resize_ref_image": true,
+  "increase_ref_index": false,
+  "control_strength": 0.9,
+  "embed_image_metadata": true,
+
+  "init_image": null,
+  "ref_images": [],
+  "mask_image": null,
+  "control_image": null,
+
+  "sample_params": {
+    "scheduler": "discrete",
+    "sample_method": "euler_a",
+    "sample_steps": 28,
+    "eta": 1.0,
+    "shifted_timestep": 0,
+    "custom_sigmas": [],
+    "flow_shift": 0.0,
+    "guidance": {
+      "txt_cfg": 7.0,
+      "img_cfg": 7.0,
+      "distilled_guidance": 3.5,
+      "slg": {
+        "layers": [7, 8, 9],
+        "layer_start": 0.01,
+        "layer_end": 0.2,
+        "scale": 0.0
+      }
+    }
+  },
+
+  "lora": [],
+
+  "vae_tiling_params": {
+    "enabled": false,
+    "tile_size_x": 0,
+    "tile_size_y": 0,
+    "target_overlap": 0.5,
+    "rel_size_x": 0.0,
+    "rel_size_y": 0.0
+  },
+
+  "cache_mode": "disabled",
+  "cache_option": "",
+  "scm_mask": "",
+  "scm_policy_dynamic": true,
+
+  "output_format": "png",
+  "output_compression": 100
+}
+```
+
+### LoRA Rules
+
+- The server only accepts explicit LoRA entries from the `lora` field.
+- Prompt-embedded `<lora:...>` tags are intentionally unsupported.
+- Clients should resolve LoRA usage through the structured `lora` array.
+
+### Image Encoding Rules
+
+Any image field accepts:
+
+- a raw base64 string, or
+- a data URL such as `data:image/png;base64,...`
+
+Channel expectations:
+
+- `init_image`: 3 channels
+- `ref_images[]`: 3 channels
+- `control_image`: 3 channels
+- `mask_image`: 1 channel
+
+If omitted or null:
+
+- single-image fields map to an empty `sd_image_t`
+- array fields map to `nullptr + count = 0`
+
+### Field Mapping Summary
+
+Top-level scalar fields:
+
+| Field | Type |
+| --- | --- |
+| `prompt` | `string` |
+| `negative_prompt` | `string` |
+| `clip_skip` | `integer` |
+| `width` | `integer` |
+| `height` | `integer` |
+| `strength` | `number` |
+| `seed` | `integer` |
+| `batch_count` | `integer` |
+| `auto_resize_ref_image` | `boolean` |
+| `increase_ref_index` | `boolean` |
+| `control_strength` | `number` |
+| `embed_image_metadata` | `boolean` |
+
+Image fields:
+
+| Field | Type |
+| --- | --- |
+| `init_image` | `string \| null` |
+| `ref_images` | `array<string>` |
+| `mask_image` | `string \| null` |
+| `control_image` | `string \| null` |
+
+LoRA fields:
+
+| Field | Type |
+| --- | --- |
+| `lora[].path` | `string` |
+| `lora[].multiplier` | `number` |
+| `lora[].is_high_noise` | `boolean` |
+
+Sampling fields:
+
+| Field | Type |
+| --- | --- |
+| `sample_params.scheduler` | `string` |
+| `sample_params.sample_method` | `string` |
+| `sample_params.sample_steps` | `integer` |
+| `sample_params.eta` | `number` |
+| `sample_params.shifted_timestep` | `integer` |
+| `sample_params.custom_sigmas` | `array<number>` |
+| `sample_params.flow_shift` | `number` |
+| `sample_params.guidance.txt_cfg` | `number` |
+| `sample_params.guidance.img_cfg` | `number` |
+| `sample_params.guidance.distilled_guidance` | `number` |
+| `sample_params.guidance.slg.layers` | `array<integer>` |
+| `sample_params.guidance.slg.layer_start` | `number` |
+| `sample_params.guidance.slg.layer_end` | `number` |
+| `sample_params.guidance.slg.scale` | `number` |
+
+Other native fields:
+
+| Field | Type |
+| --- | --- |
+| `vae_tiling_params` | `object` |
+| `cache_mode` | `string` |
+| `cache_option` | `string` |
+| `scm_mask` | `string` |
+| `scm_policy_dynamic` | `boolean` |
+
+HTTP-only output fields:
+
+| Field | Type |
+| --- | --- |
+| `output_format` | `string` |
+| `output_compression` | `integer` |
+
+### Optional Field Semantics
+
+Clients should preserve unset semantics for optional sampling fields.
+
+If a user has not explicitly provided one of these fields, the client should omit it instead of injecting a guessed fallback:
+
+- `sample_params.scheduler`
+- `sample_params.sample_method`
+- `sample_params.eta`
+- `sample_params.flow_shift`
+- `sample_params.guidance.img_cfg`
+
+### Completion Result
+
+Example completed job:
+
+```json
+{
+  "id": "job_01HTXYZABC",
+  "kind": "img_gen",
+  "status": "completed",
+  "created": 1775401200,
+  "started": 1775401203,
+  "completed": 1775401215,
+  "queue_position": 0,
+  "result": {
+    "output_format": "png",
+    "images": [
+      {
+        "index": 0,
+        "b64_json": "iVBORw0KGgoAAA..."
+      }
+    ]
+  },
+  "error": null
+}
+```
+
+### Failure Result
+
+Example failed job:
+
+```json
+{
+  "id": "job_01HTXYZABC",
+  "kind": "img_gen",
+  "status": "failed",
+  "created": 1775401200,
+  "started": 1775401203,
+  "completed": 1775401204,
+  "queue_position": 0,
+  "result": null,
+  "error": {
+    "code": "generation_failed",
+    "message": "generate_image returned empty results"
+  }
+}
+```
+
+### Cancelled Result
+
+Example cancelled job:
+
+```json
+{
+  "id": "job_01HTXYZABC",
+  "kind": "img_gen",
+  "status": "cancelled",
+  "created": 1775401200,
+  "started": null,
+  "completed": 1775401202,
+  "queue_position": 0,
+  "result": null,
+  "error": {
+    "code": "cancelled",
+    "message": "job cancelled by client"
+  }
+}
+```
+
+### Validation and Retention
+
+Recommended behavior:
+
+- malformed JSON returns `400`
+- invalid image payloads return `400`
+- invalid parameter structure returns `400`
+- queue full returns `429` or `503`
+- accepted runtime failures transition the job to `failed`
+- unsupported in-progress cancellation may return `409`
+
+Recommended retention controls:
+
+- pending job limit
+- completed job TTL
+- failed job TTL
+
+### Future `vid_gen`
+
+Future `vid_gen` should reuse the same async job model:
+
+- `POST /sdcpp/v1/vid_gen`
+- `GET /sdcpp/v1/jobs/{id}`
+- `POST /sdcpp/v1/jobs/{id}/cancel`
+
+Its request body should mirror `sd_vid_gen_params_t` in the same way that `img_gen` mirrors `sd_img_gen_params_t`.

examples/server/async_jobs.cpp

@@ -0,0 +1,275 @@
+// Extracted from main.cpp during server refactor.
+
+#include "async_jobs.h"
+
+#include <iomanip>
+#include <sstream>
+
+#include "common/log.h"
+#include "common/media_io.h"
+#include "common/resource_owners.hpp"
+
+const char* async_job_kind_name(AsyncJobKind kind) {
+    switch (kind) {
+        case AsyncJobKind::ImgGen:
+            return "img_gen";
+        case AsyncJobKind::VidGen:
+            return "vid_gen";
+        default:
+            return "img_gen";
+    }
+}
+
+const char* async_job_status_name(AsyncJobStatus status) {
+    switch (status) {
+        case AsyncJobStatus::Queued:
+            return "queued";
+        case AsyncJobStatus::Generating:
+            return "generating";
+        case AsyncJobStatus::Completed:
+            return "completed";
+        case AsyncJobStatus::Failed:
+            return "failed";
+        case AsyncJobStatus::Cancelled:
+            return "cancelled";
+        default:
+            return "failed";
+    }
+}
+
+void purge_expired_jobs(AsyncJobManager& manager) {
+    const int64_t now = unix_timestamp_now();
+
+    for (auto it = manager.expired_jobs.begin(); it != manager.expired_jobs.end();) {
+        if (it->second <= now) {
+            it = manager.expired_jobs.erase(it);
+        } else {
+            ++it;
+        }
+    }
+
+    for (auto it = manager.jobs.begin(); it != manager.jobs.end();) {
+        const auto& job = it->second;
+        if (job->completed_at == 0) {
+            ++it;
+            continue;
+        }
+
+        int64_t ttl_seconds = job->status == AsyncJobStatus::Completed
+                                  ? manager.completed_ttl_seconds
+                                  : manager.failed_ttl_seconds;
+        if (now - job->completed_at >= ttl_seconds) {
+            manager.expired_jobs[job->id] = now + std::max<int64_t>(ttl_seconds, 60);
+            it                            = manager.jobs.erase(it);
+        } else {
+            ++it;
+        }
+    }
+}
+
+size_t count_pending_jobs(const AsyncJobManager& manager) {
+    size_t pending = 0;
+    for (const auto& entry : manager.jobs) {
+        if (entry.second->status == AsyncJobStatus::Queued ||
+            entry.second->status == AsyncJobStatus::Generating) {
+            ++pending;
+        }
+    }
+    return pending;
+}
+
+std::string make_async_job_id(AsyncJobManager& manager) {
+    std::ostringstream oss;
+    oss << "job_" << std::hex << unix_timestamp_now() << "_" << std::setw(8)
+        << std::setfill('0') << manager.next_id++;
+    return oss.str();
+}
+
+bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) {
+    auto new_end = std::remove(manager.queue.begin(), manager.queue.end(), job.id);
+    if (new_end == manager.queue.end()) {
+        return false;
+    }
+
+    manager.queue.erase(new_end, manager.queue.end());
+    job.status       = AsyncJobStatus::Cancelled;
+    job.completed_at = unix_timestamp_now();
+    job.result_images_b64.clear();
+    job.error_code    = "cancelled";
+    job.error_message = "job cancelled by client";
+    return true;
+}
+
+json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJob& job) {
+    json result;
+    result["id"]             = job.id;
+    result["kind"]           = async_job_kind_name(job.kind);
+    result["status"]         = async_job_status_name(job.status);
+    result["created"]        = job.created_at;
+    result["started"]        = job.started_at == 0 ? json(nullptr) : json(job.started_at);
+    result["completed"]      = job.completed_at == 0 ? json(nullptr) : json(job.completed_at);
+    result["queue_position"] = 0;
+
+    if (job.status == AsyncJobStatus::Queued) {
+        size_t position = 1;
+        for (const auto& queued_id : manager.queue) {
+            if (queued_id == job.id) {
+                result["queue_position"] = position;
+                break;
+            }
+            ++position;
+        }
+    }
+
+    if (job.status == AsyncJobStatus::Completed) {
+        json images = json::array();
+        for (size_t i = 0; i < job.result_images_b64.size(); ++i) {
+            images.push_back({{"index", i}, {"b64_json", job.result_images_b64[i]}});
+        }
+        result["result"] = {
+            {"output_format", job.img_gen.output_format},
+            {"images", images},
+        };
+        result["error"] = nullptr;
+    } else if (job.status == AsyncJobStatus::Failed ||
+               job.status == AsyncJobStatus::Cancelled) {
+        result["result"] = nullptr;
+        result["error"]  = {
+             {"code",
+             job.error_code.empty()
+                  ? (job.status == AsyncJobStatus::Cancelled ? "cancelled" : "generation_failed")
+                  : job.error_code},
+             {"message", job.error_message},
+        };
+    } else {
+        result["result"] = nullptr;
+        result["error"]  = nullptr;
+    }
+
+    return result;
+}
+
+bool execute_img_gen_job(ServerRuntime& runtime,
+                         AsyncGenerationJob& job,
+                         std::vector<std::string>& output_images,
+                         std::string& error_message) {
+    sd_img_gen_params_t params = job.img_gen.to_sd_img_gen_params_t();
+
+    SDImageVec results;
+    int num_results = 0;
+
+    {
+        std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
+        sd_image_t* raw_results = generate_image(runtime.sd_ctx, &params);
+        num_results             = params.batch_count;
+        results.adopt(raw_results, num_results);
+    }
+
+    if (results.empty() || num_results <= 0) {
+        error_message = "generate_image returned no results";
+        return false;
+    }
+
+    EncodedImageFormat encoded_format = EncodedImageFormat::PNG;
+    if (job.img_gen.output_format == "jpeg") {
+        encoded_format = EncodedImageFormat::JPEG;
+    } else if (job.img_gen.output_format == "webp") {
+        encoded_format = EncodedImageFormat::WEBP;
+    }
+
+    for (int i = 0; i < num_results; ++i) {
+        if (results[i].data == nullptr) {
+            continue;
+        }
+
+        const std::string metadata = job.img_gen.gen_params.embed_image_metadata
+                                         ? get_image_params(*runtime.ctx_params,
+                                                            job.img_gen.gen_params,
+                                                            job.img_gen.gen_params.seed + i)
+                                         : "";
+        auto image_bytes           = encode_image_to_vector(encoded_format,
+                                                            results[i].data,
+                                                            results[i].width,
+                                                            results[i].height,
+                                                            results[i].channel,
+                                                            metadata,
+                                                            job.img_gen.output_compression);
+        if (image_bytes.empty()) {
+            continue;
+        }
+        output_images.push_back(base64_encode(image_bytes));
+    }
+
+    if (output_images.empty()) {
+        error_message = "generate_image returned empty encoded outputs";
+        return false;
+    }
+
+    return true;
+}
+
+void async_job_worker(ServerRuntime& runtime) {
+    AsyncJobManager& manager = *runtime.async_job_manager;
+
+    while (true) {
+        std::shared_ptr<AsyncGenerationJob> job;
+        {
+            std::unique_lock<std::mutex> lock(manager.mutex);
+            manager.cv.wait(lock, [&]() { return manager.stop || !manager.queue.empty(); });
+
+            if (manager.stop && manager.queue.empty()) {
+                break;
+            }
+
+            purge_expired_jobs(manager);
+            if (manager.queue.empty()) {
+                continue;
+            }
+
+            const std::string job_id = manager.queue.front();
+            manager.queue.pop_front();
+
+            auto it = manager.jobs.find(job_id);
+            if (it == manager.jobs.end()) {
+                continue;
+            }
+
+            job             = it->second;
+            job->status     = AsyncJobStatus::Generating;
+            job->started_at = unix_timestamp_now();
+        }
+
+        std::vector<std::string> output_images;
+        std::string error_message;
+        bool ok = false;
+
+        if (job->kind == AsyncJobKind::ImgGen) {
+            ok = execute_img_gen_job(runtime, *job, output_images, error_message);
+        } else {
+            error_message = "unsupported job kind";
+        }
+
+        {
+            std::lock_guard<std::mutex> lock(manager.mutex);
+            auto it = manager.jobs.find(job->id);
+            if (it == manager.jobs.end()) {
+                continue;
+            }
+
+            job->completed_at = unix_timestamp_now();
+            if (ok) {
+                job->status            = AsyncJobStatus::Completed;
+                job->result_images_b64 = std::move(output_images);
+                job->error_code.clear();
+                job->error_message.clear();
+            } else {
+                job->status        = AsyncJobStatus::Failed;
+                job->error_code    = "generation_failed";
+                job->error_message = error_message.empty() ? "unknown generation error" : error_message;
+                job->result_images_b64.clear();
+            }
+
+            purge_expired_jobs(manager);
+        }
+    }
+}

examples/server/async_jobs.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <condition_variable>
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+
+#include "runtime.h"
+
+enum class AsyncJobKind {
+    ImgGen,
+    VidGen,
+};
+
+enum class AsyncJobStatus {
+    Queued,
+    Generating,
+    Completed,
+    Failed,
+    Cancelled,
+};
+
+const char* async_job_kind_name(AsyncJobKind kind);
+const char* async_job_status_name(AsyncJobStatus status);
+
+struct AsyncGenerationJob {
+    std::string id;
+    AsyncJobKind kind     = AsyncJobKind::ImgGen;
+    AsyncJobStatus status = AsyncJobStatus::Queued;
+    int64_t created_at    = unix_timestamp_now();
+    int64_t started_at    = 0;
+    int64_t completed_at  = 0;
+    ImgGenJobRequest img_gen;
+    std::vector<std::string> result_images_b64;
+    std::string error_code;
+    std::string error_message;
+};
+
+struct AsyncJobManager {
+    std::mutex mutex;
+    std::condition_variable cv;
+    std::unordered_map<std::string, std::shared_ptr<AsyncGenerationJob>> jobs;
+    std::unordered_map<std::string, int64_t> expired_jobs;
+    std::deque<std::string> queue;
+    uint64_t next_id              = 0;
+    bool stop                     = false;
+    size_t max_pending_jobs       = 64;
+    int64_t completed_ttl_seconds = 600;
+    int64_t failed_ttl_seconds    = 600;
+};
+
+void purge_expired_jobs(AsyncJobManager& manager);
+size_t count_pending_jobs(const AsyncJobManager& manager);
+std::string make_async_job_id(AsyncJobManager& manager);
+bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job);
+json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJob& job);
+bool execute_img_gen_job(ServerRuntime& runtime,
+                         AsyncGenerationJob& job,
+                         std::vector<std::string>& output_images,
+                         std::string& error_message);
+void async_job_worker(ServerRuntime& runtime);

examples/server/routes.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <string>
+
+#include "httplib.h"
+#include "runtime.h"
+
+void register_index_endpoints(httplib::Server& svr, const SDSvrParams& svr_params, const std::string& index_html);
+void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt);
+void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt);
+void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt);

examples/server/routes_index.cpp

@@ -0,0 +1,22 @@
+#include "routes.h"
+
+#include <fstream>
+#include <iterator>
+
+void register_index_endpoints(httplib::Server& svr, const SDSvrParams& svr_params, const std::string& index_html) {
+    const std::string serve_html_path = svr_params.serve_html_path;
+    svr.Get("/", [serve_html_path, index_html](const httplib::Request&, httplib::Response& res) {
+        if (!serve_html_path.empty()) {
+            std::ifstream file(serve_html_path);
+            if (file) {
+                std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
+                res.set_content(content, "text/html");
+            } else {
+                res.status = 500;
+                res.set_content("Error: Unable to read HTML file", "text/plain");
+            }
+        } else {
+            res.set_content(index_html, "text/html");
+        }
+    });
+}

examples/server/routes_openai.cpp

@@ -0,0 +1,376 @@
+#include "routes.h"
+
+#include <algorithm>
+#include <ctime>
+#include <regex>
+
+#include "common/common.h"
+#include "common/media_io.h"
+#include "common/resource_owners.hpp"
+
+static std::string extract_and_remove_sd_cpp_extra_args(std::string& text) {
+    std::regex re("<sd_cpp_extra_args>(.*?)</sd_cpp_extra_args>");
+    std::smatch match;
+
+    std::string extracted;
+    if (std::regex_search(text, match, re)) {
+        extracted = match[1].str();
+        text      = std::regex_replace(text, re, "");
+    }
+    return extracted;
+}
+
+static bool build_openai_generation_request(const httplib::Request& req,
+                                            ServerRuntime& runtime,
+                                            ImgGenJobRequest& request,
+                                            std::string& error_message) {
+    if (req.body.empty()) {
+        error_message = "empty body";
+        return false;
+    }
+
+    json j                    = json::parse(req.body);
+    std::string prompt        = j.value("prompt", "");
+    int n                     = std::max(1, j.value("n", 1));
+    std::string size          = j.value("size", "");
+    std::string output_format = j.value("output_format", "png");
+    int output_compression    = j.value("output_compression", 100);
+    int width                 = runtime.default_gen_params->width > 0 ? runtime.default_gen_params->width : 512;
+    int height                = runtime.default_gen_params->width > 0 ? runtime.default_gen_params->height : 512;
+    if (!size.empty()) {
+        auto pos = size.find('x');
+        if (pos != std::string::npos) {
+            try {
+                width  = std::stoi(size.substr(0, pos));
+                height = std::stoi(size.substr(pos + 1));
+            } catch (...) {
+            }
+        }
+    }
+
+    if (prompt.empty()) {
+        error_message = "prompt required";
+        return false;
+    }
+
+    request.gen_params = *runtime.default_gen_params;
+    if (!assign_output_options(request, output_format, output_compression, true, error_message)) {
+        return false;
+    }
+
+    request.gen_params.prompt      = prompt;
+    request.gen_params.width       = width;
+    request.gen_params.height      = height;
+    request.gen_params.batch_count = n;
+
+    std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
+    if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
+        error_message = "invalid sd_cpp_extra_args";
+        return false;
+    }
+
+    // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+        error_message = "invalid params";
+        return false;
+    }
+    return true;
+}
+
+static bool build_openai_edit_request(const httplib::Request& req,
+                                      ServerRuntime& runtime,
+                                      ImgGenJobRequest& request,
+                                      std::string& error_message) {
+    if (!req.is_multipart_form_data()) {
+        error_message = "Content-Type must be multipart/form-data";
+        return false;
+    }
+
+    std::string prompt = req.form.get_field("prompt");
+    if (prompt.empty()) {
+        error_message = "prompt required";
+        return false;
+    }
+
+    size_t image_count    = req.form.get_file_count("image[]");
+    bool has_legacy_image = req.form.has_file("image");
+    if (image_count == 0 && !has_legacy_image) {
+        error_message = "at least one image[] required";
+        return false;
+    }
+
+    std::vector<std::vector<uint8_t>> images_bytes;
+    for (size_t i = 0; i < image_count; ++i) {
+        auto file = req.form.get_file("image[]", i);
+        images_bytes.emplace_back(file.content.begin(), file.content.end());
+    }
+    if (image_count == 0 && has_legacy_image) {
+        auto file = req.form.get_file("image");
+        images_bytes.emplace_back(file.content.begin(), file.content.end());
+    }
+
+    std::vector<uint8_t> mask_bytes;
+    if (req.form.has_file("mask")) {
+        auto file = req.form.get_file("mask");
+        mask_bytes.assign(file.content.begin(), file.content.end());
+    }
+
+    int n = 1;
+    if (req.form.has_field("n")) {
+        try {
+            n = std::stoi(req.form.get_field("n"));
+        } catch (...) {
+        }
+    }
+
+    std::string size = req.form.get_field("size");
+    int width        = -1;
+    int height       = -1;
+    if (!size.empty()) {
+        auto pos = size.find('x');
+        if (pos != std::string::npos) {
+            try {
+                width  = std::stoi(size.substr(0, pos));
+                height = std::stoi(size.substr(pos + 1));
+            } catch (...) {
+            }
+        }
+    }
+
+    std::string output_format = req.form.has_field("output_format")
+                                    ? req.form.get_field("output_format")
+                                    : "png";
+
+    int output_compression = 100;
+    try {
+        output_compression = std::stoi(req.form.get_field("output_compression"));
+    } catch (...) {
+    }
+
+    request.gen_params = *runtime.default_gen_params;
+    if (!assign_output_options(request, output_format, output_compression, false, error_message)) {
+        return false;
+    }
+
+    request.gen_params.prompt      = prompt;
+    request.gen_params.width       = width;
+    request.gen_params.height      = height;
+    request.gen_params.batch_count = n;
+
+    for (auto& bytes : images_bytes) {
+        int img_w           = 0;
+        int img_h           = 0;
+        uint8_t* raw_pixels = load_image_from_memory(
+            reinterpret_cast<const char*>(bytes.data()),
+            static_cast<int>(bytes.size()),
+            img_w, img_h,
+            width, height, 3);
+        if (raw_pixels == nullptr) {
+            continue;
+        }
+
+        SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, 3, raw_pixels});
+        request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height);
+        request.gen_params.ref_images.push_back(std::move(image_owner));
+    }
+
+    if (!request.gen_params.ref_images.empty()) {
+        request.gen_params.init_image = request.gen_params.ref_images.front();
+    }
+
+    if (!mask_bytes.empty()) {
+        int expected_width  = 0;
+        int expected_height = 0;
+        if (request.gen_params.width_and_height_are_set()) {
+            expected_width  = request.gen_params.width;
+            expected_height = request.gen_params.height;
+        }
+        int mask_w = 0;
+        int mask_h = 0;
+
+        uint8_t* mask_raw = load_image_from_memory(
+            reinterpret_cast<const char*>(mask_bytes.data()),
+            static_cast<int>(mask_bytes.size()),
+            mask_w, mask_h,
+            expected_width, expected_height, 1);
+        request.gen_params.mask_image.reset({(uint32_t)mask_w, (uint32_t)mask_h, 1, mask_raw});
+        const sd_image_t& mask_image = request.gen_params.mask_image.get();
+        request.gen_params.set_width_and_height_if_unset(mask_image.width, mask_image.height);
+    } else {
+        request.gen_params.mask_image.reset({
+            (uint32_t)request.gen_params.get_resolved_width(),
+            (uint32_t)request.gen_params.get_resolved_height(),
+            1,
+            nullptr,
+        });
+    }
+
+    std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
+    if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
+        error_message = "invalid sd_cpp_extra_args";
+        return false;
+    }
+
+    // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+        error_message = "invalid params";
+        return false;
+    }
+
+    return true;
+}
+
+static bool execute_sync_img_gen_request(ServerRuntime& runtime,
+                                         ImgGenJobRequest& request,
+                                         SDImageVec& results,
+                                         std::string& error_message) {
+    sd_img_gen_params_t img_gen_params = request.to_sd_img_gen_params_t();
+    int num_results                    = 0;
+
+    {
+        std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
+        sd_image_t* raw_results = generate_image(runtime.sd_ctx, &img_gen_params);
+        num_results             = request.gen_params.batch_count;
+        results.adopt(raw_results, num_results);
+    }
+
+    if (results.empty()) {
+        error_message = "generate_image returned no results";
+        return false;
+    }
+    return true;
+}
+
+void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
+    ServerRuntime* runtime = &rt;
+
+    svr.Get("/v1/models", [runtime](const httplib::Request&, httplib::Response& res) {
+        json r;
+        r["data"] = json::array();
+        r["data"].push_back({{"id", "sd-cpp-local"}, {"object", "model"}, {"owned_by", "local"}});
+        res.set_content(r.dump(), "application/json");
+    });
+
+    svr.Post("/v1/images/generations", [runtime](const httplib::Request& req, httplib::Response& res) {
+        try {
+            ImgGenJobRequest request;
+            std::string error_message;
+            if (!build_openai_generation_request(req, *runtime, request, error_message)) {
+                res.status = 400;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            LOG_DEBUG("%s\n", request.gen_params.to_string().c_str());
+
+            SDImageVec results;
+            if (!execute_sync_img_gen_request(*runtime, request, results, error_message)) {
+                res.status = 500;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            json out;
+            out["created"]       = static_cast<long long>(std::time(nullptr));
+            out["data"]          = json::array();
+            out["output_format"] = request.output_format;
+
+            for (int i = 0; i < request.gen_params.batch_count; ++i) {
+                if (results[i].data == nullptr) {
+                    continue;
+                }
+                std::string params = request.gen_params.embed_image_metadata
+                                         ? get_image_params(*runtime->ctx_params,
+                                                            request.gen_params,
+                                                            request.gen_params.seed + i)
+                                         : "";
+                auto image_bytes   = encode_image_to_vector(request.output_format == "jpeg"
+                                                                ? EncodedImageFormat::JPEG
+                                                            : request.output_format == "webp"
+                                                                ? EncodedImageFormat::WEBP
+                                                                : EncodedImageFormat::PNG,
+                                                          results[i].data,
+                                                          results[i].width,
+                                                          results[i].height,
+                                                          results[i].channel,
+                                                          params,
+                                                          request.output_compression);
+                if (image_bytes.empty()) {
+                    LOG_ERROR("write image to mem failed");
+                    continue;
+                }
+
+                json item;
+                item["b64_json"] = base64_encode(image_bytes);
+                out["data"].push_back(item);
+            }
+
+            res.set_content(out.dump(), "application/json");
+            res.status = 200;
+
+        } catch (const std::exception& e) {
+            res.status = 500;
+            json err;
+            err["error"]   = "server_error";
+            err["message"] = e.what();
+            res.set_content(err.dump(), "application/json");
+        }
+    });
+
+    svr.Post("/v1/images/edits", [runtime](const httplib::Request& req, httplib::Response& res) {
+        try {
+            ImgGenJobRequest request;
+            std::string error_message;
+            if (!build_openai_edit_request(req, *runtime, request, error_message)) {
+                res.status = 400;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            LOG_DEBUG("%s\n", request.gen_params.to_string().c_str());
+
+            SDImageVec results;
+            if (!execute_sync_img_gen_request(*runtime, request, results, error_message)) {
+                res.status = 500;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            json out;
+            out["created"]       = static_cast<long long>(std::time(nullptr));
+            out["data"]          = json::array();
+            out["output_format"] = request.output_format;
+
+            for (int i = 0; i < request.gen_params.batch_count; ++i) {
+                if (results[i].data == nullptr) {
+                    continue;
+                }
+                std::string params = request.gen_params.embed_image_metadata
+                                         ? get_image_params(*runtime->ctx_params,
+                                                            request.gen_params,
+                                                            request.gen_params.seed + i)
+                                         : "";
+                auto image_bytes   = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG,
+                                                          results[i].data,
+                                                          results[i].width,
+                                                          results[i].height,
+                                                          results[i].channel,
+                                                          params,
+                                                          request.output_compression);
+                json item;
+                item["b64_json"] = base64_encode(image_bytes);
+                out["data"].push_back(item);
+            }
+
+            res.set_content(out.dump(), "application/json");
+            res.status = 200;
+
+        } catch (const std::exception& e) {
+            res.status = 500;
+            json err;
+            err["error"]   = "server_error";
+            err["message"] = e.what();
+            res.set_content(err.dump(), "application/json");
+        }
+    });
+}

examples/server/routes_sdapi.cpp

@@ -0,0 +1,399 @@
+#include "routes.h"
+
+#include <algorithm>
+#include <cstring>
+#include <regex>
+#include <string_view>
+#include <unordered_map>
+
+#include "common/common.h"
+#include "common/media_io.h"
+#include "common/resource_owners.hpp"
+
+namespace fs = std::filesystem;
+
+static std::string extract_and_remove_sd_cpp_extra_args(std::string& text) {
+    std::regex re("<sd_cpp_extra_args>(.*?)</sd_cpp_extra_args>");
+    std::smatch match;
+
+    std::string extracted;
+    if (std::regex_search(text, match, re)) {
+        extracted = match[1].str();
+        text      = std::regex_replace(text, re, "");
+    }
+    return extracted;
+}
+
+static fs::path resolve_display_model_path(const ServerRuntime& runtime) {
+    const auto& ctx = *runtime.ctx_params;
+    if (!ctx.model_path.empty()) {
+        return fs::path(ctx.model_path);
+    }
+    if (!ctx.diffusion_model_path.empty()) {
+        return fs::path(ctx.diffusion_model_path);
+    }
+    return {};
+}
+
+static enum sample_method_t get_sdapi_sample_method(std::string name) {
+    enum sample_method_t result = str_to_sample_method(name.c_str());
+    if (result != SAMPLE_METHOD_COUNT) {
+        return result;
+    }
+
+    std::transform(name.begin(), name.end(), name.begin(),
+                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    static const std::unordered_map<std::string_view, sample_method_t> hardcoded{
+        {"euler a", EULER_A_SAMPLE_METHOD},
+        {"k_euler_a", EULER_A_SAMPLE_METHOD},
+        {"euler", EULER_SAMPLE_METHOD},
+        {"k_euler", EULER_SAMPLE_METHOD},
+        {"heun", HEUN_SAMPLE_METHOD},
+        {"k_heun", HEUN_SAMPLE_METHOD},
+        {"dpm2", DPM2_SAMPLE_METHOD},
+        {"k_dpm_2", DPM2_SAMPLE_METHOD},
+        {"lcm", LCM_SAMPLE_METHOD},
+        {"ddim", DDIM_TRAILING_SAMPLE_METHOD},
+        {"dpm++ 2m", DPMPP2M_SAMPLE_METHOD},
+        {"k_dpmpp_2m", DPMPP2M_SAMPLE_METHOD},
+        {"res multistep", RES_MULTISTEP_SAMPLE_METHOD},
+        {"k_res_multistep", RES_MULTISTEP_SAMPLE_METHOD},
+        {"res 2s", RES_2S_SAMPLE_METHOD},
+        {"k_res_2s", RES_2S_SAMPLE_METHOD},
+    };
+    auto it = hardcoded.find(name);
+    return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT;
+}
+
+static void assign_solid_mask(SDImageOwner& mask_owner, int width, int height) {
+    const size_t pixel_count = static_cast<size_t>(width) * static_cast<size_t>(height);
+    uint8_t* raw_mask        = static_cast<uint8_t*>(malloc(pixel_count));
+    if (raw_mask == nullptr) {
+        mask_owner.reset({0, 0, 1, nullptr});
+        return;
+    }
+    std::memset(raw_mask, 255, pixel_count);
+    mask_owner.reset({(uint32_t)width, (uint32_t)height, 1, raw_mask});
+}
+
+static bool build_sdapi_img_gen_request(const json& j,
+                                        ServerRuntime& runtime,
+                                        bool img2img,
+                                        ImgGenJobRequest& request,
+                                        std::string& error_message) {
+    std::string prompt          = j.value("prompt", "");
+    std::string negative_prompt = j.value("negative_prompt", "");
+    int width                   = j.value("width", 512);
+    int height                  = j.value("height", 512);
+    int steps                   = j.value("steps", runtime.default_gen_params->sample_params.sample_steps);
+    float cfg_scale             = j.value("cfg_scale", runtime.default_gen_params->sample_params.guidance.txt_cfg);
+    int64_t seed                = j.value("seed", -1);
+    int batch_size              = j.value("batch_size", 1);
+    int clip_skip               = j.value("clip_skip", -1);
+    std::string sampler_name    = j.value("sampler_name", "");
+    std::string scheduler_name  = j.value("scheduler", "");
+
+    if (width <= 0 || height <= 0) {
+        error_message = "width and height must be positive";
+        return false;
+    }
+
+    if (prompt.empty()) {
+        error_message = "prompt required";
+        return false;
+    }
+
+    request.gen_params = *runtime.default_gen_params;
+
+    request.gen_params.prompt                         = prompt;
+    request.gen_params.negative_prompt                = negative_prompt;
+    request.gen_params.seed                           = seed;
+    request.gen_params.sample_params.sample_steps     = steps;
+    request.gen_params.batch_count                    = batch_size;
+    request.gen_params.sample_params.guidance.txt_cfg = cfg_scale;
+    request.gen_params.width                          = j.value("width", -1);
+    request.gen_params.height                         = j.value("height", -1);
+
+    std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
+    if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
+        error_message = "invalid sd_cpp_extra_args";
+        return false;
+    }
+
+    if (clip_skip > 0) {
+        request.gen_params.clip_skip = clip_skip;
+    }
+
+    enum sample_method_t sample_method = get_sdapi_sample_method(sampler_name);
+    if (sample_method != SAMPLE_METHOD_COUNT) {
+        request.gen_params.sample_params.sample_method = sample_method;
+    }
+
+    enum scheduler_t scheduler = str_to_scheduler(scheduler_name.c_str());
+    if (scheduler != SCHEDULER_COUNT) {
+        request.gen_params.sample_params.scheduler = scheduler;
+    }
+
+    if (j.contains("lora") && j["lora"].is_array()) {
+        request.gen_params.lora_map.clear();
+        request.gen_params.high_noise_lora_map.clear();
+
+        for (const auto& item : j["lora"]) {
+            if (!item.is_object()) {
+                continue;
+            }
+
+            std::string path   = item.value("path", "");
+            float multiplier   = item.value("multiplier", 1.0f);
+            bool is_high_noise = item.value("is_high_noise", false);
+
+            if (path.empty()) {
+                error_message = "lora.path required";
+                return false;
+            }
+
+            std::string fullpath = get_lora_full_path(runtime, path);
+            if (fullpath.empty()) {
+                error_message = "invalid lora path: " + path;
+                return false;
+            }
+
+            if (is_high_noise) {
+                request.gen_params.high_noise_lora_map[fullpath] += multiplier;
+            } else {
+                request.gen_params.lora_map[fullpath] += multiplier;
+            }
+        }
+    }
+
+    if (img2img) {
+        const int expected_width  = request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0;
+        const int expected_height = request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0;
+
+        if (j.contains("init_images") && j["init_images"].is_array() && !j["init_images"].empty()) {
+            if (decode_base64_image(j["init_images"][0].get<std::string>(),
+                                    3,
+                                    expected_width,
+                                    expected_height,
+                                    request.gen_params.init_image)) {
+                const sd_image_t& image = request.gen_params.init_image.get();
+                request.gen_params.set_width_and_height_if_unset(image.width, image.height);
+            }
+        }
+
+        if (j.contains("mask") && j["mask"].is_string()) {
+            if (decode_base64_image(j["mask"].get<std::string>(),
+                                    1,
+                                    expected_width,
+                                    expected_height,
+                                    request.gen_params.mask_image)) {
+                const sd_image_t& image = request.gen_params.mask_image.get();
+                request.gen_params.set_width_and_height_if_unset(image.width, image.height);
+            }
+            sd_image_t& mask_image      = request.gen_params.mask_image.get();
+            bool inpainting_mask_invert = j.value("inpainting_mask_invert", 0) != 0;
+            if (inpainting_mask_invert && mask_image.data != nullptr) {
+                for (uint32_t i = 0; i < mask_image.width * mask_image.height; ++i) {
+                    mask_image.data[i] = 255 - mask_image.data[i];
+                }
+            }
+        } else {
+            const int resolved_width  = request.gen_params.get_resolved_width();
+            const int resolved_height = request.gen_params.get_resolved_height();
+            assign_solid_mask(request.gen_params.mask_image, resolved_width, resolved_height);
+        }
+
+        float denoising_strength = j.value("denoising_strength", -1.f);
+        if (denoising_strength >= 0.f) {
+            request.gen_params.strength = std::min(denoising_strength, 1.0f);
+        }
+    }
+
+    if (j.contains("extra_images") && j["extra_images"].is_array()) {
+        for (const auto& extra_image : j["extra_images"]) {
+            if (!extra_image.is_string()) {
+                continue;
+            }
+            SDImageOwner image_owner;
+            if (decode_base64_image(extra_image.get<std::string>(),
+                                    3,
+                                    request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0,
+                                    request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0,
+                                    image_owner)) {
+                const sd_image_t& image = image_owner.get();
+                request.gen_params.set_width_and_height_if_unset(image.width, image.height);
+                request.gen_params.ref_images.push_back(std::move(image_owner));
+            }
+        }
+    }
+
+    // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+        error_message = "invalid params";
+        return false;
+    }
+
+    return true;
+}
+
+void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
+    ServerRuntime* runtime = &rt;
+
+    auto sdapi_any2img = [runtime](const httplib::Request& req, httplib::Response& res, bool img2img) {
+        try {
+            if (req.body.empty()) {
+                res.status = 400;
+                res.set_content(R"({"error":"empty body"})", "application/json");
+                return;
+            }
+
+            json j = json::parse(req.body);
+            ImgGenJobRequest request;
+            std::string error_message;
+            if (!build_sdapi_img_gen_request(j, *runtime, img2img, request, error_message)) {
+                res.status = 400;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            LOG_DEBUG("%s\n", request.gen_params.to_string().c_str());
+
+            sd_img_gen_params_t img_gen_params = request.to_sd_img_gen_params_t();
+            SDImageVec results;
+            int num_results = 0;
+
+            {
+                std::lock_guard<std::mutex> lock(*runtime->sd_ctx_mutex);
+                sd_image_t* raw_results = generate_image(runtime->sd_ctx, &img_gen_params);
+                num_results             = request.gen_params.batch_count;
+                results.adopt(raw_results, num_results);
+            }
+
+            if (results.empty()) {
+                res.status = 500;
+                res.set_content(R"({"error":"generate_image returned no results"})", "application/json");
+                return;
+            }
+
+            json out;
+            out["images"]     = json::array();
+            out["parameters"] = j;
+            out["info"]       = "";
+
+            for (int i = 0; i < num_results; ++i) {
+                if (results[i].data == nullptr) {
+                    continue;
+                }
+
+                std::string params = request.gen_params.embed_image_metadata
+                                         ? get_image_params(*runtime->ctx_params,
+                                                            request.gen_params,
+                                                            request.gen_params.seed + i)
+                                         : "";
+                auto image_bytes   = encode_image_to_vector(EncodedImageFormat::PNG,
+                                                            results[i].data,
+                                                            results[i].width,
+                                                            results[i].height,
+                                                            results[i].channel,
+                                                            params);
+
+                if (image_bytes.empty()) {
+                    LOG_ERROR("write image to mem failed");
+                    continue;
+                }
+
+                out["images"].push_back(base64_encode(image_bytes));
+            }
+
+            res.set_content(out.dump(), "application/json");
+            res.status = 200;
+
+        } catch (const std::exception& e) {
+            res.status = 500;
+            json err;
+            err["error"]   = "server_error";
+            err["message"] = e.what();
+            res.set_content(err.dump(), "application/json");
+        }
+    };
+
+    svr.Post("/sdapi/v1/txt2img", [sdapi_any2img](const httplib::Request& req, httplib::Response& res) {
+        sdapi_any2img(req, res, false);
+    });
+
+    svr.Post("/sdapi/v1/img2img", [sdapi_any2img](const httplib::Request& req, httplib::Response& res) {
+        sdapi_any2img(req, res, true);
+    });
+
+    svr.Get("/sdapi/v1/loras", [runtime](const httplib::Request&, httplib::Response& res) {
+        refresh_lora_cache(*runtime);
+
+        json result = json::array();
+        {
+            std::lock_guard<std::mutex> lock(*runtime->lora_mutex);
+            for (const auto& e : *runtime->lora_cache) {
+                json item;
+                item["name"] = e.name;
+                item["path"] = e.path;
+                result.push_back(item);
+            }
+        }
+
+        res.set_content(result.dump(), "application/json");
+    });
+
+    svr.Get("/sdapi/v1/samplers", [runtime](const httplib::Request&, httplib::Response& res) {
+        std::vector<std::string> sampler_names;
+        sampler_names.push_back("default");
+        for (int i = 0; i < SAMPLE_METHOD_COUNT; i++) {
+            sampler_names.push_back(sd_sample_method_name((sample_method_t)i));
+        }
+        json r = json::array();
+        for (auto name : sampler_names) {
+            json entry;
+            entry["name"]    = name;
+            entry["aliases"] = json::array({name});
+            entry["options"] = json::object();
+            r.push_back(entry);
+        }
+        res.set_content(r.dump(), "application/json");
+    });
+
+    svr.Get("/sdapi/v1/schedulers", [runtime](const httplib::Request&, httplib::Response& res) {
+        std::vector<std::string> scheduler_names;
+        scheduler_names.push_back("default");
+        for (int i = 0; i < SCHEDULER_COUNT; i++) {
+            scheduler_names.push_back(sd_scheduler_name((scheduler_t)i));
+        }
+        json r = json::array();
+        for (auto name : scheduler_names) {
+            json entry;
+            entry["name"]  = name;
+            entry["label"] = name;
+            r.push_back(entry);
+        }
+        res.set_content(r.dump(), "application/json");
+    });
+
+    svr.Get("/sdapi/v1/sd-models", [runtime](const httplib::Request&, httplib::Response& res) {
+        fs::path model_path = resolve_display_model_path(*runtime);
+        json entry;
+        entry["title"]      = model_path.stem();
+        entry["model_name"] = model_path.stem();
+        entry["filename"]   = model_path.filename();
+        entry["hash"]       = "8888888888";
+        entry["sha256"]     = "8888888888888888888888888888888888888888888888888888888888888888";
+        entry["config"]     = nullptr;
+        json r              = json::array();
+        r.push_back(entry);
+        res.set_content(r.dump(), "application/json");
+    });
+
+    svr.Get("/sdapi/v1/options", [runtime](const httplib::Request&, httplib::Response& res) {
+        fs::path model_path = resolve_display_model_path(*runtime);
+        json r;
+        r["samples_format"]      = "png";
+        r["sd_model_checkpoint"] = model_path.stem();
+        res.set_content(r.dump(), "application/json");
+    });
+}

examples/server/routes_sdcpp.cpp

@@ -0,0 +1,345 @@
+#include "routes.h"
+
+#include <algorithm>
+#include <cmath>
+#include <filesystem>
+
+#include "async_jobs.h"
+#include "common/common.h"
+
+namespace fs = std::filesystem;
+
+static bool parse_cache_mode(const std::string& mode_str, sd_cache_mode_t& mode_out) {
+    if (mode_str == "disabled") {
+        mode_out = SD_CACHE_DISABLED;
+        return true;
+    }
+    if (mode_str == "easycache") {
+        mode_out = SD_CACHE_EASYCACHE;
+        return true;
+    }
+    if (mode_str == "ucache") {
+        mode_out = SD_CACHE_UCACHE;
+        return true;
+    }
+    if (mode_str == "dbcache") {
+        mode_out = SD_CACHE_DBCACHE;
+        return true;
+    }
+    if (mode_str == "taylorseer") {
+        mode_out = SD_CACHE_TAYLORSEER;
+        return true;
+    }
+    if (mode_str == "cache-dit") {
+        mode_out = SD_CACHE_CACHE_DIT;
+        return true;
+    }
+    if (mode_str == "spectrum") {
+        mode_out = SD_CACHE_SPECTRUM;
+        return true;
+    }
+    return false;
+}
+
+static json finite_number_or_null(float value) {
+    return std::isfinite(value) ? json(value) : json(nullptr);
+}
+
+static const char* capability_scheduler_name(enum scheduler_t scheduler) {
+    return scheduler < SCHEDULER_COUNT ? sd_scheduler_name(scheduler) : "default";
+}
+
+static const char* capability_sample_method_name(enum sample_method_t sample_method) {
+    return sample_method < SAMPLE_METHOD_COUNT ? sd_sample_method_name(sample_method) : "default";
+}
+
+static json make_vae_tiling_json(const sd_tiling_params_t& params) {
+    return {
+        {"enabled", params.enabled},
+        {"tile_size_x", params.tile_size_x},
+        {"tile_size_y", params.tile_size_y},
+        {"target_overlap", params.target_overlap},
+        {"rel_size_x", params.rel_size_x},
+        {"rel_size_y", params.rel_size_y},
+    };
+}
+
+static fs::path resolve_display_model_path(const ServerRuntime& runtime) {
+    const auto& ctx = *runtime.ctx_params;
+    if (!ctx.model_path.empty()) {
+        return fs::path(ctx.model_path);
+    }
+    if (!ctx.diffusion_model_path.empty()) {
+        return fs::path(ctx.diffusion_model_path);
+    }
+    return {};
+}
+
+static json make_capabilities_json(ServerRuntime& runtime) {
+    refresh_lora_cache(runtime);
+
+    AsyncJobManager& manager  = *runtime.async_job_manager;
+    const auto& defaults      = *runtime.default_gen_params;
+    const auto& sample_params = defaults.sample_params;
+    const auto& guidance      = sample_params.guidance;
+    const fs::path model_path = resolve_display_model_path(runtime);
+    json samplers             = json::array();
+    json schedulers           = json::array();
+    json output_formats       = json::array({"png", "jpeg"});
+    json available_loras      = json::array();
+
+    for (int i = 0; i < SAMPLE_METHOD_COUNT; ++i) {
+        samplers.push_back(sd_sample_method_name((sample_method_t)i));
+    }
+
+    for (int i = 0; i < SCHEDULER_COUNT; ++i) {
+        schedulers.push_back(sd_scheduler_name((scheduler_t)i));
+    }
+
+#ifdef SD_USE_WEBP
+    output_formats.push_back("webp");
+#endif
+
+    {
+        std::lock_guard<std::mutex> lock(*runtime.lora_mutex);
+        for (const auto& entry : *runtime.lora_cache) {
+            available_loras.push_back({
+                {"name", entry.name},
+                {"path", entry.path},
+            });
+        }
+    }
+
+    json result;
+    result["model"] = {
+        {"name", model_path.filename().u8string()},
+        {"stem", model_path.stem().u8string()},
+        {"path", model_path.u8string()},
+    };
+    result["defaults"] = {
+        {"prompt", defaults.prompt},
+        {"negative_prompt", defaults.negative_prompt},
+        {"clip_skip", defaults.clip_skip},
+        {"width", defaults.width > 0 ? defaults.width : 512},
+        {"height", defaults.height > 0 ? defaults.height : 512},
+        {"strength", defaults.strength},
+        {"seed", defaults.seed},
+        {"batch_count", defaults.batch_count},
+        {"auto_resize_ref_image", defaults.auto_resize_ref_image},
+        {"increase_ref_index", defaults.increase_ref_index},
+        {"control_strength", defaults.control_strength},
+        {"sample_params",
+         {
+             {"scheduler", capability_scheduler_name(sample_params.scheduler)},
+             {"sample_method", capability_sample_method_name(sample_params.sample_method)},
+             {"sample_steps", sample_params.sample_steps},
+             {"eta", finite_number_or_null(sample_params.eta)},
+             {"shifted_timestep", sample_params.shifted_timestep},
+             {"flow_shift", finite_number_or_null(sample_params.flow_shift)},
+             {"guidance",
+              {
+                  {"txt_cfg", guidance.txt_cfg},
+                  {"img_cfg", finite_number_or_null(guidance.img_cfg)},
+                  {"distilled_guidance", guidance.distilled_guidance},
+                  {"slg",
+                   {
+                       {"layers", defaults.skip_layers},
+                       {"layer_start", guidance.slg.layer_start},
+                       {"layer_end", guidance.slg.layer_end},
+                       {"scale", guidance.slg.scale},
+                   }},
+              }},
+         }},
+        {"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
+        {"cache_mode", defaults.cache_mode},
+        {"cache_option", defaults.cache_option},
+        {"scm_mask", defaults.scm_mask},
+        {"scm_policy_dynamic", defaults.scm_policy_dynamic},
+        {"output_format", "png"},
+        {"output_compression", 100},
+    };
+    result["limits"] = {
+        {"min_width", 64},
+        {"max_width", 4096},
+        {"min_height", 64},
+        {"max_height", 4096},
+        {"max_batch_count", 8},
+        {"max_queue_size", manager.max_pending_jobs},
+    };
+    result["samplers"]       = samplers;
+    result["schedulers"]     = schedulers;
+    result["output_formats"] = output_formats;
+    result["features"]       = {
+              {"init_image", true},
+              {"mask_image", true},
+              {"control_image", true},
+              {"ref_images", true},
+              {"lora", true},
+              {"vae_tiling", true},
+              {"cache", true},
+              {"cancel_queued", true},
+              {"cancel_generating", false},
+    };
+    result["loras"] = available_loras;
+    return result;
+}
+
+static bool parse_img_gen_request(const json& body,
+                                  ServerRuntime& runtime,
+                                  ImgGenJobRequest& request,
+                                  std::string& error_message) {
+    request.gen_params = *runtime.default_gen_params;
+
+    refresh_lora_cache(runtime);
+    if (!request.gen_params.from_json_str(body.dump(), [&](const std::string& path) {
+            return get_lora_full_path(runtime, path);
+        })) {
+        error_message = "invalid generation parameters";
+        return false;
+    }
+
+    std::string output_format = body.value("output_format", "png");
+    int output_compression    = body.value("output_compression", 100);
+    if (!assign_output_options(request, output_format, output_compression, true, error_message)) {
+        return false;
+    }
+    // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
+    if (!request.gen_params.resolve_and_validate(IMG_GEN, "", true)) {
+        error_message = "invalid generation parameters";
+        return false;
+    }
+    return true;
+}
+
+void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
+    ServerRuntime* runtime = &rt;
+
+    svr.Get("/sdcpp/v1/capabilities", [runtime](const httplib::Request&, httplib::Response& res) {
+        res.status = 200;
+        res.set_content(make_capabilities_json(*runtime).dump(), "application/json");
+    });
+
+    svr.Post("/sdcpp/v1/img_gen", [runtime](const httplib::Request& req, httplib::Response& res) {
+        try {
+            if (req.body.empty()) {
+                res.status = 400;
+                res.set_content(R"({"error":"empty body"})", "application/json");
+                return;
+            }
+
+            json body = json::parse(req.body);
+            ImgGenJobRequest request;
+            std::string error_message;
+            if (!parse_img_gen_request(body, *runtime, request, error_message)) {
+                res.status = 400;
+                res.set_content(json({{"error", error_message}}).dump(), "application/json");
+                return;
+            }
+
+            AsyncJobManager& manager                = *runtime->async_job_manager;
+            std::shared_ptr<AsyncGenerationJob> job = std::make_shared<AsyncGenerationJob>();
+            job->kind                               = AsyncJobKind::ImgGen;
+            job->status                             = AsyncJobStatus::Queued;
+            job->created_at                         = unix_timestamp_now();
+            job->img_gen                            = std::move(request);
+
+            {
+                std::lock_guard<std::mutex> lock(manager.mutex);
+                purge_expired_jobs(manager);
+                if (count_pending_jobs(manager) >= manager.max_pending_jobs) {
+                    res.status = 429;
+                    res.set_content(R"({"error":"job queue is full"})", "application/json");
+                    return;
+                }
+                job->id               = make_async_job_id(manager);
+                manager.jobs[job->id] = job;
+                manager.queue.push_back(job->id);
+            }
+
+            manager.cv.notify_one();
+
+            json out;
+            out["id"]       = job->id;
+            out["kind"]     = async_job_kind_name(job->kind);
+            out["status"]   = async_job_status_name(job->status);
+            out["created"]  = job->created_at;
+            out["poll_url"] = "/sdcpp/v1/jobs/" + job->id;
+
+            res.status = 202;
+            res.set_content(out.dump(), "application/json");
+        } catch (const json::parse_error& e) {
+            res.status = 400;
+            res.set_content(json({{"error", "invalid json"}, {"message", e.what()}}).dump(), "application/json");
+        } catch (const std::exception& e) {
+            res.status = 500;
+            res.set_content(json({{"error", "server_error"}, {"message", e.what()}}).dump(), "application/json");
+        }
+    });
+
+    svr.Post("/sdcpp/v1/vid_gen", [](const httplib::Request&, httplib::Response& res) {
+        res.status = 501;
+        res.set_content(R"({"error":"vid_gen is reserved and not implemented yet"})", "application/json");
+    });
+
+    svr.Get(R"(/sdcpp/v1/jobs/([A-Za-z0-9_\-]+))", [runtime](const httplib::Request& req, httplib::Response& res) {
+        AsyncJobManager& manager = *runtime->async_job_manager;
+        std::lock_guard<std::mutex> lock(manager.mutex);
+        purge_expired_jobs(manager);
+
+        std::string job_id = req.matches[1];
+        auto it            = manager.jobs.find(job_id);
+        if (it == manager.jobs.end()) {
+            if (manager.expired_jobs.find(job_id) != manager.expired_jobs.end()) {
+                res.status = 410;
+                res.set_content(R"({"error":"job expired"})", "application/json");
+            } else {
+                res.status = 404;
+                res.set_content(R"({"error":"job not found"})", "application/json");
+            }
+            return;
+        }
+
+        res.status = 200;
+        res.set_content(make_async_job_json(manager, *it->second).dump(), "application/json");
+    });
+
+    svr.Post(R"(/sdcpp/v1/jobs/([A-Za-z0-9_\-]+)/cancel)", [runtime](const httplib::Request& req, httplib::Response& res) {
+        AsyncJobManager& manager = *runtime->async_job_manager;
+        std::lock_guard<std::mutex> lock(manager.mutex);
+        purge_expired_jobs(manager);
+
+        std::string job_id = req.matches[1];
+        auto it            = manager.jobs.find(job_id);
+        if (it == manager.jobs.end()) {
+            if (manager.expired_jobs.find(job_id) != manager.expired_jobs.end()) {
+                res.status = 410;
+                res.set_content(R"({"error":"job expired"})", "application/json");
+            } else {
+                res.status = 404;
+                res.set_content(R"({"error":"job not found"})", "application/json");
+            }
+            return;
+        }
+
+        auto& job = *it->second;
+        if (job.status == AsyncJobStatus::Queued) {
+            if (!cancel_queued_job(manager, job)) {
+                res.status = 409;
+                res.set_content(R"({"error":"job queue state changed before cancellation"})", "application/json");
+                return;
+            }
+            res.status = 200;
+            res.set_content(make_async_job_json(manager, job).dump(), "application/json");
+            return;
+        }
+
+        if (job.status == AsyncJobStatus::Generating) {
+            res.status = 409;
+            res.set_content(R"({"error":"job is currently generating and cannot be interrupted yet"})", "application/json");
+            return;
+        }
+
+        res.status = 200;
+        res.set_content(make_async_job_json(manager, job).dump(), "application/json");
+    });
+}

examples/server/runtime.cpp

@@ -0,0 +1,186 @@
+#include "runtime.h"
+
+#include <algorithm>
+#include <chrono>
+#include <cstdlib>
+#include <filesystem>
+#include <mutex>
+#include <regex>
+#include <sstream>
+
+#include "common/common.h"
+#include "common/log.h"
+
+namespace fs = std::filesystem;
+
+static const std::string k_base64_chars =
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    "abcdefghijklmnopqrstuvwxyz"
+    "0123456789+/";
+
+std::string base64_encode(const std::vector<uint8_t>& bytes) {
+    std::string ret;
+    int val  = 0;
+    int valb = -6;
+    for (uint8_t c : bytes) {
+        val = (val << 8) + c;
+        valb += 8;
+        while (valb >= 0) {
+            ret.push_back(k_base64_chars[(val >> valb) & 0x3F]);
+            valb -= 6;
+        }
+    }
+    if (valb > -6) {
+        ret.push_back(k_base64_chars[((val << 8) >> (valb + 8)) & 0x3F]);
+    }
+    while (ret.size() % 4) {
+        ret.push_back('=');
+    }
+    return ret;
+}
+
+std::string normalize_output_format(std::string output_format) {
+    std::transform(output_format.begin(), output_format.end(), output_format.begin(),
+                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    return output_format;
+}
+
+bool assign_output_options(ImgGenJobRequest& request,
+                           std::string output_format,
+                           int output_compression,
+                           bool allow_webp,
+                           std::string& error_message) {
+    request.output_format      = normalize_output_format(std::move(output_format));
+    request.output_compression = std::clamp(output_compression, 0, 100);
+
+    const bool valid_format = request.output_format == "png" ||
+                              request.output_format == "jpeg" ||
+                              (allow_webp && request.output_format == "webp");
+    if (!valid_format) {
+        error_message = allow_webp
+                            ? "invalid output_format, must be one of [png, jpeg, webp]"
+                            : "invalid output_format, must be one of [png, jpeg]";
+        return false;
+    }
+
+    return true;
+}
+
+ArgOptions SDSvrParams::get_options() {
+    ArgOptions options;
+
+    options.string_options = {
+        {"-l", "--listen-ip", "server listen ip (default: 127.0.0.1)", &listen_ip},
+        {"", "--serve-html-path", "path to HTML file to serve at root (optional)", &serve_html_path},
+    };
+
+    options.int_options = {
+        {"", "--listen-port", "server listen port (default: 1234)", &listen_port},
+    };
+
+    options.bool_options = {
+        {"-v", "--verbose", "print extra info", true, &verbose},
+        {"", "--color", "colors the logging tags according to level", true, &color},
+    };
+
+    auto on_help_arg = [&](int, const char**, int) {
+        normal_exit = true;
+        return -1;
+    };
+
+    options.manual_options = {
+        {"-h", "--help", "show this help message and exit", on_help_arg},
+    };
+    return options;
+}
+
+bool SDSvrParams::validate() {
+    if (listen_ip.empty()) {
+        LOG_ERROR("error: the following arguments are required: listen_ip");
+        return false;
+    }
+
+    if (listen_port < 0 || listen_port > 65535) {
+        LOG_ERROR("error: listen_port should be in the range [0, 65535]");
+        return false;
+    }
+
+    if (!serve_html_path.empty() && !fs::exists(serve_html_path)) {
+        LOG_ERROR("error: serve_html_path file does not exist: %s", serve_html_path.c_str());
+        return false;
+    }
+    return true;
+}
+
+bool SDSvrParams::resolve_and_validate() {
+    if (!validate()) {
+        return false;
+    }
+    return true;
+}
+
+std::string SDSvrParams::to_string() const {
+    std::ostringstream oss;
+    oss << "SDSvrParams {\n"
+        << "  listen_ip: " << listen_ip << ",\n"
+        << "  listen_port: \"" << listen_port << "\",\n"
+        << "  serve_html_path: \"" << serve_html_path << "\",\n"
+        << "}";
+    return oss.str();
+}
+
+void refresh_lora_cache(ServerRuntime& rt) {
+    std::vector<LoraEntry> new_cache;
+
+    fs::path lora_dir = rt.ctx_params->lora_model_dir;
+    if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) {
+        auto is_lora_ext = [](const fs::path& p) {
+            auto ext = p.extension().string();
+            std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) {
+                return static_cast<char>(std::tolower(c));
+            });
+            return ext == ".gguf" || ext == ".pt" || ext == ".pth" || ext == ".safetensors";
+        };
+
+        for (auto& entry : fs::recursive_directory_iterator(lora_dir)) {
+            if (!entry.is_regular_file()) {
+                continue;
+            }
+            const fs::path& p = entry.path();
+            if (!is_lora_ext(p)) {
+                continue;
+            }
+
+            LoraEntry lora_entry;
+            lora_entry.name     = p.stem().u8string();
+            lora_entry.fullpath = p.u8string();
+            std::string rel     = p.lexically_relative(lora_dir).u8string();
+            std::replace(rel.begin(), rel.end(), '\\', '/');
+            lora_entry.path = rel;
+
+            new_cache.push_back(std::move(lora_entry));
+        }
+    }
+
+    std::sort(new_cache.begin(), new_cache.end(), [](const LoraEntry& a, const LoraEntry& b) {
+        return a.path < b.path;
+    });
+
+    {
+        std::lock_guard<std::mutex> lock(*rt.lora_mutex);
+        *rt.lora_cache = std::move(new_cache);
+    }
+}
+
+std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) {
+    std::lock_guard<std::mutex> lock(*rt.lora_mutex);
+    auto it = std::find_if(rt.lora_cache->begin(), rt.lora_cache->end(),
+                           [&](const LoraEntry& entry) { return entry.path == path; });
+    return it != rt.lora_cache->end() ? it->fullpath : "";
+}
+
+int64_t unix_timestamp_now() {
+    return std::chrono::duration_cast<std::chrono::seconds>(
+               std::chrono::system_clock::now().time_since_epoch())
+        .count();
+}

examples/server/runtime.h

@@ -0,0 +1,70 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include <json.hpp>
+#include "common/common.h"
+#include "common/resource_owners.hpp"
+#include "stable-diffusion.h"
+
+using json = nlohmann::json;
+
+struct ArgOptions;
+struct SDContextParams;
+struct AsyncJobManager;
+
+struct SDSvrParams {
+    std::string listen_ip = "127.0.0.1";
+    int listen_port       = 1234;
+    std::string serve_html_path;
+    bool normal_exit = false;
+    bool verbose     = false;
+    bool color       = false;
+
+    ArgOptions get_options();
+    bool validate();
+    bool resolve_and_validate();
+    std::string to_string() const;
+};
+
+struct LoraEntry {
+    std::string name;
+    std::string path;
+    std::string fullpath;
+};
+
+struct ServerRuntime {
+    sd_ctx_t* sd_ctx;
+    std::mutex* sd_ctx_mutex;
+    const SDSvrParams* svr_params;
+    const SDContextParams* ctx_params;
+    const SDGenerationParams* default_gen_params;
+    std::vector<LoraEntry>* lora_cache;
+    std::mutex* lora_mutex;
+    AsyncJobManager* async_job_manager;
+};
+
+struct ImgGenJobRequest {
+    SDGenerationParams gen_params;
+    std::string output_format = "png";
+    int output_compression    = 100;
+
+    sd_img_gen_params_t to_sd_img_gen_params_t() {
+        return gen_params.to_sd_img_gen_params_t();
+    }
+};
+
+std::string base64_encode(const std::vector<uint8_t>& bytes);
+std::string normalize_output_format(std::string output_format);
+bool assign_output_options(ImgGenJobRequest& request,
+                           std::string output_format,
+                           int output_compression,
+                           bool allow_webp,
+                           std::string& error_message);
+void refresh_lora_cache(ServerRuntime& rt);
+std::string get_lora_full_path(ServerRuntime& rt, const std::string& path);
+int64_t unix_timestamp_now();

format-code.sh

@@ -1,4 +1,6 @@
-for f in src/*.cpp src/*.h src/*.hpp src/vocab/*.h src/vocab/*.cpp examples/cli/*.cpp examples/common/*.hpp examples/cli/*.h examples/server/*.cpp; do
+for f in src/*.cpp src/*.h src/*.hpp src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \
+         examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
+         examples/common/*.hpp examples/common/*.h examples/common/*.cpp; do
   [[ "$f" == vocab* ]] && continue
   echo "formatting '$f'"
   # if [ "$f" != "stable-diffusion.h" ]; then

src/auto_encoder_kl.hpp

@@ -501,14 +501,39 @@ protected:
         bool double_z            = true;
     } dd_config;
 
+    static std::string get_tensor_name(const std::string& prefix, const std::string& name) {
+        return prefix.empty() ? name : prefix + "." + name;
+    }
+
+    void detect_decoder_ch(const String2TensorStorage& tensor_storage_map,
+                           const std::string& prefix,
+                           int& decoder_ch) {
+        auto conv_in_iter = tensor_storage_map.find(get_tensor_name(prefix, "decoder.conv_in.weight"));
+        if (conv_in_iter != tensor_storage_map.end() && conv_in_iter->second.n_dims >= 4 && conv_in_iter->second.ne[3] > 0) {
+            int last_ch_mult             = dd_config.ch_mult.back();
+            int64_t conv_in_out_channels = conv_in_iter->second.ne[3];
+            if (last_ch_mult > 0 && conv_in_out_channels % last_ch_mult == 0) {
+                decoder_ch = static_cast<int>(conv_in_out_channels / last_ch_mult);
+                LOG_INFO("vae decoder: ch = %d", decoder_ch);
+            } else {
+                LOG_WARN("vae decoder: failed to infer ch from %s (%" PRId64 " / %d)",
+                         get_tensor_name(prefix, "decoder.conv_in.weight").c_str(),
+                         conv_in_out_channels,
+                         last_ch_mult);
+            }
+        }
+    }
+
 public:
-    AutoEncoderKLModel(SDVersion version          = VERSION_SD1,
-                       bool decode_only           = true,
-                       bool use_linear_projection = false,
-                       bool use_video_decoder     = false)
+    AutoEncoderKLModel(SDVersion version                              = VERSION_SD1,
+                       bool decode_only                               = true,
+                       bool use_linear_projection                     = false,
+                       bool use_video_decoder                         = false,
+                       const String2TensorStorage& tensor_storage_map = {},
+                       const std::string& prefix                      = "")
         : version(version), decode_only(decode_only), use_video_decoder(use_video_decoder) {
         if (sd_version_is_dit(version)) {
-            if (sd_version_is_flux2(version)) {
+            if (sd_version_uses_flux2_vae(version)) {
                 dd_config.z_channels = 32;
                 embed_dim            = 32;
             } else {
@@ -519,7 +544,9 @@ public:
         if (use_video_decoder) {
             use_quant = false;
         }
-        blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder(dd_config.ch,
+        int decoder_ch = dd_config.ch;
+        detect_decoder_ch(tensor_storage_map, prefix, decoder_ch);
+        blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder(decoder_ch,
                                                                    dd_config.out_ch,
                                                                    dd_config.ch_mult,
                                                                    dd_config.num_res_blocks,
@@ -551,7 +578,7 @@ public:
 
     ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
         // z: [N, z_channels, h, w]
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             // [N, C*p*p, h, w] -> [N, C, h*p, w*p]
             int64_t p = 2;
 
@@ -590,7 +617,7 @@ public:
             auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["quant_conv"]);
             z               = quant_conv->forward(ctx, z);  // [N, 2*embed_dim, h/8, w/8]
         }
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
 
             // [N, C, H, W] -> [N, C*p*p, H/p, W/p]
@@ -613,7 +640,7 @@ public:
 
     int get_encoder_output_channels() {
         int factor = dd_config.double_z ? 2 : 1;
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             return dd_config.z_channels * 4;
         }
         return dd_config.z_channels * factor;
@@ -646,7 +673,7 @@ struct AutoEncoderKL : public VAE {
         } else if (sd_version_is_flux(version) || sd_version_is_z_image(version)) {
             scale_factor = 0.3611f;
             shift_factor = 0.1159f;
-        } else if (sd_version_is_flux2(version)) {
+        } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 1.0f;
             shift_factor = 0.f;
         }
@@ -662,7 +689,7 @@ struct AutoEncoderKL : public VAE {
                 break;
             }
         }
-        ae = AutoEncoderKLModel(version, decode_only, use_linear_projection, use_video_decoder);
+        ae = AutoEncoderKLModel(version, decode_only, use_linear_projection, use_video_decoder, tensor_storage_map, prefix);
         ae.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -720,7 +747,7 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             return vae_output;
         } else if (version == VERSION_SD1_PIX2PIX) {
             return sd::ops::chunk(vae_output, 2, 2)[0];
@@ -731,7 +758,7 @@ struct AutoEncoderKL : public VAE {
 
     std::pair<sd::Tensor<float>, sd::Tensor<float>> get_latents_mean_std(const sd::Tensor<float>& latents, int channel_dim) {
         GGML_ASSERT(channel_dim >= 0 && static_cast<size_t>(channel_dim) < static_cast<size_t>(latents.dim()));
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             GGML_ASSERT(latents.shape()[channel_dim] == 128);
             std::vector<int64_t> stats_shape(static_cast<size_t>(latents.dim()), 1);
             stats_shape[static_cast<size_t>(channel_dim)] = latents.shape()[channel_dim];
@@ -777,7 +804,7 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             int channel_dim                = 2;
             auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
             return (latents * std_tensor) / scale_factor + mean_tensor;
@@ -786,7 +813,7 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
-        if (sd_version_is_flux2(version)) {
+        if (sd_version_uses_flux2_vae(version)) {
             int channel_dim                = 2;
             auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
             return ((latents - mean_tensor) * scale_factor) / std_tensor;

src/clip.hpp

@@ -3,455 +3,7 @@
 
 #include "ggml_extend.hpp"
 #include "model.h"
-#include "tokenize_util.h"
-#include "vocab/vocab.h"
-
-/*================================================== CLIPTokenizer ===================================================*/
-
-__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) {
-        byte_set.insert(b);
-        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
-    }
-    for (int b = 161; b <= 172; ++b) {
-        byte_set.insert(b);
-        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
-    }
-    for (int b = 174; b <= 255; ++b) {
-        byte_set.insert(b);
-        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
-    }
-    int n = 0;
-    for (int b = 0; b < 256; ++b) {
-        if (byte_set.find(b) == byte_set.end()) {
-            byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
-            ++n;
-        }
-    }
-    // LOG_DEBUG("byte_unicode_pairs %d", byte_unicode_pairs.size());
-    return byte_unicode_pairs;
-}
-
-// Ref: https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py
-
-typedef std::function<bool(std::string&, std::vector<int32_t>&)> on_new_token_cb_t;
-
-class CLIPTokenizer {
-private:
-    std::map<int, std::u32string> byte_encoder;
-    std::map<std::u32string, int> byte_decoder;
-    std::map<std::u32string, int> encoder;
-    std::map<int, std::u32string> decoder;
-    std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
-    std::regex pat;
-    int encoder_len;
-    int bpe_len;
-
-    std::vector<std::string> special_tokens;
-
-public:
-    const std::string UNK_TOKEN = "<|endoftext|>";
-    const std::string BOS_TOKEN = "<|startoftext|>";
-    const std::string EOS_TOKEN = "<|endoftext|>";
-    const std::string PAD_TOKEN = "<|endoftext|>";
-
-    const int UNK_TOKEN_ID = 49407;
-    const int BOS_TOKEN_ID = 49406;
-    const int EOS_TOKEN_ID = 49407;
-    const int PAD_TOKEN_ID = 49407;
-
-private:
-    static std::string strip(const std::string& str) {
-        std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
-        std::string::size_type end   = str.find_last_not_of(" \t\n\r\v\f");
-
-        if (start == std::string::npos) {
-            // String contains only whitespace characters
-            return "";
-        }
-
-        return str.substr(start, end - start + 1);
-    }
-
-    static std::string whitespace_clean(std::string text) {
-        text = std::regex_replace(text, std::regex(R"(\s+)"), " ");
-        text = strip(text);
-        return text;
-    }
-
-    static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
-        std::set<std::pair<std::u32string, std::u32string>> pairs;
-        if (subwords.size() == 0) {
-            return pairs;
-        }
-        std::u32string prev_subword = subwords[0];
-        for (int i = 1; i < subwords.size(); i++) {
-            std::u32string subword = subwords[i];
-            std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
-            pairs.insert(pair);
-            prev_subword = subword;
-        }
-        return pairs;
-    }
-
-    bool is_special_token(const std::string& token) {
-        for (auto& special_token : special_tokens) {
-            if (special_token == token) {
-                return true;
-            }
-        }
-        return false;
-    }
-
-public:
-    CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "")
-        : PAD_TOKEN_ID(pad_token_id) {
-        if (merges_utf8_str.size() > 0) {
-            load_from_merges(merges_utf8_str);
-        } else {
-            load_from_merges(load_clip_merges());
-        }
-        add_special_token("<|startoftext|>");
-        add_special_token("<|endoftext|>");
-    }
-
-    void load_from_merges(const std::string& merges_utf8_str) {
-        auto byte_unicode_pairs = bytes_to_unicode();
-        // printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size());
-        byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-        for (auto& pair : byte_unicode_pairs) {
-            byte_decoder[pair.second] = pair.first;
-        }
-        // for (auto & pair: byte_unicode_pairs) {
-        //     std::cout << pair.first << ": " << pair.second << std::endl;
-        // }
-        std::vector<std::u32string> merges;
-        size_t start = 0;
-        size_t pos;
-        std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
-        while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
-            merges.push_back(merges_utf32_str.substr(start, pos - start));
-            start = pos + 1;
-        }
-        // LOG_DEBUG("merges size %llu", merges.size());
-        GGML_ASSERT(merges.size() == 48895);
-        merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
-        std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-        for (const auto& merge : merges) {
-            size_t space_pos = merge.find(' ');
-            merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-            // LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
-            // printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
-            //                     utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
-        }
-        std::vector<std::u32string> vocab;
-        for (const auto& pair : byte_unicode_pairs) {
-            vocab.push_back(pair.second);
-        }
-        for (const auto& pair : byte_unicode_pairs) {
-            vocab.push_back(pair.second + utf8_to_utf32("</w>"));
-        }
-        for (const auto& merge : merge_pairs) {
-            vocab.push_back(merge.first + merge.second);
-        }
-        vocab.push_back(utf8_to_utf32("<|startoftext|>"));
-        vocab.push_back(utf8_to_utf32("<|endoftext|>"));
-        LOG_DEBUG("vocab size: %llu", vocab.size());
-        int i = 0;
-        for (const auto& token : vocab) {
-            encoder[token] = i;
-            decoder[i]     = token;
-            i++;
-        }
-        encoder_len = i;
-
-        auto it = encoder.find(utf8_to_utf32("img</w>"));
-        if (it != encoder.end()) {
-            LOG_DEBUG("trigger word img already in vocab");
-        } else {
-            LOG_DEBUG("trigger word img not in vocab yet");
-        }
-
-        int rank = 0;
-        for (const auto& merge : merge_pairs) {
-            bpe_ranks[merge] = rank++;
-        }
-        bpe_len = rank;
-    };
-
-    void add_token(const std::string& text) {
-        std::u32string token = utf8_to_utf32(text);
-        auto it              = encoder.find(token);
-        if (it != encoder.end()) {
-            encoder[token]       = encoder_len;
-            decoder[encoder_len] = token;
-            encoder_len++;
-        }
-    }
-
-    void add_special_token(const std::string& token) {
-        special_tokens.push_back(token);
-    }
-
-    std::u32string bpe(const std::u32string& token) {
-        std::vector<std::u32string> word;
-
-        for (int i = 0; i < token.size() - 1; i++) {
-            word.emplace_back(1, token[i]);
-        }
-        word.push_back(token.substr(token.size() - 1) + utf8_to_utf32("</w>"));
-
-        std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
-
-        if (pairs.empty()) {
-            return token + utf8_to_utf32("</w>");
-        }
-
-        while (true) {
-            auto min_pair_iter = std::min_element(pairs.begin(),
-                                                  pairs.end(),
-                                                  [&](const std::pair<std::u32string, std::u32string>& a,
-                                                      const std::pair<std::u32string, std::u32string>& b) {
-                                                      if (bpe_ranks.find(a) == bpe_ranks.end()) {
-                                                          return false;
-                                                      } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
-                                                          return true;
-                                                      }
-                                                      return bpe_ranks.at(a) < bpe_ranks.at(b);
-                                                  });
-
-            const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
-
-            if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
-                break;
-            }
-
-            std::u32string first  = bigram.first;
-            std::u32string second = bigram.second;
-            std::vector<std::u32string> new_word;
-            int32_t i = 0;
-
-            while (i < word.size()) {
-                auto it = std::find(word.begin() + i, word.end(), first);
-                if (it == word.end()) {
-                    new_word.insert(new_word.end(), word.begin() + i, word.end());
-                    break;
-                }
-                new_word.insert(new_word.end(), word.begin() + i, it);
-                i = static_cast<int32_t>(std::distance(word.begin(), it));
-
-                if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
-                    new_word.push_back(first + second);
-                    i += 2;
-                } else {
-                    new_word.push_back(word[i]);
-                    i += 1;
-                }
-            }
-
-            word = new_word;
-
-            if (word.size() == 1) {
-                break;
-            }
-            pairs = get_pairs(word);
-        }
-
-        std::u32string result;
-        for (int i = 0; i < word.size(); i++) {
-            result += word[i];
-            if (i != word.size() - 1) {
-                result += utf8_to_utf32(" ");
-            }
-        }
-
-        return result;
-    }
-
-    std::vector<int> tokenize(std::string text,
-                              on_new_token_cb_t on_new_token_cb,
-                              size_t max_length = 0,
-                              bool padding      = false) {
-        std::vector<int32_t> tokens = encode(text, on_new_token_cb);
-
-        tokens.insert(tokens.begin(), BOS_TOKEN_ID);
-        if (max_length > 0) {
-            if (tokens.size() > max_length - 1) {
-                tokens.resize(max_length - 1);
-                tokens.push_back(EOS_TOKEN_ID);
-            } else {
-                tokens.push_back(EOS_TOKEN_ID);
-                if (padding) {
-                    tokens.insert(tokens.end(), max_length - tokens.size(), PAD_TOKEN_ID);
-                }
-            }
-        }
-
-        return tokens;
-    }
-
-    void pad_tokens(std::vector<int>& tokens,
-                    std::vector<float>& weights,
-                    size_t max_length = 0,
-                    bool padding      = false) {
-        if (max_length > 0 && padding) {
-            size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.0 / (max_length - 2)));
-            if (n == 0) {
-                n = 1;
-            }
-            size_t length = max_length * n;
-            LOG_DEBUG("token length: %llu", length);
-            std::vector<int> new_tokens;
-            std::vector<float> new_weights;
-            new_tokens.push_back(BOS_TOKEN_ID);
-            new_weights.push_back(1.0);
-            int token_idx = 0;
-            for (int i = 1; i < length; i++) {
-                if (token_idx >= tokens.size()) {
-                    break;
-                }
-                if (i % max_length == 0) {
-                    new_tokens.push_back(BOS_TOKEN_ID);
-                    new_weights.push_back(1.0);
-                } else if (i % max_length == max_length - 1) {
-                    new_tokens.push_back(EOS_TOKEN_ID);
-                    new_weights.push_back(1.0);
-                } else {
-                    new_tokens.push_back(tokens[token_idx]);
-                    new_weights.push_back(weights[token_idx]);
-                    token_idx++;
-                }
-            }
-
-            new_tokens.push_back(EOS_TOKEN_ID);
-            new_weights.push_back(1.0);
-            tokens  = new_tokens;
-            weights = new_weights;
-
-            if (padding) {
-                tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
-                weights.insert(weights.end(), length - weights.size(), 1.0);
-            }
-        }
-    }
-
-    std::string clean_up_tokenization(std::string& text) {
-        std::regex pattern(R"( ,)");
-        // Replace " ," with ","
-        std::string result = std::regex_replace(text, pattern, ",");
-        return result;
-    }
-
-    std::string decode(const std::vector<int>& tokens) {
-        std::string text = "";
-        for (int t : tokens) {
-            if (t == 49406 || t == 49407)
-                continue;
-            std::u32string ts = decoder[t];
-            // printf("%d, %s \n", t,  utf32_to_utf8(ts).c_str());
-            std::string s = utf32_to_utf8(ts);
-            if (s.length() >= 4) {
-                if (ends_with(s, "</w>")) {
-                    text += s.replace(s.length() - 4, s.length() - 1, "") + " ";
-                } else {
-                    text += s;
-                }
-            } else {
-                text += " " + s;
-            }
-        }
-        // std::vector<unsigned char> bytes;
-        // for (auto c : text){
-        //     bytes.push_back(byte_decoder[c]);
-        // }
-
-        // std::string s((char *)bytes.data());
-        // std::string s = "";
-        text = clean_up_tokenization(text);
-        return trim(text);
-    }
-
-    std::vector<std::string> token_split(const std::string& text) {
-        std::regex pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
-                       std::regex::icase);
-        std::sregex_iterator iter(text.begin(), text.end(), pat);
-        std::sregex_iterator end;
-
-        std::vector<std::string> result;
-        for (; iter != end; ++iter) {
-            result.emplace_back(iter->str());
-        }
-
-        return result;
-    }
-
-    std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb) {
-        std::string original_text = text;
-        std::vector<int32_t> bpe_tokens;
-        text = whitespace_clean(text);
-        std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return std::tolower(c); });
-
-        std::string str = text;
-        std::vector<std::string> token_strs;
-
-        auto splited_texts = split_with_special_tokens(text, special_tokens);
-
-        for (auto& splited_text : splited_texts) {
-            LOG_DEBUG("token %s", splited_text.c_str());
-            if (is_special_token(splited_text)) {
-                LOG_DEBUG("special %s", splited_text.c_str());
-                bool skip = on_new_token_cb(splited_text, bpe_tokens);
-                if (skip) {
-                    token_strs.push_back(splited_text);
-                    continue;
-                }
-                continue;
-            }
-
-            auto tokens = token_split(splited_text);
-            for (auto& token : tokens) {
-                if (on_new_token_cb != nullptr) {
-                    bool skip = on_new_token_cb(token, bpe_tokens);
-                    if (skip) {
-                        token_strs.push_back(token);
-                        continue;
-                    }
-                }
-
-                std::string token_str = token;
-                std::u32string utf32_token;
-                for (int i = 0; i < token_str.length(); i++) {
-                    unsigned char b = token_str[i];
-                    utf32_token += byte_encoder[b];
-                }
-                auto bpe_strs = bpe(utf32_token);
-                size_t start  = 0;
-                size_t pos;
-                while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {
-                    auto bpe_str = bpe_strs.substr(start, pos - start);
-                    bpe_tokens.push_back(encoder[bpe_str]);
-                    token_strs.push_back(utf32_to_utf8(bpe_str));
-
-                    start = pos + 1;
-                }
-                auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);
-                bpe_tokens.push_back(encoder[bpe_str]);
-                token_strs.push_back(utf32_to_utf8(bpe_str));
-            }
-        }
-        // std::stringstream ss;
-        // ss << "[";
-        // for (auto token : token_strs) {
-        //     ss << "\"" << token << "\", ";
-        // }
-        // ss << "]";
-        // LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
-        // printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str());
-        return bpe_tokens;
-    }
-};
+#include "tokenizers/clip_tokenizer.h"
 
 /*================================================ FrozenCLIPEmbedder ================================================*/
 

src/conditioner.hpp

@@ -256,15 +256,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         return true;
     }
 
-    std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
-    tokenize_with_trigger_token(std::string text,
-                                int num_input_imgs,
-                                int32_t image_token,
-                                bool padding = false) {
-        return tokenize_with_trigger_token(text, num_input_imgs, image_token,
-                                           text_model->model.n_token, padding);
-    }
-
     std::vector<int> convert_token_to_id(std::string text) {
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
             auto iter = embedding_map.find(str);
@@ -288,9 +279,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
     tokenize_with_trigger_token(std::string text,
                                 int num_input_imgs,
-                                int32_t image_token,
-                                size_t max_length = 0,
-                                bool padding      = false) {
+                                int32_t image_token) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -377,7 +366,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         // tokens.insert(tokens.begin(), tokenizer.BOS_TOKEN_ID);
         // weights.insert(weights.begin(), 1.0);
 
-        tokenizer.pad_tokens(tokens, weights, max_length, padding);
+        tokenizer.pad_tokens(tokens, &weights, nullptr, text_model->model.n_token, text_model->model.n_token, true);
         int offset = pm_version == PM_VERSION_2 ? 2 * num_input_imgs : num_input_imgs;
         for (int i = 0; i < tokens.size(); i++) {
             // if (class_idx + 1 <= i && i < class_idx + 1 + 2*num_input_imgs) // photomaker V2 has num_tokens(=2)*num_input_imgs
@@ -403,13 +392,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     }
 
     std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
-                                                             bool padding = false) {
-        return tokenize(text, text_model->model.n_token, padding);
-    }
-
-    std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
-                                                             size_t max_length = 0,
-                                                             bool padding      = false) {
+                                                             size_t min_length          = 0,
+                                                             size_t max_length          = 0,
+                                                             bool allow_overflow_expand = true) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -460,7 +445,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
-        tokenizer.pad_tokens(tokens, weights, max_length, padding);
+        tokenizer.pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", ";
@@ -603,8 +588,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         GGML_ASSERT(image_tokens.size() == 1);
         auto tokens_and_weights     = tokenize_with_trigger_token(conditioner_params.text,
                                                                   conditioner_params.num_input_imgs,
-                                                                  image_tokens[0],
-                                                                  true);
+                                                                  image_tokens[0]);
         std::vector<int>& tokens    = std::get<0>(tokens_and_weights);
         std::vector<float>& weights = std::get<1>(tokens_and_weights);
         std::vector<bool>& clsm     = std::get<2>(tokens_and_weights);
@@ -630,7 +614,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::string remove_trigger_from_prompt(const std::string& prompt) override {
         auto image_tokens = convert_token_to_id(trigger_word);
         GGML_ASSERT(image_tokens.size() == 1);
-        auto tokens_and_weights  = tokenize(prompt, false);
+        auto tokens_and_weights  = tokenize(prompt);
         std::vector<int>& tokens = tokens_and_weights.first;
         auto it                  = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
         GGML_ASSERT(it != tokens.end());  // prompt must have trigger word
@@ -640,7 +624,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights     = tokenize(conditioner_params.text, true);
+        auto tokens_and_weights     = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
         std::vector<int>& tokens    = tokens_and_weights.first;
         std::vector<float>& weights = tokens_and_weights.second;
         return get_learned_condition_common(n_threads,
@@ -822,8 +806,9 @@ struct SD3CLIPEmbedder : public Conditioner {
     }
 
     std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
-                                                                          size_t max_length = 0,
-                                                                          bool padding      = false) {
+                                                                          size_t min_length          = 0,
+                                                                          size_t max_length          = 0,
+                                                                          bool allow_overflow_expand = true) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -860,20 +845,20 @@ struct SD3CLIPEmbedder : public Conditioner {
                 clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (t5) {
-                std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
         }
 
         if (clip_l) {
-            clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, max_length, padding);
+            clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
         }
         if (clip_g) {
-            clip_g_tokenizer.pad_tokens(clip_g_tokens, clip_g_weights, max_length, padding);
+            clip_g_tokenizer.pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
         }
         if (t5) {
-            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding);
+            t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
         }
 
         // for (int i = 0; i < clip_l_tokens.size(); i++) {
@@ -1056,7 +1041,7 @@ struct SD3CLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights = tokenize(conditioner_params.text, 77, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, 77, 77, true);
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1158,8 +1143,8 @@ struct FluxCLIPEmbedder : public Conditioner {
     }
 
     std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
-                                                                          size_t max_length = 0,
-                                                                          bool padding      = false) {
+                                                                          size_t min_length = 0,
+                                                                          size_t max_length = 0) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -1189,17 +1174,17 @@ struct FluxCLIPEmbedder : public Conditioner {
                 clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (t5) {
-                std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
         }
 
         if (clip_l) {
-            clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, 77, padding);
+            clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
         }
         if (t5) {
-            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding);
+            t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
         }
 
         // for (int i = 0; i < clip_l_tokens.size(); i++) {
@@ -1300,7 +1285,7 @@ struct FluxCLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1377,8 +1362,8 @@ struct T5CLIPEmbedder : public Conditioner {
     }
 
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
-                                                                                  size_t max_length = 0,
-                                                                                  bool padding      = false) {
+                                                                                  size_t min_length = 0,
+                                                                                  size_t max_length = 0) {
         auto parsed_attention = parse_prompt_attention(text);
 
         {
@@ -1403,12 +1388,15 @@ struct T5CLIPEmbedder : public Conditioner {
                 const std::string& curr_text = item.first;
                 float curr_weight            = item.second;
 
-                std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
 
-            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, &t5_mask, max_length, padding);
+            t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, &t5_mask, min_length, max_length, true);
+            for (auto& mask_value : t5_mask) {
+                mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
+            }
         }
         return {t5_tokens, t5_weights, t5_mask};
     }
@@ -1496,7 +1484,7 @@ struct T5CLIPEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1505,14 +1493,14 @@ struct T5CLIPEmbedder : public Conditioner {
 };
 
 struct AnimaConditioner : public Conditioner {
-    std::shared_ptr<LLM::BPETokenizer> qwen_tokenizer;
+    std::shared_ptr<BPETokenizer> qwen_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
 
     AnimaConditioner(ggml_backend_t backend,
                      bool offload_params_to_cpu,
                      const String2TensorStorage& tensor_storage_map = {}) {
-        qwen_tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
+        qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
         llm            = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
                                                backend,
                                                offload_params_to_cpu,
@@ -1578,7 +1566,7 @@ struct AnimaConditioner : public Conditioner {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+            std::vector<int> curr_tokens = t5_tokenizer.tokenize(curr_text, nullptr, true);
             t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
             t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
         }
@@ -1620,7 +1608,7 @@ struct AnimaConditioner : public Conditioner {
 
 struct LLMEmbedder : public Conditioner {
     SDVersion version;
-    std::shared_ptr<LLM::BPETokenizer> tokenizer;
+    std::shared_ptr<BPETokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
 
     LLMEmbedder(ggml_backend_t backend,
@@ -1633,13 +1621,15 @@ struct LLMEmbedder : public Conditioner {
         LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
         if (version == VERSION_FLUX2) {
             arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
+        } else if (sd_version_is_ernie_image(version)) {
+            arch = LLM::LLMArch::MINISTRAL_3_3B;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
         }
-        if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2) {
-            tokenizer = std::make_shared<LLM::MistralTokenizer>();
+        if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
+            tokenizer = std::make_shared<MistralTokenizer>();
         } else {
-            tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
+            tokenizer = std::make_shared<Qwen2Tokenizer>();
         }
         llm = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
@@ -1677,10 +1667,10 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
-    std::tuple<std::vector<int>, std::vector<float>> tokenize(std::string text,
-                                                              const std::pair<int, int>& attn_range,
-                                                              size_t max_length = 0,
-                                                              bool padding      = false) {
+    std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
+                                                                                  const std::pair<int, int>& attn_range,
+                                                                                  size_t min_length = 0,
+                                                                                  size_t max_length = 100000000) {
         std::vector<std::pair<std::string, float>> parsed_attention;
         if (attn_range.first >= 0 && attn_range.second > 0) {
             parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
@@ -1710,39 +1700,34 @@ struct LLMEmbedder : public Conditioner {
         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);
+            std::vector<int> curr_tokens = tokenizer->encode(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);
+        std::vector<float> mask;
+        tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl;
         // }
         // std::cout << std::endl;
 
-        return {tokens, weights};
+        return {tokens, weights, mask};
     }
 
     sd::Tensor<float> encode_prompt(int n_threads,
                                     const std::string prompt,
                                     const std::pair<int, int>& prompt_attn_range,
-                                    int max_length,
                                     int min_length,
+                                    int hidden_states_min_length,
                                     const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
                                     const std::set<int>& out_layers,
                                     int prompt_template_encode_start_idx) {
-        auto tokens_and_weights = tokenize(prompt, prompt_attn_range);
-        auto& tokens            = std::get<0>(tokens_and_weights);
-        auto& weights           = std::get<1>(tokens_and_weights);
-        std::vector<float> mask;
-
-        if (max_length > 0 && tokens.size() < max_length) {
-            mask.insert(mask.end(), tokens.size(), 1.f);
-            mask.insert(mask.end(), max_length - tokens.size(), 0.f);
-            tokenizer->pad_tokens(tokens, weights, max_length, true);
-        }
+        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length);
+        auto& tokens             = std::get<0>(tokens_weights_mask);
+        auto& weights            = std::get<1>(tokens_weights_mask);
+        auto& mask               = std::get<2>(tokens_weights_mask);
 
         sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
         sd::Tensor<float> attention_mask;
@@ -1769,9 +1754,9 @@ struct LLMEmbedder : public Conditioner {
         GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
 
         int64_t zero_pad_len = 0;
-        if (min_length > 0) {
-            if (hidden_states.shape()[1] - prompt_template_encode_start_idx < min_length) {
-                zero_pad_len = min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
+        if (hidden_states_min_length > 0) {
+            if (hidden_states.shape()[1] - prompt_template_encode_start_idx < hidden_states_min_length) {
+                zero_pad_len = hidden_states_min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
             }
         }
 
@@ -1798,8 +1783,8 @@ struct LLMEmbedder : public Conditioner {
         std::vector<std::pair<int, int>> extra_prompts_attn_range;
         std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
         int prompt_template_encode_start_idx = 34;
-        int max_length                       = 0;  // pad tokens
-        int min_length                       = 0;  // zero pad hidden_states
+        int min_length                       = 0;  // pad tokens
+        int hidden_states_min_length         = 0;  // zero pad hidden_states
         std::set<int> out_layers;
 
         int64_t t0 = ggml_time_ms();
@@ -1874,7 +1859,7 @@ struct LLMEmbedder : public Conditioner {
             }
         } else if (version == VERSION_FLUX2) {
             prompt_template_encode_start_idx = 0;
-            min_length                       = 512;
+            hidden_states_min_length         = 512;
             out_layers                       = {10, 20, 30};
 
             prompt = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
@@ -1884,6 +1869,13 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "[/INST]";
+        } else if (sd_version_is_ernie_image(version)) {
+            prompt_template_encode_start_idx = 0;
+            out_layers                       = {25};  // -2
+
+            prompt_attn_range.first = 0;
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
         } else if (sd_version_is_z_image(version)) {
             prompt_template_encode_start_idx = 0;
             out_layers                       = {35};  // -2
@@ -1907,7 +1899,7 @@ struct LLMEmbedder : public Conditioner {
             }
         } else if (version == VERSION_FLUX2_KLEIN) {
             prompt_template_encode_start_idx = 0;
-            max_length                       = 512;
+            min_length                       = 512;
             out_layers                       = {9, 18, 27};
 
             prompt = "<|im_start|>user\n";
@@ -1919,7 +1911,7 @@ struct LLMEmbedder : public Conditioner {
             prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
         } else if (version == VERSION_OVIS_IMAGE) {
             prompt_template_encode_start_idx = 28;
-            max_length                       = prompt_template_encode_start_idx + 256;
+            min_length                       = prompt_template_encode_start_idx + 256;
 
             prompt = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:";
 
@@ -1935,8 +1927,8 @@ struct LLMEmbedder : public Conditioner {
         auto hidden_states = encode_prompt(n_threads,
                                            prompt,
                                            prompt_attn_range,
-                                           max_length,
                                            min_length,
+                                           hidden_states_min_length,
                                            image_embeds,
                                            out_layers,
                                            prompt_template_encode_start_idx);
@@ -1945,8 +1937,8 @@ struct LLMEmbedder : public Conditioner {
             auto extra_hidden_states = encode_prompt(n_threads,
                                                      extra_prompts[i],
                                                      extra_prompts_attn_range[i],
-                                                     max_length,
                                                      min_length,
+                                                     hidden_states_min_length,
                                                      image_embeds,
                                                      out_layers,
                                                      prompt_template_encode_start_idx);

src/denoiser.hpp

@@ -658,32 +658,22 @@ inline float time_snr_shift(float alpha, float t) {
 }
 
 struct DiscreteFlowDenoiser : public Denoiser {
-    float sigmas[TIMESTEPS];
     float shift = 3.0f;
 
-    float sigma_data = 1.0f;
-
     DiscreteFlowDenoiser(float shift = 3.0f) {
         set_shift(shift);
     }
 
-    void set_parameters() {
-        for (int i = 1; i < TIMESTEPS + 1; i++) {
-            sigmas[i - 1] = t_to_sigma(static_cast<float>(i));
-        }
-    }
-
     void set_shift(float shift) {
         this->shift = shift;
-        set_parameters();
     }
 
     float sigma_min() override {
-        return sigmas[0];
+        return t_to_sigma(0);
     }
 
     float sigma_max() override {
-        return sigmas[TIMESTEPS - 1];
+        return t_to_sigma(TIMESTEPS - 1);
     }
 
     float sigma_to_t(float sigma) override {
@@ -786,10 +776,43 @@ static std::pair<float, float> get_ancestral_step(float sigma_from,
     return {sigma_down, sigma_up};
 }
 
+static std::tuple<float, float, float> get_ancestral_step_flow(float sigma_from,
+                                                               float sigma_to,
+                                                               float eta = 1.0f) {
+    float sigma_down  = sigma_to;
+    float sigma_up    = 0.0f;
+    float alpha_scale = 1.0f;
+
+    if (eta <= 0.0f || sigma_from <= 0.0f || sigma_to <= 0.0f) {
+        return {sigma_down, sigma_up, alpha_scale};
+    }
+
+    // Flow Euler ancestral sampling becomes numerically unstable for eta > 1, so
+    // clamp to the valid maximum-noise regime instead of letting NaNs propagate.
+    eta = std::min(eta, 1.0f);
+
+    float sigma_ratio = sigma_to / sigma_from;
+    sigma_down        = sigma_to * (1.0f + (sigma_ratio - 1.0f) * eta);
+    sigma_down        = std::max(0.0f, std::min(sigma_to, sigma_down));
+
+    float denom = 1.0f - sigma_down;
+    if (denom <= 0.0f) {
+        return {sigma_to, sigma_up, alpha_scale};
+    }
+
+    alpha_scale = (1.0f - sigma_to) / denom;
+
+    float term = (sigma_down / sigma_to) * alpha_scale;
+    term       = std::max(-1.0f, std::min(1.0f, term));
+    sigma_up   = sigma_to * std::sqrt(std::max(1.0f - term * term, 0.0f));
+    return {sigma_down, sigma_up, alpha_scale};
+}
+
 static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
                                                 sd::Tensor<float> x,
                                                 const std::vector<float>& sigmas,
-                                                std::shared_ptr<RNG> rng) {
+                                                std::shared_ptr<RNG> rng,
+                                                float eta) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
@@ -799,7 +822,7 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
         }
         sd::Tensor<float> denoised  = std::move(denoised_opt);
         sd::Tensor<float> d         = (x - denoised) / sigma;
-        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1]);
+        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1], eta);
         x += d * (sigma_down - sigmas[i]);
         if (sigmas[i + 1] > 0) {
             x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
@@ -808,6 +831,30 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_euler_flow(denoise_cb_t model,
+                                           sd::Tensor<float> x,
+                                           const std::vector<float>& sigmas,
+                                           std::shared_ptr<RNG> rng,
+                                           float eta) {
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma       = sigmas[i];
+        auto denoised_opt = model(x, sigma, i + 1);
+        if (denoised_opt.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised               = std::move(denoised_opt);
+        auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step_flow(sigma, sigmas[i + 1], eta);
+        float sigma_ratio                        = sigma_down / sigma;
+        x                                        = sigma_ratio * x + (1.0f - sigma_ratio) * denoised;
+
+        if (sigma_up > 0.0f) {
+            x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+        }
+    }
+    return x;
+}
+
 static sd::Tensor<float> sample_euler(denoise_cb_t model,
                                       sd::Tensor<float> x,
                                       const std::vector<float>& sigmas) {
@@ -885,7 +932,8 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
 static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
                                                    sd::Tensor<float> x,
                                                    const std::vector<float>& sigmas,
-                                                   std::shared_ptr<RNG> rng) {
+                                                   std::shared_ptr<RNG> rng,
+                                                   float eta) {
     auto t_fn     = [](float sigma) -> float { return -log(sigma); };
     auto sigma_fn = [](float t) -> float { return exp(-t); };
 
@@ -896,7 +944,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
             return {};
         }
         sd::Tensor<float> denoised  = std::move(denoised_opt);
-        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1]);
+        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1], eta);
 
         if (sigma_down == 0) {
             x = denoised;
@@ -1368,10 +1416,14 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             sd::Tensor<float> x,
                                             std::vector<float> sigmas,
                                             std::shared_ptr<RNG> rng,
-                                            float eta) {
+                                            float eta,
+                                            bool is_flow_denoiser) {
     switch (method) {
         case EULER_A_SAMPLE_METHOD:
-            return sample_euler_ancestral(model, std::move(x), sigmas, rng);
+            if (is_flow_denoiser)
+                return sample_euler_flow(model, std::move(x), sigmas, rng, eta);
+            else
+                return sample_euler_ancestral(model, std::move(x), sigmas, rng, eta);
         case EULER_SAMPLE_METHOD:
             return sample_euler(model, std::move(x), sigmas);
         case HEUN_SAMPLE_METHOD:
@@ -1379,7 +1431,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
         case DPM2_SAMPLE_METHOD:
             return sample_dpm2(model, std::move(x), sigmas);
         case DPMPP2S_A_SAMPLE_METHOD:
-            return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng);
+            return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta);
         case DPMPP2M_SAMPLE_METHOD:
             return sample_dpmpp_2m(model, std::move(x), sigmas);
         case DPMPP2Mv2_SAMPLE_METHOD:

src/diffusion_model.hpp

@@ -3,6 +3,7 @@
 
 #include <optional>
 #include "anima.hpp"
+#include "ernie_image.hpp"
 #include "flux.hpp"
 #include "mmdit.hpp"
 #include "qwen_image.hpp"
@@ -516,4 +517,66 @@ struct ZImageModel : public DiffusionModel {
     }
 };
 
+struct ErnieImageModel : public DiffusionModel {
+    std::string prefix;
+    ErnieImage::ErnieImageRunner ernie_image;
+
+    ErnieImageModel(ggml_backend_t backend,
+                    bool offload_params_to_cpu,
+                    const String2TensorStorage& tensor_storage_map = {},
+                    const std::string prefix                       = "model.diffusion_model")
+        : prefix(prefix), ernie_image(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
+    }
+
+    std::string get_desc() override {
+        return ernie_image.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        ernie_image.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        ernie_image.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        ernie_image.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        ernie_image.get_param_tensors(tensors, prefix);
+    }
+
+    size_t get_params_buffer_size() override {
+        return ernie_image.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        ernie_image.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768;
+    }
+
+    void set_flash_attention_enabled(bool enabled) {
+        ernie_image.set_flash_attention_enabled(enabled);
+    }
+
+    void set_circular_axes(bool circular_x, bool circular_y) override {
+        ernie_image.set_circular_axes(circular_x, circular_y);
+    }
+
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        return ernie_image.compute(n_threads,
+                                   *diffusion_params.x,
+                                   *diffusion_params.timesteps,
+                                   tensor_or_empty(diffusion_params.context));
+    }
+};
+
 #endif

src/ernie_image.hpp

@@ -0,0 +1,438 @@
+#ifndef __SD_ERNIE_IMAGE_HPP__
+#define __SD_ERNIE_IMAGE_HPP__
+
+#include <memory>
+#include <vector>
+
+#include "common_dit.hpp"
+#include "flux.hpp"
+#include "qwen_image.hpp"
+#include "rope.hpp"
+
+namespace ErnieImage {
+    constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960;
+
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+                                                              ggml_tensor* timesteps,
+                                                              int dim,
+                                                              int max_period = 10000) {
+        auto emb       = ggml_ext_timestep_embedding(ctx, timesteps, dim, max_period, 1.0f);
+        int64_t half   = dim / 2;
+        auto cos_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], 0);
+        auto sin_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], half * emb->nb[0]);
+        auto sin_first = ggml_concat(ctx, sin_part, cos_part, 0);
+        return sin_first;
+    }
+
+    __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(ggml_context* ctx, ggml_tensor* x, ggml_tensor* pe) {
+        // x: [N, S, heads, head_dim]
+        // pe: [2, S, 1, head_dim], stored as ggml [head_dim, 1, S, 2].
+        int64_t head_dim = x->ne[0];
+        int64_t heads    = x->ne[1];
+        int64_t S        = x->ne[2];
+        int64_t N        = x->ne[3];
+        int64_t rot_dim  = pe->ne[0];
+        GGML_ASSERT(rot_dim <= head_dim);
+        GGML_ASSERT(rot_dim % 2 == 0);
+        GGML_ASSERT(pe->ne[1] == 1 && pe->ne[2] == S && pe->ne[3] == 2);
+
+        x           = ggml_cont(ctx, x);
+        auto x_rot  = ggml_ext_slice(ctx, x, 0, 0, rot_dim, false);
+        auto x_pass = rot_dim < head_dim ? ggml_ext_slice(ctx, x, 0, rot_dim, head_dim, false) : nullptr;
+
+        int64_t half = rot_dim / 2;
+        auto x1      = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], 0);
+        auto x2      = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], half * x_rot->nb[0]);
+        x1           = ggml_cont(ctx, x1);
+        x2           = ggml_cont(ctx, x2);
+        auto rotated = ggml_concat(ctx, ggml_neg(ctx, x2), x1, 0);
+
+        auto cos_emb = ggml_ext_slice(ctx, pe, 3, 0, 1, false);
+        auto sin_emb = ggml_ext_slice(ctx, pe, 3, 1, 2, false);
+
+        auto out = ggml_add(ctx, ggml_mul(ctx, x_rot, cos_emb), ggml_mul(ctx, rotated, sin_emb));
+        if (x_pass != nullptr) {
+            out = ggml_concat(ctx, out, x_pass, 0);
+        }
+        return out;
+    }
+
+    struct ErnieImageAttention : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        ErnieImageAttention(int64_t query_dim,
+                            int64_t heads,
+                            int64_t dim_head,
+                            float eps = 1e-6f)
+            : num_heads(heads), head_dim(dim_head) {
+            int64_t inner_dim  = heads * dim_head;
+            blocks["to_q"]     = std::make_shared<Linear>(query_dim, inner_dim, false);
+            blocks["to_k"]     = std::make_shared<Linear>(query_dim, inner_dim, false);
+            blocks["to_v"]     = std::make_shared<Linear>(query_dim, inner_dim, false);
+            blocks["norm_q"]   = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["norm_k"]   = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["to_out.0"] = std::make_shared<Linear>(inner_dim, query_dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* attention_mask = nullptr) {
+            // x: [N, S, hidden_size]
+            // pe: [S, head_dim/2, 2, 2], generated in image-token-first order.
+            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 norm_q   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
+            auto norm_k   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
+            auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+
+            int64_t S = x->ne[1];
+            int64_t N = x->ne[2];
+
+            auto q = to_q->forward(ctx, x);
+            auto k = to_k->forward(ctx, x);
+            auto v = to_v->forward(ctx, x);
+
+            q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N);  // [N, S, heads, head_dim]
+            k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);  // [N, S, heads, head_dim]
+            v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);  // [N, S, heads, head_dim]
+
+            q = norm_q->forward(ctx, q);
+            k = norm_k->forward(ctx, k);
+
+            q = apply_rotary_emb(ctx->ggml_ctx, q, pe);
+            k = apply_rotary_emb(ctx->ggml_ctx, k, pe);
+
+            q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));  // [N, heads, S, head_dim]
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
+
+            k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, heads, S, head_dim]
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
+
+            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, S, hidden_size]
+            x = to_out_0->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct ErnieImageFeedForward : public GGMLBlock {
+    public:
+        ErnieImageFeedForward(int64_t hidden_size, int64_t ffn_hidden_size) {
+            blocks["gate_proj"]  = std::make_shared<Linear>(hidden_size, ffn_hidden_size, false);
+            blocks["up_proj"]    = std::make_shared<Linear>(hidden_size, ffn_hidden_size, false);
+            blocks["linear_fc2"] = std::make_shared<Linear>(ffn_hidden_size, hidden_size, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto gate_proj  = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
+            auto up_proj    = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
+            auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
+
+            auto gate = gate_proj->forward(ctx, x);
+            gate      = ggml_ext_gelu(ctx->ggml_ctx, gate);
+            x         = up_proj->forward(ctx, x);
+            x         = ggml_mul(ctx->ggml_ctx, x, gate);
+            x         = linear_fc2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct ErnieImageSharedAdaLNBlock : public GGMLBlock {
+    public:
+        ErnieImageSharedAdaLNBlock(int64_t hidden_size,
+                                   int64_t num_heads,
+                                   int64_t ffn_hidden_size,
+                                   float eps = 1e-6f) {
+            blocks["adaLN_sa_ln"]    = std::make_shared<RMSNorm>(hidden_size, eps);
+            blocks["self_attention"] = std::make_shared<ErnieImageAttention>(hidden_size,
+                                                                             num_heads,
+                                                                             hidden_size / num_heads,
+                                                                             eps);
+            blocks["adaLN_mlp_ln"]   = std::make_shared<RMSNorm>(hidden_size, eps);
+            blocks["mlp"]            = std::make_shared<ErnieImageFeedForward>(hidden_size, ffn_hidden_size);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             const std::vector<ggml_tensor*>& temb,
+                             ggml_tensor* attention_mask = nullptr) {
+            // x: [N, image_tokens + text_tokens, hidden_size]
+            auto adaLN_sa_ln    = std::dynamic_pointer_cast<RMSNorm>(blocks["adaLN_sa_ln"]);
+            auto self_attention = std::dynamic_pointer_cast<ErnieImageAttention>(blocks["self_attention"]);
+            auto adaLN_mlp_ln   = std::dynamic_pointer_cast<RMSNorm>(blocks["adaLN_mlp_ln"]);
+            auto mlp            = std::dynamic_pointer_cast<ErnieImageFeedForward>(blocks["mlp"]);
+
+            auto shift_msa = temb[0];
+            auto scale_msa = temb[1];
+            auto gate_msa  = temb[2];
+            auto shift_mlp = temb[3];
+            auto scale_mlp = temb[4];
+            auto gate_mlp  = temb[5];
+
+            auto residual = x;
+            x             = adaLN_sa_ln->forward(ctx, x);
+            x             = Flux::modulate(ctx->ggml_ctx, x, shift_msa, scale_msa, true);
+            auto attn_out = self_attention->forward(ctx, x, pe, attention_mask);
+            x             = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
+
+            residual = x;
+            x        = adaLN_mlp_ln->forward(ctx, x);
+            x        = Flux::modulate(ctx->ggml_ctx, x, shift_mlp, scale_mlp, true);
+            x        = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, x), gate_mlp));
+            return x;
+        }
+    };
+
+    struct ErnieImageAdaLNContinuous : public GGMLBlock {
+    public:
+        ErnieImageAdaLNContinuous(int64_t hidden_size, float eps = 1e-6f) {
+            blocks["norm"]   = std::make_shared<LayerNorm>(hidden_size, eps, false);
+            blocks["linear"] = std::make_shared<Linear>(hidden_size, hidden_size * 2, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
+            auto norm   = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+
+            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, conditioning), 2, 0);
+            auto scale = mods[0];
+            auto shift = mods[1];
+
+            x = norm->forward(ctx, x);
+            x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+            return x;
+        }
+    };
+
+    struct ErnieImageParams {
+        int64_t hidden_size       = 4096;
+        int64_t num_heads         = 32;
+        int64_t num_layers        = 36;
+        int64_t ffn_hidden_size   = 12288;
+        int64_t in_channels       = 128;
+        int64_t out_channels      = 128;
+        int patch_size            = 1;
+        int64_t text_in_dim       = 3072;
+        int theta                 = 256;
+        std::vector<int> axes_dim = {32, 48, 48};
+        int axes_dim_sum          = 128;
+        float eps                 = 1e-6f;
+    };
+
+    class ErnieImageModel : public GGMLBlock {
+    public:
+        ErnieImageParams params;
+
+        ErnieImageModel() = default;
+        ErnieImageModel(ErnieImageParams params)
+            : params(params) {
+            blocks["x_embedder.proj"] = std::make_shared<Conv2d>(params.in_channels,
+                                                                 params.hidden_size,
+                                                                 std::pair<int, int>{params.patch_size, params.patch_size},
+                                                                 std::pair<int, int>{params.patch_size, params.patch_size},
+                                                                 std::pair<int, int>{0, 0},
+                                                                 std::pair<int, int>{1, 1},
+                                                                 true);
+            if (params.text_in_dim != params.hidden_size) {
+                blocks["text_proj"] = std::make_shared<Linear>(params.text_in_dim, params.hidden_size, false);
+            }
+            blocks["time_embedding"]     = std::make_shared<Qwen::TimestepEmbedding>(params.hidden_size, params.hidden_size);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(params.hidden_size, 6 * params.hidden_size, true);
+
+            for (int i = 0; i < params.num_layers; i++) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(params.hidden_size,
+                                                                                                     params.num_heads,
+                                                                                                     params.ffn_hidden_size,
+                                                                                                     params.eps);
+            }
+
+            blocks["final_norm"]   = std::make_shared<ErnieImageAdaLNContinuous>(params.hidden_size, params.eps);
+            blocks["final_linear"] = std::make_shared<Linear>(params.hidden_size,
+                                                              params.patch_size * params.patch_size * params.out_channels,
+                                                              true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe) {
+            // x: [N, C, H, W]
+            // context: [N, text_tokens, 3072]
+            // pe: [image_tokens + text_tokens, head_dim/2, 2, 2]
+            GGML_ASSERT(context != nullptr);
+            GGML_ASSERT(x->ne[1] % params.patch_size == 0 && x->ne[0] % params.patch_size == 0);
+
+            int64_t W     = x->ne[0];
+            int64_t H     = x->ne[1];
+            int64_t Hp    = H / params.patch_size;
+            int64_t Wp    = W / params.patch_size;
+            int64_t n_img = Hp * Wp;
+            int64_t N     = x->ne[3];
+
+            auto x_embedder_proj = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"]);
+            auto time_embedding  = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedding"]);
+            auto adaLN_mod       = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
+            auto final_norm      = std::dynamic_pointer_cast<ErnieImageAdaLNContinuous>(blocks["final_norm"]);
+            auto final_linear    = std::dynamic_pointer_cast<Linear>(blocks["final_linear"]);
+
+            auto img = x_embedder_proj->forward(ctx, x);                                                  // [N, hidden_size, Hp, Wp]
+            img      = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N);       // [N, hidden_size, image_tokens]
+            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));  // [N, image_tokens, hidden_size]
+
+            auto txt       = context;
+            auto text_proj = std::dynamic_pointer_cast<Linear>(blocks["text_proj"]);
+            if (text_proj) {
+                txt = text_proj->forward(ctx, txt);
+            }
+
+            auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);  // [N, image_tokens + text_tokens, hidden_size]
+
+            auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(params.hidden_size));
+            auto c      = time_embedding->forward(ctx, sample);  // [N, hidden_size]
+
+            auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 6 * hidden_size]
+            auto chunks     = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
+            std::vector<ggml_tensor*> temb;
+            temb.reserve(6);
+            for (auto chunk : chunks) {
+                temb.push_back(ggml_reshape_3d(ctx->ggml_ctx, chunk, chunk->ne[0], 1, chunk->ne[1]));  // [N, 1, hidden_size]
+            }
+
+            for (int i = 0; i < params.num_layers; i++) {
+                auto layer    = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
+                hidden_states = layer->forward(ctx, hidden_states, pe, temb);
+            }
+
+            hidden_states = final_norm->forward(ctx, hidden_states, c);
+            hidden_states = final_linear->forward(ctx, hidden_states);                  // [N, image_tokens, p*p*out_channels]
+            auto patches  = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);  // [N, image_tokens, hidden_size]
+
+            auto out = DiT::unpatchify(ctx->ggml_ctx,
+                                       patches,
+                                       Hp,
+                                       Wp,
+                                       params.patch_size,
+                                       params.patch_size,
+                                       false);  // [N, out_channels, H, W]
+            return out;
+        }
+    };
+
+    struct ErnieImageRunner : public GGMLRunner {
+        ErnieImageParams ernie_params;
+        ErnieImageModel ernie_image;
+        std::vector<float> pe_vec;
+
+        ErnieImageRunner(ggml_backend_t backend,
+                         bool offload_params_to_cpu,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "")
+            : GGMLRunner(backend, offload_params_to_cpu) {
+            ernie_params.num_layers = 0;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
+                    ernie_params.patch_size  = static_cast<int>(tensor_storage.ne[0]);
+                    ernie_params.in_channels = tensor_storage.ne[2];
+                    ernie_params.hidden_size = tensor_storage.ne[3];
+                } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
+                    ernie_params.text_in_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
+                    int64_t head_dim       = tensor_storage.ne[0];
+                    ernie_params.num_heads = ernie_params.hidden_size / head_dim;
+                } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
+                    ernie_params.ffn_hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t out_dim           = tensor_storage.ne[1];
+                    ernie_params.out_channels = out_dim / ernie_params.patch_size / ernie_params.patch_size;
+                }
+
+                size_t pos = name.find("layers.");
+                if (pos != std::string::npos) {
+                    std::string layer_name = name.substr(pos);
+                    auto items             = split_string(layer_name, '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > ernie_params.num_layers) {
+                            ernie_params.num_layers = block_index + 1;
+                        }
+                    }
+                }
+            }
+            if (ernie_params.num_layers == 0) {
+                ernie_params.num_layers = 36;
+            }
+            ernie_params.axes_dim_sum = 0;
+            for (int axis_dim : ernie_params.axes_dim) {
+                ernie_params.axes_dim_sum += axis_dim;
+            }
+
+            LOG_INFO("ernie_image: layers = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64
+                     ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                     ernie_params.num_layers,
+                     ernie_params.hidden_size,
+                     ernie_params.num_heads,
+                     ernie_params.ffn_hidden_size,
+                     ernie_params.in_channels,
+                     ernie_params.out_channels);
+
+            ernie_image = ErnieImageModel(ernie_params);
+            ernie_image.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "ernie_image";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+            ernie_image.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor) {
+            ggml_cgraph* gf        = new_graph_custom(ERNIE_IMAGE_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            pe_vec      = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
+                                                   static_cast<int>(x->ne[0]),
+                                                   ernie_params.patch_size,
+                                                   static_cast<int>(x->ne[3]),
+                                                   static_cast<int>(context->ne[1]),
+                                                   ernie_params.theta,
+                                                   circular_y_enabled,
+                                                   circular_x_enabled,
+                                                   ernie_params.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / ernie_params.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, ernie_params.axes_dim_sum, 1, pos_len, 2);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        }
+    };
+}  // namespace ErnieImage
+
+#endif  // __SD_ERNIE_IMAGE_HPP__

src/gguf_reader.hpp

@@ -59,6 +59,9 @@ private:
         if (!safe_read(fin, key_len))
             return false;
 
+        if (key_len > 4096)
+            return false;
+
         std::string key(key_len, '\0');
         if (!safe_read(fin, (char*)key.data(), key_len))
             return false;

src/llm.hpp

@@ -14,469 +14,21 @@
 #include <utility>
 #include <vector>
 
-#include "clip.hpp"
 #include "ggml_extend.hpp"
 #include "json.hpp"
 #include "rope.hpp"
-#include "tokenize_util.h"
-#include "vocab/vocab.h"
+#include "tokenizers/bpe_tokenizer.h"
+#include "tokenizers/mistral_tokenizer.h"
+#include "tokenizers/qwen2_tokenizer.h"
 
 namespace LLM {
     constexpr int LLM_GRAPH_SIZE = 10240;
 
-    class BPETokenizer {
-    protected:
-        std::map<int, std::u32string> byte_encoder;
-        std::map<std::u32string, int> byte_decoder;
-        std::map<std::u32string, int> encoder;
-        std::map<int, std::u32string> decoder;
-        std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
-        std::regex pat;
-        int encoder_len;
-        int bpe_len;
-
-        std::string UNK_TOKEN;
-        std::string BOS_TOKEN;
-        std::string EOS_TOKEN;
-        std::string PAD_TOKEN;
-
-        int UNK_TOKEN_ID;
-        int BOS_TOKEN_ID;
-        int EOS_TOKEN_ID;
-        int PAD_TOKEN_ID;
-
-        std::vector<std::string> special_tokens;
-
-        bool add_bos_token = false;
-
-    protected:
-        static std::string strip(const std::string& str) {
-            std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
-            std::string::size_type end   = str.find_last_not_of(" \t\n\r\v\f");
-
-            if (start == std::string::npos) {
-                // String contains only whitespace characters
-                return "";
-            }
-
-            return str.substr(start, end - start + 1);
-        }
-
-        static std::string whitespace_clean(std::string text) {
-            text = std::regex_replace(text, std::regex(R"(\s+)"), " ");
-            text = strip(text);
-            return text;
-        }
-
-        static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
-            std::set<std::pair<std::u32string, std::u32string>> pairs;
-            if (subwords.size() == 0) {
-                return pairs;
-            }
-            std::u32string prev_subword = subwords[0];
-            for (int i = 1; i < subwords.size(); i++) {
-                std::u32string subword = subwords[i];
-                std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
-                pairs.insert(pair);
-                prev_subword = subword;
-            }
-            return pairs;
-        }
-
-        bool is_special_token(const std::string& token) {
-            for (auto& special_token : special_tokens) {
-                if (special_token == token) {
-                    return true;
-                }
-            }
-            return false;
-        }
-
-    public:
-        BPETokenizer() = default;
-
-        std::u32string bpe(const std::u32string& token) {
-            std::vector<std::u32string> word;
-
-            for (int i = 0; i < token.size(); i++) {
-                word.emplace_back(1, token[i]);
-            }
-
-            std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
-
-            if (pairs.empty()) {
-                return token;
-            }
-
-            while (true) {
-                auto min_pair_iter = std::min_element(pairs.begin(),
-                                                      pairs.end(),
-                                                      [&](const std::pair<std::u32string, std::u32string>& a,
-                                                          const std::pair<std::u32string, std::u32string>& b) {
-                                                          if (bpe_ranks.find(a) == bpe_ranks.end()) {
-                                                              return false;
-                                                          } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
-                                                              return true;
-                                                          }
-                                                          return bpe_ranks.at(a) < bpe_ranks.at(b);
-                                                      });
-
-                const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
-
-                if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
-                    break;
-                }
-
-                std::u32string first  = bigram.first;
-                std::u32string second = bigram.second;
-                std::vector<std::u32string> new_word;
-                int32_t i = 0;
-
-                while (i < word.size()) {
-                    auto it = std::find(word.begin() + i, word.end(), first);
-                    if (it == word.end()) {
-                        new_word.insert(new_word.end(), word.begin() + i, word.end());
-                        break;
-                    }
-                    new_word.insert(new_word.end(), word.begin() + i, it);
-                    i = static_cast<int32_t>(std::distance(word.begin(), it));
-
-                    if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
-                        new_word.push_back(first + second);
-                        i += 2;
-                    } else {
-                        new_word.push_back(word[i]);
-                        i += 1;
-                    }
-                }
-
-                word = new_word;
-
-                if (word.size() == 1) {
-                    break;
-                }
-                pairs = get_pairs(word);
-            }
-
-            std::u32string result;
-            for (int i = 0; i < word.size(); i++) {
-                result += word[i];
-                if (i != word.size() - 1) {
-                    result += utf8_to_utf32(" ");
-                }
-            }
-
-            return result;
-        }
-
-        std::vector<int> tokenize(std::string text,
-                                  on_new_token_cb_t on_new_token_cb = nullptr,
-                                  size_t max_length                 = 0,
-                                  bool padding                      = false) {
-            std::vector<int32_t> tokens = encode(text, on_new_token_cb);
-
-            if (max_length > 0) {
-                if (tokens.size() < max_length) {
-                    tokens.resize(max_length);
-                } else {
-                    if (padding) {
-                        tokens.insert(tokens.end(), max_length - tokens.size(), PAD_TOKEN_ID);
-                    }
-                }
-            }
-
-            return tokens;
-        }
-
-        void pad_tokens(std::vector<int>& tokens,
-                        std::vector<float>& weights,
-                        size_t max_length = 0,
-                        bool padding      = false) {
-            if (add_bos_token) {
-                tokens.insert(tokens.begin(), BOS_TOKEN_ID);
-                weights.insert(weights.begin(), 1.f);
-            }
-            if (max_length > 0 && padding) {
-                size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.f / max_length));
-                if (n == 0) {
-                    n = 1;
-                }
-                size_t length = max_length * n;
-                LOG_DEBUG("token length: %llu", length);
-                tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
-                weights.insert(weights.end(), length - weights.size(), 1.f);
-            }
-        }
-
-        std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb = nullptr) {
-            std::string original_text = text;
-            std::vector<int32_t> bpe_tokens;
-            std::vector<std::string> token_strs;
-
-            auto splited_texts = split_with_special_tokens(text, special_tokens);
-
-            for (auto& splited_text : splited_texts) {
-                if (is_special_token(splited_text)) {
-                    bpe_tokens.push_back(encoder[utf8_to_utf32(splited_text)]);
-                    token_strs.push_back(splited_text);
-                    continue;
-                }
-                auto tokens = token_split(splited_text);
-                for (auto& token : tokens) {
-                    if (on_new_token_cb != nullptr) {
-                        bool skip = on_new_token_cb(token, bpe_tokens);
-                        if (skip) {
-                            continue;
-                        }
-                    }
-
-                    std::string token_str = token;
-                    std::u32string utf32_token;
-                    for (int i = 0; i < token_str.length(); i++) {
-                        unsigned char b = token_str[i];
-                        utf32_token += byte_encoder[b];
-                    }
-                    auto bpe_strs = bpe(utf32_token);
-                    size_t start  = 0;
-                    size_t pos;
-                    while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {
-                        auto bpe_str = bpe_strs.substr(start, pos - start);
-                        bpe_tokens.push_back(encoder[bpe_str]);
-                        token_strs.push_back(utf32_to_utf8(bpe_str));
-
-                        start = pos + 1;
-                    }
-                    auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);
-                    bpe_tokens.push_back(encoder[bpe_str]);
-                    token_strs.push_back(utf32_to_utf8(bpe_str));
-                }
-            }
-
-            std::stringstream ss;
-            ss << "[";
-            for (auto token : token_strs) {
-                ss << "\"" << token << "\", ";
-            }
-            ss << "]";
-            LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
-            // printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str());
-            return bpe_tokens;
-        }
-    };
-
-    class Qwen2Tokenizer : public BPETokenizer {
-    protected:
-        void load_from_merges(const std::string& merges_utf8_str) {
-            auto byte_unicode_pairs = bytes_to_unicode();
-            // printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size());
-            byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-            for (auto& pair : byte_unicode_pairs) {
-                byte_decoder[pair.second] = pair.first;
-            }
-            // for (auto & pair: byte_unicode_pairs) {
-            //     std::cout << pair.first << ": " << pair.second << std::endl;
-            // }
-            std::vector<std::u32string> merges;
-            size_t start = 0;
-            size_t pos;
-            std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
-            while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
-                merges.push_back(merges_utf32_str.substr(start, pos - start));
-                start = pos + 1;
-            }
-            LOG_DEBUG("merges size %llu", merges.size());
-            merges = std::vector<std::u32string>(merges.begin(), merges.end());
-            std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-            // int print_num = 10;
-            for (const auto& merge : merges) {
-                size_t space_pos = merge.find(' ');
-                merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-                // if (print_num > 0) {
-                //     print_num--;
-                //     printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
-                //                     utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
-                // }
-            }
-
-            std::vector<std::u32string> tokens;
-            for (const auto& pair : byte_unicode_pairs) {
-                tokens.push_back(pair.second);
-            }
-            for (const auto& merge : merge_pairs) {
-                tokens.push_back(merge.first + merge.second);
-            }
-            for (auto& special_token : special_tokens) {
-                tokens.push_back(utf8_to_utf32(special_token));
-            }
-
-            int i = 0;
-            for (const auto& token : tokens) {
-                encoder[token] = i;
-                decoder[i]     = token;
-                i++;
-            }
-            encoder_len = i;
-            LOG_DEBUG("vocab size: %d", encoder_len);
-
-            int rank = 0;
-            for (const auto& merge : merge_pairs) {
-                bpe_ranks[merge] = rank++;
-            }
-            bpe_len = rank;
-        };
-
-    public:
-        explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "") {
-            UNK_TOKEN = "<|endoftext|>";
-            EOS_TOKEN = "<|endoftext|>";
-            PAD_TOKEN = "<|endoftext|>";
-
-            UNK_TOKEN_ID = 151643;
-            EOS_TOKEN_ID = 151643;
-            PAD_TOKEN_ID = 151643;
-
-            special_tokens = {
-                "<|endoftext|>",
-                "<|im_start|>",
-                "<|im_end|>",
-                "<|object_ref_start|>",
-                "<|object_ref_end|>",
-                "<|box_start|>",
-                "<|box_end|>",
-                "<|quad_start|>",
-                "<|quad_end|>",
-                "<|vision_start|>",
-                "<|vision_end|>",
-                "<|vision_pad|>",
-                "<|image_pad|>",
-                "<|video_pad|>",
-                "<tool_call>",
-                "</tool_call>",
-                "<|fim_prefix|>",
-                "<|fim_middle|>",
-                "<|fim_suffix|>",
-                "<|fim_pad|>",
-                "<|repo_name|>",
-                "<|file_sep|>",
-                "<tool_response>",
-                "</tool_response>",
-                "<think>",
-                "</think>",
-            };
-
-            if (merges_utf8_str.size() > 0) {
-                load_from_merges(merges_utf8_str);
-            } else {
-                load_from_merges(load_qwen2_merges());
-            }
-        }
-    };
-
-    class MistralTokenizer : public BPETokenizer {
-    protected:
-        void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-            nlohmann::json vocab;
-
-            try {
-                vocab = nlohmann::json::parse(vocab_utf8_str);
-            } catch (const nlohmann::json::parse_error&) {
-                GGML_ABORT("invalid vocab json str");
-            }
-            for (const auto& [key, value] : vocab.items()) {
-                std::u32string token = utf8_to_utf32(key);
-                int i                = value;
-                encoder[token]       = i;
-                decoder[i]           = token;
-            }
-            encoder_len = static_cast<int>(vocab.size());
-            LOG_DEBUG("vocab size: %d", encoder_len);
-
-            auto byte_unicode_pairs = bytes_to_unicode();
-            byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-            for (auto& pair : byte_unicode_pairs) {
-                byte_decoder[pair.second] = pair.first;
-            }
-            std::vector<std::u32string> merges;
-            size_t start = 0;
-            size_t pos;
-            std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
-            while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
-                merges.push_back(merges_utf32_str.substr(start, pos - start));
-                start = pos + 1;
-            }
-            LOG_DEBUG("merges size %llu", merges.size());
-            merges = std::vector<std::u32string>(merges.begin(), merges.end());
-            std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-            // int print_num = 10;
-            for (const auto& merge : merges) {
-                size_t space_pos = merge.find(' ');
-                merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-                // if (print_num > 0) {
-                //     print_num--;
-                //     printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
-                //                     utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
-                // }
-            }
-
-            int rank = 0;
-            for (const auto& merge : merge_pairs) {
-                bpe_ranks[merge] = rank++;
-            }
-            bpe_len = rank;
-        };
-
-    public:
-        explicit MistralTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "") {
-            add_bos_token = true;
-
-            UNK_TOKEN = "<unk>";
-            BOS_TOKEN = "<s>";
-            EOS_TOKEN = "</s>";
-            PAD_TOKEN = "<pad>";
-
-            UNK_TOKEN_ID = 0;
-            BOS_TOKEN_ID = 1;
-            EOS_TOKEN_ID = 2;
-            PAD_TOKEN_ID = 11;
-
-            special_tokens = {
-                "<unk>",
-                "<s>",
-                "</s>",
-                "[INST]",
-                "[/INST]",
-                "[AVAILABLE_TOOLS]",
-                "[/AVAILABLE_TOOLS]",
-                "[TOOL_RESULTS]",
-                "[/TOOL_RESULTS]",
-                "[TOOL_CALLS]",
-                "[IMG]",
-                "<pad>",
-                "[IMG_BREAK]",
-                "[IMG_END]",
-                "[PREFIX]",
-                "[MIDDLE]",
-                "[SUFFIX]",
-                "[SYSTEM_PROMPT]",
-                "[/SYSTEM_PROMPT]",
-                "[TOOL_CONTENT]",
-            };
-            for (int i = 20; i < 1000; i++) {
-                special_tokens.push_back("<SPECIAL_" + std::to_string(i) + ">");
-            }
-
-            if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
-                load_from_merges(merges_utf8_str, vocab_utf8_str);
-            } else {
-                load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
-            }
-        }
-    };
-
     enum class LLMArch {
         QWEN2_5_VL,
         QWEN3,
         MISTRAL_SMALL_3_2,
+        MINISTRAL_3_3B,
         ARCH_COUNT,
     };
 
@@ -484,6 +36,7 @@ namespace LLM {
         "qwen2.5vl",
         "qwen3",
         "mistral_small3.2",
+        "ministral3.3b",
     };
 
     struct LLMVisionParams {
@@ -868,6 +421,9 @@ namespace LLM {
             if (arch == LLMArch::MISTRAL_SMALL_3_2) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
+            } else if (arch == LLMArch::MINISTRAL_3_3B) {
+                q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
+                k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
             } else if (arch == LLMArch::QWEN3) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -1083,7 +639,7 @@ namespace LLM {
                   bool enable_vision_ = false)
             : GGMLRunner(backend, offload_params_to_cpu), enable_vision(enable_vision_) {
             params.arch = arch;
-            if (arch == LLMArch::MISTRAL_SMALL_3_2) {
+            if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                 params.head_dim     = 128;
                 params.num_heads    = 32;
                 params.num_kv_heads = 8;
@@ -1195,7 +751,7 @@ namespace LLM {
             }
 
             int64_t n_tokens = input_ids->ne[0];
-            if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::QWEN3) {
+            if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::MINISTRAL_3_3B || params.arch == LLMArch::QWEN3) {
                 input_pos_vec.resize(n_tokens);
                 for (int i = 0; i < n_tokens; ++i) {
                     input_pos_vec[i] = i;
@@ -1431,7 +987,7 @@ namespace LLM {
                     const std::string prefix                       = "",
                     bool enable_vision                             = false)
             : model(arch, backend, offload_params_to_cpu, tensor_storage_map, prefix, enable_vision) {
-            if (arch == LLMArch::MISTRAL_SMALL_3_2) {
+            if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                 tokenizer = std::make_shared<MistralTokenizer>();
             } else {
                 tokenizer = std::make_shared<Qwen2Tokenizer>();
@@ -1479,7 +1035,7 @@ namespace LLM {
                 weights.insert(weights.end(), curr_tokens.size(), curr_weight);
             }
 
-            tokenizer->pad_tokens(tokens, weights, max_length, padding);
+            tokenizer->pad_tokens(tokens, &weights, nullptr, padding ? max_length : 0, padding ? max_length : 100000000, padding);
 
             // for (int i = 0; i < tokens.size(); i++) {
             //     std::cout << tokens[i] << ":" << weights[i] << ", ";

src/model.cpp

@@ -315,8 +315,9 @@ bool is_safetensors_file(const std::string& file_path) {
     if (!file) {
         return false;
     }
-    nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
-    if (header_.is_discarded()) {
+    try {
+        nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
+    } catch (const std::exception&) {
         return false;
     }
     return true;
@@ -511,7 +512,14 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
         return false;
     }
 
-    nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
+    nlohmann::json header_;
+    try {
+        header_ = nlohmann::json::parse(header_buf.data());
+    } catch (const std::exception&) {
+        LOG_ERROR("parsing safetensors header failed", file_path.c_str());
+        file_paths_.pop_back();
+        return false;
+    }
 
     for (auto& item : header_.items()) {
         std::string name           = item.key();
@@ -575,24 +583,29 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
 
         size_t tensor_data_size = end - begin;
 
+        bool tensor_size_ok;
         if (dtype == "F8_E4M3") {
             tensor_storage.is_f8_e4m3 = true;
             // f8 -> f16
-            GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2);
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
         } else if (dtype == "F8_E5M2") {
             tensor_storage.is_f8_e5m2 = true;
             // f8 -> f16
-            GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2);
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
         } else if (dtype == "F64") {
             tensor_storage.is_f64 = true;
             // f64 -> f32
-            GGML_ASSERT(tensor_storage.nbytes() * 2 == tensor_data_size);
+            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
         } else if (dtype == "I64") {
             tensor_storage.is_i64 = true;
             // i64 -> i32
-            GGML_ASSERT(tensor_storage.nbytes() * 2 == tensor_data_size);
+            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
         } else {
-            GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size);
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+        }
+        if (!tensor_size_ok) {
+            LOG_ERROR("size mismatch for tensor '%s' (%s)\n", name.c_str(), dtype.c_str());
+            return false;
         }
 
         add_tensor_storage(tensor_storage);
@@ -1036,6 +1049,9 @@ SDVersion ModelLoader::get_sd_version() {
             if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) {
                 return VERSION_Z_IMAGE;
             }
+            if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
+                return VERSION_ERNIE_IMAGE;
+            }
             if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
                 is_wan = true;
             }

src/model.h

@@ -50,6 +50,7 @@ enum SDVersion {
     VERSION_FLUX2_KLEIN,
     VERSION_Z_IMAGE,
     VERSION_OVIS_IMAGE,
+    VERSION_ERNIE_IMAGE,
     VERSION_COUNT,
 };
 
@@ -137,6 +138,20 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_ernie_image(SDVersion version) {
+    if (version == VERSION_ERNIE_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_uses_flux2_vae(SDVersion version) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version)) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_inpaint(SDVersion version) {
     if (version == VERSION_SD1_INPAINT ||
         version == VERSION_SD2_INPAINT ||
@@ -155,7 +170,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_wan(version) ||
         sd_version_is_qwen_image(version) ||
         sd_version_is_anima(version) ||
-        sd_version_is_z_image(version)) {
+        sd_version_is_z_image(version) ||
+        sd_version_is_ernie_image(version)) {
         return true;
     }
     return false;

src/rope.hpp

@@ -7,6 +7,11 @@
 #include "ggml_extend.hpp"
 
 namespace Rope {
+    enum class EmbedNDLayout {
+        Matrix,
+        ErnieImage,
+    };
+
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
@@ -169,7 +174,8 @@ namespace Rope {
                                                   int bs,
                                                   const std::vector<float>& axis_thetas,
                                                   const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {}) {
+                                                  const std::vector<std::vector<int>>& wrap_dims = {},
+                                                  EmbedNDLayout layout                           = EmbedNDLayout::Matrix) {
         std::vector<std::vector<float>> trans_ids = transpose(ids);
         size_t pos_len                            = ids.size() / bs;
         size_t num_axes                           = axes_dim.size();
@@ -204,6 +210,24 @@ namespace Rope {
             offset += rope_emb[0].size();
         }
 
+        if (layout == EmbedNDLayout::ErnieImage) {
+            int head_dim = emb_dim * 2;
+            std::vector<float> ernie_emb(bs * pos_len * head_dim * 2, 0.0f);
+            for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
+                for (int i = 0; i < emb_dim; ++i) {
+                    float cos_val             = emb[pos_idx][4 * i];
+                    float sin_val             = emb[pos_idx][4 * i + 2];
+                    size_t cos_offset         = pos_idx * head_dim + 2 * i;
+                    size_t sin_offset         = bs * pos_len * head_dim + cos_offset;
+                    ernie_emb[cos_offset]     = cos_val;
+                    ernie_emb[cos_offset + 1] = cos_val;
+                    ernie_emb[sin_offset]     = sin_val;
+                    ernie_emb[sin_offset + 1] = sin_val;
+                }
+            }
+            return ernie_emb;
+        }
+
         return flatten(emb);
     }
 
@@ -211,9 +235,10 @@ namespace Rope {
                                                   int bs,
                                                   float theta,
                                                   const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {}) {
+                                                  const std::vector<std::vector<int>>& wrap_dims = {},
+                                                  EmbedNDLayout layout                           = EmbedNDLayout::Matrix) {
         std::vector<float> axis_thetas(axes_dim.size(), theta);
-        return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims);
+        return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@@ -437,6 +462,74 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
+                                                                          int w,
+                                                                          int patch_size,
+                                                                          int bs,
+                                                                          int context_len) {
+        int h_len = h / patch_size;
+        int w_len = w / patch_size;
+
+        std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
+        std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
+        std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
+        for (int i = 0; i < h_len; ++i) {
+            for (int j = 0; j < w_len; ++j) {
+                img_ids[i * w_len + j][0] = static_cast<float>(context_len);
+                img_ids[i * w_len + j][1] = h_ids[i];
+                img_ids[i * w_len + j][2] = w_ids[j];
+            }
+        }
+
+        std::vector<std::vector<float>> img_ids_repeated(bs * img_ids.size(), std::vector<float>(3, 0.0f));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < static_cast<int>(img_ids.size()); ++j) {
+                img_ids_repeated[i * img_ids.size() + j] = img_ids[j];
+            }
+        }
+
+        std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < context_len; ++j) {
+                txt_ids[i * context_len + j][0] = static_cast<float>(j);
+            }
+        }
+
+        return concat_ids(img_ids_repeated, txt_ids, bs);
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
+                                                            int w,
+                                                            int patch_size,
+                                                            int bs,
+                                                            int context_len,
+                                                            int theta,
+                                                            bool circular_h,
+                                                            bool circular_w,
+                                                            const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
+        std::vector<std::vector<int>> wrap_dims;
+        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
+            int h_len = h / patch_size;
+            int w_len = w / patch_size;
+            if (h_len > 0 && w_len > 0) {
+                size_t pos_len = ids.size() / bs;
+                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
+                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
+                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
+                    if (circular_h) {
+                        wrap_dims[1][token_i] = h_len;
+                    }
+                    if (circular_w) {
+                        wrap_dims[2][token_i] = w_len;
+                    }
+                }
+            }
+        }
+
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
                                                                   int h,
                                                                   int w,

src/stable-diffusion.cpp

@@ -52,6 +52,7 @@ const char* model_version_to_str[] = {
     "Flux.2 klein",
     "Z-Image",
     "Ovis Image",
+    "Ernie Image",
 };
 
 const char* sampling_methods_str[] = {
@@ -551,6 +552,15 @@ public:
                                                                 tensor_storage_map,
                                                                 "model.diffusion_model",
                                                                 version);
+            } else if (sd_version_is_ernie_image(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(clip_backend,
+                                                                 offload_params_to_cpu,
+                                                                 tensor_storage_map,
+                                                                 version);
+                diffusion_model  = std::make_shared<ErnieImageModel>(backend,
+                                                                    offload_params_to_cpu,
+                                                                    tensor_storage_map,
+                                                                    "model.diffusion_model");
             } else {  // SD1.x SD2.x SDXL
                 std::map<std::string, std::string> embbeding_map;
                 for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
@@ -819,6 +829,10 @@ public:
         if (version == VERSION_SVD) {
             ignore_tensors.insert("conditioner.embedders.3");
         }
+        if (sd_version_is_ernie_image(version)) {
+            ignore_tensors.insert("text_encoders.llm.vision_tower.");
+            ignore_tensors.insert("text_encoders.llm.multi_modal_projector.");
+        }
         bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap);
         if (!success) {
             LOG_ERROR("load tensors from model loader failed");
@@ -922,6 +936,7 @@ public:
                            sd_version_is_wan(version) ||
                            sd_version_is_qwen_image(version) ||
                            sd_version_is_anima(version) ||
+                           sd_version_is_ernie_image(version) ||
                            sd_version_is_z_image(version)) {
                     pred_type = FLOW_PRED;
                     if (sd_version_is_wan(version)) {
@@ -1395,7 +1410,7 @@ public:
             uint32_t dim                     = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
 
             if (dim == 128) {
-                if (sd_version_is_flux2(version)) {
+                if (sd_version_uses_flux2_vae(version)) {
                     latent_rgb_proj = flux2_latent_rgb_proj;
                     latent_rgb_bias = flux2_latent_rgb_bias;
                     patch_sz        = 2;
@@ -1593,6 +1608,7 @@ public:
                              float eta,
                              int shifted_timestep,
                              sample_method_t method,
+                             bool is_flow_denoiser,
                              const std::vector<float>& sigmas,
                              int start_merge_step,
                              const std::vector<sd::Tensor<float>>& ref_latents,
@@ -1791,7 +1807,7 @@ public:
             return denoised;
         };
 
-        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta);
+        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser);
         if (x0_opt.empty()) {
             LOG_ERROR("Diffusion model sampling failed");
             if (control_net) {
@@ -1843,7 +1859,7 @@ public:
                 latent_channel = 48;
             } else if (version == VERSION_CHROMA_RADIANCE) {
                 latent_channel = 3;
-            } else if (sd_version_is_flux2(version)) {
+            } else if (sd_version_uses_flux2_vae(version)) {
                 latent_channel = 128;
             } else {
                 latent_channel = 16;
@@ -1909,6 +1925,11 @@ public:
             flow_denoiser->set_shift(flow_shift);
         }
     }
+
+    bool is_flow_denoiser() {
+        auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
+        return !!flow_denoiser;
+    }
 };
 
 /*================================================= SD API ==================================================*/
@@ -2225,6 +2246,7 @@ void sd_sample_params_init(sd_sample_params_t* sample_params) {
     sample_params->scheduler                   = SCHEDULER_COUNT;
     sample_params->sample_method               = SAMPLE_METHOD_COUNT;
     sample_params->sample_steps                = 20;
+    sample_params->eta                         = INFINITY;
     sample_params->custom_sigmas               = nullptr;
     sample_params->custom_sigmas_count         = 0;
     sample_params->flow_shift                  = INFINITY;
@@ -2438,6 +2460,26 @@ static scheduler_t resolve_scheduler(sd_ctx_t* sd_ctx,
     return scheduler;
 }
 
+static float resolve_eta(sd_ctx_t* sd_ctx,
+                         float eta,
+                         enum sample_method_t sample_method) {
+    if (eta == INFINITY) {
+        switch (sample_method) {
+            case DDIM_TRAILING_SAMPLE_METHOD:
+            case TCD_SAMPLE_METHOD:
+            case RES_MULTISTEP_SAMPLE_METHOD:
+            case RES_2S_SAMPLE_METHOD:
+                return 0.0f;
+            case EULER_A_SAMPLE_METHOD:
+            case DPMPP2S_A_SAMPLE_METHOD:
+                return 1.0f;
+            default:;
+        }
+        return 0.0f;
+    }
+    return eta;
+}
+
 struct GenerationRequest {
     std::string prompt;
     std::string negative_prompt;
@@ -2586,6 +2628,8 @@ struct GenerationRequest {
 struct SamplePlan {
     enum sample_method_t sample_method            = SAMPLE_METHOD_COUNT;
     enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT;
+    float eta                                     = 0.f;
+    float high_noise_eta                          = 0.f;
     int sample_steps                              = 0;
     int high_noise_sample_steps                   = 0;
     int total_steps                               = 0;
@@ -2597,6 +2641,7 @@ struct SamplePlan {
                const sd_img_gen_params_t* sd_img_gen_params,
                const GenerationRequest& request) {
         sample_method = sd_img_gen_params->sample_params.sample_method;
+        eta           = sd_img_gen_params->sample_params.eta;
         sample_steps  = sd_img_gen_params->sample_params.sample_steps;
         resolve(sd_ctx, &request, &sd_img_gen_params->sample_params);
     }
@@ -2605,10 +2650,12 @@ struct SamplePlan {
                const sd_vid_gen_params_t* sd_vid_gen_params,
                const GenerationRequest& request) {
         sample_method = sd_vid_gen_params->sample_params.sample_method;
+        eta           = sd_vid_gen_params->sample_params.eta;
         sample_steps  = sd_vid_gen_params->sample_params.sample_steps;
         if (sd_ctx->sd->high_noise_diffusion_model) {
             high_noise_sample_steps  = sd_vid_gen_params->high_noise_sample_params.sample_steps;
             high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method;
+            high_noise_eta           = sd_vid_gen_params->high_noise_sample_params.eta;
         }
         moe_boundary = sd_vid_gen_params->moe_boundary;
         resolve(sd_ctx, &request, &sd_vid_gen_params->sample_params);
@@ -2644,6 +2691,8 @@ struct SamplePlan {
                                                                      sd_ctx->sd->version);
         }
 
+        eta = resolve_eta(sd_ctx, eta, sample_method);
+
         if (high_noise_sample_steps < 0) {
             for (size_t i = 0; i < sigmas.size(); ++i) {
                 if (sigmas[i] < moe_boundary) {
@@ -2658,6 +2707,7 @@ struct SamplePlan {
         if (high_noise_sample_steps > 0) {
             high_noise_sample_method = resolve_sample_method(sd_ctx,
                                                              high_noise_sample_method);
+            high_noise_eta           = resolve_eta(sd_ctx, high_noise_eta, high_noise_sample_method);
             LOG_INFO("sampling(high noise) using %s method", sampling_methods_str[high_noise_sample_method]);
         }
 
@@ -2811,7 +2861,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
                                                          {request->width / request->vae_scale_factor,
                                                           request->height / request->vae_scale_factor,
                                                           1,
-                                                          1});
+                                                          1},
+                                                         sd::ops::InterpolateMode::NearestMax);
 
     sd::Tensor<float> init_latent;
     sd::Tensor<float> control_latent;
@@ -2956,8 +3007,12 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
     latents.ref_latents          = std::move(ref_latents);
 
     if (sd_version_is_inpaint(sd_ctx->sd->version)) {
-        latents.denoise_mask = std::move(latent_mask);
+        latent_mask = sd::ops::max_pool_2d(latent_mask,
+                                           {3, 3},
+                                           {1, 1},
+                                           {1, 1});
     }
+    latents.denoise_mask = std::move(latent_mask);
 
     return latents;
 }
@@ -3123,9 +3178,10 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                    latents.control_image,
                                                    request.control_strength,
                                                    request.guidance,
-                                                   request.eta,
+                                                   plan.eta,
                                                    request.shifted_timestep,
                                                    plan.sample_method,
+                                                   sd_ctx->sd->is_flow_denoiser(),
                                                    plan.sigmas,
                                                    plan.start_merge_step,
                                                    latents.ref_latents,
@@ -3482,9 +3538,10 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                            sd::Tensor<float>(),
                                                            0.f,
                                                            request.high_noise_guidance,
-                                                           sd_vid_gen_params->high_noise_sample_params.eta,
+                                                           plan.high_noise_eta,
                                                            request.shifted_timestep,
                                                            plan.high_noise_sample_method,
+                                                           sd_ctx->sd->is_flow_denoiser(),
                                                            high_noise_sigmas,
                                                            -1,
                                                            std::vector<sd::Tensor<float>>{},
@@ -3523,9 +3580,10 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                         sd::Tensor<float>(),
                                                         0.f,
                                                         sd_vid_gen_params->sample_params.guidance,
-                                                        sd_vid_gen_params->sample_params.eta,
+                                                        plan.eta,
                                                         sd_vid_gen_params->sample_params.shifted_timestep,
                                                         plan.sample_method,
+                                                        sd_ctx->sd->is_flow_denoiser(),
                                                         plan.sigmas,
                                                         -1,
                                                         std::vector<sd::Tensor<float>>{},

src/t5.hpp

@@ -1,4 +1,4 @@
-#ifndef __T5_HPP__
+#ifndef __T5_HPP__
 #define __T5_HPP__
 
 #include <cfloat>
@@ -10,452 +10,9 @@
 #include <string>
 #include <unordered_map>
 
-#include "darts.h"
 #include "ggml_extend.hpp"
-#include "json.hpp"
 #include "model.h"
-#include "vocab/vocab.h"
-
-// Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
-// and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
-// Original License: https://github.com/google/sentencepiece/blob/master/LICENSE
-//
-// Since tokenization is not the bottleneck in SD, performance was not a major consideration
-// during the migration.
-class MetaspacePreTokenizer {
-private:
-    std::string replacement;
-    bool add_prefix_space;
-
-public:
-    MetaspacePreTokenizer(const std::string replacement = " ", bool add_prefix_space = true)
-        : replacement(replacement), add_prefix_space(add_prefix_space) {}
-
-    std::string tokenize(const std::string& input) const {
-        std::string tokens;
-        std::stringstream ss(input);
-
-        if (add_prefix_space) {
-            tokens += replacement;
-        }
-
-        std::string token;
-        bool firstToken = true;
-        while (std::getline(ss, token, ' ')) {
-            if (!firstToken)
-                tokens += replacement + token;
-            else
-                tokens += token;
-
-            firstToken = false;
-        }
-
-        return tokens;
-    }
-};
-
-using EncodeResult = std::vector<std::pair<std::string, int>>;
-class T5UniGramTokenizer {
-public:
-    enum Status {
-        OK,
-        NO_PIECES_LOADED,
-        NO_ENTRY_FOUND,
-        BUILD_DOUBLE_ARRAY_FAILED,
-        PIECE_ALREADY_DEFINED,
-        INVLIAD_JSON
-    };
-
-protected:
-    MetaspacePreTokenizer pre_tokenizer;
-
-    // all <piece, score> pairs
-    std::vector<std::pair<std::string, float>> piece_score_pairs;
-
-    float min_score_ = 0.0;
-    float max_score_ = 0.0;
-    std::unique_ptr<Darts::DoubleArray> trie_;
-
-    // Maximum size of the return value of Trie, which corresponds
-    // to the maximum size of shared common prefix in the sentence pieces.
-    int trie_results_size_;
-    // unknown id.
-    int unk_id_            = 2;
-    std::string eos_token_ = "</s>";
-    int eos_id_            = 1;
-    int pad_id_            = 0;
-    // status.
-    Status status_ = OK;
-
-    float kUnkPenalty = 10.0;
-
-    std::string replacement;
-    bool add_prefix_space = true;
-
-    void InitializePieces(const std::string& json_str) {
-        nlohmann::json data;
-
-        try {
-            data = nlohmann::json::parse(json_str);
-        } catch (const nlohmann::json::parse_error&) {
-            status_ = INVLIAD_JSON;
-            return;
-        }
-        if (!data.contains("model")) {
-            status_ = INVLIAD_JSON;
-            return;
-        }
-        nlohmann::json model = data["model"];
-        if (!model.contains("vocab")) {
-            status_ = INVLIAD_JSON;
-            return;
-        }
-        if (model.contains("unk_id")) {
-            unk_id_ = model["unk_id"];
-        }
-
-        replacement      = data["pre_tokenizer"]["replacement"];
-        add_prefix_space = data["pre_tokenizer"]["add_prefix_space"];
-
-        pre_tokenizer = MetaspacePreTokenizer(replacement, add_prefix_space);
-
-        for (const auto& item : model["vocab"]) {
-            if (item.size() != 2 || !item[0].is_string() || !item[1].is_number_float()) {
-                status_ = INVLIAD_JSON;
-                return;
-            }
-            std::string piece = item[0];
-            if (piece.empty()) {
-                piece = "<empty_token>";
-            }
-            float score = item[1];
-            piece_score_pairs.emplace_back(piece, score);
-        }
-    }
-
-    // Builds a Trie index.
-    void BuildTrie(std::vector<std::pair<std::string, int>>* pieces) {
-        if (status_ != OK)
-            return;
-
-        if (pieces->empty()) {
-            status_ = NO_PIECES_LOADED;
-            return;
-        }
-
-        // sort by sentencepiece since DoubleArray::build()
-        // only accepts sorted strings.
-        sort(pieces->begin(), pieces->end());
-
-        // Makes key/value set for DoubleArrayTrie.
-        std::vector<const char*> key(pieces->size());
-        std::vector<int> value(pieces->size());
-        for (size_t i = 0; i < pieces->size(); ++i) {
-            // LOG_DEBUG("%s %d", (*pieces)[i].first.c_str(), (*pieces)[i].second);
-            key[i]   = (*pieces)[i].first.data();  // sorted piece.
-            value[i] = (*pieces)[i].second;        // vocab_id
-        }
-
-        trie_ = std::unique_ptr<Darts::DoubleArray>(new Darts::DoubleArray());
-        if (trie_->build(key.size(), const_cast<char**>(&key[0]), nullptr,
-                         &value[0]) != 0) {
-            status_ = BUILD_DOUBLE_ARRAY_FAILED;
-            return;
-        }
-
-        // Computes the maximum number of shared prefixes in the trie.
-        const int kMaxTrieResultsSize = 1024;
-        std::vector<Darts::DoubleArray::result_pair_type> results(
-            kMaxTrieResultsSize);
-        trie_results_size_ = 0;
-        for (const auto& p : *pieces) {
-            const size_t num_nodes = trie_->commonPrefixSearch(
-                p.first.data(), results.data(), results.size(), p.first.size());
-            trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
-        }
-
-        if (trie_results_size_ == 0)
-            status_ = NO_ENTRY_FOUND;
-    }
-
-    // Non-virtual (inlined) implementation for faster execution.
-    inline float GetScoreInlined(int id) const {
-        return piece_score_pairs[id].second;
-    }
-
-    inline bool IsUnusedInlined(int id) const {
-        return false;  // TODO
-    }
-
-    inline bool IsUserDefinedInlined(int id) const {
-        return false;  // TODO
-    }
-
-    inline size_t OneCharLen(const char* src) const {
-        return "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"[(*src & 0xFF) >> 4];
-    }
-
-    // The optimized Viterbi encode.
-    // Main differences from the original function:
-    // 1. Memorizes the best path at each postion so far,
-    // 2. No need to store the Lattice nodes,
-    // 3. Works in utf-8 directly,
-    // 4. Defines a new struct with fewer fields than Lattice,
-    // 5. Does not depend on `class Lattice` nor call `SetSentence()`,
-    // `PopulateNodes()`, or `Viterbi()`. It does everything in one function.
-    // For detailed explanations please see the comments inside the function body.
-    EncodeResult EncodeOptimized(const std::string& normalized) const {
-        // An optimized Viterbi algorithm for unigram language models. Benchmarking
-        // results show that it generates almost identical outputs and achieves 2.1x
-        // speedup on average for 102 languages compared to the original
-        // implementation. It's based on the following three ideas:
-        //
-        // 1. Because it uses the *unigram* model:
-        //     best_score(x1, x2, ... xt) = best_score(x1, x2, ... x{t-1}) + score(xt)
-        // Deciding the best path (and score) can be decoupled into two isolated
-        // terms: (a) the best path ended before the last token `best_score(x1, x2, ...)`
-        // x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
-        // not related to each other at all.
-        //
-        // Therefore, we can compute once and store the *best_path ending at
-        // each character position*. In this way, when we know best_path_ends_at[M],
-        // we can reuse it to compute all the best_path_ends_at_[...] where the last
-        // token starts at the same character position M.
-        //
-        // This improves the time complexity from O(n*k*k) to O(n*k) because it
-        // eliminates the extra loop of recomputing the best path ending at the same
-        // position, where n is the input length and k is the maximum number of tokens
-        // that can be recognized starting at each position.
-        //
-        // 2. Again, because it uses the *unigram* model, we don't need to actually
-        // store the lattice nodes. We still recognize all the tokens and lattice
-        // nodes from the input, but along identifying them, we use and discard them
-        // on the fly. There is no need to actually store them for best path Viterbi
-        // decoding. The only thing we need to store is the best_path ending at
-        // each character position.
-        //
-        // This improvement reduces the things needed to store in memory from O(n*k)
-        // to O(n), where n is the input length and k is the maximum number of tokens
-        // that can be recognized starting at each position.
-        //
-        // It also avoids the need of dynamic-size lattice node pool, because the
-        // number of things to store is fixed as n.
-        //
-        // 3. SentencePiece is designed to work with unicode, taking utf-8 encoding
-        // inputs. In the original implementation, the lattice positions are based on
-        // unicode positions. A mapping from unicode position to the utf-8 position is
-        // maintained to recover the utf-8 string piece.
-        //
-        // We found that it is sufficient and beneficial to directly work with utf-8
-        // positions:
-        //
-        // Firstly, it saves the conversion and mapping between unicode positions and
-        // utf-8 positions.
-        //
-        // Secondly, it reduces the number of fields we need to maintain in the
-        // node/path structure. Specifically, there are 8 fields defined in
-        // `Lattice::Node` used by the original encoder, but here in the optimized
-        // encoder we only need to define 3 fields in `BestPathNode`.
-
-        if (status() != OK || normalized.empty()) {
-            return {};
-        }
-        // Represents the last node of the best path.
-        struct BestPathNode {
-            int id = -1;  // The vocab id. (maybe -1 for UNK)
-            float best_path_score =
-                0;  // The total score of the best path ending at this node.
-            int starts_at =
-                -1;  // The starting position (in utf-8) of this node. The entire best
-                     // path can be constructed by backtracking along this link.
-        };
-        const int size        = static_cast<int>(normalized.size());
-        const float unk_score = min_score() - kUnkPenalty;
-        // The ends are exclusive.
-        std::vector<BestPathNode> best_path_ends_at(size + 1);
-        // Generate lattice on-the-fly (not stored) and update best_path_ends_at.
-        int starts_at = 0;
-        while (starts_at < size) {
-            std::size_t node_pos = 0;
-            std::size_t key_pos  = starts_at;
-            const auto best_path_score_till_here =
-                best_path_ends_at[starts_at].best_path_score;
-            bool has_single_node = false;
-            const int mblen =
-                std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)),
-                              size - starts_at);
-            while (key_pos < size) {
-                const int ret =
-                    trie_->traverse(normalized.data(), node_pos, key_pos, key_pos + 1);
-                if (ret == -2)
-                    break;
-                if (ret >= 0) {
-                    if (IsUnusedInlined(ret))
-                        continue;
-                    // Update the best path node.
-                    auto& target_node = best_path_ends_at[key_pos];
-                    const auto length = (key_pos - starts_at);
-                    // User defined symbol receives extra bonus to always be selected.
-                    const auto score = IsUserDefinedInlined(ret)
-                                           ? (length * max_score_ - 0.1)
-                                           : GetScoreInlined(ret);
-                    const auto candidate_best_path_score =
-                        score + best_path_score_till_here;
-                    if (target_node.starts_at == -1 ||
-                        candidate_best_path_score > target_node.best_path_score) {
-                        target_node.best_path_score = static_cast<float>(candidate_best_path_score);
-                        target_node.starts_at       = starts_at;
-                        target_node.id              = ret;
-                    }
-                    if (!has_single_node && length == mblen) {
-                        has_single_node = true;
-                    }
-                }
-            }
-            if (!has_single_node) {
-                auto& target_node = best_path_ends_at[starts_at + mblen];
-                const auto candidate_best_path_score =
-                    unk_score + best_path_score_till_here;
-                if (target_node.starts_at == -1 ||
-                    candidate_best_path_score > target_node.best_path_score) {
-                    target_node.best_path_score = candidate_best_path_score;
-                    target_node.starts_at       = starts_at;
-                    target_node.id              = unk_id_;
-                }
-            }
-            // Move by one unicode character.
-            starts_at += mblen;
-        }
-        // Backtrack to identify the best path.
-        EncodeResult results;
-        int ends_at = size;
-        while (ends_at > 0) {
-            const auto& node = best_path_ends_at[ends_at];
-            results.emplace_back(
-                normalized.substr(node.starts_at, ends_at - node.starts_at), node.id);
-            ends_at = node.starts_at;
-        }
-        std::reverse(results.begin(), results.end());
-        return results;
-    }
-
-public:
-    explicit T5UniGramTokenizer(bool is_umt5 = false) {
-        if (is_umt5) {
-            InitializePieces(load_umt5_tokenizer_json());
-        } else {
-            InitializePieces(load_t5_tokenizer_json());
-        }
-
-        min_score_ = FLT_MAX;
-        max_score_ = FLT_MIN;
-
-        std::vector<std::pair<std::string, int>> pieces;
-        for (int i = 0; i < piece_score_pairs.size(); i++) {
-            const auto& sp = piece_score_pairs[i];
-
-            min_score_ = std::min(min_score_, sp.second);
-            max_score_ = std::max(max_score_, sp.second);
-
-            pieces.emplace_back(sp.first, i);
-        }
-
-        BuildTrie(&pieces);
-    }
-    ~T5UniGramTokenizer(){};
-
-    std::string Normalize(const std::string& input) const {
-        // Ref: https://github.com/huggingface/tokenizers/blob/1ff56c0c70b045f0cd82da1af9ac08cd4c7a6f9f/bindings/python/py_src/tokenizers/implementations/sentencepiece_unigram.py#L29
-        // TODO: nmt-nfkc
-        std::string normalized = std::regex_replace(input, std::regex(" {2,}"), " ");
-        return normalized;
-    }
-
-    std::vector<int> Encode(const std::string& input, bool append_eos_if_not_present = true) const {
-        std::string normalized = Normalize(input);
-        normalized             = pre_tokenizer.tokenize(normalized);
-        EncodeResult result    = EncodeOptimized(normalized);
-        if (result.size() > 0 && append_eos_if_not_present) {
-            auto item = result[result.size() - 1];
-            if (item.first != eos_token_) {
-                result.emplace_back(eos_token_, eos_id_);
-            }
-        }
-        std::vector<int> tokens;
-        for (auto item : result) {
-            tokens.push_back(item.second);
-        }
-        return tokens;
-    }
-
-    void pad_tokens(std::vector<int>& tokens,
-                    std::vector<float>& weights,
-                    std::vector<float>* attention_mask,
-                    size_t max_length = 0,
-                    bool padding      = false) {
-        if (max_length > 0 && padding) {
-            size_t orig_token_num = tokens.size() - 1;
-            size_t n              = static_cast<size_t>(std::ceil(orig_token_num * 1.0 / (max_length - 1)));
-            if (n == 0) {
-                n = 1;
-            }
-            size_t length = max_length * n;
-            LOG_DEBUG("token length: %llu", length);
-            std::vector<int> new_tokens;
-            std::vector<float> new_weights;
-            std::vector<float> new_attention_mask;
-            int token_idx = 0;
-            for (int i = 0; i < length; i++) {
-                if (token_idx >= orig_token_num) {
-                    break;
-                }
-                if (attention_mask != nullptr) {
-                    new_attention_mask.push_back(0.0);
-                }
-                if (i % max_length == max_length - 1) {
-                    new_tokens.push_back(eos_id_);
-                    new_weights.push_back(1.0);
-                } else {
-                    new_tokens.push_back(tokens[token_idx]);
-                    new_weights.push_back(weights[token_idx]);
-                    token_idx++;
-                }
-            }
-
-            new_tokens.push_back(eos_id_);
-            new_weights.push_back(1.0);
-            if (attention_mask != nullptr) {
-                new_attention_mask.push_back(0.0);
-            }
-
-            tokens  = new_tokens;
-            weights = new_weights;
-            if (attention_mask != nullptr) {
-                *attention_mask = new_attention_mask;
-            }
-
-            if (padding) {
-                int pad_token_id = pad_id_;
-                tokens.insert(tokens.end(), length - tokens.size(), pad_token_id);
-                weights.insert(weights.end(), length - weights.size(), 1.0);
-                if (attention_mask != nullptr) {
-                    // maybe keep some padding tokens unmasked?
-                    attention_mask->insert(attention_mask->end(), length - attention_mask->size(), -HUGE_VALF);
-                }
-            }
-        }
-    }
-
-    // Returns the minimum score in sentence pieces.
-    // min_score() - 10 is used for the cost of unknown sentence.
-    float min_score() const { return min_score_; }
-
-    // Returns the maximum score in sentence pieces.
-    // max_score() is used for the cost of user defined symbols.
-    float max_score() const { return max_score_; }
-
-    Status status() const { return status_; }
-};
+#include "tokenizers/t5_unigram_tokenizer.h"
 
 class T5LayerNorm : public UnaryBlock {
 protected:
@@ -937,18 +494,17 @@ struct T5Embedder {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = tokenizer.Encode(curr_text, false);
+            std::vector<int> curr_tokens = tokenizer.encode(curr_text);
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
-        int EOS_TOKEN_ID = 1;
-        tokens.push_back(EOS_TOKEN_ID);
-        weights.push_back(1.0);
-
         std::vector<float> attention_mask;
 
-        tokenizer.pad_tokens(tokens, weights, &attention_mask, max_length, padding);
+        tokenizer.pad_tokens(tokens, &weights, &attention_mask, padding ? max_length : 0, padding ? max_length : 100000000, padding);
+        for (auto& mask_value : attention_mask) {
+            mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
+        }
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", ";

src/tensor.hpp

@@ -815,6 +815,9 @@ namespace sd {
     namespace ops {
         enum class InterpolateMode {
             Nearest,
+            NearestMax,
+            NearestMin,
+            NearestAvg,
         };
 
         inline int64_t normalize_slice_bound(int64_t index, int64_t dim_size) {
@@ -1012,12 +1015,16 @@ namespace sd {
                                      std::vector<int64_t> output_shape,
                                      InterpolateMode mode = InterpolateMode::Nearest,
                                      bool align_corners   = false) {
-            if (mode != InterpolateMode::Nearest) {
-                tensor_throw_invalid_argument("Only nearest interpolate mode is implemented, got mode=" +
+            const bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
+                                               mode == InterpolateMode::NearestMax ||
+                                               mode == InterpolateMode::NearestMin ||
+                                               mode == InterpolateMode::NearestAvg);
+            if (!is_nearest_like_mode) {
+                tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
                                               std::to_string(static_cast<int>(mode)));
             }
             if (align_corners) {
-                tensor_throw_invalid_argument("align_corners is not supported for nearest interpolate: input_shape=" +
+                tensor_throw_invalid_argument("align_corners is not supported for nearest-like interpolate: input_shape=" +
                                               tensor_shape_to_string(input.shape()) + ", output_shape=" +
                                               tensor_shape_to_string(output_shape));
             }
@@ -1044,14 +1051,102 @@ namespace sd {
                 }
             }
 
-            Tensor<T> output(std::move(output_shape));
-            for (int64_t flat = 0; flat < output.numel(); ++flat) {
-                std::vector<int64_t> output_coord = tensor_unravel_index(flat, output.shape());
-                std::vector<int64_t> input_coord(static_cast<size_t>(input.dim()), 0);
-                for (size_t i = 0; i < static_cast<size_t>(input.dim()); ++i) {
-                    input_coord[i] = output_coord[i] * input.shape()[i] / output.shape()[i];
+            bool has_downsampling = false;
+            for (int64_t i = 0; i < input.dim(); ++i) {
+                if (input.shape()[i] > output_shape[i]) {
+                    has_downsampling = true;
+                    break;
                 }
-                output[flat] = input.index(input_coord);
+            }
+
+            Tensor<T> output(std::move(output_shape));
+            if (mode == InterpolateMode::Nearest || !has_downsampling) {
+                for (int64_t flat = 0; flat < output.numel(); ++flat) {
+                    std::vector<int64_t> output_coord = tensor_unravel_index(flat, output.shape());
+                    std::vector<int64_t> input_coord(static_cast<size_t>(input.dim()), 0);
+                    for (size_t i = 0; i < static_cast<size_t>(input.dim()); ++i) {
+                        input_coord[i] = output_coord[i] * input.shape()[i] / output.shape()[i];
+                    }
+                    output[flat] = input.index(input_coord);
+                }
+
+                return output;
+            }
+
+            auto init_reduction = [&]() -> T {
+                switch (mode) {
+                    case InterpolateMode::NearestMax:
+                        return std::numeric_limits<T>::lowest();
+                    case InterpolateMode::NearestMin:
+                        return std::numeric_limits<T>::max();
+                    case InterpolateMode::NearestAvg:
+                        return T(0);
+                    case InterpolateMode::Nearest:
+                        return T(0);
+                }
+
+                tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
+                                              std::to_string(static_cast<int>(mode)));
+            };
+
+            auto reduce_value = [&](T& acc, const T& sample) {
+                switch (mode) {
+                    case InterpolateMode::NearestMax:
+                        acc = std::max(acc, sample);
+                        break;
+                    case InterpolateMode::NearestMin:
+                        acc = std::min(acc, sample);
+                        break;
+                    case InterpolateMode::NearestAvg:
+                        acc += sample;
+                        break;
+                    case InterpolateMode::Nearest:
+                        break;
+                }
+            };
+
+            // Reduction modes only differ from nearest mode when downsampling.
+            for (int64_t flat_out = 0; flat_out < output.numel(); ++flat_out) {
+                std::vector<int64_t> output_coord = tensor_unravel_index(flat_out, output.shape());
+
+                std::vector<int64_t> input_start(output.dim(), 0);
+                std::vector<int64_t> input_end(output.dim(), 0);
+
+                for (size_t i = 0; i < static_cast<size_t>(output.dim()); ++i) {
+                    const int64_t input_dim  = input.shape()[i];
+                    const int64_t output_dim = output.shape()[i];
+
+                    input_start[i] = std::max(int64_t(0), static_cast<int64_t>(output_coord[i] * input_dim / output_dim));
+                    input_end[i]   = std::min(input_dim, ((output_coord[i] + 1) * input_dim + output_dim - 1) / output_dim);
+                }
+
+                T value                               = init_reduction();
+                bool done_window                      = false;
+                std::vector<int64_t> current_in_coord = input_start;
+
+                while (!done_window) {
+                    reduce_value(value, input.index(current_in_coord));
+
+                    for (int d = static_cast<int>(output.dim()) - 1; d >= 0; --d) {
+                        if (++current_in_coord[d] < input_end[d]) {
+                            break;
+                        }
+                        current_in_coord[d] = input_start[d];
+                        if (d == 0) {
+                            done_window = true;
+                        }
+                    }
+                }
+
+                if (mode == InterpolateMode::NearestAvg) {
+                    int64_t window_size = 1;
+                    for (size_t i = 0; i < static_cast<size_t>(output.dim()); ++i) {
+                        window_size *= (input_end[i] - input_start[i]);
+                    }
+                    value /= static_cast<T>(window_size);
+                }
+
+                output[flat_out] = value;
             }
 
             return output;
@@ -1063,12 +1158,16 @@ namespace sd {
                                      const std::optional<std::vector<double>>& scale_factor,
                                      InterpolateMode mode = InterpolateMode::Nearest,
                                      bool align_corners   = false) {
-            if (mode != InterpolateMode::Nearest) {
-                tensor_throw_invalid_argument("Only nearest interpolate mode is implemented, got mode=" +
+            const bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
+                                               mode == InterpolateMode::NearestMax ||
+                                               mode == InterpolateMode::NearestMin ||
+                                               mode == InterpolateMode::NearestAvg);
+            if (!is_nearest_like_mode) {
+                tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
                                               std::to_string(static_cast<int>(mode)));
             }
             if (align_corners) {
-                tensor_throw_invalid_argument("align_corners is not supported for nearest interpolate: input_shape=" +
+                tensor_throw_invalid_argument("align_corners is not supported for nearest-like interpolate: input_shape=" +
                                               tensor_shape_to_string(input.shape()));
             }
             if (size.has_value() == scale_factor.has_value()) {
@@ -1128,6 +1227,80 @@ namespace sd {
                                align_corners);
         }
 
+        template <typename T>
+        inline Tensor<T> max_pool_2d(const Tensor<T>& input,
+                                     std::vector<int64_t> kernel_size,
+                                     std::vector<int64_t> stride,
+                                     std::vector<int64_t> padding) {
+            if (input.dim() < 2) {
+                tensor_throw_invalid_argument("Tensor max_pool_2d requires input_dim >= 2: input_dim=" +
+                                              std::to_string(input.dim()) + ", input_shape=" +
+                                              tensor_shape_to_string(input.shape()));
+            }
+            if (kernel_size.size() != 2 || stride.size() != 2 || padding.size() != 2) {
+                tensor_throw_invalid_argument("Tensor max_pool_2d requires kernel_size, stride, and padding to have length 2");
+            }
+            for (size_t i = 0; i < 2; ++i) {
+                if (kernel_size[i] <= 0) {
+                    tensor_throw_invalid_argument("Tensor max_pool_2d kernel_size must be positive: kernel_size=" +
+                                                  tensor_shape_to_string(kernel_size));
+                }
+                if (stride[i] <= 0) {
+                    tensor_throw_invalid_argument("Tensor max_pool_2d stride must be positive: stride=" +
+                                                  tensor_shape_to_string(stride));
+                }
+                if (padding[i] < 0) {
+                    tensor_throw_invalid_argument("Tensor max_pool_2d padding must be non-negative: padding=" +
+                                                  tensor_shape_to_string(padding));
+                }
+            }
+
+            const int64_t in_height = input.shape()[0];
+            const int64_t in_width  = input.shape()[1];
+
+            const int64_t out_height = (in_height + 2 * padding[0] - kernel_size[0]) / stride[0] + 1;
+            const int64_t out_width  = (in_width + 2 * padding[1] - kernel_size[1]) / stride[1] + 1;
+
+            if (out_height <= 0 || out_width <= 0) {
+                tensor_throw_invalid_argument("max_pool_2d results in invalid output dimensions: " +
+                                              std::to_string(out_height) + "x" + std::to_string(out_width));
+            }
+
+            std::vector<int64_t> output_shape = input.shape();
+            output_shape[0]                   = out_height;
+            output_shape[1]                   = out_width;
+
+            Tensor<T> output(std::move(output_shape));
+
+            for (int64_t flat_out = 0; flat_out < output.numel(); ++flat_out) {
+                std::vector<int64_t> output_coord = tensor_unravel_index(flat_out, output.shape());
+                std::vector<int64_t> input_coord  = output_coord;
+
+                const int64_t oh = output_coord[0];
+                const int64_t ow = output_coord[1];
+
+                T max_val            = std::numeric_limits<T>::lowest();
+                bool has_valid_input = false;
+
+                for (int64_t kh = 0; kh < kernel_size[0]; ++kh) {
+                    for (int64_t kw = 0; kw < kernel_size[1]; ++kw) {
+                        const int64_t ih = oh * stride[0] + kh - padding[0];
+                        const int64_t iw = ow * stride[1] + kw - padding[1];
+
+                        if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) {
+                            input_coord[0]  = ih;
+                            input_coord[1]  = iw;
+                            max_val         = std::max(max_val, input.index(input_coord));
+                            has_valid_input = true;
+                        }
+                    }
+                }
+
+                output[flat_out] = has_valid_input ? max_val : T(0);
+            }
+            return output;
+        }
+
         template <typename T>
         inline Tensor<T> concat(const Tensor<T>& lhs, const Tensor<T>& rhs, size_t dim) {
             if (lhs.dim() != rhs.dim()) {

src/tokenizers/bpe_tokenizer.cpp

@@ -0,0 +1,189 @@
+#include "bpe_tokenizer.h"
+
+#include <algorithm>
+#include <sstream>
+
+#include "tokenize_util.h"
+#include "util.h"
+
+std::vector<std::pair<int, std::u32string>> BPETokenizer::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) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    for (int b = 161; b <= 172; ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    for (int b = 174; b <= 255; ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    int n = 0;
+    for (int b = 0; b < 256; ++b) {
+        if (byte_set.find(b) == byte_set.end()) {
+            byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
+            ++n;
+        }
+    }
+    return byte_unicode_pairs;
+}
+
+std::vector<std::string> BPETokenizer::token_split(const std::string& text) const {
+    return ::token_split(text);
+}
+
+std::vector<std::u32string> BPETokenizer::split_utf32(const std::string& text, char32_t delimiter) {
+    std::vector<std::u32string> result;
+    size_t start              = 0;
+    size_t pos                = 0;
+    std::u32string utf32_text = utf8_to_utf32(text);
+    while ((pos = utf32_text.find(delimiter, start)) != std::u32string::npos) {
+        result.push_back(utf32_text.substr(start, pos - start));
+        start = pos + 1;
+    }
+    return result;
+}
+
+static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
+    std::set<std::pair<std::u32string, std::u32string>> pairs;
+    if (subwords.empty()) {
+        return pairs;
+    }
+
+    std::u32string prev_subword = subwords[0];
+    for (int i = 1; i < static_cast<int>(subwords.size()); i++) {
+        std::u32string subword = subwords[i];
+        std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
+        pairs.insert(pair);
+        prev_subword = subword;
+    }
+    return pairs;
+}
+
+std::vector<std::u32string> BPETokenizer::bpe(const std::u32string& token) const {
+    std::vector<std::u32string> word;
+
+    for (int i = 0; i < static_cast<int>(token.size()) - 1; i++) {
+        word.emplace_back(1, token[i]);
+    }
+    word.push_back(token.substr(token.size() - 1) + utf8_to_utf32(end_of_word_suffix));
+
+    std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
+
+    if (pairs.empty()) {
+        return {token + utf8_to_utf32(end_of_word_suffix)};
+    }
+
+    while (true) {
+        auto min_pair_iter = std::min_element(pairs.begin(),
+                                              pairs.end(),
+                                              [&](const std::pair<std::u32string, std::u32string>& a,
+                                                  const std::pair<std::u32string, std::u32string>& b) {
+                                                  if (bpe_ranks.find(a) == bpe_ranks.end()) {
+                                                      return false;
+                                                  } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
+                                                      return true;
+                                                  }
+                                                  return bpe_ranks.at(a) < bpe_ranks.at(b);
+                                              });
+
+        const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
+
+        if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
+            break;
+        }
+
+        std::u32string first  = bigram.first;
+        std::u32string second = bigram.second;
+        std::vector<std::u32string> new_word;
+        int32_t i = 0;
+
+        while (i < static_cast<int32_t>(word.size())) {
+            auto it = std::find(word.begin() + i, word.end(), first);
+            if (it == word.end()) {
+                new_word.insert(new_word.end(), word.begin() + i, word.end());
+                break;
+            }
+            new_word.insert(new_word.end(), word.begin() + i, it);
+            i = static_cast<int32_t>(std::distance(word.begin(), it));
+
+            if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
+                new_word.push_back(first + second);
+                i += 2;
+            } else {
+                new_word.push_back(word[i]);
+                i += 1;
+            }
+        }
+
+        word = new_word;
+
+        if (word.size() == 1) {
+            break;
+        }
+        pairs = get_pairs(word);
+    }
+
+    return word;
+}
+
+std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t on_new_token_cb) {
+    std::string normalized_text = normalize(text);
+    std::vector<int32_t> bpe_tokens;
+    std::vector<std::string> token_strs;
+
+    auto splited_texts = split_with_special_tokens(normalized_text, special_tokens);
+
+    for (auto& splited_text : splited_texts) {
+        if (is_special_token(splited_text)) {
+            if (on_new_token_cb != nullptr) {
+                bool skip = on_new_token_cb(splited_text, bpe_tokens);
+                if (skip) {
+                    token_strs.push_back(splited_text);
+                    continue;
+                }
+            }
+            bpe_tokens.push_back(encoder[utf8_to_utf32(splited_text)]);
+            token_strs.push_back(splited_text);
+            continue;
+        }
+        auto tokens = token_split(splited_text);
+        for (auto& token : tokens) {
+            if (on_new_token_cb != nullptr) {
+                bool skip = on_new_token_cb(token, bpe_tokens);
+                if (skip) {
+                    token_strs.push_back(splited_text);
+                    continue;
+                }
+            }
+
+            std::string token_str = token;
+            std::u32string utf32_token;
+            for (int i = 0; i < static_cast<int>(token_str.length()); i++) {
+                unsigned char b = token_str[i];
+                utf32_token += byte_encoder[b];
+            }
+            auto bpe_strs = bpe(utf32_token);
+            for (auto bpe_str : bpe_strs) {
+                bpe_tokens.push_back(encoder[bpe_str]);
+                token_strs.push_back(utf32_to_utf8(bpe_str));
+            }
+        }
+    }
+
+    std::stringstream ss;
+    ss << "[";
+    for (auto token : token_strs) {
+        ss << "\"" << token << "\", ";
+    }
+    ss << "]";
+    LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
+    return bpe_tokens;
+}
+
+std::string BPETokenizer::decode_token(int token_id) const {
+    return utf32_to_utf8(decoder.at(token_id));
+}

src/tokenizers/bpe_tokenizer.h

@@ -0,0 +1,40 @@
+#ifndef __SD_TOKENIZERS_BPE_TOKENIZER_H__
+#define __SD_TOKENIZERS_BPE_TOKENIZER_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <regex>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tokenizer.h"
+
+class BPETokenizer : public Tokenizer {
+protected:
+    std::map<int, std::u32string> byte_encoder;
+    std::map<std::u32string, int> byte_decoder;
+    std::map<std::u32string, int> encoder;
+    std::map<int, std::u32string> decoder;
+    std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
+    int encoder_len = 0;
+    int bpe_len     = 0;
+
+protected:
+    static std::vector<std::pair<int, std::u32string>> bytes_to_unicode();
+    static std::vector<std::u32string> split_utf32(const std::string& text, char32_t delimiter = U'\n');
+    virtual std::vector<std::string> token_split(const std::string& text) const;
+    std::vector<std::u32string> bpe(const std::u32string& token) const;
+    std::string decode_token(int token_id) const override;
+
+public:
+    BPETokenizer()          = default;
+    virtual ~BPETokenizer() = default;
+
+    std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) override;
+};
+
+#endif  // __SD_TOKENIZERS_BPE_TOKENIZER_H__

src/tokenizers/clip_tokenizer.cpp

@@ -0,0 +1,116 @@
+#include "clip_tokenizer.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <regex>
+#include <set>
+
+#include "ggml.h"
+#include "tokenize_util.h"
+#include "util.h"
+#include "vocab/vocab.h"
+
+CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str) {
+    UNK_TOKEN = "<|endoftext|>";
+    BOS_TOKEN = "<|startoftext|>";
+    EOS_TOKEN = "<|endoftext|>";
+    PAD_TOKEN = "<|endoftext|>";
+
+    UNK_TOKEN_ID = 49407;
+    BOS_TOKEN_ID = 49406;
+    EOS_TOKEN_ID = 49407;
+    PAD_TOKEN_ID = pad_token_id;
+
+    end_of_word_suffix = "</w>";
+    add_bos_token      = true;
+    add_eos_token      = true;
+
+    if (merges_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str);
+    } else {
+        load_from_merges(load_clip_merges());
+    }
+    add_special_token("<|startoftext|>");
+    add_special_token("<|endoftext|>");
+}
+
+void CLIPTokenizer::load_from_merges(const std::string& merges_utf8_str) {
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    GGML_ASSERT(merges.size() == 48895);
+    merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+    std::vector<std::u32string> vocab;
+    for (const auto& pair : byte_unicode_pairs) {
+        vocab.push_back(pair.second);
+    }
+    for (const auto& pair : byte_unicode_pairs) {
+        vocab.push_back(pair.second + utf8_to_utf32("</w>"));
+    }
+    for (const auto& merge : merge_pairs) {
+        vocab.push_back(merge.first + merge.second);
+    }
+    vocab.push_back(utf8_to_utf32("<|startoftext|>"));
+    vocab.push_back(utf8_to_utf32("<|endoftext|>"));
+    LOG_DEBUG("vocab size: %llu", vocab.size());
+    int i = 0;
+    for (const auto& token : vocab) {
+        encoder[token] = i;
+        decoder[i]     = token;
+        i++;
+    }
+    encoder_len = i;
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+static std::string strip(const std::string& str) {
+    std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
+    std::string::size_type end   = str.find_last_not_of(" \t\n\r\v\f");
+
+    if (start == std::string::npos) {
+        return "";
+    }
+
+    return str.substr(start, end - start + 1);
+}
+
+static std::string whitespace_clean(const std::string& text) {
+    auto result = std::regex_replace(text, std::regex(R"(\s+)"), " ");
+    result      = strip(result);
+    return result;
+}
+
+std::string CLIPTokenizer::normalize(const std::string& text) const {
+    auto normalized_text = whitespace_clean(text);
+    std::transform(normalized_text.begin(), normalized_text.end(), normalized_text.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    return normalized_text;
+}
+
+std::vector<std::string> CLIPTokenizer::token_split(const std::string& text) const {
+    std::regex clip_pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
+                        std::regex::icase);
+    std::sregex_iterator iter(text.begin(), text.end(), clip_pat);
+    std::sregex_iterator end;
+
+    std::vector<std::string> result;
+    for (; iter != end; ++iter) {
+        result.emplace_back(iter->str());
+    }
+
+    return result;
+}

src/tokenizers/clip_tokenizer.h

@@ -0,0 +1,20 @@
+#ifndef __SD_TOKENIZERS_CLIP_TOKENIZER_H__
+#define __SD_TOKENIZERS_CLIP_TOKENIZER_H__
+
+#include <cstddef>
+#include <string>
+#include <vector>
+
+#include "bpe_tokenizer.h"
+
+class CLIPTokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str);
+    std::string normalize(const std::string& text) const override;
+    std::vector<std::string> token_split(const std::string& text) const override;
+
+public:
+    explicit CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_CLIP_TOKENIZER_H__

src/tokenizers/mistral_tokenizer.cpp

@@ -0,0 +1,89 @@
+#include "mistral_tokenizer.h"
+
+#include "ggml.h"
+#include "json.hpp"
+#include "util.h"
+#include "vocab/vocab.h"
+
+void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    nlohmann::json vocab;
+
+    try {
+        vocab = nlohmann::json::parse(vocab_utf8_str);
+    } catch (const nlohmann::json::parse_error&) {
+        GGML_ABORT("invalid vocab json str");
+    }
+    for (const auto& [key, value] : vocab.items()) {
+        std::u32string token = utf8_to_utf32(key);
+        int i                = value;
+        encoder[token]       = i;
+        decoder[i]           = token;
+    }
+    encoder_len = static_cast<int>(vocab.size());
+    LOG_DEBUG("vocab size: %d", encoder_len);
+
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    LOG_DEBUG("merges size %llu", merges.size());
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+MistralTokenizer::MistralTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    add_bos_token = true;
+
+    UNK_TOKEN = "<unk>";
+    BOS_TOKEN = "<s>";
+    EOS_TOKEN = "</s>";
+    PAD_TOKEN = "<pad>";
+
+    UNK_TOKEN_ID = 0;
+    BOS_TOKEN_ID = 1;
+    EOS_TOKEN_ID = 2;
+    PAD_TOKEN_ID = 11;
+
+    special_tokens = {
+        "<unk>",
+        "<s>",
+        "</s>",
+        "[INST]",
+        "[/INST]",
+        "[AVAILABLE_TOOLS]",
+        "[/AVAILABLE_TOOLS]",
+        "[TOOL_RESULTS]",
+        "[/TOOL_RESULTS]",
+        "[TOOL_CALLS]",
+        "[IMG]",
+        "<pad>",
+        "[IMG_BREAK]",
+        "[IMG_END]",
+        "[PREFIX]",
+        "[MIDDLE]",
+        "[SUFFIX]",
+        "[SYSTEM_PROMPT]",
+        "[/SYSTEM_PROMPT]",
+        "[TOOL_CONTENT]",
+    };
+    for (int i = 20; i < 1000; i++) {
+        special_tokens.push_back("<SPECIAL_" + std::to_string(i) + ">");
+    }
+
+    if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str, vocab_utf8_str);
+    } else {
+        load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
+    }
+}

src/tokenizers/mistral_tokenizer.h

@@ -0,0 +1,16 @@
+#ifndef __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
+#define __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class MistralTokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
+
+public:
+    explicit MistralTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__

src/tokenizers/qwen2_tokenizer.cpp

@@ -0,0 +1,91 @@
+#include "qwen2_tokenizer.h"
+
+#include "util.h"
+#include "vocab/vocab.h"
+
+void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) {
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    LOG_DEBUG("merges size %llu", merges.size());
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+
+    std::vector<std::u32string> tokens;
+    for (const auto& pair : byte_unicode_pairs) {
+        tokens.push_back(pair.second);
+    }
+    for (const auto& merge : merge_pairs) {
+        tokens.push_back(merge.first + merge.second);
+    }
+    for (auto& special_token : special_tokens) {
+        tokens.push_back(utf8_to_utf32(special_token));
+    }
+
+    int i = 0;
+    for (const auto& token : tokens) {
+        encoder[token] = i;
+        decoder[i]     = token;
+        i++;
+    }
+    encoder_len = i;
+    LOG_DEBUG("vocab size: %d", encoder_len);
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
+    UNK_TOKEN = "<|endoftext|>";
+    EOS_TOKEN = "<|endoftext|>";
+    PAD_TOKEN = "<|endoftext|>";
+
+    UNK_TOKEN_ID = 151643;
+    EOS_TOKEN_ID = 151643;
+    PAD_TOKEN_ID = 151643;
+
+    special_tokens = {
+        "<|endoftext|>",
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|object_ref_start|>",
+        "<|object_ref_end|>",
+        "<|box_start|>",
+        "<|box_end|>",
+        "<|quad_start|>",
+        "<|quad_end|>",
+        "<|vision_start|>",
+        "<|vision_end|>",
+        "<|vision_pad|>",
+        "<|image_pad|>",
+        "<|video_pad|>",
+        "<tool_call>",
+        "</tool_call>",
+        "<|fim_prefix|>",
+        "<|fim_middle|>",
+        "<|fim_suffix|>",
+        "<|fim_pad|>",
+        "<|repo_name|>",
+        "<|file_sep|>",
+        "<tool_response>",
+        "</tool_response>",
+        "<think>",
+        "</think>",
+    };
+
+    if (merges_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str);
+    } else {
+        load_from_merges(load_qwen2_merges());
+    }
+}

src/tokenizers/qwen2_tokenizer.h

@@ -0,0 +1,16 @@
+#ifndef __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
+#define __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class Qwen2Tokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str);
+
+public:
+    explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_QWEN2_TOKENIZER_H__

src/tokenizers/t5_unigram_tokenizer.cpp

@@ -0,0 +1,339 @@
+#include "t5_unigram_tokenizer.h"
+
+#include <algorithm>
+#include <cfloat>
+#include <cmath>
+#include <regex>
+#include <sstream>
+
+#include "json.hpp"
+#include "tokenize_util.h"
+#include "util.h"
+#include "vocab/vocab.h"
+
+// Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
+// and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
+// Original License: https://github.com/google/sentencepiece/blob/master/LICENSE
+//
+// Since tokenization is not the bottleneck in SD, performance was not a major consideration
+// during the migration.
+
+MetaspacePreTokenizer::MetaspacePreTokenizer(const std::string replacement, bool add_prefix_space)
+    : replacement(replacement), add_prefix_space(add_prefix_space) {}
+
+std::string MetaspacePreTokenizer::tokenize(const std::string& input) const {
+    std::string tokens;
+    std::stringstream ss(input);
+
+    if (add_prefix_space) {
+        tokens += replacement;
+    }
+
+    std::string token;
+    bool first_token = true;
+    while (std::getline(ss, token, ' ')) {
+        if (!first_token) {
+            tokens += replacement + token;
+        } else {
+            tokens += token;
+        }
+
+        first_token = false;
+    }
+
+    return tokens;
+}
+
+void T5UniGramTokenizer::InitializePieces(const std::string& json_str) {
+    nlohmann::json data;
+
+    try {
+        data = nlohmann::json::parse(json_str);
+    } catch (const nlohmann::json::parse_error&) {
+        status_ = INVLIAD_JSON;
+        return;
+    }
+    if (!data.contains("model")) {
+        status_ = INVLIAD_JSON;
+        return;
+    }
+    nlohmann::json model = data["model"];
+    if (!model.contains("vocab")) {
+        status_ = INVLIAD_JSON;
+        return;
+    }
+    if (model.contains("unk_id")) {
+        UNK_TOKEN_ID = model["unk_id"];
+    }
+
+    replacement      = data["pre_tokenizer"]["replacement"];
+    add_prefix_space = data["pre_tokenizer"]["add_prefix_space"];
+
+    pre_tokenizer = MetaspacePreTokenizer(replacement, add_prefix_space);
+
+    for (const auto& item : model["vocab"]) {
+        if (item.size() != 2 || !item[0].is_string() || !item[1].is_number_float()) {
+            status_ = INVLIAD_JSON;
+            return;
+        }
+        std::string piece = item[0];
+        if (piece.empty()) {
+            piece = "<empty_token>";
+        }
+        float score = item[1];
+        piece_score_pairs.emplace_back(piece, score);
+    }
+}
+
+void T5UniGramTokenizer::BuildTrie(std::vector<std::pair<std::string, int>>* pieces) {
+    if (status_ != OK) {
+        return;
+    }
+
+    if (pieces->empty()) {
+        status_ = NO_PIECES_LOADED;
+        return;
+    }
+
+    std::sort(pieces->begin(), pieces->end());
+
+    std::vector<const char*> key(pieces->size());
+    std::vector<int> value(pieces->size());
+    for (size_t i = 0; i < pieces->size(); ++i) {
+        key[i]   = (*pieces)[i].first.data();
+        value[i] = (*pieces)[i].second;
+    }
+
+    trie_ = std::unique_ptr<Darts::DoubleArray>(new Darts::DoubleArray());
+    if (trie_->build(key.size(), const_cast<char**>(&key[0]), nullptr, &value[0]) != 0) {
+        status_ = BUILD_DOUBLE_ARRAY_FAILED;
+        return;
+    }
+
+    const int kMaxTrieResultsSize = 1024;
+    std::vector<Darts::DoubleArray::result_pair_type> results(kMaxTrieResultsSize);
+    trie_results_size_ = 0;
+    for (const auto& p : *pieces) {
+        const size_t num_nodes = trie_->commonPrefixSearch(
+            p.first.data(), results.data(), results.size(), p.first.size());
+        trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
+    }
+
+    if (trie_results_size_ == 0) {
+        status_ = NO_ENTRY_FOUND;
+    }
+}
+
+float T5UniGramTokenizer::GetScoreInlined(int id) const {
+    return piece_score_pairs[id].second;
+}
+
+bool T5UniGramTokenizer::IsUnusedInlined(int id) const {
+    (void)id;
+    return false;
+}
+
+bool T5UniGramTokenizer::IsUserDefinedInlined(int id) const {
+    (void)id;
+    return false;
+}
+
+size_t T5UniGramTokenizer::OneCharLen(const char* src) const {
+    return "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"[(*src & 0xFF) >> 4];
+}
+
+EncodeResult T5UniGramTokenizer::EncodeOptimized(const std::string& normalized) const {
+    if (status() != OK || normalized.empty()) {
+        return {};
+    }
+
+    struct BestPathNode {
+        int id                = -1;
+        float best_path_score = 0;
+        int starts_at         = -1;
+    };
+
+    const int size        = static_cast<int>(normalized.size());
+    const float unk_score = min_score() - kUnkPenalty;
+    std::vector<BestPathNode> best_path_ends_at(size + 1);
+
+    int starts_at = 0;
+    while (starts_at < size) {
+        std::size_t node_pos                 = 0;
+        std::size_t key_pos                  = starts_at;
+        const auto best_path_score_till_here = best_path_ends_at[starts_at].best_path_score;
+        bool has_single_node                 = false;
+        const int mblen                      = std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)), size - starts_at);
+        while (key_pos < static_cast<size_t>(size)) {
+            const int ret = trie_->traverse(normalized.data(), node_pos, key_pos, key_pos + 1);
+            if (ret == -2) {
+                break;
+            }
+            if (ret >= 0) {
+                if (IsUnusedInlined(ret)) {
+                    continue;
+                }
+                auto& target_node                    = best_path_ends_at[key_pos];
+                const auto length                    = static_cast<int>(key_pos - starts_at);
+                const auto score                     = IsUserDefinedInlined(ret) ? (length * max_score_ - 0.1f) : GetScoreInlined(ret);
+                const auto candidate_best_path_score = score + best_path_score_till_here;
+                if (target_node.starts_at == -1 || candidate_best_path_score > target_node.best_path_score) {
+                    target_node.best_path_score = static_cast<float>(candidate_best_path_score);
+                    target_node.starts_at       = starts_at;
+                    target_node.id              = ret;
+                }
+                if (!has_single_node && length == mblen) {
+                    has_single_node = true;
+                }
+            }
+        }
+        if (!has_single_node) {
+            auto& target_node                    = best_path_ends_at[starts_at + mblen];
+            const auto candidate_best_path_score = unk_score + best_path_score_till_here;
+            if (target_node.starts_at == -1 || candidate_best_path_score > target_node.best_path_score) {
+                target_node.best_path_score = candidate_best_path_score;
+                target_node.starts_at       = starts_at;
+                target_node.id              = UNK_TOKEN_ID;
+            }
+        }
+        starts_at += mblen;
+    }
+
+    EncodeResult results;
+    int ends_at = size;
+    while (ends_at > 0) {
+        const auto& node = best_path_ends_at[ends_at];
+        results.emplace_back(normalized.substr(node.starts_at, ends_at - node.starts_at), node.id);
+        ends_at = node.starts_at;
+    }
+    std::reverse(results.begin(), results.end());
+    return results;
+}
+
+T5UniGramTokenizer::T5UniGramTokenizer(bool is_umt5) {
+    add_bos_token = false;
+    add_eos_token = true;
+
+    if (is_umt5) {
+        PAD_TOKEN_ID = 0;
+        EOS_TOKEN_ID = 1;
+        BOS_TOKEN_ID = 2;
+        UNK_TOKEN_ID = 3;
+
+        PAD_TOKEN = "<pad>";
+        EOS_TOKEN = "</s>";
+        BOS_TOKEN = "<s>";
+        UNK_TOKEN = "<unk>";
+    } else {
+        PAD_TOKEN_ID = 0;
+        EOS_TOKEN_ID = 1;
+        UNK_TOKEN_ID = 2;
+
+        PAD_TOKEN = "<pad>";
+        EOS_TOKEN = "</s>";
+        UNK_TOKEN = "<unk>";
+    }
+
+    special_tokens = {
+        "<pad>",
+        "</s>",
+        "<unk>",
+    };
+
+    if (is_umt5) {
+        special_tokens.push_back("<s>");
+    }
+
+    if (is_umt5) {
+        InitializePieces(load_umt5_tokenizer_json());
+    } else {
+        InitializePieces(load_t5_tokenizer_json());
+    }
+
+    min_score_ = FLT_MAX;
+    max_score_ = FLT_MIN;
+
+    std::vector<std::pair<std::string, int>> pieces;
+    for (int i = 0; i < static_cast<int>(piece_score_pairs.size()); i++) {
+        const auto& sp = piece_score_pairs[i];
+
+        min_score_ = std::min(min_score_, sp.second);
+        max_score_ = std::max(max_score_, sp.second);
+
+        pieces.emplace_back(sp.first, i);
+    }
+
+    BuildTrie(&pieces);
+}
+
+T5UniGramTokenizer::~T5UniGramTokenizer() = default;
+
+std::string T5UniGramTokenizer::decode_token(int token_id) const {
+    if (token_id < 0 || token_id >= static_cast<int>(piece_score_pairs.size())) {
+        return "";
+    }
+
+    const std::string& piece = piece_score_pairs[token_id].first;
+    if (piece == "<empty_token>") {
+        return "";
+    }
+    return piece;
+}
+
+std::string T5UniGramTokenizer::normalize(const std::string& input) const {
+    // Ref: https://github.com/huggingface/tokenizers/blob/1ff56c0c70b045f0cd82da1af9ac08cd4c7a6f9f/bindings/python/py_src/tokenizers/implementations/sentencepiece_unigram.py#L29
+    // TODO: nmt-nfkc
+    std::string normalized = std::regex_replace(input, std::regex(" {2,}"), " ");
+    return normalized;
+}
+
+std::vector<int> T5UniGramTokenizer::encode(const std::string& input, on_new_token_cb_t on_new_token_cb) {
+    std::vector<int32_t> tokens;
+    std::vector<std::string> token_strs;
+    std::string normalized = normalize(input);
+    auto splited_texts     = split_with_special_tokens(normalized, special_tokens);
+    if (splited_texts.empty()) {
+        splited_texts.push_back(normalized);  // for empty string
+    }
+
+    for (auto& splited_text : splited_texts) {
+        if (is_special_token(splited_text)) {
+            if (on_new_token_cb != nullptr) {
+                bool skip = on_new_token_cb(splited_text, tokens);
+                if (skip) {
+                    token_strs.push_back(splited_text);
+                    continue;
+                }
+            }
+
+            if (splited_text == UNK_TOKEN) {
+                tokens.push_back(UNK_TOKEN_ID);
+                token_strs.push_back(UNK_TOKEN);
+            } else if (splited_text == EOS_TOKEN) {
+                tokens.push_back(EOS_TOKEN_ID);
+                token_strs.push_back(EOS_TOKEN);
+            } else if (splited_text == PAD_TOKEN) {
+                tokens.push_back(PAD_TOKEN_ID);
+                token_strs.push_back(PAD_TOKEN);
+            }
+            continue;
+        }
+
+        std::string pretokenized = pre_tokenizer.tokenize(splited_text);
+        EncodeResult result      = EncodeOptimized(pretokenized);
+        for (const auto& item : result) {
+            tokens.push_back(item.second);
+            token_strs.push_back(item.first);
+        }
+    }
+
+    std::stringstream ss;
+    ss << "[";
+    for (const auto& token_str : token_strs) {
+        ss << "\"" << token_str << "\", ";
+    }
+    ss << "]";
+    LOG_DEBUG("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str());
+
+    return tokens;
+}

src/tokenizers/t5_unigram_tokenizer.h

@@ -0,0 +1,70 @@
+#ifndef __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
+#define __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "darts.h"
+#include "tokenizer.h"
+
+class MetaspacePreTokenizer {
+private:
+    std::string replacement;
+    bool add_prefix_space;
+
+public:
+    MetaspacePreTokenizer(const std::string replacement = " ", bool add_prefix_space = true);
+
+    std::string tokenize(const std::string& input) const;
+};
+
+using EncodeResult = std::vector<std::pair<std::string, int>>;
+
+class T5UniGramTokenizer : public Tokenizer {
+public:
+    enum Status {
+        OK,
+        NO_PIECES_LOADED,
+        NO_ENTRY_FOUND,
+        BUILD_DOUBLE_ARRAY_FAILED,
+        PIECE_ALREADY_DEFINED,
+        INVLIAD_JSON
+    };
+
+protected:
+    MetaspacePreTokenizer pre_tokenizer;
+    std::vector<std::pair<std::string, float>> piece_score_pairs;
+    float min_score_ = 0.0f;
+    float max_score_ = 0.0f;
+    std::unique_ptr<Darts::DoubleArray> trie_;
+    int trie_results_size_ = 0;
+    Status status_         = OK;
+    float kUnkPenalty      = 10.0f;
+    std::string replacement;
+    bool add_prefix_space = true;
+
+    void InitializePieces(const std::string& json_str);
+    void BuildTrie(std::vector<std::pair<std::string, int>>* pieces);
+    float GetScoreInlined(int id) const;
+    bool IsUnusedInlined(int id) const;
+    bool IsUserDefinedInlined(int id) const;
+    size_t OneCharLen(const char* src) const;
+    EncodeResult EncodeOptimized(const std::string& normalized) const;
+
+    float min_score() const { return min_score_; }
+    float max_score() const { return max_score_; }
+    Status status() const { return status_; }
+    std::string decode_token(int token_id) const override;
+    std::string normalize(const std::string& input) const override;
+
+public:
+    explicit T5UniGramTokenizer(bool is_umt5 = false);
+    ~T5UniGramTokenizer();
+
+    std::vector<int> encode(const std::string& input, on_new_token_cb_t on_new_token_cb = nullptr) override;
+};
+
+#endif  // __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__

src/tokenizers/tokenize_util.cpp

@@ -1,4 +1,4 @@
-#include <algorithm>
+#include <algorithm>
 #include <iostream>
 #include <string>
 #include <vector>

src/tokenizers/tokenize_util.h

@@ -1,5 +1,5 @@
-#ifndef __TOKENIZE_UTIL__
-#define __TOKENIZE_UTIL__
+#ifndef __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
+#define __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
 
 #include <string>
 #include <vector>
@@ -7,4 +7,4 @@
 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__
+#endif  // __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__

src/tokenizers/tokenizer.cpp

@@ -0,0 +1,222 @@
+#include "tokenizer.h"
+
+#include <algorithm>
+#include <cmath>
+#include <regex>
+
+#include "util.h"
+
+void Tokenizer::add_special_token(const std::string& token) {
+    special_tokens.push_back(token);
+}
+
+bool Tokenizer::is_special_token(const std::string& token) const {
+    for (const auto& special_token : special_tokens) {
+        if (special_token == token) {
+            return true;
+        }
+    }
+    return false;
+}
+
+std::string Tokenizer::normalize(const std::string& text) const {
+    return text;
+}
+
+std::vector<int> Tokenizer::tokenize(const std::string& text,
+                                     on_new_token_cb_t on_new_token_cb,
+                                     bool padding,
+                                     size_t min_length,
+                                     size_t max_length,
+                                     bool allow_overflow_expand) {
+    std::vector<int> tokens = encode(text, on_new_token_cb);
+    if (padding) {
+        pad_tokens(tokens, nullptr, nullptr, min_length, max_length, allow_overflow_expand);
+    }
+    return tokens;
+}
+
+void Tokenizer::pad_tokens(std::vector<int>& tokens,
+                           std::vector<float>* weights,
+                           std::vector<float>* mask,
+                           size_t min_length,
+                           size_t max_length,
+                           bool allow_overflow_expand) {
+    const bool use_weights = weights != nullptr;
+    const bool use_mask    = mask != nullptr;
+
+    if (use_weights && tokens.size() != weights->size()) {
+        LOG_ERROR("tokens size != weights size");
+        return;
+    }
+
+    const size_t bos_count           = add_bos_token ? 1 : 0;
+    const size_t eos_count           = add_eos_token ? 1 : 0;
+    const size_t special_token_count = bos_count + eos_count;
+
+    auto build_sequence = [&](size_t begin,
+                              size_t count,
+                              size_t target_length,
+                              std::vector<int>& out_tokens,
+                              std::vector<float>& out_weights,
+                              std::vector<float>& out_mask) {
+        const size_t base_length  = count + special_token_count;
+        const size_t final_length = std::max(target_length, base_length);
+
+        out_tokens.clear();
+        out_weights.clear();
+        out_mask.clear();
+
+        out_tokens.reserve(final_length);
+        if (use_weights) {
+            out_weights.reserve(final_length);
+        }
+        if (use_mask) {
+            out_mask.reserve(final_length);
+        }
+
+        if (add_bos_token) {
+            out_tokens.push_back(BOS_TOKEN_ID);
+            if (use_weights) {
+                out_weights.push_back(1.0f);
+            }
+            if (use_mask) {
+                out_mask.push_back(1.0f);
+            }
+        }
+
+        for (size_t i = 0; i < count; ++i) {
+            out_tokens.push_back(tokens[begin + i]);
+            if (use_weights) {
+                out_weights.push_back((*weights)[begin + i]);
+            }
+            if (use_mask) {
+                out_mask.push_back(1.0f);
+            }
+        }
+
+        if (add_eos_token) {
+            out_tokens.push_back(EOS_TOKEN_ID);
+            if (use_weights) {
+                out_weights.push_back(1.0f);
+            }
+            if (use_mask) {
+                out_mask.push_back(1.0f);
+            }
+        }
+
+        if (final_length > out_tokens.size()) {
+            const size_t pad_count = final_length - out_tokens.size();
+            if (pad_left) {
+                out_tokens.insert(out_tokens.begin(), pad_count, PAD_TOKEN_ID);
+
+                if (use_weights) {
+                    out_weights.insert(out_weights.begin(), pad_count, 1.0f);
+                }
+                if (use_mask) {
+                    out_mask.insert(out_mask.begin(), pad_count, 0.0f);
+                }
+            } else {
+                out_tokens.insert(out_tokens.end(), pad_count, PAD_TOKEN_ID);
+
+                if (use_weights) {
+                    out_weights.insert(out_weights.end(), pad_count, 1.0f);
+                }
+                if (use_mask) {
+                    out_mask.insert(out_mask.end(), pad_count, 0.0f);
+                }
+            }
+        }
+    };
+
+    const size_t single_length    = std::max(min_length, tokens.size() + special_token_count);
+    const bool exceeds_max_length = max_length > 0 && single_length > max_length;
+
+    std::vector<int> new_tokens;
+    std::vector<float> new_weights;
+    std::vector<float> new_mask;
+
+    if (!exceeds_max_length) {
+        build_sequence(0, tokens.size(), min_length, new_tokens, new_weights, new_mask);
+    } else if (!allow_overflow_expand) {
+        build_sequence(0, tokens.size(), 0, new_tokens, new_weights, new_mask);
+
+        new_tokens.resize(max_length);
+        if (use_weights) {
+            new_weights.resize(max_length);
+        }
+        if (use_mask) {
+            new_mask.resize(max_length);
+        }
+
+        if (add_eos_token && !new_tokens.empty()) {
+            new_tokens.back() = EOS_TOKEN_ID;
+            if (use_weights) {
+                new_weights.back() = 1.0f;
+            }
+            if (use_mask) {
+                new_mask.back() = 1.0f;
+            }
+        }
+    } else if (min_length > special_token_count) {
+        const size_t tokens_per_chunk = min_length - special_token_count;
+        size_t offset                 = 0;
+
+        while (offset < tokens.size()) {
+            const size_t remaining = tokens.size() - offset;
+            const size_t take      = std::min(tokens_per_chunk, remaining);
+
+            std::vector<int> chunk_tokens;
+            std::vector<float> chunk_weights;
+            std::vector<float> chunk_mask;
+
+            build_sequence(offset, take, min_length, chunk_tokens, chunk_weights, chunk_mask);
+
+            new_tokens.insert(new_tokens.end(), chunk_tokens.begin(), chunk_tokens.end());
+            if (use_weights) {
+                new_weights.insert(new_weights.end(), chunk_weights.begin(), chunk_weights.end());
+            }
+            if (use_mask) {
+                new_mask.insert(new_mask.end(), chunk_mask.begin(), chunk_mask.end());
+            }
+
+            offset += take;
+        }
+    } else {
+        build_sequence(0, tokens.size(), min_length, new_tokens, new_weights, new_mask);
+    }
+
+    tokens = std::move(new_tokens);
+    if (use_weights) {
+        *weights = std::move(new_weights);
+    }
+    if (use_mask) {
+        *mask = std::move(new_mask);
+    }
+}
+
+static std::string clean_up_tokenization(std::string& text) {
+    std::regex pattern(R"( ,)");
+    return std::regex_replace(text, pattern, ",");
+}
+
+std::string Tokenizer::decode(const std::vector<int>& tokens) const {
+    std::string text;
+
+    for (int token_id : tokens) {
+        if (token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID || token_id == PAD_TOKEN_ID) {
+            continue;
+        }
+
+        std::string piece = decode_token(token_id);
+        if (!end_of_word_suffix.empty() && ends_with(piece, end_of_word_suffix)) {
+            piece.erase(piece.size() - end_of_word_suffix.size());
+            text += piece + " ";
+        } else {
+            text += piece;
+        }
+    }
+
+    text = clean_up_tokenization(text);
+    return trim(text);
+}

src/tokenizers/tokenizer.h

@@ -0,0 +1,53 @@
+#ifndef __SD_TOKENIZERS_TOKENIZER_H__
+#define __SD_TOKENIZERS_TOKENIZER_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <string>
+#include <vector>
+
+using on_new_token_cb_t = std::function<bool(std::string&, std::vector<int32_t>&)>;
+
+class Tokenizer {
+protected:
+    std::vector<std::string> special_tokens;
+    bool add_bos_token = false;
+    bool add_eos_token = false;
+    bool pad_left      = false;
+    std::string end_of_word_suffix;
+
+    virtual std::string decode_token(int token_id) const = 0;
+    virtual std::string normalize(const std::string& text) const;
+
+public:
+    std::string UNK_TOKEN;
+    std::string BOS_TOKEN;
+    std::string EOS_TOKEN;
+    std::string PAD_TOKEN;
+    int UNK_TOKEN_ID = 0;
+    int BOS_TOKEN_ID = 0;
+    int EOS_TOKEN_ID = 0;
+    int PAD_TOKEN_ID = 0;
+
+    virtual ~Tokenizer() = default;
+
+    void add_special_token(const std::string& token);
+    bool is_special_token(const std::string& token) const;
+    virtual std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) = 0;
+    std::vector<int> tokenize(const std::string& text,
+                              on_new_token_cb_t on_new_token_cb = nullptr,
+                              bool padding                      = false,
+                              size_t min_length                 = 0,
+                              size_t max_length                 = 100000000,
+                              bool allow_overflow_expand        = false);
+    void pad_tokens(std::vector<int>& tokens,
+                    std::vector<float>* weights,
+                    std::vector<float>* mask,
+                    size_t min_length          = 0,
+                    size_t max_length          = 100000000,
+                    bool allow_overflow_expand = false);
+    std::string decode(const std::vector<int>& tokens) const;
+};
+
+#endif  // __SD_TOKENIZERS_TOKENIZER_H__

src/tokenizers/vocab/vocab.h

@@ -1,5 +1,5 @@
-#ifndef __VOCAB_H__
-#define __VOCAB_H__
+#ifndef __SD_TOKENIZERS_VOCAB_VOCAB_H__
+#define __SD_TOKENIZERS_VOCAB_VOCAB_H__
 
 #include <string>
 
@@ -10,4 +10,4 @@ std::string load_mistral_vocab_json();
 std::string load_t5_tokenizer_json();
 std::string load_umt5_tokenizer_json();
 
-#endif  // __VOCAB_H__
+#endif  // __SD_TOKENIZERS_VOCAB_VOCAB_H__

src/vae.hpp

@@ -69,7 +69,7 @@ public:
         int scale_factor = 8;
         if (version == VERSION_WAN2_2_TI2V) {
             scale_factor = 16;
-        } else if (sd_version_is_flux2(version)) {
+        } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
         } else if (version == VERSION_CHROMA_RADIANCE) {
             scale_factor = 1;

thirdparty/CMakeLists.txt

@@ -1,3 +1,45 @@
 set(Z_TARGET zip)
 add_library(${Z_TARGET} OBJECT zip.c zip.h miniz.h)
-target_include_directories(${Z_TARGET} PUBLIC .)
+target_include_directories(${Z_TARGET} PUBLIC .)
+
+if(SD_WEBP AND NOT SD_USE_SYSTEM_WEBP)
+    set(WEBP_BUILD_ANIM_UTILS OFF)
+    set(WEBP_BUILD_CWEBP OFF)
+    set(WEBP_BUILD_DWEBP OFF)
+    set(WEBP_BUILD_GIF2WEBP OFF)
+    set(WEBP_BUILD_IMG2WEBP OFF)
+    set(WEBP_BUILD_VWEBP OFF)
+    set(WEBP_BUILD_WEBPINFO OFF)
+    set(WEBP_BUILD_WEBPMUX OFF)
+    set(WEBP_BUILD_EXTRAS OFF)
+    set(WEBP_BUILD_WEBP_JS OFF)
+    set(WEBP_BUILD_FUZZTEST OFF)
+    set(WEBP_BUILD_LIBWEBPMUX ON)
+
+    add_subdirectory(libwebp EXCLUDE_FROM_ALL)
+endif()
+
+if(SD_WEBM AND NOT SD_USE_SYSTEM_WEBM)
+    if(MSVC)
+        set(MSVC_RUNTIME dll)
+    endif()
+    set(ENABLE_WEBMTS OFF)
+    set(ENABLE_WEBMINFO OFF)
+    set(ENABLE_TESTS OFF)
+    set(ENABLE_WEBM_PARSER OFF)
+    set(ENABLE_SAMPLE_PROGRAMS OFF)
+
+    set(SD_LIBWEBM_PARENT_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
+
+    add_subdirectory(libwebm EXCLUDE_FROM_ALL)
+
+    # libwebm mutates the global CMAKE_CXX_FLAGS for non-MSVC compilers to force
+    # C++11. Restore the parent flags so the main project keeps its own C++17
+    # requirements, then pin the libwebm targets to C++17 explicitly.
+    set(CMAKE_CXX_FLAGS "${SD_LIBWEBM_PARENT_CXX_FLAGS}" CACHE STRING "" FORCE)
+    target_compile_features(mkvmuxer PRIVATE cxx_std_17)
+    target_compile_features(mkvparser PRIVATE cxx_std_17)
+    target_compile_features(webm PRIVATE cxx_std_17)
+
+    target_include_directories(webm INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/libwebm)
+endif()