feat: add beta scheduler (#811)

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**
@@ -50,12 +51,16 @@ API and command-line option may change frequently.***
     - [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/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/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

@@ -62,18 +62,22 @@ 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},
         };
 
@@ -762,8 +766,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 +806,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
             }

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",
@@ -761,6 +804,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 +884,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 +902,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},
     };
 
@@ -1364,6 +1422,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 +1475,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 +1521,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},
 
     };
 

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;

examples/server/async_jobs.cpp

@@ -173,8 +173,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();

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);
     }
 

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);
             }
 
@@ -438,6 +441,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

@@ -219,6 +219,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,7 @@ enum prediction_t {
     EDM_V_PRED,
     FLOW_PRED,
     FLUX_FLOW_PRED,
-    FLUX2_FLOW_PRED,
+    SEFI_FLOW_PRED,
     PREDICTION_COUNT
 };
 
@@ -219,6 +223,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;
@@ -450,7 +455,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 +488,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);
 

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>
@@ -142,8 +142,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                       std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
         : version(version), tokenizer(sd_version_is_sd2(version) ? 0 : 49407) {
         for (const auto& kv : orig_embedding_map) {
-            std::string name = kv.first;
-            std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return std::tolower(c); });
+            std::string name    = normalize_embedding_name(kv.first);
             embedding_map[name] = kv.second;
             tokenizer.add_special_token(name);
         }
@@ -278,17 +277,23 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         return true;
     }
 
+    static std::string normalize_embedding_name(std::string name) {
+        std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return std::tolower(c); });
+        return name;
+    }
+
+    bool append_embedding_tokens(std::string str, std::vector<int32_t>& bpe_tokens) {
+        std::string name = normalize_embedding_name(std::move(str));
+        auto iter        = embedding_map.find(name);
+        if (iter == embedding_map.end()) {
+            return false;
+        }
+        return load_embedding(name, iter->second, bpe_tokens);
+    }
+
     std::vector<int> convert_token_to_id(std::string text) {
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
-            auto iter = embedding_map.find(str);
-            if (iter == embedding_map.end()) {
-                return false;
-            }
-            std::string embedding_path = iter->second;
-            if (load_embedding(str, embedding_path, bpe_tokens)) {
-                return true;
-            }
-            return false;
+            return append_embedding_tokens(str, bpe_tokens);
         };
         std::vector<int> curr_tokens = tokenizer.encode(text, on_new_token_cb);
         return curr_tokens;
@@ -315,15 +320,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
-            auto iter = embedding_map.find(str);
-            if (iter == embedding_map.end()) {
-                return false;
-            }
-            std::string embedding_path = iter->second;
-            if (load_embedding(str, embedding_path, bpe_tokens)) {
-                return true;
-            }
-            return false;
+            return append_embedding_tokens(str, bpe_tokens);
         };
 
         std::vector<int> tokens;
@@ -1521,7 +1518,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;
@@ -1781,6 +1778,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;
 
@@ -1930,6 +1997,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/core/ggml_extend.hpp

@@ -1382,7 +1382,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);
             }
         }
     }

src/model.h

@@ -42,12 +42,15 @@ enum SDVersion {
     VERSION_LTXAV,
     VERSION_HIDREAM_O1,
     VERSION_Z_IMAGE,
+    VERSION_BOOGU_IMAGE,
     VERSION_OVIS_IMAGE,
     VERSION_ERNIE_IMAGE,
     VERSION_LENS,
     VERSION_LONGCAT,
     VERSION_PID,
     VERSION_IDEOGRAM4,
+    VERSION_SEFI_IMAGE,
+    VERSION_KREA2,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -143,6 +146,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;
@@ -178,8 +188,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 +244,14 @@ 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_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

@@ -899,10 +899,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/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/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));
@@ -1500,7 +1526,7 @@ 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)) {
+            if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
                 txt_arange_dims    = {3};
                 increase_ref_index = true;
             } else if (version == VERSION_OVIS_IMAGE) {

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/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>

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/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>

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>

src/model/vae/tae.hpp

@@ -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"

src/model/vae/wan_vae.hpp

@@ -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;
@@ -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) {
@@ -592,6 +636,7 @@ namespace WAN {
         std::vector<int> dim_mult;
         int num_res_blocks;
         std::vector<bool> temperal_downsample;
+        bool is_2D = false;
 
     public:
         Encoder3d(int64_t dim                           = 128,
@@ -599,23 +644,26 @@ namespace WAN {
                   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)
+                  bool wan2_2                           = false,
+                  bool is_2D                            = false)
             : 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}));
+            int64_t input_dim = wan2_2 ? 12 : 3;
+            if (is_2D) {
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(input_dim, 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(input_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
             }
 
             int index = 0;
@@ -630,12 +678,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 +697,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 +726,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 +784,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) {
@@ -758,6 +817,7 @@ namespace WAN {
         std::vector<int> dim_mult;
         int num_res_blocks;
         std::vector<bool> temperal_upsample;
+        bool is_2D = false;
 
     public:
         Decoder3d(int64_t dim                         = 128,
@@ -765,13 +825,15 @@ namespace WAN {
                   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)
+                  bool wan2_2                         = false,
+                  bool is_2D                          = false)
             : 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 +841,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 +865,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 +874,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 +888,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 = wan2_2 ? 12 : 3;
             // 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 +915,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 +973,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) {
@@ -935,6 +1009,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,8 +1026,8 @@ 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, bool wan2_2 = false, bool is_2D = false)
+            : decode_only(decode_only), wan2_2(wan2_2), is_2D(is_2D) {
             // attn_scales is always []
             if (wan2_2) {
                 dim     = 160;
@@ -962,12 +1037,26 @@ namespace WAN {
                 _conv_num     = 34;
                 _enc_conv_num = 26;
             }
+
+            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, 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, 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,6 +1119,10 @@ 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) {
@@ -1049,12 +1142,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 +1166,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,6 +1194,12 @@ 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) {
@@ -1129,7 +1244,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 == VERSION_WAN2_2_TI2V, is_2D);
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
 

src/model_loader.cpp

@@ -66,7 +66,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) {
@@ -453,6 +452,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;
         }
@@ -476,6 +479,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 +491,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;
         }

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,79 @@ 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 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;
-        float mu = a * n + b;
-
-        return mu;
+        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 {
-        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);
+    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());
+                }
+            }
+        }
+
+        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 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));
+
+            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;
     }
 };
 
@@ -939,6 +1441,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 +2390,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/stable-diffusion.cpp

@@ -20,11 +20,13 @@
 #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/mmdit.hpp"
@@ -87,12 +89,15 @@ const char* model_version_to_str[] = {
     "LTXAV",
     "HiDream O1",
     "Z-Image",
+    "Boogu Image",
     "Ovis Image",
     "Ernie Image",
     "Lens",
     "Longcat-Image",
     "PiD",
     "Ideogram 4",
+    "SeFi-Image",
+    "Krea2",
     "ESRGAN",
 };
 
@@ -124,7 +129,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 +202,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 +346,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 +535,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 +588,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 +648,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 +692,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,
@@ -784,6 +805,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 +893,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 +930,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",
@@ -1138,7 +1166,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 +1256,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 +1268,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 +1294,11 @@ 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 {
                     pred_type = EPS_PRED;
                 }
@@ -1288,9 +1331,9 @@ 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;
                 }
                 default: {
@@ -1598,7 +1641,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 +1743,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 +1838,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 +1966,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(),
@@ -1981,7 +2062,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
@@ -2051,7 +2132,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)) {
@@ -2151,7 +2232,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 +2276,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) {
@@ -2256,6 +2337,8 @@ public:
                 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 {
@@ -2476,6 +2559,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 +2573,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 +2590,7 @@ const char* prediction_to_str[] = {
     "edm_v",
     "sd3_flow",
     "flux_flow",
-    "flux2_flow",
+    "sefi_flow",
 };
 
 const char* sd_prediction_name(enum prediction_t prediction) {
@@ -2675,6 +2765,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 +2812,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 +2848,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),
@@ -3073,8 +3166,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;
 }
@@ -4180,7 +4279,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;
@@ -4222,11 +4322,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 +4522,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 +4556,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 +4566,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 +4576,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 +4628,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 +4636,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 +4650,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 +4663,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 +4687,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 +4698,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 +4752,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 +4760,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/tokenizers/bpe_tokenizer.cpp

@@ -134,7 +134,8 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
     std::vector<int32_t> bpe_tokens;
     std::vector<std::string> token_strs;
 
-    auto splited_texts = split_with_special_tokens(text, special_tokens);
+    std::string normalized_text = normalize_before_split ? normalize(text) : text;
+    auto splited_texts          = split_with_special_tokens(normalized_text, special_tokens);
 
     for (auto& splited_text : splited_texts) {
         if (is_special_token(splited_text)) {
@@ -159,7 +160,7 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
                 }
             }
 
-            std::string token_str = normalize(token);
+            std::string token_str = normalize_before_split ? token : normalize(token);
             std::u32string utf32_token;
             if (byte_level_bpe) {
                 for (int i = 0; i < token_str.length(); i++) {

src/tokenizers/clip_tokenizer.cpp

@@ -22,9 +22,10 @@ CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_st
     EOS_TOKEN_ID = 49407;
     PAD_TOKEN_ID = pad_token_id;
 
-    end_of_word_suffix = "</w>";
-    add_bos_token      = true;
-    add_eos_token      = true;
+    end_of_word_suffix     = "</w>";
+    add_bos_token          = true;
+    add_eos_token          = true;
+    normalize_before_split = true;
 
     if (merges_utf8_str.size() > 0) {
         load_from_merges(merges_utf8_str);

src/tokenizers/tokenizer.h

@@ -12,9 +12,10 @@ using on_new_token_cb_t = std::function<bool(std::string&, std::vector<int32_t>&
 class Tokenizer {
 protected:
     std::vector<std::string> special_tokens;
-    bool add_bos_token = false;
-    bool add_eos_token = false;
-    bool pad_left      = false;
+    bool add_bos_token          = false;
+    bool add_eos_token          = false;
+    bool pad_left               = false;
+    bool normalize_before_split = false;
     std::string end_of_word_suffix;
 
     virtual std::string decode_token(int token_id) const = 0;

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) {