fix: enable Wan/TAEHV video generation on the Metal backend (#850) (#1731)

Co-authored-by: phil2sat <phil2sat@users.noreply.github.com>

.github/workflows/build.yml

@@ -83,7 +83,7 @@ jobs:
         run: |
           mkdir build
           cd build
-          cmake .. -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON
+          cmake .. -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_INSTALL_RPATH='$ORIGIN'
           cmake --build . --config Release
 
       - name: Get commit hash
@@ -146,7 +146,7 @@ jobs:
         run: |
           mkdir build
           cd build
-          cmake .. -DSD_BUILD_SHARED_LIBS=ON -DSD_VULKAN=ON
+          cmake .. -DSD_BUILD_SHARED_LIBS=ON -DSD_VULKAN=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_INSTALL_RPATH='$ORIGIN'
           cmake --build . --config Release
 
       - name: Get commit hash
@@ -207,6 +207,7 @@ jobs:
               UBUNTU_VERSION=24.04
               CUDA_ARCHITECTURES=121
               GGML_CUDA_FA_ALL_QUANTS=ON
+              GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=OFF
 
     env:
       REGISTRY: ghcr.io
@@ -341,18 +342,12 @@ jobs:
     strategy:
       matrix:
         include:
-          - build: "noavx"
-            defines: "-DGGML_NATIVE=OFF -DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF -DSD_BUILD_SHARED_LIBS=ON"
-          - build: "avx2"
-            defines: "-DGGML_NATIVE=OFF -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON"
-          - build: "avx"
-            defines: "-DGGML_NATIVE=OFF -DGGML_AVX=ON -DGGML_AVX2=OFF -DSD_BUILD_SHARED_LIBS=ON"
-          - build: "avx512"
-            defines: "-DGGML_NATIVE=OFF -DGGML_AVX512=ON -DGGML_AVX=ON -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON"
+          - build: "cpu"
+            defines: "-DGGML_NATIVE=OFF -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
           - build: "cuda12"
-            defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"'"
+            defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
           - build: "vulkan"
-            defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON"
+            defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
     steps:
       - name: Clone
         id: checkout
@@ -400,19 +395,6 @@ jobs:
           cmake .. -DCMAKE_CXX_FLAGS='/bigobj' -G Ninja -DCMAKE_C_COMPILER=cl.exe -DCMAKE_CXX_COMPILER=cl.exe -DCMAKE_BUILD_TYPE=Release ${{ matrix.defines }}
           cmake --build .
 
-      - name: Check AVX512F support
-        id: check_avx512f
-        if: ${{ matrix.build == 'avx512' }}
-        continue-on-error: true
-        run: |
-          cd build
-          $vcdir = $(vswhere -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath)
-          $msvc = $(join-path $vcdir $('VC\Tools\MSVC\'+$(gc -raw $(join-path $vcdir 'VC\Auxiliary\Build\Microsoft.VCToolsVersion.default.txt')).Trim()))
-          $cl =  $(join-path $msvc 'bin\Hostx64\x64\cl.exe')
-          echo 'int main(void){unsigned int a[4];__cpuid(a,7);return !(a[1]&65536);}' >> avx512f.c
-          & $cl /O2 /GS- /kernel avx512f.c /link /nodefaultlib /entry:main
-          .\avx512f.exe && echo "AVX512F: YES" && ( echo HAS_AVX512F=1 >> $env:GITHUB_ENV ) || echo "AVX512F: NO"
-
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -722,6 +704,25 @@ jobs:
         run: |
           sudo apt install -y build-essential cmake wget zip ninja-build
 
+      - name: Free disk space
+        run: |
+          df -h
+
+          # Remove preinstalled SDKs and caches not needed for this job before
+          # installing ROCm. The legacy ROCm apt packages are large enough to
+          # exhaust ubuntu-latest if cleanup runs after installation.
+          sudo rm -rf /usr/share/dotnet || true
+          sudo rm -rf /usr/local/lib/android || true
+          sudo rm -rf /opt/ghc || true
+          sudo rm -rf /usr/local/.ghcup || true
+          sudo rm -rf /opt/hostedtoolcache || true
+          sudo rm -rf /usr/share/swift || true
+          sudo rm -rf /usr/local/share/boost || true
+          docker system prune -af || true
+
+          sudo apt clean
+          df -h
+
       - name: Setup Legacy ROCm
         if: matrix.ROCM_VERSION == '7.2.1'
         id: legacy_env
@@ -743,19 +744,6 @@ jobs:
           sudo apt update
           sudo apt-get install -y libssl-dev rocm-hip-sdk
 
-      - name: Free disk space
-        run: |
-          # Remove preinstalled SDKs and caches not needed for this job
-          sudo rm -rf /usr/share/dotnet || true
-          sudo rm -rf /usr/local/lib/android || true
-          sudo rm -rf /opt/ghc || true
-          sudo rm -rf /usr/local/.ghcup || true
-          sudo rm -rf /opt/hostedtoolcache || true
-
-          # Remove old package lists and caches
-          sudo rm -rf /var/lib/apt/lists/* || true
-          sudo apt clean
-
       - name: Setup TheRock
         if: matrix.ROCM_VERSION != '7.2.1'
         id: therock_env
@@ -794,6 +782,11 @@ jobs:
             -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
             -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
             -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
+            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DGGML_NATIVE=OFF \
+            -DSD_BUILD_SHARED_GGML_LIB=ON \
+            -DGGML_BACKEND_DL=ON \
+            -DGGML_CPU_ALL_VARIANTS=ON \
             -DSD_BUILD_SHARED_LIBS=ON
           cmake --build . --config Release
 

CMakeLists.txt

@@ -331,7 +331,8 @@ endif()
 
 add_subdirectory(thirdparty)
 
-target_link_libraries(${SD_LIB} PUBLIC ggml zip)
+target_sources(${SD_LIB} PRIVATE $<TARGET_OBJECTS:zip>)
+target_link_libraries(${SD_LIB} PUBLIC ggml)
 target_include_directories(${SD_LIB} PUBLIC . src include)
 target_include_directories(${SD_LIB} PRIVATE src/core)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
@@ -342,7 +343,58 @@ if (SD_BUILD_EXAMPLES)
     add_subdirectory(examples)
 endif()
 
-set(SD_PUBLIC_HEADERS include/stable-diffusion.h)
-set_target_properties(${SD_LIB} PROPERTIES PUBLIC_HEADER "${SD_PUBLIC_HEADERS}")
 
-install(TARGETS ${SD_LIB} LIBRARY PUBLIC_HEADER)
+
+#
+# install
+#
+
+include(CMakePackageConfigHelpers)
+include(GNUInstallDirs)
+
+set(SD_INSTALL_VERSION "${SDCPP_BUILD_VERSION}")
+set(SD_INSTALL_COMMIT  "${SDCPP_BUILD_COMMIT}")
+set(SD_SHARED_LIB      ${SD_BUILD_SHARED_LIBS})
+
+set(SD_INCLUDE_INSTALL_DIR ${CMAKE_INSTALL_INCLUDEDIR} CACHE PATH "Location of header  files")
+set(SD_LIB_INSTALL_DIR     ${CMAKE_INSTALL_LIBDIR}     CACHE PATH "Location of library files")
+set(SD_BIN_INSTALL_DIR     ${CMAKE_INSTALL_BINDIR}     CACHE PATH "Location of binary  files")
+
+set(SD_PUBLIC_HEADERS
+    ${CMAKE_CURRENT_SOURCE_DIR}/include/stable-diffusion.h)
+
+set_target_properties(${SD_LIB}
+    PROPERTIES
+        PUBLIC_HEADER "${SD_PUBLIC_HEADERS}")
+
+
+install(TARGETS ${SD_LIB}
+        ARCHIVE
+        LIBRARY
+        RUNTIME
+        PUBLIC_HEADER)
+
+
+configure_package_config_file(
+        ${CMAKE_CURRENT_SOURCE_DIR}/cmake/stable-diffusion-config.cmake.in
+        ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-config.cmake
+    INSTALL_DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/stable-diffusion
+    PATH_VARS SD_INCLUDE_INSTALL_DIR
+              SD_LIB_INSTALL_DIR
+              SD_BIN_INSTALL_DIR )
+
+write_basic_package_version_file(
+        ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-version.cmake
+    VERSION ${SD_INSTALL_VERSION}
+    COMPATIBILITY SameMajorVersion)
+
+install(FILES ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-config.cmake
+              ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-version.cmake
+        DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/stable-diffusion)
+
+configure_file(cmake/stable-diffusion.pc.in
+        "${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion.pc"
+        @ONLY)
+
+install(FILES "${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion.pc"
+        DESTINATION ${CMAKE_INSTALL_LIBDIR}/pkgconfig)

Dockerfile

@@ -19,7 +19,14 @@ WORKDIR /sd.cpp
 
 COPY . .
 
-RUN cmake . -B ./build
+RUN cmake . -B ./build \
+    -DSD_BUILD_SHARED_LIBS=ON \
+    -DGGML_NATIVE=OFF \
+    -DSD_BUILD_SHARED_GGML_LIB=ON \
+    -DGGML_BACKEND_DL=ON \
+    -DGGML_CPU_ALL_VARIANTS=ON \
+    -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+    -DCMAKE_INSTALL_RPATH='$ORIGIN'
 RUN cmake --build ./build --config Release --parallel
 
 FROM ubuntu:$UBUNTU_VERSION AS runtime
@@ -28,7 +35,9 @@ RUN apt-get update && \
     apt-get install --yes --no-install-recommends libgomp1 && \
     apt-get clean
 
-COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
-COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
+COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
+RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
+    printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
+    chmod +x /sd-cli /sd-server
 
 ENTRYPOINT [ "/sd-cli" ]

Dockerfile.cuda

@@ -23,11 +23,27 @@ COPY . .
 ARG CUDACXX=/usr/local/cuda/bin/nvcc
 ARG CUDA_ARCHITECTURES=""
 ARG GGML_CUDA_FA_ALL_QUANTS=""
+ARG GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=ON
 
-RUN cmake . -B ./build \
-    -DSD_CUDA=ON \
-    ${CUDA_ARCHITECTURES:+-DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}"} \
-    ${GGML_CUDA_FA_ALL_QUANTS:+-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}}
+RUN set -- \
+        -DSD_CUDA=ON; \
+    if [ "${GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS}" = "ON" ]; then \
+        set -- "$@" \
+            -DSD_BUILD_SHARED_LIBS=ON \
+            -DGGML_NATIVE=OFF \
+            -DSD_BUILD_SHARED_GGML_LIB=ON \
+            -DGGML_BACKEND_DL=ON \
+            -DGGML_CPU_ALL_VARIANTS=ON \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+            "-DCMAKE_INSTALL_RPATH=\$ORIGIN"; \
+    fi; \
+    if [ -n "${CUDA_ARCHITECTURES}" ]; then \
+        set -- "$@" "-DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}"; \
+    fi; \
+    if [ -n "${GGML_CUDA_FA_ALL_QUANTS}" ]; then \
+        set -- "$@" "-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}"; \
+    fi; \
+    cmake . -B ./build "$@"
 RUN cmake --build ./build --config Release -j$(nproc)
 
 FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime
@@ -36,7 +52,9 @@ RUN apt-get update && \
     apt-get install --yes --no-install-recommends libgomp1 && \
     apt-get clean
 
-COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
-COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
+COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
+RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
+    printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
+    chmod +x /sd-cli /sd-server
 
 ENTRYPOINT [ "/sd-cli" ]

Dockerfile.musa

@@ -24,12 +24,22 @@ RUN mkdir build && cd build && \
     cmake .. -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \
         -DCMAKE_C_FLAGS="${CMAKE_C_FLAGS} -fopenmp -I/usr/lib/llvm-14/lib/clang/14.0.0/include -L/usr/lib/llvm-14/lib" \
         -DCMAKE_CXX_FLAGS="${CMAKE_CXX_FLAGS} -fopenmp -I/usr/lib/llvm-14/lib/clang/14.0.0/include -L/usr/lib/llvm-14/lib" \
-        -DSD_MUSA=ON -DCMAKE_BUILD_TYPE=Release && \
+        -DSD_MUSA=ON \
+        -DSD_BUILD_SHARED_LIBS=ON \
+        -DGGML_NATIVE=OFF \
+        -DSD_BUILD_SHARED_GGML_LIB=ON \
+        -DGGML_BACKEND_DL=ON \
+        -DGGML_CPU_ALL_VARIANTS=ON \
+        -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+        -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+        -DCMAKE_BUILD_TYPE=Release && \
     cmake --build . --config Release
 
 FROM mthreads/musa:${MUSA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}-amd64 as runtime
 
-COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
-COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
+COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
+RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
+    printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
+    chmod +x /sd-cli /sd-server
 
 ENTRYPOINT [ "/sd-cli" ]

Dockerfile.sycl

@@ -29,4 +29,4 @@ FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS runtime
 COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
 COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
 
-ENTRYPOINT [ "/sd-cli" ]
+ENTRYPOINT [ "/sd-cli" ]

Dockerfile.vulkan

@@ -19,8 +19,16 @@ WORKDIR /sd.cpp
 
 COPY . .
 
-RUN cmake . -B ./build -DSD_VULKAN=ON
-RUN cmake --build ./build --config Release --parallel
+RUN cmake . -B ./build \
+    -DSD_VULKAN=ON \
+    -DSD_BUILD_SHARED_LIBS=ON \
+    -DGGML_NATIVE=OFF \
+    -DSD_BUILD_SHARED_GGML_LIB=ON \
+    -DGGML_BACKEND_DL=ON \
+    -DGGML_CPU_ALL_VARIANTS=ON \
+    -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+    -DCMAKE_INSTALL_RPATH='$ORIGIN'
+RUN cmake --build ./build --config Release -j$(nproc)
 
 FROM ubuntu:$UBUNTU_VERSION AS runtime
 
@@ -28,7 +36,9 @@ RUN apt-get update && \
     apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers && \
     apt-get clean
 
-COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
-COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
+COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
+RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
+    printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
+    chmod +x /sd-cli /sd-server
 
 ENTRYPOINT [ "/sd-cli" ]

README.md

@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/06/25** 🚀 stable-diffusion.cpp now supports **Krea2**
 * **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
 * **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD**
 * **2026/05/27** 🚀 stable-diffusion.cpp now supports **Lens**
@@ -47,15 +48,20 @@ API and command-line option may change frequently.***
     - [PiD](./docs/pid.md)
     - [LongCat Image](./docs/longcat_image.md)
     - [Z-Image](./docs/z_image.md)
+    - [MiniT2I](./docs/minit2i.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
     - [ERNIE-Image](./docs/ernie_image.md)
+    - [Boogu Image](./docs/boogu_image.md)
+    - [Krea2](./docs/krea2.md)
+    - [SeFi-Image](./docs/sefi_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
     - [LongCat Image Edit](./docs/longcat_image.md)
+    - [Boogu Image Edit](./docs/boogu_image.md)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [LTX-2.3](./docs/ltx2.md)

cmake/stable-diffusion-config.cmake.in

@@ -0,0 +1,37 @@
+set(SD_VERSION      "@SD_INSTALL_VERSION@")
+set(SD_BUILD_COMMIT "@SD_INSTALL_COMMIT@")
+set(SD_SHARED_LIB    @SD_SHARED_LIB@)
+
+@PACKAGE_INIT@
+
+set_and_check(SD_INCLUDE_DIR "@PACKAGE_SD_INCLUDE_INSTALL_DIR@")
+set_and_check(SD_LIB_DIR     "@PACKAGE_SD_LIB_INSTALL_DIR@")
+set(SD_BIN_DIR "@PACKAGE_SD_BIN_INSTALL_DIR@")
+
+include(CMakeFindDependencyMacro)
+find_dependency(ggml REQUIRED HINTS "${SD_LIB_DIR}/cmake")
+
+if(NOT TARGET stable-diffusion)
+    find_library(stable-diffusion_LIBRARY stable-diffusion
+        REQUIRED
+        HINTS "${SD_LIB_DIR}"
+        NO_CMAKE_FIND_ROOT_PATH
+    )
+
+    add_library(stable-diffusion UNKNOWN IMPORTED)
+    set_target_properties(stable-diffusion
+        PROPERTIES
+            INTERFACE_INCLUDE_DIRECTORIES "${SD_INCLUDE_DIR}"
+            INTERFACE_LINK_LIBRARIES "ggml::ggml"
+            IMPORTED_LINK_INTERFACE_LANGUAGES "CXX"
+            IMPORTED_LOCATION "${stable-diffusion_LIBRARY}"
+            INTERFACE_COMPILE_FEATURES "c_std_11;cxx_std_17"
+            POSITION_INDEPENDENT_CODE ON)
+
+    if(SD_SHARED_LIB)
+        target_compile_definitions(stable-diffusion
+            INTERFACE SD_BUILD_SHARED_LIB)
+    endif()
+endif()
+
+check_required_components(stable-diffusion)

cmake/stable-diffusion.pc.in

@@ -0,0 +1,11 @@
+prefix=@CMAKE_INSTALL_PREFIX@
+exec_prefix=${prefix}
+libdir=@CMAKE_INSTALL_FULL_LIBDIR@
+includedir=@CMAKE_INSTALL_FULL_INCLUDEDIR@
+
+Name: stable-diffusion
+Description: Diffusion model(SD,Flux,Wan,Qwen Image,Z-Image,...) inference in pure C/C++
+Version: @SDCPP_BUILD_VERSION@
+Libs: -L${libdir} -lstable-diffusion
+Libs.private: -lggml -lggml-base
+Cflags: -I${includedir}

docs/boogu_image.md

@@ -0,0 +1,31 @@
+# How to Use
+
+Boogu Image uses a Boogu diffusion transformer, the FLUX VAE, and Qwen3-VL as the LLM text and vision encoder.
+
+## Download weights
+
+- Download Boogu Image
+    - safetensors: https://huggingface.co/Comfy-Org/Boogu-Image/tree/main/diffusion_models
+- Download vae
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/ae.safetensors
+- Download Qwen3-VL 8B
+    - gguf: https://huggingface.co/unsloth/Qwen3-VL-8B-Instruct-GGUF/tree/main
+        - For image editing with GGUF text encoders, also download the matching mmproj file and pass it with `--llm_vision`.
+
+## Examples
+
+### Boogu Image Base
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\boogu_image_base_bf16.safetensors --llm ..\..\llm\Qwen3VL-8B-Instruct-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ae.sft -p "a lovely cat" --diffusion-fa -v --offload-to-cpu
+```
+
+<img width="256" alt="Boogu Image Base example" src="../assets/boogu/example.png" />
+
+### Boogu Image Edit
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\boogu_image_edit_bf16.safetensors --llm ..\..\llm\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\..\llm\mmproj-Qwen3VL-8B-Instruct-F16.gguf --vae ..\..\ComfyUI\models\vae\ae.sft --diffusion-fa -v --offload-to-cpu -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'boogu.cpp'"
+```
+
+<img width="256" alt="Boogu Image Edit example" src="../assets/boogu/edit_example.png" />

docs/imatrix.md

@@ -0,0 +1,59 @@
+# Importance Matrix (imatrix) Quantization
+
+## What is an Importance Matrix?
+
+Quantization reduces the precision of a model's weights, decreasing its size and computational requirements. However, this can lead to a loss of quality. An importance matrix helps mitigate this by identifying which weights are *most* important for the model's performance. During quantization, these important weights are preserved with higher precision, while less important weights are quantized more aggressively.  This allows for better overall quality at a given quantization level.
+
+This originates from work done with language models in [llama.cpp](https://github.com/ggml-org/llama.cpp/blob/master/tools/imatrix/README.md).
+
+## Usage
+
+The imatrix feature involves two main steps: *training* the matrix and *using* it during quantization.
+
+### Training the Importance Matrix
+
+To generate an imatrix, run stable-diffusion.cpp with the `--imat-out` flag, specifying the output filename.  This process runs alongside normal image generation.
+
+```bash
+sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat
+```
+
+*   **`[same exact parameters as normal generation]`**:  Use the same command-line arguments you would normally use for image generation (e.g., prompt, dimensions, sampling method, etc.).
+*   **`--imat-out imatrix.dat`**: Specifies the output file for the generated imatrix.
+
+You can generate multiple images at once using the `-b` flag to speed up the training process.
+
+### Continuing Training an Existing Matrix
+
+If you want to refine an existing imatrix, use the `--imat-in` flag *in addition* to `--imat-out`. This will load the existing matrix and continue training it.
+
+```bash
+sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat
+```
+With that, you can train and refine the imatrix while generating images like you'd normally do. 
+
+### Using Multiple Matrices
+
+You can load and merge multiple imatrices together:
+
+```bash
+sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat --imat-in imatrix2.dat
+```
+
+### Quantizing with an Importance Matrix
+
+To quantize a model using a trained imatrix, use the `-M convert` option (or equivalent quantization command) and the `--imat-in` flag, specifying the imatrix file.
+
+```bash
+sd.exe -M convert [same exact parameters as normal quantization] --imat-in imatrix.dat
+```
+
+*   **`[same exact parameters as normal quantization]`**: Use the same command-line arguments you would normally use for quantization (e.g., target quantization method, input/output filenames).
+*   **`--imat-in imatrix.dat`**: Specifies the imatrix file to use during quantization.  You can specify multiple `--imat-in` flags to combine multiple matrices.
+
+## Important Considerations
+
+*   The quality of the imatrix depends on the prompts and settings used during training. Use prompts and settings representative of the types of images you intend to generate for the best results.
+*   Experiment with different training parameters (e.g., number of images, prompt variations) to optimize the imatrix for your specific use case.
+*   The performance impact of training an imatrix during image generation or using an imatrix for quantization is negligible.
+*   Using already quantized models to train the imatrix seems to be working fine.

docs/krea2.md

@@ -0,0 +1,27 @@
+# How to Use
+
+Krea2 uses a Krea2 diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the LLM text encoder.
+
+## Download weights
+
+- Download Krea2 Raw
+    - safetensors: https://huggingface.co/krea/Krea-2-Raw/tree/main
+    - gguf: https://huggingface.co/realrebelai/KREA-2_GGUFs/tree/main/BASE
+- Download Krea2 Turbo
+    - safetensors: https://huggingface.co/krea/Krea-2-Turbo/tree/main
+    - gguf: https://huggingface.co/realrebelai/KREA-2_GGUFs/tree/main/TURBO
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
+- Download Qwen3-VL 4B
+    - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
+
+## Examples
+
+### Krea2
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\Krea-2-Raw-Q8_0.gguf --llm ..\..\ComfyUI\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors -p "a lovely cat holding a sign says 'krea2.cpp'"  --diffusion-fa -v --offload-to-cpu
+```
+
+<img width="256" alt="Krea2 Raw example" src="../assets/krea2/example.png" />

docs/minit2i.md

@@ -0,0 +1,48 @@
+# How to Use
+
+MiniT2I uses a MiniT2I diffusion transformer and `google/flan-t5-large` as the text encoder.
+
+## Download weights
+
+- Download MiniT2I diffusion model
+    - safetensors: https://huggingface.co/MiniT2I/MiniT2I/tree/main/minit2i-b-16/transformer (`diffusion_pytorch_model.safetensors`)
+- Download flan-t5-large text encoder
+    - safetensors: https://huggingface.co/google/flan-t5-large/tree/main (`model.safetensors`)
+
+## Examples
+
+### Mac Metal
+
+```
+./bin/sd-cli \
+  --backend metal \
+  --diffusion-model ../models/minit2i/diffusion_pytorch_model.safetensors \
+  --t5xxl ../models/flan-t5-large/model.safetensors \
+  --prompt "a cat" \
+  --steps 100 \
+  --cfg-scale 6 \
+  --width 512 \
+  --height 512 \
+  --seed 42 \
+  --sampling-method euler \
+  --rng cpu \
+  --output minit2i_metal.png \
+  --threads 8
+```
+
+### CUDA with diffusion flash attention
+
+```
+./bin/sd-cli \
+  --diffusion-model ../models/minit2i/diffusion_pytorch_model.safetensors \
+  --t5xxl ../models/flan-t5-large/model.safetensors \
+  --prompt "a cat" \
+  --steps 100 \
+  --cfg-scale 6 \
+  --width 512 \
+  --height 512 \
+  --seed 42 \
+  --sampling-method euler \
+  --diffusion-fa \
+  --output minit2i_cuda.png
+```

docs/sefi_image.md

@@ -0,0 +1,50 @@
+# How to Use
+
+SeFi-Image uses a Flux2-style dual-time transformer (semantic + texture streams), the standard Flux2 VAE, and Qwen3-VL as the LLM text encoder. Tech report: [arXiv:2606.22568](https://arxiv.org/abs/2606.22568).
+
+## Download weights
+
+The SeFi-Image family ships in three scales (1B / 2B / 5B) and three families (Base / RL / turbo), all gated on Hugging Face under https://huggingface.co/SeFi-Image.
+
+- 1B and 2B variants pair with Qwen3-VL-2B-Instruct.
+- 5B variants pair with Qwen3-VL-4B-Instruct.
+- All variants use the standard Flux2 VAE (`flux2_ae.safetensors` from https://huggingface.co/black-forest-labs/FLUX.2-dev).
+
+Convert the transformer and text encoder to sd.cpp safetensors:
+
+```bash
+python3 script/convert_sefi.py     <hf_repo_dir>                          <out_dir>/sefi_<scale>_<family>.safetensors
+python3 script/convert_qwen3_vl.py  <hf_repo_dir>/Qwen3-VL-XB-Instruct    <out_dir>/qwen3_vl_<X>b.safetensors
+```
+
+## Variant defaults
+
+| Family | timestep_shift_alpha | steps | cfg-scale |
+|---|---|---|---|
+| Base | 0.3 | 50 | 4.0 |
+| RL | 0.3 | 50 | 4.0 |
+| turbo | 1.0 | 4 | 1.0 |
+
+The dispatcher picks `alpha` from the filename (`turbo` substring => 1.0, otherwise 0.3). Override via `--extra-sample-args sefi_alpha=<value>` or `sefi_delta_t=<value>`.
+
+## Examples
+
+### 1B / 2B turbo
+
+```
+./build/bin/sd-cli --diffusion-model /path/to/sefi_1b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_2b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --offload-to-cpu -o out.png
+```
+
+### 1B / 2B base
+
+```
+./build/bin/sd-cli --diffusion-model /path/to/sefi_1b_base.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_2b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 4.0 --steps 50 -W 1024 -H 1024 -s 42 --diffusion-fa --offload-to-cpu -o out.png
+```
+
+### 5B (needs streaming on 12 GiB VRAM)
+
+```
+./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --stream-layers --offload-to-cpu -o out.png
+```
+
+<img alt="SeFi-Image 5B turbo example" src="../assets/sefi_image/example.png" />

examples/cli/main.cpp

@@ -1,6 +1,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <time.h>
+#include <algorithm>
 #include <cctype>
 #include <filesystem>
 #include <functional>
@@ -53,6 +54,9 @@ struct SDCliParams {
     bool metadata_brief      = false;
     bool metadata_all        = false;
 
+    std::string imatrix_out;
+    std::vector<std::string> imatrix_in;
+
     bool normal_exit = false;
 
     ArgOptions get_options() {
@@ -62,19 +66,28 @@ struct SDCliParams {
             {"-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). Single-file video outputs support .avi, .webm, and animated .webp",
+             0,
              &output_path},
             {"",
              "--image",
              "path to the image to inspect (for metadata mode)",
+             0,
              &image_path},
             {"",
              "--metadata-format",
              "metadata output format, one of [text, json] (default: text)",
+             0,
              &metadata_format},
             {"",
              "--preview-path",
              "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
+             0,
              &preview_path},
+            {"",
+             "--imat-out",
+             "compute the imatrix for this run and save it to the provided path",
+             0,
+             &imatrix_out},
         };
 
         options.int_options = {
@@ -175,6 +188,14 @@ struct SDCliParams {
             return -1;
         };
 
+        auto on_imatrix_in_arg = [&](int argc, const char** argv, int index) {
+            if (++index >= argc) {
+                return -1;
+            }
+            imatrix_in.push_back(argv[index]);
+            return 1;
+        };
+
         options.manual_options = {
             {"-M",
              "--mode",
@@ -188,6 +209,10 @@ struct SDCliParams {
              "--help",
              "show this help message and exit",
              on_help_arg},
+            {"",
+             "--imat-in",
+             "load an imatrix file for quantization or continued collection; can be specified multiple times",
+             on_imatrix_in_arg},
         };
 
         return options;
@@ -249,6 +274,7 @@ struct SDCliParams {
             << "  preview_fps: " << preview_fps << ",\n"
             << "  taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n"
             << "  preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n"
+            << "  imatrix_out: \"" << imatrix_out << "\",\n"
             << "  metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n"
             << "  metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n"
             << "  metadata_all: " << (metadata_all ? "true" : "false") << "\n"
@@ -455,7 +481,8 @@ bool save_results(const SDCliParams& cli_params,
         if (!img.data)
             return false;
 
-        const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx;
+        int images_per_batch        = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1;
+        const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch;
         std::string params          = gen_params.embed_image_metadata
                                           ? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
                                           : "";
@@ -601,13 +628,32 @@ int main(int argc, const char* argv[]) {
     LOG_DEBUG("%s", ctx_params.to_string().c_str());
     LOG_DEBUG("%s", gen_params.to_string().c_str());
 
+    if (!cli_params.imatrix_out.empty()) {
+        if (fs::exists(cli_params.imatrix_out) &&
+            std::find(cli_params.imatrix_in.begin(), cli_params.imatrix_in.end(), cli_params.imatrix_out) == cli_params.imatrix_in.end()) {
+            LOG_WARN("imatrix file '%s' already exists and will be overwritten", cli_params.imatrix_out.c_str());
+        }
+        enable_imatrix_collection();
+    }
+
+    for (const auto& in_file : cli_params.imatrix_in) {
+        LOG_INFO("loading imatrix from '%s'", in_file.c_str());
+        if (!load_imatrix(in_file.c_str())) {
+            LOG_WARN("failed to load imatrix from '%s'", in_file.c_str());
+        }
+    }
+
     if (cli_params.mode == CONVERT) {
-        bool success = convert(ctx_params.model_path.c_str(),
-                               ctx_params.vae_path.c_str(),
-                               cli_params.output_path.c_str(),
-                               ctx_params.wtype,
-                               ctx_params.tensor_type_rules.c_str(),
-                               cli_params.convert_name);
+        bool success = convert_with_components(ctx_params.model_path.c_str(),
+                                               ctx_params.clip_l_path.c_str(),
+                                               ctx_params.clip_g_path.c_str(),
+                                               ctx_params.t5xxl_path.c_str(),
+                                               ctx_params.diffusion_model_path.c_str(),
+                                               ctx_params.vae_path.c_str(),
+                                               cli_params.output_path.c_str(),
+                                               ctx_params.wtype,
+                                               ctx_params.tensor_type_rules.c_str(),
+                                               cli_params.convert_name);
         if (!success) {
             LOG_ERROR("convert '%s'/'%s' to '%s' failed",
                       ctx_params.model_path.c_str(),
@@ -762,8 +808,12 @@ int main(int argc, const char* argv[]) {
         if (cli_params.mode == IMG_GEN) {
             sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
 
-            num_results = gen_params.batch_count;
-            results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
+            sd_image_t* generated_images = nullptr;
+            if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
+                generated_images = nullptr;
+                num_results      = 0;
+            }
+            results.adopt(generated_images, num_results);
         } else if (cli_params.mode == VID_GEN) {
             sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
             sd_image_t* generated_video        = nullptr;
@@ -798,12 +848,22 @@ int main(int argc, const char* argv[]) {
                 SDImageOwner current_image(results[i]);
                 results[i] = {0, 0, 0, nullptr};
                 for (int u = 0; u < gen_params.upscale_repeats; ++u) {
-                    SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor));
-                    if (upscaled_image.get().data == nullptr) {
+                    sd_image_t* upscaled_images = nullptr;
+                    int upscaled_count          = 0;
+                    bool upscale_ok             = upscale(upscaler_ctx.get(),
+                                                          current_image.get(),
+                                                          upscale_factor,
+                                                          &upscaled_images,
+                                                          &upscaled_count);
+                    if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) {
+                        free_sd_images(upscaled_images, upscaled_count);
                         LOG_ERROR("upscale failed");
                         break;
                     }
-                    current_image = std::move(upscaled_image);
+                    sd_image_t upscaled_image = upscaled_images[0];
+                    upscaled_images[0]        = {0, 0, 0, nullptr};
+                    free_sd_images(upscaled_images, upscaled_count);
+                    current_image.reset(upscaled_image);
                 }
                 results[i] = current_image.release();  // Set the final upscaled image as the result
             }
@@ -815,6 +875,11 @@ int main(int argc, const char* argv[]) {
         return 1;
     }
 
+    if (!cli_params.imatrix_out.empty()) {
+        LOG_INFO("saving imatrix to '%s'", cli_params.imatrix_out.c_str());
+        save_imatrix(cli_params.imatrix_out.c_str());
+    }
+
     free_sd_audio(generated_audio);
 
     return 0;

examples/common/common.cpp

@@ -6,6 +6,7 @@
 #include <cstdlib>
 #include <ctime>
 #include <filesystem>
+#include <fstream>
 #include <iomanip>
 #include <iostream>
 #include <regex>
@@ -260,8 +261,15 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
                         invalid_arg = true;
                         return;
                     }
-                    *option.target = argv_to_utf8(i, argv);
-                    found_arg      = true;
+                    if (option.concat && !option.target->empty()) {
+                        if (option.concat > 0 && option.concat <= 0xff) {
+                            *option.target += static_cast<char>(option.concat);
+                        }
+                        *option.target += argv_to_utf8(i, argv);
+                    } else {
+                        *option.target = argv_to_utf8(i, argv);
+                    }
+                    found_arg = true;
                 }))
                 break;
 
@@ -324,120 +332,151 @@ ArgOptions SDContextParams::get_options() {
         {"-m",
          "--model",
          "path to full model",
+         0,
          &model_path},
         {"",
          "--clip_l",
-         "path to the clip-l text encoder", &clip_l_path},
+         "path to the clip-l text encoder",
+         0,
+         &clip_l_path},
         {"", "--clip_g",
          "path to the clip-g text encoder",
+         0,
          &clip_g_path},
         {"",
          "--clip_vision",
          "path to the clip-vision encoder",
+         0,
          &clip_vision_path},
         {"",
          "--t5xxl",
          "path to the t5xxl text encoder",
+         0,
          &t5xxl_path},
         {"",
          "--llm",
          "path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)",
+         0,
          &llm_path},
         {"",
          "--llm_vision",
          "path to the llm vit",
+         0,
          &llm_vision_path},
         {"",
          "--qwen2vl",
          "alias of --llm. Deprecated.",
+         0,
          &llm_path},
         {"",
          "--qwen2vl_vision",
          "alias of --llm_vision. Deprecated.",
+         0,
          &llm_vision_path},
         {"",
          "--diffusion-model",
          "path to the standalone diffusion model",
+         0,
          &diffusion_model_path},
         {"",
          "--high-noise-diffusion-model",
          "path to the standalone high noise diffusion model",
+         0,
          &high_noise_diffusion_model_path},
         {"",
          "--uncond-diffusion-model",
          "path to the standalone unconditional diffusion model, currently used by Ideogram4 CFG",
+         0,
          &uncond_diffusion_model_path},
         {"",
          "--embeddings-connectors",
          "path to LTXAV embeddings connectors",
+         0,
          &embeddings_connectors_path},
         {"",
          "--vae",
          "path to standalone vae model",
+         0,
          &vae_path},
         {"",
          "--vae-format",
          "VAE latent format override: auto, flux, sd3, or flux2 (default: auto)",
+         0,
          &vae_format},
         {"",
          "--audio-vae",
          "path to standalone LTX audio vae model",
+         0,
          &audio_vae_path},
         {"",
          "--taesd",
          "path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)",
+         0,
          &taesd_path},
         {"",
          "--tae",
          "alias of --taesd",
+         0,
          &taesd_path},
         {"",
          "--control-net",
          "path to control net model",
+         0,
          &control_net_path},
         {"",
          "--embd-dir",
          "embeddings directory",
+         0,
          &embedding_dir},
         {"",
          "--lora-model-dir",
          "lora model directory",
+         0,
          &lora_model_dir},
         {"",
          "--hires-upscalers-dir",
          "highres fix upscaler model directory",
+         0,
          &hires_upscalers_dir},
         {"",
          "--tensor-type-rules",
          "weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
+         (int)',',
          &tensor_type_rules},
         {"",
          "--photo-maker",
          "path to PHOTOMAKER model",
+         0,
          &photo_maker_path},
         {"",
          "--pulid-weights",
          "path to PuLID Flux weights",
+         0,
          &pulid_weights_path},
         {"",
          "--upscale-model",
          "path to esrgan model.",
+         0,
          &esrgan_path},
         {"",
          "--backend",
          "runtime backend assignment, e.g. cpu or clip=cpu,vae=cuda0,diffusion=vulkan0",
+         (int)',',
          &backend},
         {"",
          "--params-backend",
          "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu",
+         (int)',',
          &params_backend},
         {"",
          "--rpc-servers",
          "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052",
+         (int)',',
          &rpc_servers},
         {"",
          "--max-vram",
          "maximum VRAM budget in GiB for graph-cut segmented execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value",
+         0,
          &max_vram},
     };
 
@@ -458,6 +497,10 @@ ArgOptions SDContextParams::get_options() {
          "--stream-layers",
          "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)",
          true, &stream_layers},
+        {"",
+         "--eager-load",
+         "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
+         true, &eager_load},
         {"",
          "--force-sdxl-vae-conv-scale",
          "force use of conv scale on sdxl vae",
@@ -610,7 +653,7 @@ ArgOptions SDContextParams::get_options() {
          on_sampler_rng_arg},
         {"",
          "--prediction",
-         "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]",
+         "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]",
          on_prediction_arg},
         {"",
          "--lora-apply-mode",
@@ -667,7 +710,18 @@ bool SDContextParams::resolve(SDMode mode) {
 }
 
 bool SDContextParams::validate(SDMode mode) {
-    if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
+    if (mode == CONVERT) {
+        const bool has_convert_input = model_path.length() != 0 ||
+                                       clip_l_path.length() != 0 ||
+                                       clip_g_path.length() != 0 ||
+                                       t5xxl_path.length() != 0 ||
+                                       diffusion_model_path.length() != 0 ||
+                                       vae_path.length() != 0;
+        if (!has_convert_input) {
+            LOG_ERROR("error: convert mode needs at least one model input path\n");
+            return false;
+        }
+    } else if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
         LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n");
         return false;
     }
@@ -761,6 +815,7 @@ std::string SDContextParams::to_string() const {
         << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
         << "  max_vram: \"" << max_vram << "\",\n"
         << "  stream_layers: " << (stream_layers ? "true" : "false") << ",\n"
+        << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
@@ -840,6 +895,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.vae_format                      = str_to_vae_format(vae_format);
     sd_ctx_params.max_vram                        = max_vram.c_str();
     sd_ctx_params.stream_layers                   = stream_layers;
+    sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();
     sd_ctx_params.rpc_servers                     = rpc_servers.c_str();
@@ -857,58 +913,71 @@ ArgOptions SDGenerationParams::get_options() {
         {"-p",
          "--prompt",
          "the prompt to render",
+         0,
          &prompt},
         {"-n",
          "--negative-prompt",
          "the negative prompt (default: \"\")",
+         0,
          &negative_prompt},
         {"-i",
          "--init-img",
          "path to the init image",
+         0,
          &init_image_path},
         {"",
          "--end-img",
          "path to the end image, required by flf2v",
+         0,
          &end_image_path},
         {"",
          "--mask",
          "path to the mask image",
+         0,
          &mask_image_path},
         {"",
          "--control-image",
          "path to control image, control net",
+         0,
          &control_image_path},
         {"",
          "--control-video",
          "path to control video frames, It must be a directory path. The video frames inside should be stored as images in "
          "lexicographical (character) order. For example, if the control video path is `frames`, the directory contain images "
          "such as 00.png, 01.png, ... etc.",
+         0,
          &control_video_path},
         {"",
          "--pm-id-images-dir",
          "path to PHOTOMAKER input id images dir",
+         0,
          &pm_id_images_dir},
         {"",
          "--pm-id-embed-path",
          "path to PHOTOMAKER v2 id embed",
+         0,
          &pm_id_embed_path},
         {"",
          "--pulid-id-embedding",
          "path to PuLID id embedding",
+         0,
          &pulid_id_embedding_path},
         {"",
          "--hires-upscaler",
          "highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent (nearest-exact), "
          "Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name "
          "under --hires-upscalers-dir (default: Latent)",
+         0,
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma",
+         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
+         (int)',',
          &extra_sample_args},
         {"",
          "--extra-tiling-args",
          "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
+         (int)',',
          &extra_tiling_args},
     };
 
@@ -938,6 +1007,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--batch-count",
          "batch count",
          &batch_count},
+        {"",
+         "--qwen-image-layers",
+         "number of Qwen Image Layered layers; latent/output count is layers + 1 (default: 3)",
+         &qwen_image_layers},
         {"",
          "--video-frames",
          "video frames (default: 1)",
@@ -1364,6 +1437,42 @@ ArgOptions SDGenerationParams::get_options() {
         return 1;
     };
 
+    auto on_prompt_file_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            return -1;
+        }
+        const char* arg = argv[index];
+        std::ifstream f(arg, std::ios::binary);
+        try {
+            prompt = std::string(std::istreambuf_iterator<char>{f}, {});
+        } catch (const std::ios_base::failure&) {
+            f.setstate(std::ios_base::failbit);
+        }
+        if (f.fail()) {
+            LOG_ERROR("error: failed to read prompt file '%s'\n", arg);
+            return -1;
+        }
+        return 1;
+    };
+
+    auto on_negative_prompt_file_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            return -1;
+        }
+        const char* arg = argv[index];
+        std::ifstream f(arg, std::ios::binary);
+        try {
+            negative_prompt = std::string(std::istreambuf_iterator<char>{f}, {});
+        } catch (const std::ios_base::failure&) {
+            f.setstate(std::ios_base::failbit);
+        }
+        if (f.fail()) {
+            LOG_ERROR("error: failed to read negative prompt file '%s'\n", arg);
+            return -1;
+        }
+        return 1;
+    };
+
     options.manual_options = {
         {"-s",
          "--seed",
@@ -1381,7 +1490,7 @@ ArgOptions SDGenerationParams::get_options() {
          on_high_noise_sample_method_arg},
         {"",
          "--scheduler",
-         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default: model-specific",
+         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific",
          on_scheduler_arg},
         {"",
          "--sigmas",
@@ -1427,6 +1536,14 @@ ArgOptions SDGenerationParams::get_options() {
          "--vae-relative-tile-size",
          "relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)",
          on_relative_tile_size_arg},
+        {"",
+         "--prompt-file",
+         "path to the file containing the prompt to render",
+         on_prompt_file_arg},
+        {"",
+         "--negative-prompt-file",
+         "path to the file containing the negative prompt",
+         on_negative_prompt_file_arg},
 
     };
 
@@ -1714,6 +1831,7 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("width", width);
     load_if_exists("height", height);
     load_if_exists("batch_count", batch_count);
+    load_if_exists("qwen_image_layers", qwen_image_layers);
     load_if_exists("video_frames", video_frames);
     load_if_exists("fps", fps);
     load_if_exists("upscale_repeats", upscale_repeats);
@@ -2138,6 +2256,11 @@ bool SDGenerationParams::validate(SDMode mode) {
         return false;
     }
 
+    if (qwen_image_layers < 0) {
+        LOG_ERROR("error: qwen_image_layers must be non-negative");
+        return false;
+    }
+
     if (sample_params.sample_steps <= 0) {
         LOG_ERROR("error: the sample_steps must be greater than 0\n");
         return false;
@@ -2304,6 +2427,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.strength              = strength;
     params.seed                  = seed;
     params.batch_count           = batch_count;
+    params.qwen_image_layers     = qwen_image_layers;
     params.control_image         = control_image.get();
     params.control_strength      = control_strength;
     params.pm_params             = pm_params;
@@ -2429,6 +2553,7 @@ std::string SDGenerationParams::to_string() const {
         << "  width: " << width << ",\n"
         << "  height: " << height << ",\n"
         << "  batch_count: " << batch_count << ",\n"
+        << "  qwen_image_layers: " << qwen_image_layers << ",\n"
         << "  init_image_path: \"" << init_image_path << "\",\n"
         << "  end_image_path: \"" << end_image_path << "\",\n"
         << "  mask_image_path: \"" << mask_image_path << "\",\n"

examples/common/common.h

@@ -31,6 +31,7 @@ struct StringOption {
     std::string short_name;
     std::string long_name;
     std::string desc;
+    int concat;
     std::string* target;
 };
 
@@ -147,6 +148,7 @@ struct SDContextParams {
     bool offload_params_to_cpu  = false;
     std::string max_vram        = "0";
     bool stream_layers          = false;
+    bool eager_load             = false;
     std::string backend;
     std::string params_backend;
     std::string rpc_servers;
@@ -195,6 +197,7 @@ struct SDGenerationParams {
     int width                  = -1;
     int height                 = -1;
     int batch_count            = 1;
+    int qwen_image_layers      = 3;
     int64_t seed               = 42;
     float strength             = 0.75f;
     float control_strength     = 0.9f;

examples/server/async_jobs.cpp

@@ -2,6 +2,7 @@
 
 #include "async_jobs.h"
 
+#include <algorithm>
 #include <iomanip>
 #include <sstream>
 
@@ -173,8 +174,13 @@ bool execute_img_gen_job(ServerRuntime& runtime,
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
-        sd_image_t* raw_results = generate_image(runtime.sd_ctx, &params);
-        results.adopt(raw_results, params.batch_count);
+        sd_image_t* raw_results = nullptr;
+        int num_results         = 0;
+        if (!generate_image(runtime.sd_ctx, &params, &raw_results, &num_results)) {
+            raw_results = nullptr;
+            num_results = 0;
+        }
+        results.adopt(raw_results, num_results);
     }
 
     const int num_results = results.count();
@@ -190,6 +196,8 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         encoded_format = EncodedImageFormat::WEBP;
     }
 
+    int batch_count      = job.img_gen.gen_params.batch_count;
+    int images_per_batch = batch_count > 0 ? std::max(1, num_results / batch_count) : 1;
     for (int i = 0; i < num_results; ++i) {
         if (results[i].data == nullptr) {
             continue;
@@ -198,7 +206,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         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)
+                                                            job.img_gen.gen_params.seed + i / images_per_batch)
                                          : "";
         auto image_bytes           = encode_image_to_vector(encoded_format,
                                                             results[i].data,

examples/server/routes_openai.cpp

@@ -229,8 +229,11 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime,
 
     {
         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;
+        sd_image_t* raw_results = nullptr;
+        if (!generate_image(runtime.sd_ctx, &img_gen_params, &raw_results, &num_results)) {
+            raw_results = nullptr;
+            num_results = 0;
+        }
         results.adopt(raw_results, num_results);
     }
 
@@ -281,14 +284,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
             out["data"]          = json::array();
             out["output_format"] = request.output_format;
 
-            for (int i = 0; i < request.gen_params.batch_count; ++i) {
+            int result_count     = results.count();
+            int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1;
+            for (int i = 0; i < result_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)
+                                                            request.gen_params.seed + i / images_per_batch)
                                          : "";
                 auto image_bytes   = encode_image_to_vector(request.output_format == "jpeg"
                                                                 ? EncodedImageFormat::JPEG
@@ -353,14 +358,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
             out["data"]          = json::array();
             out["output_format"] = request.output_format;
 
-            for (int i = 0; i < request.gen_params.batch_count; ++i) {
+            int result_count     = results.count();
+            int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1;
+            for (int i = 0; i < result_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)
+                                                            request.gen_params.seed + i / images_per_batch)
                                          : "";
                 auto image_bytes   = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG,
                                                           results[i].data,

examples/server/routes_sdapi.cpp

@@ -292,8 +292,11 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
 
             {
                 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;
+                sd_image_t* raw_results = nullptr;
+                if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) {
+                    raw_results = nullptr;
+                    num_results = 0;
+                }
                 results.adopt(raw_results, num_results);
             }
 
@@ -308,6 +311,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
             out["parameters"] = j;
             out["info"]       = "";
 
+            int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, num_results / request.gen_params.batch_count) : 1;
             for (int i = 0; i < num_results; ++i) {
                 if (results[i].data == nullptr) {
                     continue;
@@ -316,7 +320,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 std::string params = request.gen_params.embed_image_metadata
                                          ? get_image_params(*runtime->ctx_params,
                                                             request.gen_params,
-                                                            request.gen_params.seed + i)
+                                                            request.gen_params.seed + i / images_per_batch)
                                          : "";
                 auto image_bytes   = encode_image_to_vector(EncodedImageFormat::PNG,
                                                             results[i].data,
@@ -438,6 +442,9 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
         scheduler_names.push_back("default");
         for (int i = 0; i < SCHEDULER_COUNT; i++) {
             scheduler_names.push_back(sd_scheduler_name((scheduler_t)i));
+            if (i == DISCRETE_SCHEDULER) {
+                scheduler_names.push_back("normal");
+            }
         }
         json r = json::array();
         for (auto name : scheduler_names) {

examples/server/routes_sdcpp.cpp

@@ -126,6 +126,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"strength", defaults.strength},
         {"seed", defaults.seed},
         {"batch_count", defaults.batch_count},
+        {"qwen_image_layers", defaults.qwen_image_layers},
         {"auto_resize_ref_image", defaults.auto_resize_ref_image},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
@@ -219,6 +220,9 @@ static json make_capabilities_json(ServerRuntime& runtime) {
 
     for (int i = 0; i < SCHEDULER_COUNT; ++i) {
         schedulers.push_back(sd_scheduler_name((scheduler_t)i));
+        if (i == DISCRETE_SCHEDULER) {
+            schedulers.push_back("normal");
+        }
     }
 
     {

examples/server/runtime.cpp

@@ -190,8 +190,8 @@ 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},
+        {"-l", "--listen-ip", "server listen ip (default: 127.0.0.1)", 0, &listen_ip},
+        {"", "--serve-html-path", "path to HTML file to serve at root (optional)", 0, &serve_html_path},
     };
 
     options.int_options = {

include/stable-diffusion.h

@@ -70,6 +70,10 @@ enum scheduler_t {
     LCM_SCHEDULER,
     BONG_TANGENT_SCHEDULER,
     LTX2_SCHEDULER,
+    LOGIT_NORMAL_SCHEDULER,
+    FLUX2_SCHEDULER,
+    FLUX_SCHEDULER,
+    BETA_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -79,7 +83,8 @@ enum prediction_t {
     EDM_V_PRED,
     FLOW_PRED,
     FLUX_FLOW_PRED,
-    FLUX2_FLOW_PRED,
+    SEFI_FLOW_PRED,
+    MINIT2I_FLOW_PRED,
     PREDICTION_COUNT
 };
 
@@ -219,6 +224,7 @@ typedef struct {
     enum sd_vae_format_t vae_format;
     const char* max_vram;  // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
     bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
+    bool eager_load;  // Load all params into the params backend at model-load time instead of lazily on first use
     const char* backend;
     const char* params_backend;
     const char* rpc_servers;
@@ -374,6 +380,7 @@ typedef struct {
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
     sd_hires_params_t hires;
+    int qwen_image_layers;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -402,14 +409,17 @@ typedef struct {
 } sd_vid_gen_params_t;
 
 typedef struct sd_ctx_t sd_ctx_t;
+struct ggml_tensor;
 
 typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
 typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
 typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
+typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data);
 
 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
 SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
 SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
+SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
 SD_API int32_t sd_get_num_physical_cores();
 SD_API const char* sd_get_system_info();
 SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
@@ -450,7 +460,10 @@ SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sa
 
 SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params);
 SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params);
-SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
+SD_API bool generate_image(sd_ctx_t* sd_ctx,
+                           const sd_img_gen_params_t* sd_img_gen_params,
+                           sd_image_t** images_out,
+                           int* num_images_out);
 
 enum sd_cancel_mode_t {
     // Stop the current generation as soon as possible.
@@ -480,9 +493,11 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
                                         const char* params_backend);
 SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
-SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,
-                          sd_image_t input_image,
-                          uint32_t upscale_factor);
+SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
+                    sd_image_t input_image,
+                    uint32_t upscale_factor,
+                    sd_image_t** images_out,
+                    int* num_images_out);
 
 SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
 
@@ -493,6 +508,17 @@ SD_API bool convert(const char* input_path,
                     const char* tensor_type_rules,
                     bool convert_name);
 
+SD_API bool convert_with_components(const char* model_path,
+                                    const char* clip_l_path,
+                                    const char* clip_g_path,
+                                    const char* t5xxl_path,
+                                    const char* diffusion_model_path,
+                                    const char* vae_path,
+                                    const char* output_path,
+                                    enum sd_type_t output_type,
+                                    const char* tensor_type_rules,
+                                    bool convert_name);
+
 SD_API bool preprocess_canny(sd_image_t image,
                              float high_threshold,
                              float low_threshold,
@@ -500,6 +526,11 @@ SD_API bool preprocess_canny(sd_image_t image,
                              float strong,
                              bool inverse);
 
+SD_API bool load_imatrix(const char* imatrix_path);
+SD_API void save_imatrix(const char* imatrix_path);
+SD_API void enable_imatrix_collection(void);
+SD_API void disable_imatrix_collection(void);
+
 SD_API const char* sd_commit(void);
 SD_API const char* sd_version(void);
 

script/convert_qwen3_vl.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""Convert a Qwen3-VL HF safetensors checkpoint into a sd.cpp-loadable form.
+
+The HF dump prefixes text-tower keys with ``model.language_model.`` and
+vision-tower keys with ``model.visual.``. sd.cpp expects ``model.<rest>`` for
+the text side; the vision side is converted by sd.cpp's own
+``convert_qwen3_vl_vision_name`` and is left as-is here.
+
+Operates on raw safetensors bytes so any dtype (BF16/F16/F32) is preserved.
+
+Usage:
+    python3 script/convert_qwen3_vl.py <hf_qwen3_vl_dir_or_safetensors> <output.safetensors>
+"""
+
+import argparse
+import json
+import os
+import struct
+import sys
+
+
+def rewrite_key(key: str) -> str:
+    if key.startswith("model.language_model."):
+        return "model." + key[len("model.language_model."):]
+    return key
+
+
+def read_safetensors_header(path: str):
+    with open(path, "rb") as f:
+        hdr_len = struct.unpack("<Q", f.read(8))[0]
+        hdr_bytes = f.read(hdr_len)
+    return json.loads(hdr_bytes), 8 + hdr_len
+
+
+def collect_shard_paths(path: str):
+    if os.path.isdir(path):
+        index_path = os.path.join(path, "model.safetensors.index.json")
+        if os.path.isfile(index_path):
+            with open(index_path) as f:
+                idx = json.load(f)
+            return sorted({os.path.join(path, n) for n in idx["weight_map"].values()})
+        single = os.path.join(path, "model.safetensors")
+        if os.path.isfile(single):
+            return [single]
+        raise FileNotFoundError(f"No Qwen3-VL safetensors in {path}")
+    if os.path.isfile(path):
+        return [path]
+    raise FileNotFoundError(path)
+
+
+def stage_tensors(input_path: str):
+    entries = []
+    for shard_path in collect_shard_paths(input_path):
+        hdr, data_off = read_safetensors_header(shard_path)
+        for key, info in hdr.items():
+            if key == "__metadata__":
+                continue
+            entries.append((rewrite_key(key), shard_path, data_off, info))
+    return entries
+
+
+def write_consolidated(out_path: str, entries):
+    entries = sorted(entries, key=lambda e: e[0])
+
+    new_header = {}
+    cur_offset = 0
+    for new_key, shard_path, data_off, info in entries:
+        start, end = info["data_offsets"]
+        size = end - start
+        new_header[new_key] = {
+            "dtype": info["dtype"],
+            "shape": info["shape"],
+            "data_offsets": [cur_offset, cur_offset + size],
+        }
+        cur_offset += size
+
+    header_json = json.dumps(new_header, separators=(",", ":")).encode("utf-8")
+    pad = (-len(header_json)) % 8
+    header_json = header_json + (b" " * pad)
+
+    with open(out_path, "wb") as out:
+        out.write(struct.pack("<Q", len(header_json)))
+        out.write(header_json)
+        for new_key, shard_path, data_off, info in entries:
+            start, end = info["data_offsets"]
+            with open(shard_path, "rb") as src:
+                src.seek(data_off + start)
+                remaining = end - start
+                while remaining > 0:
+                    chunk = src.read(min(8 * 1024 * 1024, remaining))
+                    if not chunk:
+                        raise IOError(f"Truncated tensor in {shard_path}")
+                    out.write(chunk)
+                    remaining -= len(chunk)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("input", help="HF Qwen3-VL directory or single safetensors file")
+    parser.add_argument("output", help="Output single safetensors path")
+    args = parser.parse_args()
+
+    entries = stage_tensors(args.input)
+    print(f"Tensors: {len(entries)}")
+    print(f"Writing -> {args.output}")
+    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+    write_consolidated(args.output, entries)
+    print(f"Done. Output size: {os.path.getsize(args.output) / 1e9:.2f} GB")
+
+
+if __name__ == "__main__":
+    main()

script/convert_sefi.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+"""Convert a SeFi-Image diffusers checkpoint into a single sd.cpp-compatible safetensors.
+
+Operates on raw safetensors bytes so any dtype (BF16, F32, ...) is preserved exactly.
+No numpy or torch dependency required.
+
+Usage:
+    python3 script/convert_sefi.py <sefi_diffusers_dir> <output.safetensors>
+"""
+
+import argparse
+import json
+import os
+import re
+import struct
+import sys
+
+
+_LINEAR_TO_LIN = re.compile(r"\.linear\.")
+_SHARED_MOD_PREFIXES = (
+    "double_stream_modulation_img",
+    "double_stream_modulation_txt",
+    "single_stream_modulation",
+)
+
+
+def rewrite_transformer_key(key: str) -> str:
+    if key.startswith("backbone."):
+        key = key[len("backbone."):]
+    elif key.startswith("dual_time_embed."):
+        return key
+
+    if any(key.startswith(prefix + ".") for prefix in _SHARED_MOD_PREFIXES):
+        key = _LINEAR_TO_LIN.sub(".lin.", key, count=1)
+
+    if key == "context_embedder.weight":
+        return "txt_in.weight"
+    if key == "context_embedder.bias":
+        return "txt_in.bias"
+    if key == "x_embedder.weight":
+        return "img_in.weight"
+    if key == "x_embedder.bias":
+        return "img_in.bias"
+
+    if key == "proj_out.weight":
+        return "final_layer.linear.weight"
+    if key == "proj_out.bias":
+        return "final_layer.linear.bias"
+    if key == "norm_out.linear.weight":
+        return "final_layer.adaLN_modulation.1.weight"
+    if key == "norm_out.linear.bias":
+        return "final_layer.adaLN_modulation.1.bias"
+
+    m = re.match(r"transformer_blocks\.(\d+)\.(.*)$", key)
+    if m:
+        return _rewrite_double_stream(m.group(1), m.group(2))
+    m = re.match(r"single_transformer_blocks\.(\d+)\.(.*)$", key)
+    if m:
+        return _rewrite_single_stream(m.group(1), m.group(2))
+
+    return key
+
+
+def _rewrite_double_stream(idx: str, tail: str) -> str:
+    dst = f"double_blocks.{idx}."
+    mapping = {
+        "norm1.linear.weight":          "img_mod.lin.weight",
+        "norm1_context.linear.weight":  "txt_mod.lin.weight",
+        "attn.norm_q.weight":           "img_attn.norm.query_norm.scale",
+        "attn.norm_k.weight":           "img_attn.norm.key_norm.scale",
+        "attn.norm_added_q.weight":     "txt_attn.norm.query_norm.scale",
+        "attn.norm_added_k.weight":     "txt_attn.norm.key_norm.scale",
+        "attn.to_out.0.weight":         "img_attn.proj.weight",
+        "attn.to_add_out.weight":       "txt_attn.proj.weight",
+        "ff.net.0.proj.weight":         "img_mlp.0.weight",
+        "ff.net.2.weight":              "img_mlp.2.weight",
+        "ff_context.net.0.proj.weight": "txt_mlp.0.weight",
+        "ff_context.net.2.weight":      "txt_mlp.2.weight",
+        "ff.linear_in.weight":          "img_mlp.0.weight",
+        "ff.linear_out.weight":         "img_mlp.2.weight",
+        "ff_context.linear_in.weight":  "txt_mlp.0.weight",
+        "ff_context.linear_out.weight": "txt_mlp.2.weight",
+    }
+    return dst + mapping.get(tail, tail)
+
+
+# QKV triplets to fuse on output: source tails -> target fused tail.
+# Each tuple is (q_tail, k_tail, v_tail, fused_target_tail).
+QKV_DOUBLE_TRIPLETS = [
+    ("attn.to_q.weight",       "attn.to_k.weight",       "attn.to_v.weight",       "img_attn.qkv.weight"),
+    ("attn.add_q_proj.weight", "attn.add_k_proj.weight", "attn.add_v_proj.weight", "txt_attn.qkv.weight"),
+]
+
+
+def _rewrite_single_stream(idx: str, tail: str) -> str:
+    dst = f"single_blocks.{idx}."
+    mapping = {
+        "norm.linear.weight":          "modulation.lin.weight",
+        "attn.norm_q.weight":          "norm.query_norm.scale",
+        "attn.norm_k.weight":          "norm.key_norm.scale",
+        "attn.to_qkv_mlp_proj.weight": "linear1.weight",
+        "attn.to_out.weight":          "linear2.weight",
+    }
+    return dst + mapping.get(tail, tail)
+
+
+
+
+def read_safetensors_header(path: str):
+    """Return (header dict, data start byte offset)."""
+    with open(path, "rb") as f:
+        hdr_len = struct.unpack("<Q", f.read(8))[0]
+        hdr_bytes = f.read(hdr_len)
+    return json.loads(hdr_bytes), 8 + hdr_len
+
+
+def collect_shard_paths(directory: str, weight_pattern: str):
+    index_path = os.path.join(directory, f"{weight_pattern}.safetensors.index.json")
+    if os.path.isfile(index_path):
+        with open(index_path) as f:
+            idx = json.load(f)
+        return sorted({os.path.join(directory, n) for n in idx["weight_map"].values()})
+    single = os.path.join(directory, f"{weight_pattern}.safetensors")
+    if not os.path.isfile(single):
+        raise FileNotFoundError(f"No checkpoint at {directory}: missing {weight_pattern}")
+    return [single]
+
+
+def stage_tensors_for_section(section_dir: str, rewrite_fn):
+    """Return a list of (new_key, shard_path, data_start_offset, info_dict) entries.
+
+    A "qkv_fuse" pseudo-entry with three source descriptors is emitted when a
+    transformer_blocks.* split q/k/v triplet is found, so the writer can fuse
+    them into a single output tensor.
+    """
+    entries = []
+    # First, index all raw keys per shard so we can detect qkv triplets.
+    raw_by_block = {}  # block_idx -> {tail: (key, shard_path, data_off, info)}
+    raw_others = []
+    for shard_path in collect_shard_paths(section_dir, "diffusion_pytorch_model"):
+        hdr, data_off = read_safetensors_header(shard_path)
+        for key, info in hdr.items():
+            if key == "__metadata__":
+                continue
+            m = re.match(r"backbone\.transformer_blocks\.(\d+)\.(.*)$", key)
+            if m and any(m.group(2) in trip[:3] for trip in QKV_DOUBLE_TRIPLETS):
+                idx = m.group(1)
+                raw_by_block.setdefault(idx, {})[m.group(2)] = (key, shard_path, data_off, info)
+            else:
+                raw_others.append((key, shard_path, data_off, info))
+
+    for key, shard_path, data_off, info in raw_others:
+        new_key = rewrite_fn(key)
+        # Swap the (scale, shift) halves to (shift, scale) at conversion time so
+        # the on-disk weight matches BFL flux ordering and the runtime stays
+        # version-agnostic. norm_out.linear weight shape is [2*dim, dim] and bias
+        # is [2*dim]; both split along axis 0 (outermost == row-major outer).
+        if new_key in ("final_layer.adaLN_modulation.1.weight",
+                       "final_layer.adaLN_modulation.1.bias"):
+            info = dict(info)
+            info["_chunk_swap_halves"] = True
+        entries.append((new_key, shard_path, data_off, info))
+
+    for block_idx, tails in raw_by_block.items():
+        for q_tail, k_tail, v_tail, fused_tail in QKV_DOUBLE_TRIPLETS:
+            if q_tail in tails and k_tail in tails and v_tail in tails:
+                q = tails[q_tail]; k = tails[k_tail]; v = tails[v_tail]
+                # Validate shapes match.
+                q_shape = q[3]["shape"]; k_shape = k[3]["shape"]; v_shape = v[3]["shape"]
+                if q_shape != k_shape or q_shape != v_shape:
+                    raise ValueError(f"qkv shape mismatch at block {block_idx} {q_tail}: q={q_shape} k={k_shape} v={v_shape}")
+                fused_shape = [q_shape[0] * 3] + list(q_shape[1:])
+                fused_info = {
+                    "dtype": q[3]["dtype"],
+                    "shape": fused_shape,
+                    "_qkv_sources": [q, k, v],  # pseudo field consumed by writer
+                }
+                entries.append((f"double_blocks.{block_idx}.{fused_tail}",
+                                None, None, fused_info))
+                del tails[q_tail]; del tails[k_tail]; del tails[v_tail]
+        # Anything left in tails was an unmatched single - pass through.
+        for tail, payload in tails.items():
+            entries.append((rewrite_fn(payload[0]),) + payload[1:])
+    return entries
+
+
+_DTYPE_BYTES = {
+    "BF16": 2, "F16": 2, "F32": 4, "F64": 8,
+    "U8": 1, "I8": 1, "I16": 2, "I32": 4, "I64": 8,
+    "BOOL": 1,
+}
+
+
+def _total_bytes(info: dict) -> int:
+    if "_qkv_sources" in info:
+        elems = 1
+        for d in info["shape"]:
+            elems *= d
+        return elems * _DTYPE_BYTES[info["dtype"]]
+    start, end = info["data_offsets"]
+    return end - start
+
+
+def write_consolidated(out_path: str, entries):
+    """Write a single safetensors file by streaming raw bytes from each shard.
+
+    For qkv-fused entries, q/k/v are concatenated along axis 0 (row-major), so a
+    simple byte-level concatenation produces the correct fused layout for any
+    standard dtype.
+    """
+    entries = sorted(entries, key=lambda e: e[0])
+
+    new_header = {}
+    cur_offset = 0
+    for new_key, shard_path, data_off, info in entries:
+        size = _total_bytes(info)
+        new_header[new_key] = {
+            "dtype": info["dtype"],
+            "shape": info["shape"],
+            "data_offsets": [cur_offset, cur_offset + size],
+        }
+        cur_offset += size
+
+    header_json = json.dumps(new_header, separators=(",", ":")).encode("utf-8")
+    pad = (-len(header_json)) % 8
+    header_json = header_json + (b" " * pad)
+
+    def copy_range(src_path, src_data_off, src_info, out, byte_range=None):
+        start, end = src_info["data_offsets"]
+        if byte_range is not None:
+            sub_start, sub_end = byte_range
+            start, end = start + sub_start, start + sub_end
+        with open(src_path, "rb") as src:
+            src.seek(src_data_off + start)
+            remaining = end - start
+            while remaining > 0:
+                chunk = src.read(min(8 * 1024 * 1024, remaining))
+                if not chunk:
+                    raise IOError(f"Truncated tensor in {src_path}")
+                out.write(chunk)
+                remaining -= len(chunk)
+
+    with open(out_path, "wb") as out:
+        out.write(struct.pack("<Q", len(header_json)))
+        out.write(header_json)
+        for new_key, shard_path, data_off, info in entries:
+            if "_qkv_sources" in info:
+                for q_entry in info["_qkv_sources"]:
+                    _, src_path, src_data_off, src_info = q_entry
+                    copy_range(src_path, src_data_off, src_info, out)
+            elif info.get("_chunk_swap_halves"):
+                size = _total_bytes(info)
+                half = size // 2
+                if size != half * 2:
+                    raise ValueError(f"{new_key}: odd byte size {size} cannot be split into halves")
+                copy_range(shard_path, data_off, info, out, byte_range=(half, size))
+                copy_range(shard_path, data_off, info, out, byte_range=(0, half))
+            else:
+                copy_range(shard_path, data_off, info, out)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("input_dir", help="SeFi diffusers checkpoint directory")
+    parser.add_argument("output", help="Output transformer safetensors path (load via --diffusion-model)")
+    args = parser.parse_args()
+
+    transformer_entries = stage_tensors_for_section(
+        os.path.join(args.input_dir, "transformer"), rewrite_transformer_key)
+
+    print(f"Transformer tensors: {len(transformer_entries)}")
+    print(f"Writing {len(transformer_entries)} tensors -> {args.output}")
+    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+    write_consolidated(args.output, transformer_entries)
+    print(f"Done. Output size: {os.path.getsize(args.output) / 1e9:.2f} GB")
+
+
+if __name__ == "__main__":
+    main()

script/remove_utf8_bom.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""Remove UTF-8 BOMs from files under a directory.
+
+By default this scans the current working directory recursively and skips
+repository areas that should not be touched by ordinary maintenance scripts.
+Only files whose first three bytes are the UTF-8 BOM are rewritten.
+"""
+
+import argparse
+import os
+import shutil
+import sys
+import tempfile
+from pathlib import Path
+
+
+UTF8_BOM = b"\xef\xbb\xbf"
+
+DEFAULT_EXCLUDED_DIR_NAMES = {
+    ".git",
+    ".hg",
+    ".svn",
+    ".mypy_cache",
+    ".pytest_cache",
+    "__pycache__",
+    "test",
+}
+
+DEFAULT_EXCLUDED_DIR_PREFIXES = {
+    "build",
+}
+
+DEFAULT_EXCLUDED_REL_DIRS = {
+    "examples/server/frontend",
+    "ggml",
+    "models",
+    "src/vocab",
+    "thirdparty",
+}
+
+
+def rel_posix(path: Path, root: Path) -> str:
+    try:
+        return path.relative_to(root).as_posix()
+    except ValueError:
+        return path.as_posix()
+
+
+def should_skip_dir(
+    path: Path,
+    root: Path,
+    excluded_rel_dirs: set[str],
+    excluded_names: set[str],
+    excluded_prefixes: set[str],
+) -> bool:
+    rel = rel_posix(path, root)
+    return (
+        path.name in excluded_names
+        or rel in excluded_rel_dirs
+        or any(path.name.startswith(prefix) for prefix in excluded_prefixes)
+    )
+
+
+def iter_files(
+    root: Path,
+    recursive: bool,
+    excluded_rel_dirs: set[str],
+    excluded_names: set[str],
+    excluded_prefixes: set[str],
+    follow_symlinks: bool,
+):
+    if recursive:
+        for dirpath, dirnames, filenames in os.walk(root, followlinks=follow_symlinks):
+            current_dir = Path(dirpath)
+            dirnames[:] = [
+                name
+                for name in dirnames
+                if not should_skip_dir(
+                    current_dir / name,
+                    root,
+                    excluded_rel_dirs,
+                    excluded_names,
+                    excluded_prefixes,
+                )
+            ]
+            for filename in filenames:
+                path = current_dir / filename
+                if path.is_symlink() and not follow_symlinks:
+                    continue
+                yield path
+    else:
+        for path in root.iterdir():
+            if path.is_file() and (follow_symlinks or not path.is_symlink()):
+                yield path
+
+
+def has_utf8_bom(path: Path) -> bool:
+    with path.open("rb") as f:
+        return f.read(len(UTF8_BOM)) == UTF8_BOM
+
+
+def strip_utf8_bom(path: Path) -> None:
+    tmp_path = None
+    try:
+        with path.open("rb") as src:
+            if src.read(len(UTF8_BOM)) != UTF8_BOM:
+                return
+
+            fd, tmp_name = tempfile.mkstemp(
+                prefix=f".{path.name}.",
+                suffix=".tmp",
+                dir=str(path.parent),
+            )
+            tmp_path = Path(tmp_name)
+            with os.fdopen(fd, "wb") as dst:
+                shutil.copyfileobj(src, dst, length=1024 * 1024)
+
+        shutil.copystat(path, tmp_path, follow_symlinks=False)
+        os.replace(tmp_path, path)
+        tmp_path = None
+    finally:
+        if tmp_path is not None:
+            try:
+                tmp_path.unlink()
+            except FileNotFoundError:
+                pass
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Scan files and convert UTF-8 BOM files to UTF-8 without BOM.",
+    )
+    parser.add_argument(
+        "root",
+        nargs="?",
+        default=".",
+        help="Directory to scan. Defaults to the current directory.",
+    )
+    parser.add_argument(
+        "-n",
+        "--dry-run",
+        action="store_true",
+        help="Only list files that would be converted.",
+    )
+    parser.add_argument(
+        "--no-recursive",
+        action="store_true",
+        help="Only scan files directly under root.",
+    )
+    parser.add_argument(
+        "--include-repo-excluded",
+        action="store_true",
+        help="Do not skip default repository excluded directories.",
+    )
+    parser.add_argument(
+        "--exclude-dir",
+        action="append",
+        default=[],
+        metavar="DIR",
+        help="Additional directory name or root-relative path to skip. Can be used multiple times.",
+    )
+    parser.add_argument(
+        "--follow-symlinks",
+        action="store_true",
+        help="Follow symlinked directories and files.",
+    )
+    parser.add_argument(
+        "-q",
+        "--quiet",
+        action="store_true",
+        help="Only print the final summary.",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    root = Path(args.root).resolve()
+
+    if not root.is_dir():
+        print(f"error: not a directory: {root}", file=sys.stderr)
+        return 2
+
+    excluded_names = set()
+    excluded_rel_dirs = set()
+    excluded_prefixes = set()
+    if not args.include_repo_excluded:
+        excluded_names.update(DEFAULT_EXCLUDED_DIR_NAMES)
+        excluded_rel_dirs.update(DEFAULT_EXCLUDED_REL_DIRS)
+        excluded_prefixes.update(DEFAULT_EXCLUDED_DIR_PREFIXES)
+
+    for item in args.exclude_dir:
+        normalized = Path(item).as_posix().strip("/")
+        if "/" in normalized:
+            excluded_rel_dirs.add(normalized)
+        else:
+            excluded_names.add(normalized)
+
+    scanned = 0
+    converted = 0
+    errors = 0
+
+    for path in iter_files(
+        root=root,
+        recursive=not args.no_recursive,
+        excluded_rel_dirs=excluded_rel_dirs,
+        excluded_names=excluded_names,
+        excluded_prefixes=excluded_prefixes,
+        follow_symlinks=args.follow_symlinks,
+    ):
+        scanned += 1
+        try:
+            if not has_utf8_bom(path):
+                continue
+            converted += 1
+            rel = rel_posix(path, root)
+            if args.dry_run:
+                if not args.quiet:
+                    print(f"would convert: {rel}")
+            else:
+                strip_utf8_bom(path)
+                if not args.quiet:
+                    print(f"converted: {rel}")
+        except OSError as exc:
+            errors += 1
+            print(f"error: {rel_posix(path, root)}: {exc}", file=sys.stderr)
+
+    action = "would convert" if args.dry_run else "converted"
+    print(f"scanned {scanned} file(s), {action} {converted}, errors {errors}")
+    return 1 if errors else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/conditioning/conditioner.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
+#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
 #define __SD_CONDITIONING_CONDITIONER_HPP__
 
 #include <cmath>
@@ -1378,6 +1378,101 @@ struct T5CLIPEmbedder : public Conditioner {
     }
 };
 
+struct MiniT2IConditioner : public Conditioner {
+    T5UniGramTokenizer tokenizer;
+    std::shared_ptr<T5Runner> t5;
+    size_t prompt_length = 256;
+
+    MiniT2IConditioner(ggml_backend_t backend,
+                       const String2TensorStorage& tensor_storage_map      = {},
+                       std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
+        bool use_t5 = false;
+        for (const auto& pair : tensor_storage_map) {
+            if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
+                use_t5 = true;
+                break;
+            }
+        }
+        if (!use_t5) {
+            LOG_WARN("IMPORTANT NOTICE: No MiniT2I T5 text encoder provided, cannot process prompts!");
+            return;
+        }
+        t5 = std::make_shared<T5Runner>(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager);
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        if (t5) {
+            t5->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
+    void set_stream_layers_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_stream_layers_enabled(enabled);
+        }
+    }
+
+    void set_flash_attention_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_flash_attention_enabled(enabled);
+        }
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        if (t5) {
+            t5->set_weight_adapter(adapter);
+        }
+    }
+
+    void runner_done() override {
+        if (t5) {
+            t5->runner_done();
+        }
+    }
+
+    SDCondition get_learned_condition(int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        SDCondition result;
+        if (!t5) {
+            result.c_crossattn = sd::Tensor<float>::zeros({1024, static_cast<int64_t>(prompt_length)});
+            result.c_vector    = sd::Tensor<float>::zeros({static_cast<int64_t>(prompt_length)});
+            return result;
+        }
+
+        std::vector<int> tokens = tokenizer.encode(conditioner_params.text);
+        if (tokens.size() > prompt_length) {
+            tokens.resize(prompt_length);
+        }
+        std::vector<float> mask(tokens.size(), 1.0f);
+        while (tokens.size() < prompt_length) {
+            tokens.push_back(tokenizer.PAD_TOKEN_ID);
+            mask.push_back(0.0f);
+        }
+
+        sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
+        std::vector<float> t5_mask(mask.size(), 0.0f);
+        for (size_t i = 0; i < mask.size(); ++i) {
+            t5_mask[i] = mask[i] > 0.0f ? 0.0f : -HUGE_VALF;
+        }
+        sd::Tensor<float> hidden_states = t5->compute(n_threads,
+                                                      input_ids,
+                                                      sd::Tensor<float>::from_vector(t5_mask),
+                                                      false,
+                                                      true,
+                                                      true);
+        GGML_ASSERT(!hidden_states.empty());
+        result.c_crossattn = std::move(hidden_states);
+        result.c_vector    = sd::Tensor<float>::from_vector(mask);
+        return result;
+    }
+};
+
 struct AnimaConditioner : public Conditioner {
     std::shared_ptr<BPETokenizer> qwen_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
@@ -1518,7 +1613,7 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_ideogram4(version)) {
+        } else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
@@ -1778,6 +1873,76 @@ struct LLMEmbedder : public Conditioner {
 
                 prompt += "<|im_end|>\n<|im_start|>assistant\n";
             }
+        } else if (sd_version_is_boogu_image(version)) {
+            prompt_template_encode_start_idx = 0;
+
+            const std::string t2i_system_prompt =
+                "You are a helpful assistant that generates high-quality images based on user instructions. The instructions are as follows.";
+            const std::string edit_system_prompt =
+                "Describe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.";
+            const bool has_ref_images = llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty();
+            const bool text_empty     = conditioner_params.text.find_first_not_of(" \t\r\n") == std::string::npos;
+
+            if (has_ref_images) {
+                LOG_INFO("BooguImageEditPipeline");
+                const std::string prompt_prefix = "<|im_start|>system\n" + edit_system_prompt + "<|im_end|>\n<|im_start|>user\n";
+                std::string img_prompt;
+                const std::string placeholder = "<|image_pad|>";
+
+                for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
+                    const auto& image = (*conditioner_params.ref_images)[i];
+                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    int height        = static_cast<int>(image.shape()[1]);
+                    int width         = static_cast<int>(image.shape()[0]);
+                    double beta       = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
+                    int h_bar         = std::max(static_cast<int>(factor),
+                                                 static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                    int w_bar         = std::max(static_cast<int>(factor),
+                                                 static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+
+                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+
+                    auto resized_image = clip_preprocess(image, w_bar, h_bar);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    GGML_ASSERT(!image_embed.empty());
+
+                    std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
+                    int image_embed_idx      = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
+                    image_embeds.emplace_back(image_embed_idx, image_embed);
+
+                    img_prompt += "<|vision_start|>";
+                    int64_t num_image_tokens = image_embed.shape()[1];
+                    img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
+                    for (int j = 0; j < num_image_tokens; j++) {
+                        img_prompt += placeholder;
+                    }
+                    img_prompt += "<|vision_end|>";
+                }
+
+                prompt                  = prompt_prefix + img_prompt;
+                prompt_attn_range.first = static_cast<int>(prompt.size());
+                prompt += conditioner_params.text;
+                prompt_attn_range.second = static_cast<int>(prompt.size());
+                prompt += "<|im_end|>\n";
+            } else {
+                const std::string& system_prompt = text_empty ? edit_system_prompt : t2i_system_prompt;
+                prompt                           = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user\n";
+                prompt_attn_range.first          = static_cast<int>(prompt.size());
+                prompt += conditioner_params.text;
+                prompt_attn_range.second = static_cast<int>(prompt.size());
+                prompt += "<|im_end|>\n";
+            }
+        } else if (sd_version_is_krea2(version)) {
+            prompt_template_encode_start_idx = 34;
+            out_layers                       = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
+
+            prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (sd_version_is_longcat(version)) {
             spell_quotes = true;
 
@@ -1927,6 +2092,18 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
+        } else if (sd_version_is_sefi_image(version)) {
+            prompt_template_encode_start_idx = 0;
+            min_length                       = 1024;
+            out_layers                       = {9, 18, 27};
+
+            prompt = "<|im_start|>user\n";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (version == VERSION_OVIS_IMAGE) {
             prompt_template_encode_start_idx = 28;
             min_length                       = prompt_template_encode_start_idx + 256;

src/convert.cpp

@@ -76,29 +76,22 @@ static bool load_tensors_for_export(ModelLoader& model_loader,
     return success;
 }
 
-bool convert(const char* input_path,
-             const char* vae_path,
-             const char* output_path,
-             sd_type_t output_type,
-             const char* tensor_type_rules,
-             bool convert_name) {
-    ModelLoader model_loader;
-
-    if (!model_loader.init_from_file(input_path)) {
-        LOG_ERROR("init model loader from file failed: '%s'", input_path);
+static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
+    if (path == nullptr || strlen(path) == 0) {
+        return true;
+    }
+    if (!model_loader.init_from_file(path, prefix)) {
+        LOG_ERROR("init model loader from file failed: '%s'", path);
         return false;
     }
+    loaded_any = true;
+    return true;
+}
 
-    if (vae_path != nullptr && strlen(vae_path) > 0) {
-        if (!model_loader.init_from_file(vae_path, "vae.")) {
-            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
-            return false;
-        }
-    }
-    if (convert_name) {
-        model_loader.convert_tensors_name();
-    }
-
+static bool export_loaded_model(ModelLoader& model_loader,
+                                const char* output_path,
+                                sd_type_t output_type,
+                                const char* tensor_type_rules) {
     ggml_type type             = sd_type_to_ggml_type(output_type);
     bool output_is_safetensors = ends_with(output_path, ".safetensors");
     TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
@@ -136,3 +129,55 @@ bool convert(const char* input_path,
     ggml_free(ggml_ctx);
     return success;
 }
+
+bool convert_with_components(const char* model_path,
+                             const char* clip_l_path,
+                             const char* clip_g_path,
+                             const char* t5xxl_path,
+                             const char* diffusion_model_path,
+                             const char* vae_path,
+                             const char* output_path,
+                             sd_type_t output_type,
+                             const char* tensor_type_rules,
+                             bool convert_name) {
+    ModelLoader model_loader;
+    bool loaded_any = false;
+
+    if (!init_convert_path(model_loader, model_path, "", loaded_any) ||
+        !init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) ||
+        !init_convert_path(model_loader, vae_path, "vae.", loaded_any)) {
+        return false;
+    }
+
+    if (!loaded_any) {
+        LOG_ERROR("no input model path provided for convert");
+        return false;
+    }
+
+    if (convert_name) {
+        model_loader.convert_tensors_name();
+    }
+
+    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules);
+}
+
+bool convert(const char* input_path,
+             const char* vae_path,
+             const char* output_path,
+             sd_type_t output_type,
+             const char* tensor_type_rules,
+             bool convert_name) {
+    return convert_with_components(input_path,
+                                   nullptr,
+                                   nullptr,
+                                   nullptr,
+                                   nullptr,
+                                   vae_path,
+                                   output_path,
+                                   output_type,
+                                   tensor_type_rules,
+                                   convert_name);
+}

src/core/ggml_extend.hpp

@@ -391,7 +391,7 @@ __STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t*
     int64_t width    = input->ne[0];
     int64_t height   = input->ne[1];
     int64_t channels = input->ne[2];
-    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    GGML_ASSERT(input->type == GGML_TYPE_F32);
     if (image_data == nullptr) {
         image_data = (uint8_t*)malloc(width * height * channels);
     }
@@ -1038,6 +1038,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx,
 }
 
 __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                                                ggml_backend_t backend,
                                                 ggml_tensor* x,
                                                 int lp0,
                                                 int rp0,
@@ -1063,7 +1064,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
     }
 
     if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
-        x = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
+            x = padded;
+        } else {
+            // Some backends (e.g. Metal) only implement right-padding for
+            // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
+            // the padding around to the left. shift < ne always holds because
+            // ne grew by lp+rp.
+            x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
+            x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
+        }
     }
     return x;
 }
@@ -1076,7 +1087,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
                                             int p3          = 0,
                                             bool circular_x = false,
                                             bool circular_y = false) {
-    return ggml_ext_pad_ext(ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
+    return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
 }
 
 // w: [OC，IC, KH, KW]
@@ -1105,7 +1116,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
     }
 
     if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
-        x  = ggml_ext_pad_ext(ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
+        x  = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
         p0 = 0;
         p1 = 0;
     }
@@ -1130,6 +1141,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
 // b: [OC,]
 // result: [N*OC, OD, OH, OW]
 __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                                                ggml_backend_t backend,
                                                 ggml_tensor* x,
                                                 ggml_tensor* w,
                                                 ggml_tensor* b,
@@ -1159,7 +1171,21 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
         x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
     } else {
-        x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
+        // (e.g. Metal, see #850) do not implement. Fall back to
+        // GGML_OP_CONV_3D on those backends.
+        bool im2col_3d_supported = true;
+        if (backend != nullptr) {
+            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
+        }
+        if (im2col_3d_supported) {
+            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        } else {
+            int64_t OC = w->ne[3] / IC;
+            int64_t N  = x->ne[3] / IC;
+            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+        }
     }
 
     if (b != nullptr) {
@@ -1382,7 +1408,16 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
             if (!ggml_backend_supports_op(backend, kqv)) {
                 kqv = nullptr;
             } else {
-                kqv = ggml_view_3d(ctx, kqv, d_head, n_head, L_q, kqv->nb[1], kqv->nb[2], 0);
+                kqv = ggml_view_4d(ctx,
+                                   kqv,
+                                   d_head,
+                                   n_head,
+                                   L_q,
+                                   N,
+                                   kqv->nb[1],
+                                   kqv->nb[2],
+                                   kqv->nb[1] * n_head,
+                                   0);
             }
         }
     }
@@ -2464,7 +2499,10 @@ protected:
             sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
         }
 
-        ggml_status status = ggml_backend_graph_compute(runtime_backend, gf);
+        ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
+                                                                         gf,
+                                                                         sd_get_backend_eval_callback(),
+                                                                         sd_get_backend_eval_callback_data());
         if (status != GGML_STATUS_SUCCESS) {
             LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
             return std::nullopt;
@@ -3356,7 +3394,7 @@ public:
                 b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
             }
         }
-        return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
+        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                 std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                 std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
                                 std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),

src/core/ggml_extend_backend.cpp

@@ -9,6 +9,7 @@
 #include <vector>
 
 #include "core/util.h"
+#include "ggml/src/ggml-impl.h"
 #include "stable-diffusion.h"
 
 static std::string trim_copy(const std::string& value) {
@@ -110,7 +111,67 @@ static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type)
     if (dev == nullptr) {
         return "";
     }
-    return ggml_backend_dev_name(dev);
+    const char* dev_name = ggml_backend_dev_name(dev);
+    if (dev_name != nullptr && dev_name[0] != '\0') {
+        return dev_name;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    const char* reg_name   = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
+    return reg_name != nullptr ? reg_name : "";
+}
+
+static ggml_backend_dev_t resolve_first_device_by_registry_name(const std::string& name) {
+    std::string lower = lower_copy(trim_copy(name));
+    if (lower == "metal") {
+        lower = "mtl";
+    }
+    if (lower.empty()) {
+        return nullptr;
+    }
+
+    const size_t device_count = ggml_backend_dev_count();
+    for (size_t i = 0; i < device_count; ++i) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+        if (reg == nullptr) {
+            continue;
+        }
+        const char* reg_name = ggml_backend_reg_name(reg);
+        if (reg_name != nullptr && lower_copy(reg_name) == lower) {
+            return dev;
+        }
+    }
+    return nullptr;
+}
+
+static ggml_backend_dev_t resolve_device_by_name(const std::string& name) {
+    const std::string lower = lower_copy(trim_copy(name));
+    if (lower.empty()) {
+        return nullptr;
+    }
+
+    const size_t device_count = ggml_backend_dev_count();
+    for (size_t i = 0; i < device_count; ++i) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        const char* dev_name   = ggml_backend_dev_name(dev);
+        if (dev_name != nullptr && lower_copy(dev_name) == lower) {
+            return dev;
+        }
+    }
+    return nullptr;
+}
+
+static std::string backend_device_name(ggml_backend_dev_t dev) {
+    if (dev == nullptr) {
+        return "";
+    }
+    const char* name = ggml_backend_dev_name(dev);
+    if (name != nullptr && name[0] != '\0') {
+        return name;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    const char* reg_name   = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
+    return reg_name != nullptr ? reg_name : "";
 }
 
 static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) {
@@ -296,6 +357,10 @@ std::string sd_backend_resolve_name(const std::string& name) {
         return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
     }
 
+    if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(requested)) {
+        return backend_device_name(dev);
+    }
+
     const size_t device_count = ggml_backend_dev_count();
     for (size_t i = 0; i < device_count; ++i) {
         ggml_backend_dev_t dev = ggml_backend_dev_get(i);
@@ -328,7 +393,20 @@ static ggml_backend_t init_named_backend(const std::string& name) {
         return ggml_backend_init_best();
     }
 
+    if (ggml_backend_dev_t dev = resolve_device_by_name(name)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
+    if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(name)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
+
     std::string resolved = sd_backend_resolve_name(name);
+    if (ggml_backend_dev_t dev = resolve_device_by_name(resolved)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
+    if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(resolved)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
     if (resolved.empty()) {
         return nullptr;
     }
@@ -364,6 +442,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
     return false;
 }
 
+static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
+    ggml_cgraph cgraph = {
+        /*.size             =*/0,
+        /*.n_nodes          =*/i1 - i0,
+        /*.n_leafs          =*/0,
+        /*.nodes            =*/cgraph0->nodes + i0,
+        /*.grads            =*/nullptr,
+        /*.grad_accs        =*/nullptr,
+        /*.leafs            =*/nullptr,
+        /*.use_counts       =*/cgraph0->use_counts,
+        /*.visited_hash_set =*/cgraph0->visited_hash_set,
+        /*.order            =*/cgraph0->order,
+        /*.uid              =*/0,
+    };
+    return cgraph;
+}
+
+ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
+                                                        ggml_cgraph* gf,
+                                                        sd_graph_eval_callback_t callback_eval,
+                                                        void* callback_eval_user_data) {
+    if (callback_eval == nullptr) {
+        return ggml_backend_graph_compute(backend, gf);
+    }
+
+    ggml_status status = GGML_STATUS_SUCCESS;
+    const int n_nodes  = ggml_graph_n_nodes(gf);
+    bool stopped       = false;
+
+    for (int j0 = 0; j0 < n_nodes; ++j0) {
+        ggml_tensor* t = ggml_graph_node(gf, j0);
+        bool need      = callback_eval(t, true, callback_eval_user_data);
+        int j1         = j0;
+
+        while (!need && j1 < n_nodes - 1) {
+            t    = ggml_graph_node(gf, ++j1);
+            need = callback_eval(t, true, callback_eval_user_data);
+        }
+
+        ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
+        status         = ggml_backend_graph_compute_async(backend, &gv);
+        if (status != GGML_STATUS_SUCCESS) {
+            break;
+        }
+
+        ggml_backend_synchronize(backend);
+
+        if (need && !callback_eval(t, false, callback_eval_user_data)) {
+            stopped = true;
+            break;
+        }
+
+        j0 = j1;
+    }
+
+    ggml_backend_synchronize(backend);
+    if (stopped && status == GGML_STATUS_SUCCESS) {
+        status = GGML_STATUS_ABORTED;
+    }
+    return status;
+}
+
 const char* sd_get_system_info() {
     static std::string cache_info = []() -> std::string {
         ggml_backend_load_all_once();
@@ -599,7 +739,7 @@ bool SDBackendManager::validate(std::string* error) const {
             }
             return false;
         }
-        if (!sd_backend_resolve_name(name).empty()) {
+        if (!sd_backend_resolve_name(name).empty() || resolve_first_device_by_registry_name(name) != nullptr) {
             return true;
         }
         if (error != nullptr) {

src/core/ggml_extend_backend.h

@@ -9,6 +9,7 @@
 
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "stable-diffusion.h"
 
 enum class SDBackendModule {
     DIFFUSION,
@@ -71,6 +72,10 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
+ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
+                                                        ggml_cgraph* gf,
+                                                        sd_graph_eval_callback_t callback_eval,
+                                                        void* callback_eval_user_data);
 std::string sd_backend_resolve_name(const std::string& name);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);

src/core/util.cpp

@@ -346,6 +346,9 @@ int sd_preview_interval              = 1;
 bool sd_preview_denoised             = true;
 bool sd_preview_noisy                = false;
 
+static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr;
+static void* sd_backend_eval_cb_data               = nullptr;
+
 std::u32string utf8_to_utf32(const std::string& utf8_str) {
     std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
     return converter.from_bytes(utf8_str);
@@ -629,6 +632,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b
     sd_preview_noisy    = noisy;
 }
 
+void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) {
+    sd_backend_eval_cb      = cb;
+    sd_backend_eval_cb_data = data;
+}
+
 sd_preview_cb_t sd_get_preview_callback() {
     return sd_preview_cb;
 }
@@ -649,6 +657,14 @@ bool sd_should_preview_noisy() {
     return sd_preview_noisy;
 }
 
+sd_graph_eval_callback_t sd_get_backend_eval_callback() {
+    return sd_backend_eval_cb;
+}
+
+void* sd_get_backend_eval_callback_data() {
+    return sd_backend_eval_cb_data;
+}
+
 sd_progress_cb_t sd_get_progress_callback() {
     return sd_progress_cb;
 }

src/core/util.h

@@ -98,6 +98,9 @@ int sd_get_preview_interval();
 bool sd_should_preview_denoised();
 bool sd_should_preview_noisy();
 
+sd_graph_eval_callback_t sd_get_backend_eval_callback();
+void* sd_get_backend_eval_callback_data();
+
 // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 

src/model.h

@@ -36,18 +36,23 @@ enum SDVersion {
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
     VERSION_QWEN_IMAGE,
+    VERSION_QWEN_IMAGE_LAYERED,
     VERSION_ANIMA,
     VERSION_FLUX2,
     VERSION_FLUX2_KLEIN,
     VERSION_LTXAV,
     VERSION_HIDREAM_O1,
     VERSION_Z_IMAGE,
+    VERSION_BOOGU_IMAGE,
     VERSION_OVIS_IMAGE,
     VERSION_ERNIE_IMAGE,
     VERSION_LENS,
+    VERSION_MINIT2I,
     VERSION_LONGCAT,
     VERSION_PID,
     VERSION_IDEOGRAM4,
+    VERSION_SEFI_IMAGE,
+    VERSION_KREA2,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -123,7 +128,7 @@ static inline bool sd_version_is_wan(SDVersion version) {
 }
 
 static inline bool sd_version_is_qwen_image(SDVersion version) {
-    if (version == VERSION_QWEN_IMAGE) {
+    if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
         return true;
     }
     return false;
@@ -143,6 +148,13 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_boogu_image(SDVersion version) {
+    if (version == VERSION_BOOGU_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_longcat(SDVersion version) {
     if (version == VERSION_LONGCAT) {
         return true;
@@ -164,6 +176,13 @@ static inline bool sd_version_is_lens(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_minit2i(SDVersion version) {
+    if (version == VERSION_MINIT2I) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_pid(SDVersion version) {
     if (version == VERSION_PID) {
         return true;
@@ -178,8 +197,36 @@ static inline bool sd_version_is_ideogram4(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_sefi_image(SDVersion version) {
+    if (version == VERSION_SEFI_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_is_krea2(SDVersion version) {
+    if (version == VERSION_KREA2) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_uses_flux_vae(SDVersion version) {
+    if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
+        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) || sd_version_is_lens(version) || sd_version_is_ideogram4(version)) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_uses_wan_vae(SDVersion version) {
+    if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
         return true;
     }
     return false;
@@ -206,11 +253,15 @@ static inline bool sd_version_is_dit(SDVersion version) {
         version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||
         sd_version_is_z_image(version) ||
+        sd_version_is_boogu_image(version) ||
         sd_version_is_ernie_image(version) ||
         sd_version_is_lens(version) ||
+        sd_version_is_minit2i(version) ||
         sd_version_is_longcat(version) ||
         sd_version_is_pid(version) ||
-        sd_version_is_ideogram4(version)) {
+        sd_version_is_ideogram4(version) ||
+        sd_version_is_sefi_image(version) ||
+        sd_version_is_krea2(version)) {
         return true;
     }
     return false;

src/model/common/rope.hpp

@@ -12,6 +12,16 @@ namespace Rope {
         ErnieImage,
     };
 
+    enum class RefIndexMode {
+        FIXED,
+        INCREASE,
+        DECREASE,
+    };
+
+    __STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
+        return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
+    }
+
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
@@ -346,7 +356,7 @@ namespace Rope {
                                                                    int axes_dim_num,
                                                                    int start_index,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
-                                                                   bool increase_ref_index,
+                                                                   RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool scale_rope,
                                                                    int base_offset = 0) {
@@ -357,13 +367,15 @@ namespace Rope {
         for (ggml_tensor* ref : ref_latents) {
             int h_offset = 0;
             int w_offset = 0;
-            if (!increase_ref_index) {
+            if (ref_index_mode == RefIndexMode::FIXED) {
                 if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
                     w_offset = curr_w_offset;
                 } else {
                     h_offset = curr_h_offset;
                 }
                 scale_rope = false;
+            } else if (ref_index_mode == RefIndexMode::DECREASE) {
+                index--;
             }
 
             auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
@@ -377,7 +389,7 @@ namespace Rope {
                                             scale_rope);
             ids          = concat_ids(ids, ref_ids, bs);
 
-            if (increase_ref_index) {
+            if (ref_index_mode == RefIndexMode::INCREASE) {
                 index++;
             }
 
@@ -395,7 +407,7 @@ namespace Rope {
                                                                    int context_len,
                                                                    std::set<int> txt_arange_dims,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
-                                                                   bool increase_ref_index,
+                                                                   RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool is_longcat) {
         int x_index = is_longcat ? 1 : 0;
@@ -406,7 +418,7 @@ namespace Rope {
 
         auto ids = concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -420,7 +432,7 @@ namespace Rope {
                                                      int context_len,
                                                      std::set<int> txt_arange_dims,
                                                      const std::vector<ggml_tensor*>& ref_latents,
-                                                     bool increase_ref_index,
+                                                     RefIndexMode ref_index_mode,
                                                      float ref_index_scale,
                                                      int theta,
                                                      bool circular_h,
@@ -435,7 +447,7 @@ namespace Rope {
                                                            context_len,
                                                            txt_arange_dims,
                                                            ref_latents,
-                                                           increase_ref_index,
+                                                           ref_index_mode,
                                                            ref_index_scale,
                                                            is_longcat);
         std::vector<std::vector<int>> wrap_dims;
@@ -481,17 +493,64 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
+                                                                  int h,
+                                                                  int w,
+                                                                  int pt,
+                                                                  int ph,
+                                                                  int pw,
+                                                                  int bs,
+                                                                  int t_offset    = 0,
+                                                                  int h_offset    = 0,
+                                                                  int w_offset    = 0,
+                                                                  bool scale_rope = false) {
+        int t_len = (t + (pt / 2)) / pt;
+        int h_len = (h + (ph / 2)) / ph;
+        int w_len = (w + (pw / 2)) / pw;
+
+        std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
+
+        if (scale_rope) {
+            h_offset -= h_len / 2;
+            w_offset -= w_len / 2;
+        }
+
+        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
+        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
+        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
+
+        for (int i = 0; i < t_len; ++i) {
+            for (int j = 0; j < h_len; ++j) {
+                for (int k = 0; k < w_len; ++k) {
+                    int idx         = i * h_len * w_len + j * w_len + k;
+                    vid_ids[idx][0] = t_ids[i];
+                    vid_ids[idx][1] = h_ids[j];
+                    vid_ids[idx][2] = w_ids[k];
+                }
+            }
+        }
+
+        std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < vid_ids.size(); ++j) {
+                vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
+            }
+        }
+        return vid_ids_repeated;
+    }
+
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
+                                                                         int h,
                                                                          int w,
                                                                          int patch_size,
                                                                          int bs,
                                                                          int context_len,
                                                                          const std::vector<ggml_tensor*>& ref_latents,
-                                                                         bool increase_ref_index) {
+                                                                         RefIndexMode ref_index_mode) {
         int h_len        = (h + (patch_size / 2)) / patch_size;
         int w_len        = (w + (patch_size / 2)) / patch_size;
-        int txt_id_start = std::max(h_len, w_len);
-        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
+        int txt_id_start = std::max(h_len, w_len) / 2;
+        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * txt_id_start + context_len - 1, context_len);
         std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
         for (int i = 0; i < bs; ++i) {
             for (int j = 0; j < txt_ids.size(); ++j) {
@@ -499,28 +558,30 @@ namespace Rope {
             }
         }
         int axes_dim_num = 3;
-        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
+        auto img_ids     = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
         auto ids         = concat_ids(txt_ids_repeated, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
-            ids           = concat_ids(ids, refs_ids, bs);
+            int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
+            auto refs_ids       = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
+            ids                 = concat_ids(ids, refs_ids, bs);
         }
         return ids;
     }
 
     // Generate qwen_image positional embeddings
-    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
+    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
+                                                           int h,
                                                            int w,
                                                            int patch_size,
                                                            int bs,
                                                            int context_len,
                                                            const std::vector<ggml_tensor*>& ref_latents,
-                                                           bool increase_ref_index,
+                                                           RefIndexMode ref_index_mode,
                                                            int theta,
                                                            bool circular_h,
                                                            bool circular_w,
                                                            const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
+        std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
         std::vector<std::vector<int>> wrap_dims;
         // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
@@ -533,7 +594,7 @@ namespace Rope {
                 // Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
                 wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
                 size_t cursor           = context_len;  // ignore text tokens
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
+                const size_t img_tokens = static_cast<size_t>(t) * 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][cursor + token_i] = h_len;
@@ -684,46 +745,6 @@ namespace Rope {
         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,
-                                                                  int pt,
-                                                                  int ph,
-                                                                  int pw,
-                                                                  int bs,
-                                                                  int t_offset = 0,
-                                                                  int h_offset = 0,
-                                                                  int w_offset = 0) {
-        int t_len = (t + (pt / 2)) / pt;
-        int h_len = (h + (ph / 2)) / ph;
-        int w_len = (w + (pw / 2)) / pw;
-
-        std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
-
-        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
-        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
-        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
-
-        for (int i = 0; i < t_len; ++i) {
-            for (int j = 0; j < h_len; ++j) {
-                for (int k = 0; k < w_len; ++k) {
-                    int idx         = i * h_len * w_len + j * w_len + k;
-                    vid_ids[idx][0] = t_ids[i];
-                    vid_ids[idx][1] = h_ids[j];
-                    vid_ids[idx][2] = w_ids[k];
-                }
-            }
-        }
-
-        std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
-        for (int i = 0; i < bs; ++i) {
-            for (int j = 0; j < vid_ids.size(); ++j) {
-                vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
-            }
-        }
-        return vid_ids_repeated;
-    }
-
     // Generate wan positional embeddings
     __STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
                                                     int h,
@@ -785,7 +806,8 @@ namespace Rope {
                                                                       int context_len,
                                                                       int seq_multi_of,
                                                                       const std::vector<ggml_tensor*>& ref_latents,
-                                                                      bool increase_ref_index) {
+                                                                      RefIndexMode ref_index_mode) {
+        SD_UNUSED(ref_index_mode);
         int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
         auto txt_ids           = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
         for (int i = 0; i < bs * padded_context_len; i++) {
@@ -816,12 +838,12 @@ namespace Rope {
                                                         int context_len,
                                                         int seq_multi_of,
                                                         const std::vector<ggml_tensor*>& ref_latents,
-                                                        bool increase_ref_index,
+                                                        RefIndexMode ref_index_mode,
                                                         int theta,
                                                         bool circular_h,
                                                         bool circular_w,
                                                         const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
+        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
         std::vector<std::vector<int>> wrap_dims;
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
             int pad_h = (patch_size - (h % patch_size)) % patch_size;
@@ -899,10 +921,12 @@ namespace Rope {
         // q,k,v: [N, L, n_head, d_head]
         // pe: [L, d_head/2, 2, 2]
         // return: [N, L, n_head*d_head]
+        int64_t n_head = q->ne[1];
+
         q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved);  // [N*n_head, L, d_head]
         k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved);  // [N*n_head, L, d_head]
 
-        auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, v->ne[1], mask, true, ctx->flash_attn_enabled, kv_scale);  // [N, L, n_head*d_head]
+        auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale);  // [N, L, n_head*d_head]
         return x;
     }
 };  // namespace Rope

src/model/diffusion/anima.hpp

@@ -227,7 +227,6 @@ namespace Anima {
             k4 = k_norm->forward(ctx, k4);
 
             ggml_tensor* attn_out = nullptr;
-            float scale           = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
             if (pe_q != nullptr || pe_k != nullptr) {
                 if (pe_q == nullptr) {
                     pe_q = pe_k;
@@ -245,8 +244,7 @@ namespace Anima {
                                                      num_heads,
                                                      nullptr,
                                                      true,
-                                                     ctx->flash_attn_enabled,
-                                                     scale);
+                                                     ctx->flash_attn_enabled);
             } else {
                 auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
                 auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
@@ -258,8 +256,7 @@ namespace Anima {
                                                      num_heads,
                                                      nullptr,
                                                      false,
-                                                     ctx->flash_attn_enabled,
-                                                     scale);
+                                                     ctx->flash_attn_enabled);
             }
 
             return out_proj->forward(ctx, attn_out);
@@ -615,7 +612,7 @@ namespace Anima {
                                           0,
                                           {},
                                           empty_ref_latents,
-                                          false,
+                                          Rope::RefIndexMode::FIXED,
                                           1.0f,
                                           false);
 

src/model/diffusion/boogu.hpp

@@ -0,0 +1,835 @@
+#ifndef __SD_MODEL_DIFFUSION_BOOGU_HPP__
+#define __SD_MODEL_DIFFUSION_BOOGU_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <tuple>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/qwen_image.hpp"
+#include "model_loader.h"
+
+namespace Boogu {
+    constexpr int BOOGU_GRAPH_SIZE = 65536;
+
+    struct BooguConfig {
+        int patch_size                   = 2;
+        int64_t in_channels              = 16;
+        int64_t out_channels             = 16;
+        int64_t hidden_size              = 3360;
+        int64_t num_layers               = 32;
+        int64_t num_double_stream_layers = 8;
+        int64_t num_refiner_layers       = 2;
+        int64_t num_attention_heads      = 28;
+        int64_t num_kv_heads             = 7;
+        int64_t head_dim                 = 120;
+        int64_t multiple_of              = 256;
+        int64_t instruction_feat_dim     = 4096;
+        int64_t timestep_embed_dim       = 1024;
+        int theta                        = 10000;
+        float timestep_scale             = 1000.0f;
+        float norm_eps                   = 1e-5f;
+        std::vector<int> axes_dim        = {40, 40, 40};
+        int64_t axes_dim_sum             = 120;
+
+        static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
+                                    const std::string& prefix,
+                                    const std::string& block_prefix) {
+            int64_t count = 0;
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                size_t pos = name.find(block_prefix);
+                if (pos == std::string::npos) {
+                    continue;
+                }
+                auto items = split_string(name.substr(pos), '.');
+                if (items.size() > 1) {
+                    count = std::max<int64_t>(count, atoi(items[1].c_str()) + 1);
+                }
+            }
+            return count;
+        }
+
+        static BooguConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            BooguConfig config;
+            int64_t detected_head_dim = 0;
+            int64_t detected_kv_dim   = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area = config.patch_size * config.patch_size;
+                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) {
+                    config.instruction_feat_dim = tensor_storage.ne[0];
+                    config.hidden_size          = tensor_storage.ne[1];
+                } else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
+                    detected_head_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "double_stream_layers.0.img_self_attn.norm_q.weight") && tensor_storage.n_dims == 1) {
+                    detected_head_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "single_stream_layers.0.attn.to_k.weight") && tensor_storage.n_dims == 2) {
+                    detected_kv_dim = tensor_storage.ne[1];
+                } else if (ends_with(name, "double_stream_layers.0.img_instruct_attn.processor.img_to_k.weight") && tensor_storage.n_dims == 2) {
+                    detected_kv_dim = tensor_storage.ne[1];
+                } else if (ends_with(name, "norm_out.linear_2.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area  = config.patch_size * config.patch_size;
+                    config.out_channels = tensor_storage.ne[1] / patch_area;
+                }
+            }
+
+            config.num_layers               = std::max<int64_t>(1, count_blocks(tensor_storage_map, prefix, "single_stream_layers."));
+            config.num_double_stream_layers = std::max<int64_t>(0, count_blocks(tensor_storage_map, prefix, "double_stream_layers."));
+            int64_t noise_refiner_layers    = count_blocks(tensor_storage_map, prefix, "noise_refiner.");
+            int64_t ref_refiner_layers      = count_blocks(tensor_storage_map, prefix, "ref_image_refiner.");
+            int64_t context_refiner_layers  = count_blocks(tensor_storage_map, prefix, "context_refiner.");
+            config.num_refiner_layers       = std::max<int64_t>(1, std::max(noise_refiner_layers, std::max(ref_refiner_layers, context_refiner_layers)));
+
+            if (detected_head_dim > 0) {
+                config.head_dim            = detected_head_dim;
+                config.num_attention_heads = config.hidden_size / config.head_dim;
+                config.axes_dim_sum        = config.head_dim;
+                if (detected_kv_dim > 0) {
+                    config.num_kv_heads = detected_kv_dim / config.head_dim;
+                }
+                if (config.axes_dim_sum == 120) {
+                    config.axes_dim = {40, 40, 40};
+                } else if (config.axes_dim_sum % 3 == 0) {
+                    int axis        = static_cast<int>(config.axes_dim_sum / 3);
+                    config.axes_dim = {axis, axis, axis};
+                }
+            }
+            config.timestep_embed_dim = std::min<int64_t>(config.hidden_size, 1024);
+
+            LOG_DEBUG("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64,
+                      config.num_layers,
+                      config.num_double_stream_layers,
+                      config.num_refiner_layers,
+                      config.hidden_size,
+                      config.num_attention_heads,
+                      config.num_kv_heads,
+                      config.head_dim,
+                      config.in_channels,
+                      config.out_channels);
+            return config;
+        }
+    };
+
+    __STATIC_INLINE__ ggml_tensor* scale_modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
+        scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);
+        return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
+    }
+
+    __STATIC_INLINE__ ggml_tensor* gate_residual(ggml_context* ctx, ggml_tensor* residual, ggml_tensor* x, ggml_tensor* gate) {
+        gate = ggml_tanh(ctx, gate);
+        gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]);
+        x    = ggml_mul(ctx, x, gate);
+        return ggml_add(ctx, residual, x);
+    }
+
+    struct LuminaCombinedTimestepCaptionEmbedding : public GGMLBlock {
+        int64_t frequency_embedding_size;
+        float timestep_scale;
+
+        LuminaCombinedTimestepCaptionEmbedding(int64_t hidden_size,
+                                               int64_t instruction_feat_dim,
+                                               int64_t frequency_embedding_size,
+                                               float norm_eps,
+                                               float timestep_scale)
+            : frequency_embedding_size(frequency_embedding_size),
+              timestep_scale(timestep_scale) {
+            blocks["timestep_embedder"]  = std::make_shared<Qwen::TimestepEmbedding>(frequency_embedding_size, std::min<int64_t>(hidden_size, 1024));
+            blocks["caption_embedder.0"] = std::make_shared<RMSNorm>(instruction_feat_dim, norm_eps);
+            blocks["caption_embedder.1"] = std::make_shared<Linear>(instruction_feat_dim, hidden_size, true);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* timestep, ggml_tensor* text_hidden_states) {
+            auto timestep_embedder  = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
+            auto caption_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["caption_embedder.0"]);
+            auto caption_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["caption_embedder.1"]);
+
+            auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(frequency_embedding_size), 10000, timestep_scale);
+            auto time_embed    = timestep_embedder->forward(ctx, timestep_proj);
+            auto caption_embed = caption_embedder_1->forward(ctx, caption_embedder_0->forward(ctx, text_hidden_states));
+            return {time_embed, caption_embed};
+        }
+    };
+
+    struct LuminaRMSNormZero : public GGMLBlock {
+        LuminaRMSNormZero(int64_t embedding_dim, int64_t conditioning_embedding_dim, float norm_eps) {
+            blocks["linear"] = std::make_shared<Linear>(conditioning_embedding_dim, 4 * embedding_dim, true);
+            blocks["norm"]   = std::make_shared<RMSNorm>(embedding_dim, norm_eps);
+        }
+
+        std::tuple<ggml_tensor*, ggml_tensor*, ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb) {
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            auto norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
+
+            emb       = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));
+            auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 4, 0);
+
+            auto scale_msa = mods[0];
+            auto gate_msa  = mods[1];
+            auto scale_mlp = mods[2];
+            auto gate_mlp  = mods[3];
+
+            x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), scale_msa);
+            return {x, gate_msa, scale_mlp, gate_mlp};
+        }
+    };
+
+    struct LuminaFeedForward : public GGMLBlock {
+        LuminaFeedForward(int64_t dim, int64_t inner_dim, int64_t multiple_of) {
+            inner_dim          = multiple_of * ((inner_dim + multiple_of - 1) / multiple_of);
+            blocks["linear_1"] = std::make_shared<Linear>(dim, inner_dim, false);
+            blocks["linear_2"] = std::make_shared<Linear>(inner_dim, dim, false);
+            blocks["linear_3"] = std::make_shared<Linear>(dim, inner_dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
+            auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
+
+            if (sd_backend_is(ctx->backend, "Vulkan")) {
+                linear_2->set_force_prec_f32(true);
+            }
+
+            auto h1 = linear_1->forward(ctx, x);
+            auto h2 = linear_3->forward(ctx, x);
+            x       = ggml_swiglu_split(ctx->ggml_ctx, h1, h2);
+            x       = linear_2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct LuminaLayerNormContinuous : public GGMLBlock {
+        LuminaLayerNormContinuous(int64_t embedding_dim,
+                                  int64_t conditioning_embedding_dim,
+                                  int64_t out_dim) {
+            blocks["linear_1"] = std::make_shared<Linear>(conditioning_embedding_dim, embedding_dim, true);
+            blocks["norm"]     = std::make_shared<LayerNorm>(embedding_dim, 1e-6f, false);
+            blocks["linear_2"] = std::make_shared<Linear>(embedding_dim, out_dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning_embedding) {
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
+            auto norm     = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
+
+            auto emb = linear_1->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning_embedding));
+            x        = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), emb);
+            x        = linear_2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct Attention : public GGMLBlock {
+        int64_t dim_head;
+        int64_t heads;
+        int64_t kv_heads;
+
+        Attention(int64_t query_dim, int64_t dim_head, int64_t heads, int64_t kv_heads, float eps = 1e-5f)
+            : dim_head(dim_head), heads(heads), kv_heads(kv_heads) {
+            blocks["to_q"]     = std::make_shared<Linear>(query_dim, heads * dim_head, false);
+            blocks["to_k"]     = std::make_shared<Linear>(query_dim, kv_heads * dim_head, false);
+            blocks["to_v"]     = std::make_shared<Linear>(query_dim, kv_heads * dim_head, 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>(heads * dim_head, query_dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* hidden_states,
+                             ggml_tensor* encoder_hidden_states,
+                             ggml_tensor* rotary_emb,
+                             ggml_tensor* attention_mask = nullptr) {
+            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"]);
+
+            if (sd_backend_is(ctx->backend, "Vulkan")) {
+                to_out_0->set_force_prec_f32(true);
+            }
+
+            int64_t N  = hidden_states->ne[2];
+            int64_t Lq = hidden_states->ne[1];
+            int64_t Lk = encoder_hidden_states->ne[1];
+
+            auto q = to_q->forward(ctx, hidden_states);
+            q      = ggml_reshape_4d(ctx->ggml_ctx, q, dim_head, heads, Lq, N);
+            auto k = to_k->forward(ctx, encoder_hidden_states);
+            k      = ggml_reshape_4d(ctx->ggml_ctx, k, dim_head, kv_heads, Lk, N);
+            auto v = to_v->forward(ctx, encoder_hidden_states);
+            v      = ggml_reshape_4d(ctx->ggml_ctx, v, dim_head, kv_heads, Lk, N);
+
+            q = norm_q->forward(ctx, q);
+            k = norm_k->forward(ctx, k);
+
+            auto out = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask);
+            out      = to_out_0->forward(ctx, out);
+            return out;
+        }
+    };
+
+    struct BooguImageTransformerBlock : public GGMLBlock {
+        bool modulation;
+
+        BooguImageTransformerBlock(int64_t dim,
+                                   int64_t num_attention_heads,
+                                   int64_t num_kv_heads,
+                                   int64_t multiple_of,
+                                   float norm_eps,
+                                   bool modulation)
+            : modulation(modulation) {
+            int64_t head_dim       = dim / num_attention_heads;
+            blocks["attn"]         = std::make_shared<Attention>(dim, head_dim, num_attention_heads, num_kv_heads, 1e-5f);
+            blocks["feed_forward"] = std::make_shared<LuminaFeedForward>(dim, 4 * dim, multiple_of);
+            if (modulation) {
+                blocks["norm1"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
+            } else {
+                blocks["norm1"] = std::make_shared<RMSNorm>(dim, norm_eps);
+            }
+            blocks["ffn_norm1"] = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["norm2"]     = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["ffn_norm2"] = std::make_shared<RMSNorm>(dim, norm_eps);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* hidden_states,
+                             ggml_tensor* rotary_emb,
+                             ggml_tensor* temb           = nullptr,
+                             ggml_tensor* attention_mask = nullptr) {
+            auto attn         = std::dynamic_pointer_cast<Attention>(blocks["attn"]);
+            auto feed_forward = std::dynamic_pointer_cast<LuminaFeedForward>(blocks["feed_forward"]);
+            auto ffn_norm1    = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
+            auto norm2        = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
+            auto ffn_norm2    = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
+
+            if (modulation) {
+                auto norm1 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["norm1"]);
+                auto mods  = norm1->forward(ctx, hidden_states, temb);
+
+                auto norm_hidden_states = std::get<0>(mods);
+                auto gate_msa           = std::get<1>(mods);
+                auto scale_mlp          = std::get<2>(mods);
+                auto gate_mlp           = std::get<3>(mods);
+
+                auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask);
+                hidden_states    = gate_residual(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output), gate_msa);
+
+                auto mlp_input  = scale_modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp);
+                auto mlp_output = feed_forward->forward(ctx, mlp_input);
+                hidden_states   = gate_residual(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output), gate_mlp);
+            } else {
+                auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
+
+                auto norm_hidden_states = norm1->forward(ctx, hidden_states);
+                auto attn_output        = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask);
+                hidden_states           = ggml_add(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output));
+
+                auto mlp_output = feed_forward->forward(ctx, ffn_norm1->forward(ctx, hidden_states));
+                hidden_states   = ggml_add(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output));
+            }
+            return hidden_states;
+        }
+    };
+
+    struct BooguImageJointAttention : public GGMLBlock {
+        int64_t dim_head;
+        int64_t heads;
+        int64_t kv_heads;
+
+        BooguImageJointAttention(int64_t dim, int64_t dim_head, int64_t heads, int64_t kv_heads)
+            : dim_head(dim_head), heads(heads), kv_heads(kv_heads) {
+            blocks["norm_q"]                  = std::make_shared<RMSNorm>(dim_head, 1e-5f);
+            blocks["norm_k"]                  = std::make_shared<RMSNorm>(dim_head, 1e-5f);
+            blocks["to_out.0"]                = std::make_shared<Linear>(heads * dim_head, dim, false);
+            blocks["processor.img_to_q"]      = std::make_shared<Linear>(dim, heads * dim_head, false);
+            blocks["processor.img_to_k"]      = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
+            blocks["processor.img_to_v"]      = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
+            blocks["processor.instruct_to_q"] = std::make_shared<Linear>(dim, heads * dim_head, false);
+            blocks["processor.instruct_to_k"] = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
+            blocks["processor.instruct_to_v"] = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
+            blocks["processor.instruct_out"]  = std::make_shared<Linear>(heads * dim_head, dim, false);
+            blocks["processor.img_out"]       = std::make_shared<Linear>(heads * dim_head, dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* img_hidden_states,
+                             ggml_tensor* instruct_hidden_states,
+                             ggml_tensor* rotary_emb,
+                             ggml_tensor* attention_mask = nullptr) {
+            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"]);
+            auto img_to_q      = std::dynamic_pointer_cast<Linear>(blocks["processor.img_to_q"]);
+            auto img_to_k      = std::dynamic_pointer_cast<Linear>(blocks["processor.img_to_k"]);
+            auto img_to_v      = std::dynamic_pointer_cast<Linear>(blocks["processor.img_to_v"]);
+            auto instruct_to_q = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_to_q"]);
+            auto instruct_to_k = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_to_k"]);
+            auto instruct_to_v = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_to_v"]);
+            auto instruct_out  = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
+            auto img_out       = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
+
+            if (sd_backend_is(ctx->backend, "Vulkan")) {
+                to_out_0->set_force_prec_f32(true);
+            }
+
+            int64_t N          = img_hidden_states->ne[2];
+            int64_t L_img      = img_hidden_states->ne[1];
+            int64_t L_instruct = instruct_hidden_states->ne[1];
+
+            auto img_q = img_to_q->forward(ctx, img_hidden_states);
+            img_q      = ggml_reshape_4d(ctx->ggml_ctx, img_q, dim_head, heads, L_img, N);
+            auto img_k = img_to_k->forward(ctx, img_hidden_states);
+            img_k      = ggml_reshape_4d(ctx->ggml_ctx, img_k, dim_head, kv_heads, L_img, N);
+            auto img_v = img_to_v->forward(ctx, img_hidden_states);
+            img_v      = ggml_reshape_4d(ctx->ggml_ctx, img_v, dim_head, kv_heads, L_img, N);
+
+            auto instruct_q = instruct_to_q->forward(ctx, instruct_hidden_states);
+            instruct_q      = ggml_reshape_4d(ctx->ggml_ctx, instruct_q, dim_head, heads, L_instruct, N);
+            auto instruct_k = instruct_to_k->forward(ctx, instruct_hidden_states);
+            instruct_k      = ggml_reshape_4d(ctx->ggml_ctx, instruct_k, dim_head, kv_heads, L_instruct, N);
+            auto instruct_v = instruct_to_v->forward(ctx, instruct_hidden_states);
+            instruct_v      = ggml_reshape_4d(ctx->ggml_ctx, instruct_v, dim_head, kv_heads, L_instruct, N);
+
+            auto q = ggml_concat(ctx->ggml_ctx, instruct_q, img_q, 2);
+            auto k = ggml_concat(ctx->ggml_ctx, instruct_k, img_k, 2);
+            auto v = ggml_concat(ctx->ggml_ctx, instruct_v, img_v, 2);
+            q      = norm_q->forward(ctx, q);
+            k      = norm_k->forward(ctx, k);
+
+            auto hidden_states = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask);
+            auto instruct_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, L_instruct);
+            auto img_attn      = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, L_instruct, L_instruct + L_img);
+
+            instruct_attn = instruct_out->forward(ctx, instruct_attn);
+            img_attn      = img_out->forward(ctx, img_attn);
+            hidden_states = ggml_concat(ctx->ggml_ctx, instruct_attn, img_attn, 1);
+            hidden_states = to_out_0->forward(ctx, hidden_states);
+            return hidden_states;
+        }
+    };
+
+    struct BooguImageDoubleStreamBlock : public GGMLBlock {
+        BooguImageDoubleStreamBlock(int64_t dim,
+                                    int64_t num_attention_heads,
+                                    int64_t num_kv_heads,
+                                    int64_t multiple_of,
+                                    float norm_eps) {
+            int64_t head_dim                = dim / num_attention_heads;
+            blocks["img_instruct_attn"]     = std::make_shared<BooguImageJointAttention>(dim, head_dim, num_attention_heads, num_kv_heads);
+            blocks["img_self_attn"]         = std::make_shared<Attention>(dim, head_dim, num_attention_heads, num_kv_heads, 1e-5f);
+            blocks["img_feed_forward"]      = std::make_shared<LuminaFeedForward>(dim, 4 * dim, multiple_of);
+            blocks["instruct_feed_forward"] = std::make_shared<LuminaFeedForward>(dim, 4 * dim, multiple_of);
+            blocks["img_norm1"]             = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
+            blocks["img_norm2"]             = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
+            blocks["img_norm3"]             = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
+            blocks["instruct_norm1"]        = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
+            blocks["instruct_norm2"]        = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
+            blocks["img_attn_norm"]         = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["img_self_attn_norm"]    = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["img_ffn_norm1"]         = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["img_ffn_norm2"]         = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["instruct_attn_norm"]    = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["instruct_ffn_norm1"]    = std::make_shared<RMSNorm>(dim, norm_eps);
+            blocks["instruct_ffn_norm2"]    = std::make_shared<RMSNorm>(dim, norm_eps);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* img_hidden_states,
+                                                      ggml_tensor* instruct_hidden_states,
+                                                      ggml_tensor* joint_rotary_emb,
+                                                      ggml_tensor* img_rotary_emb,
+                                                      ggml_tensor* temb) {
+            auto img_instruct_attn     = std::dynamic_pointer_cast<BooguImageJointAttention>(blocks["img_instruct_attn"]);
+            auto img_self_attn         = std::dynamic_pointer_cast<Attention>(blocks["img_self_attn"]);
+            auto img_feed_forward      = std::dynamic_pointer_cast<LuminaFeedForward>(blocks["img_feed_forward"]);
+            auto instruct_feed_forward = std::dynamic_pointer_cast<LuminaFeedForward>(blocks["instruct_feed_forward"]);
+            auto img_norm1             = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["img_norm1"]);
+            auto img_norm2             = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["img_norm2"]);
+            auto img_norm3             = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["img_norm3"]);
+            auto instruct_norm1        = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["instruct_norm1"]);
+            auto instruct_norm2        = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["instruct_norm2"]);
+            auto img_attn_norm         = std::dynamic_pointer_cast<RMSNorm>(blocks["img_attn_norm"]);
+            auto img_self_attn_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["img_self_attn_norm"]);
+            auto img_ffn_norm1         = std::dynamic_pointer_cast<RMSNorm>(blocks["img_ffn_norm1"]);
+            auto img_ffn_norm2         = std::dynamic_pointer_cast<RMSNorm>(blocks["img_ffn_norm2"]);
+            auto instruct_attn_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["instruct_attn_norm"]);
+            auto instruct_ffn_norm1    = std::dynamic_pointer_cast<RMSNorm>(blocks["instruct_ffn_norm1"]);
+            auto instruct_ffn_norm2    = std::dynamic_pointer_cast<RMSNorm>(blocks["instruct_ffn_norm2"]);
+
+            int64_t L_instruct = instruct_hidden_states->ne[1];
+
+            auto img_norm1_out_vec      = img_norm1->forward(ctx, img_hidden_states, temb);
+            auto img_norm2_out_vec      = img_norm2->forward(ctx, img_hidden_states, temb);
+            auto img_norm3_out_vec      = img_norm3->forward(ctx, img_hidden_states, temb);
+            auto instruct_norm1_out_vec = instruct_norm1->forward(ctx, instruct_hidden_states, temb);
+            auto instruct_norm2_out_vec = instruct_norm2->forward(ctx, instruct_hidden_states, temb);
+
+            auto img_norm1_out = std::get<0>(img_norm1_out_vec);
+            auto img_gate_msa  = std::get<1>(img_norm1_out_vec);
+            auto img_scale_mlp = std::get<2>(img_norm1_out_vec);
+            auto img_gate_mlp  = std::get<3>(img_norm1_out_vec);
+
+            auto img_norm2_out = std::get<0>(img_norm2_out_vec);
+            auto img_shift_mlp = std::get<1>(img_norm2_out_vec);
+
+            auto img_norm3_out = std::get<0>(img_norm3_out_vec);
+            auto img_gate_self = std::get<1>(img_norm3_out_vec);
+
+            auto instruct_norm1_out = std::get<0>(instruct_norm1_out_vec);
+            auto instruct_gate_msa  = std::get<1>(instruct_norm1_out_vec);
+            auto instruct_scale_mlp = std::get<2>(instruct_norm1_out_vec);
+            auto instruct_gate_mlp  = std::get<3>(instruct_norm1_out_vec);
+
+            auto instruct_norm2_out = std::get<0>(instruct_norm2_out_vec);
+            auto instruct_shift_mlp = std::get<1>(instruct_norm2_out_vec);
+
+            auto joint_attn_out    = img_instruct_attn->forward(ctx, img_norm1_out, instruct_norm1_out, joint_rotary_emb);
+            auto instruct_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, 0, L_instruct);
+            auto img_attn_out      = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]);
+
+            auto img_self_attn_out = img_self_attn->forward(ctx, img_norm3_out, img_norm3_out, img_rotary_emb);
+
+            img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_attn_norm->forward(ctx, img_attn_out), img_gate_msa);
+            img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_self_attn_norm->forward(ctx, img_self_attn_out), img_gate_self);
+
+            auto img_mlp_input = scale_modulate(ctx->ggml_ctx, img_norm2_out, img_scale_mlp);
+            img_shift_mlp      = ggml_reshape_3d(ctx->ggml_ctx, img_shift_mlp, img_shift_mlp->ne[0], 1, img_shift_mlp->ne[1]);
+            img_mlp_input      = ggml_add(ctx->ggml_ctx, img_mlp_input, img_shift_mlp);
+            auto img_mlp_out   = img_feed_forward->forward(ctx, img_ffn_norm1->forward(ctx, img_mlp_input));
+            img_hidden_states  = gate_residual(ctx->ggml_ctx, img_hidden_states, img_ffn_norm2->forward(ctx, img_mlp_out), img_gate_mlp);
+
+            instruct_hidden_states  = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_attn_norm->forward(ctx, instruct_attn_out), instruct_gate_msa);
+            auto instruct_mlp_input = scale_modulate(ctx->ggml_ctx, instruct_norm2_out, instruct_scale_mlp);
+            instruct_shift_mlp      = ggml_reshape_3d(ctx->ggml_ctx, instruct_shift_mlp, instruct_shift_mlp->ne[0], 1, instruct_shift_mlp->ne[1]);
+            instruct_mlp_input      = ggml_add(ctx->ggml_ctx, instruct_mlp_input, instruct_shift_mlp);
+            auto instruct_mlp_out   = instruct_feed_forward->forward(ctx, instruct_ffn_norm1->forward(ctx, instruct_mlp_input));
+            instruct_hidden_states  = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_ffn_norm2->forward(ctx, instruct_mlp_out), instruct_gate_mlp);
+
+            return {img_hidden_states, instruct_hidden_states};
+        }
+    };
+
+    struct BooguImageModel : public GGMLBlock {
+        BooguConfig config;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            GGML_UNUSED(tensor_storage_map);
+            GGML_UNUSED(prefix);
+            params["image_index_embedding"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, 5);
+        }
+
+        BooguImageModel() = default;
+        BooguImageModel(BooguConfig config)
+            : config(std::move(config)) {
+            blocks["x_embedder"]               = std::make_shared<Linear>(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true);
+            blocks["ref_image_patch_embedder"] = std::make_shared<Linear>(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true);
+            blocks["time_caption_embed"]       = std::make_shared<LuminaCombinedTimestepCaptionEmbedding>(this->config.hidden_size,
+                                                                                                    this->config.instruction_feat_dim,
+                                                                                                    256,
+                                                                                                    this->config.norm_eps,
+                                                                                                    this->config.timestep_scale);
+
+            for (int i = 0; i < this->config.num_refiner_layers; i++) {
+                blocks["noise_refiner." + std::to_string(i)]     = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
+                                                                                                            this->config.num_attention_heads,
+                                                                                                            this->config.num_kv_heads,
+                                                                                                            this->config.multiple_of,
+                                                                                                            this->config.norm_eps,
+                                                                                                            true);
+                blocks["ref_image_refiner." + std::to_string(i)] = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
+                                                                                                                this->config.num_attention_heads,
+                                                                                                                this->config.num_kv_heads,
+                                                                                                                this->config.multiple_of,
+                                                                                                                this->config.norm_eps,
+                                                                                                                true);
+                blocks["context_refiner." + std::to_string(i)]   = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
+                                                                                                              this->config.num_attention_heads,
+                                                                                                              this->config.num_kv_heads,
+                                                                                                              this->config.multiple_of,
+                                                                                                              this->config.norm_eps,
+                                                                                                              false);
+            }
+
+            for (int i = 0; i < this->config.num_double_stream_layers; i++) {
+                blocks["double_stream_layers." + std::to_string(i)] = std::make_shared<BooguImageDoubleStreamBlock>(this->config.hidden_size,
+                                                                                                                    this->config.num_attention_heads,
+                                                                                                                    this->config.num_kv_heads,
+                                                                                                                    this->config.multiple_of,
+                                                                                                                    this->config.norm_eps);
+            }
+
+            for (int i = 0; i < this->config.num_layers; i++) {
+                blocks["single_stream_layers." + std::to_string(i)] = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
+                                                                                                                   this->config.num_attention_heads,
+                                                                                                                   this->config.num_kv_heads,
+                                                                                                                   this->config.multiple_of,
+                                                                                                                   this->config.norm_eps,
+                                                                                                                   true);
+            }
+
+            blocks["norm_out"] = std::make_shared<LuminaLayerNormContinuous>(this->config.hidden_size,
+                                                                             this->config.timestep_embed_dim,
+                                                                             this->config.patch_size * this->config.patch_size * this->config.out_channels);
+        }
+
+        ggml_tensor* image_index_embedding(GGMLRunnerContext* ctx, int index) {
+            GGML_ASSERT(index >= 0 && index < 5);
+            auto embedding = params["image_index_embedding"];
+            auto out       = ggml_view_1d(ctx->ggml_ctx,
+                                          embedding,
+                                          config.hidden_size,
+                                          index * config.hidden_size * ggml_element_size(embedding));
+            out            = ggml_reshape_3d(ctx->ggml_ctx, out, config.hidden_size, 1, 1);
+            return out;
+        }
+
+        ggml_tensor* embed_refs(GGMLRunnerContext* ctx, const std::vector<ggml_tensor*>& ref_latents) {
+            if (ref_latents.empty()) {
+                return nullptr;
+            }
+            auto ref_image_patch_embedder = std::dynamic_pointer_cast<Linear>(blocks["ref_image_patch_embedder"]);
+
+            ggml_tensor* ref_img = nullptr;
+            for (int i = 0; i < static_cast<int>(ref_latents.size()); i++) {
+                auto ref = DiT::pad_and_patchify(ctx, ref_latents[i], config.patch_size, config.patch_size, false);
+                ref      = ref_image_patch_embedder->forward(ctx, ref);
+                ref      = ggml_add(ctx->ggml_ctx, ref, image_index_embedding(ctx, std::min(i, 4)));
+                ref_img  = ref_img == nullptr ? ref : ggml_concat(ctx->ggml_ctx, ref_img, ref, 1);
+            }
+            return ref_img;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* context,
+                             ggml_tensor* pe,
+                             std::vector<ggml_tensor*> ref_latents = {}) {
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t N = x->ne[3];
+            GGML_ASSERT(N == 1);
+
+            auto x_embedder         = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
+            auto time_caption_embed = std::dynamic_pointer_cast<LuminaCombinedTimestepCaptionEmbedding>(blocks["time_caption_embed"]);
+            auto norm_out           = std::dynamic_pointer_cast<LuminaLayerNormContinuous>(blocks["norm_out"]);
+
+            auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, timesteps), timesteps);
+            auto embeds   = time_caption_embed->forward(ctx, timestep, context);
+            auto temb     = embeds.first;
+            auto txt      = embeds.second;
+
+            auto img        = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, false);
+            int64_t img_len = img->ne[1];
+            img             = x_embedder->forward(ctx, img);
+            auto ref_img    = embed_refs(ctx, ref_latents);
+            int64_t ref_len = ref_img != nullptr ? ref_img->ne[1] : 0;
+            int64_t txt_len = txt->ne[1];
+
+            GGML_ASSERT(pe->ne[3] == txt_len + ref_len + img_len);
+            auto txt_pe   = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt_len);
+            auto noise_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len);
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
+                txt        = block->forward(ctx, txt, txt_pe);
+                sd::ggml_graph_cut::mark_graph_cut(txt, "boogu.context_refiner." + std::to_string(i), "txt");
+            }
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
+                img        = block->forward(ctx, img, noise_pe, temb);
+                sd::ggml_graph_cut::mark_graph_cut(img, "boogu.noise_refiner." + std::to_string(i), "img");
+            }
+
+            ggml_tensor* combined_img = img;
+            if (ref_img != nullptr) {
+                auto ref_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + ref_len);
+                for (int i = 0; i < config.num_refiner_layers; i++) {
+                    auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["ref_image_refiner." + std::to_string(i)]);
+                    ref_img    = block->forward(ctx, ref_img, ref_pe, temb);
+                    sd::ggml_graph_cut::mark_graph_cut(ref_img, "boogu.ref_image_refiner." + std::to_string(i), "ref_img");
+                }
+                combined_img = ggml_concat(ctx->ggml_ctx, ref_img, img, 1);
+            }
+
+            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]);
+            for (int i = 0; i < config.num_double_stream_layers; i++) {
+                auto block   = std::dynamic_pointer_cast<BooguImageDoubleStreamBlock>(blocks["double_stream_layers." + std::to_string(i)]);
+                auto result  = block->forward(ctx, combined_img, txt, pe, img_pe, temb);
+                combined_img = result.first;
+                txt          = result.second;
+                sd::ggml_graph_cut::mark_graph_cut(combined_img, "boogu.double_stream_layers." + std::to_string(i), "img");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "boogu.double_stream_layers." + std::to_string(i), "txt");
+            }
+
+            auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, combined_img, 1);
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block    = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["single_stream_layers." + std::to_string(i)]);
+                hidden_states = block->forward(ctx, hidden_states, pe, temb);
+                sd::ggml_graph_cut::mark_graph_cut(hidden_states, "boogu.single_stream_layers." + std::to_string(i), "hidden_states");
+            }
+
+            hidden_states = norm_out->forward(ctx, hidden_states, temb);
+            hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]);
+            hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, false);
+            hidden_states = ggml_ext_scale(ctx->ggml_ctx, hidden_states, -1.f);
+            return hidden_states;
+        }
+    };
+
+    __STATIC_INLINE__ int patched_token_count(int64_t size, int patch_size) {
+        int pad = (patch_size - (static_cast<int>(size) % patch_size)) % patch_size;
+        return (static_cast<int>(size) + pad) / patch_size;
+    }
+
+    __STATIC_INLINE__ void append_spatial_ids(std::vector<std::vector<float>>& ids,
+                                              int bs,
+                                              int pe_shift,
+                                              int h_tokens,
+                                              int w_tokens) {
+        std::vector<std::vector<float>> image_ids(h_tokens * w_tokens, std::vector<float>(3, 0.0f));
+        for (int h = 0; h < h_tokens; h++) {
+            for (int w = 0; w < w_tokens; w++) {
+                image_ids[h * w_tokens + w][0] = static_cast<float>(pe_shift);
+                image_ids[h * w_tokens + w][1] = static_cast<float>(h);
+                image_ids[h * w_tokens + w][2] = static_cast<float>(w);
+            }
+        }
+        for (int b = 0; b < bs; b++) {
+            ids.insert(ids.end(), image_ids.begin(), image_ids.end());
+        }
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
+                                                      int w,
+                                                      int patch_size,
+                                                      int bs,
+                                                      int context_len,
+                                                      const std::vector<ggml_tensor*>& ref_latents,
+                                                      int theta,
+                                                      const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids;
+        ids.reserve(static_cast<size_t>(bs) * context_len);
+        for (int b = 0; b < bs; b++) {
+            for (int i = 0; i < context_len; i++) {
+                float pos = static_cast<float>(i);
+                ids.push_back({pos, pos, pos});
+            }
+        }
+
+        int pe_shift = context_len;
+        for (ggml_tensor* ref : ref_latents) {
+            int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
+            int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
+            append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
+            pe_shift += std::max(ref_h_tokens, ref_w_tokens);
+        }
+
+        int h_tokens = patched_token_count(h, patch_size);
+        int w_tokens = patched_token_count(w, patch_size);
+        append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
+
+        return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+    }
+
+    struct BooguImageRunner : public DiffusionModelRunner {
+        BooguConfig config;
+        BooguImageModel boogu;
+        std::vector<float> pe_vec;
+
+        BooguImageRunner(ggml_backend_t backend,
+                         const String2TensorStorage& tensor_storage_map      = {},
+                         const std::string prefix                            = "",
+                         SDVersion version                                   = VERSION_BOOGU_IMAGE,
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(BooguConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            boogu = BooguImageModel(config);
+            boogu.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "boogu_image";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            boogu.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,
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
+            ggml_cgraph* gf        = new_graph_custom(BOOGU_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);
+
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
+            }
+
+            pe_vec      = gen_boogu_pe(static_cast<int>(x->ne[1]),
+                                       static_cast<int>(x->ne[0]),
+                                       config.patch_size,
+                                       static_cast<int>(x->ne[3]),
+                                       static_cast<int>(context->ne[1]),
+                                       ref_latents,
+                                       config.theta,
+                                       config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
+            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,
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {}) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, ref_latents);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
+        }
+    };
+}  // namespace Boogu
+
+#endif  // __SD_MODEL_DIFFUSION_BOOGU_HPP__

src/model/diffusion/control.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"

src/model/diffusion/dit.hpp

@@ -135,23 +135,23 @@ namespace DiT {
         return x;
     }
 
-    inline ggml_tensor* unpatchify(ggml_context* ctx,
-                                   ggml_tensor* x,
-                                   int64_t t_len,
-                                   int64_t h_len,
-                                   int64_t w_len,
-                                   int pt,
-                                   int ph,
-                                   int pw) {
-        // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
+    inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
+                                      ggml_tensor* x,
+                                      int64_t t_len,
+                                      int64_t h_len,
+                                      int64_t w_len,
+                                      int pt,
+                                      int ph,
+                                      int pw) {
+        // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
         // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
-        int64_t N = x->ne[3];
+        int64_t N = x->ne[2];
         int64_t C = x->ne[0] / pt / ph / pw;
 
         GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-        x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, C, pt*ph*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
         x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
         x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
         x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]

src/model/diffusion/ernie_image.hpp

@@ -162,8 +162,6 @@ namespace ErnieImage {
             int64_t S = x->ne[1];
             int64_t N = x->ne[2];
 
-            float scale = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
-
             auto q = to_q->forward(ctx, x);
             auto k = to_k->forward(ctx, x);
             auto v = to_v->forward(ctx, x);
@@ -184,7 +182,7 @@ namespace ErnieImage {
             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, scale);  // [N, S, hidden_size]
+            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;
         }

src/model/diffusion/flux.hpp

@@ -8,6 +8,7 @@
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
+#include "model/diffusion/sefi_image.hpp"
 #include "model_loader.h"
 
 #define FLUX_GRAPH_SIZE 10240
@@ -26,6 +27,9 @@ namespace Flux {
     struct FluxConfig {
         SDVersion version         = VERSION_FLUX;
         bool is_chroma            = false;
+        bool is_sefi              = false;
+        int64_t semantic_channels = 0;
+        float sefi_delta_t        = 0.1f;
         int patch_size            = 2;
         int64_t in_channels       = 64;
         int64_t out_channels      = 64;
@@ -88,6 +92,21 @@ namespace Flux {
                 config.share_modulation = true;
                 config.ref_index_scale  = 10.f;
                 config.use_mlp_silu_act = true;
+            } else if (sd_version_is_sefi_image(version)) {
+                config.is_sefi           = true;
+                config.semantic_channels = 16;
+                config.in_channels       = 128 + config.semantic_channels;
+                config.patch_size        = 1;
+                config.out_channels      = 128 + config.semantic_channels;
+                config.mlp_ratio         = 3.f;
+                config.theta             = 2000;
+                config.axes_dim          = {32, 32, 32, 32};
+                config.vec_in_dim        = 0;
+                config.qkv_bias          = false;
+                config.disable_bias      = true;
+                config.share_modulation  = true;
+                config.ref_index_scale   = 10.f;
+                config.use_mlp_silu_act  = true;
             } else if (sd_version_is_longcat(version)) {
                 config.context_in_dim = 3584;
                 config.vec_in_dim     = 0;
@@ -723,8 +742,8 @@ namespace Flux {
 
                 auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 2 * hidden_size]
                 auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0);
-                shift      = m_vec[0];  // [N, hidden_size]
-                scale      = m_vec[1];  // [N, hidden_size]
+                shift      = m_vec[0];
+                scale      = m_vec[1];
             }
 
             x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
@@ -902,6 +921,8 @@ namespace Flux {
             }
             if (config.is_chroma) {
                 blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(config.in_dim, config.hidden_size);
+            } else if (config.is_sefi) {
+                blocks["dual_time_embed"] = std::make_shared<SefiImage::SefiDualTimestepEmbeddings>(256, config.hidden_size);
             } else {
                 blocks["time_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
                 if (config.vec_in_dim > 0) {
@@ -1027,6 +1048,11 @@ namespace Flux {
                 if (y != nullptr) {
                     txt_img_mask = ggml_pad(ctx->ggml_ctx, y, static_cast<int>(img->ne[1]), 0, 0, 0);
                 }
+            } else if (config.is_sefi) {
+                auto dual_time_embed = std::dynamic_pointer_cast<SefiImage::SefiDualTimestepEmbeddings>(blocks["dual_time_embed"]);
+                auto timestep_sem    = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, 0);
+                auto timestep_tex    = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, ggml_element_size(timesteps));
+                vec                  = dual_time_embed->forward(ctx, timestep_sem, timestep_tex);
             } else {
                 auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
                 vec          = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f));
@@ -1459,7 +1485,7 @@ namespace Flux {
                                  const sd::Tensor<float>& y_tensor                        = {},
                                  const sd::Tensor<float>& guidance_tensor                 = {},
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false,
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::FIXED,
                                  std::vector<int> skip_layers                             = {},
                                  const sd::Tensor<float>& pulid_id_tensor                 = {},
                                  float pulid_id_weight                                    = 1.0f) {
@@ -1500,9 +1526,9 @@ namespace Flux {
                 set_backend_tensor_data(mod_index_arange, mod_index_arange_vec.data());
             }
             std::set<int> txt_arange_dims;
-            if (sd_version_is_flux2(version)) {
-                txt_arange_dims    = {3};
-                increase_ref_index = true;
+            if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
+                txt_arange_dims = {3};
+                ref_index_mode  = Rope::RefIndexMode::INCREASE;
             } else if (version == VERSION_OVIS_IMAGE) {
                 txt_arange_dims = {1, 2};
             }
@@ -1513,7 +1539,7 @@ namespace Flux {
                                             static_cast<int>(context->ne[1]),
                                             txt_arange_dims,
                                             ref_latents,
-                                            increase_ref_index,
+                                            ref_index_mode,
                                             config.ref_index_scale,
                                             config.theta,
                                             circular_y_enabled,
@@ -1573,7 +1599,7 @@ namespace Flux {
                                   const sd::Tensor<float>& y                        = {},
                                   const sd::Tensor<float>& guidance                 = {},
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false,
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED,
                                   std::vector<int> skip_layers                      = std::vector<int>(),
                                   const sd::Tensor<float>& pulid_id                 = {},
                                   float pulid_id_weight                             = 1.0f) {
@@ -1584,7 +1610,7 @@ namespace Flux {
             // guidance: [N, ]
             // pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers, pulid_id, pulid_id_weight);
+                return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
             };
 
             auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -1606,7 +1632,7 @@ namespace Flux {
                            tensor_or_empty(diffusion_params.y),
                            tensor_or_empty(extra->guidance),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index,
+                           diffusion_params.ref_index_mode,
                            extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
                            tensor_or_empty(extra->pulid_id),
                            extra->pulid_id_weight);
@@ -1657,7 +1683,7 @@ namespace Flux {
                                        {},
                                        guidance,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());

src/model/diffusion/hidream_o1.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
+#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 #define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 
 #include <algorithm>

src/model/diffusion/krea2.hpp

@@ -0,0 +1,683 @@
+#ifndef __SD_MODEL_DIFFUSION_KREA2_HPP__
+#define __SD_MODEL_DIFFUSION_KREA2_HPP__
+
+#include <inttypes.h>
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "core/ggml_graph_cut.h"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
+
+namespace Krea2 {
+    constexpr int KREA2_GRAPH_SIZE = 65536;
+
+    struct Krea2Config {
+        int patch_size            = 2;
+        int64_t in_channels       = 16;
+        int64_t out_channels      = 16;
+        int64_t features          = 6144;
+        int64_t timestep_dim      = 256;
+        int64_t text_dim          = 2560;
+        int64_t text_layers       = 12;
+        int64_t layers            = 28;
+        int64_t heads             = 48;
+        int64_t kv_heads          = 12;
+        int64_t text_heads        = 20;
+        int64_t text_kv_heads     = 20;
+        int64_t mlp_multiplier    = 4;
+        float theta               = 1000.f;
+        float norm_eps            = 1e-5f;
+        std::vector<int> axes_dim = {32, 48, 48};
+        int axes_dim_sum          = 128;
+
+        int64_t head_dim() const {
+            return features / heads;
+        }
+
+        static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
+                                    const std::string& prefix,
+                                    const std::string& block_prefix) {
+            int64_t count           = 0;
+            std::string full_prefix = prefix.empty() ? block_prefix : prefix + "." + block_prefix;
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (!starts_with(name, full_prefix)) {
+                    continue;
+                }
+                std::string tail = name.substr(full_prefix.size());
+                size_t dot       = tail.find('.');
+                if (dot == std::string::npos) {
+                    continue;
+                }
+                int block_index = std::atoi(tail.substr(0, dot).c_str());
+                count           = std::max<int64_t>(count, block_index + 1);
+            }
+            return count;
+        }
+
+        void update_axes_dim() {
+            int64_t dim_head = head_dim();
+            int64_t unit     = dim_head / 16;
+            axes_dim         = {
+                        static_cast<int>(dim_head - 12 * unit),
+                        static_cast<int>(6 * unit),
+                        static_cast<int>(6 * unit),
+            };
+            axes_dim_sum = axes_dim[0] + axes_dim[1] + axes_dim[2];
+        }
+
+        static Krea2Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                               const std::string& prefix) {
+            Krea2Config config;
+            int64_t detected_head_dim      = 0;
+            int64_t detected_text_head_dim = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "first.weight") && tensor_storage.n_dims == 2) {
+                    config.in_channels  = tensor_storage.ne[0] / (config.patch_size * config.patch_size);
+                    config.out_channels = config.in_channels;
+                    config.features     = tensor_storage.ne[1];
+                } else if (ends_with(name, "blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) {
+                    detected_head_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) {
+                    if (detected_head_dim > 0) {
+                        config.heads = tensor_storage.ne[1] / detected_head_dim;
+                    }
+                } else if (ends_with(name, "blocks.0.attn.wk.weight") && tensor_storage.n_dims == 2) {
+                    if (detected_head_dim > 0) {
+                        config.kv_heads = tensor_storage.ne[1] / detected_head_dim;
+                    }
+                } else if (ends_with(name, "txtfusion.projector.weight") && tensor_storage.n_dims == 2) {
+                    config.text_layers = tensor_storage.ne[0];
+                } else if (ends_with(name, "txtfusion.layerwise_blocks.0.prenorm.scale") && tensor_storage.n_dims == 1) {
+                    config.text_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) {
+                    detected_text_head_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) {
+                    if (detected_text_head_dim > 0) {
+                        config.text_heads = tensor_storage.ne[1] / detected_text_head_dim;
+                    }
+                } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wk.weight") && tensor_storage.n_dims == 2) {
+                    if (detected_text_head_dim > 0) {
+                        config.text_kv_heads = tensor_storage.ne[1] / detected_text_head_dim;
+                    }
+                } else if (ends_with(name, "last.linear.weight") && tensor_storage.n_dims == 2) {
+                    config.out_channels = tensor_storage.ne[1] / (config.patch_size * config.patch_size);
+                }
+            }
+
+            config.layers = std::max<int64_t>(1, count_blocks(tensor_storage_map, prefix, "blocks."));
+            if (detected_head_dim > 0 && config.features > 0) {
+                config.heads = config.features / detected_head_dim;
+            }
+            if (detected_head_dim > 0) {
+                std::string wk_name = prefix.empty() ? "blocks.0.attn.wk.weight" : prefix + ".blocks.0.attn.wk.weight";
+                auto it             = tensor_storage_map.find(wk_name);
+                if (it != tensor_storage_map.end() && it->second.n_dims == 2) {
+                    config.kv_heads = it->second.ne[1] / detected_head_dim;
+                }
+            }
+            if (detected_text_head_dim > 0 && config.text_dim > 0) {
+                config.text_heads = config.text_dim / detected_text_head_dim;
+            }
+            if (detected_text_head_dim > 0) {
+                std::string wk_name = prefix.empty() ? "txtfusion.layerwise_blocks.0.attn.wk.weight" : prefix + ".txtfusion.layerwise_blocks.0.attn.wk.weight";
+                auto it             = tensor_storage_map.find(wk_name);
+                if (it != tensor_storage_map.end() && it->second.n_dims == 2) {
+                    config.text_kv_heads = it->second.ne[1] / detected_text_head_dim;
+                }
+            }
+            config.update_axes_dim();
+
+            LOG_DEBUG("krea2: layers=%" PRId64 ", features=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", text_dim=%" PRId64 ", text_layers=%" PRId64 ", text_heads=%" PRId64 ", text_kv_heads=%" PRId64 ", channels=%" PRId64,
+                      config.layers,
+                      config.features,
+                      config.heads,
+                      config.kv_heads,
+                      config.text_dim,
+                      config.text_layers,
+                      config.text_heads,
+                      config.text_kv_heads,
+                      config.in_channels);
+            return config;
+        }
+    };
+
+    __STATIC_INLINE__ int64_t ceil_to_multiple(int64_t value, int64_t multiple) {
+        return ((value + multiple - 1) / multiple) * multiple;
+    }
+
+    class KreaRMSNorm : public UnaryBlock {
+    protected:
+        int64_t hidden_size;
+        float eps;
+        std::string prefix;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            GGML_UNUSED(tensor_storage_map);
+            this->prefix    = prefix;
+            params["scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);
+        }
+
+    public:
+        KreaRMSNorm(int64_t hidden_size, float eps = 1e-5f)
+            : hidden_size(hidden_size),
+              eps(eps) {}
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            ggml_tensor* scale = params["scale"];
+            scale              = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
+            x                  = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+            x                  = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
+            return x;
+        }
+    };
+
+    class KreaSwiGLU : public UnaryBlock {
+    public:
+        KreaSwiGLU(int64_t features, int64_t multiplier) {
+            int64_t mlp_dim = ceil_to_multiple(((2 * features) / 3) * multiplier, 128);
+            blocks["gate"]  = std::make_shared<Linear>(features, mlp_dim, false);
+            blocks["up"]    = std::make_shared<Linear>(features, mlp_dim, false);
+            blocks["down"]  = std::make_shared<Linear>(mlp_dim, features, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto gate = std::dynamic_pointer_cast<Linear>(blocks["gate"]);
+            auto up   = std::dynamic_pointer_cast<Linear>(blocks["up"]);
+            auto down = std::dynamic_pointer_cast<Linear>(blocks["down"]);
+
+            auto gated = ggml_silu(ctx->ggml_ctx, gate->forward(ctx, x));
+            auto up_x  = up->forward(ctx, x);
+            x          = ggml_mul(ctx->ggml_ctx, gated, up_x);
+            return down->forward(ctx, x);
+        }
+    };
+
+    class KreaAttention : public GGMLBlock {
+    protected:
+        int64_t features;
+        int64_t heads;
+        int64_t kv_heads;
+        int64_t head_dim_;
+
+        ggml_tensor* attention_no_rope(GGMLRunnerContext* ctx,
+                                       ggml_tensor* q,
+                                       ggml_tensor* k,
+                                       ggml_tensor* v,
+                                       ggml_tensor* mask) {
+            int64_t Lq = q->ne[2];
+            int64_t Lk = k->ne[2];
+            int64_t N  = q->ne[3];
+            q          = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim_ * heads, Lq, N);
+            k          = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim_ * kv_heads, Lk, N);
+            v          = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim_ * kv_heads, Lk, N);
+            return ggml_ext_attention_ext(ctx->ggml_ctx,
+                                          ctx->backend,
+                                          q,
+                                          k,
+                                          v,
+                                          heads,
+                                          mask,
+                                          false,
+                                          ctx->flash_attn_enabled);
+        }
+
+    public:
+        KreaAttention(int64_t features,
+                      int64_t heads,
+                      int64_t kv_heads,
+                      float eps = 1e-5f)
+            : features(features),
+              heads(heads),
+              kv_heads(kv_heads),
+              head_dim_(features / heads) {
+            blocks["wq"]           = std::make_shared<Linear>(features, heads * head_dim_, false);
+            blocks["wk"]           = std::make_shared<Linear>(features, kv_heads * head_dim_, false);
+            blocks["wv"]           = std::make_shared<Linear>(features, kv_heads * head_dim_, false);
+            blocks["gate"]         = std::make_shared<Linear>(features, features, false);
+            blocks["qknorm.qnorm"] = std::make_shared<KreaRMSNorm>(head_dim_, eps);
+            blocks["qknorm.knorm"] = std::make_shared<KreaRMSNorm>(head_dim_, eps);
+            blocks["wo"]           = std::make_shared<Linear>(features, features, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe   = nullptr,
+                             ggml_tensor* mask = nullptr) {
+            auto wq    = std::dynamic_pointer_cast<Linear>(blocks["wq"]);
+            auto wk    = std::dynamic_pointer_cast<Linear>(blocks["wk"]);
+            auto wv    = std::dynamic_pointer_cast<Linear>(blocks["wv"]);
+            auto gate  = std::dynamic_pointer_cast<Linear>(blocks["gate"]);
+            auto qnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.qnorm"]);
+            auto knorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.knorm"]);
+            auto wo    = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
+
+            if (sd_backend_is(ctx->backend, "Vulkan")) {
+                wo->set_force_prec_f32(true);
+            }
+
+            int64_t L = x->ne[1];
+            int64_t N = x->ne[2];
+
+            auto q = wq->forward(ctx, x);
+            q      = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim_, heads, L, N);
+            auto k = wk->forward(ctx, x);
+            k      = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim_, kv_heads, L, N);
+            auto v = wv->forward(ctx, x);
+            v      = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim_, kv_heads, L, N);
+
+            q = qnorm->forward(ctx, q);
+            k = knorm->forward(ctx, k);
+
+            auto out = pe != nullptr ? Rope::attention(ctx, q, k, v, pe, mask)
+                                     : attention_no_rope(ctx, q, k, v, mask);
+            out      = ggml_mul(ctx->ggml_ctx, out, ggml_sigmoid(ctx->ggml_ctx, gate->forward(ctx, x)));
+            out      = wo->forward(ctx, out);
+            return out;
+        }
+    };
+
+    class KreaDoubleSharedModulation : public GGMLBlock {
+    protected:
+        int64_t dim;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            GGML_UNUSED(tensor_storage_map);
+            GGML_UNUSED(prefix);
+            params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
+        }
+
+    public:
+        KreaDoubleSharedModulation(int64_t dim)
+            : dim(dim) {}
+
+        std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
+            auto lin = ggml_repeat(ctx->ggml_ctx, params["lin"], vec);
+            auto out = ggml_add(ctx->ggml_ctx, vec, lin);
+            return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
+        }
+    };
+
+    class KreaFinalModulation : public GGMLBlock {
+    protected:
+        int64_t dim;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            GGML_UNUSED(tensor_storage_map);
+            GGML_UNUSED(prefix);
+            params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
+        }
+
+    public:
+        KreaFinalModulation(int64_t dim)
+            : dim(dim) {}
+
+        std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
+            auto out = ggml_add(ctx->ggml_ctx, params["lin"], vec);
+            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+        }
+    };
+
+    class KreaTextFusionBlock : public UnaryBlock {
+    public:
+        KreaTextFusionBlock(int64_t dim,
+                            int64_t heads,
+                            int64_t kv_heads,
+                            int64_t multiplier,
+                            float eps) {
+            blocks["prenorm"]  = std::make_shared<KreaRMSNorm>(dim, eps);
+            blocks["postnorm"] = std::make_shared<KreaRMSNorm>(dim, eps);
+            blocks["attn"]     = std::make_shared<KreaAttention>(dim, heads, kv_heads, eps);
+            blocks["mlp"]      = std::make_shared<KreaSwiGLU>(dim, multiplier);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto prenorm  = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["prenorm"]);
+            auto postnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["postnorm"]);
+            auto attn     = std::dynamic_pointer_cast<KreaAttention>(blocks["attn"]);
+            auto mlp      = std::dynamic_pointer_cast<KreaSwiGLU>(blocks["mlp"]);
+
+            x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, prenorm->forward(ctx, x)));
+            x = ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, postnorm->forward(ctx, x)));
+            return x;
+        }
+    };
+
+    class KreaTextFusionTransformer : public UnaryBlock {
+    protected:
+        Krea2Config config;
+
+    public:
+        explicit KreaTextFusionTransformer(Krea2Config config)
+            : config(std::move(config)) {
+            for (int i = 0; i < 2; ++i) {
+                blocks["layerwise_blocks." + std::to_string(i)] = std::make_shared<KreaTextFusionBlock>(this->config.text_dim,
+                                                                                                        this->config.text_heads,
+                                                                                                        this->config.text_kv_heads,
+                                                                                                        this->config.mlp_multiplier,
+                                                                                                        this->config.norm_eps);
+                blocks["refiner_blocks." + std::to_string(i)]   = std::make_shared<KreaTextFusionBlock>(this->config.text_dim,
+                                                                                                      this->config.text_heads,
+                                                                                                      this->config.text_kv_heads,
+                                                                                                      this->config.mlp_multiplier,
+                                                                                                      this->config.norm_eps);
+            }
+            blocks["projector"] = std::make_shared<Linear>(this->config.text_layers, 1, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* context) override {
+            int64_t text_tokens = context->ne[1];
+            int64_t batch       = context->ne[2];
+
+            context = ggml_reshape_3d(ctx->ggml_ctx,
+                                      context,
+                                      config.text_dim,
+                                      config.text_layers,
+                                      text_tokens * batch);
+
+            for (int i = 0; i < 2; ++i) {
+                auto block = std::dynamic_pointer_cast<KreaTextFusionBlock>(blocks["layerwise_blocks." + std::to_string(i)]);
+                context    = block->forward(ctx, context);
+            }
+
+            context        = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3));
+            auto projector = std::dynamic_pointer_cast<Linear>(blocks["projector"]);
+            context        = projector->forward(ctx, context);
+            context        = ggml_reshape_3d(ctx->ggml_ctx, context, config.text_dim, text_tokens, batch);
+
+            for (int i = 0; i < 2; ++i) {
+                auto block = std::dynamic_pointer_cast<KreaTextFusionBlock>(blocks["refiner_blocks." + std::to_string(i)]);
+                context    = block->forward(ctx, context);
+            }
+            return context;
+        }
+    };
+
+    class KreaSingleStreamBlock : public UnaryBlock {
+    public:
+        explicit KreaSingleStreamBlock(Krea2Config config) {
+            blocks["mod"]      = std::make_shared<KreaDoubleSharedModulation>(config.features);
+            blocks["prenorm"]  = std::make_shared<KreaRMSNorm>(config.features, config.norm_eps);
+            blocks["postnorm"] = std::make_shared<KreaRMSNorm>(config.features, config.norm_eps);
+            blocks["attn"]     = std::make_shared<KreaAttention>(config.features, config.heads, config.kv_heads, config.norm_eps);
+            blocks["mlp"]      = std::make_shared<KreaSwiGLU>(config.features, config.mlp_multiplier);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* vec,
+                             ggml_tensor* pe) {
+            auto mod      = std::dynamic_pointer_cast<KreaDoubleSharedModulation>(blocks["mod"]);
+            auto prenorm  = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["prenorm"]);
+            auto postnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["postnorm"]);
+            auto attn     = std::dynamic_pointer_cast<KreaAttention>(blocks["attn"]);
+            auto mlp      = std::dynamic_pointer_cast<KreaSwiGLU>(blocks["mlp"]);
+
+            auto mods       = mod->forward(ctx, vec);
+            auto attn_input = Flux::modulate(ctx->ggml_ctx,
+                                             prenorm->forward(ctx, x),
+                                             mods[1],
+                                             mods[0],
+                                             true);
+            auto attn_out   = attn->forward(ctx, attn_input, pe);
+            x               = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
+
+            auto mlp_input = Flux::modulate(ctx->ggml_ctx,
+                                            postnorm->forward(ctx, x),
+                                            mods[4],
+                                            mods[3],
+                                            true);
+            auto mlp_out   = mlp->forward(ctx, mlp_input);
+            x              = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
+            return x;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            GGML_UNUSED(ctx);
+            GGML_UNUSED(x);
+            GGML_ABORT("KreaSingleStreamBlock requires conditioning");
+            return nullptr;
+        }
+    };
+
+    class KreaTimeMLP : public UnaryBlock {
+    public:
+        explicit KreaTimeMLP(Krea2Config config) {
+            blocks["0"] = std::make_shared<Linear>(config.timestep_dim, config.features, true);
+            blocks["2"] = std::make_shared<Linear>(config.features, config.features, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
+            x             = linear_0->forward(ctx, x);
+            x             = ggml_ext_gelu(ctx->ggml_ctx, x, false);
+            x             = linear_2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class KreaTProj : public UnaryBlock {
+    public:
+        explicit KreaTProj(Krea2Config config) {
+            blocks["1"] = std::make_shared<Linear>(config.features, config.features * 6, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
+            x             = ggml_ext_gelu(ctx->ggml_ctx, x, false);
+            x             = linear_1->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class KreaTextMLP : public UnaryBlock {
+    public:
+        explicit KreaTextMLP(Krea2Config config) {
+            blocks["0"] = std::make_shared<KreaRMSNorm>(config.text_dim, config.norm_eps);
+            blocks["1"] = std::make_shared<Linear>(config.text_dim, config.features, true);
+            blocks["3"] = std::make_shared<Linear>(config.features, config.features, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto norm     = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["0"]);
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
+            auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["3"]);
+            x             = norm->forward(ctx, x);
+            x             = linear_1->forward(ctx, x);
+            x             = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x             = linear_3->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class KreaLastLayer : public GGMLBlock {
+    public:
+        explicit KreaLastLayer(Krea2Config config) {
+            blocks["norm"]       = std::make_shared<KreaRMSNorm>(config.features, config.norm_eps);
+            blocks["linear"]     = std::make_shared<Linear>(config.features, config.patch_size * config.patch_size * config.out_channels, true);
+            blocks["modulation"] = std::make_shared<KreaFinalModulation>(config.features);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* vec) {
+            auto norm       = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["norm"]);
+            auto linear     = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            auto modulation = std::dynamic_pointer_cast<KreaFinalModulation>(blocks["modulation"]);
+
+            auto mods = modulation->forward(ctx, vec);
+            x         = Flux::modulate(ctx->ggml_ctx,
+                                       norm->forward(ctx, x),
+                                       mods[1],
+                                       mods[0],
+                                       true);
+            x         = linear->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class Krea2Model : public GGMLBlock {
+    protected:
+        Krea2Config config;
+
+    public:
+        Krea2Model() = default;
+        explicit Krea2Model(Krea2Config config)
+            : config(std::move(config)) {
+            blocks["first"]     = std::make_shared<Linear>(this->config.patch_size * this->config.patch_size * this->config.in_channels,
+                                                       this->config.features,
+                                                       true);
+            blocks["tmlp"]      = std::make_shared<KreaTimeMLP>(this->config);
+            blocks["txtfusion"] = std::make_shared<KreaTextFusionTransformer>(this->config);
+            blocks["txtmlp"]    = std::make_shared<KreaTextMLP>(this->config);
+            blocks["tproj"]     = std::make_shared<KreaTProj>(this->config);
+            for (int i = 0; i < this->config.layers; ++i) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<KreaSingleStreamBlock>(this->config);
+            }
+            blocks["last"] = std::make_shared<KreaLastLayer>(this->config);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe) {
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t N = x->ne[3];
+            GGML_ASSERT(N == 1);
+
+            auto first     = std::dynamic_pointer_cast<Linear>(blocks["first"]);
+            auto tmlp      = std::dynamic_pointer_cast<KreaTimeMLP>(blocks["tmlp"]);
+            auto txtfusion = std::dynamic_pointer_cast<KreaTextFusionTransformer>(blocks["txtfusion"]);
+            auto txtmlp    = std::dynamic_pointer_cast<KreaTextMLP>(blocks["txtmlp"]);
+            auto tproj     = std::dynamic_pointer_cast<KreaTProj>(blocks["tproj"]);
+            auto last      = std::dynamic_pointer_cast<KreaLastLayer>(blocks["last"]);
+
+            auto img        = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true);
+            int64_t img_len = img->ne[1];
+            img             = first->forward(ctx, img);
+
+            auto t    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
+            t         = tmlp->forward(ctx, t);
+            t         = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
+            auto tvec = tproj->forward(ctx, t);
+
+            auto txt        = txtfusion->forward(ctx, context);
+            txt             = txtmlp->forward(ctx, txt);
+            int64_t txt_len = txt->ne[1];
+
+            auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1);
+            for (int i = 0; i < config.layers; ++i) {
+                auto block    = std::dynamic_pointer_cast<KreaSingleStreamBlock>(blocks["blocks." + std::to_string(i)]);
+                hidden_states = block->forward(ctx, hidden_states, tvec, pe);
+                sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
+            }
+
+            hidden_states = last->forward(ctx, hidden_states, t);
+            hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
+            hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
+            return hidden_states;
+        }
+    };
+
+    __STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h,
+                                                      int w,
+                                                      int patch_size,
+                                                      int bs,
+                                                      int context_len,
+                                                      float theta,
+                                                      const std::vector<int>& axes_dim) {
+        auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
+        auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
+        auto ids     = Rope::concat_ids(txt_ids, img_ids, bs);
+        return Rope::embed_nd(ids, bs, theta, axes_dim);
+    }
+
+    struct Krea2Runner : public DiffusionModelRunner {
+        Krea2Config config;
+        Krea2Model model;
+        std::vector<float> pe_vec;
+
+        Krea2Runner(ggml_backend_t backend,
+                    const String2TensorStorage& tensor_storage_map      = {},
+                    const std::string prefix                            = "",
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(Krea2Config::detect_from_weights(tensor_storage_map, prefix)) {
+            model = Krea2Model(config);
+            model.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "krea2";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.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(KREA2_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      = gen_krea2_pe(static_cast<int>(x->ne[1]),
+                                       static_cast<int>(x->ne[0]),
+                                       config.patch_size,
+                                       static_cast<int>(x->ne[3]),
+                                       static_cast<int>(context->ne[1]),
+                                       config.theta,
+                                       config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = model.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, false, false), x.dim());
+        }
+
+        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 compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context));
+        }
+    };
+}  // namespace Krea2
+
+#endif  // __SD_MODEL_DIFFUSION_KREA2_HPP__

src/model/diffusion/minit2i.hpp

@@ -0,0 +1,611 @@
+#ifndef __SD_MODEL_DIFFUSION_MINIT2I_HPP__
+#define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
+
+namespace MiniT2I {
+    constexpr int MINIT2I_GRAPH_SIZE = 196608;
+
+    struct MiniT2IConfig {
+        int64_t image_size         = 512;
+        int64_t patch_size         = 16;
+        int64_t in_channels        = 3;
+        int64_t txt_input_size     = 1024;
+        int64_t hidden_size        = 768;
+        int64_t txt_hidden_size    = 768;
+        int64_t cond_vec_size      = 768;
+        int64_t depth_double       = 17;
+        int64_t txt_preamble_depth = 2;
+        int64_t num_heads          = 12;
+        int64_t head_dim           = 64;
+        float mlp_ratio            = 2.6667f;
+        int64_t pca_channels       = 128;
+        int64_t prompt_length      = 256;
+        int64_t n_T                = 100;
+        float cfg_interval_start   = 0.0f;
+        float cfg_interval_end     = 1.0f;
+
+        static MiniT2IConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            MiniT2IConfig config;
+            config.depth_double       = 0;
+            config.txt_preamble_depth = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "img_embedder.proj1.weight") && tensor_storage.n_dims == 4) {
+                    config.patch_size   = tensor_storage.ne[0];
+                    config.in_channels  = tensor_storage.ne[2];
+                    config.pca_channels = tensor_storage.ne[3];
+                } else if (ends_with(name, "img_embedder.proj2.weight") && tensor_storage.n_dims == 4) {
+                    config.pca_channels = tensor_storage.ne[2];
+                    config.hidden_size  = tensor_storage.ne[3];
+                } else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) {
+                    config.txt_input_size  = tensor_storage.ne[0];
+                    config.txt_hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) {
+                    config.cond_vec_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) {
+                    int64_t inner3     = tensor_storage.ne[1];
+                    int64_t inner      = inner3 / 3;
+                    config.hidden_size = tensor_storage.ne[0];
+                    if (config.hidden_size == 768) {
+                        config.num_heads = 12;
+                        config.head_dim  = 64;
+                    } else if (config.hidden_size == 1248) {
+                        config.num_heads = 24;
+                        config.head_dim  = 52;
+                    } else if (inner > 0) {
+                        config.head_dim  = 64;
+                        config.num_heads = std::max<int64_t>(1, inner / config.head_dim);
+                    }
+                } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area = config.patch_size * config.patch_size;
+                    config.hidden_size = tensor_storage.ne[0];
+                    config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels;
+                } else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) {
+                    config.prompt_length = tensor_storage.ne[1];
+                }
+
+                size_t pos = name.find("double_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int64_t idx         = atoi(items[1].c_str());
+                        config.depth_double = std::max<int64_t>(config.depth_double, idx + 1);
+                    }
+                }
+                pos = name.find("txt_preamble_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int64_t idx               = atoi(items[1].c_str());
+                        config.txt_preamble_depth = std::max<int64_t>(config.txt_preamble_depth, idx + 1);
+                    }
+                }
+            }
+
+            if (config.depth_double <= 0) {
+                config.depth_double = config.hidden_size == 1248 ? 23 : 17;
+            }
+            if (config.txt_preamble_depth <= 0) {
+                config.txt_preamble_depth = 2;
+            }
+            if (config.head_dim <= 0 || config.num_heads <= 0) {
+                config.head_dim  = config.hidden_size == 1248 ? 52 : 64;
+                config.num_heads = config.hidden_size / config.head_dim;
+            }
+            LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64,
+                      config.hidden_size,
+                      config.txt_hidden_size,
+                      config.num_heads,
+                      config.head_dim,
+                      config.depth_double,
+                      config.txt_preamble_depth,
+                      config.patch_size,
+                      config.in_channels);
+            return config;
+        }
+    };
+
+    inline std::vector<float> make_2d_sincos_pos_embed(int grid_size, int dim) {
+        GGML_ASSERT(dim % 4 == 0);
+        int half_dim = dim / 2;
+        int quarter  = half_dim / 2;
+        std::vector<float> out(static_cast<size_t>(grid_size) * grid_size * dim);
+        std::vector<float> omega(quarter);
+        for (int i = 0; i < quarter; ++i) {
+            omega[i] = 1.0f / std::pow(10000.0f, static_cast<float>(i) / static_cast<float>(quarter));
+        }
+        for (int y = 0; y < grid_size; ++y) {
+            for (int x = 0; x < grid_size; ++x) {
+                size_t base = static_cast<size_t>(y * grid_size + x) * dim;
+                for (int i = 0; i < quarter; ++i) {
+                    float ay                           = y * omega[i];
+                    float ax                           = x * omega[i];
+                    out[base + i]                      = std::sin(ax);
+                    out[base + quarter + i]            = std::cos(ax);
+                    out[base + half_dim + i]           = std::sin(ay);
+                    out[base + half_dim + quarter + i] = std::cos(ay);
+                }
+            }
+        }
+        return out;
+    }
+
+    inline std::vector<float> make_text_rope(int length, int head_dim) {
+        return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
+    }
+
+    inline std::vector<float> make_vision_rope(int side, int head_dim) {
+        GGML_ASSERT(head_dim % 4 == 0);
+        int dim     = head_dim / 2;
+        int quarter = dim / 2;
+        int length  = side * side;
+        std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
+        std::vector<float> freqs(quarter);
+        for (int i = 0; i < quarter; ++i) {
+            freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
+        }
+        for (int y = 0; y < side; ++y) {
+            for (int x = 0; x < side; ++x) {
+                int pos     = y * side + x;
+                size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
+                for (int i = 0; i < quarter; ++i) {
+                    float ay        = y * freqs[i];
+                    float ax        = x * freqs[i];
+                    float angles[2] = {ay, ax};
+                    for (int axis = 0; axis < 2; ++axis) {
+                        int j                 = axis * quarter + i;
+                        out[base + 4 * j]     = std::cos(angles[axis]);
+                        out[base + 4 * j + 1] = -std::sin(angles[axis]);
+                        out[base + 4 * j + 2] = std::sin(angles[axis]);
+                        out[base + 4 * j + 3] = std::cos(angles[axis]);
+                    }
+                }
+            }
+        }
+        return out;
+    }
+
+    struct SwiGLUMlp : public GGMLBlock {
+        SwiGLUMlp(int64_t in_features, int64_t hidden_features) {
+            int64_t hidden_dim = ((hidden_features + 7) / 8) * 8;
+            blocks["w1"]       = std::make_shared<Linear>(in_features, hidden_dim, false);
+            blocks["w3"]       = std::make_shared<Linear>(in_features, hidden_dim, false);
+            blocks["w2"]       = std::make_shared<Linear>(hidden_dim, in_features, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto w1   = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
+            auto w3   = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
+            auto w2   = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
+            auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto up   = w3->forward(ctx, x);
+            return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up));
+        }
+    };
+
+    struct BottleneckPatchEmbed : public GGMLBlock {
+        int64_t patch_size;
+
+        BottleneckPatchEmbed(int64_t patch_size, int64_t in_channels, int64_t pca_channels, int64_t hidden_size)
+            : patch_size(patch_size) {
+            blocks["proj1"] = std::make_shared<Conv2d>(in_channels,
+                                                       pca_channels,
+                                                       std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
+                                                       std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
+                                                       std::pair<int, int>{0, 0},
+                                                       std::pair<int, int>{1, 1},
+                                                       false);
+            blocks["proj2"] = std::make_shared<Conv2d>(pca_channels,
+                                                       hidden_size,
+                                                       std::pair<int, int>{1, 1},
+                                                       std::pair<int, int>{1, 1},
+                                                       std::pair<int, int>{0, 0},
+                                                       std::pair<int, int>{1, 1},
+                                                       true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto proj1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj1"]);
+            auto proj2 = std::dynamic_pointer_cast<Conv2d>(blocks["proj2"]);
+            x          = proj1->forward(ctx, x);
+            x          = proj2->forward(ctx, x);
+            x          = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+            x          = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            return x;
+        }
+    };
+
+    struct TimestepEmbedder : public GGMLBlock {
+        int frequency_embedding_size;
+
+        TimestepEmbedder(int64_t hidden_size, int frequency_embedding_size = 256)
+            : frequency_embedding_size(frequency_embedding_size) {
+            blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
+            blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
+            auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f);
+            t_emb      = mlp_0->forward(ctx, t_emb);
+            t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
+            return mlp_2->forward(ctx, t_emb);
+        }
+    };
+
+    inline std::vector<ggml_tensor*> split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) {
+        int64_t N = qkv->ne[2];
+        int64_t L = qkv->ne[1];
+        auto q    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
+                                 qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
+        auto k    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
+                                 qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads);
+        auto v    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
+                                 qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads * 2);
+        return {q, k, v};
+    }
+
+    struct PlainTextTransformerBlock : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        PlainTextTransformerBlock(int64_t hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
+            : num_heads(num_heads), head_dim(head_dim) {
+            int64_t inner_dim   = num_heads * head_dim;
+            blocks["norm1"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["norm2"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["qkv"]       = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
+            blocks["attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
+            blocks["mlp"]       = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
+            blocks["q_norm"]    = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm"]    = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* pe) {
+            auto norm1     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
+            auto norm2     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
+            auto qkv_proj  = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
+            auto attn_proj = std::dynamic_pointer_cast<Linear>(blocks["attn_proj"]);
+            auto mlp       = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["mlp"]);
+            auto q_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+
+            auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
+            auto q   = q_norm->forward(ctx, qkv[0]);
+            auto k   = k_norm->forward(ctx, qkv[1]);
+            auto v   = qkv[2];
+            auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
+            txt      = ggml_add(ctx->ggml_ctx, txt, attn_proj->forward(ctx, out));
+            txt      = ggml_add(ctx->ggml_ctx, txt, mlp->forward(ctx, norm2->forward(ctx, txt)));
+            return txt;
+        }
+    };
+
+    struct DoubleStreamDiTBlock : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        DoubleStreamDiTBlock(int64_t hidden_size, int64_t txt_hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
+            : num_heads(num_heads), head_dim(head_dim) {
+            int64_t inner_dim       = num_heads * head_dim;
+            blocks["img_norm1"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["img_norm2"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["txt_norm1"]     = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
+            blocks["txt_norm2"]     = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
+            blocks["img_qkv"]       = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
+            blocks["txt_qkv"]       = std::make_shared<Linear>(txt_hidden_size, inner_dim * 3, true);
+            blocks["q_norm"]        = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm"]        = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["img_attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
+            blocks["txt_attn_proj"] = std::make_shared<Linear>(inner_dim, txt_hidden_size, true);
+            blocks["img_mlp"]       = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
+            blocks["txt_mlp"]       = std::make_shared<SwiGLUMlp>(txt_hidden_size, static_cast<int64_t>(txt_hidden_size * mlp_ratio));
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* img,
+                                                      ggml_tensor* txt,
+                                                      ggml_tensor* pe) {
+            auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
+            auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
+            auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
+            auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
+            auto img_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
+            auto txt_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
+            auto q_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+            auto img_proj  = std::dynamic_pointer_cast<Linear>(blocks["img_attn_proj"]);
+            auto txt_proj  = std::dynamic_pointer_cast<Linear>(blocks["txt_attn_proj"]);
+            auto img_mlp   = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["img_mlp"]);
+            auto txt_mlp   = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["txt_mlp"]);
+
+            int64_t li = img->ne[1];
+            int64_t lt = txt->ne[1];
+
+            auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
+            auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
+
+            auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2);
+            auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2);
+            auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2);
+
+            auto out     = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
+            auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt);
+            auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li);
+
+            img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img));
+            txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt));
+            img = ggml_add(ctx->ggml_ctx, img, img_mlp->forward(ctx, img_norm2->forward(ctx, img)));
+            txt = ggml_add(ctx->ggml_ctx, txt, txt_mlp->forward(ctx, txt_norm2->forward(ctx, txt)));
+            return {img, txt};
+        }
+    };
+
+    struct FinalLayer : public GGMLBlock {
+        FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels) {
+            blocks["norm_final"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["linear"]     = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm_final = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_final"]);
+            auto linear     = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            return linear->forward(ctx, norm_final->forward(ctx, x));
+        }
+    };
+
+    struct MMJiT : public GGMLBlock {
+        MiniT2IConfig config;
+
+        MMJiT(const MiniT2IConfig& config)
+            : config(config) {
+            blocks["img_embedder"]    = std::make_shared<BottleneckPatchEmbed>(config.patch_size, config.in_channels, config.pca_channels, config.hidden_size);
+            blocks["txt_embedder"]    = std::make_shared<Linear>(config.txt_input_size, config.txt_hidden_size, false);
+            blocks["t_embedder"]      = std::make_shared<TimestepEmbedder>(config.cond_vec_size);
+            blocks["pooled_embedder"] = std::make_shared<Linear>(config.txt_input_size, config.cond_vec_size, false);
+            for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
+                blocks["txt_preamble_blocks." + std::to_string(i)] = std::make_shared<PlainTextTransformerBlock>(config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
+            }
+            for (int64_t i = 0; i < config.depth_double; ++i) {
+                blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamDiTBlock>(config.hidden_size, config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
+            }
+            blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.in_channels);
+        }
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
+            enum ggml_type wtype = get_type(prefix + "mask_token", tensor_storage_map, GGML_TYPE_F32);
+            params["mask_token"] = ggml_new_tensor_3d(ctx, wtype, config.txt_input_size, 1, 1);
+        }
+
+        ggml_tensor* apply_text_mask(GGMLRunnerContext* ctx, ggml_tensor* context, ggml_tensor* mask) {
+            if (mask == nullptr) {
+                return context;
+            }
+            mask            = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]);
+            mask            = ggml_repeat(ctx->ggml_ctx, mask, context);
+            auto keep       = ggml_mul(ctx->ggml_ctx, context, mask);
+            auto inv        = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask);
+            auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context);
+            return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv));
+        }
+
+        ggml_tensor* pool_context(GGMLRunnerContext* ctx, ggml_tensor* context) {
+            int64_t dim = context->ne[0];
+            int64_t len = context->ne[1];
+            int64_t N   = context->ne[2];
+            auto x      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3));
+            x           = ggml_reshape_3d(ctx->ggml_ctx, x, len, dim, N);
+            x           = ggml_mean(ctx->ggml_ctx, x);
+            x           = ggml_reshape_2d(ctx->ggml_ctx, x, dim, N);
+            return x;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* img,
+                             ggml_tensor* context,
+                             ggml_tensor* mask,
+                             ggml_tensor* pos_embed,
+                             ggml_tensor* txt_pe,
+                             ggml_tensor* joint_pe) {
+            auto img_embedder = std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["img_embedder"]);
+            auto txt_embedder = std::dynamic_pointer_cast<Linear>(blocks["txt_embedder"]);
+            auto final_layer  = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
+
+            int64_t W  = img->ne[0];
+            int64_t H  = img->ne[1];
+            int64_t hp = H / config.patch_size;
+            int64_t wp = W / config.patch_size;
+
+            context = apply_text_mask(ctx, context, mask);
+            auto x  = img_embedder->forward(ctx, img);
+            x       = ggml_add(ctx->ggml_ctx, x, pos_embed);
+
+            auto txt = txt_embedder->forward(ctx, context);
+            for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
+                auto block = std::dynamic_pointer_cast<PlainTextTransformerBlock>(blocks["txt_preamble_blocks." + std::to_string(i)]);
+                txt        = block->forward(ctx, txt, txt_pe);
+                sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.txt_preamble_blocks." + std::to_string(i), "txt");
+            }
+            for (int64_t i = 0; i < config.depth_double; ++i) {
+                auto block = std::dynamic_pointer_cast<DoubleStreamDiTBlock>(blocks["double_blocks." + std::to_string(i)]);
+                auto out   = block->forward(ctx, x, txt, joint_pe);
+                x          = out.first;
+                txt        = out.second;
+                sd::ggml_graph_cut::mark_graph_cut(x, "minit2i.double_blocks." + std::to_string(i), "x");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.double_blocks." + std::to_string(i), "txt");
+            }
+            auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1);
+            auto out      = final_layer->forward(ctx, combined);
+            auto img_out  = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
+            return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
+        }
+    };
+
+    struct MiniT2IRunner : public DiffusionModelRunner {
+        MiniT2IConfig config;
+        MMJiT model;
+        ggml_context* position_cache_ctx            = nullptr;
+        ggml_backend_buffer_t position_cache_buffer = nullptr;
+        ggml_tensor* cached_pos_embed               = nullptr;
+        ggml_tensor* cached_txt_pe                  = nullptr;
+        ggml_tensor* cached_joint_pe                = nullptr;
+        int64_t cached_img_side                     = -1;
+        int64_t cached_txt_len                      = -1;
+        int64_t cached_hidden_size                  = -1;
+        int64_t cached_head_dim                     = -1;
+
+        MiniT2IRunner(ggml_backend_t backend,
+                      const String2TensorStorage& tensor_storage_map      = {},
+                      const std::string prefix                            = "",
+                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(MiniT2IConfig::detect_from_weights(tensor_storage_map, this->prefix)),
+              model(config) {
+            model.init(params_ctx, tensor_storage_map, this->prefix);
+        }
+
+        ~MiniT2IRunner() override {
+            free_position_cache();
+        }
+
+        std::string get_desc() override {
+            return "MiniT2I";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        void free_position_cache() {
+            if (position_cache_buffer != nullptr) {
+                ggml_backend_buffer_free(position_cache_buffer);
+                position_cache_buffer = nullptr;
+            }
+            if (position_cache_ctx != nullptr) {
+                ggml_free(position_cache_ctx);
+                position_cache_ctx = nullptr;
+            }
+            cached_pos_embed   = nullptr;
+            cached_txt_pe      = nullptr;
+            cached_joint_pe    = nullptr;
+            cached_img_side    = -1;
+            cached_txt_len     = -1;
+            cached_hidden_size = -1;
+            cached_head_dim    = -1;
+        }
+
+        void ensure_position_cache(int64_t img_side, int64_t txt_len) {
+            if (cached_img_side == img_side &&
+                cached_txt_len == txt_len &&
+                cached_hidden_size == config.hidden_size &&
+                cached_head_dim == config.head_dim &&
+                cached_pos_embed != nullptr &&
+                cached_txt_pe != nullptr &&
+                cached_joint_pe != nullptr) {
+                return;
+            }
+
+            free_position_cache();
+
+            auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
+            auto txt_pe_vec    = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
+            auto img_pe_vec    = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
+            auto joint_pe_vec  = txt_pe_vec;
+            joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
+
+            ggml_init_params params;
+            params.mem_size    = static_cast<size_t>(3 * ggml_tensor_overhead());
+            params.mem_buffer  = nullptr;
+            params.no_alloc    = true;
+            position_cache_ctx = ggml_init(params);
+            GGML_ASSERT(position_cache_ctx != nullptr);
+
+            cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1);
+            ggml_set_name(cached_pos_embed, "minit2i.pos_embed");
+            cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len);
+            ggml_set_name(cached_txt_pe, "minit2i.txt_pe");
+            cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side);
+            ggml_set_name(cached_joint_pe, "minit2i.joint_pe");
+
+            position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend);
+            GGML_ASSERT(position_cache_buffer != nullptr);
+            ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+            ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed));
+            ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe));
+            ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe));
+            ggml_backend_synchronize(runtime_backend);
+
+            cached_img_side    = img_side;
+            cached_txt_len     = txt_len;
+            cached_hidden_size = config.hidden_size;
+            cached_head_dim    = config.head_dim;
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const sd::Tensor<float>& mask_tensor) {
+            ggml_cgraph* gf      = new_graph_custom(MINIT2I_GRAPH_SIZE);
+            ggml_tensor* x       = make_input(x_tensor);
+            ggml_tensor* context = make_input(context_tensor);
+            ggml_tensor* mask    = make_input(mask_tensor);
+            SD_UNUSED(timesteps_tensor);
+
+            int64_t W        = x->ne[0];
+            int64_t H        = x->ne[1];
+            int64_t img_side = H / config.patch_size;
+            int64_t txt_len  = context->ne[1];
+            ensure_position_cache(img_side, txt_len);
+
+            auto runner_ctx = get_context();
+            auto out        = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_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,
+                                  const sd::Tensor<float>& mask) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, mask);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            GGML_ASSERT(diffusion_params.context != nullptr);
+            const auto* extra = diffusion_extra_as<MiniT2IDiffusionExtra>(diffusion_params);
+            GGML_ASSERT(extra->mask != nullptr);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           *diffusion_params.context,
+                           *extra->mask);
+        }
+    };
+}  // namespace MiniT2I
+
+#endif  // __SD_MODEL_DIFFUSION_MINIT2I_HPP__

src/model/diffusion/model.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
 #define __SD_MODEL_DIFFUSION_MODEL_HPP__
 
 #include <string>
@@ -7,6 +7,7 @@
 
 #include "core/ggml_extend.hpp"
 #include "core/tensor_ggml.hpp"
+#include "model/common/rope.hpp"
 #include "model_manager.h"
 
 struct UNetDiffusionExtra {
@@ -52,6 +53,10 @@ struct LTXAVDiffusionExtra {
     const sd::Tensor<float>* video_positions = nullptr;
 };
 
+struct MiniT2IDiffusionExtra {
+    const sd::Tensor<float>* mask = nullptr;
+};
+
 using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
@@ -59,7 +64,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           AnimaDiffusionExtra,
                                           WanDiffusionExtra,
                                           HiDreamO1DiffusionExtra,
-                                          LTXAVDiffusionExtra>;
+                                          LTXAVDiffusionExtra,
+                                          MiniT2IDiffusionExtra>;
 
 struct DiffusionParams {
     const sd::Tensor<float>* x                        = nullptr;
@@ -68,7 +74,7 @@ struct DiffusionParams {
     const sd::Tensor<float>* c_concat                 = nullptr;
     const sd::Tensor<float>* y                        = nullptr;
     const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
-    bool increase_ref_index                           = false;
+    Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED;
     DiffusionExtraParams extra                        = std::monostate{};
 };
 

src/model/diffusion/qwen_image.hpp

@@ -4,6 +4,7 @@
 #include <memory>
 
 #include "model/common/block.hpp"
+#include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/model.hpp"
 #include "model_loader.h"
@@ -23,6 +24,7 @@ namespace Qwen {
         std::vector<int> axes_dim   = {16, 56, 56};
         int axes_dim_sum            = 128;
         bool zero_cond_t            = false;
+        bool use_additional_t_cond  = false;
 
         static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
             QwenImageConfig config;
@@ -88,19 +90,33 @@ namespace Qwen {
     };
 
     struct QwenTimestepProjEmbeddings : public GGMLBlock {
+    protected:
+        bool use_additional_t_cond = false;
+
     public:
-        QwenTimestepProjEmbeddings(int64_t embedding_dim) {
+        QwenTimestepProjEmbeddings(int64_t embedding_dim, bool use_additional_t_cond = false)
+            : use_additional_t_cond(use_additional_t_cond) {
             blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
+            if (use_additional_t_cond) {
+                blocks["addition_t_embedding"] = std::shared_ptr<GGMLBlock>(new Embedding(2, embedding_dim));
+            }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* timesteps) {
+                             ggml_tensor* timesteps,
+                             ggml_tensor* addition_t_cond = nullptr) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
 
             auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
             auto timesteps_emb  = timestep_embedder->forward(ctx, timesteps_proj);
+            if (use_additional_t_cond) {
+                GGML_ASSERT(addition_t_cond != nullptr);
+                auto addition_t_embedding = std::dynamic_pointer_cast<Embedding>(blocks["addition_t_embedding"]);
+                auto addition_t_emb       = addition_t_embedding->forward(ctx, addition_t_cond);
+                timesteps_emb             = ggml_add(ctx->ggml_ctx, timesteps_emb, addition_t_emb);
+            }
             return timesteps_emb;
         }
     };
@@ -402,7 +418,7 @@ namespace Qwen {
         QwenImageModel(QwenImageConfig config)
             : config(config) {
             int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
-            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
+            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
             blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
             blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
             blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
@@ -424,6 +440,7 @@ namespace Qwen {
         ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
                                   ggml_tensor* x,
                                   ggml_tensor* timestep,
+                                  ggml_tensor* addition_t_cond,
                                   ggml_tensor* context,
                                   ggml_tensor* pe,
                                   ggml_tensor* modulate_index = nullptr) {
@@ -434,9 +451,9 @@ namespace Qwen {
             auto norm_out        = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
             auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
 
-            auto t_emb = time_text_embed->forward(ctx, timestep);
+            auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
             if (config.zero_cond_t) {
-                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
+                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
                 t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
             }
             auto img = img_in->forward(ctx, x);
@@ -469,33 +486,50 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* x,
                              ggml_tensor* timestep,
+                             ggml_tensor* addition_t_cond,
                              ggml_tensor* context,
                              ggml_tensor* pe,
                              std::vector<ggml_tensor*> ref_latents = {},
                              ggml_tensor* modulate_index           = nullptr) {
             // Forward pass of DiT.
-            // x: [N, C, H, W]
+            // x: [N, C, H, W] or [N*C, T, H, W]
             // timestep: [N,]
             // context: [N, L, D]
             // pe: [L, d_head/2, 2, 2]
-            // return: [N, C, H, W]
+            // return: [N, C, H, W] or [N*C, T, H, W]
 
-            int64_t W = x->ne[0];
-            int64_t H = x->ne[1];
-            int64_t C = x->ne[2];
-            int64_t N = x->ne[3];
+            int64_t W          = x->ne[0];
+            int64_t H          = x->ne[1];
+            int64_t T          = 1;
+            int64_t N          = addition_t_cond != nullptr ? addition_t_cond->ne[0] : x->ne[3];
+            bool has_time_axis = false;
+            if (x->ne[3] != 1) {
+                T             = x->ne[2];
+                has_time_axis = true;
+            }
 
-            auto img           = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
+            auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
+                input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
+                if (!has_time_axis) {
+                    return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
+                }
+                if (input->ne[3] == 1) {
+                    input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
+                }
+                return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
+            };
+
+            auto img           = patchify_input(x);
             int64_t img_tokens = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
+                    ref = patchify_input(ref);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
 
-            auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
+            auto out = forward_orig(ctx, img, timestep, addition_t_cond, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
 
             if (out->ne[1] > img_tokens) {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [num_tokens, N, C * patch_size * patch_size]
@@ -503,7 +537,17 @@ namespace Qwen {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [N, h*w, C * patch_size * patch_size]
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+            if (has_time_axis) {
+                int pad_h = (config.patch_size - H % config.patch_size) % config.patch_size;
+                int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
+                int h_len = static_cast<int>((H + pad_h) / config.patch_size);
+                int w_len = static_cast<int>((W + pad_w) / config.patch_size);
+                out       = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
+                out       = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+                out       = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            } else {
+                out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+            }
 
             return out;
         }
@@ -515,6 +559,7 @@ namespace Qwen {
         QwenImageModel qwen_image;
         std::vector<float> pe_vec;
         std::vector<float> modulate_index_vec;
+        std::vector<int32_t> additional_t_cond_vec;
         SDVersion version;
 
         QwenImageRunner(ggml_backend_t backend,
@@ -524,9 +569,13 @@ namespace Qwen {
                         bool zero_cond_t                                    = false,
                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
-              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
+              version(version) {
             config.zero_cond_t = config.zero_cond_t || zero_cond_t;
-            qwen_image         = QwenImageModel(config);
+            if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                config.use_additional_t_cond = true;
+            }
+            qwen_image = QwenImageModel(config);
             qwen_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -542,11 +591,11 @@ namespace Qwen {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false) {
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::INCREASE) {
             ggml_cgraph* gf        = new_graph_custom(QWEN_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(x->ne[3] == 1 || x_tensor.dim() == 5);
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
             std::vector<ggml_tensor*> ref_latents;
@@ -555,13 +604,29 @@ namespace Qwen {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
+            int batch_size = static_cast<int>(x->ne[3]);
+            int time_len   = 1;
+            if (x_tensor.dim() == 5) {
+                time_len   = static_cast<int>(x_tensor.shape()[2]);
+                batch_size = static_cast<int>(x_tensor.shape()[4]);
+            }
+
+            ggml_tensor* addition_t_cond = nullptr;
+            if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                additional_t_cond_vec.assign(static_cast<size_t>(batch_size), 0);
+                addition_t_cond = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, batch_size);
+                set_backend_tensor_data(addition_t_cond, additional_t_cond_vec.data());
+                ref_index_mode = Rope::RefIndexMode::DECREASE;
+            }
+
+            pe_vec      = Rope::gen_qwen_image_pe(time_len,
+                                                  static_cast<int>(x->ne[1]),
                                                   static_cast<int>(x->ne[0]),
                                                   config.patch_size,
-                                                  static_cast<int>(x->ne[3]),
+                                                  batch_size,
                                                   static_cast<int>(context->ne[1]),
                                                   ref_latents,
-                                                  increase_ref_index,
+                                                  ref_index_mode,
                                                   config.theta,
                                                   circular_y_enabled,
                                                   circular_x_enabled,
@@ -604,6 +669,7 @@ namespace Qwen {
             ggml_tensor* out = qwen_image.forward(&runner_ctx,
                                                   x,
                                                   timesteps,
+                                                  addition_t_cond,
                                                   context,
                                                   pe,
                                                   ref_latents,
@@ -619,12 +685,12 @@ namespace Qwen {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context,
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false) {
-            // x: [N, in_channels, h, w]
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::INCREASE) {
+            // x: [N, C, H, W] or [N*C, T, H, W]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
+                return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -640,7 +706,7 @@ namespace Qwen {
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index);
+                           diffusion_params.ref_index_mode);
         }
 
         void test() {
@@ -674,7 +740,7 @@ namespace Qwen {
                                        timesteps,
                                        context,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());

src/model/diffusion/sefi_image.hpp

@@ -0,0 +1,91 @@
+#ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
+
+#include <memory>
+
+#include "model/common/block.hpp"
+
+namespace SefiImage {
+    struct SefiImageConfig {
+        int64_t semantic_channels        = 16;
+        int64_t texture_latent_channels  = 32;
+        int64_t timestep_guidance_in_dim = 256;
+        int64_t hidden_size              = 3072;
+        float timestep_shift_alpha       = 0.3f;
+        float delta_t                    = 0.1f;
+
+        int64_t packed_texture_channels(int patch_size) const {
+            return texture_latent_channels * patch_size * patch_size;
+        }
+
+        int64_t packed_input_channels(int patch_size) const {
+            return semantic_channels + packed_texture_channels(patch_size);
+        }
+
+        static SefiImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                   const std::string& prefix) {
+            SefiImageConfig config;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
+                    config.timestep_guidance_in_dim = tensor_storage.ne[0];
+                    config.hidden_size              = tensor_storage.ne[1] * 2;
+                }
+            }
+            LOG_DEBUG("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64,
+                      config.semantic_channels,
+                      config.texture_latent_channels,
+                      config.hidden_size);
+            return config;
+        }
+    };
+
+    struct SefiTimestepEmbedding : public GGMLBlock {
+    public:
+        SefiTimestepEmbedding(int64_t in_channels, int64_t time_embed_dim) {
+            blocks["linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, time_embed_dim, false));
+            blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim, false));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* sample) {
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
+
+            sample = linear_1->forward(ctx, sample);
+            sample = ggml_silu_inplace(ctx->ggml_ctx, sample);
+            sample = linear_2->forward(ctx, sample);
+            return sample;
+        }
+    };
+
+    struct SefiDualTimestepEmbeddings : public GGMLBlock {
+    public:
+        SefiDualTimestepEmbeddings(int64_t in_channels, int64_t embedding_dim) {
+            GGML_ASSERT(embedding_dim % 2 == 0);
+            int64_t half_dim            = embedding_dim / 2;
+            blocks["semantic_embedder"] = std::make_shared<SefiTimestepEmbedding>(in_channels, half_dim);
+            blocks["texture_embedder"]  = std::make_shared<SefiTimestepEmbedding>(in_channels, half_dim);
+            timestep_guidance_in_dim    = in_channels;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* timestep_sem,
+                             ggml_tensor* timestep_tex) {
+            auto semantic_embedder = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["semantic_embedder"]);
+            auto texture_embedder  = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["texture_embedder"]);
+
+            auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto sem_emb  = semantic_embedder->forward(ctx, sem_proj);
+            auto tex_emb  = texture_embedder->forward(ctx, tex_proj);
+            return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0);
+        }
+
+    private:
+        int64_t timestep_guidance_in_dim = 256;
+    };
+}  // namespace SefiImage
+
+#endif  // __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__

src/model/diffusion/z_image.hpp

@@ -575,7 +575,7 @@ namespace ZImage {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false) {
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::FIXED) {
             ggml_cgraph* gf        = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
             ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -595,7 +595,7 @@ namespace ZImage {
                                                static_cast<int>(context->ne[1]),
                                                SEQ_MULTI_OF,
                                                ref_latents,
-                                               increase_ref_index,
+                                               ref_index_mode,
                                                config.theta,
                                                circular_y_enabled,
                                                circular_x_enabled,
@@ -626,12 +626,12 @@ namespace ZImage {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context,
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false) {
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED) {
             // x: [N, in_channels, h, w]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
+                return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -647,7 +647,7 @@ namespace ZImage {
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index);
+                           diffusion_params.ref_index_mode);
         }
 
         void test() {
@@ -681,7 +681,7 @@ namespace ZImage {
                                        timesteps,
                                        context,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());

src/model/te/clip.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_TE_CLIP_HPP__
+#ifndef __SD_MODEL_TE_CLIP_HPP__
 #define __SD_MODEL_TE_CLIP_HPP__
 
 #include "core/ggml_extend.hpp"

src/model/te/llm.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_TE_LLM_HPP__
+#ifndef __SD_MODEL_TE_LLM_HPP__
 #define __SD_MODEL_TE_LLM_HPP__
 
 #include <algorithm>
@@ -79,6 +79,7 @@ namespace LLM {
         int window_size                     = 112;
         int num_position_embeddings         = 0;
         std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
+        bool split_patch_embed              = false;
     };
 
     struct LLMConfig {
@@ -179,7 +180,8 @@ namespace LLM {
                 config.num_experts_per_tok     = 4;
             }
 
-            config.num_layers = 0;
+            config.num_layers          = 0;
+            int detected_vision_layers = 0;
             for (const auto& [name, tensor_storage] : tensor_storage_map) {
                 if (!starts_with(name, prefix)) {
                     continue;
@@ -190,6 +192,38 @@ namespace LLM {
                     if (contains(name, "attn.q_proj")) {
                         config.llama_cpp_style = true;
                     }
+                    if (contains(name, "visual.patch_embed.proj.1.weight")) {
+                        config.vision.split_patch_embed = true;
+                    }
+                    if (contains(name, "visual.patch_embed.proj.0.weight")) {
+                        config.vision.patch_size  = static_cast<int>(tensor_storage.ne[0]);
+                        config.vision.in_channels = tensor_storage.ne[2];
+                        config.vision.hidden_size = tensor_storage.ne[3];
+                    }
+                    if (contains(name, "visual.patch_embed.bias")) {
+                        config.vision.hidden_size = tensor_storage.ne[0];
+                    }
+                    if (contains(name, "visual.pos_embed.weight")) {
+                        config.vision.hidden_size             = tensor_storage.ne[0];
+                        config.vision.num_position_embeddings = static_cast<int>(tensor_storage.ne[1]);
+                    }
+                    if (contains(name, "visual.blocks.")) {
+                        auto items = split_string(name.substr(pos), '.');
+                        if (items.size() > 2) {
+                            int block_index = atoi(items[2].c_str());
+                            if (block_index + 1 > detected_vision_layers) {
+                                detected_vision_layers = block_index + 1;
+                            }
+                        }
+                    }
+                    if (contains(name, "visual.blocks.0.mlp.linear_fc1.weight") ||
+                        contains(name, "visual.blocks.0.mlp.gate_proj.weight")) {
+                        config.vision.intermediate_size = tensor_storage.ne[1];
+                    }
+                    if (contains(name, "visual.merger.linear_fc2.weight") ||
+                        contains(name, "visual.merger.mlp.2.weight")) {
+                        config.vision.out_hidden_size = tensor_storage.ne[1];
+                    }
                     continue;
                 }
                 pos = name.find("layers.");
@@ -216,9 +250,12 @@ namespace LLM {
                     config.intermediate_size = tensor_storage.ne[1];
                 }
             }
-            if (arch == LLMArch::QWEN3 && config.num_layers == 28) {
+            if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
                 config.num_heads = 16;
             }
+            if (detected_vision_layers > 0) {
+                config.vision.num_layers = detected_vision_layers;
+            }
             LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
                       config.num_layers,
                       config.vocab_size,
@@ -539,40 +576,51 @@ namespace LLM {
 
     struct VisionPatchEmbed : public GGMLBlock {
     protected:
-        bool llama_cpp_style;
+        bool split_patch_embed;
+        bool bias;
         int patch_size;
         int temporal_patch_size;
         int64_t in_channels;
         int64_t embed_dim;
 
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            GGML_UNUSED(tensor_storage_map);
+            GGML_UNUSED(prefix);
+            if (split_patch_embed && bias) {
+                params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim);
+            }
+        }
+
     public:
-        VisionPatchEmbed(bool llama_cpp_style,
+        VisionPatchEmbed(bool split_patch_embed,
                          LLMVisionArch arch,
                          int patch_size          = 14,
                          int temporal_patch_size = 2,
                          int64_t in_channels     = 3,
                          int64_t embed_dim       = 1152)
-            : llama_cpp_style(llama_cpp_style),
+            : split_patch_embed(split_patch_embed),
+              bias(arch == LLMVisionArch::QWEN3_VL),
               patch_size(patch_size),
               temporal_patch_size(temporal_patch_size),
               in_channels(in_channels),
               embed_dim(embed_dim) {
-            bool bias = arch == LLMVisionArch::QWEN3_VL;
-            if (llama_cpp_style) {
+            if (split_patch_embed) {
                 blocks["proj.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
                                                                          embed_dim,
                                                                          {patch_size, patch_size},
                                                                          {patch_size, patch_size},
                                                                          {0, 0},
                                                                          {1, 1},
-                                                                         bias));
+                                                                         false));
                 blocks["proj.1"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
                                                                          embed_dim,
                                                                          {patch_size, patch_size},
                                                                          {patch_size, patch_size},
                                                                          {0, 0},
                                                                          {1, 1},
-                                                                         bias));
+                                                                         false));
             } else {
                 std::tuple<int, int, int> kernel_size = {(int)temporal_patch_size, (int)patch_size, (int)patch_size};
                 blocks["proj"]                        = std::shared_ptr<GGMLBlock>(new Conv3d(in_channels,
@@ -593,7 +641,7 @@ namespace LLM {
                                 temporal_patch_size,
                                 ggml_nelements(x) / (temporal_patch_size * patch_size * patch_size));
 
-            if (llama_cpp_style) {
+            if (split_patch_embed) {
                 auto proj_0 = std::dynamic_pointer_cast<Conv2d>(blocks["proj.0"]);
                 auto proj_1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj.1"]);
 
@@ -606,6 +654,10 @@ namespace LLM {
                 x1      = proj_1->forward(ctx, x1);
 
                 x = ggml_add(ctx->ggml_ctx, x0, x1);
+                if (bias) {
+                    auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, embed_dim, 1);
+                    x      = ggml_add_inplace(ctx->ggml_ctx, x, b);
+                }
             } else {
                 auto proj = std::dynamic_pointer_cast<Conv3d>(blocks["proj"]);
 
@@ -798,7 +850,7 @@ namespace LLM {
               spatial_merge_size(vision_params.spatial_merge_size),
               num_grid_per_side(vision_params.num_position_embeddings > 0 ? static_cast<int>(std::sqrt(vision_params.num_position_embeddings)) : 0),
               fullatt_block_indexes(vision_params.fullatt_block_indexes) {
-            blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(llama_cpp_style,
+            blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(vision_params.split_patch_embed,
                                                                                     arch_,
                                                                                     vision_params.patch_size,
                                                                                     vision_params.temporal_patch_size,

src/model/te/t5.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_TE_T5_HPP__
+#ifndef __SD_MODEL_TE_T5_HPP__
 #define __SD_MODEL_TE_T5_HPP__
 
 #include <cfloat>
@@ -26,13 +26,66 @@ struct T5Config {
     static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
                                         const std::string& prefix,
                                         bool is_umt5 = false) {
-        (void)tensor_storage_map;
-        (void)prefix;
         T5Config config;
         if (is_umt5) {
             config.vocab_size         = 256384;
             config.relative_attention = false;
         }
+        auto find_tensor = [&](const std::string& suffix) -> const TensorStorage* {
+            auto it = tensor_storage_map.find(prefix + "." + suffix);
+            if (it != tensor_storage_map.end()) {
+                return &it->second;
+            }
+            it = tensor_storage_map.find(prefix + suffix);
+            if (it != tensor_storage_map.end()) {
+                return &it->second;
+            }
+            return nullptr;
+        };
+
+        if (const TensorStorage* shared = find_tensor("shared.weight")) {
+            if (shared->n_dims == 2) {
+                config.vocab_size = shared->ne[1];
+                config.model_dim  = shared->ne[0];
+            }
+        }
+        if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) {
+            if (q->n_dims == 2) {
+                config.model_dim  = q->ne[0];
+                int64_t inner_dim = q->ne[1];
+                // Flan-T5/T5 uses d_kv=64 for common sizes.
+                if (inner_dim % 64 == 0) {
+                    config.num_heads = inner_dim / 64;
+                }
+            }
+        }
+        if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) {
+            if (wi->n_dims == 2) {
+                config.model_dim = wi->ne[0];
+                config.ff_dim    = wi->ne[1];
+            }
+        }
+        int64_t detected_layers = 0;
+        for (const auto& [name, _] : tensor_storage_map) {
+            std::string base = prefix;
+            if (!base.empty() && base.back() != '.') {
+                base += ".";
+            }
+            std::string layer_prefix = base + "encoder.block.";
+            if (!starts_with(name, layer_prefix)) {
+                continue;
+            }
+            size_t pos = layer_prefix.size();
+            size_t dot = name.find('.', pos);
+            if (dot == std::string::npos) {
+                continue;
+            }
+            int64_t layer   = atoi(name.substr(pos, dot - pos).c_str());
+            detected_layers = std::max(detected_layers, layer + 1);
+        }
+        if (detected_layers > 0) {
+            config.num_layers = detected_layers;
+        }
         return config;
     }
 };

src/model/upscaler/esrgan.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
+#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
 #define __SD_MODEL_UPSCALER_ESRGAN_HPP__
 
 #include <algorithm>

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
+#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
 #define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
 
 #include <algorithm>

src/model/vae/auto_encoder_kl.hpp

@@ -682,7 +682,7 @@ struct AutoEncoderKL : public VAE {
         } else if (sd_version_is_sd3(version)) {
             scale_factor = 1.5305f;
             shift_factor = 0.0609f;
-        } else if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_longcat(version)) {
+        } else if (sd_version_uses_flux_vae(version)) {
             scale_factor = 0.3611f;
             shift_factor = 0.1159f;
         } else if (sd_version_uses_flux2_vae(version)) {
@@ -816,12 +816,13 @@ struct AutoEncoderKL : public VAE {
     }
 
     sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+        auto latents_ = sd_version_is_sefi_image(version) ? sd::ops::slice(latents, 2, 16, 144) : latents;
         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;
+            auto [mean_tensor, std_tensor] = get_latents_mean_std(latents_, channel_dim);
+            return (latents_ * std_tensor) / scale_factor + mean_tensor;
         }
-        return (latents / scale_factor) + shift_factor;
+        return (latents_ / scale_factor) + shift_factor;
     }
 
     sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {

src/model/vae/ltx_audio_vae.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
+#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 #define __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 
 #include <cmath>
@@ -214,7 +214,7 @@ namespace LTXV {
 
         auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
         if (left_pad > 0) {
-            x = ggml_pad_ext(ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
+            x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
         }
 
         auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
@@ -451,6 +451,7 @@ namespace LTXV {
             int pad_h = kernel_size.first - 1;
             int pad_w = kernel_size.second - 1;
             x         = ggml_ext_pad_ext(ctx->ggml_ctx,
+                                         ctx->backend,
                                          x,
                                          pad_w / 2,
                                          pad_w - pad_w / 2,

src/model/vae/tae.hpp

@@ -408,7 +408,7 @@ public:
             h = conv->forward(ctx, h);
             for (int j = 0; j < num_blocks; j++) {
                 auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
-                auto mem   = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                auto mem   = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
                 mem        = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
                 h          = block->forward(ctx, h, mem);
             }
@@ -479,7 +479,7 @@ public:
         int index = 3;
         for (int i = 0; i < num_layers; i++) {
             for (int j = 0; j < num_blocks; j++) {
-                auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
                 mem      = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
                 if (is_wide) {
                     auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
@@ -548,7 +548,7 @@ public:
         }
         auto result = decoder->forward(ctx, z);
         if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
-            // (W, H, C, T) -> (W, H, T, C)
+            // (W, H, T, C) -> (W, H, C, T)
             result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
         }
         return result;
@@ -556,8 +556,10 @@ public:
 
     ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
         auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
-        // (W, H, T, C) -> (W, H, C, T)
-        x                  = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+            // (W, H, T, C) -> (W, H, C, T)
+            x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        }
         int64_t num_frames = x->ne[3];
         if (num_frames % encoder->t_downscale) {
             // pad to multiple of encoder->t_downscale at the end
@@ -567,7 +569,10 @@ public:
             }
         }
         x = encoder->forward(ctx, x);
-        x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+            // (W, H, C, T) -> (W, H, T, C)
+            x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        }
         return x;
     }
 };

src/model/vae/vae.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_VAE_VAE_HPP__
+#ifndef __SD_MODEL_VAE_VAE_HPP__
 #define __SD_MODEL_VAE_VAE_HPP__
 
 #include "core/tensor_ggml.hpp"
@@ -78,7 +78,7 @@ public:
             scale_factor = 16;
         } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
-        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
+        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
             scale_factor = 1;
         }
         return scale_factor;

src/model/vae/wan_vae.hpp

@@ -72,8 +72,8 @@ namespace WAN {
                 lp2 -= (int)cache_x->ne[2];
             }
 
-            x = ggml_ext_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-            return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
+            x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
                                     std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation));
@@ -113,6 +113,24 @@ namespace WAN {
         }
     };
 
+    class Conv2dBut3d : public Conv2d {
+    public:
+        using Conv2d::Conv2d;
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            ggml_tensor* x_swapped = ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2);
+            x_swapped              = ggml_cont(ctx->ggml_ctx, x_swapped);
+
+            ggml_tensor* out = Conv2d::forward(ctx, x_swapped);
+
+            ggml_tensor* out_swapped = ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2);
+
+            out_swapped = ggml_cont(ctx->ggml_ctx, out_swapped);
+
+            return out_swapped;
+        }
+    };
+
     class Resample : public GGMLBlock {
     protected:
         int64_t dim;
@@ -177,7 +195,7 @@ namespace WAN {
                                                   2);
                         }
                         if (chunk_idx == 1 && cache_x->ne[2] < 2) {  // Rep
-                            cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
+                            cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
                             // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2)
                         }
                         if (chunk_idx == 1) {
@@ -265,7 +283,7 @@ namespace WAN {
 
             int pad_t = (factor_t - T % factor_t) % factor_t;
 
-            x = ggml_pad_ext(ctx->ggml_ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
+            x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
             T = x->ne[2];
 
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C);                                                  // [C, T/factor_t, factor_t, H*W]
@@ -338,19 +356,32 @@ namespace WAN {
     protected:
         int64_t in_dim;
         int64_t out_dim;
+        bool is_2D;
 
     public:
-        ResidualBlock(int64_t in_dim, int64_t out_dim)
-            : in_dim(in_dim), out_dim(out_dim) {
+        ResidualBlock(int64_t in_dim, int64_t out_dim, bool is_2D = false)
+            : in_dim(in_dim), out_dim(out_dim), is_2D(is_2D) {
             blocks["residual.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(in_dim));
             // residual.1 is nn.SiLU()
-            blocks["residual.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            if (is_2D) {
+                blocks["residual.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_dim, out_dim, {3, 3}, {1, 1}, {1, 1}));
+            } else {
+                blocks["residual.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            }
             blocks["residual.3"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
             // residual.4 is nn.SiLU()
             // residual.5 is nn.Dropout()
-            blocks["residual.6"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            if (is_2D) {
+                blocks["residual.6"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, out_dim, {3, 3}, {1, 1}, {1, 1}));
+            } else {
+                blocks["residual.6"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            }
             if (in_dim != out_dim) {
-                blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {1, 1, 1}));
+                if (is_2D) {
+                    blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_dim, out_dim, {1, 1}));
+                } else {
+                    blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_dim, out_dim, {1, 1, 1}));
+                }
             }
         }
 
@@ -363,9 +394,15 @@ namespace WAN {
             GGML_ASSERT(b == 1);
             ggml_tensor* h = x;
             if (in_dim != out_dim) {
-                auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
+                if (is_2D) {
+                    auto shortcut = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["shortcut"]);
 
-                h = shortcut->forward(ctx, x);
+                    h = shortcut->forward(ctx, x);
+                } else {
+                    auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
+
+                    h = shortcut->forward(ctx, x);
+                }
             }
 
             for (int i = 0; i < 7; i++) {
@@ -385,8 +422,13 @@ namespace WAN {
                                                   cache_x,
                                                   2);
                         }
+                        if (is_2D) {
+                            auto layer = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["residual." + std::to_string(i)]);
 
-                        x               = layer->forward(ctx, x, feat_cache[idx]);
+                            x = layer->forward(ctx, x);
+                        } else {
+                            x = layer->forward(ctx, x, feat_cache[idx]);
+                        }
                         feat_cache[idx] = cache_x;
                         feat_idx += 1;
                     }
@@ -412,13 +454,14 @@ namespace WAN {
                            int64_t out_dim,
                            int mult,
                            bool temperal_downsample = false,
-                           bool down_flag           = false)
+                           bool down_flag           = false,
+                           bool is_2D               = false)
             : mult(mult), down_flag(down_flag) {
             blocks["avg_shortcut"] = std::shared_ptr<GGMLBlock>(new AvgDown3D(in_dim, out_dim, temperal_downsample ? 2 : 1, down_flag ? 2 : 1));
 
             int i = 0;
             for (; i < mult; i++) {
-                blocks["downsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
+                blocks["downsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
                 in_dim                                     = out_dim;
             }
             if (down_flag) {
@@ -472,7 +515,8 @@ namespace WAN {
                          int64_t out_dim,
                          int mult,
                          bool temperal_upsample = false,
-                         bool up_flag           = false)
+                         bool up_flag           = false,
+                         bool is_2D             = false)
             : mult(mult), up_flag(up_flag) {
             if (up_flag) {
                 blocks["avg_shortcut"] = std::shared_ptr<GGMLBlock>(new DupUp3D(in_dim, out_dim, temperal_upsample ? 2 : 1, up_flag ? 2 : 1));
@@ -480,7 +524,7 @@ namespace WAN {
 
             int i = 0;
             for (; i < mult; i++) {
-                blocks["upsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
+                blocks["upsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
                 in_dim                                   = out_dim;
             }
             if (up_flag) {
@@ -587,35 +631,41 @@ namespace WAN {
     class Encoder3d : public GGMLBlock {
     protected:
         bool wan2_2;
+        int64_t in_channels;
         int64_t dim;
         int64_t z_dim;
         std::vector<int> dim_mult;
         int num_res_blocks;
         std::vector<bool> temperal_downsample;
+        bool is_2D = false;
 
     public:
         Encoder3d(int64_t dim                           = 128,
                   int64_t z_dim                         = 4,
+                  int64_t in_channels                   = 3,
                   std::vector<int> dim_mult             = {1, 2, 4, 4},
                   int num_res_blocks                    = 2,
                   std::vector<bool> temperal_downsample = {false, true, true},
-                  bool wan2_2                           = false)
-            : dim(dim),
+                  bool wan2_2                           = false,
+                  bool is_2D                            = false)
+            : in_channels(in_channels),
+              dim(dim),
               z_dim(z_dim),
               dim_mult(dim_mult),
               num_res_blocks(num_res_blocks),
               temperal_downsample(temperal_downsample),
-              wan2_2(wan2_2) {
+              wan2_2(wan2_2),
+              is_2D(is_2D) {
             // attn_scales is always []
             std::vector<int64_t> dims = {dim};
             for (int u : dim_mult) {
                 dims.push_back(dim * u);
             }
 
-            if (wan2_2) {
-                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(12, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            if (is_2D) {
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_channels, dims[0], {3, 3}, {1, 1}, {1, 1}));
             } else {
-                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(3, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_channels, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
             }
 
             int index = 0;
@@ -630,12 +680,13 @@ namespace WAN {
                                                                                          out_dim,
                                                                                          num_res_blocks,
                                                                                          t_down_flag,
-                                                                                         i != dim_mult.size() - 1));
+                                                                                         i != dim_mult.size() - 1,
+                                                                                         is_2D));
 
                     blocks["downsamples." + std::to_string(index++)] = block;
                 } else {
                     for (int j = 0; j < num_res_blocks; j++) {
-                        auto block                                       = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
+                        auto block                                       = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
                         blocks["downsamples." + std::to_string(index++)] = block;
                         in_dim                                           = out_dim;
                     }
@@ -648,13 +699,17 @@ namespace WAN {
                 }
             }
 
-            blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim));
+            blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim, is_2D));
             blocks["middle.1"] = std::shared_ptr<GGMLBlock>(new AttentionBlock(out_dim));
-            blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim));
+            blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(out_dim, out_dim, is_2D));
 
             blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
             // head.1 is nn.SiLU()
-            blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            if (is_2D) {
+                blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, z_dim, {3, 3}, {1, 1}, {1, 1}));
+            } else {
+                blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -673,7 +728,10 @@ namespace WAN {
             auto head_2   = std::dynamic_pointer_cast<CausalConv3d>(blocks["head.2"]);
 
             // conv1
-            if (feat_cache.size() > 0) {
+            if (is_2D) {
+                auto conv1 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv1"]);
+                x          = conv1->forward(ctx, x);
+            } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
                 auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@@ -728,7 +786,10 @@ namespace WAN {
             // head
             x = head_0->forward(ctx, x);
             x = ggml_silu(ctx->ggml_ctx, x);
-            if (feat_cache.size() > 0) {
+            if (is_2D) {
+                auto head_2 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["head.2"]);
+                x           = head_2->forward(ctx, x);
+            } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
                 auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@@ -753,25 +814,31 @@ namespace WAN {
     class Decoder3d : public GGMLBlock {
     protected:
         bool wan2_2;
+        int64_t out_channels;
         int64_t dim;
         int64_t z_dim;
         std::vector<int> dim_mult;
         int num_res_blocks;
         std::vector<bool> temperal_upsample;
+        bool is_2D = false;
 
     public:
         Decoder3d(int64_t dim                         = 128,
                   int64_t z_dim                       = 4,
+                  int64_t out_channels                = 3,
                   std::vector<int> dim_mult           = {1, 2, 4, 4},
                   int num_res_blocks                  = 2,
                   std::vector<bool> temperal_upsample = {true, true, false},
-                  bool wan2_2                         = false)
-            : dim(dim),
+                  bool wan2_2                         = false,
+                  bool is_2D                          = false)
+            : out_channels(out_channels),
+              dim(dim),
               z_dim(z_dim),
               dim_mult(dim_mult),
               num_res_blocks(num_res_blocks),
               temperal_upsample(temperal_upsample),
-              wan2_2(wan2_2) {
+              wan2_2(wan2_2),
+              is_2D(is_2D) {
             // attn_scales is always []
             std::vector<int64_t> dims = {dim_mult[dim_mult.size() - 1] * dim};
             for (int i = static_cast<int>(dim_mult.size()) - 1; i >= 0; i--) {
@@ -779,12 +846,16 @@ namespace WAN {
             }
 
             // init block
-            blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            if (is_2D) {
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, dims[0], {3, 3}, {1, 1}, {1, 1}));
+            } else {
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            }
 
             // middle blocks
-            blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0]));
+            blocks["middle.0"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0], is_2D));
             blocks["middle.1"] = std::shared_ptr<GGMLBlock>(new AttentionBlock(dims[0]));
-            blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0]));
+            blocks["middle.2"] = std::shared_ptr<GGMLBlock>(new ResidualBlock(dims[0], dims[0], is_2D));
 
             // upsample blocks
             int index = 0;
@@ -799,7 +870,8 @@ namespace WAN {
                                                                                      out_dim,
                                                                                      num_res_blocks + 1,
                                                                                      t_up_flag,
-                                                                                     i != dim_mult.size() - 1));
+                                                                                     i != dim_mult.size() - 1,
+                                                                                     is_2D));
 
                     blocks["upsamples." + std::to_string(index++)] = block;
                 } else {
@@ -807,7 +879,7 @@ namespace WAN {
                         in_dim = in_dim / 2;
                     }
                     for (int j = 0; j < num_res_blocks + 1; j++) {
-                        auto block                                     = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim));
+                        auto block                                     = std::shared_ptr<GGMLBlock>(new ResidualBlock(in_dim, out_dim, is_2D));
                         blocks["upsamples." + std::to_string(index++)] = block;
                         in_dim                                         = out_dim;
                     }
@@ -821,13 +893,14 @@ namespace WAN {
             }
 
             // output blocks
-            blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
+            blocks["head.0"]  = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
+            int64_t final_dim = out_channels;
             // head.1 is nn.SiLU()
-            if (wan2_2) {
-                blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, 12, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+            if (is_2D) {
+                blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1}));
 
             } else {
-                blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, 3, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+                blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, final_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
             }
         }
 
@@ -847,7 +920,10 @@ namespace WAN {
             auto head_2   = std::dynamic_pointer_cast<CausalConv3d>(blocks["head.2"]);
 
             // conv1
-            if (feat_cache.size() > 0) {
+            if (is_2D) {
+                auto conv1 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv1"]);
+                x          = conv1->forward(ctx, x);
+            } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
                 auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@@ -902,7 +978,10 @@ namespace WAN {
             // head
             x = head_0->forward(ctx, x);
             x = ggml_silu(ctx->ggml_ctx, x);
-            if (feat_cache.size() > 0) {
+            if (is_2D) {
+                auto head_2 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["head.2"]);
+                x           = head_2->forward(ctx, x);
+            } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
                 auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
@@ -928,6 +1007,8 @@ namespace WAN {
     public:
         bool wan2_2                           = false;
         bool decode_only                      = true;
+        int64_t input_channels                = 3;
+        int patch_size                        = 1;
         int64_t dim                           = 96;
         int64_t dec_dim                       = 96;
         int64_t z_dim                         = 16;
@@ -935,6 +1016,7 @@ namespace WAN {
         int num_res_blocks                    = 2;
         std::vector<bool> temperal_upsample   = {true, true, false};
         std::vector<bool> temperal_downsample = {false, true, true};
+        bool is_2D                            = false;
 
         int _conv_num = 33;
         int _conv_idx = 0;
@@ -951,23 +1033,43 @@ namespace WAN {
         }
 
     public:
-        WanVAE(bool decode_only = true, bool wan2_2 = false)
-            : decode_only(decode_only), wan2_2(wan2_2) {
+        WanVAE(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_2D = false)
+            : decode_only(decode_only),
+              wan2_2(version == VERSION_WAN2_2_TI2V),
+              is_2D(is_2D) {
             // attn_scales is always []
             if (wan2_2) {
-                dim     = 160;
-                dec_dim = 256;
-                z_dim   = 48;
+                dim            = 160;
+                dec_dim        = 256;
+                z_dim          = 48;
+                input_channels = 12;
+                patch_size     = 2;
 
                 _conv_num     = 34;
                 _enc_conv_num = 26;
+            } else if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                input_channels = 4;
             }
+
+            if (is_2D) {
+                temperal_upsample   = {false, false, false};
+                temperal_downsample = {false, false, false};
+            }
+
             if (!decode_only) {
-                blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, temperal_downsample, wan2_2));
-                blocks["conv1"]   = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1}));
+                blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, input_channels, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D));
+                if (is_2D) {
+                    blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1}));
+                } else {
+                    blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1}));
+                }
+            }
+            blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, input_channels, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D));
+            if (is_2D) {
+                blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, z_dim, {1, 1}));
+            } else {
+                blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
             }
-            blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2));
-            blocks["conv2"]   = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
         }
 
         static ggml_tensor* patchify(ggml_context* ctx,
@@ -1030,11 +1132,13 @@ namespace WAN {
             GGML_ASSERT(b == 1);
             GGML_ASSERT(decode_only == false);
 
+            if (x->ne[2] > 1 && is_2D) {
+                LOG_WARN("Using 2D VAE to encode video, expect poor results");
+            }
+
             clear_cache();
 
-            if (wan2_2) {
-                x = patchify(ctx->ggml_ctx, x, 2, b);
-            }
+            x = patchify(ctx->ggml_ctx, x, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
 
             auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
@@ -1049,12 +1153,18 @@ namespace WAN {
                     auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);  // [b*c, 1, h, w]
                     out     = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
                 } else {
-                    auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), 1 + 4 * i);  // [b*c, 4, h, w]
+                    // if is_2D, drop 3 out of 4 frames
+                    auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i);  // [b*c, 4, h, w]
                     auto out_ = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
                     out       = ggml_concat(ctx->ggml_ctx, out, out_, 2);
                 }
             }
-            out     = conv1->forward(ctx, out);
+            if (is_2D) {
+                auto conv1 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv1"]);
+                out        = conv1->forward(ctx, out);
+            } else {
+                out = conv1->forward(ctx, out);
+            }
             auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0];
             // sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu");
             clear_cache();
@@ -1067,13 +1177,23 @@ namespace WAN {
             // z: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
 
+            if (z->ne[2] > 1 && is_2D) {
+                LOG_WARN("Using 2D VAE to decode video, expect poor results");
+            }
+
             clear_cache();
 
             auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
             auto conv2   = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
 
             int64_t iter_ = z->ne[2];
-            auto x        = conv2->forward(ctx, z);
+            auto x        = z;
+            if (is_2D) {
+                auto conv2 = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv2"]);
+                x          = conv2->forward(ctx, z);
+            } else {
+                x = conv2->forward(ctx, z);
+            }
             // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode.prelude", "x");
             ggml_tensor* out;
             for (int i = 0; i < iter_; i++) {
@@ -1085,11 +1205,15 @@ namespace WAN {
                     auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
                     auto out_ = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
                     out       = ggml_concat(ctx->ggml_ctx, out, out_, 2);
+                    if (is_2D) {
+                        // repeat frames to avoid mismatch
+                        for (int j = 0; j < 4 - 1; j++) {
+                            out = ggml_concat(ctx->ggml_ctx, out, out_, 2);
+                        }
+                    }
                 }
             }
-            if (wan2_2) {
-                out = unpatchify(ctx->ggml_ctx, out, 2, b);
-            }
+            out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
             clear_cache();
             return out;
@@ -1110,9 +1234,7 @@ namespace WAN {
             auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
             _conv_idx = 0;
             auto out  = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
-            if (wan2_2) {
-                out = unpatchify(ctx->ggml_ctx, out, 2, b);
-            }
+            out       = unpatchify(ctx->ggml_ctx, out, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
             return out;
         }
@@ -1129,7 +1251,20 @@ namespace WAN {
                      bool decode_only                                    = false,
                      SDVersion version                                   = VERSION_WAN2,
                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-            : VAE(version, backend, prefix, weight_manager), decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V) {
+            : VAE(version, backend, prefix, weight_manager), decode_only(decode_only) {
+            bool is_2D = false;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (ends_with(name, "decoder.conv1.weight")) {
+                    if (tensor_storage.ne[2] > 3) {
+                        is_2D = true;
+                    }
+                    break;
+                }
+            }
+            if (is_2D) {
+                LOG_DEBUG("USING 2D VAE");
+            }
+            ae = WanVAE(decode_only, version, is_2D);
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
 

src/model_loader.cpp

@@ -20,6 +20,7 @@
 #include "model_io/torch_legacy_io.h"
 #include "model_io/torch_zip_io.h"
 #include "model_loader.h"
+#include "runtime/imatrix.h"
 #include "stable-diffusion.h"
 
 #include "core/ggml_extend_backend.h"
@@ -66,7 +67,6 @@ const char* unused_tensors[] = {
     // "v_pred", // Used to detect SDXL vpred models
     "text_encoders.llm.output.weight",
     "text_encoders.llm.lm_head.",
-    "first_stage_model.bn.",
 };
 
 bool is_unused_tensor(const std::string& name) {
@@ -157,7 +157,8 @@ void convert_tensor(void* src,
                     void* dst,
                     ggml_type dst_type,
                     int nrows,
-                    int n_per_row) {
+                    int n_per_row,
+                    std::vector<float> imatrix = {}) {
     int n = nrows * n_per_row;
     if (src_type == dst_type) {
         size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
@@ -166,7 +167,7 @@ void convert_tensor(void* src,
         if (dst_type == GGML_TYPE_F16) {
             ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
         } else {
-            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            imatrix.resize(n_per_row, 1.0f);
             const float* im = imatrix.data();
             ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im);
         }
@@ -196,7 +197,7 @@ void convert_tensor(void* src,
         if (dst_type == GGML_TYPE_F16) {
             ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
         } else {
-            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            imatrix.resize(n_per_row, 1.0f);
             const float* im = imatrix.data();
             ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im);
         }
@@ -453,6 +454,10 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("embed_image_indicator.weight") != std::string::npos) {
             return VERSION_IDEOGRAM4;
         }
+        if (tensor_storage.name.find("model.diffusion_model.txtfusion.projector.weight") != std::string::npos ||
+            tensor_storage.name.find("model.diffusion_model.text_fusion.projector.weight") != std::string::npos) {
+            return VERSION_KREA2;
+        }
         if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
             return VERSION_CHROMA_RADIANCE;
         }
@@ -467,7 +472,13 @@ SDVersion ModelLoader::get_sd_version() {
             tensor_storage_map.find("model.diffusion_model.transformer_blocks.0.img_mlp.w1.weight") != tensor_storage_map.end()) {
             return VERSION_LENS;
         }
+        if (tensor_storage.name.find("net.img_embedder.proj1.weight") != std::string::npos) {
+            return VERSION_MINIT2I;
+        }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
+            if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
+                return VERSION_QWEN_IMAGE_LAYERED;
+            }
             return VERSION_QWEN_IMAGE;
         }
         if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
@@ -476,6 +487,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) {
             is_flux2 = true;
         }
+        if (tensor_storage.name.find("dual_time_embed.semantic_embedder.linear_1.weight") != std::string::npos) {
+            return VERSION_SEFI_IMAGE;
+        }
         if (tensor_storage.name.find("single_blocks.47.linear1.weight") != std::string::npos) {
             has_single_block_47 = true;
         }
@@ -485,6 +499,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("double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") != std::string::npos) {
+            return VERSION_BOOGU_IMAGE;
+        }
         if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
             return VERSION_ERNIE_IMAGE;
         }
@@ -961,6 +978,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
     size_t total_tensors_processed = 0;
     const int64_t t_start          = start_time;
     int last_n_threads             = 1;
+    SDVersion imatrix_version      = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
 
     for (size_t file_index = 0; file_index < file_data.size(); ++file_index) {
         auto& fdata                  = file_data[file_index];
@@ -1145,12 +1163,15 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                             failed = true;
                             return;
                         }
+                        std::string processed_name = convert_tensor_name(tensor_storage.name, imatrix_version);
+                        std::vector<float> imatrix = get_imatrix_collector().get_values(processed_name);
                         convert_tensor((void*)target_buf,
                                        tensor_storage.type,
                                        convert_buf,
                                        dst_tensor->type,
                                        (int)tensor_storage.nelements() / (int)tensor_storage.ne[0],
-                                       (int)tensor_storage.ne[0]);
+                                       (int)tensor_storage.ne[0],
+                                       std::move(imatrix));
                     } else {
                         convert_buf = read_buf;
                     }

src/model_manager.cpp

@@ -147,6 +147,17 @@ bool ModelManager::register_param_tensors(const std::string& desc,
     return true;
 }
 
+bool ModelManager::load_all_params_eagerly() {
+    std::vector<TensorState*> all_states;
+    all_states.reserve(tensor_states_.size());
+    for (const auto& s : tensor_states_) {
+        if (s != nullptr) {
+            all_states.push_back(s.get());
+        }
+    }
+    return load_tensors_to_params_backend(all_states);
+}
+
 bool ModelManager::validate_registered_tensors() {
     bool ok = true;
     for (const auto& state : tensor_states_) {
@@ -469,7 +480,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
         return true;
     }
 
-    auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, true);
+    auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
     if (mmap_store.empty()) {
         return true;
     }

src/model_manager.h

@@ -69,6 +69,7 @@ private:
     uint64_t current_lora_epoch_ = 0;
     int n_threads_               = 0;
     bool enable_mmap_            = false;
+    bool writable_mmap_          = false;
 
     void finish_compute_backend_usage(const std::vector<TensorState*>& states);
     void release_all();
@@ -110,6 +111,7 @@ public:
         model_loader_.set_n_threads(n_threads);
     }
     void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; }
+    void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
     void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
     void set_loras(std::vector<LoraSpec> loras, SDVersion version);
 
@@ -158,6 +160,7 @@ public:
     }
 
     bool validate_registered_tensors();
+    bool load_all_params_eagerly();
 
     bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;

src/name_conversion.cpp

@@ -184,6 +184,27 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
     return name;
 }
 
+std::string convert_qwen3_vl_vision_name(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> qwen3_vl_vision_name_map{
+        {"mm.0.", "merger.linear_fc1."},
+        {"mm.2.", "merger.linear_fc2."},
+        {"v.post_ln.", "merger.norm."},
+        {"v.position_embd.weight", "pos_embed.weight"},
+        {"v.patch_embd.weight.1", "patch_embed.proj.1.weight"},
+        {"v.patch_embd.weight", "patch_embed.proj.0.weight"},
+        {"v.patch_embd.bias", "patch_embed.bias"},
+        {"v.blk.", "blocks."},
+        {"attn_qkv.", "attn.qkv."},
+        {"attn_out.", "attn.proj."},
+        {"ffn_up.", "mlp.linear_fc1."},
+        {"ffn_down.", "mlp.linear_fc2."},
+        {"ln1.", "norm1."},
+        {"ln2.", "norm2."},
+    };
+    replace_with_name_map(name, qwen3_vl_vision_name_map);
+    return name;
+}
+
 // ref: https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
 std::string convert_diffusers_unet_to_original_sd1(std::string name) {
     // (stable-diffusion, HF Diffusers)
@@ -683,6 +704,38 @@ std::string convert_other_dit_to_original_anima(std::string name) {
     return name;
 }
 
+std::string convert_diffusers_dit_to_original_krea2(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"img_in.", "first."},
+        {"time_embed.linear_1.", "tmlp.0."},
+        {"time_embed.linear_2.", "tmlp.2."},
+        {"time_mod_proj.", "tproj.1."},
+        {"txt_in.linear_1.", "txtmlp.1."},
+        {"txt_in.linear_2.", "txtmlp.3."},
+        {"text_fusion.", "txtfusion."},
+        {"transformer_blocks.", "blocks."},
+        {"final_layer.", "last."},
+    };
+    static const std::vector<std::pair<std::string, std::string>> name_map = {
+        {"attn.to_out.0.", "attn.wo."},
+        {"attn.to_out.", "attn.wo."},
+        {"attn.to_gate.", "attn.gate."},
+        {"attn.to_q.", "attn.wq."},
+        {"attn.to_k.", "attn.wk."},
+        {"attn.to_v.", "attn.wv."},
+        {"ff.gate.", "mlp.gate."},
+        {"ff.up.", "mlp.up."},
+        {"ff.down.", "mlp.down."},
+        {"txt_in.norm.", "txtmlp.0."},
+        {"last.norm.weight", "last.norm.scale"},
+        {"last.modulation.weight", "last.modulation.lin"},
+    };
+
+    replace_with_prefix_map(name, prefix_map);
+    replace_with_name_map(name, name_map);
+    return name;
+}
+
 std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
     if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
         name = convert_diffusers_unet_to_original_sd1(name);
@@ -690,12 +743,14 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_diffusers_unet_to_original_sdxl(name);
     } else if (sd_version_is_sd3(version)) {
         name = convert_diffusers_dit_to_original_sd3(name);
-    } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) {
+    } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
         name = convert_diffusers_dit_to_original_flux(name);
     } else if (sd_version_is_z_image(version)) {
         name = convert_diffusers_dit_to_original_lumina2(name);
     } else if (sd_version_is_anima(version)) {
         name = convert_other_dit_to_original_anima(name);
+    } else if (sd_version_is_krea2(version)) {
+        name = convert_diffusers_dit_to_original_krea2(name);
     }
     return name;
 }
@@ -795,7 +850,77 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) {
     return result;
 }
 
-std::string convert_first_stage_model_name(std::string name, std::string prefix) {
+std::string convert_diffusers_to_original_wan_vae(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"quant_conv.", "conv1."},
+        {"post_quant_conv.", "conv2."},
+
+        {"decoder.up_blocks.0.resnets.0.", "decoder.upsamples.0.residual."},
+        {"decoder.up_blocks.0.resnets.1.", "decoder.upsamples.1.residual."},
+        {"decoder.up_blocks.0.resnets.2.", "decoder.upsamples.2.residual."},
+        {"decoder.up_blocks.0.upsamplers.0.", "decoder.upsamples.3."},
+
+        {"decoder.up_blocks.1.resnets.0.conv_shortcut.", "decoder.upsamples.4.shortcut."},
+        {"decoder.up_blocks.1.resnets.0.", "decoder.upsamples.4.residual."},
+        {"decoder.up_blocks.1.resnets.1.", "decoder.upsamples.5.residual."},
+        {"decoder.up_blocks.1.resnets.2.", "decoder.upsamples.6.residual."},
+        {"decoder.up_blocks.1.upsamplers.0.", "decoder.upsamples.7."},
+        {"decoder.up_blocks.2.resnets.0.", "decoder.upsamples.8.residual."},
+        {"decoder.up_blocks.2.resnets.1.", "decoder.upsamples.9.residual."},
+        {"decoder.up_blocks.2.resnets.2.", "decoder.upsamples.10.residual."},
+        {"decoder.up_blocks.2.upsamplers.0.", "decoder.upsamples.11."},
+        {"decoder.up_blocks.3.resnets.0.", "decoder.upsamples.12.residual."},
+        {"decoder.up_blocks.3.resnets.1.", "decoder.upsamples.13.residual."},
+        {"decoder.up_blocks.3.resnets.2.", "decoder.upsamples.14.residual."},
+
+        {"encoder.down_blocks.0.", "encoder.downsamples.0.residual."},
+        {"encoder.down_blocks.1.", "encoder.downsamples.1.residual."},
+        {"encoder.down_blocks.2.", "encoder.downsamples.2."},
+        {"encoder.down_blocks.3.conv_shortcut.", "encoder.downsamples.3.shortcut."},
+        {"encoder.down_blocks.3.", "encoder.downsamples.3.residual."},
+        {"encoder.down_blocks.4.", "encoder.downsamples.4.residual."},
+        {"encoder.down_blocks.5.", "encoder.downsamples.5."},
+        {"encoder.down_blocks.6.conv_shortcut.", "encoder.downsamples.6.shortcut."},
+        {"encoder.down_blocks.6.", "encoder.downsamples.6.residual."},
+        {"encoder.down_blocks.7.", "encoder.downsamples.7.residual."},
+        {"encoder.down_blocks.8.", "encoder.downsamples.8."},
+        {"encoder.down_blocks.9.", "encoder.downsamples.9.residual."},
+        {"encoder.down_blocks.10.", "encoder.downsamples.10.residual."},
+    };
+
+    static const std::vector<std::pair<std::string, std::string>> shared_name_map = {
+        {".conv_in.", ".conv1."},
+        {".norm_out.", ".head.0."},
+        {".conv_out.", ".head.2."},
+
+        {".mid_block.attentions.0.", ".middle.1."},
+        {".mid_block.resnets.0.", ".middle.0.residual."},
+        {".mid_block.resnets.1.", ".middle.2.residual."},
+    };
+
+    static const std::vector<std::pair<std::string, std::string>> resnet_name_map = {
+        {".norm1.", ".0."},
+        {".conv1.", ".2."},
+        {".norm2.", ".3."},
+        {".conv2.", ".6."},
+    };
+
+    replace_with_name_map(name, shared_name_map);
+    replace_with_prefix_map(name, prefix_map);
+
+    // Only apply the ResNet-specific renaming if the tensor belongs to a ResNet block.
+    // This prevents generic ".conv1." or ".conv2." matching on top-level encoder/decoder convolutions.
+    if (name.find(".residual.") != std::string::npos) {
+        replace_with_name_map(name, resnet_name_map);
+    }
+
+    return name;
+}
+
+std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
+    if (sd_version_uses_wan_vae(version)) {
+        return convert_diffusers_to_original_wan_vae(name);
+    }
     static std::unordered_map<std::string, std::string> vae_name_map = {
         {"decoder.post_quant_conv.", "post_quant_conv."},
         {"encoder.quant_conv.", "quant_conv."},
@@ -1154,6 +1279,10 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
 
     replace_with_prefix_map(name, prefix_map);
 
+    if ((sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) && starts_with(name, "text_encoders.llm.visual.")) {
+        name = convert_qwen3_vl_vision_name(std::move(name));
+    }
+
     // diffusion model
     {
         for (const auto& prefix : diffuison_model_prefix_vec) {
@@ -1180,7 +1309,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
     {
         for (const auto& prefix : first_stage_model_prefix_vec) {
             if (starts_with(name, prefix)) {
-                name = convert_first_stage_model_name(name.substr(prefix.size()), prefix);
+                name = convert_first_stage_model_name(name.substr(prefix.size()), prefix, version);
                 if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
                     name = "tae." + name;
                 } else {

src/runtime/denoiser.hpp

@@ -302,6 +302,137 @@ struct KarrasScheduler : SigmaScheduler {
     }
 };
 
+struct BetaScheduler : SigmaScheduler {
+    static constexpr double alpha = 0.6;
+    static constexpr double beta  = 0.6;
+
+    static double log_beta(double a, double b) {
+        return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
+    }
+
+    static double incbeta(double x, double a, double b) {
+        if (x <= 0.0) {
+            return 0.0;
+        }
+        if (x >= 1.0) {
+            return 1.0;
+        }
+
+        // Continued fraction approximation using Lentz's method.
+        const int max_iter   = 200;
+        const double epsilon = 3.0e-7;
+        const double tiny    = 1e-30;
+
+        const double qab = a + b;
+        const double qap = a + 1.0;
+        const double qam = a - 1.0;
+
+        double c = 1.0;
+        double d = 1.0 - qab * x / qap;
+        if (std::abs(d) < tiny) {
+            d = tiny;
+        }
+        d        = 1.0 / d;
+        double h = d;
+
+        for (int m = 1; m <= max_iter; m++) {
+            const int m2 = 2 * m;
+
+            double aa = m * (b - m) * x / ((qam + m2) * (a + m2));
+            d         = 1.0 + aa * d;
+            if (std::abs(d) < tiny) {
+                d = tiny;
+            }
+            c = 1.0 + aa / c;
+            if (std::abs(c) < tiny) {
+                c = tiny;
+            }
+            d = 1.0 / d;
+            h *= d * c;
+
+            aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
+            d  = 1.0 + aa * d;
+            if (std::abs(d) < tiny) {
+                d = tiny;
+            }
+            c = 1.0 + aa / c;
+            if (std::abs(c) < tiny) {
+                c = tiny;
+            }
+            d                = 1.0 / d;
+            const double del = d * c;
+            h *= del;
+
+            if (std::abs(del - 1.0) < epsilon) {
+                break;
+            }
+        }
+
+        return std::exp(a * std::log(x) + b * std::log(1.0 - x) - log_beta(a, b)) / a * h;
+    }
+
+    static double beta_cdf(double x, double a, double b) {
+        if (x == 0.0) {
+            return 0.0;
+        }
+        if (x == 1.0) {
+            return 1.0;
+        }
+        if (x < (a + 1.0) / (a + b + 2.0)) {
+            return incbeta(x, a, b);
+        }
+        return 1.0 - incbeta(1.0 - x, b, a);
+    }
+
+    static double beta_ppf(double u, double a, double b, int max_iter = 30) {
+        double x = 0.5;
+        for (int i = 0; i < max_iter; i++) {
+            const double f = beta_cdf(x, a, b) - u;
+            if (std::abs(f) < 1e-10) {
+                break;
+            }
+            const double df = std::exp((a - 1.0) * std::log(x) + (b - 1.0) * std::log(1.0 - x) - log_beta(a, b));
+            x -= f / df;
+            if (x <= 0.0) {
+                x = 1e-10;
+            }
+            if (x >= 1.0) {
+                x = 1.0 - 1e-10;
+            }
+        }
+        return x;
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t t_to_sigma) override {
+        std::vector<float> result;
+        result.reserve(n + 1);
+
+        const int t_max = TIMESTEPS - 1;
+        if (n == 0) {
+            return result;
+        } else if (n == 1) {
+            result.push_back(t_to_sigma(static_cast<float>(t_max)));
+            result.push_back(0.f);
+            return result;
+        }
+
+        int last_t = -1;
+        for (uint32_t i = 0; i < n; i++) {
+            const double u      = 1.0 - static_cast<double>(i) / static_cast<double>(n);
+            const double t_cont = beta_ppf(u, alpha, beta) * t_max;
+            const int t         = static_cast<int>(std::lround(t_cont));
+
+            if (t != last_t) {
+                result.push_back(t_to_sigma(static_cast<float>(t)));
+                last_t = t;
+            }
+        }
+
+        result.push_back(0.f);
+        return result;
+    }
+};
+
 struct SimpleScheduler : SigmaScheduler {
     std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
         std::vector<float> result_sigmas;
@@ -559,6 +690,318 @@ struct LTX2Scheduler : SigmaScheduler {
     }
 };
 
+inline float flux_time_shift(float mu, float sigma, float t) {
+    return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
+}
+
+// https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
+struct FluxScheduler : SigmaScheduler {
+    int image_seq_len = 0;
+    float base_shift  = 0.5f;
+    float max_shift   = 1.15f;
+
+    explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
+        : image_seq_len(image_seq_len) {
+        parse_extra_sample_args(extra_sample_args);
+    }
+
+    void parse_extra_sample_args(const char* extra_sample_args) {
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "flux scheduler arg")) {
+            if (key == "base_shift") {
+                if (!parse_strict_float(value, base_shift)) {
+                    LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "max_shift") {
+                if (!parse_strict_float(value, max_shift)) {
+                    LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
+        }
+    }
+
+    float compute_mu() const {
+        constexpr float base_shift_anchor = 256.0f;
+        constexpr float max_shift_anchor  = 4096.0f;
+        float m                           = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
+        float b                           = base_shift - m * base_shift_anchor;
+        return static_cast<float>(image_seq_len) * m + b;
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
+        std::vector<float> sigmas;
+        sigmas.reserve(n + 1);
+
+        float mu = compute_mu();
+        LOG_DEBUG("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
+
+        if (n == 0) {
+            sigmas.push_back(1.0f);
+            return sigmas;
+        }
+
+        for (uint32_t i = 0; i <= n; ++i) {
+            float t = 1.0f - static_cast<float>(i) / static_cast<float>(n);
+            if (t <= 0.0f) {
+                sigmas.push_back(0.0f);
+            } else {
+                sigmas.push_back(flux_time_shift(mu, 1.0f, t));
+            }
+        }
+
+        sigmas[n] = 0.0f;
+        return sigmas;
+    }
+};
+
+// https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244
+struct Flux2Scheduler : SigmaScheduler {
+    int image_seq_len = 0;
+
+    explicit Flux2Scheduler(int image_seq_len)
+        : image_seq_len(image_seq_len) {}
+
+    static float compute_empirical_mu(int image_seq_len, uint32_t num_steps) {
+        const float a1 = 8.73809524e-05f;
+        const float b1 = 1.89833333f;
+        const float a2 = 0.00016927f;
+        const float b2 = 0.45666666f;
+
+        if (image_seq_len > 4300) {
+            return a2 * image_seq_len + b2;
+        }
+
+        float m_200 = a2 * image_seq_len + b2;
+        float m_10  = a1 * image_seq_len + b1;
+
+        float a = (m_200 - m_10) / 190.0f;
+        float b = m_200 - 200.0f * a;
+        return a * num_steps + b;
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
+        std::vector<float> sigmas;
+        sigmas.reserve(n + 1);
+
+        float mu = compute_empirical_mu(image_seq_len, n);
+        LOG_DEBUG("Flux2 scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
+
+        if (n == 0) {
+            sigmas.push_back(1.0f);
+            return sigmas;
+        }
+
+        for (uint32_t i = 0; i <= n; ++i) {
+            float t = 1.0f - static_cast<float>(i) / static_cast<float>(n);
+            if (t <= 0.0f) {
+                sigmas.push_back(0.0f);
+            } else if (t >= 1.0f) {
+                sigmas.push_back(1.0f);
+            } else {
+                sigmas.push_back(flux_time_shift(mu, 1.0f, t));
+            }
+        }
+
+        sigmas[n] = 0.0f;
+        return sigmas;
+    }
+};
+
+/*
+ * Logit-Normal Scheduler
+ * Based on: https://github.com/ideogram-oss/ideogram4/blob/main/src/ideogram4/scheduler.py
+ */
+struct LogitNormalScheduler : SigmaScheduler {
+    float mean       = 0.0f;
+    float std        = 1.75f;
+    float logsnr_min = -15.0f;
+    float logsnr_max = 18.0f;
+
+    bool resolution_aware = true;
+
+    float one_minus_t_min, one_minus_t_max;
+
+    void parse_extra_sample_args(int image_seq_len = 0, const char* extra_sample_args = nullptr) {
+        const int known_seq_len = (512 * 512) / (16 * 16);
+        if (extra_sample_args) {
+            for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "logit-normal scheduler arg")) {
+                if (key == "mu") {
+                    if (!parse_strict_float(value, mean)) {
+                        LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                } else if (key == "std") {
+                    if (!parse_strict_float(value, std)) {
+                        LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+                if (key == "logsnr_min") {
+                    if (!parse_strict_float(value, logsnr_min)) {
+                        LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                } else if (key == "logsnr_max") {
+                    if (!parse_strict_float(value, logsnr_max)) {
+                        LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                } else if (key == "resolution_aware") {
+                    if (!parse_strict_bool(value, resolution_aware)) {
+                        LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
+        }
+        if (image_seq_len > 0 && resolution_aware) {
+            mean += 0.5f * std::log(static_cast<float>(image_seq_len) / static_cast<float>(known_seq_len));
+        }
+    }
+
+    float sigmoid(float x) {
+        return 1.0f / (1.0f + std::exp(-x));
+    }
+
+    LogitNormalScheduler(float mean = 0.0f, float std = 1.75f, float logsnr_min = -18.0f, float logsnr_max = 15.0f)
+        : mean(mean), std(std), logsnr_min(logsnr_min), logsnr_max(logsnr_max) {
+        // t_min = 1.0f / (1.0f + std::exp(0.5f * logsnr_max));
+        one_minus_t_min = sigmoid(0.5f * logsnr_max);
+        // t_max = 1.0f / (1.0f + std::exp(0.5f * logsnr_min));
+        one_minus_t_max = sigmoid(0.5f * logsnr_min);
+    }
+
+    LogitNormalScheduler(int image_seq_len = 0, const char* extra_sample_args = nullptr) {
+        mean       = 0.0f;
+        std        = 1.75f;
+        logsnr_min = -15.0f;
+        logsnr_max = 18.0f;
+
+        parse_extra_sample_args(image_seq_len, extra_sample_args);
+        // t_min = 1.0f / (1.0f + std::exp(0.5f * logsnr_max));
+        one_minus_t_min = sigmoid(0.5f * logsnr_max);
+        // t_max = 1.0f / (1.0f + std::exp(0.5f * logsnr_min));
+        one_minus_t_max = sigmoid(0.5f * logsnr_min);
+    }
+
+    // https://stackedboxes.org/2017/05/01/acklams-normal-quantile-function/
+    double ndtri(double p) {
+        if (p <= 0.0) {
+            return -std::numeric_limits<double>::infinity();
+        } else if (p >= 1.0) {
+            return std::numeric_limits<double>::infinity();
+        }
+
+        static const double p_low  = 0.02425;
+        static const double p_high = 1.0 - p_low;
+
+        static const double c[6] = {-7.784894002430293e-03,
+                                    -3.223964580411365e-01,
+                                    -2.400758277161838e+00,
+                                    -2.549732539343734e+00,
+                                    4.374664141464968e+00,
+                                    2.938163982698783e+00};
+
+        static const double d[5] = {7.784695709041462e-03,
+                                    3.224671290700398e-01,
+                                    2.445134137142996e+00,
+                                    3.754408661907416e+00,
+                                    1.0};
+
+        // Coefficients for the central region
+        static const double a[6] = {-3.969683028665376e+01,
+                                    2.209460984245205e+02,
+                                    -2.759285104469687e+02,
+                                    1.383577518672690e+02,
+                                    -3.066479806614716e+01,
+                                    2.506628277459239e+00};
+
+        static const double b[6] = {-5.447609879822406e+01,
+                                    1.615858368580409e+02,
+                                    -1.556989798598866e+02,
+                                    6.680131188771972e+01,
+                                    -1.328068155288572e+01,
+                                    1.0};
+
+        double x = 0.0;
+
+        if (p < p_low) {
+            // Lower region
+            double q = std::sqrt(-2.0 * std::log(p));
+
+            // Numerator: c[0]*q^5 + c[1]*q^4 + ... + c[5]
+            double numerator = c[0];
+            for (int i = 1; i < 6; ++i) {
+                numerator = numerator * q + c[i];
+            }
+
+            // Denominator: d[0]*q^4 + d[1]*q^3 + ... + d[3]*q + 1
+            double denominator = d[0];
+            for (int i = 1; i < 5; ++i) {
+                denominator = denominator * q + d[i];
+            }
+
+            x = numerator / denominator;
+        } else if (p > p_high) {
+            // Upper region
+            double q = std::sqrt(-2.0 * std::log(1.0 - p));
+
+            double numerator = c[0];
+            for (int i = 1; i < 6; ++i) {
+                numerator = numerator * q + c[i];
+            }
+
+            double denominator = d[0];
+            for (int i = 1; i < 5; ++i) {
+                denominator = denominator * q + d[i];
+            }
+
+            x = -(numerator / denominator);
+        } else {
+            // Central region
+            double q = p - 0.5;
+            double r = q * q;
+
+            // Numerator: (a[0]*r^5 + a[1]*r^4 + ... + a[5])*q
+            double numerator = a[0];
+            for (int i = 1; i < 6; ++i) {
+                numerator = numerator * r + a[i];
+            }
+            numerator *= q;
+
+            // Denominator: b[0]*r^4 + b[1]*r^3 + ... + b[4]*r + 1
+            double denominator = b[0];
+            for (int i = 1; i < 6; ++i) {
+                denominator = denominator * r + b[i];
+            }
+
+            x = numerator / denominator;
+        }
+        return x;
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
+        std::vector<float> sigmas;
+        LOG_INFO("LOGIT_NORMAL_SCHEDULER using mean=%.4f, std=%.4f, logsnr_min=%.4f, logsnr_max=%.4f", mean, std, logsnr_min, logsnr_max);
+        sigmas.reserve(n + 1);
+        for (uint32_t i = 0; i <= n; ++i) {
+            float t = static_cast<float>(i) / static_cast<float>(n);
+
+            // ndtri(1-t) == -ndtri(t)
+            float z = static_cast<float>(-ndtri(t));
+
+            float y = mean + std * z;
+
+            float timestep = sigmoid(y);
+
+            if (timestep > one_minus_t_min)
+                timestep = one_minus_t_min;
+            if (timestep < one_minus_t_max)
+                timestep = one_minus_t_max;
+
+            float sigma = timestep;
+
+            sigmas.push_back(sigma);
+        }
+        sigmas[n] = 0.0f;
+        return sigmas;
+    }
+};
+
 struct Denoiser {
     virtual float sigma_min()                                                        = 0;
     virtual float sigma_max()                                                        = 0;
@@ -583,6 +1026,10 @@ struct Denoiser {
                 LOG_INFO("get_sigmas with Karras scheduler");
                 scheduler = std::make_shared<KarrasScheduler>();
                 break;
+            case BETA_SCHEDULER:
+                LOG_INFO("get_sigmas with Beta scheduler");
+                scheduler = std::make_shared<BetaScheduler>();
+                break;
             case EXPONENTIAL_SCHEDULER:
                 LOG_INFO("get_sigmas exponential scheduler");
                 scheduler = std::make_shared<ExponentialScheduler>();
@@ -623,6 +1070,21 @@ struct Denoiser {
                 LOG_INFO("get_sigmas with LTX2 scheduler");
                 scheduler = std::make_shared<LTX2Scheduler>(image_seq_len, extra_sample_args);
                 break;
+            case LOGIT_NORMAL_SCHEDULER: {
+                LOG_INFO("get_sigmas with Logit-Normal scheduler");
+                scheduler = std::make_shared<LogitNormalScheduler>(image_seq_len, extra_sample_args);
+                break;
+            }
+            case FLUX2_SCHEDULER: {
+                LOG_INFO("get_sigmas with Flux2 scheduler");
+                scheduler = std::make_shared<Flux2Scheduler>(image_seq_len);
+                break;
+            }
+            case FLUX_SCHEDULER: {
+                LOG_INFO("get_sigmas with Flux scheduler");
+                scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
+                break;
+            }
             default:
                 LOG_INFO("get_sigmas with discrete scheduler (default)");
                 scheduler = std::make_shared<DiscreteScheduler>();
@@ -787,10 +1249,6 @@ struct DiscreteFlowDenoiser : public Denoiser {
     }
 };
 
-inline float flux_time_shift(float mu, float sigma, float t) {
-    return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
-}
-
 struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
     FluxFlowDenoiser() = default;
 
@@ -804,35 +1262,141 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
     }
 };
 
-struct Flux2FlowDenoiser : public FluxFlowDenoiser {
-    Flux2FlowDenoiser() = default;
+struct SefiFlowDenoiser;
 
-    float compute_empirical_mu(uint32_t n, int image_seq_len) {
-        const float a1 = 8.73809524e-05f;
-        const float b1 = 1.89833333f;
-        const float a2 = 0.00016927f;
-        const float b2 = 0.45666666f;
+struct SefiFlowDenoiser : public FluxFlowDenoiser {
+    static constexpr int kNumTrainTimesteps = 1000;
+    static constexpr int kSemChannels       = 16;
+    static constexpr int kTotalChannels     = 144;
 
-        if (image_seq_len > 4300) {
-            float mu = a2 * image_seq_len + b2;
-            return mu;
+    float delta_t              = 0.1f;
+    float timestep_shift_alpha = 1.0f;
+
+    std::vector<float> sem_sigmas;
+    std::vector<float> tex_sigmas;
+    std::vector<float> sem_timesteps;
+    std::vector<float> tex_timesteps;
+
+    SefiFlowDenoiser() = default;
+
+    static float apply_alpha_shift(float u_unit, float alpha) {
+        if (alpha == 1.0f) {
+            return u_unit;
+        }
+        float denom = 1.0f + (alpha - 1.0f) * u_unit;
+        return (alpha * u_unit) / denom;
+    }
+
+    std::vector<float> get_sigmas(uint32_t n,
+                                  int image_seq_len,
+                                  scheduler_t scheduler_type,
+                                  SDVersion version,
+                                  const char* extra_sample_args = nullptr) override {
+        sem_sigmas.clear();
+        tex_sigmas.clear();
+        sem_timesteps.clear();
+        tex_timesteps.clear();
+
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "sefi scheduler arg")) {
+            if (key == "sefi_alpha") {
+                if (!parse_strict_float(value, timestep_shift_alpha)) {
+                    LOG_WARN("ignoring invalid sefi scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "sefi_delta_t") {
+                if (!parse_strict_float(value, delta_t)) {
+                    LOG_WARN("ignoring invalid sefi scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
         }
 
-        float m_200 = a2 * image_seq_len + b2;
-        float m_10  = a1 * image_seq_len + b1;
+        for (uint32_t i = 0; i <= n; ++i) {
+            float u_base    = static_cast<float>(i) / static_cast<float>(n);
+            float u_shifted = apply_alpha_shift(u_base, timestep_shift_alpha);
+            float u_sem_raw = u_shifted * (1.0f + delta_t);
 
-        float a  = (m_200 - m_10) / 190.0f;
-        float b  = m_200 - 200.0f * a;
-        float mu = a * n + b;
+            float u_sem = std::min(u_sem_raw, 1.0f);
+            float u_tex = std::max(0.0f, std::min(u_sem_raw - delta_t, 1.0f));
 
-        return mu;
+            int idx_sem = std::min(kNumTrainTimesteps - 1,
+                                   std::max(0, static_cast<int>(u_sem * (kNumTrainTimesteps - 1))));
+            int idx_tex = std::min(kNumTrainTimesteps - 1,
+                                   std::max(0, static_cast<int>(u_tex * (kNumTrainTimesteps - 1))));
+
+            float t_sem     = static_cast<float>(kNumTrainTimesteps - idx_sem);
+            float t_tex     = static_cast<float>(kNumTrainTimesteps - idx_tex);
+            float sigma_sem = t_sem / static_cast<float>(kNumTrainTimesteps);
+            float sigma_tex = t_tex / static_cast<float>(kNumTrainTimesteps);
+
+            sem_timesteps.push_back(t_sem);
+            tex_timesteps.push_back(t_tex);
+            sem_sigmas.push_back(sigma_sem);
+            tex_sigmas.push_back(sigma_tex);
+        }
+        LOG_DEBUG("SefiFlowDenoiser: built %u-step dual schedule (alpha=%.2f delta_t=%.2f)",
+                  n, timestep_shift_alpha, delta_t);
+        return tex_sigmas;
+    }
+};
+
+// MiniT2I predicts x0 directly and integrates a linear flow ODE:
+//   x_{t+dt} = x_t + (x0 - x_t)/(1 - t) * dt,  t in [0, 1), x0 = start = noise * 2.
+// Mapping sigma = 1 - t makes the generic Euler update
+//   x += (x - denoised)/sigma * (sigma_next - sigma)
+// exactly reproduce that step when denoised == x0. To make the generic
+// `denoised = pred * c_out + x * c_skip` yield x0 from the model's raw x0
+// prediction we use c_skip = 0, c_out = 1, c_in = 1. Sigmas run linearly 1 -> 0.
+struct MiniT2IFlowDenoiser : public Denoiser {
+    float sigma_min() override {
+        return 0.0f;
+    }
+
+    float sigma_max() override {
+        return 1.0f;
+    }
+
+    float sigma_to_t(float sigma) override {
+        return 1.0f - sigma;
+    }
+
+    float t_to_sigma(float t) override {
+        return 1.0f - t;
+    }
+
+    std::vector<float> get_scalings(float sigma) override {
+        SD_UNUSED(sigma);
+        float c_skip = 0.0f;
+        float c_out  = 1.0f;
+        float c_in   = 1.0f;
+        return {c_skip, c_out, c_in};
+    }
+
+    sd::Tensor<float> noise_scaling(float sigma,
+                                    const sd::Tensor<float>& noise,
+                                    const sd::Tensor<float>& latent) override {
+        SD_UNUSED(sigma);
+        SD_UNUSED(latent);
+        // Sampling starts from x0_init = noise * 2 (see MiniT2I reference).
+        return noise * 2.0f;
+    }
+
+    sd::Tensor<float> inverse_noise_scaling(float sigma, const sd::Tensor<float>& latent) override {
+        SD_UNUSED(sigma);
+        return latent;
     }
 
     std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) override {
-        float mu = compute_empirical_mu(n, image_seq_len);
-        LOG_DEBUG("Flux2FlowDenoiser: set shift to %.3f", mu);
-        set_shift(mu);
-        return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version, extra_sample_args);
+        SD_UNUSED(image_seq_len);
+        SD_UNUSED(scheduler_type);
+        SD_UNUSED(version);
+        SD_UNUSED(extra_sample_args);
+        // Uniform t schedule 0 -> 1 => sigma 1 -> 0, matching the reference loop.
+        std::vector<float> sigmas;
+        sigmas.reserve(n + 1);
+        for (uint32_t i = 0; i < n; ++i) {
+            sigmas.push_back(1.0f - static_cast<float>(i) / static_cast<float>(n));
+        }
+        sigmas.push_back(0.0f);
+        return sigmas;
     }
 };
 
@@ -939,6 +1503,40 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_sefi_euler(SefiFlowDenoiser* sefi,
+                                           denoise_cb_t model,
+                                           sd::Tensor<float> x) {
+    const std::vector<float>& sigma_tex_vec = sefi->tex_sigmas;
+    const std::vector<float>& sigma_sem_vec = sefi->sem_sigmas;
+    int steps                               = static_cast<int>(sigma_tex_vec.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma_tex_cur  = sigma_tex_vec[i];
+        float sigma_tex_next = sigma_tex_vec[i + 1];
+        float sigma_sem_cur  = sigma_sem_vec[i];
+        float sigma_sem_next = sigma_sem_vec[i + 1];
+        if (sigma_tex_cur <= 1e-9f) {
+            continue;
+        }
+        auto denoised_opt = model(x, sigma_tex_cur, i + 1);
+        if (denoised_opt.pred.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised = std::move(denoised_opt.pred);
+        sd::Tensor<float> velocity = (x - denoised) / sigma_tex_cur;
+
+        auto x_sem      = sd::ops::slice(x, 2, 0, SefiFlowDenoiser::kSemChannels);
+        auto x_tex      = sd::ops::slice(x, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels);
+        auto vel_sem    = sd::ops::slice(velocity, 2, 0, SefiFlowDenoiser::kSemChannels);
+        auto vel_tex    = sd::ops::slice(velocity, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels);
+        auto x_sem_next = x_sem + vel_sem * (sigma_sem_next - sigma_sem_cur);
+        auto x_tex_next = x_tex + vel_tex * (sigma_tex_next - sigma_tex_cur);
+
+        sd::ops::slice_assign(&x, 2, 0, SefiFlowDenoiser::kSemChannels, x_sem_next);
+        sd::ops::slice_assign(&x, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels, x_tex_next);
+    }
+    return x;
+}
+
 static sd::Tensor<float> sample_euler(denoise_cb_t model,
                                       sd::Tensor<float> x,
                                       const std::vector<float>& sigmas) {
@@ -1854,7 +2452,13 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             std::shared_ptr<RNG> rng,
                                             float eta,
                                             bool is_flow_denoiser,
-                                            const char* extra_sample_args) {
+                                            const char* extra_sample_args,
+                                            std::shared_ptr<Denoiser> denoiser_for_dispatch = nullptr) {
+    if (denoiser_for_dispatch) {
+        if (auto sefi = std::dynamic_pointer_cast<SefiFlowDenoiser>(denoiser_for_dispatch)) {
+            return sample_sefi_euler(sefi.get(), model, std::move(x));
+        }
+    }
     SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
     switch (method) {
         case EULER_A_SAMPLE_METHOD:

src/runtime/guidance.cpp

@@ -3,6 +3,7 @@
 #include <algorithm>
 #include <cmath>
 #include <cstdlib>
+#include <optional>
 #include <string>
 #include <utility>
 
@@ -63,6 +64,82 @@ namespace sd::guidance {
         return uncond;
     }
 
+    std::vector<float> parse_guidance_schedule_from_spec(std::string spec) {
+        std::vector<float> schedule;
+
+        while (!spec.empty()) {
+            auto sep     = spec.find('+');
+            auto segment = spec.substr(0, sep);
+
+            auto x = segment.find('x');
+            if (x == std::string::npos) {
+                LOG_ERROR("Invalid guidance schedule segment: '%s' (expected <guidance>x<count>)", segment.c_str());
+                return {};
+            }
+
+            float guidance;
+            int count;
+
+            auto guidance_str = segment.substr(0, x);
+            auto count_str    = segment.substr(x + 1);
+
+            try {
+                size_t idx = 0;
+                guidance   = std::stof(guidance_str, &idx);
+                if (idx != guidance_str.size()) {
+                    LOG_ERROR("Invalid guidance value in guidance schedule: '%s'", guidance_str.c_str());
+                    return {};
+                }
+            } catch (const std::exception&) {
+                LOG_ERROR("Invalid guidance value in guidance schedule: '%s'", guidance_str.c_str());
+                return {};
+            }
+
+            try {
+                size_t idx = 0;
+                count      = std::stoi(count_str, &idx);
+                if (idx != count_str.size()) {
+                    LOG_ERROR("Invalid count in guidance schedule: '%s'", count_str.c_str());
+                    return {};
+                }
+            } catch (const std::exception&) {
+                LOG_ERROR("Invalid count in guidance schedule: '%s'", count_str.c_str());
+                return {};
+            }
+
+            if (count <= 0) {
+                LOG_ERROR("Guidance schedule count must be positive");
+                return {};
+            }
+
+            schedule.insert(schedule.end(), count, guidance);
+
+            if (sep == std::string::npos) {
+                break;
+            }
+
+            spec = spec.substr(sep + 1);
+        }
+
+        return schedule;
+    }
+
+    std::vector<float> parse_guidance_schedule(const char* extra_sample_args) {
+        std::vector<float> guidance_schedule;
+        std::string guidance_schedule_str = "";
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) {
+            float parsed = 0.0f;
+            if (key == "guidance_schedule") {
+                guidance_schedule_str = value;
+            }
+        }
+
+        if (!guidance_schedule_str.empty()) {
+            guidance_schedule = parse_guidance_schedule_from_spec(guidance_schedule_str);
+        }
+        return guidance_schedule;
+    }
+
     ClassifierFreeGuidance::ClassifierFreeGuidance(float guidance_scale,
                                                    float image_guidance_scale)
         : guidance_scale_(guidance_scale),
@@ -70,8 +147,10 @@ namespace sd::guidance {
     }
 
     GuiderOutput ClassifierFreeGuidance::forward(const GuidanceInput& input,
-                                                 GuiderOutput previous) const {
+                                                 GuiderOutput previous,
+                                                 std::optional<float> scale_override) const {
         (void)previous;
+        float guidance_scale = scale_override.value_or(guidance_scale_);
 
         GuiderOutput output;
         if (!has_tensor(input.pred_cond)) {
@@ -86,14 +165,14 @@ namespace sd::guidance {
                 const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
                 output.pred                              = pred_img_uncond +
                               image_guidance_scale_ * (pred_uncond - pred_img_uncond) +
-                              guidance_scale_ * (pred_cond - pred_uncond);
+                              guidance_scale * (pred_cond - pred_uncond);
 
             } else {
-                output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond);
+                output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond);
             }
         } else if (has_tensor(input.pred_img_uncond)) {
             const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
-            output.pred                              = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond);
+            output.pred                              = pred_img_uncond + guidance_scale * (pred_cond - pred_img_uncond);
         }
 
         return output;
@@ -128,8 +207,10 @@ namespace sd::guidance {
     }
 
     GuiderOutput AdaptiveProjectedGuidance::forward(const GuidanceInput& input,
-                                                    GuiderOutput previous) const {
+                                                    GuiderOutput previous,
+                                                    std::optional<float> scale_override) const {
         (void)previous;
+        float guidance_scale = scale_override.value_or(guidance_scale_);
 
         GuiderOutput output;
         if (!has_tensor(input.pred_cond)) {
@@ -144,13 +225,13 @@ namespace sd::guidance {
                 const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
                 output.pred                              = pred_img_uncond +
                               image_guidance_scale_ * (pred_uncond - pred_img_uncond) +
-                              guidance_scale_ * (pred_cond - pred_uncond);
+                              guidance_scale * (pred_cond - pred_uncond);
             } else {
-                output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond);
+                output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond);
             }
         } else if (has_tensor(input.pred_img_uncond)) {
             const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
-            output.pred                              = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond);
+            output.pred                              = pred_img_uncond + guidance_scale * (pred_cond - pred_img_uncond);
         }
         if (!has_tensor(input.pred_uncond) && !has_tensor(input.pred_img_uncond)) {
             return output;
@@ -162,7 +243,7 @@ namespace sd::guidance {
         sd::Tensor<float> deltas = calculate_guidance_delta(pred_cond,
                                                             pred_uncond,
                                                             pred_img_uncond,
-                                                            guidance_scale_,
+                                                            guidance_scale,
                                                             image_guidance_scale_);
         if (params_.momentum != 0.0f) {
             if (momentum_buffer_.shape() != deltas.shape()) {
@@ -239,7 +320,8 @@ namespace sd::guidance {
     }
 
     GuiderOutput SkipLayerGuidance::forward(const GuidanceInput& input,
-                                            GuiderOutput output) const {
+                                            GuiderOutput output,
+                                            std::optional<float> /*scale_override*/) const {
         if (scale_ == 0.0f || !is_enabled_for_step(input) || !input.predict_skip_layer) {
             return output;
         }

src/runtime/guidance.h

@@ -3,6 +3,7 @@
 
 #include <cstddef>
 #include <functional>
+#include <optional>
 #include <vector>
 
 #include "core/tensor.hpp"
@@ -27,6 +28,7 @@ namespace sd::guidance {
     AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args);
     bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params);
     bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args);
+    std::vector<float> parse_guidance_schedule(const char* extra_sample_args);
 
     struct GuidanceInput {
         int step                                 = 0;
@@ -40,9 +42,10 @@ namespace sd::guidance {
 
     class BaseGuidance {
     public:
-        virtual ~BaseGuidance()                                   = default;
+        virtual ~BaseGuidance()                                                                = default;
         virtual GuiderOutput forward(const GuidanceInput& input,
-                                     GuiderOutput previous) const = 0;
+                                     GuiderOutput previous,
+                                     std::optional<float> scale_override = std::nullopt) const = 0;
     };
 
     class ClassifierFreeGuidance : public BaseGuidance {
@@ -54,7 +57,8 @@ namespace sd::guidance {
                                float image_guidance_scale);
 
         GuiderOutput forward(const GuidanceInput& input,
-                             GuiderOutput previous) const override;
+                             GuiderOutput previous,
+                             std::optional<float> scale_override = std::nullopt) const override;
     };
 
     class AdaptiveProjectedGuidance : public BaseGuidance {
@@ -69,7 +73,8 @@ namespace sd::guidance {
                                   AdaptiveProjectedGuidanceParams params);
 
         GuiderOutput forward(const GuidanceInput& input,
-                             GuiderOutput previous) const override;
+                             GuiderOutput previous,
+                             std::optional<float> scale_override = std::nullopt) const override;
     };
 
     class SkipLayerGuidance : public BaseGuidance {
@@ -88,7 +93,8 @@ namespace sd::guidance {
         const std::vector<int>& layers() const;
 
         GuiderOutput forward(const GuidanceInput& input,
-                             GuiderOutput previous) const override;
+                             GuiderOutput previous,
+                             std::optional<float> scale_override = std::nullopt) const override;
     };
 
 }  // namespace sd::guidance

src/runtime/imatrix.cpp

@@ -0,0 +1,308 @@
+#include "runtime/imatrix.h"
+
+/* Adapted from llama.cpp (credits: Kawrakow). */
+
+#include "core/util.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+#include "stable-diffusion.h"
+
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+
+static IMatrixCollector imatrix_collector;
+
+IMatrixCollector& get_imatrix_collector() {
+    return imatrix_collector;
+}
+
+// remove any prefix and suffixes from the name
+// CUDA0#blk.0.attn_k.weight#0 => blk.0.attn_k.weight
+static std::string filter_tensor_name(const char* name) {
+    std::string wname;
+    const char* p = strchr(name, '#');
+    if (p != NULL) {
+        p             = p + 1;
+        const char* q = strchr(p, '#');
+        if (q != NULL) {
+            wname = std::string(p, q - p);
+        } else {
+            wname = p;
+        }
+    } else {
+        wname = name;
+    }
+    return wname;
+}
+
+bool IMatrixCollector::collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
+    GGML_UNUSED(user_data);
+    if (t == nullptr) {
+        return false;
+    }
+    if (t->op != GGML_OP_MUL_MAT && t->op != GGML_OP_MUL_MAT_ID) {
+        return false;
+    }
+
+    const struct ggml_tensor* src0 = t->src[0];
+    const struct ggml_tensor* src1 = t->src[1];
+    if (src0 == nullptr || src1 == nullptr) {
+        return false;
+    }
+    std::string wname = filter_tensor_name(src0->name);
+
+    // when ask is true, the scheduler wants to know if we are interested in data from this tensor
+    // if we return true, a follow-up call will be made with ask=false in which we can do the actual collection
+    if (ask) {
+        if (t->op == GGML_OP_MUL_MAT_ID) {
+            return true;  // collect all indirect matrix multiplications
+        }
+        // why are small batches ignored (<16 tokens)?
+        // if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false;
+        if (!(wname.substr(0, 6) == "model." || wname.substr(0, 17) == "cond_stage_model." || wname.substr(0, 14) == "text_encoders.")) {
+            return false;
+        }
+        return true;
+    }
+    std::lock_guard<std::mutex> lock(mutex_);
+
+    // copy the data from the GPU memory if needed
+    const bool is_host = src1->buffer == NULL || ggml_backend_buffer_is_host(src1->buffer);
+
+    if (!is_host) {
+        src1_data_.resize(ggml_nelements(src1));
+        ggml_backend_tensor_get(src1, src1_data_.data(), 0, ggml_nbytes(src1));
+    }
+
+    const float* data = is_host ? (const float*)src1->data : src1_data_.data();
+
+    // this has been adapted to the new format of storing merged experts in a single 3d tensor
+    // ref: https://github.com/ggml-org/llama.cpp/pull/6387
+    if (t->op == GGML_OP_MUL_MAT_ID) {
+        //   ids  -> [n_experts_used, n_tokens]
+        //   src1 -> [cols, n_expert_used, n_tokens]
+        const ggml_tensor* ids = t->src[2];
+        const int n_as         = static_cast<int>(src0->ne[2]);
+        const int n_ids        = static_cast<int>(ids->ne[0]);
+
+        // the top-k selected expert ids are stored in the ids tensor
+        // for simplicity, always copy ids to host, because it is small
+        // take into account that ids is not contiguous!
+
+        GGML_ASSERT(ids->ne[1] == src1->ne[2]);
+
+        ids_.resize(ggml_nbytes(ids));
+        ggml_backend_tensor_get(ids, ids_.data(), 0, ggml_nbytes(ids));
+
+        auto& e = stats_[wname];
+
+        ++e.ncall;
+
+        if (e.values.empty()) {
+            e.values.resize(src1->ne[0] * n_as, 0);
+            e.counts.resize(src1->ne[0] * n_as, 0);
+        } else if (e.values.size() != (size_t)src1->ne[0] * n_as) {
+            LOG_ERROR("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0] * n_as);
+            exit(1);  // GGML_ABORT("fatal error");
+        }
+        // loop over all possible experts, regardless if they are used or not in the batch
+        for (int ex = 0; ex < n_as; ++ex) {
+            size_t e_start = ex * src1->ne[0];
+
+            for (int idx = 0; idx < n_ids; ++idx) {
+                for (int row = 0; row < (int)src1->ne[2]; ++row) {
+                    const int excur = *(const int32_t*)(ids_.data() + row * ids->nb[1] + idx * ids->nb[0]);
+
+                    GGML_ASSERT(excur >= 0 && excur < n_as);  // sanity check
+
+                    if (excur != ex)
+                        continue;
+
+                    const int64_t i11 = idx % src1->ne[1];
+                    const int64_t i12 = row;
+                    const float* x    = (const float*)((const char*)data + i11 * src1->nb[1] + i12 * src1->nb[2]);
+
+                    for (int j = 0; j < (int)src1->ne[0]; ++j) {
+                        e.values[e_start + j] += x[j] * x[j];
+                        e.counts[e_start + j]++;
+                        if (!std::isfinite(e.values[e_start + j])) {
+                            LOG_ERROR("%f detected in %s\n", e.values[e_start + j], wname.c_str());
+                            exit(1);
+                        }
+                    }
+                }
+            }
+        }
+    } else {
+        auto& e = stats_[wname];
+        if (e.values.empty()) {
+            e.values.resize(src1->ne[0], 0);
+            e.counts.resize(src1->ne[0], 0);
+        } else if (e.values.size() != (size_t)src1->ne[0]) {
+            LOG_WARN("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]);
+            exit(1);  // GGML_ABORT("fatal error");
+        }
+
+        ++e.ncall;
+        for (int row = 0; row < (int)src1->ne[1]; ++row) {
+            const float* x = data + row * src1->ne[0];
+            for (int j = 0; j < (int)src1->ne[0]; ++j) {
+                if (std::isfinite(x[j])) {
+                    e.values[j] += x[j] * x[j];
+                    e.counts[j]++;
+                    if (!std::isfinite(e.values[j])) {
+                        LOG_WARN("%f detected in %s\n", e.values[j], wname.c_str());
+                        exit(1);
+                    }
+                } else {
+                    // Likely something from an attention mask?
+                }
+            }
+        }
+    }
+    return true;
+}
+
+bool load_imatrix(const char* imatrix_path) {
+    return imatrix_collector.load_imatrix(imatrix_path);
+}
+
+void save_imatrix(const char* imatrix_path) {
+    imatrix_collector.save_imatrix(imatrix_path);
+}
+
+static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
+    return imatrix_collector.collect_imatrix(t, ask, user_data);
+}
+
+void enable_imatrix_collection() {
+    sd_set_backend_eval_callback(collect_imatrix, nullptr);
+}
+
+void disable_imatrix_collection() {
+    sd_set_backend_eval_callback(nullptr, nullptr);
+}
+
+void IMatrixCollector::save_imatrix(std::string fname, int ncall) const {
+    if (ncall > 0) {
+        fname += ".at_";
+        fname += std::to_string(ncall);
+    }
+    // avoid writing imatrix entries that do not have full data
+    // this can happen with MoE models where some of the experts end up not being exercised by the provided training data
+
+    int n_entries = 0;
+    std::vector<std::string> to_store;
+
+    for (const auto& kv : stats_) {
+        const int n_all = static_cast<int>(kv.second.counts.size());
+
+        if (n_all == 0) {
+            continue;
+        }
+
+        int n_zeros = 0;
+        for (const int c : kv.second.counts) {
+            if (c == 0) {
+                n_zeros++;
+            }
+        }
+
+        if (n_zeros == n_all) {
+            LOG_WARN("entry '%40s' has no data - skipping\n", kv.first.c_str());
+            continue;
+        }
+
+        if (n_zeros > 0) {
+            LOG_WARN("entry '%40s' has partial data (%.2f%%) - skipping\n", kv.first.c_str(), 100.0f * (n_all - n_zeros) / n_all);
+            continue;
+        }
+
+        n_entries++;
+        to_store.push_back(kv.first);
+    }
+
+    if (to_store.size() < stats_.size()) {
+        LOG_WARN("storing only %zu out of %zu entries\n", to_store.size(), stats_.size());
+    }
+
+    std::ofstream out(fname, std::ios::binary);
+    out.write((const char*)&n_entries, sizeof(n_entries));
+    for (const auto& name : to_store) {
+        const auto& stat = stats_.at(name);
+        int len          = static_cast<int>(name.size());
+        out.write((const char*)&len, sizeof(len));
+        out.write(name.c_str(), len);
+        out.write((const char*)&stat.ncall, sizeof(stat.ncall));
+        int nval = static_cast<int>(stat.values.size());
+        out.write((const char*)&nval, sizeof(nval));
+        if (nval > 0) {
+            std::vector<float> tmp(nval);
+            for (int i = 0; i < nval; i++) {
+                tmp[i] = (stat.values[i] / static_cast<float>(stat.counts[i])) * static_cast<float>(stat.ncall);
+            }
+            out.write((const char*)tmp.data(), nval * sizeof(float));
+        }
+    }
+
+    // Write the number of call the matrix was computed with
+    out.write((const char*)&last_call_, sizeof(last_call_));
+}
+
+bool IMatrixCollector::load_imatrix(const char* fname) {
+    std::ifstream in(fname, std::ios::binary);
+    if (!in) {
+        LOG_ERROR("failed to open %s\n", fname);
+        return false;
+    }
+    int n_entries;
+    in.read((char*)&n_entries, sizeof(n_entries));
+    if (in.fail() || n_entries < 1) {
+        LOG_ERROR("no data in file %s\n", fname);
+        return false;
+    }
+    for (int i = 0; i < n_entries; ++i) {
+        int len;
+        in.read((char*)&len, sizeof(len));
+        std::vector<char> name_as_vec(len + 1);
+        in.read((char*)name_as_vec.data(), len);
+        if (in.fail()) {
+            LOG_ERROR("failed reading name for entry %d from %s\n", i + 1, fname);
+            return false;
+        }
+        name_as_vec[len] = 0;
+        std::string name{name_as_vec.data()};
+        auto& e = stats_[std::move(name)];
+        int ncall;
+        in.read((char*)&ncall, sizeof(ncall));
+        int nval;
+        in.read((char*)&nval, sizeof(nval));
+        if (in.fail() || nval < 1) {
+            LOG_ERROR("failed reading number of values for entry %d\n", i);
+            stats_ = {};
+            return false;
+        }
+
+        if (e.values.empty()) {
+            e.values.resize(nval, 0);
+            e.counts.resize(nval, 0);
+        }
+
+        std::vector<float> tmp(nval);
+        in.read((char*)tmp.data(), nval * sizeof(float));
+        if (in.fail()) {
+            LOG_ERROR("failed reading data for entry %d\n", i);
+            stats_ = {};
+            return false;
+        }
+
+        // Recreate the state as expected by save_imatrix(), and correct for weighted sum.
+        for (int i = 0; i < nval; i++) {
+            e.values[i] += tmp[i];
+            e.counts[i] += ncall;
+        }
+        e.ncall += ncall;
+    }
+    return true;
+}

src/runtime/imatrix.h

@@ -0,0 +1,45 @@
+#ifndef __SD_RUNTIME_IMATRIX_H__
+#define __SD_RUNTIME_IMATRIX_H__
+
+#include <fstream>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+/* Adapted from llama.cpp (credits: Kawrakow). */
+
+struct ggml_tensor;
+
+struct IMatrixStats {
+    std::vector<float> values{};
+    std::vector<int> counts{};
+    int ncall = 0;
+};
+
+class IMatrixCollector {
+private:
+    std::unordered_map<std::string, IMatrixStats> stats_ = {};
+    std::mutex mutex_;
+    int last_call_ = 0;
+    std::vector<float> src1_data_;
+    std::vector<char> ids_;  // the expert ids from ggml_mul_mat_id
+
+public:
+    IMatrixCollector() = default;
+    bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data);
+    void save_imatrix(std::string fname, int ncall = -1) const;
+    bool load_imatrix(const char* fname);
+    std::vector<float> get_values(const std::string& key) const {
+        auto it = stats_.find(key);
+        if (it != stats_.end()) {
+            return it->second.values;
+        } else {
+            return {};
+        }
+    }
+};
+
+IMatrixCollector& get_imatrix_collector();
+
+#endif  // __SD_RUNTIME_IMATRIX_H__

src/stable-diffusion.cpp

@@ -20,13 +20,16 @@
 #include "extensions/generation_extension.h"
 #include "model/adapter/lora.hpp"
 #include "model/diffusion/anima.hpp"
+#include "model/diffusion/boogu.hpp"
 #include "model/diffusion/control.hpp"
 #include "model/diffusion/ernie_image.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/hidream_o1.hpp"
 #include "model/diffusion/ideogram4.hpp"
+#include "model/diffusion/krea2.hpp"
 #include "model/diffusion/lens.hpp"
 #include "model/diffusion/ltxv.hpp"
+#include "model/diffusion/minit2i.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
 #include "model/diffusion/pid.hpp"
@@ -81,18 +84,23 @@ const char* model_version_to_str[] = {
     "Wan 2.2 I2V",
     "Wan 2.2 TI2V",
     "Qwen Image",
+    "Qwen Image Layered",
     "Anima",
     "Flux.2",
     "Flux.2 klein",
     "LTXAV",
     "HiDream O1",
     "Z-Image",
+    "Boogu Image",
     "Ovis Image",
     "Ernie Image",
     "Lens",
+    "MiniT2I",
     "Longcat-Image",
     "PiD",
     "Ideogram 4",
+    "SeFi-Image",
+    "Krea2",
     "ESRGAN",
 };
 
@@ -124,7 +132,8 @@ static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
            sd_version_is_flux2(version) ||
            sd_version_is_qwen_image(version) ||
            sd_version_is_longcat(version) ||
-           sd_version_is_z_image(version);
+           sd_version_is_z_image(version) ||
+           sd_version_is_boogu_image(version);
 }
 
 static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) {
@@ -196,6 +205,7 @@ public:
     bool enable_mmap                     = false;
     sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
     bool stream_layers = false;
+    bool eager_load    = false;
     std::string backend_spec;
     std::string params_backend_spec;
 
@@ -339,6 +349,7 @@ public:
         n_threads           = sd_ctx_params->n_threads;
         enable_mmap         = sd_ctx_params->enable_mmap;
         stream_layers       = sd_ctx_params->stream_layers;
+        eager_load          = sd_ctx_params->eager_load;
         backend_spec        = SAFE_STR(sd_ctx_params->backend);
         params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
         max_vram_assignment.reset(0.f);
@@ -527,7 +538,6 @@ public:
         if (wtype != GGML_TYPE_COUNT || tensor_type_rules.size() > 0) {
             model_loader.set_wtype_override(wtype, tensor_type_rules);
         }
-        model_loader.process_model_files(enable_mmap, true);
 
         std::map<ggml_type, uint32_t> wtype_stat                 = model_loader.get_wtype_stat();
         std::map<ggml_type, uint32_t> conditioner_wtype_stat     = model_loader.get_conditioner_wtype_stat();
@@ -581,9 +591,12 @@ public:
             apply_lora_immediately = false;
         }
 
+        bool needs_writable_mmap = enable_mmap && apply_lora_immediately;
+        model_manager->set_writable_mmap(needs_writable_mmap);
         if (enable_mmap && apply_lora_immediately) {
             LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap");
         }
+        model_loader.process_model_files(enable_mmap, needs_writable_mmap);
         load_alphas_cumprod(model_loader);
 
         size_t text_encoder_params_mem_size = 0;
@@ -638,6 +651,17 @@ public:
                                                                                tensor_storage_map,
                                                                                "model.diffusion_model",
                                                                                model_manager);
+            } else if (sd_version_is_krea2(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 false,
+                                                                 model_manager);
+                diffusion_model  = std::make_shared<Krea2::Krea2Runner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                       tensor_storage_map,
+                                                                       "model.diffusion_model",
+                                                                       model_manager);
             } else if (sd_version_is_flux(version)) {
                 bool is_chroma = false;
                 for (auto pair : tensor_storage_map) {
@@ -671,7 +695,7 @@ public:
                                                                      version,
                                                                      sd_ctx_params->chroma_use_dit_mask,
                                                                      model_manager);
-            } else if (sd_version_is_flux2(version)) {
+            } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
                 bool is_chroma   = false;
                 cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
@@ -731,13 +755,14 @@ public:
                     }
                 }
             } else if (sd_version_is_qwen_image(version)) {
-                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
+                cond_stage_model   = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
                                                                  version,
                                                                  "",
-                                                                 true,
+                                                                 enable_vision,
                                                                  model_manager);
-                diffusion_model  = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
+                diffusion_model    = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
                                                                           tensor_storage_map,
                                                                           "model.diffusion_model",
                                                                           version,
@@ -764,6 +789,14 @@ public:
                                                                                tensor_storage_map,
                                                                                "model",
                                                                                model_manager);
+            } else if (sd_version_is_minit2i(version)) {
+                cond_stage_model = std::make_shared<MiniT2IConditioner>(backend_for(SDBackendModule::TE),
+                                                                        tensor_storage_map,
+                                                                        model_manager);
+                diffusion_model  = std::make_shared<MiniT2I::MiniT2IRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                           tensor_storage_map,
+                                                                           "model.diffusion_model.model.net",
+                                                                           model_manager);
             } else if (sd_version_is_anima(version)) {
                 cond_stage_model = std::make_shared<AnimaConditioner>(backend_for(SDBackendModule::TE),
                                                                       tensor_storage_map,
@@ -784,6 +817,18 @@ public:
                                                                          "model.diffusion_model",
                                                                          version,
                                                                          model_manager);
+            } else if (sd_version_is_boogu_image(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 true,
+                                                                 model_manager);
+                diffusion_model  = std::make_shared<Boogu::BooguImageRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model.diffusion_model",
+                                                                            version,
+                                                                            model_manager);
             } else if (sd_version_is_ernie_image(version)) {
                 cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
@@ -860,10 +905,7 @@ public:
             }
 
             auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
-                if (sd_version_is_wan(version) ||
-                    sd_version_is_qwen_image(version) ||
-                    sd_version_is_anima(version) ||
-                    sd_version_is_ltxav(version)) {
+                if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) {
                     return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),
                                                                   tensor_storage_map,
                                                                   "decoder",
@@ -900,9 +942,7 @@ public:
                                                          false,
                                                          version,
                                                          model_manager);
-                } else if (sd_version_is_wan(version) ||
-                           sd_version_is_qwen_image(version) ||
-                           sd_version_is_anima(version)) {
+                } else if (sd_version_uses_wan_vae(version)) {
                     return std::make_shared<WAN::WanVAERunner>(backend_for(SDBackendModule::VAE),
                                                                tensor_storage_map,
                                                                "first_stage_model",
@@ -930,7 +970,7 @@ public:
                 }
             };
 
-            if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
+            if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
                 LOG_INFO("using FakeVAE");
                 first_stage_model = std::make_shared<FakeVAE>(version,
                                                               backend_for(SDBackendModule::VAE),
@@ -1138,7 +1178,15 @@ public:
             return false;
         }
 
-        LOG_DEBUG("model metadata validated; weights will be prepared lazily");
+        if (eager_load) {
+            if (!model_manager->load_all_params_eagerly()) {
+                LOG_ERROR("model params eager load failed");
+                return false;
+            }
+            LOG_DEBUG("model metadata validated; weights pre-loaded to params backend");
+        } else {
+            LOG_DEBUG("model metadata validated; weights will be prepared lazily");
+        }
 
         {
             size_t total_params_ram_size  = 0;
@@ -1220,6 +1268,7 @@ public:
                            sd_version_is_anima(version) ||
                            sd_version_is_ernie_image(version) ||
                            sd_version_is_z_image(version) ||
+                           sd_version_is_boogu_image(version) ||
                            sd_version_is_pid(version) ||
                            sd_version_is_ideogram4(version)) {
                     pred_type = FLOW_PRED;
@@ -1231,13 +1280,17 @@ public:
                         default_flow_shift = 1.5f;
                     } else if (sd_version_is_ideogram4(version)) {
                         default_flow_shift = 1.0f;
+                    } else if (sd_version_is_boogu_image(version)) {
+                        default_flow_shift = 3.16f;
                     } else {
                         default_flow_shift = 3.f;
                     }
                 } else if (sd_version_is_flux(version) ||
+                           sd_version_is_flux2(version) ||
                            sd_version_is_longcat(version) ||
                            sd_version_is_lens(version) ||
-                           sd_version_is_ltxav(version)) {
+                           sd_version_is_ltxav(version) ||
+                           sd_version_is_krea2(version)) {
                     pred_type = FLUX_FLOW_PRED;
 
                     default_flow_shift = 1.0f;  // TODO: validate
@@ -1253,9 +1306,13 @@ public:
                         default_flow_shift = 1.83f;
                     } else if (sd_version_is_ltxav(version)) {
                         default_flow_shift = 2.37f;
+                    } else if (sd_version_is_krea2(version)) {
+                        default_flow_shift = 1.15f;
                     }
-                } else if (sd_version_is_flux2(version)) {
-                    pred_type = FLUX2_FLOW_PRED;
+                } else if (sd_version_is_sefi_image(version)) {
+                    pred_type = SEFI_FLOW_PRED;
+                } else if (sd_version_is_minit2i(version)) {
+                    pred_type = MINIT2I_FLOW_PRED;
                 } else {
                     pred_type = EPS_PRED;
                 }
@@ -1288,9 +1345,14 @@ public:
                     denoiser = std::make_shared<FluxFlowDenoiser>();
                     break;
                 }
-                case FLUX2_FLOW_PRED: {
-                    LOG_INFO("running in Flux2 FLOW mode");
-                    denoiser = std::make_shared<Flux2FlowDenoiser>();
+                case SEFI_FLOW_PRED: {
+                    LOG_INFO("running in SeFi-Image dual-time FLOW mode");
+                    denoiser = std::make_shared<SefiFlowDenoiser>();
+                    break;
+                }
+                case MINIT2I_FLOW_PRED: {
+                    LOG_INFO("running in MiniT2I FLOW mode");
+                    denoiser = std::make_shared<MiniT2IFlowDenoiser>();
                     break;
                 }
                 default: {
@@ -1598,7 +1660,16 @@ public:
 
     std::vector<float> process_timesteps(const std::vector<float>& timesteps,
                                          const sd::Tensor<float>& init_latent,
-                                         const sd::Tensor<float>& denoise_mask) {
+                                         const sd::Tensor<float>& denoise_mask,
+                                         int step) {
+        if (auto sefi_denoiser = std::dynamic_pointer_cast<SefiFlowDenoiser>(denoiser)) {
+            int sched_idx = step > 0 ? step - 1 : 0;
+            if (sched_idx >= static_cast<int>(sefi_denoiser->tex_timesteps.size())) {
+                sched_idx = static_cast<int>(sefi_denoiser->tex_timesteps.size()) - 1;
+            }
+            return {sefi_denoiser->sem_timesteps[sched_idx],
+                    sefi_denoiser->tex_timesteps[sched_idx]};
+        }
         if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") {
             int64_t frame_count = init_latent.shape()[2];
             auto new_timesteps  = std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
@@ -1691,10 +1762,10 @@ public:
                 if (sd_version_is_sd3(version)) {
                     latent_rgb_proj = sd3_latent_rgb_proj;
                     latent_rgb_bias = sd3_latent_rgb_bias;
-                } else if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_longcat(version)) {
+                } else if (sd_version_uses_flux_vae(version)) {
                     latent_rgb_proj = flux_latent_rgb_proj;
                     latent_rgb_bias = flux_latent_rgb_bias;
-                } else if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_anima(version)) {
+                } else if (sd_version_uses_wan_vae(version)) {
                     latent_rgb_proj = wan_21_latent_rgb_proj;
                     latent_rgb_bias = wan_21_latent_rgb_bias;
                 } else {
@@ -1786,6 +1857,9 @@ public:
         if (sd_version_is_anima(version)) {
             return std::vector<float>{t / static_cast<float>(TIMESTEPS)};
         }
+        if (sd_version_is_boogu_image(version)) {
+            return std::vector<float>{t / static_cast<float>(TIMESTEPS)};
+        }
         if (version == VERSION_HIDREAM_O1) {
             return std::vector<float>{1.0f - (t / static_cast<float>(TIMESTEPS))};
         }
@@ -1911,6 +1985,32 @@ public:
         float slg_scale     = guidance.slg.scale;
         bool slg_uncond     = sd::guidance::parse_skip_layer_guidance_uncond_arg(extra_sample_args);
 
+        std::vector<float> guidance_schedule = sd::guidance::parse_guidance_schedule(extra_sample_args);
+        if (!guidance_schedule.empty() && guidance_schedule.size() != sigmas.size() - 1) {
+            if (guidance_schedule.size() > sigmas.size()) {
+                LOG_WARN("guidance_schedule length (%zu) is greater than number of steps (%zu)", guidance_schedule.size(), sigmas.size() - 1);
+                LOG_WARN("truncating guidance_schedule to match step count");
+                guidance_schedule.resize(sigmas.size() - 1);
+            } else {
+                LOG_INFO("padding guidance_schedule with cfg_scale");
+                while (guidance_schedule.size() < sigmas.size() - 1) {
+                    guidance_schedule.push_back(cfg_scale);
+                }
+            }
+        }
+
+        if (!guidance_schedule.empty()) {
+            std::string schedule_str = "[";
+            for (size_t i = 0; i < guidance_schedule.size(); ++i) {
+                schedule_str += std::to_string(guidance_schedule[i]);
+                if (i < guidance_schedule.size() - 1) {
+                    schedule_str += ", ";
+                }
+            }
+            schedule_str += "]";
+            LOG_DEBUG("using guidance schedule: %s", schedule_str.c_str());
+        }
+
         sd_sample::SampleCacheRuntime cache_runtime = sd_sample::init_sample_cache_runtime(version,
                                                                                            cache_params,
                                                                                            denoiser.get(),
@@ -1951,12 +2051,13 @@ public:
         }
 
         int64_t last_progress_us     = ggml_time_us();
-        sd::Tensor<float> x_t        = !noise.empty()
-                                           ? denoiser->noise_scaling(sigmas[0], noise, init_latent)
-                                           : init_latent;
-        sd::Tensor<float> denoised   = x_t;
         SamplePreviewContext preview = prepare_sample_preview_context();
 
+        sd::Tensor<float> x_t      = !noise.empty()
+                                         ? denoiser->noise_scaling(sigmas[0], noise, init_latent)
+                                         : init_latent;
+        sd::Tensor<float> denoised = x_t;
+
         auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
             if (get_cancel_flag() == SD_CANCEL_ALL) {
                 LOG_DEBUG("cancelling generation");
@@ -1981,7 +2082,7 @@ public:
                 timesteps_vec          = process_ltxav_video_timesteps(base_timesteps_vec, init_latent, denoise_mask);
                 audio_timesteps_tensor = sd::Tensor<float>({static_cast<int64_t>(base_timesteps_vec.size())}, base_timesteps_vec);
             } else {
-                timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask);
+                timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask, step);
             }
             const std::vector<float>& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty())
                                                                   ? base_timesteps_vec
@@ -2019,9 +2120,9 @@ public:
             sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
             std::vector<sd::Tensor<float>> controls;
             DiffusionParams diffusion_params;
-            diffusion_params.x                  = &noised_input;
-            diffusion_params.timesteps          = &timesteps_tensor;
-            diffusion_params.increase_ref_index = increase_ref_index;
+            diffusion_params.x              = &noised_input;
+            diffusion_params.timesteps      = &timesteps_tensor;
+            diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
             sd::guidance::GuidanceInput step_guidance_input;
             step_guidance_input.step          = step;
             step_guidance_input.schedule_size = sigmas.size();
@@ -2051,7 +2152,7 @@ public:
                     diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength};
                 } else if (sd_version_is_sd3(version)) {
                     diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers};
-                } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) {
+                } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
                     diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
                                                                 local_skip_layers};
                 } else if (sd_version_is_anima(version)) {
@@ -2074,6 +2175,9 @@ public:
                         audio_length,
                         frame_rate,
                         video_positions.empty() ? nullptr : &video_positions};
+                } else if (sd_version_is_minit2i(version)) {
+                    diffusion_params.extra = MiniT2IDiffusionExtra{
+                        condition.c_vector.empty() ? nullptr : &condition.c_vector};
                 } else {
                     diffusion_params.extra = std::monostate{};
                 }
@@ -2151,7 +2255,7 @@ public:
             guidance_input.pred_uncond     = uncond_out.empty() ? nullptr : &uncond_out;
             guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out;
 
-            sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {});
+            sd::guidance::GuiderOutput guided = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]);
             if (guided.pred.empty()) {
                 return {};
             }
@@ -2195,7 +2299,7 @@ public:
             return output;
         };
 
-        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args);
+        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser);
         if (x0_opt.empty()) {
             LOG_ERROR("Diffusion model sampling failed");
             if (control_net) {
@@ -2254,8 +2358,12 @@ public:
                 latent_channel = 3;
             } else if (version == VERSION_CHROMA_RADIANCE) {
                 latent_channel = 3;
+            } else if (sd_version_is_minit2i(version)) {
+                latent_channel = 3;
             } else if (sd_version_is_pid(version)) {
                 latent_channel = 3;
+            } else if (sd_version_is_sefi_image(version)) {
+                latent_channel = 144;
             } else if (sd_version_uses_flux2_vae(version)) {
                 latent_channel = 128;
             } else {
@@ -2265,6 +2373,10 @@ public:
         return latent_channel;
     }
 
+    int get_image_channels() const {
+        return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3;
+    }
+
     int get_image_seq_len(int h, int w) {
         int vae_scale_factor = get_vae_scale_factor();
         return (h / vae_scale_factor) * (w / vae_scale_factor);
@@ -2333,7 +2445,7 @@ public:
     }
 
     sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false) {
-        if (sd_version_is_pid(version)) {
+        if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) {
             return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
         }
         auto latents = first_stage_model->diffusion_to_vae_latents(x);
@@ -2476,6 +2588,10 @@ const char* scheduler_to_str[] = {
     "lcm",
     "bong_tangent",
     "ltx2",
+    "logit_normal",
+    "flux2",
+    "flux",
+    "beta",
 };
 
 const char* sd_scheduler_name(enum scheduler_t scheduler) {
@@ -2486,6 +2602,9 @@ const char* sd_scheduler_name(enum scheduler_t scheduler) {
 }
 
 enum scheduler_t str_to_scheduler(const char* str) {
+    if (!strcmp(str, "normal")) {
+        return DISCRETE_SCHEDULER;
+    }
     for (int i = 0; i < SCHEDULER_COUNT; i++) {
         if (!strcmp(str, scheduler_to_str[i])) {
             return (enum scheduler_t)i;
@@ -2500,7 +2619,8 @@ const char* prediction_to_str[] = {
     "edm_v",
     "sd3_flow",
     "flux_flow",
-    "flux2_flow",
+    "sefi_flow",
+    "minit2i_flow",
 };
 
 const char* sd_prediction_name(enum prediction_t prediction) {
@@ -2675,6 +2795,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->lora_apply_mode      = LORA_APPLY_AUTO;
     sd_ctx_params->max_vram             = nullptr;
     sd_ctx_params->stream_layers        = false;
+    sd_ctx_params->eager_load           = false;
     sd_ctx_params->enable_mmap          = false;
     sd_ctx_params->diffusion_flash_attn = false;
     sd_ctx_params->circular_x           = false;
@@ -2721,6 +2842,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "prediction: %s\n"
              "max_vram: %s\n"
              "stream_layers: %s\n"
+             "eager_load: %s\n"
              "backend: %s\n"
              "params_backend: %s\n"
              "flash_attn: %s\n"
@@ -2756,6 +2878,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              sd_prediction_name(sd_ctx_params->prediction),
              SAFE_STR(sd_ctx_params->max_vram),
              BOOL_STR(sd_ctx_params->stream_layers),
+             BOOL_STR(sd_ctx_params->eager_load),
              SAFE_STR(sd_ctx_params->backend),
              SAFE_STR(sd_ctx_params->params_backend),
              BOOL_STR(sd_ctx_params->flash_attn),
@@ -2841,6 +2964,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_img_gen_params->seed              = -1;
     sd_img_gen_params->batch_count       = 1;
     sd_img_gen_params->control_strength  = 0.9f;
+    sd_img_gen_params->qwen_image_layers = 3;
     sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
     sd_img_gen_params->pulid_params      = {nullptr, 1.0f};
     sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
@@ -2867,6 +2991,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "seed: %" PRId64
              "\n"
              "batch_count: %d\n"
+             "qwen_image_layers: %d\n"
              "ref_images_count: %d\n"
              "auto_resize_ref_image: %s\n"
              "increase_ref_index: %s\n"
@@ -2883,6 +3008,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->strength,
              sd_img_gen_params->seed,
              sd_img_gen_params->batch_count,
+             sd_img_gen_params->qwen_image_layers,
              sd_img_gen_params->ref_images_count,
              BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
              BOOL_STR(sd_img_gen_params->increase_ref_index),
@@ -3073,8 +3199,14 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
         return LCM_SCHEDULER;
     } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
         return SIMPLE_SCHEDULER;
+    } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux(sd_ctx->sd->version)) {
+        return FLUX_SCHEDULER;
+    } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux2(sd_ctx->sd->version)) {
+        return FLUX2_SCHEDULER;
     } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) {
         return LTX2_SCHEDULER;
+    } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) {
+        return LOGIT_NORMAL_SCHEDULER;
     }
     return DISCRETE_SCHEDULER;
 }
@@ -3144,6 +3276,7 @@ struct GenerationRequest {
     bool has_ref_images                      = false;
     const sd_cache_params_t* cache_params    = nullptr;
     int batch_count                          = 1;
+    int qwen_image_layers                    = 3;
     int shifted_timestep                     = 0;
     float strength                           = 1.f;
     float control_strength                   = 0.f;
@@ -3169,6 +3302,7 @@ struct GenerationRequest {
         diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor();
         seed                        = sd_img_gen_params->seed;
         batch_count                 = sd_img_gen_params->batch_count;
+        qwen_image_layers           = std::max(0, sd_img_gen_params->qwen_image_layers);
         clip_skip                   = sd_img_gen_params->clip_skip;
         shifted_timestep            = sd_img_gen_params->sample_params.shifted_timestep;
         strength                    = sd_img_gen_params->strength;
@@ -3873,6 +4007,34 @@ public:
     ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete;
 };
 
+static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels) {
+    if (image.empty()) {
+        return image;
+    }
+    GGML_ASSERT(image.dim() == 4);
+    int64_t current_channels = image.shape()[2];
+    if (current_channels == channels) {
+        return image;
+    }
+    if (channels == 4) {
+        sd::Tensor<float> alpha = sd::full<float>({image.shape()[0], image.shape()[1], 1, image.shape()[3]}, 1.f);
+        if (current_channels == 3) {
+            return sd::ops::concat(image, alpha, 2);
+        }
+        if (current_channels == 1) {
+            sd::Tensor<float> rgb = sd::ops::concat(image, image, 2);
+            rgb                   = sd::ops::concat(rgb, image, 2);
+            return sd::ops::concat(rgb, alpha, 2);
+        }
+    }
+    if (channels == 3 && current_channels >= 3) {
+        return sd::ops::slice(image, 2, 0, 3);
+    }
+    GGML_ABORT("cannot convert image tensor from %lld to %d channels",
+               (long long)current_channels,
+               channels);
+}
+
 static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
                                                                               const sd_img_gen_params_t* sd_img_gen_params,
                                                                               GenerationRequest* request,
@@ -3882,6 +4044,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
     sd::Tensor<float> init_image_tensor;
     sd::Tensor<float> control_image_tensor;
     sd::Tensor<float> mask_image_tensor;
+    int image_channels = sd_ctx->sd->get_image_channels();
 
     if (sd_img_gen_params->init_image.data != nullptr) {
         LOG_INFO("IMG2IMG");
@@ -3898,7 +4061,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             plan->sample_steps = static_cast<int>(plan->sigmas.size() - 1);
         }
 
-        init_image_tensor = sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height);
+        init_image_tensor = ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height),
+                                                         image_channels);
     }
 
     if (sd_img_gen_params->mask_image.data != nullptr) {
@@ -3930,7 +4094,11 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
     sd::Tensor<float> init_latent;
     sd::Tensor<float> control_latent;
     if (init_image_tensor.empty()) {
-        init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
+        if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
+            init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->qwen_image_layers + 1, true);
+        } else {
+            init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
+        }
     } else {
         init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor);
         if (init_latent.empty()) {
@@ -3949,12 +4117,13 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
 
     std::vector<sd::Tensor<float>> ref_images;
     for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) {
-        ref_images.push_back(sd_image_to_tensor(sd_img_gen_params->ref_images[i]));
+        ref_images.push_back(ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->ref_images[i]),
+                                                          image_channels));
     }
 
     if (ref_images.empty() && sd_version_is_unet_edit(sd_ctx->sd->version)) {
         LOG_WARN("This model needs at least one reference image; using an empty reference");
-        ref_images.push_back(sd::zeros<float>({request->width, request->height, 3, 1}));
+        ref_images.push_back(sd::zeros<float>({request->width, request->height, image_channels, 1}));
         request->guidance.img_cfg = request->guidance.txt_cfg;
         request->use_img_uncond   = false;
     }
@@ -3980,7 +4149,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             vae_width  = round(vae_width / factor) * factor;
 
             auto resized_ref_img = sd::ops::interpolate(ref_images[i],
-                                                        {static_cast<int>(vae_width), static_cast<int>(vae_height), 3, 1});
+                                                        {static_cast<int>(vae_width),
+                                                         static_cast<int>(vae_height),
+                                                         ref_images[i].shape()[2],
+                                                         ref_images[i].shape()[3]});
 
             LOG_DEBUG("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64,
                       static_cast<int>(i),
@@ -4125,6 +4297,11 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
     if (request->use_uncond || request->use_high_noise_uncond) {
         if (sd_version_is_ideogram4(sd_ctx->sd->version)) {
             uncond.c_vector = sd::Tensor<float>::from_vector({1.0f});
+        } else if (sd_version_is_minit2i(sd_ctx->sd->version)) {
+            // MiniT2I derives the unconditional signal from the same T5 hidden
+            // states with a zeroed prompt mask, so no extra text encode is needed.
+            uncond.c_crossattn = cond.c_crossattn;
+            uncond.c_vector    = sd::Tensor<float>::zeros_like(cond.c_vector);
         } else {
             bool zero_out_masked = false;
             if (sd_version_is_sdxl(sd_ctx->sd->version) &&
@@ -4180,7 +4357,8 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
 
 static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
                                         const GenerationRequest& request,
-                                        const std::vector<sd::Tensor<float>>& final_latents) {
+                                        const std::vector<sd::Tensor<float>>& final_latents,
+                                        int* num_images_out) {
     if (final_latents.empty()) {
         LOG_ERROR("no latent images to decode");
         return nullptr;
@@ -4204,13 +4382,41 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
             cancelled = true;
             break;
         }
-        int64_t t1              = ggml_time_ms();
-        sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
-        if (image.empty()) {
-            LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
-            return nullptr;
+        int64_t t1 = ggml_time_ms();
+        if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
+            int qwen_image_latent_layers = request.qwen_image_layers + 1;
+            if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) {
+                LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d",
+                          qwen_image_latent_layers,
+                          final_latents[i].dim());
+                return nullptr;
+            }
+            for (int layer_index = 0; layer_index < qwen_image_latent_layers; layer_index++) {
+                if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
+                    LOG_ERROR("cancelling latent decodings");
+                    cancelled = true;
+                    break;
+                }
+                sd::Tensor<float> layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1);
+                layer_latent.squeeze_(2);
+                sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(layer_latent);
+                if (image.empty()) {
+                    LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1);
+                    return nullptr;
+                }
+                decoded_images.push_back(std::move(image));
+            }
+            if (cancelled) {
+                break;
+            }
+        } else {
+            sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
+            if (image.empty()) {
+                LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
+                return nullptr;
+            }
+            decoded_images.push_back(std::move(image));
         }
-        decoded_images.push_back(std::move(image));
         int64_t t2 = ggml_time_ms();
         LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
     }
@@ -4222,11 +4428,14 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
         return nullptr;
     }
 
-    sd_image_t* result_images = (sd_image_t*)calloc(request.batch_count, sizeof(sd_image_t));
+    int image_count           = static_cast<int>(decoded_images.size());
+    sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t));
     if (result_images == nullptr) {
         return nullptr;
     }
-    memset(result_images, 0, request.batch_count * sizeof(sd_image_t));
+    if (num_images_out != nullptr) {
+        *num_images_out = image_count;
+    }
 
     for (size_t i = 0; i < decoded_images.size(); i++) {
         result_images[i] = tensor_to_sd_image(decoded_images[i]);
@@ -4419,9 +4628,18 @@ static std::vector<float> make_hires_sigma_schedule(sd_ctx_t* sd_ctx,
                               sigmas.end());
 }
 
-SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
+SD_API bool generate_image(sd_ctx_t* sd_ctx,
+                           const sd_img_gen_params_t* sd_img_gen_params,
+                           sd_image_t** images_out,
+                           int* num_images_out) {
+    if (images_out != nullptr) {
+        *images_out = nullptr;
+    }
+    if (num_images_out != nullptr) {
+        *num_images_out = 0;
+    }
     if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
-        return nullptr;
+        return false;
     }
 
     sd_ctx->sd->reset_cancel_flag();
@@ -4444,7 +4662,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                         &request,
                                                         &plan);
     if (!latents_opt.has_value()) {
-        return nullptr;
+        return false;
     }
     ImageGenerationLatents latents = std::move(*latents_opt);
 
@@ -4454,7 +4672,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                       &plan,
                                                       &latents);
     if (!embeds_opt.has_value()) {
-        return nullptr;
+        return false;
     }
     ImageGenerationEmbeds embeds = std::move(*embeds_opt);
 
@@ -4464,7 +4682,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         sd_cancel_mode_t cancel = sd_ctx->sd->get_cancel_flag();
         if (cancel == SD_CANCEL_ALL) {
             LOG_ERROR("cancelling generation");
-            return nullptr;
+            return false;
         }
         if (cancel == SD_CANCEL_NEW_LATENTS) {
             LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents",
@@ -4516,7 +4734,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                   b + 1,
                   request.batch_count,
                   (sampling_end - sampling_start) * 1.0f / 1000);
-        return nullptr;
+        return false;
     }
     int64_t denoise_end = ggml_time_ms();
     LOG_INFO("generating %zu latent images completed, taking %.2fs",
@@ -4524,13 +4742,13 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
              (denoise_end - denoise_start) * 1.0f / 1000);
     if (final_latents.empty()) {
         LOG_ERROR("no latent images generated");
-        return nullptr;
+        return false;
     }
 
     if (request.hires.enabled && request.hires.target_width > 0) {
         if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
             LOG_ERROR("cancelling generation before hires fix");
-            return nullptr;
+            return false;
         }
         LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
 
@@ -4538,7 +4756,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
             if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
                 LOG_ERROR("cancelling generation before hires model load");
-                return nullptr;
+                return false;
             }
             LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
             hires_upscaler                    = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
@@ -4551,7 +4769,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
             if (!hires_upscaler->load_from_file(request.hires.model_path,
                                                 sd_ctx->sd->n_threads)) {
                 LOG_ERROR("load hires model upscaler failed");
-                return nullptr;
+                return false;
             }
         }
 
@@ -4575,7 +4793,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         for (int b = 0; b < (int)final_latents.size(); b++) {
             if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
                 LOG_ERROR("cancelling generation during hires fix");
-                return nullptr;
+                return false;
             }
             int64_t cur_seed = request.seed + b;
             sd_ctx->sd->rng->manual_seed(cur_seed);
@@ -4586,7 +4804,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                               request,
                                                               hires_upscaler.get());
             if (upscaled.empty()) {
-                return nullptr;
+                return false;
             }
 
             sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
@@ -4640,7 +4858,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                       b + 1,
                       (int)final_latents.size(),
                       (hires_sample_end - hires_sample_start) * 1.0f / 1000);
-            return nullptr;
+            return false;
         }
         int64_t hires_denoise_end = ggml_time_ms();
         LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
@@ -4648,16 +4866,25 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         final_latents = std::move(hires_final_latents);
     }
 
-    auto result = decode_image_outputs(sd_ctx, request, final_latents);
+    int num_images = 0;
+    auto result    = decode_image_outputs(sd_ctx, request, final_latents, &num_images);
     if (result == nullptr) {
-        return nullptr;
+        return false;
     }
 
     sd_ctx->sd->lora_stat();
 
     int64_t t1 = ggml_time_ms();
     LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000);
-    return result;
+    if (num_images_out != nullptr) {
+        *num_images_out = num_images;
+    }
+    if (images_out != nullptr) {
+        *images_out = result;
+    } else {
+        free_sd_images(result, num_images);
+    }
+    return true;
 }
 
 static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd_ctx_t* sd_ctx,

src/upscaler.cpp

@@ -4,6 +4,7 @@
 #include "model_loader.h"
 #include "stable-diffusion.h"
 
+#include <cstdlib>
 #include <utility>
 
 UpscalerGGML::UpscalerGGML(int n_threads,
@@ -198,8 +199,41 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
     return upscaler_ctx;
 }
 
-sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor) {
-    return upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
+bool upscale(upscaler_ctx_t* upscaler_ctx,
+             sd_image_t input_image,
+             uint32_t upscale_factor,
+             sd_image_t** images_out,
+             int* num_images_out) {
+    if (images_out != nullptr) {
+        *images_out = nullptr;
+    }
+    if (num_images_out != nullptr) {
+        *num_images_out = 0;
+    }
+    if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) {
+        return false;
+    }
+
+    sd_image_t* result_images = (sd_image_t*)calloc(1, sizeof(sd_image_t));
+    if (result_images == nullptr) {
+        return false;
+    }
+
+    result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
+    if (result_images[0].data == nullptr) {
+        free(result_images);
+        return false;
+    }
+
+    if (num_images_out != nullptr) {
+        *num_images_out = 1;
+    }
+    if (images_out != nullptr) {
+        *images_out = result_images;
+    } else {
+        free_sd_images(result_images, 1);
+    }
+    return true;
 }
 
 int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {