.dockerignore

@@ -1,5 +1,4 @@
 build*/
-docs/
 test/
 
 .cache/

.github/workflows/build.yml

@@ -21,13 +21,11 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
-        "examples/server/frontend/**",
       ]
   pull_request:
     types: [opened, synchronize, reopened]
     paths:
       [
-        ".github/workflows/**",
         "**/CMakeLists.txt",
         "**/Makefile",
         "**/*.h",
@@ -35,16 +33,11 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
-        "examples/server/frontend/**",
       ]
 
 env:
   BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
 
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
-  cancel-in-progress: true
-
 jobs:
   ubuntu-latest-cmake:
     runs-on: ubuntu-latest
@@ -56,16 +49,6 @@ jobs:
         with:
           submodules: recursive
 
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
       - name: Dependencies
         id: depends
         run: |
@@ -83,7 +66,7 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        uses: pr-mpt/actions-commit-hash@v2
 
       - name: Fetch system info
         id: system-info
@@ -109,143 +92,6 @@ jobs:
           path: |
             sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
 
-  ubuntu-latest-cmake-vulkan:
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: recursive
-
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
-      - name: Dependencies
-        id: depends
-        run: |
-          sudo apt-get update
-          sudo apt-get install build-essential libvulkan-dev glslc
-
-      - name: Build
-        id: cmake_build
-        run: |
-          mkdir build
-          cd build
-          cmake .. -DSD_BUILD_SHARED_LIBS=ON -DSD_VULKAN=ON
-          cmake --build . --config Release
-
-      - name: Get commit hash
-        id: commit
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
-
-      - name: Fetch system info
-        id: system-info
-        run: |
-          echo "CPU_ARCH=`uname -m`" >> "$GITHUB_OUTPUT"
-          echo "OS_NAME=`lsb_release -s -i`" >> "$GITHUB_OUTPUT"
-          echo "OS_VERSION=`lsb_release -s -r`" >> "$GITHUB_OUTPUT"
-          echo "OS_TYPE=`uname -s`" >> "$GITHUB_OUTPUT"
-
-      - name: Pack artifacts
-        id: pack_artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        run: |
-          cp ggml/LICENSE ./build/bin/ggml.txt
-          cp LICENSE ./build/bin/stable-diffusion.cpp.txt
-          zip -j sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-vulkan.zip ./build/bin/*
-
-      - name: Upload artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-vulkan.zip
-          path: |
-            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-vulkan.zip
-
-  build-and-push-docker-images:
-    name: Build and push container images
-    runs-on: ubuntu-latest
-
-    permissions:
-      contents: read
-      packages: write
-      id-token: write
-      attestations: write
-      artifact-metadata: write
-
-    strategy:
-      matrix:
-        variant: [musa, sycl, vulkan, cuda]
-
-    env:
-      REGISTRY: ghcr.io
-      IMAGE_NAME: ${{ github.repository }}
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v6
-        with:
-          submodules: recursive
-
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
-      - name: Get commit hash
-        id: commit
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Log in to the container registry
-        uses: docker/login-action@v3
-        with:
-          registry: ${{ env.REGISTRY }}
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Extract metadata for Docker
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
-
-      - name: Free Disk Space (Ubuntu)
-        uses: jlumbroso/free-disk-space@v1.3.1
-        with:
-          # this might remove tools that are actually needed,
-          # if set to "true" but frees about 6 GB
-          tool-cache: false
-
-      - name: Build and push Docker image
-        id: build-push
-        uses: docker/build-push-action@v6
-        with:
-          platforms: linux/amd64
-          push: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-          file: Dockerfile.${{ matrix.variant }}
-          tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ env.BRANCH_NAME }}-${{ matrix.variant }}
-          labels: ${{ steps.meta.outputs.labels }}
-          annotations: ${{ steps.meta.outputs.annotations }}
-
   macOS-latest-cmake:
     runs-on: macos-latest
 
@@ -256,16 +102,6 @@ jobs:
         with:
           submodules: recursive
 
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
       - name: Dependencies
         id: depends
         run: |
@@ -283,7 +119,7 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        uses: pr-mpt/actions-commit-hash@v2
 
       - name: Fetch system info
         id: system-info
@@ -310,7 +146,7 @@ jobs:
             sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
 
   windows-latest-cmake:
-    runs-on: windows-2022
+    runs-on: windows-2025
 
     env:
       VULKAN_VERSION: 1.4.328.1
@@ -327,8 +163,8 @@ jobs:
           - build: "avx512"
             defines: "-DGGML_NATIVE=OFF -DGGML_AVX512=ON -DGGML_AVX=ON -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=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\"'"
-          - build: "vulkan"
+            defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120'"
+          - build: 'vulkan'
             defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON"
     steps:
       - name: Clone
@@ -337,16 +173,6 @@ jobs:
         with:
           submodules: recursive
 
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
       - name: Install cuda-toolkit
         id: cuda-toolkit
         if: ${{ matrix.build == 'cuda12' }}
@@ -365,17 +191,13 @@ jobs:
           Add-Content $env:GITHUB_ENV "VULKAN_SDK=C:\VulkanSDK\${env:VULKAN_VERSION}"
           Add-Content $env:GITHUB_PATH "C:\VulkanSDK\${env:VULKAN_VERSION}\bin"
 
-      - name: Activate MSVC environment
-        id: msvc_dev_cmd
-        uses: ilammy/msvc-dev-cmd@v1
-
       - name: Build
         id: cmake_build
         run: |
           mkdir build
           cd build
-          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 .
+          cmake .. ${{ matrix.defines }}
+          cmake --build . --config Release
 
       - name: Check AVX512F support
         id: check_avx512f
@@ -393,7 +215,7 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        uses: pr-mpt/actions-commit-hash@v2
 
       - name: Pack artifacts
         id: pack_artifacts
@@ -452,16 +274,6 @@ jobs:
         with:
           submodules: recursive
 
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
       - name: Cache ROCm Installation
         id: cache-rocm
         uses: actions/cache@v4
@@ -526,7 +338,7 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        uses: pr-mpt/actions-commit-hash@v2
 
       - name: Pack artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -548,156 +360,6 @@ jobs:
           path: |
             sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-x64.zip
 
-  ubuntu-latest-rocm:
-    runs-on: ubuntu-latest
-    container: rocm/dev-ubuntu-24.04:7.2
-
-    env:
-      ROCM_VERSION: "7.2"
-      UBUNTU_VERSION: "24.04"
-      GPU_TARGETS: "gfx1151;gfx1150;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
-
-    steps:
-      - run: apt-get update && apt-get install -y git
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-        with:
-          submodules: recursive
-
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 9
-
-      - 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: Dependencies
-        id: depends
-        run: |
-          sudo apt-get update
-          sudo apt install -y \
-            cmake \
-            hip-dev \
-            hipblas-dev \
-            ninja-build \
-            rocm-dev \
-            zip
-          # Clean apt caches to recover disk space
-          sudo apt clean
-          sudo rm -rf /var/lib/apt/lists/* || true
-
-      - name: Setup ROCm Environment
-        run: |
-          # Add ROCm to PATH for current session
-          echo "/opt/rocm/bin" >> $GITHUB_PATH
-
-          # Build regex pattern from ${{ env.GPU_TARGETS }} (match target as substring)
-          TARGET_REGEX="($(printf '%s' "${{ env.GPU_TARGETS }}" | sed 's/;/|/g'))"
-
-          # Remove library files for architectures we're not building for to save disk space
-          echo "Cleaning up unneeded architecture files..."
-          cd /opt/rocm/lib/rocblas/library
-          # Keep only our target architectures
-          for file in *; do
-            if printf '%s' "$file" | grep -q 'gfx'; then
-              if ! printf '%s' "$file" | grep -Eq "$TARGET_REGEX"; then
-                echo "Removing $file" &&
-                sudo rm -f "$file";
-              fi
-            fi
-          done
-
-          cd /opt/rocm/lib/hipblaslt/library
-          for file in *; do
-            if printf '%s' "$file" | grep -q 'gfx'; then
-              if ! printf '%s' "$file" | grep -Eq "$TARGET_REGEX"; then
-                echo "Removing $file" &&
-                sudo rm -f "$file";
-              fi
-            fi
-          done
-
-      - name: Build
-        id: cmake_build
-        run: |
-          mkdir build
-          cd build
-          cmake .. -G Ninja \
-            -DCMAKE_CXX_COMPILER=amdclang++ \
-            -DCMAKE_C_COMPILER=amdclang \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DSD_HIPBLAS=ON \
-            -DGPU_TARGETS="${{ env.GPU_TARGETS }}" \
-            -DAMDGPU_TARGETS="${{ env.GPU_TARGETS }}" \
-            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-            -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
-            -DSD_BUILD_SHARED_LIBS=ON
-          cmake --build . --config Release
-
-      - name: Get commit hash
-        id: commit
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
-
-      - name: Prepare artifacts
-        id: prepare_artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        run: |
-          # Copy licenses
-          cp ggml/LICENSE ./build/bin/ggml.txt
-          cp LICENSE ./build/bin/stable-diffusion.cpp.txt
-
-          # Move ROCm runtime libraries (to avoid double space consumption)
-          sudo mv /opt/rocm/lib/librocsparse.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libhsa-runtime64.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libamdhip64.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libhipblas.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libhipblaslt.so* ./build/bin/
-          sudo mv /opt/rocm/lib/librocblas.so* ./build/bin/
-          sudo mv /opt/rocm/lib/rocblas/ ./build/bin/
-          sudo mv /opt/rocm/lib/hipblaslt/ ./build/bin/
-
-      - name: Fetch system info
-        id: system-info
-        run: |
-          echo "CPU_ARCH=`uname -m`" >> "$GITHUB_OUTPUT"
-          echo "OS_NAME=`lsb_release -s -i`" >> "$GITHUB_OUTPUT"
-          echo "OS_VERSION=`lsb_release -s -r`" >> "$GITHUB_OUTPUT"
-          echo "OS_TYPE=`uname -s`" >> "$GITHUB_OUTPUT"
-
-      - name: Pack artifacts
-        id: pack_artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        run: |
-          cp ggml/LICENSE ./build/bin/ggml.txt
-          cp LICENSE ./build/bin/stable-diffusion.cpp.txt
-          zip -y -r sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm.zip ./build/bin
-
-      - name: Upload artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm.zip
-          path: |
-            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm.zip
-
   release:
     if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
 
@@ -705,9 +367,6 @@ jobs:
 
     needs:
       - ubuntu-latest-cmake
-      - ubuntu-latest-cmake-vulkan
-      - ubuntu-latest-rocm
-      - build-and-push-docker-images
       - macOS-latest-cmake
       - windows-latest-cmake
       - windows-latest-cmake-hip
@@ -733,7 +392,7 @@ jobs:
 
       - name: Get commit hash
         id: commit
-        uses: prompt/actions-commit-hash@v2
+        uses: pr-mpt/actions-commit-hash@v2
 
       - name: Create release
         id: create_release

.gitmodules

@@ -1,6 +1,3 @@
 [submodule "ggml"]
     path = ggml
 	url = https://github.com/ggml-org/ggml.git
-[submodule "examples/server/frontend"]
-	path = examples/server/frontend
-	url = https://github.com/leejet/stable-ui.git

CMakeLists.txt

@@ -8,11 +8,6 @@ if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
     set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
 endif()
 
-if (MSVC)
-    add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
-    add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
-endif()
-
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 
@@ -36,6 +31,7 @@ option(SD_VULKAN                     "sd: vulkan backend" OFF)
 option(SD_OPENCL                     "sd: opencl backend" OFF)
 option(SD_SYCL                       "sd: sycl backend" OFF)
 option(SD_MUSA                       "sd: musa backend" OFF)
+option(SD_FAST_SOFTMAX               "sd: x1.5 faster softmax, indeterministic (sometimes, same seed don't generate same image), cuda only" OFF)
 option(SD_BUILD_SHARED_LIBS          "sd: build shared libs" OFF)
 option(SD_BUILD_SHARED_GGML_LIB      "sd: build ggml as a separate shared lib" OFF)
 option(SD_USE_SYSTEM_GGML            "sd: use system-installed GGML library" OFF)
@@ -69,22 +65,26 @@ if (SD_HIPBLAS)
     message("-- Use HIPBLAS as backend stable-diffusion")
     set(GGML_HIP ON)
     add_definitions(-DSD_USE_CUDA)
+    if(SD_FAST_SOFTMAX)
+        set(GGML_CUDA_FAST_SOFTMAX ON)
+    endif()
 endif ()
 
 if(SD_MUSA)
     message("-- Use MUSA as backend stable-diffusion")
     set(GGML_MUSA ON)
     add_definitions(-DSD_USE_CUDA)
+    if(SD_FAST_SOFTMAX)
+        set(GGML_CUDA_FAST_SOFTMAX ON)
+    endif()
 endif()
 
 set(SD_LIB stable-diffusion)
 
 file(GLOB SD_LIB_SOURCES
-    "src/*.h"
-    "src/*.cpp"
-    "src/*.hpp"
-    "src/vocab/*.h"
-    "src/vocab/*.cpp"
+    "*.h"
+    "*.cpp"
+    "*.hpp"
 )
 
 find_program(GIT_EXE NAMES git git.exe NO_CMAKE_FIND_ROOT_PATH)
@@ -114,7 +114,7 @@ endif()
 message(STATUS "stable-diffusion.cpp commit ${SDCPP_BUILD_COMMIT}")
 
 set_property(
-  SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/src/version.cpp
+  SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp
   APPEND PROPERTY COMPILE_DEFINITIONS
   SDCPP_BUILD_COMMIT=${SDCPP_BUILD_COMMIT} SDCPP_BUILD_VERSION=${SDCPP_BUILD_VERSION}
 )
@@ -177,7 +177,6 @@ endif()
 add_subdirectory(thirdparty)
 
 target_link_libraries(${SD_LIB} PUBLIC ggml zip)
-target_include_directories(${SD_LIB} PUBLIC . include)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
 target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17)
 
@@ -186,7 +185,7 @@ if (SD_BUILD_EXAMPLES)
     add_subdirectory(examples)
 endif()
 
-set(SD_PUBLIC_HEADERS include/stable-diffusion.h)
+set(SD_PUBLIC_HEADERS stable-diffusion.h)
 set_target_properties(${SD_LIB} PROPERTIES PUBLIC_HEADER "${SD_PUBLIC_HEADERS}")
 
 install(TARGETS ${SD_LIB} LIBRARY PUBLIC_HEADER)

Dockerfile

@@ -1,4 +1,4 @@
-ARG UBUNTU_VERSION=24.04
+ARG UBUNTU_VERSION=22.04
 
 FROM ubuntu:$UBUNTU_VERSION AS build
 
@@ -18,6 +18,5 @@ RUN apt-get update && \
     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
 
 ENTRYPOINT [ "/sd-cli" ]

Dockerfile.cuda

@@ -1,25 +0,0 @@
-ARG CUDA_VERSION=12.6.3
-ARG UBUNTU_VERSION=24.04
-
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS build
-
-RUN apt-get update && apt-get install -y --no-install-recommends build-essential git ccache cmake
-
-WORKDIR /sd.cpp
-
-COPY . .
-
-ARG CUDACXX=/usr/local/cuda/bin/nvcc
-RUN cmake . -B ./build -DSD_CUDA=ON
-RUN cmake --build ./build --config Release -j$(nproc)
-
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime
-
-RUN apt-get update && \
-    apt-get install --yes --no-install-recommends libgomp1 && \
-    apt-get clean
-
-COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
-COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
-
-ENTRYPOINT [ "/sd-cli" ]

Dockerfile.musa

@@ -19,6 +19,5 @@ RUN mkdir build && cd build && \
 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
 
 ENTRYPOINT [ "/sd-cli" ]

Dockerfile.sycl

@@ -15,6 +15,5 @@ RUN mkdir build && cd build && \
 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" ]

Dockerfile.vulkan

@@ -1,23 +0,0 @@
-ARG UBUNTU_VERSION=24.04
-
-FROM ubuntu:$UBUNTU_VERSION AS build
-
-RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake libvulkan-dev glslc
-
-WORKDIR /sd.cpp
-
-COPY . .
-
-RUN cmake . -B ./build -DSD_VULKAN=ON
-RUN cmake --build ./build --config Release --parallel
-
-FROM ubuntu:$UBUNTU_VERSION AS runtime
-
-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
-
-ENTRYPOINT [ "/sd-cli" ]

README.md

@@ -15,9 +15,6 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
-* **2026/01/18** 🚀 stable-diffusion.cpp now supports **FLUX.2-klein**  
-  👉 Details: [PR #1193](https://github.com/leejet/stable-diffusion.cpp/pull/1193)
-
 * **2025/12/01** 🚀 stable-diffusion.cpp now supports **Z-Image**  
   👉 Details: [PR #1020](https://github.com/leejet/stable-diffusion.cpp/pull/1020)
 
@@ -46,17 +43,16 @@ API and command-line option may change frequently.***
     - SDXL, [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo)
     - [Some SD1.x and SDXL distilled models](./docs/distilled_sd.md)
     - [SD3/SD3.5](./docs/sd3.md)
-    - [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
-    - [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md)
+    - [FlUX.1-dev/FlUX.1-schnell](./docs/flux.md)
+    - [FLUX.2-dev](./docs/flux2.md)
     - [Chroma](./docs/chroma.md)
     - [Chroma1-Radiance](./docs/chroma_radiance.md)
     - [Qwen Image](./docs/qwen_image.md)
     - [Z-Image](./docs/z_image.md)
     - [Ovis-Image](./docs/ovis_image.md)
-    - [Anima](./docs/anima.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
-    - [Qwen Image Edit series](./docs/qwen_image_edit.md)
+    - [Qwen Image Edit/Qwen Image Edit 2509](./docs/qwen_image_edit.md)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
   - [PhotoMaker](https://github.com/TencentARC/PhotoMaker) support.
@@ -74,7 +70,7 @@ API and command-line option may change frequently.***
   - SYCL
 - Supported weight formats
   - Pytorch checkpoint (`.ckpt` or `.pth`)
-  - Safetensors (`.safetensors`)
+  - Safetensors (`./safetensors`)
   - GGUF (`.gguf`)
 - Supported platforms
     - Linux
@@ -131,16 +127,15 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 
 - [SD1.x/SD2.x/SDXL](./docs/sd.md)
 - [SD3/SD3.5](./docs/sd3.md)
-- [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
-- [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md)
+- [FlUX.1-dev/FlUX.1-schnell](./docs/flux.md)
+- [FLUX.2-dev](./docs/flux2.md)
 - [FLUX.1-Kontext-dev](./docs/kontext.md)
 - [Chroma](./docs/chroma.md)
 - [🔥Qwen Image](./docs/qwen_image.md)
-- [🔥Qwen Image Edit series](./docs/qwen_image_edit.md)
+- [🔥Qwen Image Edit/Qwen Image Edit 2509](./docs/qwen_image_edit.md)
 - [🔥Wan2.1/Wan2.2](./docs/wan.md)
 - [🔥Z-Image](./docs/z_image.md)
 - [Ovis-Image](./docs/ovis_image.md)
-- [Anima](./docs/anima.md)
 - [LoRA](./docs/lora.md)
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)
@@ -148,7 +143,6 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 - [Using TAESD to faster decoding](./docs/taesd.md)
 - [Docker](./docs/docker.md)
 - [Quantization and GGUF](./docs/quantization_and_gguf.md)
-- [Inference acceleration via caching](./docs/caching.md)
 
 ## Bindings
 

clip.hpp

@@ -4,7 +4,6 @@
 #include "ggml_extend.hpp"
 #include "model.h"
 #include "tokenize_util.h"
-#include "vocab/vocab.h"
 
 /*================================================== CLIPTokenizer ===================================================*/
 
@@ -111,7 +110,7 @@ public:
         if (merges_utf8_str.size() > 0) {
             load_from_merges(merges_utf8_str);
         } else {
-            load_from_merges(load_clip_merges());
+            load_from_merges(ModelLoader::load_merges());
         }
         add_special_token("<|startoftext|>");
         add_special_token("<|endoftext|>");
@@ -297,7 +296,7 @@ public:
                     size_t max_length = 0,
                     bool padding      = false) {
         if (max_length > 0 && padding) {
-            size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.0 / (max_length - 2)));
+            size_t n = std::ceil(tokens.size() * 1.0 / (max_length - 2));
             if (n == 0) {
                 n = 1;
             }
@@ -473,16 +472,16 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, n_token, d_model]
         auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
         auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
 
         x = fc1->forward(ctx, x);
         if (use_gelu) {
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_gelu_inplace(ctx->ggml_ctx, x);
         } else {
-            x = ggml_ext_gelu_quick(ctx->ggml_ctx, x, true);
+            x = ggml_gelu_quick_inplace(ctx->ggml_ctx, x);
         }
         x = fc2->forward(ctx, x);
         return x;
@@ -511,7 +510,7 @@ public:
         blocks["mlp"] = std::shared_ptr<GGMLBlock>(new CLIPMLP(d_model, intermediate_size));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask = nullptr) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, bool mask = true) {
         // x: [N, n_token, d_model]
         auto self_attn   = std::dynamic_pointer_cast<MultiheadAttention>(blocks["self_attn"]);
         auto layer_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm1"]);
@@ -526,10 +525,10 @@ public:
 
 struct CLIPEncoder : public GGMLBlock {
 protected:
-    int n_layer;
+    int64_t n_layer;
 
 public:
-    CLIPEncoder(int n_layer,
+    CLIPEncoder(int64_t n_layer,
                 int64_t d_model,
                 int64_t n_head,
                 int64_t intermediate_size,
@@ -541,10 +540,10 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* mask = nullptr,
-                         int clip_skip     = -1) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                int clip_skip = -1,
+                                bool mask     = true) {
         // x: [N, n_token, d_model]
         int layer_idx = n_layer - 1;
         // LOG_DEBUG("clip_skip %d", clip_skip);
@@ -573,7 +572,7 @@ protected:
     int64_t num_positions;
     bool force_clip_f32;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         enum ggml_type token_wtype = GGML_TYPE_F32;
         if (!force_clip_f32) {
             token_wtype = get_type(prefix + "token_embedding.weight", tensor_storage_map, GGML_TYPE_F32);
@@ -597,13 +596,13 @@ public:
           force_clip_f32(force_clip_f32) {
     }
 
-    ggml_tensor* get_token_embed_weight() {
+    struct ggml_tensor* get_token_embed_weight() {
         return params["token_embedding.weight"];
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids,
-                         ggml_tensor* custom_embed_weight) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* input_ids,
+                                struct ggml_tensor* custom_embed_weight) {
         // input_ids: [N, n_token]
         auto token_embed_weight    = params["token_embedding.weight"];
         auto position_embed_weight = params["position_embedding.weight"];
@@ -624,13 +623,13 @@ public:
 class CLIPVisionEmbeddings : public GGMLBlock {
 protected:
     int64_t embed_dim;
-    int num_channels;
-    int patch_size;
-    int image_size;
-    int num_patches;
+    int64_t num_channels;
+    int64_t patch_size;
+    int64_t image_size;
+    int64_t num_patches;
     int64_t num_positions;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         enum ggml_type patch_wtype    = GGML_TYPE_F16;
         enum ggml_type class_wtype    = GGML_TYPE_F32;
         enum ggml_type position_wtype = GGML_TYPE_F32;
@@ -642,9 +641,9 @@ protected:
 
 public:
     CLIPVisionEmbeddings(int64_t embed_dim,
-                         int num_channels = 3,
-                         int patch_size   = 14,
-                         int image_size   = 224)
+                         int64_t num_channels = 3,
+                         int64_t patch_size   = 14,
+                         int64_t image_size   = 224)
         : embed_dim(embed_dim),
           num_channels(num_channels),
           patch_size(patch_size),
@@ -653,7 +652,7 @@ public:
         num_positions = num_patches + 1;
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* pixel_values) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* pixel_values) {
         // pixel_values: [N, num_channels, image_size, image_size]
         // return: [N, num_positions, embed_dim]
         GGML_ASSERT(pixel_values->ne[0] == image_size && pixel_values->ne[1] == image_size && pixel_values->ne[2] == num_channels);
@@ -663,20 +662,20 @@ public:
         auto position_embed_weight = params["position_embedding.weight"];
 
         // concat(patch_embedding, class_embedding) + position_embedding
-        ggml_tensor* patch_embedding;
+        struct ggml_tensor* patch_embedding;
         int64_t N       = pixel_values->ne[3];
         patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size);  // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
         patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N);                          // [N, embed_dim, num_patches]
         patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3));                  // [N, num_patches, embed_dim]
         patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N);                       // [N, num_patches, embed_dim, 1]
 
-        ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
-        class_embedding              = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding);      // [N, embed_dim]
-        class_embedding              = ggml_reshape_4d(ctx->ggml_ctx, class_embedding, 1, embed_dim, 1, N);  // [N, 1, embed_dim, 1]
+        struct ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
+        class_embedding                     = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding);      // [N, embed_dim]
+        class_embedding                     = ggml_reshape_4d(ctx->ggml_ctx, class_embedding, 1, embed_dim, 1, N);  // [N, 1, embed_dim, 1]
 
-        ggml_tensor* x = ggml_concat(ctx->ggml_ctx, class_embedding, patch_embedding, 2);  // [N, num_positions, embed_dim, 1]
-        x              = ggml_reshape_3d(ctx->ggml_ctx, x, embed_dim, num_positions, N);   // [N, num_positions, embed_dim]
-        x              = ggml_add(ctx->ggml_ctx, x, position_embed_weight);
+        struct ggml_tensor* x = ggml_concat(ctx->ggml_ctx, class_embedding, patch_embedding, 2);  // [N, num_positions, embed_dim, 1]
+        x                     = ggml_reshape_3d(ctx->ggml_ctx, x, embed_dim, num_positions, N);   // [N, num_positions, embed_dim]
+        x                     = ggml_add(ctx->ggml_ctx, x, position_embed_weight);
         return x;  // [N, num_positions, embed_dim]
     }
 };
@@ -693,7 +692,7 @@ enum CLIPVersion {
 
 class CLIPTextModel : public GGMLBlock {
 protected:
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         if (version == OPEN_CLIP_VIT_BIGG_14) {
             enum ggml_type wtype      = GGML_TYPE_F32;
             params["text_projection"] = ggml_new_tensor_2d(ctx, wtype, projection_dim, hidden_size);
@@ -734,25 +733,24 @@ public:
         blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
     }
 
-    ggml_tensor* get_token_embed_weight() {
+    struct ggml_tensor* get_token_embed_weight() {
         auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
         return embeddings->get_token_embed_weight();
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids,
-                         ggml_tensor* tkn_embeddings,
-                         ggml_tensor* mask    = nullptr,
-                         size_t max_token_idx = 0,
-                         bool return_pooled   = false,
-                         int clip_skip        = -1) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* input_ids,
+                                struct ggml_tensor* tkn_embeddings,
+                                size_t max_token_idx = 0,
+                                bool return_pooled   = false,
+                                int clip_skip        = -1) {
         // input_ids: [N, n_token]
         auto embeddings       = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
         auto encoder          = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
         auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
 
         auto x = embeddings->forward(ctx, input_ids, tkn_embeddings);  // [N, n_token, hidden_size]
-        x      = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip);
+        x      = encoder->forward(ctx, x, return_pooled ? -1 : clip_skip, true);
         if (return_pooled || with_final_ln) {
             x = final_layer_norm->forward(ctx, x);
         }
@@ -804,10 +802,10 @@ public:
         blocks["post_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* pixel_values,
-                         bool return_pooled = true,
-                         int clip_skip      = -1) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* pixel_values,
+                                bool return_pooled = true,
+                                int clip_skip      = -1) {
         // pixel_values: [N, num_channels, image_size, image_size]
         auto embeddings     = std::dynamic_pointer_cast<CLIPVisionEmbeddings>(blocks["embeddings"]);
         auto pre_layernorm  = std::dynamic_pointer_cast<LayerNorm>(blocks["pre_layernorm"]);
@@ -816,11 +814,10 @@ public:
 
         auto x = embeddings->forward(ctx, pixel_values);  // [N, num_positions, embed_dim]
         x      = pre_layernorm->forward(ctx, x);
-        x      = encoder->forward(ctx, x, nullptr, clip_skip);
-
+        x      = encoder->forward(ctx, x, clip_skip, false);
+        // print_ggml_tensor(x, true, "ClipVisionModel x: ");
         auto last_hidden_state = x;
-
-        x = post_layernorm->forward(ctx, x);  // [N, n_token, hidden_size]
+        x                      = post_layernorm->forward(ctx, x);  // [N, n_token, hidden_size]
 
         GGML_ASSERT(x->ne[3] == 1);
         if (return_pooled) {
@@ -839,7 +836,7 @@ protected:
     int64_t out_features;
     bool transpose_weight;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
         if (transpose_weight) {
             params["weight"] = ggml_new_tensor_2d(ctx, wtype, out_features, in_features);
@@ -856,8 +853,8 @@ public:
           out_features(out_features),
           transpose_weight(transpose_weight) {}
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-        ggml_tensor* w = params["weight"];
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+        struct ggml_tensor* w = params["weight"];
         if (transpose_weight) {
             w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
         }
@@ -886,10 +883,10 @@ public:
         blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new CLIPProjection(hidden_size, projection_dim, transpose_proj_w));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* pixel_values,
-                         bool return_pooled = true,
-                         int clip_skip      = -1) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* pixel_values,
+                                bool return_pooled = true,
+                                int clip_skip      = -1) {
         // pixel_values: [N, num_channels, image_size, image_size]
         // return: [N, projection_dim] if return_pooled else [N, n_token, hidden_size]
         auto vision_model      = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
@@ -908,8 +905,6 @@ public:
 struct CLIPTextModelRunner : public GGMLRunner {
     CLIPTextModel model;
 
-    std::vector<float> attention_mask_vec;
-
     CLIPTextModelRunner(ggml_backend_t backend,
                         bool offload_params_to_cpu,
                         const String2TensorStorage& tensor_storage_map,
@@ -936,17 +931,16 @@ struct CLIPTextModelRunner : public GGMLRunner {
         return "clip";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids,
-                         ggml_tensor* embeddings,
-                         ggml_tensor* mask,
-                         size_t max_token_idx = 0,
-                         bool return_pooled   = false,
-                         int clip_skip        = -1) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* input_ids,
+                                struct ggml_tensor* embeddings,
+                                size_t max_token_idx = 0,
+                                bool return_pooled   = false,
+                                int clip_skip        = -1) {
         size_t N       = input_ids->ne[1];
         size_t n_token = input_ids->ne[0];
         if (input_ids->ne[0] > model.n_token) {
@@ -954,19 +948,20 @@ struct CLIPTextModelRunner : public GGMLRunner {
             input_ids = ggml_reshape_2d(ctx->ggml_ctx, input_ids, model.n_token, input_ids->ne[0] / model.n_token);
         }
 
-        return model.forward(ctx, input_ids, embeddings, mask, max_token_idx, return_pooled, clip_skip);
+        return model.forward(ctx, input_ids, embeddings, max_token_idx, return_pooled, clip_skip);
     }
 
-    ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
-                             int num_custom_embeddings    = 0,
-                             void* custom_embeddings_data = nullptr,
-                             size_t max_token_idx         = 0,
-                             bool return_pooled           = false,
-                             int clip_skip                = -1) {
-        ggml_cgraph* gf        = new_graph_custom(2048);
-        ggml_tensor* input_ids = make_input(input_ids_tensor);
+    struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
+                                    int num_custom_embeddings    = 0,
+                                    void* custom_embeddings_data = nullptr,
+                                    size_t max_token_idx         = 0,
+                                    bool return_pooled           = false,
+                                    int clip_skip                = -1) {
+        struct ggml_cgraph* gf = new_graph_custom(2048);
 
-        ggml_tensor* embeddings = nullptr;
+        input_ids = to_backend(input_ids);
+
+        struct ggml_tensor* embeddings = nullptr;
 
         if (num_custom_embeddings > 0 && custom_embeddings_data != nullptr) {
             auto token_embed_weight = model.get_token_embed_weight();
@@ -980,44 +975,28 @@ struct CLIPTextModelRunner : public GGMLRunner {
             embeddings = ggml_concat(compute_ctx, token_embed_weight, custom_embeddings, 1);
         }
 
-        int n_tokens = static_cast<int>(input_ids->ne[0]);
-        attention_mask_vec.resize(n_tokens * n_tokens);
-        for (int i0 = 0; i0 < n_tokens; i0++) {
-            for (int i1 = 0; i1 < n_tokens; i1++) {
-                float value = 0.f;
-                if (i0 > i1) {
-                    value = -INFINITY;
-                }
-                attention_mask_vec[i1 * n_tokens + i0] = value;
-            }
-        }
-        auto attention_mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, n_tokens, n_tokens);
-        set_backend_tensor_data(attention_mask, attention_mask_vec.data());
-
         auto runner_ctx = get_context();
 
-        ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, attention_mask, max_token_idx, return_pooled, clip_skip);
+        struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, max_token_idx, return_pooled, clip_skip);
 
         ggml_build_forward_expand(gf, hidden_states);
 
         return gf;
     }
 
-    sd::Tensor<float> compute(const int n_threads,
-                              const sd::Tensor<int32_t>& input_ids,
-                              int num_custom_embeddings,
-                              void* custom_embeddings_data,
-                              size_t max_token_idx,
-                              bool return_pooled,
-                              int clip_skip) {
-        auto get_graph = [&]() -> ggml_cgraph* {
+    bool compute(const int n_threads,
+                 struct ggml_tensor* input_ids,
+                 int num_custom_embeddings,
+                 void* custom_embeddings_data,
+                 size_t max_token_idx,
+                 bool return_pooled,
+                 int clip_skip,
+                 ggml_tensor** output,
+                 ggml_context* output_ctx = nullptr) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, true);
-        if (return_pooled) {
-            return take_or_empty(std::move(result));
-        }
-        return restore_trailing_singleton_dims(std::move(result), 3);
+        return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
     }
 };
 

common.hpp

@@ -1,5 +1,5 @@
-#ifndef __COMMON_BLOCK_HPP__
-#define __COMMON_BLOCK_HPP__
+#ifndef __COMMON_HPP__
+#define __COMMON_HPP__
 
 #include "ggml_extend.hpp"
 
@@ -23,12 +23,12 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
         if (vae_downsample) {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
-            x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            x = ggml_pad(ctx->ggml_ctx, x, 1, 1, 0, 0);
             x = conv->forward(ctx, x);
         } else {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["op"]);
@@ -52,7 +52,7 @@ public:
         blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
         auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
@@ -80,7 +80,7 @@ protected:
                                        std::pair<int, int> padding) {
         GGML_ASSERT(dims == 2 || dims == 3);
         if (dims == 3) {
-            return std::shared_ptr<GGMLBlock>(new Conv3d(in_channels, out_channels, {kernel_size.first, 1, 1}, {1, 1, 1}, {padding.first, 0, 0}));
+            return std::shared_ptr<GGMLBlock>(new Conv3dnx1x1(in_channels, out_channels, kernel_size.first, 1, padding.first));
         } else {
             return std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, out_channels, kernel_size, {1, 1}, padding));
         }
@@ -121,7 +121,7 @@ public:
         }
     }
 
-    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb = nullptr) {
+    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* emb = nullptr) {
         // For dims==3, we reduce dimension from 5d to 4d by merging h and w, in order not to change ggml
         // [N, c, t, h, w] => [N, c, t, h * w]
         // x: [N, channels, h, w] if dims == 2 else [N, channels, t, h, w]
@@ -188,19 +188,17 @@ public:
         blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out * 2));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [ne3, ne2, ne1, dim_in]
         // return: [ne3, ne2, ne1, dim_out]
         auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
         x          = proj->forward(ctx, x);  // [ne3, ne2, ne1, dim_out*2]
-        auto x_vec = ggml_ext_chunk(ctx->ggml_ctx, x, 2, 0, false);
+        auto x_vec = ggml_ext_chunk(ctx->ggml_ctx, x, 2, 0);
         x          = x_vec[0];  // [ne3, ne2, ne1, dim_out]
         auto gate  = x_vec[1];  // [ne3, ne2, ne1, dim_out]
 
-        gate = ggml_cont(ctx->ggml_ctx, gate);
-
-        gate = ggml_ext_gelu(ctx->ggml_ctx, gate, true);
+        gate = ggml_gelu_inplace(ctx->ggml_ctx, gate);
 
         x = ggml_mul(ctx->ggml_ctx, x, gate);  // [ne3, ne2, ne1, dim_out]
 
@@ -214,13 +212,13 @@ public:
         blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out, bias));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [ne3, ne2, ne1, dim_in]
         // return: [ne3, ne2, ne1, dim_out]
         auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
         x = proj->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_gelu_inplace(ctx->ggml_ctx, x);
         return x;
     }
 };
@@ -258,7 +256,7 @@ public:
         blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [ne3, ne2, ne1, dim]
         // return: [ne3, ne2, ne1, dim_out]
 
@@ -297,9 +295,9 @@ public:
         // to_out_1 is nn.Dropout(), skip for inference
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* context) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* context) {
         // x: [N, n_token, query_dim]
         // context: [N, n_context, context_dim]
         // return: [N, n_token, query_dim]
@@ -317,7 +315,7 @@ public:
         auto k = to_k->forward(ctx, context);  // [N, n_context, inner_dim]
         auto v = to_v->forward(ctx, context);  // [N, n_context, inner_dim]
 
-        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
+        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
 
         x = to_out_0->forward(ctx, x);  // [N, n_token, query_dim]
         return x;
@@ -355,9 +353,9 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* context) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* context) {
         // x: [N, n_token, query_dim]
         // context: [N, n_context, context_dim]
         // return: [N, n_token, query_dim]
@@ -406,7 +404,7 @@ protected:
     int64_t context_dim = 768;  // hidden_size, 1024 for VERSION_SD2
     bool use_linear     = false;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
         auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
         if (iter != tensor_storage_map.end()) {
             int64_t inner_dim = n_head * d_head;
@@ -456,9 +454,9 @@ public:
         }
     }
 
-    virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                 ggml_tensor* x,
-                                 ggml_tensor* context) {
+    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                        struct ggml_tensor* x,
+                                        struct ggml_tensor* context) {
         // x: [N, in_channels, h, w]
         // context: [N, max_position(aka n_token), hidden_size(aka context_dim)]
         auto norm     = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
@@ -510,7 +508,7 @@ public:
 
 class AlphaBlender : public GGMLBlock {
 protected:
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
         // Get the type of the "mix_factor" tensor from the input tensors map with the specified prefix
         enum ggml_type wtype = GGML_TYPE_F32;
         params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
@@ -530,23 +528,23 @@ public:
         // since mix_factor.shape is [1,], we don't need rearrange using  rearrange_pattern
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x_spatial,
-                         ggml_tensor* x_temporal) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x_spatial,
+                                struct ggml_tensor* x_temporal) {
         // image_only_indicator is always tensor([0.])
         float alpha = get_alpha();
         auto x      = ggml_add(ctx->ggml_ctx,
-                               ggml_ext_scale(ctx->ggml_ctx, x_spatial, alpha),
-                               ggml_ext_scale(ctx->ggml_ctx, x_temporal, 1.0f - alpha));
+                               ggml_scale(ctx->ggml_ctx, x_spatial, alpha),
+                               ggml_scale(ctx->ggml_ctx, x_temporal, 1.0f - alpha));
         return x;
     }
 };
 
 class VideoResBlock : public ResBlock {
 public:
-    VideoResBlock(int64_t channels,
-                  int64_t emb_channels,
-                  int64_t out_channels,
+    VideoResBlock(int channels,
+                  int emb_channels,
+                  int out_channels,
                   std::pair<int, int> kernel_size = {3, 3},
                   int64_t video_kernel_size       = 3,
                   int dims                        = 2)  // always 2
@@ -555,10 +553,10 @@ public:
         blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* emb,
-                         int num_video_frames) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* emb,
+                                int num_video_frames) {
         // x: [N, channels, h, w] aka [b*t, channels, h, w]
         // emb: [N, emb_channels] aka [b*t, emb_channels]
         // image_only_indicator is always tensor([0.])
@@ -590,4 +588,4 @@ public:
     }
 };
 
-#endif  // __COMMON_BLOCK_HPP__
+#endif  // __COMMON_HPP__

control.hpp

@@ -1,7 +1,8 @@
 #ifndef __CONTROL_HPP__
 #define __CONTROL_HPP__
 
-#include "common_block.hpp"
+#include "common.hpp"
+#include "ggml_extend.hpp"
 #include "model.h"
 
 #define CONTROL_NET_GRAPH_SIZE 1536
@@ -164,26 +165,26 @@ public:
         blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
     }
 
-    ggml_tensor* resblock_forward(std::string name,
-                                  GGMLRunnerContext* ctx,
-                                  ggml_tensor* x,
-                                  ggml_tensor* emb) {
+    struct ggml_tensor* resblock_forward(std::string name,
+                                         GGMLRunnerContext* ctx,
+                                         struct ggml_tensor* x,
+                                         struct ggml_tensor* emb) {
         auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
         return block->forward(ctx, x, emb);
     }
 
-    ggml_tensor* attention_layer_forward(std::string name,
-                                         GGMLRunnerContext* ctx,
-                                         ggml_tensor* x,
-                                         ggml_tensor* context) {
+    struct ggml_tensor* attention_layer_forward(std::string name,
+                                                GGMLRunnerContext* ctx,
+                                                struct ggml_tensor* x,
+                                                struct ggml_tensor* context) {
         auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
         return block->forward(ctx, x, context);
     }
 
-    ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
-                                          ggml_tensor* hint,
-                                          ggml_tensor* emb,
-                                          ggml_tensor* context) {
+    struct ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
+                                                 struct ggml_tensor* hint,
+                                                 struct ggml_tensor* emb,
+                                                 struct ggml_tensor* context) {
         int num_input_blocks = 15;
         auto h               = hint;
         for (int i = 0; i < num_input_blocks; i++) {
@@ -198,13 +199,13 @@ public:
         return h;
     }
 
-    std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                      ggml_tensor* x,
-                                      ggml_tensor* hint,
-                                      ggml_tensor* guided_hint,
-                                      ggml_tensor* timesteps,
-                                      ggml_tensor* context,
-                                      ggml_tensor* y = nullptr) {
+    std::vector<struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                             struct ggml_tensor* x,
+                                             struct ggml_tensor* hint,
+                                             struct ggml_tensor* guided_hint,
+                                             struct ggml_tensor* timesteps,
+                                             struct ggml_tensor* context,
+                                             struct ggml_tensor* y = nullptr) {
         // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
         // timesteps: [N,]
         // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
@@ -246,7 +247,7 @@ public:
             emb = ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
         }
 
-        std::vector<ggml_tensor*> outs;
+        std::vector<struct ggml_tensor*> outs;
 
         if (guided_hint == nullptr) {
             guided_hint = input_hint_block_forward(ctx, hint, emb, context);
@@ -310,13 +311,11 @@ struct ControlNet : public GGMLRunner {
     SDVersion version = VERSION_SD1;
     ControlNetBlock control_net;
 
-    ggml_backend_buffer_t control_buffer = nullptr;
+    ggml_backend_buffer_t control_buffer = nullptr;  // keep control output tensors in backend memory
     ggml_context* control_ctx            = nullptr;
-    std::vector<ggml_tensor*> control_outputs_ggml;
-    ggml_tensor* guided_hint_output_ggml = nullptr;
-    std::vector<sd::Tensor<float>> controls;
-    sd::Tensor<float> guided_hint;
-    bool guided_hint_cached = false;
+    std::vector<struct ggml_tensor*> controls;  // (12 input block outputs, 1 middle block output) SD 1.5
+    struct ggml_tensor* guided_hint = nullptr;  // guided_hint cache, for faster inference
+    bool guided_hint_cached         = false;
 
     ControlNet(ggml_backend_t backend,
                bool offload_params_to_cpu,
@@ -330,23 +329,23 @@ struct ControlNet : public GGMLRunner {
         free_control_ctx();
     }
 
-    void alloc_control_ctx(std::vector<ggml_tensor*> outs) {
-        ggml_init_params params;
+    void alloc_control_ctx(std::vector<struct ggml_tensor*> outs) {
+        struct ggml_init_params params;
         params.mem_size   = static_cast<size_t>(outs.size() * ggml_tensor_overhead()) + 1024 * 1024;
         params.mem_buffer = nullptr;
         params.no_alloc   = true;
         control_ctx       = ggml_init(params);
 
-        control_outputs_ggml.resize(outs.size() - 1);
+        controls.resize(outs.size() - 1);
 
         size_t control_buffer_size = 0;
 
-        guided_hint_output_ggml = ggml_dup_tensor(control_ctx, outs[0]);
-        control_buffer_size += ggml_nbytes(guided_hint_output_ggml);
+        guided_hint = ggml_dup_tensor(control_ctx, outs[0]);
+        control_buffer_size += ggml_nbytes(guided_hint);
 
         for (int i = 0; i < outs.size() - 1; i++) {
-            control_outputs_ggml[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
-            control_buffer_size += ggml_nbytes(control_outputs_ggml[i]);
+            controls[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
+            control_buffer_size += ggml_nbytes(controls[i]);
         }
 
         control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, runtime_backend);
@@ -363,10 +362,8 @@ struct ControlNet : public GGMLRunner {
             ggml_free(control_ctx);
             control_ctx = nullptr;
         }
-        guided_hint_output_ggml = nullptr;
-        guided_hint_cached      = false;
-        guided_hint             = {};
-        control_outputs_ggml.clear();
+        guided_hint        = nullptr;
+        guided_hint_cached = false;
         controls.clear();
     }
 
@@ -374,37 +371,33 @@ struct ControlNet : public GGMLRunner {
         return "control_net";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         control_net.get_param_tensors(tensors, prefix);
     }
 
-    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
-                             const sd::Tensor<float>& hint_tensor,
-                             const sd::Tensor<float>& timesteps_tensor,
-                             const sd::Tensor<float>& context_tensor = {},
-                             const sd::Tensor<float>& y_tensor       = {}) {
-        ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                    struct ggml_tensor* hint,
+                                    struct ggml_tensor* timesteps,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* y = nullptr) {
+        struct ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
 
-        ggml_tensor* x         = make_input(x_tensor);
-        ggml_tensor* hint      = nullptr;
-        ggml_tensor* timesteps = make_input(timesteps_tensor);
-        ggml_tensor* context   = make_optional_input(context_tensor);
-        ggml_tensor* y         = make_optional_input(y_tensor);
-
-        ggml_tensor* guided_hint_input = nullptr;
-        if (guided_hint_cached && !guided_hint.empty()) {
-            guided_hint_input = make_input(guided_hint);
-            hint              = nullptr;
+        x = to_backend(x);
+        if (guided_hint_cached) {
+            hint = nullptr;
         } else {
-            hint = make_input(hint_tensor);
+            hint = to_backend(hint);
         }
+        context   = to_backend(context);
+        y         = to_backend(y);
+        timesteps = to_backend(timesteps);
 
         auto runner_ctx = get_context();
 
         auto outs = control_net.forward(&runner_ctx,
                                         x,
                                         hint,
-                                        guided_hint_input,
+                                        guided_hint_cached ? guided_hint : nullptr,
                                         timesteps,
                                         context,
                                         y);
@@ -413,46 +406,36 @@ struct ControlNet : public GGMLRunner {
             alloc_control_ctx(outs);
         }
 
-        ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint_output_ggml));
+        ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint));
         for (int i = 0; i < outs.size() - 1; i++) {
-            ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], control_outputs_ggml[i]));
+            ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], controls[i]));
         }
 
         return gf;
     }
 
-    std::optional<std::vector<sd::Tensor<float>>> compute(int n_threads,
-                                                          const sd::Tensor<float>& x,
-                                                          const sd::Tensor<float>& hint,
-                                                          const sd::Tensor<float>& timesteps,
-                                                          const sd::Tensor<float>& context = {},
-                                                          const sd::Tensor<float>& y       = {}) {
+    bool compute(int n_threads,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* hint,
+                 struct ggml_tensor* timesteps,
+                 struct ggml_tensor* context,
+                 struct ggml_tensor* y,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) {
         // x: [N, in_channels, h, w]
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
-        auto get_graph = [&]() -> ggml_cgraph* {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(x, hint, timesteps, context, y);
         };
 
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false);
-        if (!compute_result.has_value()) {
-            return std::nullopt;
+        bool res = GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        if (res) {
+            // cache guided_hint
+            guided_hint_cached = true;
         }
-
-        if (guided_hint_output_ggml != nullptr) {
-            guided_hint = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(guided_hint_output_ggml),
-                                                          4);
-        }
-        controls.clear();
-        controls.reserve(control_outputs_ggml.size());
-        for (ggml_tensor* control : control_outputs_ggml) {
-            auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
-            GGML_ASSERT(!control_host.empty());
-            controls.push_back(std::move(control_host));
-        }
-        guided_hint_cached = true;
-        return controls;
+        return res;
     }
 
     bool load_from_file(const std::string& file_path, int n_threads) {

diffusion_model.hpp

@@ -0,0 +1,424 @@
+#ifndef __DIFFUSION_MODEL_H__
+#define __DIFFUSION_MODEL_H__
+
+#include "flux.hpp"
+#include "mmdit.hpp"
+#include "qwen_image.hpp"
+#include "unet.hpp"
+#include "wan.hpp"
+#include "z_image.hpp"
+
+struct DiffusionParams {
+    struct ggml_tensor* x                     = nullptr;
+    struct ggml_tensor* timesteps             = nullptr;
+    struct ggml_tensor* context               = nullptr;
+    struct ggml_tensor* c_concat              = nullptr;
+    struct ggml_tensor* y                     = nullptr;
+    struct ggml_tensor* guidance              = nullptr;
+    std::vector<ggml_tensor*> ref_latents     = {};
+    bool increase_ref_index                   = false;
+    int num_video_frames                      = -1;
+    std::vector<struct ggml_tensor*> controls = {};
+    float control_strength                    = 0.f;
+    struct ggml_tensor* vace_context          = nullptr;
+    float vace_strength                       = 1.f;
+    std::vector<int> skip_layers              = {};
+};
+
+struct DiffusionModel {
+    virtual std::string get_desc()                                                      = 0;
+    virtual bool compute(int n_threads,
+                         DiffusionParams diffusion_params,
+                         struct ggml_tensor** output     = nullptr,
+                         struct ggml_context* output_ctx = nullptr)                     = 0;
+    virtual void alloc_params_buffer()                                                  = 0;
+    virtual void free_params_buffer()                                                   = 0;
+    virtual void free_compute_buffer()                                                  = 0;
+    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) = 0;
+    virtual size_t get_params_buffer_size()                                             = 0;
+    virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
+    virtual int64_t get_adm_in_channels()             = 0;
+    virtual void set_flash_attn_enabled(bool enabled) = 0;
+};
+
+struct UNetModel : public DiffusionModel {
+    UNetModelRunner unet;
+
+    UNetModel(ggml_backend_t backend,
+              bool offload_params_to_cpu,
+              const String2TensorStorage& tensor_storage_map = {},
+              SDVersion version                              = VERSION_SD1)
+        : unet(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", version) {
+    }
+
+    std::string get_desc() override {
+        return unet.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        unet.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        unet.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        unet.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        unet.get_param_tensors(tensors, "model.diffusion_model");
+    }
+
+    size_t get_params_buffer_size() override {
+        return unet.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        unet.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return unet.unet.adm_in_channels;
+    }
+
+    void set_flash_attn_enabled(bool enabled) {
+        unet.set_flash_attention_enabled(enabled);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) override {
+        return unet.compute(n_threads,
+                            diffusion_params.x,
+                            diffusion_params.timesteps,
+                            diffusion_params.context,
+                            diffusion_params.c_concat,
+                            diffusion_params.y,
+                            diffusion_params.num_video_frames,
+                            diffusion_params.controls,
+                            diffusion_params.control_strength, output, output_ctx);
+    }
+};
+
+struct MMDiTModel : public DiffusionModel {
+    MMDiTRunner mmdit;
+
+    MMDiTModel(ggml_backend_t backend,
+               bool offload_params_to_cpu,
+               const String2TensorStorage& tensor_storage_map = {})
+        : mmdit(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model") {
+    }
+
+    std::string get_desc() override {
+        return mmdit.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        mmdit.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        mmdit.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        mmdit.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        mmdit.get_param_tensors(tensors, "model.diffusion_model");
+    }
+
+    size_t get_params_buffer_size() override {
+        return mmdit.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        mmdit.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768 + 1280;
+    }
+
+    void set_flash_attn_enabled(bool enabled) {
+        mmdit.set_flash_attention_enabled(enabled);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) override {
+        return mmdit.compute(n_threads,
+                             diffusion_params.x,
+                             diffusion_params.timesteps,
+                             diffusion_params.context,
+                             diffusion_params.y,
+                             output,
+                             output_ctx,
+                             diffusion_params.skip_layers);
+    }
+};
+
+struct FluxModel : public DiffusionModel {
+    Flux::FluxRunner flux;
+
+    FluxModel(ggml_backend_t backend,
+              bool offload_params_to_cpu,
+              const String2TensorStorage& tensor_storage_map = {},
+              SDVersion version                              = VERSION_FLUX,
+              bool use_mask                                  = false)
+        : flux(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", version, use_mask) {
+    }
+
+    std::string get_desc() override {
+        return flux.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        flux.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        flux.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        flux.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        flux.get_param_tensors(tensors, "model.diffusion_model");
+    }
+
+    size_t get_params_buffer_size() override {
+        return flux.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        flux.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768;
+    }
+
+    void set_flash_attn_enabled(bool enabled) {
+        flux.set_flash_attention_enabled(enabled);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) override {
+        return flux.compute(n_threads,
+                            diffusion_params.x,
+                            diffusion_params.timesteps,
+                            diffusion_params.context,
+                            diffusion_params.c_concat,
+                            diffusion_params.y,
+                            diffusion_params.guidance,
+                            diffusion_params.ref_latents,
+                            diffusion_params.increase_ref_index,
+                            output,
+                            output_ctx,
+                            diffusion_params.skip_layers);
+    }
+};
+
+struct WanModel : public DiffusionModel {
+    std::string prefix;
+    WAN::WanRunner wan;
+
+    WanModel(ggml_backend_t backend,
+             bool offload_params_to_cpu,
+             const String2TensorStorage& tensor_storage_map = {},
+             const std::string prefix                       = "model.diffusion_model",
+             SDVersion version                              = VERSION_WAN2)
+        : prefix(prefix), wan(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) {
+    }
+
+    std::string get_desc() override {
+        return wan.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        wan.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        wan.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        wan.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        wan.get_param_tensors(tensors, prefix);
+    }
+
+    size_t get_params_buffer_size() override {
+        return wan.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        wan.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768;
+    }
+
+    void set_flash_attn_enabled(bool enabled) {
+        wan.set_flash_attention_enabled(enabled);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) override {
+        return wan.compute(n_threads,
+                           diffusion_params.x,
+                           diffusion_params.timesteps,
+                           diffusion_params.context,
+                           diffusion_params.y,
+                           diffusion_params.c_concat,
+                           nullptr,
+                           diffusion_params.vace_context,
+                           diffusion_params.vace_strength,
+                           output,
+                           output_ctx);
+    }
+};
+
+struct QwenImageModel : public DiffusionModel {
+    std::string prefix;
+    Qwen::QwenImageRunner qwen_image;
+
+    QwenImageModel(ggml_backend_t backend,
+                   bool offload_params_to_cpu,
+                   const String2TensorStorage& tensor_storage_map = {},
+                   const std::string prefix                       = "model.diffusion_model",
+                   SDVersion version                              = VERSION_QWEN_IMAGE)
+        : prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) {
+    }
+
+    std::string get_desc() override {
+        return qwen_image.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        qwen_image.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        qwen_image.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        qwen_image.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        qwen_image.get_param_tensors(tensors, prefix);
+    }
+
+    size_t get_params_buffer_size() override {
+        return qwen_image.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        qwen_image.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768;
+    }
+
+    void set_flash_attn_enabled(bool enabled) {
+        qwen_image.set_flash_attention_enabled(enabled);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) override {
+        return qwen_image.compute(n_threads,
+                                  diffusion_params.x,
+                                  diffusion_params.timesteps,
+                                  diffusion_params.context,
+                                  diffusion_params.ref_latents,
+                                  true,  // increase_ref_index
+                                  output,
+                                  output_ctx);
+    }
+};
+
+struct ZImageModel : public DiffusionModel {
+    std::string prefix;
+    ZImage::ZImageRunner z_image;
+
+    ZImageModel(ggml_backend_t backend,
+                bool offload_params_to_cpu,
+                const String2TensorStorage& tensor_storage_map = {},
+                const std::string prefix                       = "model.diffusion_model",
+                SDVersion version                              = VERSION_Z_IMAGE)
+        : prefix(prefix), z_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) {
+    }
+
+    std::string get_desc() override {
+        return z_image.get_desc();
+    }
+
+    void alloc_params_buffer() override {
+        z_image.alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        z_image.free_params_buffer();
+    }
+
+    void free_compute_buffer() override {
+        z_image.free_compute_buffer();
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        z_image.get_param_tensors(tensors, prefix);
+    }
+
+    size_t get_params_buffer_size() override {
+        return z_image.get_params_buffer_size();
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        z_image.set_weight_adapter(adapter);
+    }
+
+    int64_t get_adm_in_channels() override {
+        return 768;
+    }
+
+    void set_flash_attn_enabled(bool enabled) {
+        z_image.set_flash_attention_enabled(enabled);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr) override {
+        return z_image.compute(n_threads,
+                               diffusion_params.x,
+                               diffusion_params.timesteps,
+                               diffusion_params.context,
+                               diffusion_params.ref_latents,
+                               true,  // increase_ref_index
+                               output,
+                               output_ctx);
+    }
+};
+
+#endif

docs/anima.md

@@ -1,21 +0,0 @@
-# How to Use
-
-## Download weights
-
-- Download Anima
-    - safetensors: https://huggingface.co/circlestone-labs/Anima/tree/main/split_files/diffusion_models
-    - gguf: https://huggingface.co/Bedovyy/Anima-GGUF/tree/main
-    - gguf Anima2: https://huggingface.co/JusteLeo/Anima2-GGUF/tree/main
-- Download vae
-    - safetensors: https://huggingface.co/circlestone-labs/Anima/tree/main/split_files/vae
-- Download Qwen3-0.6B-Base
-    - safetensors: https://huggingface.co/circlestone-labs/Anima/tree/main/split_files/text_encoders
-    - gguf: https://huggingface.co/mradermacher/Qwen3-0.6B-Base-GGUF/tree/main
-
-## Examples
-
-```sh
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\anima-preview.safetensors --vae ..\..\ComfyUI\models\vae\qwen_image_vae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_06b_base.safetensors  -p "a lovely cat holding a sign says 'anima.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --diffusion-fa
-```
-
-<img alt="anima image example" src="../assets/anima/example.png" />

docs/caching.md

@@ -1,141 +0,0 @@
-## Caching
-
-Caching methods accelerate diffusion inference by reusing intermediate computations when changes between steps are small.
-
-### Cache Modes
-
-| Mode | Target | Description |
-|------|--------|-------------|
-| `ucache` | UNET models | Condition-level caching with error tracking |
-| `easycache` | DiT models | Condition-level cache |
-| `dbcache` | DiT models | Block-level L1 residual threshold |
-| `taylorseer` | DiT models | Taylor series approximation |
-| `cache-dit` | DiT models | Combined DBCache + TaylorSeer |
-| `spectrum` | UNET and DiT models | Chebyshev + Taylor output forecasting |
-
-### UCache (UNET Models)
-
-UCache caches the residual difference (output - input) and reuses it when input changes are below threshold.
-
-```bash
-sd-cli -m model.safetensors -p "a cat" --cache-mode ucache --cache-option "threshold=1.5"
-```
-
-#### Parameters
-
-| Parameter | Description | Default |
-|-----------|-------------|---------|
-| `threshold` | Error threshold for reuse decision | 1.0 |
-| `start` | Start caching at this percent of steps | 0.15 |
-| `end` | Stop caching at this percent of steps | 0.95 |
-| `decay` | Error decay rate (0-1) | 1.0 |
-| `relative` | Scale threshold by output norm (0/1) | 1 |
-| `reset` | Reset error after computing (0/1) | 1 |
-
-#### Reset Parameter
-
-The `reset` parameter controls error accumulation behavior:
-
-- `reset=1` (default): Resets accumulated error after each computed step. More aggressive caching, works well with most samplers.
-- `reset=0`: Keeps error accumulated. More conservative, recommended for `euler_a` sampler.
-
-### EasyCache (DiT Models)
-
-Condition-level caching for DiT models. Caches and reuses outputs when input changes are below threshold.
-
-```bash
---cache-mode easycache --cache-option "threshold=0.3"
-```
-
-#### Parameters
-
-| Parameter | Description | Default |
-|-----------|-------------|---------|
-| `threshold` | Input change threshold for reuse | 0.2 |
-| `start` | Start caching at this percent of steps | 0.15 |
-| `end` | Stop caching at this percent of steps | 0.95 |
-
-### Cache-DIT (DiT Models)
-
-For DiT models like FLUX and QWEN, use block-level caching modes.
-
-#### DBCache
-
-Caches blocks based on L1 residual difference threshold:
-
-```bash
---cache-mode dbcache --cache-option "threshold=0.25,warmup=4"
-```
-
-#### TaylorSeer
-
-Uses Taylor series approximation to predict block outputs:
-
-```bash
---cache-mode taylorseer
-```
-
-#### Cache-DIT (Combined)
-
-Combines DBCache and TaylorSeer:
-
-```bash
---cache-mode cache-dit
-```
-
-#### Parameters
-
-| Parameter | Description | Default |
-|-----------|-------------|---------|
-| `Fn` | Front blocks to always compute | 8 |
-| `Bn` | Back blocks to always compute | 0 |
-| `threshold` | L1 residual difference threshold | 0.08 |
-| `warmup` | Steps before caching starts | 8 |
-
-#### SCM Options
-
-Steps Computation Mask controls which steps can be cached:
-
-```bash
---scm-mask "1,1,1,1,0,0,1,0,0,0,1,0,0,0,1,0,0,0,1,1"
-```
-
-Mask values: `1` = compute, `0` = can cache.
-
-| Policy | Description |
-|--------|-------------|
-| `dynamic` | Check threshold before caching |
-| `static` | Always cache on cacheable steps |
-
-```bash
---scm-policy dynamic
-```
-
-### Spectrum (UNET and DiT Models)
-
-Spectrum uses Chebyshev polynomial fitting blended with Taylor extrapolation to predict denoised outputs, skipping entire forward passes. Based on the paper [Spectrum: Adaptive Spectral Feature Forecasting for Efficient Diffusion Sampling](https://github.com/tingyu215/Spectrum).
-
-```bash
-sd-cli -m model.safetensors -p "a cat" --cache-mode spectrum
-```
-
-#### Parameters
-
-| Parameter | Description | Default |
-|-----------|-------------|---------|
-| `w` | Chebyshev vs Taylor blend weight (0=Taylor, 1=Chebyshev) | 0.40 |
-| `m` | Chebyshev polynomial degree | 3 |
-| `lam` | Ridge regression regularization | 1.0 |
-| `window` | Initial window size (compute every N steps) | 2 |
-| `flex` | Window growth per computed step after warmup | 0.50 |
-| `warmup` | Steps to always compute before caching starts | 4 |
-| `stop` | Stop caching at this fraction of total steps | 0.9 |
-
-```
-
-### Performance Tips
-
-- Start with default thresholds and adjust based on output quality
-- Lower threshold = better quality, less speedup
-- Higher threshold = more speedup, potential quality loss
-- More steps generally means more caching opportunities

docs/distilled_sd.md

@@ -1,8 +1,8 @@
-# Running distilled models: SSD1B, Vega and SDx.x with tiny U-Nets
+# Running distilled models: SSD1B and SDx.x with tiny U-Nets
 
 ## Preface 
 
-These models feature a reduced U-Net architecture. Unlike standard SDXL models, the SSD-1B and Vega U-Net contains only one middle block and fewer attention layers in its up- and down-blocks, resulting in significantly smaller file sizes. Using these models can reduce inference time by more than 33%. For more details, refer to Segmind's paper: https://arxiv.org/abs/2401.02677v1.
+These models feature a reduced U-Net architecture. Unlike standard SDXL models, the SSD-1B U-Net contains only one middle block and fewer attention layers in its up- and down-blocks, resulting in significantly smaller file sizes. Using these models can reduce inference time by more than 33%. For more details, refer to Segmind's paper: https://arxiv.org/abs/2401.02677v1.
 Similarly, SD1.x- and SD2.x-style models with a tiny U-Net consist of only 6 U-Net blocks, leading to very small files and time savings of up to 50%. For more information, see the paper: https://arxiv.org/pdf/2305.15798.pdf.
 
 ## SSD1B
@@ -17,17 +17,7 @@ Useful LoRAs are also available:
  * https://huggingface.co/seungminh/lora-swarovski-SSD-1B/resolve/main/pytorch_lora_weights.safetensors
  * https://huggingface.co/kylielee505/mylcmlorassd/resolve/main/pytorch_lora_weights.safetensors
 
-## Vega
-
-Segmind's Vega model is available online here:
-
- * https://huggingface.co/segmind/Segmind-Vega/resolve/main/segmind-vega.safetensors
- 
-VegaRT is an example for an LCM-LoRA:
-
- * https://huggingface.co/segmind/Segmind-VegaRT/resolve/main/pytorch_lora_weights.safetensors
-
-Both files can be used out-of-the-box, unlike the models described in next sections.
+These files can be used out-of-the-box, unlike the models described in the next section.
 
 
 ## SD1.x, SD2.x with tiny U-Nets
@@ -93,7 +83,7 @@ python convert_diffusers_to_original_stable_diffusion.py \
 The file segmind_tiny-sd.ckpt will be generated and is now ready for use with sd.cpp. You can follow a similar process for the other models mentioned above.
 
 
-##### Another available .ckpt file:
+### Another available .ckpt file:
 
  * https://huggingface.co/ClashSAN/small-sd/resolve/main/tinySDdistilled.ckpt
 
@@ -107,31 +97,3 @@ for key, value in ckpt['state_dict'].items():
         ckpt['state_dict'][key] = value.contiguous()
 torch.save(ckpt, "tinySDdistilled_fixed.ckpt")
 ```
-
-
-### SDXS-512
-
-Another very tiny and **incredibly fast**  model is SDXS by IDKiro et al.  The authors refer to it as *"Real-Time One-Step Latent Diffusion Models with Image Conditions"*. For details read the paper: https://arxiv.org/pdf/2403.16627 . Once again the authors removed some more blocks of U-Net part and unlike other SD1 models they use an adjusted _AutoEncoderTiny_ instead of default _AutoEncoderKL_ for the VAE part.
-
-##### 1. Download the diffusers model from  Hugging Face using Python:
-
-```python
-from diffusers import StableDiffusionPipeline
-pipe = StableDiffusionPipeline.from_pretrained("IDKiro/sdxs-512-dreamshaper")
-pipe.save_pretrained(save_directory="sdxs")
-```
-##### 2. Create a safetensors file
-
-```bash
-python convert_diffusers_to_original_stable_diffusion.py \
-    --model_path  sdxs  --checkpoint_path sdxs.safetensors --half --use_safetensors
-```
-
-##### 3. Run the model as follows:
-
-```bash
-~/stable-diffusion.cpp/build/bin/sd-cli -m sdxs.safetensors -p "portrait of a lovely cat" \
-  --cfg-scale 1 --steps 1
-```
-
-Both options: ``` --cfg-scale 1 ``` and  ``` --steps 1 ``` are mandatory here.                                                 

docs/docker.md

@@ -1,39 +1,15 @@
-# Docker
+## Docker
 
-## Run CLI
-
-```shell
-docker run --rm -v /path/to/models:/models -v /path/to/output/:/output ghcr.io/leejet/stable-diffusion.cpp:master [args...]
-# For example
-# docker run --rm -v ./models:/models -v ./build:/output ghcr.io/leejet/stable-diffusion.cpp:master -m /models/sd-v1-4.ckpt -p "a lovely cat" -v -o /output/output.png
-```
-
-## Run server
-
-```shell
-docker run --rm --init -v /path/to/models:/models -v /path/to/output/:/output -p "1234:1234" --entrypoint "/sd-server" ghcr.io/leejet/stable-diffusion.cpp:master [args...]
-# For example
-# docker run --rm --init -v ./models:/models -v ./build:/output -p "1234:1234" --entrypoint "/sd-server" ghcr.io/leejet/stable-diffusion.cpp:master -m /models/sd-v1-4.ckpt -p "a lovely cat" -v -o /output/output.png
-```
-
-## Building using Docker
+### Building using Docker
 
 ```shell
 docker build -t sd .
 ```
 
-## Building variants using Docker
-
-Vulkan:
+### Run
 
 ```shell
-docker build -f Dockerfile.vulkan -t sd .
-```
-
-## Run locally built image's CLI
-
-```shell
-docker run --rm -v /path/to/models:/models -v /path/to/output/:/output sd [args...]
+docker run -v /path/to/models:/models -v /path/to/output/:/output sd-cli [args...]
 # For example
-# docker run --rm -v ./models:/models -v ./build:/output sd -m /models/sd-v1-4.ckpt -p "a lovely cat" -v -o /output/output.png
+# docker run -v ./models:/models -v ./build:/output sd-cli -m /models/sd-v1-4.ckpt -p "a lovely cat" -v -o /output/output.png
 ```

docs/esrgan.md

@@ -1,6 +1,6 @@
 ## Using ESRGAN to upscale results
 
-You can use ESRGAN—such as the model [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)—to upscale the generated images and improve their overall resolution and clarity.
+You can use ESRGAN to upscale the generated images. At the moment, only the [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) model is supported. Support for more models of this architecture will be added soon.
 
 - Specify the model path using the `--upscale-model PATH` parameter. example:
 

docs/flux2.md

@@ -1,8 +1,6 @@
 # How to Use
 
-## Flux.2-dev
-
-### Download weights
+## Download weights
 
 - Download FLUX.2-dev
     - gguf: https://huggingface.co/city96/FLUX.2-dev-gguf/tree/main
@@ -11,7 +9,7 @@
 - Download Mistral-Small-3.2-24B-Instruct-2506-GGUF
     - gguf: https://huggingface.co/unsloth/Mistral-Small-3.2-24B-Instruct-2506-GGUF/tree/main
 
-### Examples
+## Examples
 
 ```
 .\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux2-dev-Q4_K_S.gguf --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\Mistral-Small-3.2-24B-Instruct-2506-Q4_K_M.gguf -r .\kontext_input.png -p "change 'flux.cpp' to 'flux2-dev.cpp'" --cfg-scale 1.0 --sampling-method euler -v --diffusion-fa --offload-to-cpu
@@ -19,74 +17,5 @@
 
 <img alt="flux2 example" src="../assets/flux2/example.png" />
 
-## Flux.2 klein 4B / Flux.2 klein base 4B
 
-### Download weights
 
-- Download FLUX.2-klein-4B
-    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-klein-4B
-    - gguf: https://huggingface.co/leejet/FLUX.2-klein-4B-GGUF/tree/main
-- Download FLUX.2-klein-base-4B
-    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B
-    - gguf: https://huggingface.co/leejet/FLUX.2-klein-base-4B-GGUF/tree/main
-- Download vae
-    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
-- Download Qwen3 4b
-    - safetensors: https://huggingface.co/Comfy-Org/flux2-klein-4B/tree/main/split_files/text_encoders
-    - gguf: https://huggingface.co/unsloth/Qwen3-4B-GGUF/tree/main
-
-### Examples
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux-2-klein-4b.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_4b.safetensors -p "a lovely cat" --cfg-scale 1.0 --steps 4 -v --offload-to-cpu --diffusion-fa
-```
-
-<img alt="flux2-klein-4b" src="../assets/flux2/flux2-klein-4b.png" />
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux-2-klein-4b.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_4b.safetensors -r .\kontext_input.png -p "change 'flux.cpp' to 'klein.cpp'" --cfg-scale 1.0 --sampling-method euler -v --diffusion-fa --offload-to-cpu --steps 4
-```
-
-<img alt="flux2-klein-4b-edit" src="../assets/flux2/flux2-klein-4b-edit.png" />
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux-2-klein-base-4b.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_4b.safetensors -p "a lovely cat" --cfg-scale 4.0 --steps 20 -v --offload-to-cpu --diffusion-fa
-```
-
-<img alt="flux2-klein-base-4b" src="../assets/flux2/flux2-klein-base-4b.png" />
-
-## Flux.2 klein 9B / Flux.2 klein base 9B
-
-### Download weights
-
-- Download FLUX.2-klein-9B
-    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-klein-9B
-    - gguf: https://huggingface.co/leejet/FLUX.2-klein-9B-GGUF/tree/main
-- Download FLUX.2-klein-base-9B
-    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B
-    - gguf: https://huggingface.co/leejet/FLUX.2-klein-base-9B-GGUF/tree/main
-- Download vae
-    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
-- Download Qwen3 8B
-    - safetensors: https://huggingface.co/Comfy-Org/flux2-klein-9B/tree/main/split_files/text_encoders
-    - gguf: https://huggingface.co/unsloth/Qwen3-8B-GGUF/tree/main
-
-### Examples
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux-2-klein-9b.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_8b.safetensors -p "a lovely cat" --cfg-scale 1.0 --steps 4 -v --offload-to-cpu --diffusion-fa
-```
-
-<img alt="flux2-klein-9b" src="../assets/flux2/flux2-klein-9b.png" />
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux-2-klein-9b.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_8b.safetensors -r .\kontext_input.png -p "change 'flux.cpp' to 'klein.cpp'" --cfg-scale 1.0 --sampling-method euler -v --diffusion-fa --offload-to-cpu --steps 4
-```
-
-<img alt="flux2-klein-9b-edit" src="../assets/flux2/flux2-klein-9b-edit.png" />
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\flux-2-klein-base-9b.safetensors --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_8b.safetensors -p "a lovely cat" --cfg-scale 4.0 --steps 20 -v --offload-to-cpu --diffusion-fa
-```
-
-<img alt="flux2-klein-base-9b" src="../assets/flux2/flux2-klein-base-9b.png" />

docs/qwen_image_edit.md

@@ -9,9 +9,6 @@
     - Qwen Image Edit 2509
         - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image-Edit_ComfyUI/tree/main/split_files/diffusion_models
         - gguf: https://huggingface.co/QuantStack/Qwen-Image-Edit-2509-GGUF/tree/main
-    - Qwen Image Edit 2511
-        - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image-Edit_ComfyUI/tree/main/split_files/diffusion_models
-        - gguf: https://huggingface.co/unsloth/Qwen-Image-Edit-2511-GGUF/tree/main
 - Download vae
     - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/vae
 - Download qwen_2.5_vl 7b
@@ -36,13 +33,3 @@
 ```
 
 <img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2509.png" />
-
-### Qwen Image Edit 2511
-
-To use the new Qwen Image Edit 2511 mode, the  `--qwen-image-zero-cond-t` flag must be enabled; otherwise, image editing quality will degrade significantly.
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\qwen-image-edit-2511-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\qwen_image_vae.safetensors  --llm ..\..\ComfyUI\models\text_encoders\qwen_2.5_vl_7b.safetensors --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'edit.cpp'"  --qwen-image-zero-cond-t
-```
-
-<img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2511.png" />

docs/taesd.md

@@ -15,25 +15,3 @@ curl -L -O https://huggingface.co/madebyollin/taesd/resolve/main/diffusion_pytor
 ```bash
 sd-cli -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --taesd ../models/diffusion_pytorch_model.safetensors
 ```
-
-### Qwen-Image and wan (TAEHV)
-
-sd.cpp also supports [TAEHV](https://github.com/madebyollin/taehv) (#937), which can be used for Qwen-Image and wan.
-
-- For **Qwen-Image and wan2.1 and wan2.2-A14B**, download the wan2.1 tae [safetensors weights](https://github.com/madebyollin/taehv/blob/main/safetensors/taew2_1.safetensors)
-  
-  Or curl
-  
-  ```bash
-  curl -L -O https://github.com/madebyollin/taehv/raw/refs/heads/main/safetensors/taew2_1.safetensors
-  ```
-
-- For **wan2.2-TI2V-5B**, use the wan2.2 tae [safetensors weights](https://github.com/madebyollin/taehv/blob/main/safetensors/taew2_2.safetensors)
-  
-  Or curl
-  
-  ```bash
-  curl -L -O https://github.com/madebyollin/taehv/raw/refs/heads/main/safetensors/taew2_2.safetensors
-  ```
-
-Then simply replace the `--vae xxx.safetensors` with `--tae xxx.safetensors` in the commands. If it still out of VRAM, add `--vae-conv-direct` to your command though might be slower.

docs/wan.md

@@ -39,9 +39,6 @@
         - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
     - wan_2.2_vae (for Wan2.2 TI2V 5B only)
         - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/blob/main/split_files/vae/wan2.2_vae.safetensors
-
-    > Wan models vae requires really much VRAM! If you do not have enough VRAM, please try tae instead, though the results may be poorer. For tae usage, please refer to [taesd](taesd.md)
-
 - Download umt5_xxl
     - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/text_encoders/umt5_xxl_fp16.safetensors
     - gguf: https://huggingface.co/city96/umt5-xxl-encoder-gguf/tree/main

docs/z_image.md

@@ -7,9 +7,6 @@ You can run Z-Image with stable-diffusion.cpp on GPUs with 4GB of VRAM — or ev
 - Download Z-Image-Turbo
     - safetensors: https://huggingface.co/Comfy-Org/z_image_turbo/tree/main/split_files/diffusion_models
     - gguf: https://huggingface.co/leejet/Z-Image-Turbo-GGUF/tree/main
-- Download Z-Image
-    - safetensors: https://huggingface.co/Comfy-Org/z_image/tree/main/split_files/diffusion_models
-    - gguf: https://huggingface.co/unsloth/Z-Image-GGUF/tree/main
 - Download vae
     - safetensors: https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main
 - Download Qwen3 4b
@@ -18,22 +15,12 @@ You can run Z-Image with stable-diffusion.cpp on GPUs with 4GB of VRAM — or ev
 
 ## Examples
 
-### Z-Image-Turbo
-
 ```
 .\bin\Release\sd-cli.exe --diffusion-model  z_image_turbo-Q3_K.gguf --vae ..\..\ComfyUI\models\vae\ae.sft  --llm ..\..\ComfyUI\models\text_encoders\Qwen3-4B-Instruct-2507-Q4_K_M.gguf -p "A cinematic, melancholic photograph of a solitary hooded figure walking through a sprawling, rain-slicked metropolis at night. The city lights are a chaotic blur of neon orange and cool blue, reflecting on the wet asphalt. The scene evokes a sense of being a single component in a vast machine. Superimposed over the image in a sleek, modern, slightly glitched font is the philosophical quote: 'THE CITY IS A CIRCUIT BOARD, AND I AM A BROKEN TRANSISTOR.' -- moody, atmospheric, profound, dark academic" --cfg-scale 1.0 -v --offload-to-cpu --diffusion-fa -H 1024 -W 512
 ```
 
 <img width="256" alt="z-image example" src="../assets/z_image/q3_K.png" />
 
-### Z-Image-Base
-
-```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\z_image_bf16.safetensors --vae ..\..\ComfyUI\models\vae\ae.sft  --llm ..\..\ComfyUI\models\text_encoders\qwen_3_4b.safetensors -p "A cinematic, melancholic photograph of a solitary hooded figure walking through a sprawling, rain-slicked metropolis at night. The city lights are a chaotic blur of neon orange and cool blue, reflecting on the wet asphalt. The scene evokes a sense of being a single component in a vast machine. Superimposed over the image in a sleek, modern, slightly glitched font is the philosophical quote: 'THE CITY IS A CIRCUIT BOARD, AND I AM A BROKEN TRANSISTOR.' -- moody, atmospheric, profound, dark academic" --cfg-scale 5.0 -v --offload-to-cpu --diffusion-fa -H 1024 -W 512
-```
-
-<img width="256" alt="z-image example" src="../assets/z_image/base_bf16.png" />
-
 ## Comparison of Different Quantization Types
 
 | bf16 | q8_0 | q6_K | q5_0 | q4_K | q4_0 | q3_K | q2_K|

easycache.hpp

@@ -1,15 +1,10 @@
-#ifndef __EASYCACHE_HPP__
-#define __EASYCACHE_HPP__
-
 #include <cmath>
 #include <limits>
 #include <unordered_map>
 #include <vector>
 
-#include "condition_cache_utils.hpp"
 #include "denoiser.hpp"
 #include "ggml_extend.hpp"
-#include "tensor.hpp"
 
 struct EasyCacheConfig {
     bool enabled          = false;
@@ -24,15 +19,15 @@ struct EasyCacheCacheEntry {
 
 struct EasyCacheState {
     EasyCacheConfig config;
-    Denoiser* denoiser           = nullptr;
-    float start_sigma            = std::numeric_limits<float>::max();
-    float end_sigma              = 0.0f;
-    bool initialized             = false;
-    bool initial_step            = true;
-    bool skip_current_step       = false;
-    bool step_active             = false;
-    const void* anchor_condition = nullptr;
-    std::unordered_map<const void*, EasyCacheCacheEntry> cache_diffs;
+    Denoiser* denoiser                  = nullptr;
+    float start_sigma                   = std::numeric_limits<float>::max();
+    float end_sigma                     = 0.0f;
+    bool initialized                    = false;
+    bool initial_step                   = true;
+    bool skip_current_step              = false;
+    bool step_active                    = false;
+    const SDCondition* anchor_condition = nullptr;
+    std::unordered_map<const SDCondition*, EasyCacheCacheEntry> cache_diffs;
     std::vector<float> prev_input;
     std::vector<float> prev_output;
     float output_prev_norm                = 0.0f;
@@ -125,30 +120,41 @@ struct EasyCacheState {
         return enabled() && step_active && skip_current_step;
     }
 
-    bool has_cache(const void* cond) const {
+    bool has_cache(const SDCondition* cond) const {
         auto it = cache_diffs.find(cond);
         return it != cache_diffs.end() && !it->second.diff.empty();
     }
 
-    void update_cache(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
+    void update_cache(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) {
         EasyCacheCacheEntry& entry = cache_diffs[cond];
-        sd::store_condition_cache_diff(&entry.diff, input, output);
+        size_t ne                  = static_cast<size_t>(ggml_nelements(output));
+        entry.diff.resize(ne);
+        float* out_data = (float*)output->data;
+        float* in_data  = (float*)input->data;
+        for (size_t i = 0; i < ne; ++i) {
+            entry.diff[i] = out_data[i] - in_data[i];
+        }
     }
 
-    void apply_cache(const void* cond, const sd::Tensor<float>& input, sd::Tensor<float>* output) {
+    void apply_cache(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) {
         auto it = cache_diffs.find(cond);
         if (it == cache_diffs.end() || it->second.diff.empty()) {
             return;
         }
-        sd::apply_condition_cache_diff(it->second.diff, input, output);
+        copy_ggml_tensor(output, input);
+        float* out_data                = (float*)output->data;
+        const std::vector<float>& diff = it->second.diff;
+        for (size_t i = 0; i < diff.size(); ++i) {
+            out_data[i] += diff[i];
+        }
     }
 
-    bool before_condition(const void* cond,
-                          const sd::Tensor<float>& input,
-                          sd::Tensor<float>* output,
+    bool before_condition(const SDCondition* cond,
+                          ggml_tensor* input,
+                          ggml_tensor* output,
                           float sigma,
                           int step_index) {
-        if (!enabled() || step_index < 0 || output == nullptr) {
+        if (!enabled() || step_index < 0) {
             return false;
         }
         if (step_index != current_step_index) {
@@ -175,12 +181,12 @@ struct EasyCacheState {
         if (!has_prev_input || !has_prev_output || !has_cache(cond)) {
             return false;
         }
-        size_t ne = static_cast<size_t>(input.numel());
+        size_t ne = static_cast<size_t>(ggml_nelements(input));
         if (prev_input.size() != ne) {
             return false;
         }
-        const float* input_data = input.data();
-        last_input_change       = 0.0f;
+        float* input_data = (float*)input->data;
+        last_input_change = 0.0f;
         for (size_t i = 0; i < ne; ++i) {
             last_input_change += std::fabs(input_data[i] - prev_input[i]);
         }
@@ -205,7 +211,7 @@ struct EasyCacheState {
         return false;
     }
 
-    void after_condition(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
+    void after_condition(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) {
         if (!step_is_active()) {
             return;
         }
@@ -214,16 +220,16 @@ struct EasyCacheState {
             return;
         }
 
-        size_t ne            = static_cast<size_t>(input.numel());
-        const float* in_data = input.data();
+        size_t ne      = static_cast<size_t>(ggml_nelements(input));
+        float* in_data = (float*)input->data;
         prev_input.resize(ne);
         for (size_t i = 0; i < ne; ++i) {
             prev_input[i] = in_data[i];
         }
         has_prev_input = true;
 
-        const float* out_data = output.data();
-        float output_change   = 0.0f;
+        float* out_data     = (float*)output->data;
+        float output_change = 0.0f;
         if (has_prev_output && prev_output.size() == ne) {
             for (size_t i = 0; i < ne; ++i) {
                 output_change += std::fabs(out_data[i] - prev_output[i]);
@@ -257,5 +263,3 @@ struct EasyCacheState {
         has_last_input_change  = false;
     }
 };
-
-#endif

esrgan.hpp

@@ -27,11 +27,11 @@ public:
         blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* lrelu(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [n, num_feat, h, w]
         // return: [n, num_feat, h, w]
 
@@ -51,7 +51,7 @@ public:
         x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
         auto x5    = conv5->forward(ctx, x_cat);
 
-        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
+        x5 = ggml_add(ctx->ggml_ctx, ggml_scale(ctx->ggml_ctx, x5, 0.2f), x);
         return x5;
     }
 };
@@ -64,7 +64,7 @@ public:
         blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [n, num_feat, h, w]
         // return: [n, num_feat, h, w]
 
@@ -76,7 +76,7 @@ public:
         out      = rdb2->forward(ctx, out);
         out      = rdb3->forward(ctx, out);
 
-        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
+        out = ggml_add(ctx->ggml_ctx, ggml_scale(ctx->ggml_ctx, out, 0.2f), x);
         return out;
     }
 };
@@ -112,11 +112,11 @@ public:
     int get_scale() { return scale; }
     int get_num_block() { return num_block; }
 
-    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* lrelu(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [n, num_in_ch, h, w]
         // return: [n, num_out_ch, h*scale, w*scale]
         auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
@@ -341,24 +341,27 @@ struct ESRGAN : public GGMLRunner {
         return success;
     }
 
-    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x) {
         if (!rrdb_net)
             return nullptr;
         constexpr int kGraphNodes = 1 << 16;  // 65k
-        ggml_cgraph* gf           = new_graph_custom(kGraphNodes);
-        ggml_tensor* x            = make_input(x_tensor);
+        struct ggml_cgraph* gf    = new_graph_custom(kGraphNodes);
+        x                         = to_backend(x);
 
-        auto runner_ctx  = get_context();
-        ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
+        auto runner_ctx         = get_context();
+        struct ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
         ggml_build_forward_expand(gf, out);
         return gf;
     }
 
-    sd::Tensor<float> compute(const int n_threads,
-                              const sd::Tensor<float>& x) {
-        auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
-        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
-        return result;
+    bool compute(const int n_threads,
+                 struct ggml_tensor* x,
+                 ggml_tensor** output,
+                 ggml_context* output_ctx = nullptr) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(x);
+        };
+        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
     }
 };
 

examples/cli/README.md

@@ -4,14 +4,11 @@
 usage: ./bin/sd-cli  [options]
 
 CLI Options:
-  -o, --output <string>       path to write result image to. you can use printf-style %d format specifiers for image sequences (default:
-                              ./output.png) (eg. output_%03d.png)
+  -o, --output <string>       path to write result image to (default: ./output.png)
   --preview-path <string>     path to write preview image to (default: ./preview.png)
   --preview-interval <int>    interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at
                               every step)
-  --output-begin-idx <int>    starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)
   --canny                     apply canny preprocessor (edge detection)
-  --convert-name              convert tensor name (for convert mode)
   -v, --verbose               print extra info
   --color                     colors the logging tags according to level
   --taesd-preview-only        prevents usage of taesd for decoding the final image. (for use with --preview tae)
@@ -34,7 +31,6 @@ Context Options:
   --high-noise-diffusion-model <string>    path to the standalone high noise diffusion model
   --vae <string>                           path to standalone vae model
   --taesd <string>                         path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
-  --tae <string>                           alias of --taesd
   --control-net <string>                   path to control net model
   --embd-dir <string>                      embeddings directory
   --lora-model-dir <string>                lora model directory
@@ -45,22 +41,17 @@ Context Options:
                                            CPU physical cores
   --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
   --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
+  --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --vae-tiling                             process vae in tiles to reduce memory usage
   --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
   --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM when needed
-  --mmap                                   whether to memory-map model
   --control-net-cpu                        keep controlnet in cpu (for low vram)
   --clip-on-cpu                            keep clip in cpu (for low vram)
   --vae-on-cpu                             keep vae in cpu (for low vram)
-  --fa                                     use flash attention
-  --diffusion-fa                           use flash attention in the diffusion model only
+  --diffusion-fa                           use flash attention in the diffusion model
   --diffusion-conv-direct                  use ggml_conv2d_direct in the diffusion model
   --vae-conv-direct                        use ggml_conv2d_direct in the vae model
-  --circular                               enable circular padding for convolutions
-  --circularx                              enable circular RoPE wrapping on x-axis (width) only
-  --circulary                              enable circular RoPE wrapping on y-axis (height) only
   --chroma-disable-dit-mask                disable dit mask for chroma
-  --qwen-image-zero-cond-t                 enable zero_cond_t for qwen image
   --chroma-enable-t5-mask                  enable t5 mask for chroma
   --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). If not specified, the default is the
                                            type of the weight file
@@ -101,7 +92,6 @@ Generation Options:
   --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for NitroSD-Realism around 250 and 500 for
                                            NitroSD-Vibrant
   --upscale-repeats <int>                  Run the ESRGAN upscaler this many times (default: 1)
-  --upscale-tile-size <int>                tile size for ESRGAN upscaling (default: 128)
   --cfg-scale <float>                      unconditional guidance scale: (default: 7.0)
   --img-cfg-scale <float>                  image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)
   --guidance <float>                       distilled guidance scale for models with guidance input (default: 3.5)
@@ -110,7 +100,6 @@ Generation Options:
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
   --eta <float>                            eta in DDIM, only for DDIM and TCD (default: 0)
-  --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
   --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input (default: 3.5)
@@ -127,23 +116,13 @@ Generation Options:
   --disable-auto-resize-ref-image          disable auto resize of ref images
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing,
-                                           tcd, res_multistep, res_2s] (default: euler for Flux/SD3/Wan, euler_a
-                                           otherwise)
+                                           tcd] (default: euler for Flux/SD3/Wan, euler_a otherwise)
   --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm,
-                                           ddim_trailing, tcd, res_multistep, res_2s] default: euler for Flux/SD3/Wan,
-                                           euler_a otherwise
-  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple,
-                                           kl_optimal, lcm, bong_tangent], default: discrete
-  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g., "14.61,7.8,3.5,0.0").
+                                           ddim_trailing, tcd] default: euler for Flux/SD3/Wan, euler_a otherwise
+  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm],
+                                           default: discrete
   --skip-layers                            layers to skip for SLG steps (default: [7,8,9])
   --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
   -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
-  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level),
-                                           'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)
-  --cache-option                           named cache params (key=value format, comma-separated). easycache/ucache:
-                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit: Fn=,Bn=,threshold=,warmup=;
-                                           spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=. Examples:
-                                           "threshold=0.25" or "threshold=1.5,reset=0" or "w=0.4,window=2"
-  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g., "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
-  --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
+  --easycache                              enable EasyCache for DiT models with optional "threshold,start_percent,end_percent" (default: 0.2,0.15,0.95)
 ```

examples/cli/avi_writer.h

@@ -172,9 +172,9 @@ int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int
 
         // Write '00dc' chunk (video frame)
         fwrite("00dc", 4, 1, f);
-        write_u32_le(f, (uint32_t)jpeg_data.size);
+        write_u32_le(f, jpeg_data.size);
         index[i].offset = ftell(f) - 8;
-        index[i].size   = (uint32_t)jpeg_data.size;
+        index[i].size   = jpeg_data.size;
         fwrite(jpeg_data.buf, 1, jpeg_data.size, f);
 
         // Align to even byte size

examples/cli/main.cpp

@@ -26,16 +26,13 @@ const char* previews_str[] = {
     "vae",
 };
 
-std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
-
 struct SDCliParams {
     SDMode mode             = IMG_GEN;
     std::string output_path = "output.png";
-    int output_begin_idx    = -1;
 
     bool verbose          = false;
+    bool version          = false;
     bool canny_preprocess = false;
-    bool convert_name     = false;
 
     preview_t preview_method = PREVIEW_NONE;
     int preview_interval     = 1;
@@ -53,7 +50,7 @@ struct SDCliParams {
         options.string_options = {
             {"-o",
              "--output",
-             "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png)",
+             "path to write result image to (default: ./output.png)",
              &output_path},
             {"",
              "--preview-path",
@@ -66,10 +63,6 @@ struct SDCliParams {
              "--preview-interval",
              "interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at every step)",
              &preview_interval},
-            {"",
-             "--output-begin-idx",
-             "starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)",
-             &output_begin_idx},
         };
 
         options.bool_options = {
@@ -77,14 +70,14 @@ struct SDCliParams {
              "--canny",
              "apply canny preprocessor (edge detection)",
              true, &canny_preprocess},
-            {"",
-             "--convert-name",
-             "convert tensor name (for convert mode)",
-             true, &convert_name},
             {"-v",
              "--verbose",
              "print extra info",
              true, &verbose},
+            {"",
+             "--version",
+             "print stable-diffusion.cpp version",
+             true, &version},
             {"",
              "--color",
              "colors the logging tags according to level",
@@ -113,8 +106,9 @@ struct SDCliParams {
                     }
                 }
                 if (mode_found == -1) {
-                    LOG_ERROR("error: invalid mode %s, must be one of [%s]\n",
-                              mode_c_str, SD_ALL_MODES_STR);
+                    fprintf(stderr,
+                            "error: invalid mode %s, must be one of [%s]\n",
+                            mode_c_str, SD_ALL_MODES_STR);
                     exit(1);
                 }
                 mode = (SDMode)mode_found;
@@ -134,7 +128,8 @@ struct SDCliParams {
                 }
             }
             if (preview_found == -1) {
-                LOG_ERROR("error: preview method %s", preview);
+                fprintf(stderr, "error: preview method %s\n",
+                        preview);
                 return -1;
             }
             preview_method = (preview_t)preview_found;
@@ -166,7 +161,7 @@ struct SDCliParams {
 
     bool process_and_check() {
         if (output_path.length() == 0) {
-            LOG_ERROR("error: the following arguments are required: output_path");
+            fprintf(stderr, "error: the following arguments are required: output_path\n");
             return false;
         }
 
@@ -186,7 +181,6 @@ struct SDCliParams {
             << "  verbose: " << (verbose ? "true" : "false") << ",\n"
             << "  color: " << (color ? "true" : "false") << ",\n"
             << "  canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
-            << "  convert_name: " << (convert_name ? "true" : "false") << ",\n"
             << "  preview_method: " << previews_str[preview_method] << ",\n"
             << "  preview_interval: " << preview_interval << ",\n"
             << "  preview_path: \"" << preview_path << "\",\n"
@@ -225,8 +219,20 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
     }
 }
 
+static std::string sd_basename(const std::string& path) {
+    size_t pos = path.find_last_of('/');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    pos = path.find_last_of('\\');
+    if (pos != std::string::npos) {
+        return path.substr(pos + 1);
+    }
+    return path;
+}
+
 std::string get_image_params(const SDCliParams& cli_params, const SDContextParams& ctx_params, const SDGenerationParams& gen_params, int64_t seed) {
-    std::string parameter_string = gen_params.prompt_with_lora + "\n";
+    std::string parameter_string = gen_params.prompt + "\n";
     if (gen_params.negative_prompt.size() != 0) {
         parameter_string += "Negative prompt: " + gen_params.negative_prompt + "\n";
     }
@@ -245,22 +251,14 @@ std::string get_image_params(const SDCliParams& cli_params, const SDContextParam
     parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", ";
     parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", ";
     parameter_string += "Seed: " + std::to_string(seed) + ", ";
-    parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
+    parameter_string += "Size: " + std::to_string(gen_params.width) + "x" + std::to_string(gen_params.height) + ", ";
     parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";
     parameter_string += "RNG: " + std::string(sd_rng_type_name(ctx_params.rng_type)) + ", ";
     if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) {
         parameter_string += "Sampler RNG: " + std::string(sd_rng_type_name(ctx_params.sampler_rng_type)) + ", ";
     }
     parameter_string += "Sampler: " + std::string(sd_sample_method_name(gen_params.sample_params.sample_method));
-    if (!gen_params.custom_sigmas.empty()) {
-        parameter_string += ", Custom Sigmas: [";
-        for (size_t i = 0; i < gen_params.custom_sigmas.size(); ++i) {
-            std::ostringstream oss;
-            oss << std::fixed << std::setprecision(4) << gen_params.custom_sigmas[i];
-            parameter_string += oss.str() + (i == gen_params.custom_sigmas.size() - 1 ? "" : ", ");
-        }
-        parameter_string += "]";
-    } else if (gen_params.sample_params.scheduler != SCHEDULER_COUNT) {  // Only show schedule if not using custom sigmas
+    if (gen_params.sample_params.scheduler != SCHEDULER_COUNT) {
         parameter_string += " " + std::string(sd_scheduler_name(gen_params.sample_params.scheduler));
     }
     parameter_string += ", ";
@@ -282,9 +280,47 @@ std::string get_image_params(const SDCliParams& cli_params, const SDContextParam
     return parameter_string;
 }
 
+/* Enables Printing the log level tag in color using ANSI escape codes */
 void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
     SDCliParams* cli_params = (SDCliParams*)data;
-    log_print(level, log, cli_params->verbose, cli_params->color);
+    int tag_color;
+    const char* level_str;
+    FILE* out_stream = (level == SD_LOG_ERROR) ? stderr : stdout;
+
+    if (!log || (!cli_params->verbose && level <= SD_LOG_DEBUG)) {
+        return;
+    }
+
+    switch (level) {
+        case SD_LOG_DEBUG:
+            tag_color = 37;
+            level_str = "DEBUG";
+            break;
+        case SD_LOG_INFO:
+            tag_color = 34;
+            level_str = "INFO";
+            break;
+        case SD_LOG_WARN:
+            tag_color = 35;
+            level_str = "WARN";
+            break;
+        case SD_LOG_ERROR:
+            tag_color = 31;
+            level_str = "ERROR";
+            break;
+        default: /* Potential future-proofing */
+            tag_color = 33;
+            level_str = "?????";
+            break;
+    }
+
+    if (cli_params->color == true) {
+        fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str);
+    } else {
+        fprintf(out_stream, "[%-5s] ", level_str);
+    }
+    fputs(log, out_stream);
+    fflush(out_stream);
 }
 
 bool load_images_from_dir(const std::string dir,
@@ -294,7 +330,7 @@ bool load_images_from_dir(const std::string dir,
                           int max_image_num   = 0,
                           bool verbose        = false) {
     if (!fs::exists(dir) || !fs::is_directory(dir)) {
-        LOG_ERROR("'%s' is not a valid directory\n", dir.c_str());
+        fprintf(stderr, "'%s' is not a valid directory\n", dir.c_str());
         return false;
     }
 
@@ -316,12 +352,14 @@ bool load_images_from_dir(const std::string dir,
         std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
 
         if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp") {
-            LOG_DEBUG("load image %zu from '%s'", images.size(), path.c_str());
+            if (verbose) {
+                printf("load image %zu from '%s'\n", images.size(), path.c_str());
+            }
             int width             = 0;
             int height            = 0;
             uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height);
             if (image_buffer == nullptr) {
-                LOG_ERROR("load image from '%s' failed", path.c_str());
+                fprintf(stderr, "load image from '%s' failed\n", path.c_str());
                 return false;
             }
 
@@ -351,129 +389,6 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy,
     }
 }
 
-std::string format_frame_idx(std::string pattern, int frame_idx) {
-    std::smatch match;
-    std::string result = pattern;
-    while (std::regex_search(result, match, format_specifier_regex)) {
-        std::string specifier = match.str(1);
-        char buffer[32];
-        snprintf(buffer, sizeof(buffer), specifier.c_str(), frame_idx);
-        result.replace(match.position(1), match.length(1), buffer);
-    }
-
-    // Then replace all '%%' with '%'
-    size_t pos = 0;
-    while ((pos = result.find("%%", pos)) != std::string::npos) {
-        result.replace(pos, 2, "%");
-        pos += 1;
-    }
-    return result;
-}
-
-bool save_results(const SDCliParams& cli_params,
-                  const SDContextParams& ctx_params,
-                  const SDGenerationParams& gen_params,
-                  sd_image_t* results,
-                  int num_results) {
-    if (results == nullptr || num_results <= 0) {
-        return false;
-    }
-
-    namespace fs      = std::filesystem;
-    fs::path out_path = cli_params.output_path;
-
-    if (!out_path.parent_path().empty()) {
-        std::error_code ec;
-        fs::create_directories(out_path.parent_path(), ec);
-        if (ec) {
-            LOG_ERROR("failed to create directory '%s': %s",
-                      out_path.parent_path().string().c_str(), ec.message().c_str());
-            return false;
-        }
-    }
-
-    fs::path base_path = out_path;
-    fs::path ext       = out_path.has_extension() ? out_path.extension() : fs::path{};
-
-    std::string ext_lower = ext.string();
-    std::transform(ext_lower.begin(), ext_lower.end(), ext_lower.begin(), ::tolower);
-    bool is_jpg = (ext_lower == ".jpg" || ext_lower == ".jpeg" || ext_lower == ".jpe");
-    if (!ext.empty()) {
-        if (is_jpg || ext_lower == ".png") {
-            base_path.replace_extension();
-        }
-    }
-
-    int output_begin_idx = cli_params.output_begin_idx;
-    if (output_begin_idx < 0) {
-        output_begin_idx = 0;
-    }
-
-    auto write_image = [&](const fs::path& path, int idx) {
-        const sd_image_t& img = results[idx];
-        if (!img.data)
-            return false;
-
-        std::string params = get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + idx);
-        int ok             = 0;
-        if (is_jpg) {
-            ok = stbi_write_jpg(path.string().c_str(), img.width, img.height, img.channel, img.data, 90, params.c_str());
-        } else {
-            ok = stbi_write_png(path.string().c_str(), img.width, img.height, img.channel, img.data, 0, params.c_str());
-        }
-        LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
-        return ok != 0;
-    };
-
-    int sucessful_reults = 0;
-
-    if (std::regex_search(cli_params.output_path, format_specifier_regex)) {
-        if (!is_jpg && ext_lower != ".png")
-            ext = ".png";
-        fs::path pattern = base_path;
-        pattern += ext;
-
-        for (int i = 0; i < num_results; ++i) {
-            fs::path img_path = format_frame_idx(pattern.string(), output_begin_idx + i);
-            if (write_image(img_path, i)) {
-                sucessful_reults++;
-            }
-        }
-        LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
-        return sucessful_reults != 0;
-    }
-
-    if (cli_params.mode == VID_GEN && num_results > 1) {
-        if (ext_lower != ".avi")
-            ext = ".avi";
-        fs::path video_path = base_path;
-        video_path += ext;
-        if (create_mjpg_avi_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
-            LOG_INFO("save result MJPG AVI video to '%s'", video_path.string().c_str());
-            return true;
-        } else {
-            LOG_ERROR("Failed to save result MPG AVI video to '%s'", video_path.string().c_str());
-            return false;
-        }
-    }
-
-    if (!is_jpg && ext_lower != ".png")
-        ext = ".png";
-
-    for (int i = 0; i < num_results; ++i) {
-        fs::path img_path = base_path;
-        if (num_results > 1) {
-            img_path += "_" + std::to_string(output_begin_idx + i);
-        }
-        img_path += ext;
-        if (write_image(img_path, i)) {
-            sucessful_reults++;
-        }
-    }
-    LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
-    return sucessful_reults != 0;
-}
-
 int main(int argc, const char* argv[]) {
     if (argc > 1 && std::string(argv[1]) == "--version") {
         std::cout << version_string() << "\n";
@@ -485,6 +400,9 @@ int main(int argc, const char* argv[]) {
     SDGenerationParams gen_params;
 
     parse_args(argc, argv, cli_params, ctx_params, gen_params);
+    if (cli_params.verbose || cli_params.version) {
+        std::cout << version_string() << "\n";
+    }
     if (gen_params.video_frames > 4) {
         size_t last_dot_pos   = cli_params.preview_path.find_last_of(".");
         std::string base_path = cli_params.preview_path;
@@ -503,8 +421,6 @@ int main(int argc, const char* argv[]) {
         cli_params.preview_fps /= 4;
 
     sd_set_log_callback(sd_log_cb, (void*)&cli_params);
-    log_verbose = cli_params.verbose;
-    log_color   = cli_params.color;
     sd_set_preview_callback(step_callback,
                             cli_params.preview_method,
                             cli_params.preview_interval,
@@ -512,39 +428,40 @@ int main(int argc, const char* argv[]) {
                             cli_params.preview_noisy,
                             (void*)&cli_params);
 
-    LOG_DEBUG("version: %s", version_string().c_str());
-    LOG_DEBUG("%s", sd_get_system_info());
-    LOG_DEBUG("%s", cli_params.to_string().c_str());
-    LOG_DEBUG("%s", ctx_params.to_string().c_str());
-    LOG_DEBUG("%s", gen_params.to_string().c_str());
+    if (cli_params.verbose) {
+        printf("%s", sd_get_system_info());
+        printf("%s\n", cli_params.to_string().c_str());
+        printf("%s\n", ctx_params.to_string().c_str());
+        printf("%s\n", gen_params.to_string().c_str());
+    }
 
     if (cli_params.mode == CONVERT) {
         bool success = convert(ctx_params.model_path.c_str(),
                                ctx_params.vae_path.c_str(),
                                cli_params.output_path.c_str(),
                                ctx_params.wtype,
-                               ctx_params.tensor_type_rules.c_str(),
-                               cli_params.convert_name);
+                               ctx_params.tensor_type_rules.c_str());
         if (!success) {
-            LOG_ERROR("convert '%s'/'%s' to '%s' failed",
-                      ctx_params.model_path.c_str(),
-                      ctx_params.vae_path.c_str(),
-                      cli_params.output_path.c_str());
+            fprintf(stderr,
+                    "convert '%s'/'%s' to '%s' failed\n",
+                    ctx_params.model_path.c_str(),
+                    ctx_params.vae_path.c_str(),
+                    cli_params.output_path.c_str());
             return 1;
         } else {
-            LOG_INFO("convert '%s'/'%s' to '%s' success",
-                     ctx_params.model_path.c_str(),
-                     ctx_params.vae_path.c_str(),
-                     cli_params.output_path.c_str());
+            printf("convert '%s'/'%s' to '%s' success\n",
+                   ctx_params.model_path.c_str(),
+                   ctx_params.vae_path.c_str(),
+                   cli_params.output_path.c_str());
             return 0;
         }
     }
 
     bool vae_decode_only     = true;
-    sd_image_t init_image    = {0, 0, 3, nullptr};
-    sd_image_t end_image     = {0, 0, 3, nullptr};
-    sd_image_t control_image = {0, 0, 3, nullptr};
-    sd_image_t mask_image    = {0, 0, 1, nullptr};
+    sd_image_t init_image    = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
+    sd_image_t end_image     = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
+    sd_image_t control_image = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
+    sd_image_t mask_image    = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 1, nullptr};
     std::vector<sd_image_t> ref_images;
     std::vector<sd_image_t> pmid_images;
     std::vector<sd_image_t> control_frames;
@@ -571,80 +488,58 @@ int main(int argc, const char* argv[]) {
         control_frames.clear();
     };
 
-    auto load_image_and_update_size = [&](const std::string& path,
-                                          sd_image_t& image,
-                                          bool resize_image    = true,
-                                          int expected_channel = 3) -> bool {
-        int expected_width  = 0;
-        int expected_height = 0;
-        if (resize_image && gen_params.width_and_height_are_set()) {
-            expected_width  = gen_params.width;
-            expected_height = gen_params.height;
-        }
-
-        if (!load_sd_image_from_file(&image, path.c_str(), expected_width, expected_height, expected_channel)) {
-            LOG_ERROR("load image from '%s' failed", path.c_str());
-            release_all_resources();
-            return false;
-        }
-
-        gen_params.set_width_and_height_if_unset(image.width, image.height);
-        return true;
-    };
-
     if (gen_params.init_image_path.size() > 0) {
         vae_decode_only = false;
-        if (!load_image_and_update_size(gen_params.init_image_path, init_image)) {
+
+        int width       = 0;
+        int height      = 0;
+        init_image.data = load_image_from_file(gen_params.init_image_path.c_str(), width, height, gen_params.width, gen_params.height);
+        if (init_image.data == nullptr) {
+            fprintf(stderr, "load image from '%s' failed\n", gen_params.init_image_path.c_str());
+            release_all_resources();
             return 1;
         }
     }
 
     if (gen_params.end_image_path.size() > 0) {
         vae_decode_only = false;
-        if (!load_image_and_update_size(gen_params.end_image_path, end_image)) {
-            return 1;
-        }
-    }
 
-    if (gen_params.ref_image_paths.size() > 0) {
-        vae_decode_only = false;
-        for (auto& path : gen_params.ref_image_paths) {
-            sd_image_t ref_image = {0, 0, 3, nullptr};
-            if (!load_image_and_update_size(path, ref_image, false)) {
-                return 1;
-            }
-            ref_images.push_back(ref_image);
+        int width      = 0;
+        int height     = 0;
+        end_image.data = load_image_from_file(gen_params.end_image_path.c_str(), width, height, gen_params.width, gen_params.height);
+        if (end_image.data == nullptr) {
+            fprintf(stderr, "load image from '%s' failed\n", gen_params.end_image_path.c_str());
+            release_all_resources();
+            return 1;
         }
     }
 
     if (gen_params.mask_image_path.size() > 0) {
-        if (!load_sd_image_from_file(&mask_image,
-                                     gen_params.mask_image_path.c_str(),
-                                     gen_params.get_resolved_width(),
-                                     gen_params.get_resolved_height(),
-                                     1)) {
-            LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
+        int c           = 0;
+        int width       = 0;
+        int height      = 0;
+        mask_image.data = load_image_from_file(gen_params.mask_image_path.c_str(), width, height, gen_params.width, gen_params.height, 1);
+        if (mask_image.data == nullptr) {
+            fprintf(stderr, "load image from '%s' failed\n", gen_params.mask_image_path.c_str());
             release_all_resources();
             return 1;
         }
     } else {
-        mask_image.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        mask_image.data = (uint8_t*)malloc(gen_params.width * gen_params.height);
+        memset(mask_image.data, 255, gen_params.width * gen_params.height);
         if (mask_image.data == nullptr) {
-            LOG_ERROR("malloc mask image failed");
+            fprintf(stderr, "malloc mask image failed\n");
             release_all_resources();
             return 1;
         }
-        mask_image.width  = gen_params.get_resolved_width();
-        mask_image.height = gen_params.get_resolved_height();
-        memset(mask_image.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
     }
 
     if (gen_params.control_image_path.size() > 0) {
-        if (!load_sd_image_from_file(&control_image,
-                                     gen_params.control_image_path.c_str(),
-                                     gen_params.get_resolved_width(),
-                                     gen_params.get_resolved_height())) {
-            LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
+        int width          = 0;
+        int height         = 0;
+        control_image.data = load_image_from_file(gen_params.control_image_path.c_str(), width, height, gen_params.width, gen_params.height);
+        if (control_image.data == nullptr) {
+            fprintf(stderr, "load image from '%s' failed\n", gen_params.control_image_path.c_str());
             release_all_resources();
             return 1;
         }
@@ -658,11 +553,29 @@ int main(int argc, const char* argv[]) {
         }
     }
 
+    if (gen_params.ref_image_paths.size() > 0) {
+        vae_decode_only = false;
+        for (auto& path : gen_params.ref_image_paths) {
+            int width             = 0;
+            int height            = 0;
+            uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height);
+            if (image_buffer == nullptr) {
+                fprintf(stderr, "load image from '%s' failed\n", path.c_str());
+                release_all_resources();
+                return 1;
+            }
+            ref_images.push_back({(uint32_t)width,
+                                  (uint32_t)height,
+                                  3,
+                                  image_buffer});
+        }
+    }
+
     if (!gen_params.control_video_path.empty()) {
         if (!load_images_from_dir(gen_params.control_video_path,
                                   control_frames,
-                                  gen_params.get_resolved_width(),
-                                  gen_params.get_resolved_height(),
+                                  gen_params.width,
+                                  gen_params.height,
                                   gen_params.video_frames,
                                   cli_params.verbose)) {
             release_all_resources();
@@ -695,7 +608,7 @@ int main(int argc, const char* argv[]) {
         num_results = 1;
         results     = (sd_image_t*)calloc(num_results, sizeof(sd_image_t));
         if (results == nullptr) {
-            LOG_INFO("failed to allocate results array");
+            printf("failed to allocate results array\n");
             release_all_resources();
             return 1;
         }
@@ -706,7 +619,7 @@ int main(int argc, const char* argv[]) {
         sd_ctx_t* sd_ctx = new_sd_ctx(&sd_ctx_params);
 
         if (sd_ctx == nullptr) {
-            LOG_INFO("new_sd_ctx_t failed");
+            printf("new_sd_ctx_t failed\n");
             release_all_resources();
             return 1;
         }
@@ -720,7 +633,7 @@ int main(int argc, const char* argv[]) {
         }
 
         if (gen_params.sample_params.scheduler == SCHEDULER_COUNT) {
-            gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx, gen_params.sample_params.sample_method);
+            gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx);
         }
 
         if (cli_params.mode == IMG_GEN) {
@@ -736,8 +649,8 @@ int main(int argc, const char* argv[]) {
                 gen_params.auto_resize_ref_image,
                 gen_params.increase_ref_index,
                 mask_image,
-                gen_params.get_resolved_width(),
-                gen_params.get_resolved_height(),
+                gen_params.width,
+                gen_params.height,
                 gen_params.sample_params,
                 gen_params.strength,
                 gen_params.seed,
@@ -750,8 +663,8 @@ int main(int argc, const char* argv[]) {
                     gen_params.pm_id_embed_path.c_str(),
                     gen_params.pm_style_strength,
                 },  // pm_params
-                gen_params.vae_tiling_params,
-                gen_params.cache_params,
+                ctx_params.vae_tiling_params,
+                gen_params.easycache_params,
             };
 
             results     = generate_image(sd_ctx, &img_gen_params);
@@ -767,8 +680,8 @@ int main(int argc, const char* argv[]) {
                 end_image,
                 control_frames.data(),
                 (int)control_frames.size(),
-                gen_params.get_resolved_width(),
-                gen_params.get_resolved_height(),
+                gen_params.width,
+                gen_params.height,
                 gen_params.sample_params,
                 gen_params.high_noise_sample_params,
                 gen_params.moe_boundary,
@@ -776,15 +689,14 @@ int main(int argc, const char* argv[]) {
                 gen_params.seed,
                 gen_params.video_frames,
                 gen_params.vace_strength,
-                gen_params.vae_tiling_params,
-                gen_params.cache_params,
+                gen_params.easycache_params,
             };
 
             results = generate_video(sd_ctx, &vid_gen_params, &num_results);
         }
 
         if (results == nullptr) {
-            LOG_ERROR("generate failed");
+            printf("generate failed\n");
             free_sd_ctx(sd_ctx);
             return 1;
         }
@@ -801,7 +713,7 @@ int main(int argc, const char* argv[]) {
                                                         gen_params.upscale_tile_size);
 
         if (upscaler_ctx == nullptr) {
-            LOG_ERROR("new_upscaler_ctx failed");
+            printf("new_upscaler_ctx failed\n");
         } else {
             for (int i = 0; i < num_results; i++) {
                 if (results[i].data == nullptr) {
@@ -811,7 +723,7 @@ int main(int argc, const char* argv[]) {
                 for (int u = 0; u < gen_params.upscale_repeats; ++u) {
                     sd_image_t upscaled_image = upscale(upscaler_ctx, current_image, upscale_factor);
                     if (upscaled_image.data == nullptr) {
-                        LOG_ERROR("upscale failed");
+                        printf("upscale failed\n");
                         break;
                     }
                     free(current_image.data);
@@ -822,8 +734,67 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    if (!save_results(cli_params, ctx_params, gen_params, results, num_results)) {
-        return 1;
+    // create directory if not exists
+    {
+        const fs::path out_path = cli_params.output_path;
+        if (const fs::path out_dir = out_path.parent_path(); !out_dir.empty()) {
+            std::error_code ec;
+            fs::create_directories(out_dir, ec);  // OK if already exists
+            if (ec) {
+                fprintf(stderr, "failed to create directory '%s': %s\n",
+                        out_dir.string().c_str(), ec.message().c_str());
+                return 1;
+            }
+        }
+    }
+
+    std::string base_path;
+    std::string file_ext;
+    std::string file_ext_lower;
+    bool is_jpg;
+    size_t last_dot_pos   = cli_params.output_path.find_last_of(".");
+    size_t last_slash_pos = std::min(cli_params.output_path.find_last_of("/"),
+                                     cli_params.output_path.find_last_of("\\"));
+    if (last_dot_pos != std::string::npos && (last_slash_pos == std::string::npos || last_dot_pos > last_slash_pos)) {  // filename has extension
+        base_path = cli_params.output_path.substr(0, last_dot_pos);
+        file_ext = file_ext_lower = cli_params.output_path.substr(last_dot_pos);
+        std::transform(file_ext.begin(), file_ext.end(), file_ext_lower.begin(), ::tolower);
+        is_jpg = (file_ext_lower == ".jpg" || file_ext_lower == ".jpeg" || file_ext_lower == ".jpe");
+    } else {
+        base_path = cli_params.output_path;
+        file_ext = file_ext_lower = "";
+        is_jpg                    = false;
+    }
+
+    if (cli_params.mode == VID_GEN && num_results > 1) {
+        std::string vid_output_path = cli_params.output_path;
+        if (file_ext_lower == ".png") {
+            vid_output_path = base_path + ".avi";
+        }
+        create_mjpg_avi_from_sd_images(vid_output_path.c_str(), results, num_results, gen_params.fps);
+        printf("save result MJPG AVI video to '%s'\n", vid_output_path.c_str());
+    } else {
+        // appending ".png" to absent or unknown extension
+        if (!is_jpg && file_ext_lower != ".png") {
+            base_path += file_ext;
+            file_ext = ".png";
+        }
+        for (int i = 0; i < num_results; i++) {
+            if (results[i].data == nullptr) {
+                continue;
+            }
+            int write_ok;
+            std::string final_image_path = i > 0 ? base_path + "_" + std::to_string(i + 1) + file_ext : base_path + file_ext;
+            if (is_jpg) {
+                write_ok = stbi_write_jpg(final_image_path.c_str(), results[i].width, results[i].height, results[i].channel,
+                                          results[i].data, 90, get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + i).c_str());
+                printf("save result JPEG image to '%s' (%s)\n", final_image_path.c_str(), write_ok == 0 ? "failure" : "success");
+            } else {
+                write_ok = stbi_write_png(final_image_path.c_str(), results[i].width, results[i].height, results[i].channel,
+                                          results[i].data, 0, get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + i).c_str());
+                printf("save result PNG image to '%s' (%s)\n", final_image_path.c_str(), write_ok == 0 ? "failure" : "success");
+            }
+        }
     }
 
     for (int i = 0; i < num_results; i++) {

examples/server/CMakeLists.txt

@@ -1,73 +1,6 @@
 set(TARGET sd-server)
 
-option(SD_SERVER_BUILD_FRONTEND "Build server frontend with pnpm" ON)
-
-set(FRONTEND_DIR "${CMAKE_CURRENT_SOURCE_DIR}/frontend")
-set(GENERATED_HTML_HEADER "${FRONTEND_DIR}/dist/gen_index_html.h")
-
-set(HAVE_FRONTEND_BUILD OFF)
-
-if(SD_SERVER_BUILD_FRONTEND AND EXISTS "${FRONTEND_DIR}")
-    if(WIN32)
-        find_program(PNPM_EXECUTABLE NAMES pnpm.cmd pnpm)
-    else()
-        find_program(PNPM_EXECUTABLE NAMES pnpm)
-    endif()
-
-    if(PNPM_EXECUTABLE)
-        message(STATUS "Frontend dir found: ${FRONTEND_DIR}")
-        message(STATUS "pnpm found: ${PNPM_EXECUTABLE}")
-
-        set(HAVE_FRONTEND_BUILD ON)
-
-        add_custom_target(${TARGET}_frontend_install
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" install
-            WORKING_DIRECTORY "${FRONTEND_DIR}"
-            COMMENT "Installing frontend dependencies"
-            VERBATIM
-        )
-
-        add_custom_target(${TARGET}_frontend_build
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build
-            WORKING_DIRECTORY "${FRONTEND_DIR}"
-            COMMENT "Building frontend"
-            VERBATIM
-        )
-
-        add_custom_target(${TARGET}_frontend_header
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build:header
-            WORKING_DIRECTORY "${FRONTEND_DIR}"
-            COMMENT "Generating gen_index_html.h"
-            VERBATIM
-        )
-
-        add_dependencies(${TARGET}_frontend_build ${TARGET}_frontend_install)
-        add_dependencies(${TARGET}_frontend_header ${TARGET}_frontend_build)
-
-        add_custom_target(${TARGET}_frontend
-            DEPENDS ${TARGET}_frontend_header
-        )
-
-        set_source_files_properties("${GENERATED_HTML_HEADER}" PROPERTIES GENERATED TRUE)
-    else()
-        message(WARNING "pnpm not found, frontend build disabled")
-    endif()
-else()
-    message(STATUS "Frontend disabled or directory not found: ${FRONTEND_DIR}")
-endif()
-
 add_executable(${TARGET} main.cpp)
-
-if(HAVE_FRONTEND_BUILD)
-    add_dependencies(${TARGET} ${TARGET}_frontend)
-    target_sources(${TARGET} PRIVATE "${GENERATED_HTML_HEADER}")
-    target_include_directories(${TARGET} PRIVATE "${FRONTEND_DIR}/dist")
-    target_compile_definitions(${TARGET} PRIVATE HAVE_INDEX_HTML)
-    message(STATUS "HAVE_INDEX_HTML enabled")
-else()
-    message(STATUS "HAVE_INDEX_HTML disabled")
-endif()
-
 install(TARGETS ${TARGET} RUNTIME)
 target_link_libraries(${TARGET} PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
 target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17)

examples/server/README.md

@@ -1,104 +1,14 @@
-# Frontend
-
-## Build with Frontend
-
-The server can optionally build the web frontend and embed it into the binary as `gen_index_html.h`.
-
-### Requirements
-
-Install the following tools:
-
-* **Node.js** ≥ 22.18
-  https://nodejs.org/
-
-* **pnpm** ≥ 10
-  Install via npm:
-
-```bash
-npm install -g pnpm
-```
-
-Verify installation:
-
-```bash
-node -v
-pnpm -v
-```
-
-### Install frontend dependencies
-
-Go to the frontend directory and install dependencies:
-
-```bash
-cd examples/server/frontend
-pnpm install
-```
-
-### Build the server with CMake
-
-Enable the frontend build option when configuring CMake:
-
-```bash
-cmake -B build -DSD_SERVER_BUILD_FRONTEND=ON
-cmake --build build --config Release
-```
-
-If `pnpm` is available, the build system will automatically run:
-
-```
-pnpm run build
-pnpm run build:header
-```
-
-and embed the generated frontend into the server binary.
-
-## Frontend Repository
-
-The web frontend is maintained in a **separate repository**, https://github.com/leejet/stable-ui.
-
-If you want to modify the UI or frontend logic, please submit pull requests to the **frontend repository**.
-
-This repository (`stable-diffusion.cpp`) only vendors the frontend periodically. Changes from the frontend repo are synchronized:
-
-* approximately **every 1–2 weeks**, or
-* when there are **major frontend updates**
-
-Because of this, frontend changes will **not appear here immediately** after being merged upstream.
-
-## Using an external frontend
-
-By default, the server uses the **embedded frontend** generated during the build (`gen_index_html.h`).
-
-You can also serve a custom frontend file instead of the embedded one by using:
-
-```bash
---serve-html-path <path-to-index.html>
-```
-
-For example:
-
-```bash
-sd-server --serve-html-path ./index.html
-```
-
-In this case, the server will load and serve the specified `index.html` file instead of the embedded frontend. This is useful when:
-
-* developing or testing frontend changes
-* using a custom UI
-* avoiding rebuilding the binary after frontend modifications
-
 # Run
 
 ```
 usage: ./bin/sd-server  [options]
 
 Svr Options:
-  -l, --listen-ip <string>      server listen ip (default: 127.0.0.1)        
-  --serve-html-path <string>    path to HTML file to serve at root (optional)
-  --listen-port <int>           server listen port (default: 1234)
-  -v, --verbose                 print extra info
-  --color                       colors the logging tags according to level   
-  -h, --help                    show this help message and exit
+  -l, --listen-ip <string>    server listen ip (default: 127.0.0.1)
+  --listen-port <int>         server listen port (default: 1234)
+  -v, --verbose               print extra info
+  --color                     colors the logging tags according to level
+  -h, --help                  show this help message and exit
 
 Context Options:
   -m, --model <string>                     path to full model
@@ -114,7 +24,6 @@ Context Options:
   --high-noise-diffusion-model <string>    path to the standalone high noise diffusion model
   --vae <string>                           path to standalone vae model
   --taesd <string>                         path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
-  --tae <string>                           alias of --taesd
   --control-net <string>                   path to control net model
   --embd-dir <string>                      embeddings directory
   --lora-model-dir <string>                lora model directory
@@ -125,22 +34,17 @@ Context Options:
                                            CPU physical cores
   --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
   --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
+  --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --vae-tiling                             process vae in tiles to reduce memory usage
   --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
   --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM when needed
-  --mmap                                   whether to memory-map model
   --control-net-cpu                        keep controlnet in cpu (for low vram)
   --clip-on-cpu                            keep clip in cpu (for low vram)
   --vae-on-cpu                             keep vae in cpu (for low vram)
-  --fa                                     use flash attention
-  --diffusion-fa                           use flash attention in the diffusion model only
+  --diffusion-fa                           use flash attention in the diffusion model
   --diffusion-conv-direct                  use ggml_conv2d_direct in the diffusion model
   --vae-conv-direct                        use ggml_conv2d_direct in the vae model
-  --circular                               enable circular padding for convolutions
-  --circularx                              enable circular RoPE wrapping on x-axis (width) only
-  --circulary                              enable circular RoPE wrapping on y-axis (height) only
   --chroma-disable-dit-mask                disable dit mask for chroma
-  --qwen-image-zero-cond-t                 enable zero_cond_t for qwen image
   --chroma-enable-t5-mask                  enable t5 mask for chroma
   --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). If not specified, the default is the
                                            type of the weight file
@@ -190,7 +94,6 @@ Default Generation Options:
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
   --eta <float>                            eta in DDIM, only for DDIM and TCD (default: 0)
-  --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
   --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input (default: 3.5)
@@ -207,21 +110,13 @@ Default Generation Options:
   --disable-auto-resize-ref-image          disable auto resize of ref images
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing,
-                                           tcd, res_multistep, res_2s] (default: euler for Flux/SD3/Wan, euler_a
-                                           otherwise)
+                                           tcd] (default: euler for Flux/SD3/Wan, euler_a otherwise)
   --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm,
-                                           ddim_trailing, tcd, res_multistep, res_2s] default: euler for Flux/SD3/Wan,
-                                           euler_a otherwise
-  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple,
-                                           kl_optimal, lcm, bong_tangent], default: discrete
-  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g., "14.61,7.8,3.5,0.0").
+                                           ddim_trailing, tcd] default: euler for Flux/SD3/Wan, euler_a otherwise
+  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm],
+                                           default: discrete
   --skip-layers                            layers to skip for SLG steps (default: [7,8,9])
   --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
   -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
-  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)
-  --cache-option                           named cache params (key=value format, comma-separated). easycache/ucache:
-                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit: Fn=,Bn=,threshold=,warmup=. Examples:
-                                           "threshold=0.25" or "threshold=1.5,reset=0"
-  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g., "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
-  --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
+  --easycache                              enable EasyCache for DiT models with optional "threshold,start_percent,end_percent" (default: 0.2,0.15,0.95)
 ```

examples/server/frontend

@@ -1 +0,0 @@
-Subproject commit 1a34176cd6d39ad3a226b2b69047e71f6797f6bc

format-code.sh

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

ggml

@@ -1 +1 @@
-Subproject commit 404fcb9d7c96989569e68c9e7881ee3465a05c50
+Subproject commit 2d3876d554551d35c06dccc5852be50d5fd2a275

gguf_reader.hpp

@@ -151,7 +151,7 @@ private:
         }
 
         if (n_dims > GGML_MAX_DIMS) {
-            for (uint32_t i = GGML_MAX_DIMS; i < n_dims; i++) {
+            for (int i = GGML_MAX_DIMS; i < n_dims; i++) {
                 info.shape[GGML_MAX_DIMS - 1] *= info.shape[i];  // stack to last dim;
             }
             info.shape.resize(GGML_MAX_DIMS);

latent-preview.h

@@ -1,8 +1,6 @@
-#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include "ggml.h"
-#include "tensor.hpp"
 
 const float wan_21_latent_rgb_proj[16][3] = {
     {0.015123f, -0.148418f, 0.479828f},
@@ -165,15 +163,15 @@ const float sd_latent_rgb_proj[4][3] = {
     {-0.178022f, -0.200862f, -0.678514f}};
 float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f};
 
-void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
+void preview_latent_video(uint8_t* buffer, struct ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
     size_t buffer_head = 0;
 
-    uint32_t latent_width  = static_cast<uint32_t>(latents->ne[0]);
-    uint32_t latent_height = static_cast<uint32_t>(latents->ne[1]);
-    uint32_t dim           = static_cast<uint32_t>(latents->ne[ggml_n_dims(latents) - 1]);
+    uint32_t latent_width  = latents->ne[0];
+    uint32_t latent_height = latents->ne[1];
+    uint32_t dim           = latents->ne[ggml_n_dims(latents) - 1];
     uint32_t frames        = 1;
     if (ggml_n_dims(latents) == 4) {
-        frames = static_cast<uint32_t>(latents->ne[2]);
+        frames = latents->ne[2];
     }
 
     uint32_t rgb_width  = latent_width * patch_size;
@@ -181,9 +179,9 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
 
     uint32_t unpatched_dim = dim / (patch_size * patch_size);
 
-    for (uint32_t k = 0; k < frames; k++) {
-        for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
-            for (uint32_t rgb_y = 0; rgb_y < rgb_height; rgb_y++) {
+    for (int k = 0; k < frames; k++) {
+        for (int rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
+            for (int rgb_y = 0; rgb_y < rgb_height; rgb_y++) {
                 int latent_x = rgb_x / patch_size;
                 int latent_y = rgb_y / patch_size;
 
@@ -199,7 +197,7 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
 
                 float r = 0, g = 0, b = 0;
                 if (latent_rgb_proj != nullptr) {
-                    for (uint32_t d = 0; d < unpatched_dim; d++) {
+                    for (int d = 0; d < unpatched_dim; d++) {
                         float value = *(float*)((char*)latents->data + latent_id + (d * patch_size * patch_size + channel_offset) * latents->nb[ggml_n_dims(latents) - 1]);
                         r += value * latent_rgb_proj[d][0];
                         g += value * latent_rgb_proj[d][1];
@@ -234,67 +232,3 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
         }
     }
 }
-
-static inline bool preview_latent_tensor_is_video(const sd::Tensor<float>& latents) {
-    return latents.dim() == 5;
-}
-
-void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
-    uint32_t latent_width  = static_cast<uint32_t>(latents.shape()[0]);
-    uint32_t latent_height = static_cast<uint32_t>(latents.shape()[1]);
-    bool is_video          = preview_latent_tensor_is_video(latents);
-    uint32_t frames        = is_video ? static_cast<uint32_t>(latents.shape()[2]) : 1;
-    uint32_t dim           = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
-
-    uint32_t rgb_width     = latent_width * patch_size;
-    uint32_t rgb_height    = latent_height * patch_size;
-    uint32_t unpatched_dim = dim / (patch_size * patch_size);
-
-    for (uint32_t k = 0; k < frames; k++) {
-        for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
-            for (uint32_t rgb_y = 0; rgb_y < rgb_height; rgb_y++) {
-                uint32_t latent_x = rgb_x / patch_size;
-                uint32_t latent_y = rgb_y / patch_size;
-
-                uint32_t channel_offset = 0;
-                if (patch_size > 1) {
-                    channel_offset = ((rgb_y % patch_size) * patch_size + (rgb_x % patch_size));
-                }
-
-                size_t pixel_id   = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x;
-                auto latent_value = [&](uint32_t latent_channel) -> float {
-                    return is_video
-                               ? latents.values()[latent_x + latent_width * (latent_y + latent_height * (k + frames * latent_channel))]
-                               : latents.values()[latent_x + latent_width * (latent_y + latent_height * latent_channel)];
-                };
-
-                float r = 0.f, g = 0.f, b = 0.f;
-                if (latent_rgb_proj != nullptr) {
-                    for (uint32_t d = 0; d < unpatched_dim; d++) {
-                        uint32_t latent_channel = d * patch_size * patch_size + channel_offset;
-                        float value             = latent_value(latent_channel);
-                        r += value * latent_rgb_proj[d][0];
-                        g += value * latent_rgb_proj[d][1];
-                        b += value * latent_rgb_proj[d][2];
-                    }
-                } else {
-                    r = latent_value(0);
-                    g = latent_value(1);
-                    b = latent_value(2);
-                }
-                if (latent_rgb_bias != nullptr) {
-                    r += latent_rgb_bias[0];
-                    g += latent_rgb_bias[1];
-                    b += latent_rgb_bias[2];
-                }
-                r = std::min(1.0f, std::max(0.0f, r * .5f + .5f));
-                g = std::min(1.0f, std::max(0.0f, g * .5f + .5f));
-                b = std::min(1.0f, std::max(0.0f, b * .5f + .5f));
-
-                buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255);
-                buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255);
-                buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255);
-            }
-        }
-    }
-}

llm.hpp

@@ -19,7 +19,6 @@
 #include "json.hpp"
 #include "rope.hpp"
 #include "tokenize_util.h"
-#include "vocab/vocab.h"
 
 namespace LLM {
     constexpr int LLM_GRAPH_SIZE = 10240;
@@ -194,17 +193,16 @@ namespace LLM {
                         bool padding      = false) {
             if (add_bos_token) {
                 tokens.insert(tokens.begin(), BOS_TOKEN_ID);
-                weights.insert(weights.begin(), 1.f);
             }
             if (max_length > 0 && padding) {
-                size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.f / max_length));
+                size_t n = std::ceil(tokens.size() * 1.0 / max_length);
                 if (n == 0) {
                     n = 1;
                 }
                 size_t length = max_length * n;
                 LOG_DEBUG("token length: %llu", length);
                 tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
-                weights.insert(weights.end(), length - weights.size(), 1.f);
+                weights.insert(weights.end(), length - weights.size(), 1.0);
             }
         }
 
@@ -367,7 +365,7 @@ namespace LLM {
             if (merges_utf8_str.size() > 0) {
                 load_from_merges(merges_utf8_str);
             } else {
-                load_from_merges(load_qwen2_merges());
+                load_from_merges(ModelLoader::load_qwen2_merges());
             }
         }
     };
@@ -379,7 +377,7 @@ namespace LLM {
 
             try {
                 vocab = nlohmann::json::parse(vocab_utf8_str);
-            } catch (const nlohmann::json::parse_error&) {
+            } catch (const nlohmann::json::parse_error& e) {
                 GGML_ABORT("invalid vocab json str");
             }
             for (const auto& [key, value] : vocab.items()) {
@@ -388,7 +386,7 @@ namespace LLM {
                 encoder[token]       = i;
                 decoder[i]           = token;
             }
-            encoder_len = static_cast<int>(vocab.size());
+            encoder_len = vocab.size();
             LOG_DEBUG("vocab size: %d", encoder_len);
 
             auto byte_unicode_pairs = bytes_to_unicode();
@@ -468,7 +466,7 @@ namespace LLM {
             if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
                 load_from_merges(merges_utf8_str, vocab_utf8_str);
             } else {
-                load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
+                load_from_merges(ModelLoader::load_mistral_merges(), ModelLoader::load_mistral_vocab_json());
             }
         }
     };
@@ -487,16 +485,16 @@ namespace LLM {
     };
 
     struct LLMVisionParams {
-        int num_layers                      = 32;
+        int64_t num_layers                  = 32;
         int64_t hidden_size                 = 1280;
         int64_t intermediate_size           = 3420;
-        int num_heads                       = 16;
+        int64_t num_heads                   = 16;
         int64_t in_channels                 = 3;
         int64_t out_hidden_size             = 3584;
-        int temporal_patch_size             = 2;
-        int patch_size                      = 14;
-        int spatial_merge_size              = 2;
-        int window_size                     = 112;
+        int64_t temporal_patch_size         = 2;
+        int64_t patch_size                  = 14;
+        int64_t spatial_merge_size          = 2;
+        int64_t window_size                 = 112;
         std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
     };
 
@@ -505,9 +503,9 @@ namespace LLM {
         int64_t num_layers        = 28;
         int64_t hidden_size       = 3584;
         int64_t intermediate_size = 18944;
-        int num_heads             = 28;
-        int num_kv_heads          = 4;
-        int head_dim              = 128;
+        int64_t num_heads         = 28;
+        int64_t num_kv_heads      = 4;
+        int64_t head_dim          = 128;
         bool qkv_bias             = true;
         bool qk_norm              = false;
         int64_t vocab_size        = 152064;
@@ -523,7 +521,7 @@ namespace LLM {
             blocks["down_proj"] = std::shared_ptr<GGMLBlock>(new Linear(intermediate_size, hidden_size, bias));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
             // x: [N, n_token, hidden_size]
             auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
             auto up_proj   = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
@@ -583,7 +581,7 @@ namespace LLM {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
             // x: [N*grid_t*grid_h*grid_w, in_channels, temporal_patch_size*patch_size*patch_size]
             // return: [N*grid_t*grid_h*grid_w, embed_dim]
             x = ggml_reshape_4d(ctx->ggml_ctx,
@@ -632,7 +630,7 @@ namespace LLM {
             blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, dim));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
             auto ln_q  = std::dynamic_pointer_cast<RMSNorm>(blocks["ln_q"]);
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
@@ -640,7 +638,7 @@ namespace LLM {
             x = ln_q->forward(ctx, x);
             x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size);
             x = mlp_0->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_gelu(ctx->ggml_ctx, x);
             x = mlp_2->forward(ctx, x);
             return x;
         }
@@ -649,15 +647,15 @@ namespace LLM {
     struct VisionAttention : public GGMLBlock {
     protected:
         bool llama_cpp_style;
-        int head_dim;
-        int num_heads;
+        int64_t head_dim;
+        int64_t num_heads;
 
     public:
         VisionAttention(bool llama_cpp_style,
                         int64_t hidden_size,
-                        int num_heads)
+                        int64_t num_heads)
             : llama_cpp_style(llama_cpp_style), num_heads(num_heads) {
-            head_dim = static_cast<int>(hidden_size / num_heads);
+            head_dim = hidden_size / num_heads;
             GGML_ASSERT(num_heads * head_dim == hidden_size);
             if (llama_cpp_style) {
                 blocks["q_proj"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size));
@@ -669,10 +667,10 @@ namespace LLM {
             blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* pe,
-                             ggml_tensor* mask = nullptr) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* pe,
+                                    struct ggml_tensor* mask = nullptr) {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
@@ -711,7 +709,7 @@ namespace LLM {
         VisionBlock(bool llama_cpp_style,
                     int64_t hidden_size,
                     int64_t intermediate_size,
-                    int num_heads,
+                    int64_t num_heads,
                     float eps = 1e-6f) {
             blocks["attn"]  = std::shared_ptr<GGMLBlock>(new VisionAttention(llama_cpp_style, hidden_size, num_heads));
             blocks["mlp"]   = std::shared_ptr<GGMLBlock>(new MLP(hidden_size, intermediate_size, true));
@@ -719,10 +717,10 @@ namespace LLM {
             blocks["norm2"] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_size, eps));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* pe,
-                             ggml_tensor* mask = nullptr) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* pe,
+                                    struct ggml_tensor* mask = nullptr) {
             // x: [N, n_token, hidden_size]
             auto attn  = std::dynamic_pointer_cast<VisionAttention>(blocks["attn"]);
             auto mlp   = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
@@ -745,22 +743,22 @@ namespace LLM {
 
     struct VisionModel : public GGMLBlock {
     protected:
-        int num_layers;
-        int spatial_merge_size;
+        int64_t num_layers;
+        int64_t spatial_merge_size;
         std::set<int> fullatt_block_indexes;
 
     public:
         VisionModel(bool llama_cpp_style,
-                    int num_layers,
+                    int64_t num_layers,
                     int64_t in_channels,
                     int64_t hidden_size,
                     int64_t out_hidden_size,
                     int64_t intermediate_size,
-                    int num_heads,
-                    int spatial_merge_size,
-                    int patch_size,
-                    int temporal_patch_size,
-                    int window_size,
+                    int64_t num_heads,
+                    int64_t spatial_merge_size,
+                    int64_t patch_size,
+                    int64_t temporal_patch_size,
+                    int64_t window_size,
                     std::set<int> fullatt_block_indexes = {7, 15, 23, 31},
                     float eps                           = 1e-6f)
             : num_layers(num_layers), fullatt_block_indexes(std::move(fullatt_block_indexes)), spatial_merge_size(spatial_merge_size) {
@@ -779,12 +777,12 @@ namespace LLM {
             blocks["merger"] = std::shared_ptr<GGMLBlock>(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* pixel_values,
-                             ggml_tensor* pe,
-                             ggml_tensor* window_index,
-                             ggml_tensor* window_inverse_index,
-                             ggml_tensor* window_mask) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* pixel_values,
+                                    struct ggml_tensor* pe,
+                                    struct ggml_tensor* window_index,
+                                    struct ggml_tensor* window_inverse_index,
+                                    struct ggml_tensor* window_mask) {
             // pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
             // window_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
             // window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
@@ -819,7 +817,7 @@ namespace LLM {
     struct Attention : public GGMLBlock {
     protected:
         LLMArch arch;
-        int head_dim;
+        int64_t head_dim;
         int64_t num_heads;
         int64_t num_kv_heads;
         bool qk_norm;
@@ -837,10 +835,9 @@ namespace LLM {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* input_pos,
-                             ggml_tensor* attention_mask = nullptr) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* input_pos) {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
@@ -883,7 +880,7 @@ namespace LLM {
             k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, num_kv_heads, n_token, head_dim]
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);      // [N*num_kv_heads, n_token, head_dim]
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
+            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, true, true, false);  // [N, n_token, hidden_size]
 
             x = out_proj->forward(ctx, x);  // [N, n_token, hidden_size]
             return x;
@@ -899,10 +896,9 @@ namespace LLM {
             blocks["post_attention_layernorm"] = std::make_shared<RMSNorm>(params.hidden_size, params.rms_norm_eps);
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* input_pos,
-                             ggml_tensor* attention_mask = nullptr) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* input_pos) {
             // x: [N, n_token, hidden_size]
             auto self_attn                = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
             auto mlp                      = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
@@ -911,7 +907,7 @@ namespace LLM {
 
             auto residual = x;
             x             = input_layernorm->forward(ctx, x);
-            x             = self_attn->forward(ctx, x, input_pos, attention_mask);
+            x             = self_attn->forward(ctx, x, input_pos);
             x             = ggml_add_inplace(ctx->ggml_ctx, x, residual);
 
             residual = x;
@@ -937,12 +933,11 @@ namespace LLM {
             blocks["norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.hidden_size, params.rms_norm_eps));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* input_ids,
-                             ggml_tensor* input_pos,
-                             ggml_tensor* attention_mask,
-                             std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                             std::set<int> out_layers) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* input_ids,
+                                    struct ggml_tensor* input_pos,
+                                    std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                                    std::set<int> out_layers) {
             // input_ids: [N, n_token]
             // return: [N, n_token, hidden_size]
 
@@ -995,7 +990,7 @@ namespace LLM {
             for (int i = 0; i < num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
 
-                x = block->forward(ctx, x, input_pos, attention_mask);
+                x = block->forward(ctx, x, input_pos);
                 if (out_layers.find(i + 1) != out_layers.end()) {
                     intermediate_outputs.push_back(x);
                 }
@@ -1038,25 +1033,24 @@ namespace LLM {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* input_ids,
-                             ggml_tensor* input_pos,
-                             ggml_tensor* attention_mask,
-                             std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                             std::set<int> out_layers) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* input_ids,
+                                    struct ggml_tensor* input_pos,
+                                    std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                                    std::set<int> out_layers) {
             // input_ids: [N, n_token]
             auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
 
-            auto x = model->forward(ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);
+            auto x = model->forward(ctx, input_ids, input_pos, image_embeds, out_layers);
             return x;
         }
 
-        ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
-                                    ggml_tensor* pixel_values,
-                                    ggml_tensor* pe,
-                                    ggml_tensor* window_index,
-                                    ggml_tensor* window_inverse_index,
-                                    ggml_tensor* window_mask) {
+        struct ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
+                                           struct ggml_tensor* pixel_values,
+                                           struct ggml_tensor* pe,
+                                           struct ggml_tensor* window_index,
+                                           struct ggml_tensor* window_inverse_index,
+                                           struct ggml_tensor* window_mask) {
             GGML_ASSERT(enable_vision);
             auto vision_model = std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
             return vision_model->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
@@ -1069,7 +1063,6 @@ namespace LLM {
         LLM model;
 
         std::vector<int> input_pos_vec;
-        std::vector<float> attention_mask_vec;
         std::vector<float> window_mask_vec;
         std::vector<int> window_index_vec;
         std::vector<int> window_inverse_index_vec;
@@ -1157,41 +1150,38 @@ namespace LLM {
             return llm_arch_to_str[static_cast<int>(params.arch)];
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
             model.get_param_tensors(tensors, prefix);
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* input_ids,
-                             ggml_tensor* input_pos,
-                             ggml_tensor* attention_mask,
-                             std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                             std::set<int> out_layers) {
-            auto hidden_states = model.forward(ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);  // [N, n_token, hidden_size]
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* input_ids,
+                                    struct ggml_tensor* input_pos,
+                                    std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                                    std::set<int> out_layers) {
+            auto hidden_states = model.forward(ctx, input_ids, input_pos, image_embeds, out_layers);  // [N, n_token, hidden_size]
             return hidden_states;
         }
 
-        ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
-                                    ggml_tensor* pixel_values,
-                                    ggml_tensor* input_pos,
-                                    ggml_tensor* window_index,
-                                    ggml_tensor* window_inverse_index,
-                                    ggml_tensor* window_mask) {
+        struct ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
+                                           struct ggml_tensor* pixel_values,
+                                           struct ggml_tensor* input_pos,
+                                           struct ggml_tensor* window_index,
+                                           struct ggml_tensor* window_inverse_index,
+                                           struct ggml_tensor* window_mask) {
             auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask);
             return hidden_states;
         }
 
-        ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
-                                 const sd::Tensor<float>& attention_mask_tensor,
-                                 const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
-                                 std::set<int> out_layers) {
-            ggml_cgraph* gf        = ggml_new_graph(compute_ctx);
-            ggml_tensor* input_ids = make_input(input_ids_tensor);
-            std::vector<std::pair<int, ggml_tensor*>> image_embeds;
-            image_embeds.reserve(image_embeds_tensor.size());
-            for (const auto& [idx, embed_tensor] : image_embeds_tensor) {
-                ggml_tensor* embed = make_input(embed_tensor);
-                image_embeds.emplace_back(idx, embed);
+        struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
+                                        std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                                        std::set<int> out_layers) {
+            struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+
+            input_ids = to_backend(input_ids);
+
+            for (auto& image_embed : image_embeds) {
+                image_embed.second = to_backend(image_embed.second);
             }
 
             int64_t n_tokens = input_ids->ne[0];
@@ -1215,54 +1205,37 @@ namespace LLM {
                                                 input_pos_vec.size());
             set_backend_tensor_data(input_pos, input_pos_vec.data());
 
-            ggml_tensor* attention_mask = nullptr;
-            if (!attention_mask_tensor.empty()) {
-                attention_mask = make_input(attention_mask_tensor);
-            } else {
-                attention_mask_vec.resize(n_tokens * n_tokens);
-                for (int i0 = 0; i0 < n_tokens; i0++) {
-                    for (int i1 = 0; i1 < n_tokens; i1++) {
-                        float value = 0.f;
-                        if (i0 > i1) {
-                            value = -INFINITY;
-                        }
-                        attention_mask_vec[i1 * n_tokens + i0] = value;
-                    }
-                }
-                attention_mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, n_tokens, n_tokens);
-                set_backend_tensor_data(attention_mask, attention_mask_vec.data());
-            }
-
             auto runner_ctx = get_context();
 
-            ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);
+            struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, image_embeds, out_layers);
 
             ggml_build_forward_expand(gf, hidden_states);
 
             return gf;
         }
 
-        sd::Tensor<float> compute(const int n_threads,
-                                  const sd::Tensor<int32_t>& input_ids,
-                                  const sd::Tensor<float>& attention_mask,
-                                  const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
-                                  std::set<int> out_layers) {
-            auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(input_ids, attention_mask, image_embeds, out_layers);
+        bool compute(const int n_threads,
+                     struct ggml_tensor* input_ids,
+                     std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                     std::set<int> out_layers,
+                     ggml_tensor** output,
+                     ggml_context* output_ctx = nullptr) {
+            auto get_graph = [&]() -> struct ggml_cgraph* {
+                return build_graph(input_ids, image_embeds, out_layers);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+            return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
         }
 
         int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
-            int64_t grid_t     = 1;
-            int64_t grid_h     = h / params.vision.patch_size;
-            int64_t grid_w     = w / params.vision.patch_size;
-            int64_t llm_grid_h = grid_h / params.vision.spatial_merge_size;
-            int64_t llm_grid_w = grid_w / params.vision.spatial_merge_size;
+            int grid_t     = 1;
+            int grid_h     = h / params.vision.patch_size;
+            int grid_w     = w / params.vision.patch_size;
+            int llm_grid_h = grid_h / params.vision.spatial_merge_size;
+            int llm_grid_w = grid_w / params.vision.spatial_merge_size;
             return grid_t * grid_h * grid_w;
         }
 
-        ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
+        struct ggml_tensor* process_image(struct ggml_context* ctx, struct ggml_tensor* image) {
             // image: [C, H, W]
             // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
             int64_t C  = image->ne[2];
@@ -1289,20 +1262,21 @@ namespace LLM {
             return image;
         }
 
-        ggml_cgraph* build_encode_image_graph(const sd::Tensor<float>& image_tensor) {
-            ggml_cgraph* gf    = new_graph_custom(LLM_GRAPH_SIZE);
-            ggml_tensor* image = make_input(image_tensor);
+        struct ggml_cgraph* build_encode_image_graph(struct ggml_tensor* image) {
+            struct ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
 
             GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
             GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
 
             int grid_t                 = 1;
-            int grid_h                 = static_cast<int>(image->ne[1]) / params.vision.patch_size;
-            int grid_w                 = static_cast<int>(image->ne[0]) / params.vision.patch_size;
+            int grid_h                 = image->ne[1] / params.vision.patch_size;
+            int grid_w                 = image->ne[0] / params.vision.patch_size;
             int llm_grid_h             = grid_h / params.vision.spatial_merge_size;
             int llm_grid_w             = grid_w / params.vision.spatial_merge_size;
             int vit_merger_window_size = params.vision.window_size / params.vision.patch_size / params.vision.spatial_merge_size;
 
+            image = to_backend(image);
+
             auto pixel_values = process_image(compute_ctx, image);
 
             // window index
@@ -1384,14 +1358,14 @@ namespace LLM {
             set_backend_tensor_data(window_mask, window_mask_vec.data());
 
             // pe
-            int head_dim = static_cast<int>(params.vision.hidden_size / params.vision.num_heads);
+            int head_dim = params.vision.hidden_size / params.vision.num_heads;
             pe_vec       = Rope::gen_qwen2vl_pe(grid_h,
                                                 grid_w,
                                                 params.vision.spatial_merge_size,
                                                 window_inverse_index_vec,
-                                                10000,
+                                                10000.f,
                                                 {head_dim / 2, head_dim / 2});
-            int pos_len  = static_cast<int>(pe_vec.size() / head_dim / 2);
+            int pos_len  = pe_vec.size() / head_dim / 2;
             // LOG_DEBUG("pos_len %d", pos_len);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
             // pe->data = pe_vec.data();
@@ -1399,24 +1373,26 @@ namespace LLM {
             // pe->data = nullptr;
             set_backend_tensor_data(pe, pe_vec.data());
 
-            auto runnter_ctx           = get_context();
-            ggml_tensor* hidden_states = vision_forward(&runnter_ctx,
-                                                        pixel_values,
-                                                        pe,
-                                                        window_index,
-                                                        window_inverse_index,
-                                                        window_mask);
+            auto runnter_ctx                  = get_context();
+            struct ggml_tensor* hidden_states = vision_forward(&runnter_ctx,
+                                                               pixel_values,
+                                                               pe,
+                                                               window_index,
+                                                               window_inverse_index,
+                                                               window_mask);
             ggml_build_forward_expand(gf, hidden_states);
 
             return gf;
         }
 
-        sd::Tensor<float> encode_image(const int n_threads,
-                                       const sd::Tensor<float>& image) {
-            auto get_graph = [&]() -> ggml_cgraph* {
+        void encode_image(const int n_threads,
+                          struct ggml_tensor* image,
+                          ggml_tensor** output,
+                          ggml_context* output_ctx = nullptr) {
+            auto get_graph = [&]() -> struct ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
+            GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
         }
     };
 
@@ -1438,7 +1414,7 @@ namespace LLM {
             }
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
             model.get_param_tensors(tensors, prefix);
         }
 
@@ -1490,46 +1466,44 @@ namespace LLM {
         }
 
         void test() {
-            ggml_init_params params;
+            struct ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            ggml_context* ctx = ggml_init(params);
-            GGML_ASSERT(ctx != nullptr);
+            struct ggml_context* work_ctx = ggml_init(params);
+            GGML_ASSERT(work_ctx != nullptr);
             bool test_mistral          = false;
             bool test_qwen3            = true;
             bool test_vit              = false;
             bool test_decoder_with_vit = false;
 
             if (test_decoder_with_vit) {
-                sd::Tensor<float> image_embed;
+                ggml_tensor* image_embed = nullptr;
                 {
-                    auto image = sd::load_tensor_from_file_as_tensor<float>("qwen2vl_normalized.bin");
-                    print_sd_tensor(image, false, "image");
-                    sd::Tensor<float> out;
+                    auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
+                    print_ggml_tensor(image, false, "image");
+                    struct ggml_tensor* out = nullptr;
 
-                    int64_t t0   = ggml_time_ms();
-                    auto out_opt = model.encode_image(8, image);
-                    int64_t t1   = ggml_time_ms();
+                    int t0 = ggml_time_ms();
+                    model.encode_image(8, image, &out, work_ctx);
+                    int t1 = ggml_time_ms();
 
-                    GGML_ASSERT(!out_opt.empty());
-                    out = std::move(out_opt);
-                    print_sd_tensor(out, false, "image_embed");
+                    print_ggml_tensor(out, false, "image_embed");
                     image_embed = out;
-                    LOG_DEBUG("llm encode_image test done in %lldms", t1 - t0);
+                    LOG_DEBUG("llm encode_image test done in %dms", t1 - t0);
                 }
 
                 std::string placeholder  = "<|image_pad|>";
                 std::string img_prompt   = "Picture 1: <|vision_start|>";  // [24669, 220, 16, 25, 220, 151652]
-                int64_t num_image_tokens = image_embed.shape()[1];
+                int64_t num_image_tokens = image_embed->ne[1];
                 img_prompt.reserve(num_image_tokens * placeholder.size());
                 for (int i = 0; i < num_image_tokens; i++) {
                     img_prompt += placeholder;
                 }
                 img_prompt += "<|vision_end|>";
 
-                std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
+                std::vector<std::pair<int, ggml_tensor*>> image_embeds;
                 image_embeds.emplace_back(64, image_embed);
 
                 std::pair<int, int> prompt_attn_range;
@@ -1547,36 +1521,32 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids = sd::Tensor<int32_t>::from_vector(tokens);
-                sd::Tensor<float> out;
+                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                struct ggml_tensor* out = nullptr;
 
-                int64_t t0   = ggml_time_ms();
-                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), image_embeds, {});
-                int64_t t1   = ggml_time_ms();
+                int t0 = ggml_time_ms();
+                model.compute(8, input_ids, image_embeds, {}, &out, work_ctx);
+                int t1 = ggml_time_ms();
 
-                GGML_ASSERT(!out_opt.empty());
-                out = std::move(out_opt);
-                print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                print_ggml_tensor(out);
+                LOG_DEBUG("llm test done in %dms", t1 - t0);
             } else if (test_vit) {
-                // auto image = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 280, 280, 3);
+                // auto image = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 280, 280, 3);
                 // ggml_set_f32(image, 0.f);
-                auto image = sd::load_tensor_from_file_as_tensor<float>("qwen2vl_normalized.bin");
-                print_sd_tensor(image, false, "image");
-                sd::Tensor<float> out;
+                auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
+                print_ggml_tensor(image, false, "image");
+                struct ggml_tensor* out = nullptr;
 
-                int64_t t0   = ggml_time_ms();
-                auto out_opt = model.encode_image(8, image);
-                int64_t t1   = ggml_time_ms();
+                int t0 = ggml_time_ms();
+                model.encode_image(8, image, &out, work_ctx);
+                int t1 = ggml_time_ms();
 
-                GGML_ASSERT(!out_opt.empty());
-                out = std::move(out_opt);
-                print_sd_tensor(out, false, "out");
+                print_ggml_tensor(out, false, "out");
 
-                // auto ref_out = load_tensor_from_file(ctx, "qwen2vl.bin");
+                // auto ref_out = load_tensor_from_file(work_ctx, "qwen2vl.bin");
                 // ggml_ext_tensor_diff(ref_out, out, 0.01f);
 
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                LOG_DEBUG("llm test done in %dms", t1 - t0);
             } else if (test_mistral) {
                 std::pair<int, int> prompt_attn_range;
                 std::string text        = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
@@ -1591,17 +1561,15 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids = sd::Tensor<int32_t>::from_vector(tokens);
-                sd::Tensor<float> out;
+                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                struct ggml_tensor* out = nullptr;
 
-                int64_t t0   = ggml_time_ms();
-                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), {}, {10, 20, 30});
-                int64_t t1   = ggml_time_ms();
+                int t0 = ggml_time_ms();
+                model.compute(8, input_ids, {}, {10, 20, 30}, &out, work_ctx);
+                int t1 = ggml_time_ms();
 
-                GGML_ASSERT(!out_opt.empty());
-                out = std::move(out_opt);
-                print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                print_ggml_tensor(out);
+                LOG_DEBUG("llm test done in %dms", t1 - t0);
             } else if (test_qwen3) {
                 std::pair<int, int> prompt_attn_range;
                 std::string text        = "<|im_start|>user\n";
@@ -1616,17 +1584,15 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids = sd::Tensor<int32_t>::from_vector(tokens);
-                sd::Tensor<float> out;
+                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                struct ggml_tensor* out = nullptr;
 
-                int64_t t0   = ggml_time_ms();
-                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), {}, {35});
-                int64_t t1   = ggml_time_ms();
+                int t0 = ggml_time_ms();
+                model.compute(8, input_ids, {}, {35}, &out, work_ctx);
+                int t1 = ggml_time_ms();
 
-                GGML_ASSERT(!out_opt.empty());
-                out = std::move(out_opt);
-                print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                print_ggml_tensor(out);
+                LOG_DEBUG("llm test done in %dms", t1 - t0);
             } else {
                 std::pair<int, int> prompt_attn_range;
                 std::string text        = "<|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";
@@ -1641,17 +1607,15 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids = sd::Tensor<int32_t>::from_vector(tokens);
-                sd::Tensor<float> out;
+                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                struct ggml_tensor* out = nullptr;
 
-                int64_t t0   = ggml_time_ms();
-                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), {}, {});
-                int64_t t1   = ggml_time_ms();
+                int t0 = ggml_time_ms();
+                model.compute(8, input_ids, {}, {}, &out, work_ctx);
+                int t1 = ggml_time_ms();
 
-                GGML_ASSERT(!out_opt.empty());
-                out = std::move(out_opt);
-                print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                print_ggml_tensor(out);
+                LOG_DEBUG("llm test done in %dms", t1 - t0);
             }
         }
 

lora.hpp

@@ -9,7 +9,7 @@
 struct LoraModel : public GGMLRunner {
     std::string lora_id;
     float multiplier = 1.0f;
-    std::unordered_map<std::string, ggml_tensor*> lora_tensors;
+    std::unordered_map<std::string, struct ggml_tensor*> lora_tensors;
     std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
     std::set<std::string> applied_lora_tensors;
     std::string file_path;
@@ -76,13 +76,13 @@ struct LoraModel : public GGMLRunner {
         }
 
         for (const auto& pair : tensors_to_create) {
-            const auto& name   = pair.first;
-            const auto& ts     = pair.second;
-            ggml_tensor* real  = ggml_new_tensor(params_ctx,
-                                                 ts.type,
-                                                 ts.n_dims,
-                                                 ts.ne);
-            lora_tensors[name] = real;
+            const auto& name         = pair.first;
+            const auto& ts           = pair.second;
+            struct ggml_tensor* real = ggml_new_tensor(params_ctx,
+                                                       ts.type,
+                                                       ts.n_dims,
+                                                       ts.ne);
+            lora_tensors[name]       = real;
         }
 
         alloc_params_buffer();
@@ -195,7 +195,7 @@ struct LoraModel : public GGMLRunner {
             scale_value *= multiplier;
 
             auto curr_updown = ggml_ext_merge_lora(ctx, lora_down, lora_up, lora_mid);
-            curr_updown      = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown      = ggml_scale_inplace(ctx, curr_updown, scale_value);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -235,7 +235,7 @@ struct LoraModel : public GGMLRunner {
             float scale_value = 1.0f;
             scale_value *= multiplier;
 
-            curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown = ggml_scale_inplace(ctx, curr_updown, scale_value);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -337,10 +337,10 @@ struct LoraModel : public GGMLRunner {
             }
             scale_value *= multiplier;
 
-            ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
-            ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
-            auto curr_updown      = ggml_mul_inplace(ctx, updown_1, updown_2);
-            curr_updown           = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            struct ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
+            struct ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
+            auto curr_updown             = ggml_mul_inplace(ctx, updown_1, updown_2);
+            curr_updown                  = ggml_scale_inplace(ctx, curr_updown, scale_value);
             if (updown == nullptr) {
                 updown = curr_updown;
             } else {
@@ -456,7 +456,7 @@ struct LoraModel : public GGMLRunner {
             scale_value *= multiplier;
 
             auto curr_updown = ggml_ext_kronecker(ctx, lokr_w1, lokr_w2);
-            curr_updown      = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown      = ggml_scale_inplace(ctx, curr_updown, scale_value);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -468,10 +468,10 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) {
+    ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora = true) {
         // lora
         ggml_tensor* diff = nullptr;
-        if (with_lora_and_lokr) {
+        if (with_lora) {
             diff = get_lora_weight_diff(model_tensor_name, ctx);
         }
         // diff
@@ -483,7 +483,7 @@ struct LoraModel : public GGMLRunner {
             diff = get_loha_weight_diff(model_tensor_name, ctx);
         }
         // lokr
-        if (diff == nullptr && with_lora_and_lokr) {
+        if (diff == nullptr) {
             diff = get_lokr_weight_diff(model_tensor_name, ctx);
         }
         if (diff != nullptr) {
@@ -514,108 +514,6 @@ struct LoraModel : public GGMLRunner {
             } else {
                 key = model_tensor_name + "." + std::to_string(index);
             }
-            bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
-
-            std::string lokr_w1_name   = "lora." + key + ".lokr_w1";
-            std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
-            // if either of these is found, then we have a lokr lora
-            auto iter   = lora_tensors.find(lokr_w1_name);
-            auto iter_a = lora_tensors.find(lokr_w1_a_name);
-            if (iter != lora_tensors.end() || iter_a != lora_tensors.end()) {
-                std::string lokr_w1_b_name = "lora." + key + ".lokr_w1_b";
-                std::string lokr_w2_name   = "lora." + key + ".lokr_w2";
-                std::string lokr_w2_a_name = "lora." + key + ".lokr_w2_a";
-                std::string lokr_w2_b_name = "lora." + key + ".lokr_w2_b";
-                std::string alpha_name     = "lora." + key + ".alpha";
-
-                ggml_tensor* lokr_w1   = nullptr;
-                ggml_tensor* lokr_w1_a = nullptr;
-                ggml_tensor* lokr_w1_b = nullptr;
-                ggml_tensor* lokr_w2   = nullptr;
-                ggml_tensor* lokr_w2_a = nullptr;
-                ggml_tensor* lokr_w2_b = nullptr;
-
-                if (iter != lora_tensors.end()) {
-                    lokr_w1 = iter->second;
-                }
-                iter = iter_a;
-                if (iter != lora_tensors.end()) {
-                    lokr_w1_a = iter->second;
-                }
-                iter = lora_tensors.find(lokr_w1_b_name);
-                if (iter != lora_tensors.end()) {
-                    lokr_w1_b = iter->second;
-                }
-
-                iter = lora_tensors.find(lokr_w2_name);
-                if (iter != lora_tensors.end()) {
-                    lokr_w2 = iter->second;
-                    if (is_conv2d && lokr_w2->type != GGML_TYPE_F16) {
-                        lokr_w2 = ggml_cast(ctx, lokr_w2, GGML_TYPE_F16);
-                    }
-                }
-                iter = lora_tensors.find(lokr_w2_a_name);
-                if (iter != lora_tensors.end()) {
-                    lokr_w2_a = iter->second;
-                    if (is_conv2d && lokr_w2_a->type != GGML_TYPE_F16) {
-                        lokr_w2_a = ggml_cast(ctx, lokr_w2_a, GGML_TYPE_F16);
-                    }
-                }
-                iter = lora_tensors.find(lokr_w2_b_name);
-                if (iter != lora_tensors.end()) {
-                    lokr_w2_b = iter->second;
-                    if (is_conv2d && lokr_w2_b->type != GGML_TYPE_F16) {
-                        lokr_w2_b = ggml_cast(ctx, lokr_w2_b, GGML_TYPE_F16);
-                    }
-                }
-
-                int rank = 1;
-                if (lokr_w1_b) {
-                    rank = (int)lokr_w1_b->ne[ggml_n_dims(lokr_w1_b) - 1];
-                }
-                if (lokr_w2_b) {
-                    rank = (int)lokr_w2_b->ne[ggml_n_dims(lokr_w2_b) - 1];
-                }
-
-                float scale_value = 1.0f;
-                iter              = lora_tensors.find(alpha_name);
-                if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
-                    scale_value = alpha / rank;
-                    applied_lora_tensors.insert(alpha_name);
-                }
-
-                if (rank == 1) {
-                    scale_value = 1.0f;
-                }
-                scale_value *= multiplier;
-
-                auto curr_out_diff = ggml_ext_lokr_forward(ctx, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
-                if (out_diff == nullptr) {
-                    out_diff = curr_out_diff;
-                } else {
-                    out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
-                }
-
-                if (lokr_w1)
-                    applied_lora_tensors.insert(lokr_w1_name);
-                if (lokr_w1_a)
-                    applied_lora_tensors.insert(lokr_w1_a_name);
-                if (lokr_w1_b)
-                    applied_lora_tensors.insert(lokr_w1_b_name);
-                if (lokr_w2)
-                    applied_lora_tensors.insert(lokr_w2_name);
-                if (lokr_w2_a)
-                    applied_lora_tensors.insert(lokr_w2_name);
-                if (lokr_w2_b)
-                    applied_lora_tensors.insert(lokr_w2_b_name);
-                applied_lora_tensors.insert(alpha_name);
-
-                index++;
-                continue;
-            }
-
-            // not a lokr, normal lora path
 
             std::string lora_down_name = "lora." + key + ".lora_down";
             std::string lora_up_name   = "lora." + key + ".lora_up";
@@ -627,7 +525,9 @@ struct LoraModel : public GGMLRunner {
             ggml_tensor* lora_mid  = nullptr;
             ggml_tensor* lora_down = nullptr;
 
-            iter = lora_tensors.find(lora_up_name);
+            bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+
+            auto iter = lora_tensors.find(lora_up_name);
             if (iter != lora_tensors.end()) {
                 lora_up = iter->second;
                 if (is_conv2d && lora_up->type != GGML_TYPE_F16) {
@@ -699,8 +599,6 @@ struct LoraModel : public GGMLRunner {
                                       forward_params.conv2d.d0,
                                       forward_params.conv2d.d1,
                                       forward_params.conv2d.direct,
-                                      forward_params.conv2d.circular_x,
-                                      forward_params.conv2d.circular_y,
                                       forward_params.conv2d.scale);
                 if (lora_mid) {
                     lx = ggml_ext_conv_2d(ctx,
@@ -714,8 +612,6 @@ struct LoraModel : public GGMLRunner {
                                           1,
                                           1,
                                           forward_params.conv2d.direct,
-                                          forward_params.conv2d.circular_x,
-                                          forward_params.conv2d.circular_y,
                                           forward_params.conv2d.scale);
                 }
                 lx = ggml_ext_conv_2d(ctx,
@@ -729,12 +625,10 @@ struct LoraModel : public GGMLRunner {
                                       1,
                                       1,
                                       forward_params.conv2d.direct,
-                                      forward_params.conv2d.circular_x,
-                                      forward_params.conv2d.circular_y,
                                       forward_params.conv2d.scale);
             }
 
-            auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
+            auto curr_out_diff = ggml_scale_inplace(ctx, lx, scale_value);
 
             if (out_diff == nullptr) {
                 out_diff = curr_out_diff;
@@ -747,9 +641,9 @@ struct LoraModel : public GGMLRunner {
         return out_diff;
     }
 
-    ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
+    struct ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
         size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
-        ggml_cgraph* gf        = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
+        struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
 
         preprocess_lora_tensors(model_tensors);
 
@@ -788,11 +682,11 @@ struct LoraModel : public GGMLRunner {
         return gf;
     }
 
-    void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
-        auto get_graph = [&]() -> ggml_cgraph* {
+    void apply(std::map<std::string, struct ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_lora_graph(model_tensors, version);
         };
-        GGMLRunner::compute<float>(get_graph, n_threads, false, true);
+        GGMLRunner::compute(get_graph, n_threads, false);
         stat();
         for (auto item : original_tensor_to_final_tensor) {
             ggml_tensor* original_tensor = item.first;
@@ -841,9 +735,9 @@ public:
         : lora_models(lora_models) {
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) {
+    ggml_tensor* patch_weight(ggml_context* ctx, ggml_tensor* weight, const std::string& weight_name, bool with_lora) {
         for (auto& lora_model : lora_models) {
-            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, ctx, weight, with_lora_and_lokr);
+            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, ctx, weight, with_lora);
             if (diff == nullptr) {
                 continue;
             }
@@ -885,8 +779,6 @@ public:
                                    forward_params.conv2d.d0,
                                    forward_params.conv2d.d1,
                                    forward_params.conv2d.direct,
-                                   forward_params.conv2d.circular_x,
-                                   forward_params.conv2d.circular_y,
                                    forward_params.conv2d.scale);
         }
         for (auto& lora_model : lora_models) {

ltxv.hpp

@@ -1,7 +1,8 @@
 #ifndef __LTXV_HPP__
 #define __LTXV_HPP__
 
-#include "common_block.hpp"
+#include "common.hpp"
+#include "ggml_extend.hpp"
 
 namespace LTXV {
 
@@ -26,9 +27,9 @@ namespace LTXV {
                                                                      bias));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             bool causal = true) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    bool causal = true) {
             // x: [N*IC, ID, IH, IW]
             // result: [N*OC, OD, OH, OW]
             auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);

mmdit.hpp

@@ -27,13 +27,13 @@ public:
         blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, n_token, in_features]
         auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
         auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
 
         x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_gelu_inplace(ctx->ggml_ctx, x);
         x = fc2->forward(ctx, x);
         return x;
     }
@@ -72,7 +72,7 @@ public:
                                                                bias));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, C, H, W]
         // return: [N, H*W, embed_dim]
         auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
@@ -97,12 +97,12 @@ public:
 struct TimestepEmbedder : public GGMLBlock {
     // Embeds scalar timesteps into vector representations.
 protected:
-    int frequency_embedding_size;
+    int64_t frequency_embedding_size;
 
 public:
     TimestepEmbedder(int64_t hidden_size,
-                     int frequency_embedding_size = 256,
-                     int64_t out_channels         = 0)
+                     int64_t frequency_embedding_size = 256,
+                     int64_t out_channels             = 0)
         : frequency_embedding_size(frequency_embedding_size) {
         if (out_channels <= 0) {
             out_channels = hidden_size;
@@ -111,7 +111,7 @@ public:
         blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_channels, true, true));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* t) {
         // t: [N, ]
         // return: [N, hidden_size]
         auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
@@ -135,7 +135,7 @@ public:
         blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size, true, true));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, input_dim]
         // return: [N, hidden_size]
         auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
@@ -167,15 +167,15 @@ public:
             blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim));
         }
         if (qk_norm == "rms") {
-            blocks["ln_q"] = std::shared_ptr<GGMLBlock>(new RMSNorm(d_head, 1.0e-6f));
-            blocks["ln_k"] = std::shared_ptr<GGMLBlock>(new RMSNorm(d_head, 1.0e-6f));
+            blocks["ln_q"] = std::shared_ptr<GGMLBlock>(new RMSNorm(d_head, 1.0e-6));
+            blocks["ln_k"] = std::shared_ptr<GGMLBlock>(new RMSNorm(d_head, 1.0e-6));
         } else if (qk_norm == "ln") {
-            blocks["ln_q"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_head, 1.0e-6f));
-            blocks["ln_k"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_head, 1.0e-6f));
+            blocks["ln_q"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_head, 1.0e-6));
+            blocks["ln_k"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_head, 1.0e-6));
         }
     }
 
-    std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    std::vector<struct ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
         auto qkv         = qkv_proj->forward(ctx, x);
@@ -198,7 +198,7 @@ public:
         return {q, k, v};
     }
 
-    ggml_tensor* post_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* post_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         GGML_ASSERT(!pre_only);
 
         auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
@@ -208,19 +208,19 @@ public:
     }
 
     // x: [N, n_token, dim]
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x) {
         auto qkv = pre_attention(ctx, x);
-        x        = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
-        x        = post_attention(ctx, x);                                                                                                           // [N, n_token, dim]
+        x        = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+        x        = post_attention(ctx, x);                                                                                                                  // [N, n_token, dim]
         return x;
     }
 };
 
-__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
-                                        ggml_tensor* x,
-                                        ggml_tensor* shift,
-                                        ggml_tensor* scale) {
+__STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
+                                               struct ggml_tensor* x,
+                                               struct ggml_tensor* shift,
+                                               struct ggml_tensor* scale) {
     // x: [N, L, C]
     // scale: [N, C]
     // shift: [N, C]
@@ -274,8 +274,8 @@ public:
     }
 
     std::tuple<std::vector<ggml_tensor*>, std::vector<ggml_tensor*>, std::vector<ggml_tensor*>> pre_attention_x(GGMLRunnerContext* ctx,
-                                                                                                                ggml_tensor* x,
-                                                                                                                ggml_tensor* c) {
+                                                                                                                struct ggml_tensor* x,
+                                                                                                                struct ggml_tensor* c) {
         GGML_ASSERT(self_attn);
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
@@ -284,19 +284,23 @@ public:
         auto attn2              = std::dynamic_pointer_cast<SelfAttention>(blocks["attn2"]);
         auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
-        int n_mods = 9;
-        auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, n_mods * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0);
+        int64_t n_mods = 9;
+        auto m         = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));         // [N, n_mods * hidden_size]
+        m              = ggml_reshape_3d(ctx->ggml_ctx, m, c->ne[0], n_mods, c->ne[1]);         // [N, n_mods, hidden_size]
+        m              = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, m, 0, 2, 1, 3));  // [n_mods, N, hidden_size]
 
-        auto shift_msa  = m_vec[0];  // [N, hidden_size]
-        auto scale_msa  = m_vec[1];  // [N, hidden_size]
-        auto gate_msa   = m_vec[2];  // [N, hidden_size]
-        auto shift_mlp  = m_vec[3];  // [N, hidden_size]
-        auto scale_mlp  = m_vec[4];  // [N, hidden_size]
-        auto gate_mlp   = m_vec[5];  // [N, hidden_size]
-        auto shift_msa2 = m_vec[6];  // [N, hidden_size]
-        auto scale_msa2 = m_vec[7];  // [N, hidden_size]
-        auto gate_msa2  = m_vec[8];  // [N, hidden_size]
+        int64_t offset = m->nb[1] * m->ne[1];
+        auto shift_msa = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 0);  // [N, hidden_size]
+        auto scale_msa = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 1);  // [N, hidden_size]
+        auto gate_msa  = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 2);  // [N, hidden_size]
+
+        auto shift_mlp = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 3);  // [N, hidden_size]
+        auto scale_mlp = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 4);  // [N, hidden_size]
+        auto gate_mlp  = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 5);  // [N, hidden_size]
+
+        auto shift_msa2 = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 6);  // [N, hidden_size]
+        auto scale_msa2 = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 7);  // [N, hidden_size]
+        auto gate_msa2  = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 8);  // [N, hidden_size]
 
         auto x_norm = norm1->forward(ctx, x);
 
@@ -309,29 +313,31 @@ public:
         return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}};
     }
 
-    std::pair<std::vector<ggml_tensor*>, std::vector<ggml_tensor*>> pre_attention(GGMLRunnerContext* ctx,
-                                                                                  ggml_tensor* x,
-                                                                                  ggml_tensor* c) {
+    std::pair<std::vector<struct ggml_tensor*>, std::vector<struct ggml_tensor*>> pre_attention(GGMLRunnerContext* ctx,
+                                                                                                struct ggml_tensor* x,
+                                                                                                struct ggml_tensor* c) {
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
         auto norm1              = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
         auto attn               = std::dynamic_pointer_cast<SelfAttention>(blocks["attn"]);
         auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
-        int n_mods = 6;
+        int64_t n_mods = 6;
         if (pre_only) {
             n_mods = 2;
         }
-        auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, n_mods * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0);
+        auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));         // [N, n_mods * hidden_size]
+        m      = ggml_reshape_3d(ctx->ggml_ctx, m, c->ne[0], n_mods, c->ne[1]);         // [N, n_mods, hidden_size]
+        m      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, m, 0, 2, 1, 3));  // [n_mods, N, hidden_size]
 
-        auto shift_msa = m_vec[0];  // [N, hidden_size]
-        auto scale_msa = m_vec[1];  // [N, hidden_size]
+        int64_t offset = m->nb[1] * m->ne[1];
+        auto shift_msa = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 0);  // [N, hidden_size]
+        auto scale_msa = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 1);  // [N, hidden_size]
         if (!pre_only) {
-            auto gate_msa  = m_vec[2];  // [N, hidden_size]
-            auto shift_mlp = m_vec[3];  // [N, hidden_size]
-            auto scale_mlp = m_vec[4];  // [N, hidden_size]
-            auto gate_mlp  = m_vec[5];  // [N, hidden_size]
+            auto gate_msa  = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 2);  // [N, hidden_size]
+            auto shift_mlp = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 3);  // [N, hidden_size]
+            auto scale_mlp = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 4);  // [N, hidden_size]
+            auto gate_mlp  = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 5);  // [N, hidden_size]
 
             auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
 
@@ -346,15 +352,15 @@ public:
         }
     }
 
-    ggml_tensor* post_attention_x(GGMLRunnerContext* ctx,
-                                  ggml_tensor* attn_out,
-                                  ggml_tensor* attn2_out,
-                                  ggml_tensor* x,
-                                  ggml_tensor* gate_msa,
-                                  ggml_tensor* shift_mlp,
-                                  ggml_tensor* scale_mlp,
-                                  ggml_tensor* gate_mlp,
-                                  ggml_tensor* gate_msa2) {
+    struct ggml_tensor* post_attention_x(GGMLRunnerContext* ctx,
+                                         struct ggml_tensor* attn_out,
+                                         struct ggml_tensor* attn2_out,
+                                         struct ggml_tensor* x,
+                                         struct ggml_tensor* gate_msa,
+                                         struct ggml_tensor* shift_mlp,
+                                         struct ggml_tensor* scale_mlp,
+                                         struct ggml_tensor* gate_mlp,
+                                         struct ggml_tensor* gate_msa2) {
         // attn_out: [N, n_token, hidden_size]
         // x: [N, n_token, hidden_size]
         // gate_msa: [N, hidden_size]
@@ -384,13 +390,13 @@ public:
         return x;
     }
 
-    ggml_tensor* post_attention(GGMLRunnerContext* ctx,
-                                ggml_tensor* attn_out,
-                                ggml_tensor* x,
-                                ggml_tensor* gate_msa,
-                                ggml_tensor* shift_mlp,
-                                ggml_tensor* scale_mlp,
-                                ggml_tensor* gate_mlp) {
+    struct ggml_tensor* post_attention(GGMLRunnerContext* ctx,
+                                       struct ggml_tensor* attn_out,
+                                       struct ggml_tensor* x,
+                                       struct ggml_tensor* gate_msa,
+                                       struct ggml_tensor* shift_mlp,
+                                       struct ggml_tensor* scale_mlp,
+                                       struct ggml_tensor* gate_mlp) {
         // attn_out: [N, n_token, hidden_size]
         // x: [N, n_token, hidden_size]
         // gate_msa: [N, hidden_size]
@@ -416,9 +422,9 @@ public:
         return x;
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* c) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* c) {
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
         // return: [N, n_token, hidden_size]
@@ -433,8 +439,8 @@ public:
             auto qkv2          = std::get<1>(qkv_intermediates);
             auto intermediates = std::get<2>(qkv_intermediates);
 
-            auto attn_out  = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);     // [N, n_token, dim]
-            auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            auto attn_out  = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, false, ctx->flash_attn_enabled);     // [N, n_token, dim]
+            auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
             x              = post_attention_x(ctx,
                                               attn_out,
                                               attn2_out,
@@ -450,7 +456,7 @@ public:
             auto qkv               = qkv_intermediates.first;
             auto intermediates     = qkv_intermediates.second;
 
-            auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
             x             = post_attention(ctx,
                                            attn_out,
                                            intermediates[0],
@@ -463,11 +469,11 @@ public:
     }
 };
 
-__STATIC_INLINE__ std::pair<ggml_tensor*, ggml_tensor*>
+__STATIC_INLINE__ std::pair<struct ggml_tensor*, struct ggml_tensor*>
 block_mixing(GGMLRunnerContext* ctx,
-             ggml_tensor* context,
-             ggml_tensor* x,
-             ggml_tensor* c,
+             struct ggml_tensor* context,
+             struct ggml_tensor* x,
+             struct ggml_tensor* c,
              std::shared_ptr<DismantledBlock> context_block,
              std::shared_ptr<DismantledBlock> x_block) {
     // context: [N, n_context, hidden_size]
@@ -489,29 +495,31 @@ block_mixing(GGMLRunnerContext* ctx,
         x_qkv                    = x_qkv_intermediates.first;
         x_intermediates          = x_qkv_intermediates.second;
     }
-    std::vector<ggml_tensor*> qkv;
+    std::vector<struct ggml_tensor*> qkv;
     for (int i = 0; i < 3; i++) {
         qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
     }
 
-    auto attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_context + n_token, hidden_size]
-
+    auto attn         = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, false, ctx->flash_attn_enabled);  // [N, n_context + n_token, hidden_size]
+    attn              = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, attn, 0, 2, 1, 3));                                                                          // [n_context + n_token, N, hidden_size]
     auto context_attn = ggml_view_3d(ctx->ggml_ctx,
                                      attn,
                                      attn->ne[0],
+                                     attn->ne[1],
                                      context->ne[1],
-                                     attn->ne[2],
                                      attn->nb[1],
                                      attn->nb[2],
-                                     0);  // [N, n_context, hidden_size]
+                                     0);                                                                  // [n_context, N, hidden_size]
+    context_attn      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context_attn, 0, 2, 1, 3));  // [N, n_context, hidden_size]
     auto x_attn       = ggml_view_3d(ctx->ggml_ctx,
                                      attn,
                                      attn->ne[0],
+                                     attn->ne[1],
                                      x->ne[1],
-                                     attn->ne[2],
                                      attn->nb[1],
                                      attn->nb[2],
-                                     context->ne[1] * attn->nb[1]);  // [N, n_token, hidden_size]
+                                     attn->nb[2] * context->ne[1]);                                 // [n_token, N, hidden_size]
+    x_attn            = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x_attn, 0, 2, 1, 3));  // [N, n_token, hidden_size]
 
     if (!context_block->pre_only) {
         context = context_block->post_attention(ctx,
@@ -526,7 +534,7 @@ block_mixing(GGMLRunnerContext* ctx,
     }
 
     if (x_block->self_attn) {
-        auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
+        auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, false, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
 
         x = x_block->post_attention_x(ctx,
                                       x_attn,
@@ -563,10 +571,10 @@ public:
         blocks["x_block"]       = std::shared_ptr<GGMLBlock>(new DismantledBlock(hidden_size, num_heads, mlp_ratio, qk_norm, qkv_bias, false, self_attn_x));
     }
 
-    std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                  ggml_tensor* context,
-                                                  ggml_tensor* x,
-                                                  ggml_tensor* c) {
+    std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                                struct ggml_tensor* context,
+                                                                struct ggml_tensor* x,
+                                                                struct ggml_tensor* c) {
         auto context_block = std::dynamic_pointer_cast<DismantledBlock>(blocks["context_block"]);
         auto x_block       = std::dynamic_pointer_cast<DismantledBlock>(blocks["x_block"]);
 
@@ -586,9 +594,9 @@ public:
         blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* c) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* c) {
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
         // return: [N, n_token, patch_size * patch_size * out_channels]
@@ -596,10 +604,13 @@ public:
         auto linear             = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
         auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
-        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);
-        auto shift = m_vec[0];  // [N, hidden_size]
-        auto scale = m_vec[1];  // [N, hidden_size]
+        auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));         // [N, 2 * hidden_size]
+        m      = ggml_reshape_3d(ctx->ggml_ctx, m, c->ne[0], 2, c->ne[1]);              // [N, 2, hidden_size]
+        m      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, m, 0, 2, 1, 3));  // [2, N, hidden_size]
+
+        int64_t offset = m->nb[1] * m->ne[1];
+        auto shift     = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 0);  // [N, hidden_size]
+        auto scale     = ggml_view_2d(ctx->ggml_ctx, m, m->ne[0], m->ne[1], m->nb[1], offset * 1);  // [N, hidden_size]
 
         x = modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
         x = linear->forward(ctx, x);
@@ -612,7 +623,7 @@ struct MMDiT : public GGMLBlock {
     // Diffusion model with a Transformer backbone.
 protected:
     int64_t input_size               = -1;
-    int patch_size                   = 2;
+    int64_t patch_size               = 2;
     int64_t in_channels              = 16;
     int64_t d_self                   = -1;  // >=0 for MMdiT-X
     int64_t depth                    = 24;
@@ -626,7 +637,7 @@ protected:
     int64_t hidden_size;
     std::string qk_norm;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
         enum ggml_type wtype = GGML_TYPE_F32;
         params["pos_embed"]  = ggml_new_tensor_3d(ctx, wtype, hidden_size, num_patchs, 1);
     }
@@ -705,8 +716,8 @@ public:
         blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(hidden_size, patch_size, out_channels));
     }
 
-    ggml_tensor*
-    cropped_pos_embed(ggml_context* ctx,
+    struct ggml_tensor*
+    cropped_pos_embed(struct ggml_context* ctx,
                       int64_t h,
                       int64_t w) {
         auto pos_embed = params["pos_embed"];
@@ -745,11 +756,33 @@ public:
         return spatial_pos_embed;
     }
 
-    ggml_tensor* forward_core_with_concat(GGMLRunnerContext* ctx,
-                                          ggml_tensor* x,
-                                          ggml_tensor* c_mod,
-                                          ggml_tensor* context,
-                                          std::vector<int> skip_layers = std::vector<int>()) {
+    struct ggml_tensor* unpatchify(struct ggml_context* ctx,
+                                   struct ggml_tensor* x,
+                                   int64_t h,
+                                   int64_t w) {
+        // x: [N, H*W, patch_size * patch_size * C]
+        // return: [N, C, H, W]
+        int64_t n = x->ne[2];
+        int64_t c = out_channels;
+        int64_t p = patch_size;
+        h         = (h + 1) / p;
+        w         = (w + 1) / p;
+
+        GGML_ASSERT(h * w == x->ne[1]);
+
+        x = ggml_reshape_4d(ctx, x, c, p * p, w * h, n);       // [N, H*W, P*P, C]
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3));  // [N, C, H*W, P*P]
+        x = ggml_reshape_4d(ctx, x, p, p, w, h * c * n);       // [N*C*H, W, P, P]
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*H, P, W, P]
+        x = ggml_reshape_4d(ctx, x, p * w, p * h, c, n);       // [N, C, H*P, W*P]
+        return x;
+    }
+
+    struct ggml_tensor* forward_core_with_concat(GGMLRunnerContext* ctx,
+                                                 struct ggml_tensor* x,
+                                                 struct ggml_tensor* c_mod,
+                                                 struct ggml_tensor* context,
+                                                 std::vector<int> skip_layers = std::vector<int>()) {
         // x: [N, H*W, hidden_size]
         // context: [N, n_context, d_context]
         // c: [N, hidden_size]
@@ -774,12 +807,12 @@ public:
         return x;
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* t,
-                         ggml_tensor* y               = nullptr,
-                         ggml_tensor* context         = nullptr,
-                         std::vector<int> skip_layers = std::vector<int>()) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* t,
+                                struct ggml_tensor* y        = nullptr,
+                                struct ggml_tensor* context  = nullptr,
+                                std::vector<int> skip_layers = std::vector<int>()) {
         // Forward pass of DiT.
         // x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
         // t: (N,) tensor of diffusion timesteps
@@ -789,11 +822,11 @@ public:
         auto x_embedder = std::dynamic_pointer_cast<PatchEmbed>(blocks["x_embedder"]);
         auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
 
-        int64_t W = x->ne[0];
-        int64_t H = x->ne[1];
+        int64_t w = x->ne[0];
+        int64_t h = x->ne[1];
 
         auto patch_embed = x_embedder->forward(ctx, x);                      // [N, H*W, hidden_size]
-        auto pos_embed   = cropped_pos_embed(ctx->ggml_ctx, H, W);           // [1, H*W, hidden_size]
+        auto pos_embed   = cropped_pos_embed(ctx->ggml_ctx, h, w);           // [1, H*W, hidden_size]
         x                = ggml_add(ctx->ggml_ctx, patch_embed, pos_embed);  // [N, H*W, hidden_size]
 
         auto c = t_embedder->forward(ctx, t);  // [N, hidden_size]
@@ -812,7 +845,7 @@ public:
 
         x = forward_core_with_concat(ctx, x, c, context, skip_layers);  // (N, H*W, patch_size ** 2 * out_channels)
 
-        x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, /*patch_last*/ false);  // [N, C, H, W]
+        x = unpatchify(ctx->ggml_ctx, x, h, w);  // [N, C, H, W]
 
         return x;
     }
@@ -832,94 +865,90 @@ struct MMDiTRunner : public GGMLRunner {
         return "mmdit";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         mmdit.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 sd::Tensor<float>& y_tensor       = {},
-                             std::vector<int> skip_layers            = std::vector<int>()) {
-        ggml_cgraph* gf = new_graph_custom(MMDIT_GRAPH_SIZE);
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                    struct ggml_tensor* timesteps,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* y,
+                                    std::vector<int> skip_layers = std::vector<int>()) {
+        struct ggml_cgraph* gf = new_graph_custom(MMDIT_GRAPH_SIZE);
 
-        ggml_tensor* x         = make_input(x_tensor);
-        ggml_tensor* timesteps = make_input(timesteps_tensor);
-        ggml_tensor* context   = make_optional_input(context_tensor);
-        ggml_tensor* y         = make_optional_input(y_tensor);
+        x         = to_backend(x);
+        context   = to_backend(context);
+        y         = to_backend(y);
+        timesteps = to_backend(timesteps);
 
-        auto runner_ctx  = get_context();
-        ggml_tensor* out = mmdit.forward(&runner_ctx,
-                                         x,
-                                         timesteps,
-                                         y,
-                                         context,
-                                         skip_layers);
+        auto runner_ctx         = get_context();
+        struct ggml_tensor* out = mmdit.forward(&runner_ctx,
+                                                x,
+                                                timesteps,
+                                                y,
+                                                context,
+                                                skip_layers);
 
         ggml_build_forward_expand(gf, out);
 
         return gf;
     }
 
-    sd::Tensor<float> compute(int n_threads,
-                              const sd::Tensor<float>& x,
-                              const sd::Tensor<float>& timesteps,
-                              const sd::Tensor<float>& context = {},
-                              const sd::Tensor<float>& y       = {},
-                              std::vector<int> skip_layers     = std::vector<int>()) {
+    bool compute(int n_threads,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* timesteps,
+                 struct ggml_tensor* context,
+                 struct ggml_tensor* y,
+                 struct ggml_tensor** output     = nullptr,
+                 struct ggml_context* output_ctx = nullptr,
+                 std::vector<int> skip_layers    = std::vector<int>()) {
         // x: [N, in_channels, h, w]
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 154, 4096]) or [1, max_position, hidden_size]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
-        auto get_graph = [&]() -> ggml_cgraph* {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
     }
 
     void test() {
-        ggml_init_params params;
+        struct ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        ggml_context* ctx = ggml_init(params);
-        GGML_ASSERT(ctx != nullptr);
+        struct ggml_context* work_ctx = ggml_init(params);
+        GGML_ASSERT(work_ctx != nullptr);
 
         {
             // cpu f16: pass
             // cpu f32: pass
             // cuda f16: pass
             // cuda f32: pass
-            sd::Tensor<float> x({128, 128, 16, 1});
+            auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 128, 128, 16, 1);
             std::vector<float> timesteps_vec(1, 999.f);
-            auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
-            x.fill_(0.01f);
+            auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
+            ggml_set_f32(x, 0.01f);
             // print_ggml_tensor(x);
 
-            sd::Tensor<float> context({4096, 154, 1});
-            context.fill_(0.01f);
+            auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 4096, 154, 1);
+            ggml_set_f32(context, 0.01f);
             // print_ggml_tensor(context);
 
-            sd::Tensor<float> y({2048, 1});
-            y.fill_(0.01f);
+            auto y = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 2048, 1);
+            ggml_set_f32(y, 0.01f);
             // print_ggml_tensor(y);
 
-            sd::Tensor<float> out;
+            struct ggml_tensor* out = nullptr;
 
-            int64_t t0   = ggml_time_ms();
-            auto out_opt = compute(8,
-                                   x,
-                                   timesteps,
-                                   context,
-                                   y);
-            int64_t t1   = ggml_time_ms();
+            int t0 = ggml_time_ms();
+            compute(8, x, timesteps, context, y, &out, work_ctx);
+            int t1 = ggml_time_ms();
 
-            GGML_ASSERT(!out_opt.empty());
-            out = std::move(out_opt);
-            print_sd_tensor(out);
-            LOG_DEBUG("mmdit test done in %lldms", t1 - t0);
+            print_ggml_tensor(out);
+            LOG_DEBUG("mmdit test done in %dms", t1 - t0);
         }
     }
 

model.cpp

@@ -16,6 +16,10 @@
 #include "model.h"
 #include "stable-diffusion.h"
 #include "util.h"
+#include "vocab.hpp"
+#include "vocab_mistral.hpp"
+#include "vocab_qwen.hpp"
+#include "vocab_umt5.hpp"
 
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
@@ -162,7 +166,43 @@ uint16_t f8_e4m3_to_f16(uint8_t f8) {
 }
 
 uint16_t f8_e5m2_to_f16(uint8_t fp8) {
-    return static_cast<uint16_t>(fp8) << 8;
+    uint8_t sign     = (fp8 >> 7) & 0x1;
+    uint8_t exponent = (fp8 >> 2) & 0x1F;
+    uint8_t mantissa = fp8 & 0x3;
+
+    uint16_t fp16_sign = sign << 15;
+    uint16_t fp16_exponent;
+    uint16_t fp16_mantissa;
+
+    if (exponent == 0 && mantissa == 0) {  // zero
+        return fp16_sign;
+    }
+
+    if (exponent == 0x1F) {  // NAN and INF
+        fp16_exponent = 0x1F;
+        fp16_mantissa = mantissa ? (mantissa << 8) : 0;
+        return fp16_sign | (fp16_exponent << 10) | fp16_mantissa;
+    }
+
+    if (exponent == 0) {  // subnormal numbers
+        fp16_mantissa = (mantissa << 8);
+        return fp16_sign | fp16_mantissa;
+    }
+
+    // normal numbers
+    int16_t true_exponent = (int16_t)exponent - 15 + 15;
+    if (true_exponent <= 0) {
+        fp16_exponent = 0;
+        fp16_mantissa = (mantissa << 8);
+    } else if (true_exponent >= 0x1F) {
+        fp16_exponent = 0x1F;
+        fp16_mantissa = 0;
+    } else {
+        fp16_exponent = (uint16_t)true_exponent;
+        fp16_mantissa = mantissa << 8;
+    }
+
+    return fp16_sign | (fp16_exponent << 10) | fp16_mantissa;
 }
 
 void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
@@ -251,7 +291,7 @@ void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
 }
 
 bool is_zip_file(const std::string& file_path) {
-    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
     if (zip == nullptr) {
         return false;
     }
@@ -336,11 +376,7 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
         LOG_INFO("load %s using checkpoint format", file_path.c_str());
         return init_from_ckpt_file(file_path, prefix);
     } else {
-        if (file_exists(file_path)) {
-            LOG_WARN("unknown format %s", file_path.c_str());
-        } else {
-            LOG_WARN("file %s not found", file_path.c_str());
-        }
+        LOG_WARN("unknown format %s", file_path.c_str());
         return false;
     }
 }
@@ -400,7 +436,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
                 name,
                 gguf_tensor_info.type,
                 gguf_tensor_info.shape.data(),
-                static_cast<int>(gguf_tensor_info.shape.size()),
+                gguf_tensor_info.shape.size(),
                 file_index,
                 data_offset + gguf_tensor_info.offset);
 
@@ -412,14 +448,14 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
         return true;
     }
 
-    int n_tensors = static_cast<int>(gguf_get_n_tensors(ctx_gguf_));
+    int n_tensors = gguf_get_n_tensors(ctx_gguf_);
 
     size_t total_size  = 0;
     size_t data_offset = gguf_get_data_offset(ctx_gguf_);
     for (int i = 0; i < n_tensors; i++) {
-        std::string name   = gguf_get_tensor_name(ctx_gguf_, i);
-        ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
-        size_t offset      = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
+        std::string name          = gguf_get_tensor_name(ctx_gguf_, i);
+        struct ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
+        size_t offset             = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
 
         // LOG_DEBUG("%s", name.c_str());
 
@@ -776,7 +812,7 @@ struct PickleTensorReader {
         }
     }
 
-    void read_string(const std::string& str, zip_t* zip, std::string dir) {
+    void read_string(const std::string& str, struct zip_t* zip, std::string dir) {
         if (str == "storage") {
             read_global_type = true;
         } else if (str != "state_dict") {
@@ -959,7 +995,7 @@ bool ModelLoader::init_from_ckpt_file(const std::string& file_path, const std::s
     file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
 
-    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
     if (zip == nullptr) {
         LOG_ERROR("failed to open '%s'", file_path.c_str());
         return false;
@@ -998,14 +1034,10 @@ SDVersion ModelLoader::get_sd_version() {
 
     bool is_xl                       = false;
     bool is_flux                     = false;
-    bool is_flux2                    = false;
-    bool has_single_block_47         = false;
     bool is_wan                      = false;
     int64_t patch_embedding_channels = 0;
     bool has_img_emb                 = false;
     bool has_middle_block_1          = false;
-    bool has_output_block_311        = false;
-    bool has_output_block_71         = false;
 
     for (auto& [name, tensor_storage] : tensor_storage_map) {
         if (!(is_xl)) {
@@ -1021,14 +1053,8 @@ SDVersion ModelLoader::get_sd_version() {
             if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
                 return VERSION_QWEN_IMAGE;
             }
-            if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
-                return VERSION_ANIMA;
-            }
             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("single_blocks.47.linear1.weight") != std::string::npos) {
-                has_single_block_47 = true;
+                return VERSION_FLUX2;
             }
             if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) {
                 return VERSION_OVIS_IMAGE;
@@ -1068,14 +1094,6 @@ SDVersion ModelLoader::get_sd_version() {
             tensor_storage.name.find("unet.mid_block.resnets.1.") != std::string::npos) {
             has_middle_block_1 = true;
         }
-        if (tensor_storage.name.find("model.diffusion_model.output_blocks.3.1.transformer_blocks.1") != std::string::npos ||
-            tensor_storage.name.find("unet.up_blocks.1.attentions.0.transformer_blocks.1") != std::string::npos) {
-            has_output_block_311 = true;
-        }
-        if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1") != std::string::npos ||
-            tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) {
-            has_output_block_71 = true;
-        }
         if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
             tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
             tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
@@ -1111,15 +1129,12 @@ SDVersion ModelLoader::get_sd_version() {
             return VERSION_SDXL_PIX2PIX;
         }
         if (!has_middle_block_1) {
-            if (!has_output_block_311) {
-                return VERSION_SDXL_VEGA;
-            }
             return VERSION_SDXL_SSD1B;
         }
         return VERSION_SDXL;
     }
 
-    if (is_flux && !is_flux2) {
+    if (is_flux) {
         if (input_block_weight.ne[0] == 384) {
             return VERSION_FLUX_FILL;
         }
@@ -1132,13 +1147,6 @@ SDVersion ModelLoader::get_sd_version() {
         return VERSION_FLUX;
     }
 
-    if (is_flux2) {
-        if (has_single_block_47) {
-            return VERSION_FLUX2;
-        }
-        return VERSION_FLUX2_KLEIN;
-    }
-
     if (token_embedding_weight.ne[0] == 768) {
         if (is_inpaint) {
             return VERSION_SD1_INPAINT;
@@ -1147,9 +1155,6 @@ SDVersion ModelLoader::get_sd_version() {
             return VERSION_SD1_PIX2PIX;
         }
         if (!has_middle_block_1) {
-            if (!has_output_block_71) {
-                return VERSION_SDXS;
-            }
             return VERSION_SD1_TINY_UNET;
         }
         return VERSION_SD1;
@@ -1305,7 +1310,37 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
     }
 }
 
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
+std::string ModelLoader::load_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(merges_utf8_c_str), sizeof(merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
+std::string ModelLoader::load_qwen2_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(qwen2_merges_utf8_c_str), sizeof(qwen2_merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
+std::string ModelLoader::load_mistral_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(mistral_merges_utf8_c_str), sizeof(mistral_merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
+std::string ModelLoader::load_mistral_vocab_json() {
+    std::string json_str(reinterpret_cast<const char*>(mistral_vocab_json_utf8_c_str), sizeof(mistral_vocab_json_utf8_c_str));
+    return json_str;
+}
+
+std::string ModelLoader::load_t5_tokenizer_json() {
+    std::string json_str(reinterpret_cast<const char*>(t5_tokenizer_json_str), sizeof(t5_tokenizer_json_str));
+    return json_str;
+}
+
+std::string ModelLoader::load_umt5_tokenizer_json() {
+    std::string json_str(reinterpret_cast<const char*>(umt5_tokenizer_json_str), sizeof(umt5_tokenizer_json_str));
+    return json_str;
+}
+
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p) {
     int64_t process_time_ms = 0;
     std::atomic<int64_t> read_time_ms(0);
     std::atomic<int64_t> memcpy_time_ms(0);
@@ -1355,15 +1390,6 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
             }
         }
 
-        std::unique_ptr<MmapWrapper> mmapped;
-        if (enable_mmap && !is_zip) {
-            LOG_DEBUG("using mmap for I/O");
-            mmapped = MmapWrapper::create(file_path);
-            if (!mmapped) {
-                LOG_WARN("failed to memory-map '%s'", file_path.c_str());
-            }
-        }
-
         int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size());
         if (n_threads < 1) {
             n_threads = 1;
@@ -1377,7 +1403,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
         for (int i = 0; i < n_threads; ++i) {
             workers.emplace_back([&, file_path, is_zip]() {
                 std::ifstream file;
-                zip_t* zip = nullptr;
+                struct zip_t* zip = nullptr;
                 if (is_zip) {
                     zip = zip_open(file_path.c_str(), 0, 'r');
                     if (zip == nullptr) {
@@ -1385,7 +1411,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         failed = true;
                         return;
                     }
-                } else if (!mmapped) {
+                } else {
                     file.open(file_path, std::ios::binary);
                     if (!file.is_open()) {
                         LOG_ERROR("failed to open '%s'", file_path.c_str());
@@ -1438,11 +1464,6 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                                 zip_entry_noallocread(zip, (void*)buf, n);
                             }
                             zip_entry_close(zip);
-                        } else if (mmapped) {
-                            if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
-                                LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                                failed = true;
-                            }
                         } else {
                             file.seekg(tensor_storage.offset);
                             file.read(buf, n);
@@ -1499,11 +1520,6 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         i64_to_i32_vec((int64_t*)read_buf, (int32_t*)target_buf, tensor_storage.nelements());
                     }
                     if (tensor_storage.type != dst_tensor->type) {
-                        if (convert_buf == nullptr) {
-                            LOG_ERROR("read tensor data failed: too less memory for conversion");
-                            failed = true;
-                            return;
-                        }
                         convert_tensor((void*)target_buf,
                                        tensor_storage.type,
                                        convert_buf,
@@ -1535,7 +1551,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                 break;
             }
             size_t curr_num = total_tensors_processed + current_idx;
-            pretty_progress(static_cast<int>(curr_num), static_cast<int>(total_tensors_to_process), (ggml_time_ms() - t_start) / 1000.0f / (curr_num + 1e-6f));
+            pretty_progress(curr_num, total_tensors_to_process, (ggml_time_ms() - t_start) / 1000.0f / (curr_num + 1e-6f));
             std::this_thread::sleep_for(std::chrono::milliseconds(200));
         }
 
@@ -1548,7 +1564,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
             break;
         }
         total_tensors_processed += file_tensors.size();
-        pretty_progress(static_cast<int>(total_tensors_processed), static_cast<int>(total_tensors_to_process), (ggml_time_ms() - t_start) / 1000.0f / (total_tensors_processed + 1e-6f));
+        pretty_progress(total_tensors_processed, total_tensors_to_process, (ggml_time_ms() - t_start) / 1000.0f / (total_tensors_processed + 1e-6f));
         if (total_tensors_processed < total_tensors_to_process) {
             printf("\n");
         }
@@ -1565,10 +1581,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     return success;
 }
 
-bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
+bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
                                std::set<std::string> ignore_tensors,
-                               int n_threads,
-                               bool enable_mmap) {
+                               int n_threads) {
     std::set<std::string> tensor_names_in_file;
     std::mutex tensor_names_mutex;
     auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
@@ -1579,7 +1594,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
             tensor_names_in_file.insert(name);
         }
 
-        ggml_tensor* real;
+        struct ggml_tensor* real;
         if (tensors.find(name) != tensors.end()) {
             real = tensors[name];
         } else {
@@ -1611,7 +1626,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
         return true;
     };
 
-    bool success = load_tensors(on_new_tensor_cb, n_threads, enable_mmap);
+    bool success = load_tensors(on_new_tensor_cb, n_threads);
     if (!success) {
         LOG_ERROR("load tensors from file failed");
         return false;
@@ -1717,13 +1732,6 @@ bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type
         // tensor_storage.ne[0], tensor_storage.ne[1], tensor_storage.ne[2], tensor_storage.ne[3],
         // tensor->n_dims, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
 
-        if (!tensor->data) {
-            GGML_ASSERT(ggml_nelements(tensor) == 0);
-            // avoid crashing the gguf writer by setting a dummy pointer for zero-sized tensors
-            LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
-            tensor->data = ggml_get_mem_buffer(ggml_ctx);
-        }
-
         *dst_tensor = tensor;
 
         gguf_add_tensor(gguf_ctx, tensor);
@@ -1763,12 +1771,7 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type)
     return mem_size;
 }
 
-bool convert(const char* input_path,
-             const char* vae_path,
-             const char* output_path,
-             sd_type_t output_type,
-             const char* tensor_type_rules,
-             bool convert_name) {
+bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type, const char* tensor_type_rules) {
     ModelLoader model_loader;
 
     if (!model_loader.init_from_file(input_path)) {
@@ -1782,9 +1785,7 @@ bool convert(const char* input_path,
             return false;
         }
     }
-    if (convert_name) {
-        model_loader.convert_tensors_name();
-    }
+    model_loader.convert_tensors_name();
     bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type, tensor_type_rules);
     return success;
 }

model.h

@@ -28,11 +28,9 @@ enum SDVersion {
     VERSION_SD2,
     VERSION_SD2_INPAINT,
     VERSION_SD2_TINY_UNET,
-    VERSION_SDXS,
     VERSION_SDXL,
     VERSION_SDXL_INPAINT,
     VERSION_SDXL_PIX2PIX,
-    VERSION_SDXL_VEGA,
     VERSION_SDXL_SSD1B,
     VERSION_SVD,
     VERSION_SD3,
@@ -45,16 +43,14 @@ enum SDVersion {
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
     VERSION_QWEN_IMAGE,
-    VERSION_ANIMA,
     VERSION_FLUX2,
-    VERSION_FLUX2_KLEIN,
     VERSION_Z_IMAGE,
     VERSION_OVIS_IMAGE,
     VERSION_COUNT,
 };
 
 static inline bool sd_version_is_sd1(SDVersion version) {
-    if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS) {
+    if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET) {
         return true;
     }
     return false;
@@ -68,7 +64,7 @@ static inline bool sd_version_is_sd2(SDVersion version) {
 }
 
 static inline bool sd_version_is_sdxl(SDVersion version) {
-    if (version == VERSION_SDXL || version == VERSION_SDXL_INPAINT || version == VERSION_SDXL_PIX2PIX || version == VERSION_SDXL_SSD1B || version == VERSION_SDXL_VEGA) {
+    if (version == VERSION_SDXL || version == VERSION_SDXL_INPAINT || version == VERSION_SDXL_PIX2PIX || version == VERSION_SDXL_SSD1B) {
         return true;
     }
     return false;
@@ -103,7 +99,7 @@ static inline bool sd_version_is_flux(SDVersion version) {
 }
 
 static inline bool sd_version_is_flux2(SDVersion version) {
-    if (version == VERSION_FLUX2 || version == VERSION_FLUX2_KLEIN) {
+    if (version == VERSION_FLUX2) {
         return true;
     }
     return false;
@@ -123,13 +119,6 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
     return false;
 }
 
-static inline bool sd_version_is_anima(SDVersion version) {
-    if (version == VERSION_ANIMA) {
-        return true;
-    }
-    return false;
-}
-
 static inline bool sd_version_is_z_image(SDVersion version) {
     if (version == VERSION_Z_IMAGE) {
         return true;
@@ -154,7 +143,6 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_sd3(version) ||
         sd_version_is_wan(version) ||
         sd_version_is_qwen_image(version) ||
-        sd_version_is_anima(version) ||
         sd_version_is_z_image(version)) {
         return true;
     }
@@ -322,11 +310,10 @@ public:
     std::map<ggml_type, uint32_t> get_vae_wtype_stat();
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
-    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
-    bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
+    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0);
+    bool load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
                       std::set<std::string> ignore_tensors = {},
-                      int n_threads                        = 0,
-                      bool use_mmap                        = false);
+                      int n_threads                        = 0);
 
     std::vector<std::string> get_tensor_names() const {
         std::vector<std::string> names;
@@ -340,6 +327,13 @@ public:
     bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
     int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
     ~ModelLoader() = default;
+
+    static std::string load_merges();
+    static std::string load_qwen2_merges();
+    static std::string load_mistral_merges();
+    static std::string load_mistral_vocab_json();
+    static std::string load_t5_tokenizer_json();
+    static std::string load_umt5_tokenizer_json();
 };
 
 #endif  // __MODEL_H__

name_conversion.cpp

@@ -653,14 +653,6 @@ std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
     return name;
 }
 
-std::string convert_other_dit_to_original_anima(std::string name) {
-    static const std::string anima_net_prefix = "net.";
-    if (!starts_with(name, anima_net_prefix)) {
-        name = anima_net_prefix + name;
-    }
-    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);
@@ -672,8 +664,6 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         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);
     }
     return name;
 }
@@ -845,14 +835,12 @@ std::string convert_sep_to_dot(std::string name) {
         "proj_out",
         "transformer_blocks",
         "single_transformer_blocks",
-        "single_blocks",
         "diffusion_model",
         "cond_stage_model",
         "first_stage_model",
         "conv_in",
         "conv_out",
         "lora_down",
-        "lora_mid",
         "lora_up",
         "diff_b",
         "hada_w1_a",
@@ -888,18 +876,7 @@ std::string convert_sep_to_dot(std::string name) {
         "ff_context",
         "norm_added_q",
         "norm_added_v",
-        "to_add_out",
-        "txt_mod",
-        "img_mod",
-        "txt_mlp",
-        "img_mlp",
-        "proj_mlp",
-        "wi_0",
-        "wi_1",
-        "norm1_context",
-        "ff_context",
-        "x_embedder",
-    };
+        "to_add_out"};
 
     // record the positions of underscores that should NOT be replaced
     std::unordered_set<size_t> protected_positions;
@@ -971,7 +948,6 @@ bool is_first_stage_model_name(const std::string& name) {
 std::string convert_tensor_name(std::string name, SDVersion version) {
     bool is_lora                             = false;
     bool is_lycoris_underline                = false;
-    bool is_underline                        = false;
     std::vector<std::string> lora_prefix_vec = {
         "lora.lora.",
         "lora.lora_",
@@ -979,27 +955,12 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
         "lora.lycoris.",
         "lora.",
     };
-    std::vector<std::string> underline_lora_prefix_vec = {
-        "unet_",
-        "te_",
-        "te1_",
-        "te2_",
-        "te3_",
-        "vae_",
-    };
     for (const auto& prefix : lora_prefix_vec) {
         if (starts_with(name, prefix)) {
             is_lora = true;
             name    = name.substr(prefix.size());
             if (contains(prefix, "lycoris_")) {
                 is_lycoris_underline = true;
-            } else {
-                for (const auto& underline_lora_prefix : underline_lora_prefix_vec) {
-                    if (starts_with(name, underline_lora_prefix)) {
-                        is_underline = true;
-                        break;
-                    }
-                }
             }
             break;
         }
@@ -1008,13 +969,10 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
     if (is_lora) {
         std::map<std::string, std::string> lora_suffix_map = {
             {".lora_down.weight", ".weight.lora_down"},
-            {".lora_mid.weight", ".weight.lora_mid"},
             {".lora_up.weight", ".weight.lora_up"},
             {".lora.down.weight", ".weight.lora_down"},
-            {".lora.mid.weight", ".weight.lora_mid"},
             {".lora.up.weight", ".weight.lora_up"},
             {"_lora.down.weight", ".weight.lora_down"},
-            {"_lora.mid.weight", ".weight.lora_mid"},
             {"_lora.up.weight", ".weight.lora_up"},
             {".lora_A.weight", ".weight.lora_down"},
             {".lora_B.weight", ".weight.lora_up"},
@@ -1062,14 +1020,12 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
             }
         }
 
-        // LOG_DEBUG("name %s %d", name.c_str(), version);
-
-        if (sd_version_is_unet(version) || is_underline || is_lycoris_underline) {
+        if (sd_version_is_unet(version) || is_lycoris_underline) {
             name = convert_sep_to_dot(name);
         }
     }
 
-    std::unordered_map<std::string, std::string> prefix_map = {
+    std::vector<std::pair<std::string, std::string>> prefix_map = {
         {"diffusion_model.", "model.diffusion_model."},
         {"unet.", "model.diffusion_model."},
         {"transformer.", "model.diffusion_model."},  // dit
@@ -1084,13 +1040,8 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
         // {"te2.text_model.encoder.layers.", "cond_stage_model.1.model.transformer.resblocks."},
         {"te2.", "cond_stage_model.1.transformer."},
         {"te1.", "cond_stage_model.transformer."},
-        {"te3.", "text_encoders.t5xxl.transformer."},
     };
 
-    if (sd_version_is_flux(version)) {
-        prefix_map["te1."] = "text_encoders.clip_l.transformer.";
-    }
-
     replace_with_prefix_map(name, prefix_map);
 
     // diffusion model
@@ -1120,11 +1071,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);
-                if (version == VERSION_SDXS) {
-                    name = "tae." + name;
-                } else {
-                    name = prefix + name;
-                }
+                name = prefix + name;
                 break;
             }
         }

pmid.hpp

@@ -21,19 +21,19 @@ public:
         blocks["layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(in_dim));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
 
         auto fc1        = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
         auto fc2        = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
         auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
 
-        ggml_tensor* r = x;
+        struct ggml_tensor* r = x;
         // x = ggml_ext_layer_norm(ctx, x, ln_w, ln_b);
         x = layer_norm->forward(ctx, x);
         // x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x),  fc1_b);
         x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_gelu_inplace(ctx->ggml_ctx, x);
         x = fc2->forward(ctx, x);
         // x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x),  fc2_b);
         if (use_residue)
@@ -54,8 +54,8 @@ public:
         blocks["1"]   = std::shared_ptr<GGMLBlock>(new Mlp(dim, inner_dim, dim, false));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x) {
         auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["0"]);
         auto ff   = std::dynamic_pointer_cast<Mlp>(blocks["1"]);
 
@@ -72,7 +72,7 @@ struct PerceiverAttention : public GGMLBlock {
     int heads;     // = heads
 public:
     PerceiverAttention(int dim, int dim_h = 64, int h = 8)
-        : scale(powf(static_cast<float>(dim_h), -0.5f)), dim_head(dim_h), heads(h) {
+        : scale(powf(dim_h, -0.5)), dim_head(dim_h), heads(h) {
         int inner_dim    = dim_head * heads;
         blocks["norm1"]  = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
         blocks["norm2"]  = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
@@ -81,9 +81,9 @@ public:
         blocks["to_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim, false));
     }
 
-    ggml_tensor* reshape_tensor(ggml_context* ctx,
-                                ggml_tensor* x,
-                                int heads) {
+    struct ggml_tensor* reshape_tensor(struct ggml_context* ctx,
+                                       struct ggml_tensor* x,
+                                       int heads) {
         int64_t ne[4];
         for (int i = 0; i < 4; ++i)
             ne[i] = x->ne[i];
@@ -92,17 +92,17 @@ public:
         return x;
     }
 
-    std::vector<ggml_tensor*> chunk_half(ggml_context* ctx,
-                                         ggml_tensor* x) {
+    std::vector<struct ggml_tensor*> chunk_half(struct ggml_context* ctx,
+                                                struct ggml_tensor* x) {
         auto tlo = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0);
         auto tli = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], x->nb[0] * x->ne[0] / 2);
         return {ggml_cont(ctx, tlo),
                 ggml_cont(ctx, tli)};
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* latents) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* latents) {
         // x (torch.Tensor): image features
         //     shape (b, n1, D)
         // latent (torch.Tensor): latent features
@@ -129,8 +129,8 @@ public:
         k             = reshape_tensor(ctx->ggml_ctx, k, heads);
         v             = reshape_tensor(ctx->ggml_ctx, v, heads);
         scale         = 1.f / sqrt(sqrt((float)dim_head));
-        k             = ggml_ext_scale(ctx->ggml_ctx, k, scale, true);
-        q             = ggml_ext_scale(ctx->ggml_ctx, q, scale, true);
+        k             = ggml_scale_inplace(ctx->ggml_ctx, k, scale);
+        q             = ggml_scale_inplace(ctx->ggml_ctx, q, scale);
         // auto weight = ggml_mul_mat(ctx, q, k);
         auto weight = ggml_mul_mat(ctx->ggml_ctx, k, q);  // NOTE order of mul is opposite to pytorch
 
@@ -176,9 +176,9 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* latents,
-                         ggml_tensor* x) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* latents,
+                                struct ggml_tensor* x) {
         // x: [N, channels, h, w]
         auto proj_in  = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
         auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
@@ -225,19 +225,19 @@ public:
             4));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* last_hidden_state) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* last_hidden_state) {
         // x: [N, channels, h, w]
         auto token_proj          = std::dynamic_pointer_cast<Mlp>(blocks["token_proj"]);
         auto token_norm          = std::dynamic_pointer_cast<LayerNorm>(blocks["token_norm"]);
         auto perceiver_resampler = std::dynamic_pointer_cast<FacePerceiverResampler>(blocks["perceiver_resampler"]);
 
-        x                = token_proj->forward(ctx, x);
-        int64_t nel      = ggml_nelements(x);
-        x                = ggml_reshape_3d(ctx->ggml_ctx, x, cross_attention_dim, num_tokens, nel / (cross_attention_dim * num_tokens));
-        x                = token_norm->forward(ctx, x);
-        ggml_tensor* out = perceiver_resampler->forward(ctx, x, last_hidden_state);
+        x                       = token_proj->forward(ctx, x);
+        int64_t nel             = ggml_nelements(x);
+        x                       = ggml_reshape_3d(ctx->ggml_ctx, x, cross_attention_dim, num_tokens, nel / (cross_attention_dim * num_tokens));
+        x                       = token_norm->forward(ctx, x);
+        struct ggml_tensor* out = perceiver_resampler->forward(ctx, x, last_hidden_state);
         if (use_residul)
             out = ggml_add(ctx->ggml_ctx, x, out);
         return out;
@@ -256,9 +256,9 @@ public:
         blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(embed_dim));
     }
 
-    ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
-                         ggml_tensor* prompt_embeds,
-                         ggml_tensor* id_embeds) {
+    struct ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* id_embeds) {
         auto mlp1       = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp1"]);
         auto mlp2       = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp2"]);
         auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
@@ -273,24 +273,24 @@ public:
         return stacked_id_embeds;
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* prompt_embeds,
-                         ggml_tensor* id_embeds,
-                         ggml_tensor* class_tokens_mask,
-                         ggml_tensor* class_tokens_mask_pos,
-                         ggml_tensor* left,
-                         ggml_tensor* right) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* id_embeds,
+                                struct ggml_tensor* class_tokens_mask,
+                                struct ggml_tensor* class_tokens_mask_pos,
+                                struct ggml_tensor* left,
+                                struct ggml_tensor* right) {
         // x: [N, channels, h, w]
 
-        ggml_tensor* valid_id_embeds = id_embeds;
+        struct ggml_tensor* valid_id_embeds = id_embeds;
         // # slice out the image token embeddings
         ggml_set_name(class_tokens_mask_pos, "class_tokens_mask_pos");
         ggml_set_name(prompt_embeds, "prompt_embeds");
-        ggml_tensor* image_token_embeds = ggml_get_rows(ctx->ggml_ctx, prompt_embeds, class_tokens_mask_pos);
+        struct ggml_tensor* image_token_embeds = ggml_get_rows(ctx->ggml_ctx, prompt_embeds, class_tokens_mask_pos);
         ggml_set_name(image_token_embeds, "image_token_embeds");
-        valid_id_embeds                = ggml_reshape_2d(ctx->ggml_ctx, valid_id_embeds, valid_id_embeds->ne[0],
-                                                         ggml_nelements(valid_id_embeds) / valid_id_embeds->ne[0]);
-        ggml_tensor* stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds);
+        valid_id_embeds                       = ggml_reshape_2d(ctx->ggml_ctx, valid_id_embeds, valid_id_embeds->ne[0],
+                                                                ggml_nelements(valid_id_embeds) / valid_id_embeds->ne[0]);
+        struct ggml_tensor* stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds);
 
         if (left && right) {
             stacked_id_embeds = ggml_concat(ctx->ggml_ctx, left, stacked_id_embeds, 1);
@@ -301,10 +301,10 @@ public:
             stacked_id_embeds = ggml_concat(ctx->ggml_ctx, stacked_id_embeds, right, 1);
         }
 
-        class_tokens_mask                  = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, class_tokens_mask));
-        class_tokens_mask                  = ggml_repeat(ctx->ggml_ctx, class_tokens_mask, prompt_embeds);
-        prompt_embeds                      = ggml_mul(ctx->ggml_ctx, prompt_embeds, class_tokens_mask);
-        ggml_tensor* updated_prompt_embeds = ggml_add(ctx->ggml_ctx, prompt_embeds, stacked_id_embeds);
+        class_tokens_mask                         = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, class_tokens_mask));
+        class_tokens_mask                         = ggml_repeat(ctx->ggml_ctx, class_tokens_mask, prompt_embeds);
+        prompt_embeds                             = ggml_mul(ctx->ggml_ctx, prompt_embeds, class_tokens_mask);
+        struct ggml_tensor* updated_prompt_embeds = ggml_add(ctx->ggml_ctx, prompt_embeds, stacked_id_embeds);
         ggml_set_name(updated_prompt_embeds, "updated_prompt_embeds");
         return updated_prompt_embeds;
     }
@@ -317,22 +317,22 @@ struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
         blocks["fuse_module"]         = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* id_pixel_values,
-                         ggml_tensor* prompt_embeds,
-                         ggml_tensor* class_tokens_mask,
-                         ggml_tensor* class_tokens_mask_pos,
-                         ggml_tensor* left,
-                         ggml_tensor* right) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* id_pixel_values,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* class_tokens_mask,
+                                struct ggml_tensor* class_tokens_mask_pos,
+                                struct ggml_tensor* left,
+                                struct ggml_tensor* right) {
         // x: [N, channels, h, w]
         auto vision_model        = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
         auto visual_projection   = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);
         auto visual_projection_2 = std::dynamic_pointer_cast<Linear>(blocks["visual_projection_2"]);
         auto fuse_module         = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
 
-        ggml_tensor* shared_id_embeds = vision_model->forward(ctx, id_pixel_values);          // [N, hidden_size]
-        ggml_tensor* id_embeds        = visual_projection->forward(ctx, shared_id_embeds);    // [N, proj_dim(768)]
-        ggml_tensor* id_embeds_2      = visual_projection_2->forward(ctx, shared_id_embeds);  // [N, 1280]
+        struct ggml_tensor* shared_id_embeds = vision_model->forward(ctx, id_pixel_values);          // [N, hidden_size]
+        struct ggml_tensor* id_embeds        = visual_projection->forward(ctx, shared_id_embeds);    // [N, proj_dim(768)]
+        struct ggml_tensor* id_embeds_2      = visual_projection_2->forward(ctx, shared_id_embeds);  // [N, 1280]
 
         id_embeds   = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 2, 0, 1, 3));
         id_embeds_2 = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds_2, 2, 0, 1, 3));
@@ -340,12 +340,12 @@ struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
         id_embeds = ggml_concat(ctx->ggml_ctx, id_embeds, id_embeds_2, 2);  // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
         id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 1, 2, 0, 3));
 
-        ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
-                                                                  prompt_embeds,
-                                                                  id_embeds,
-                                                                  class_tokens_mask,
-                                                                  class_tokens_mask_pos,
-                                                                  left, right);
+        struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
+                                                                         prompt_embeds,
+                                                                         id_embeds,
+                                                                         class_tokens_mask,
+                                                                         class_tokens_mask_pos,
+                                                                         left, right);
         return updated_prompt_embeds;
     }
 };
@@ -365,29 +365,29 @@ struct PhotoMakerIDEncoder_CLIPInsightfaceExtendtokenBlock : public CLIPVisionMo
                                                                                         num_tokens));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* id_pixel_values,
-                         ggml_tensor* prompt_embeds,
-                         ggml_tensor* class_tokens_mask,
-                         ggml_tensor* class_tokens_mask_pos,
-                         ggml_tensor* id_embeds,
-                         ggml_tensor* left,
-                         ggml_tensor* right) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* id_pixel_values,
+                                struct ggml_tensor* prompt_embeds,
+                                struct ggml_tensor* class_tokens_mask,
+                                struct ggml_tensor* class_tokens_mask_pos,
+                                struct ggml_tensor* id_embeds,
+                                struct ggml_tensor* left,
+                                struct ggml_tensor* right) {
         // x: [N, channels, h, w]
         auto vision_model      = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
         auto fuse_module       = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
         auto qformer_perceiver = std::dynamic_pointer_cast<QFormerPerceiver>(blocks["qformer_perceiver"]);
 
-        // ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values);          // [N, hidden_size]
-        ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values, false);  // [N, hidden_size]
-        id_embeds                      = qformer_perceiver->forward(ctx, id_embeds, last_hidden_state);
+        // struct ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values);          // [N, hidden_size]
+        struct ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values, false);  // [N, hidden_size]
+        id_embeds                             = qformer_perceiver->forward(ctx, id_embeds, last_hidden_state);
 
-        ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
-                                                                  prompt_embeds,
-                                                                  id_embeds,
-                                                                  class_tokens_mask,
-                                                                  class_tokens_mask_pos,
-                                                                  left, right);
+        struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
+                                                                         prompt_embeds,
+                                                                         id_embeds,
+                                                                         class_tokens_mask,
+                                                                         class_tokens_mask_pos,
+                                                                         left, right);
         return updated_prompt_embeds;
     }
 };
@@ -436,17 +436,18 @@ public:
         return pm_version;
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         if (pm_version == PM_VERSION_1)
             id_encoder.get_param_tensors(tensors, prefix);
         else if (pm_version == PM_VERSION_2)
             id_encoder2.get_param_tensors(tensors, prefix);
     }
 
-    ggml_cgraph* build_graph(const sd::Tensor<float>& id_pixel_values_tensor,
-                             const sd::Tensor<float>& prompt_embeds_tensor,
-                             std::vector<bool>& class_tokens_mask,
-                             const sd::Tensor<float>& id_embeds_tensor = {}) {
+    struct ggml_cgraph* build_graph(  // struct ggml_allocr* allocr,
+        struct ggml_tensor* id_pixel_values,
+        struct ggml_tensor* prompt_embeds,
+        std::vector<bool>& class_tokens_mask,
+        struct ggml_tensor* id_embeds) {
         ctm.clear();
         ctmf16.clear();
         ctmpos.clear();
@@ -457,20 +458,20 @@ public:
 
         auto runner_ctx = get_context();
 
-        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
-
-        ggml_tensor* id_pixel_values = make_input(id_pixel_values_tensor);
-        ggml_tensor* prompt_embeds   = make_input(prompt_embeds_tensor);
-        ggml_tensor* id_embeds       = make_optional_input(id_embeds_tensor);
+        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
         int64_t hidden_size = prompt_embeds->ne[0];
         int64_t seq_length  = prompt_embeds->ne[1];
         ggml_type type      = GGML_TYPE_F32;
 
-        ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
+        struct ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
 
-        ggml_tensor* left  = nullptr;
-        ggml_tensor* right = nullptr;
+        struct ggml_tensor* id_pixel_values_d = to_backend(id_pixel_values);
+        struct ggml_tensor* prompt_embeds_d   = to_backend(prompt_embeds);
+        struct ggml_tensor* id_embeds_d       = to_backend(id_embeds);
+
+        struct ggml_tensor* left  = nullptr;
+        struct ggml_tensor* right = nullptr;
         for (int i = 0; i < class_tokens_mask.size(); i++) {
             if (class_tokens_mask[i]) {
                 // printf(" 1,");
@@ -494,7 +495,7 @@ public:
             right = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type,
                                        hidden_size, seq_length - ctmpos[ctmpos.size() - 1] - 1, 1);
         }
-        ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
+        struct ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
 
         {
             if (type == GGML_TYPE_F16)
@@ -525,21 +526,21 @@ public:
                 }
             }
         }
-        ggml_tensor* updated_prompt_embeds = nullptr;
+        struct ggml_tensor* updated_prompt_embeds = nullptr;
         if (pm_version == PM_VERSION_1)
             updated_prompt_embeds = id_encoder.forward(&runner_ctx,
-                                                       id_pixel_values,
-                                                       prompt_embeds,
+                                                       id_pixel_values_d,
+                                                       prompt_embeds_d,
                                                        class_tokens_mask_d,
                                                        class_tokens_mask_pos,
                                                        left, right);
         else if (pm_version == PM_VERSION_2)
             updated_prompt_embeds = id_encoder2.forward(&runner_ctx,
-                                                        id_pixel_values,
-                                                        prompt_embeds,
+                                                        id_pixel_values_d,
+                                                        prompt_embeds_d,
                                                         class_tokens_mask_d,
                                                         class_tokens_mask_pos,
-                                                        id_embeds,
+                                                        id_embeds_d,
                                                         left, right);
 
         ggml_build_forward_expand(gf, updated_prompt_embeds);
@@ -547,21 +548,25 @@ public:
         return gf;
     }
 
-    sd::Tensor<float> compute(const int n_threads,
-                              const sd::Tensor<float>& id_pixel_values,
-                              const sd::Tensor<float>& prompt_embeds,
-                              const sd::Tensor<float>& id_embeds,
-                              std::vector<bool>& class_tokens_mask) {
-        auto get_graph = [&]() -> ggml_cgraph* {
+    bool compute(const int n_threads,
+                 struct ggml_tensor* id_pixel_values,
+                 struct ggml_tensor* prompt_embeds,
+                 struct ggml_tensor* id_embeds,
+                 std::vector<bool>& class_tokens_mask,
+                 struct ggml_tensor** updated_prompt_embeds,
+                 ggml_context* output_ctx) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            // return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask);
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+        // GGMLRunner::compute(get_graph, n_threads, updated_prompt_embeds);
+        return GGMLRunner::compute(get_graph, n_threads, true, updated_prompt_embeds, output_ctx);
     }
 };
 
 struct PhotoMakerIDEmbed : public GGMLRunner {
-    std::map<std::string, ggml_tensor*> tensors;
+    std::map<std::string, struct ggml_tensor*> tensors;
     std::string file_path;
     ModelLoader* model_loader;
     bool load_failed = false;
@@ -601,11 +606,11 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             }
             if (dry_run) {
                 std::lock_guard<std::mutex> lock(tensor_mutex);
-                ggml_tensor* real = ggml_new_tensor(params_ctx,
-                                                    tensor_storage.type,
-                                                    tensor_storage.n_dims,
-                                                    tensor_storage.ne);
-                tensors[name]     = real;
+                struct ggml_tensor* real = ggml_new_tensor(params_ctx,
+                                                           tensor_storage.type,
+                                                           tensor_storage.n_dims,
+                                                           tensor_storage.ne);
+                tensors[name]            = real;
             } else {
                 auto real   = tensors[name];
                 *dst_tensor = real;
@@ -624,8 +629,8 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
         return true;
     }
 
-    ggml_tensor* get() {
-        std::map<std::string, ggml_tensor*>::iterator pos;
+    struct ggml_tensor* get() {
+        std::map<std::string, struct ggml_tensor*>::iterator pos;
         pos = tensors.find("pmid.id_embeds");
         if (pos != tensors.end())
             return pos->second;

preprocessing.hpp

@@ -0,0 +1,226 @@
+#ifndef __PREPROCESSING_HPP__
+#define __PREPROCESSING_HPP__
+
+#include "ggml_extend.hpp"
+#define M_PI_ 3.14159265358979323846
+
+void convolve(struct ggml_tensor* input, struct ggml_tensor* output, struct ggml_tensor* kernel, int padding) {
+    struct ggml_init_params params;
+    params.mem_size                 = 80 * input->ne[0] * input->ne[1];  // 20M for 512x512
+    params.mem_buffer               = nullptr;
+    params.no_alloc                 = false;
+    struct ggml_context* ctx0       = ggml_init(params);
+    struct ggml_tensor* kernel_fp16 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, kernel->ne[0], kernel->ne[1], 1, 1);
+    ggml_fp32_to_fp16_row((float*)kernel->data, (ggml_fp16_t*)kernel_fp16->data, ggml_nelements(kernel));
+    ggml_tensor* h  = ggml_conv_2d(ctx0, kernel_fp16, input, 1, 1, padding, padding, 1, 1);
+    ggml_cgraph* gf = ggml_new_graph(ctx0);
+    ggml_build_forward_expand(gf, ggml_cpy(ctx0, h, output));
+    ggml_graph_compute_with_ctx(ctx0, gf, 1);
+    ggml_free(ctx0);
+}
+
+void gaussian_kernel(struct ggml_tensor* kernel) {
+    int ks_mid   = kernel->ne[0] / 2;
+    float sigma  = 1.4f;
+    float normal = 1.f / (2.0f * M_PI_ * powf(sigma, 2.0f));
+    for (int y = 0; y < kernel->ne[0]; y++) {
+        float gx = -ks_mid + y;
+        for (int x = 0; x < kernel->ne[1]; x++) {
+            float gy = -ks_mid + x;
+            float k_ = expf(-((gx * gx + gy * gy) / (2.0f * powf(sigma, 2.0f)))) * normal;
+            ggml_ext_tensor_set_f32(kernel, k_, x, y);
+        }
+    }
+}
+
+void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
+    for (int iy = 0; iy < rgb_img->ne[1]; iy++) {
+        for (int ix = 0; ix < rgb_img->ne[0]; ix++) {
+            float r    = ggml_ext_tensor_get_f32(rgb_img, ix, iy);
+            float g    = ggml_ext_tensor_get_f32(rgb_img, ix, iy, 1);
+            float b    = ggml_ext_tensor_get_f32(rgb_img, ix, iy, 2);
+            float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b;
+            ggml_ext_tensor_set_f32(grayscale, gray, ix, iy);
+        }
+    }
+}
+
+void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
+    int n_elements = ggml_nelements(h);
+    float* dx      = (float*)x->data;
+    float* dy      = (float*)y->data;
+    float* dh      = (float*)h->data;
+    for (int i = 0; i < n_elements; i++) {
+        dh[i] = sqrtf(dx[i] * dx[i] + dy[i] * dy[i]);
+    }
+}
+
+void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
+    int n_elements = ggml_nelements(h);
+    float* dx      = (float*)x->data;
+    float* dy      = (float*)y->data;
+    float* dh      = (float*)h->data;
+    for (int i = 0; i < n_elements; i++) {
+        dh[i] = atan2f(dy[i], dx[i]);
+    }
+}
+
+void normalize_tensor(struct ggml_tensor* g) {
+    int n_elements = ggml_nelements(g);
+    float* dg      = (float*)g->data;
+    float max      = -INFINITY;
+    for (int i = 0; i < n_elements; i++) {
+        max = dg[i] > max ? dg[i] : max;
+    }
+    max = 1.0f / max;
+    for (int i = 0; i < n_elements; i++) {
+        dg[i] *= max;
+    }
+}
+
+void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struct ggml_tensor* D) {
+    for (int iy = 1; iy < result->ne[1] - 1; iy++) {
+        for (int ix = 1; ix < result->ne[0] - 1; ix++) {
+            float angle = ggml_ext_tensor_get_f32(D, ix, iy) * 180.0f / M_PI_;
+            angle       = angle < 0.0f ? angle += 180.0f : angle;
+            float q     = 1.0f;
+            float r     = 1.0f;
+
+            // angle 0
+            if ((0 >= angle && angle < 22.5f) || (157.5f >= angle && angle <= 180)) {
+                q = ggml_ext_tensor_get_f32(G, ix, iy + 1);
+                r = ggml_ext_tensor_get_f32(G, ix, iy - 1);
+            }
+            // angle 45
+            else if (22.5f >= angle && angle < 67.5f) {
+                q = ggml_ext_tensor_get_f32(G, ix + 1, iy - 1);
+                r = ggml_ext_tensor_get_f32(G, ix - 1, iy + 1);
+            }
+            // angle 90
+            else if (67.5f >= angle && angle < 112.5) {
+                q = ggml_ext_tensor_get_f32(G, ix + 1, iy);
+                r = ggml_ext_tensor_get_f32(G, ix - 1, iy);
+            }
+            // angle 135
+            else if (112.5 >= angle && angle < 157.5f) {
+                q = ggml_ext_tensor_get_f32(G, ix - 1, iy - 1);
+                r = ggml_ext_tensor_get_f32(G, ix + 1, iy + 1);
+            }
+
+            float cur = ggml_ext_tensor_get_f32(G, ix, iy);
+            if ((cur >= q) && (cur >= r)) {
+                ggml_ext_tensor_set_f32(result, cur, ix, iy);
+            } else {
+                ggml_ext_tensor_set_f32(result, 0.0f, ix, iy);
+            }
+        }
+    }
+}
+
+void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float low_threshold, float weak, float strong) {
+    int n_elements = ggml_nelements(img);
+    float* imd     = (float*)img->data;
+    float max      = -INFINITY;
+    for (int i = 0; i < n_elements; i++) {
+        max = imd[i] > max ? imd[i] : max;
+    }
+    float ht = max * high_threshold;
+    float lt = ht * low_threshold;
+    for (int i = 0; i < n_elements; i++) {
+        float img_v = imd[i];
+        if (img_v >= ht) {  // strong pixel
+            imd[i] = strong;
+        } else if (img_v <= ht && img_v >= lt) {  // strong pixel
+            imd[i] = weak;
+        }
+    }
+
+    for (int iy = 0; iy < img->ne[1]; iy++) {
+        for (int ix = 0; ix < img->ne[0]; ix++) {
+            if (ix >= 3 && ix <= img->ne[0] - 3 && iy >= 3 && iy <= img->ne[1] - 3) {
+                ggml_ext_tensor_set_f32(img, ggml_ext_tensor_get_f32(img, ix, iy), ix, iy);
+            } else {
+                ggml_ext_tensor_set_f32(img, 0.0f, ix, iy);
+            }
+        }
+    }
+
+    // hysteresis
+    for (int iy = 1; iy < img->ne[1] - 1; iy++) {
+        for (int ix = 1; ix < img->ne[0] - 1; ix++) {
+            float imd_v = ggml_ext_tensor_get_f32(img, ix, iy);
+            if (imd_v == weak) {
+                if (ggml_ext_tensor_get_f32(img, ix + 1, iy - 1) == strong || ggml_ext_tensor_get_f32(img, ix + 1, iy) == strong ||
+                    ggml_ext_tensor_get_f32(img, ix, iy - 1) == strong || ggml_ext_tensor_get_f32(img, ix, iy + 1) == strong ||
+                    ggml_ext_tensor_get_f32(img, ix - 1, iy - 1) == strong || ggml_ext_tensor_get_f32(img, ix - 1, iy) == strong) {
+                    ggml_ext_tensor_set_f32(img, strong, ix, iy);
+                } else {
+                    ggml_ext_tensor_set_f32(img, 0.0f, ix, iy);
+                }
+            }
+        }
+    }
+}
+
+bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
+    struct ggml_init_params params;
+    params.mem_size               = static_cast<size_t>(40 * img.width * img.height);  // 10MB for 512x512
+    params.mem_buffer             = nullptr;
+    params.no_alloc               = false;
+    struct ggml_context* work_ctx = ggml_init(params);
+
+    if (!work_ctx) {
+        LOG_ERROR("ggml_init() failed");
+        return false;
+    }
+
+    float kX[9] = {
+        -1, 0, 1,
+        -2, 0, 2,
+        -1, 0, 1};
+
+    float kY[9] = {
+        1, 2, 1,
+        0, 0, 0,
+        -1, -2, -1};
+
+    // generate kernel
+    int kernel_size             = 5;
+    struct ggml_tensor* gkernel = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, 1);
+    struct ggml_tensor* sf_kx   = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
+    memcpy(sf_kx->data, kX, ggml_nbytes(sf_kx));
+    struct ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
+    memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
+    gaussian_kernel(gkernel);
+    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
+    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
+    struct ggml_tensor* iX         = ggml_dup_tensor(work_ctx, image_gray);
+    struct ggml_tensor* iY         = ggml_dup_tensor(work_ctx, image_gray);
+    struct ggml_tensor* G          = ggml_dup_tensor(work_ctx, image_gray);
+    struct ggml_tensor* tetha      = ggml_dup_tensor(work_ctx, image_gray);
+    sd_image_to_ggml_tensor(img, image);
+    grayscale(image, image_gray);
+    convolve(image_gray, image_gray, gkernel, 2);
+    convolve(image_gray, iX, sf_kx, 1);
+    convolve(image_gray, iY, sf_ky, 1);
+    prop_hypot(iX, iY, G);
+    normalize_tensor(G);
+    prop_arctan2(iX, iY, tetha);
+    non_max_supression(image_gray, G, tetha);
+    threshold_hystersis(image_gray, high_threshold, low_threshold, weak, strong);
+    // to RGB channels
+    for (int iy = 0; iy < img.height; iy++) {
+        for (int ix = 0; ix < img.width; ix++) {
+            float gray = ggml_ext_tensor_get_f32(image_gray, ix, iy);
+            gray       = inverse ? 1.0f - gray : gray;
+            ggml_ext_tensor_set_f32(image, gray, ix, iy);
+            ggml_ext_tensor_set_f32(image, gray, ix, iy, 1);
+            ggml_ext_tensor_set_f32(image, gray, ix, iy, 2);
+        }
+    }
+    ggml_tensor_to_sd_image(image, img.data);
+    ggml_free(work_ctx);
+    return true;
+}
+
+#endif  // __PREPROCESSING_HPP__

qwen_image.hpp

@@ -3,8 +3,9 @@
 
 #include <memory>
 
-#include "common_block.hpp"
+#include "common.hpp"
 #include "flux.hpp"
+#include "ggml_extend.hpp"
 
 namespace Qwen {
     constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480;
@@ -26,9 +27,9 @@ namespace Qwen {
             blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, out_dim, sample_proj_bias));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* sample,
-                             ggml_tensor* condition = nullptr) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* sample,
+                                    struct ggml_tensor* condition = nullptr) {
             if (condition != nullptr) {
                 auto cond_proj = std::dynamic_pointer_cast<Linear>(blocks["cond_proj"]);
                 sample         = ggml_add(ctx->ggml_ctx, sample, cond_proj->forward(ctx, condition));
@@ -49,8 +50,8 @@ namespace Qwen {
             blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* timesteps) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* timesteps) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
@@ -107,10 +108,10 @@ namespace Qwen {
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                      ggml_tensor* img,
-                                                      ggml_tensor* txt,
-                                                      ggml_tensor* pe,
-                                                      ggml_tensor* mask = nullptr) {
+                                                      struct ggml_tensor* img,
+                                                      struct ggml_tensor* txt,
+                                                      struct ggml_tensor* pe,
+                                                      struct ggml_tensor* mask = nullptr) {
             // img: [N, n_img_token, hidden_size]
             // txt: [N, n_txt_token, hidden_size]
             // pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
@@ -161,25 +162,26 @@ namespace Qwen {
             auto k = ggml_concat(ctx->ggml_ctx, txt_k, img_k, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
             auto v = ggml_concat(ctx->ggml_ctx, txt_v, img_v, 2);  // [N, n_txt_token + n_img_token, n_head, d_head]
 
-            auto attn         = Rope::attention(ctx, q, k, v, pe, mask, (1.0f / 128.f));  // [N, n_txt_token + n_img_token, n_head*d_head]
+            auto attn         = Rope::attention(ctx, q, k, v, pe, mask, (1.0f / 128.f));                  // [N, n_txt_token + n_img_token, n_head*d_head]
+            attn              = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, attn, 0, 2, 1, 3));  // [n_txt_token + n_img_token, N, hidden_size]
             auto txt_attn_out = ggml_view_3d(ctx->ggml_ctx,
                                              attn,
                                              attn->ne[0],
+                                             attn->ne[1],
                                              txt->ne[1],
-                                             attn->ne[2],
                                              attn->nb[1],
                                              attn->nb[2],
-                                             0);  // [N, n_txt_token, n_head*d_head]
+                                             0);                                                                  // [n_txt_token, N, hidden_size]
+            txt_attn_out      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, txt_attn_out, 0, 2, 1, 3));  // [N, n_txt_token, hidden_size]
             auto img_attn_out = ggml_view_3d(ctx->ggml_ctx,
                                              attn,
                                              attn->ne[0],
+                                             attn->ne[1],
                                              img->ne[1],
-                                             attn->ne[2],
                                              attn->nb[1],
                                              attn->nb[2],
-                                             txt->ne[1] * attn->nb[1]);  // [N, n_img_token, n_head*d_head]
-            img_attn_out      = ggml_cont(ctx->ggml_ctx, img_attn_out);
-            txt_attn_out      = ggml_cont(ctx->ggml_ctx, txt_attn_out);
+                                             attn->nb[2] * txt->ne[1]);                                           // [n_img_token, N, hidden_size]
+            img_attn_out      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, img_attn_out, 0, 2, 1, 3));  // [N, n_img_token, hidden_size]
 
             img_attn_out = to_out_0->forward(ctx, img_attn_out);
             txt_attn_out = to_add_out->forward(ctx, txt_attn_out);
@@ -189,16 +191,11 @@ namespace Qwen {
     };
 
     class QwenImageTransformerBlock : public GGMLBlock {
-    protected:
-        bool zero_cond_t;
-
     public:
         QwenImageTransformerBlock(int64_t dim,
                                   int64_t num_attention_heads,
                                   int64_t attention_head_dim,
-                                  float eps        = 1e-6,
-                                  bool zero_cond_t = false)
-            : zero_cond_t(zero_cond_t) {
+                                  float eps = 1e-6) {
             // img_mod.0 is nn.SiLU()
             blocks["img_mod.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, 6 * dim, true));
 
@@ -211,7 +208,7 @@ namespace Qwen {
 
             blocks["txt_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
             blocks["txt_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
-            blocks["txt_mlp"]   = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim, 4, FeedForward::Activation::GELU, true));
+            blocks["txt_mlp"]   = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim, 4, FeedForward::Activation::GELU));
 
             blocks["attn"] = std::shared_ptr<GGMLBlock>(new QwenImageAttention(dim,
                                                                                attention_head_dim,
@@ -223,37 +220,11 @@ namespace Qwen {
                                                                                eps));
         }
 
-        std::vector<ggml_tensor*> get_mod_params_vec(ggml_context* ctx, ggml_tensor* mod_params, ggml_tensor* index = nullptr) {
-            // index: [N, n_img_token]
-            // mod_params: [N, hidden_size * 12]
-            if (index == nullptr) {
-                return ggml_ext_chunk(ctx, mod_params, 6, 0);
-            }
-            mod_params          = ggml_reshape_1d(ctx, mod_params, ggml_nelements(mod_params));
-            auto mod_params_vec = ggml_ext_chunk(ctx, mod_params, 12, 0);
-            index               = ggml_reshape_3d(ctx, index, 1, index->ne[0], index->ne[1]);                                      // [N, n_img_token, 1]
-            index               = ggml_repeat_4d(ctx, index, mod_params_vec[0]->ne[0], index->ne[1], index->ne[2], index->ne[3]);  // [N, n_img_token, hidden_size]
-            std::vector<ggml_tensor*> mod_results;
-            for (int i = 0; i < 6; i++) {
-                auto mod_0 = mod_params_vec[i];
-                auto mod_1 = mod_params_vec[i + 6];
-
-                // mod_result = torch.where(index == 0, mod_0, mod_1)
-                // mod_result = (1 - index)*mod_0 + index*mod_1
-                mod_0           = ggml_sub(ctx, ggml_repeat(ctx, mod_0, index), ggml_mul(ctx, index, mod_0));  // [N, n_img_token, hidden_size]
-                mod_1           = ggml_mul(ctx, index, mod_1);                                                 // [N, n_img_token, hidden_size]
-                auto mod_result = ggml_add(ctx, mod_0, mod_1);
-                mod_results.push_back(mod_result);
-            }
-            return mod_results;
-        }
-
         virtual std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                              ggml_tensor* img,
-                                                              ggml_tensor* txt,
-                                                              ggml_tensor* t_emb,
-                                                              ggml_tensor* pe,
-                                                              ggml_tensor* modulate_index = nullptr) {
+                                                              struct ggml_tensor* img,
+                                                              struct ggml_tensor* txt,
+                                                              struct ggml_tensor* t_emb,
+                                                              struct ggml_tensor* pe) {
             // img: [N, n_img_token, hidden_size]
             // txt: [N, n_txt_token, hidden_size]
             // pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
@@ -273,18 +244,14 @@ namespace Qwen {
 
             auto img_mod_params    = ggml_silu(ctx->ggml_ctx, t_emb);
             img_mod_params         = img_mod_1->forward(ctx, img_mod_params);
-            auto img_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, img_mod_params, modulate_index);
-
-            if (zero_cond_t) {
-                t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
-            }
+            auto img_mod_param_vec = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0);
 
             auto txt_mod_params    = ggml_silu(ctx->ggml_ctx, t_emb);
             txt_mod_params         = txt_mod_1->forward(ctx, txt_mod_params);
-            auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params);
+            auto txt_mod_param_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
 
             auto img_normed    = img_norm1->forward(ctx, img);
-            auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
+            auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1]);
             auto img_gate1     = img_mod_param_vec[2];
 
             auto txt_normed    = txt_norm1->forward(ctx, txt);
@@ -297,7 +264,7 @@ namespace Qwen {
             txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_gate1));
 
             auto img_normed2    = img_norm2->forward(ctx, img);
-            auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr);
+            auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4]);
             auto img_gate2      = img_mod_param_vec[5];
 
             auto txt_normed2    = txt_norm2->forward(ctx, txt);
@@ -325,9 +292,9 @@ namespace Qwen {
             blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(conditioning_embedding_dim, embedding_dim * 2, bias));
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* c) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* c) {
             // x: [N, n_token, hidden_size]
             // c: [N, hidden_size]
             // return: [N, n_token, patch_size * patch_size * out_channels]
@@ -348,17 +315,16 @@ namespace Qwen {
     };
 
     struct QwenImageParams {
-        int patch_size              = 2;
+        int64_t patch_size          = 2;
         int64_t in_channels         = 64;
         int64_t out_channels        = 16;
-        int num_layers              = 60;
+        int64_t num_layers          = 60;
         int64_t attention_head_dim  = 128;
         int64_t num_attention_heads = 24;
         int64_t joint_attention_dim = 3584;
-        int theta                   = 10000;
+        float theta                 = 10000;
         std::vector<int> axes_dim   = {16, 56, 56};
-        int axes_dim_sum            = 128;
-        bool zero_cond_t            = false;
+        int64_t axes_dim_sum        = 128;
     };
 
     class QwenImageModel : public GGMLBlock {
@@ -380,8 +346,7 @@ namespace Qwen {
                 auto block                                        = std::shared_ptr<GGMLBlock>(new QwenImageTransformerBlock(inner_dim,
                                                                                                                              params.num_attention_heads,
                                                                                                                              params.attention_head_dim,
-                                                                                                                             1e-6f,
-                                                                                                                             params.zero_cond_t));
+                                                                                                                             1e-6f));
                 blocks["transformer_blocks." + std::to_string(i)] = block;
             }
 
@@ -389,12 +354,74 @@ namespace Qwen {
             blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
         }
 
-        ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
-                                  ggml_tensor* x,
-                                  ggml_tensor* timestep,
-                                  ggml_tensor* context,
-                                  ggml_tensor* pe,
-                                  ggml_tensor* modulate_index = nullptr) {
+        struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
+                                              struct ggml_tensor* x) {
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+
+            int pad_h = (params.patch_size - H % params.patch_size) % params.patch_size;
+            int pad_w = (params.patch_size - W % params.patch_size) % params.patch_size;
+            x         = ggml_pad(ctx, x, pad_w, pad_h, 0, 0);  // [N, C, H + pad_h, W + pad_w]
+            return x;
+        }
+
+        struct ggml_tensor* patchify(struct ggml_context* ctx,
+                                     struct ggml_tensor* x) {
+            // x: [N, C, H, W]
+            // return: [N, h*w, C * patch_size * patch_size]
+            int64_t N = x->ne[3];
+            int64_t C = x->ne[2];
+            int64_t H = x->ne[1];
+            int64_t W = x->ne[0];
+            int64_t p = params.patch_size;
+            int64_t h = H / params.patch_size;
+            int64_t w = W / params.patch_size;
+
+            GGML_ASSERT(h * p == H && w * p == W);
+
+            x = ggml_reshape_4d(ctx, x, p, w, p, h * C * N);       // [N*C*h, p, w, p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*h, w, p, p]
+            x = ggml_reshape_4d(ctx, x, p * p, w * h, C, N);       // [N, C, h*w, p*p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, h*w, C, p*p]
+            x = ggml_reshape_3d(ctx, x, p * p * C, w * h, N);      // [N, h*w, C*p*p]
+            return x;
+        }
+
+        struct ggml_tensor* process_img(struct ggml_context* ctx,
+                                        struct ggml_tensor* x) {
+            x = pad_to_patch_size(ctx, x);
+            x = patchify(ctx, x);
+            return x;
+        }
+
+        struct ggml_tensor* unpatchify(struct ggml_context* ctx,
+                                       struct ggml_tensor* x,
+                                       int64_t h,
+                                       int64_t w) {
+            // x: [N, h*w, C*patch_size*patch_size]
+            // return: [N, C, H, W]
+            int64_t N = x->ne[2];
+            int64_t C = x->ne[0] / params.patch_size / params.patch_size;
+            int64_t H = h * params.patch_size;
+            int64_t W = w * params.patch_size;
+            int64_t p = params.patch_size;
+
+            GGML_ASSERT(C * p * p == x->ne[0]);
+
+            x = ggml_reshape_4d(ctx, x, p * p, C, w * h, N);       // [N, h*w, C, p*p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, C, h*w, p*p]
+            x = ggml_reshape_4d(ctx, x, p, p, w, h * C * N);       // [N*C*h, w, p, p]
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*h, p, w, p]
+            x = ggml_reshape_4d(ctx, x, W, H, C, N);               // [N, C, h*p, w*p]
+
+            return x;
+        }
+
+        struct ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
+                                         struct ggml_tensor* x,
+                                         struct ggml_tensor* timestep,
+                                         struct ggml_tensor* context,
+                                         struct ggml_tensor* pe) {
             auto time_text_embed = std::dynamic_pointer_cast<QwenTimestepProjEmbeddings>(blocks["time_text_embed"]);
             auto txt_norm        = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
             auto img_in          = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
@@ -403,39 +430,30 @@ namespace Qwen {
             auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
 
             auto t_emb = time_text_embed->forward(ctx, timestep);
-            if (params.zero_cond_t) {
-                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
-                t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
-            }
-            auto img = img_in->forward(ctx, x);
-            auto txt = txt_norm->forward(ctx, context);
-            txt      = txt_in->forward(ctx, txt);
+            auto img   = img_in->forward(ctx, x);
+            auto txt   = txt_norm->forward(ctx, context);
+            txt        = txt_in->forward(ctx, txt);
 
             for (int i = 0; i < params.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
 
-                auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
+                auto result = block->forward(ctx, img, txt, t_emb, pe);
                 img         = result.first;
                 txt         = result.second;
             }
 
-            if (params.zero_cond_t) {
-                t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
-            }
-
             img = norm_out->forward(ctx, img, t_emb);
             img = proj_out->forward(ctx, img);
 
             return img;
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* timestep,
-                             ggml_tensor* context,
-                             ggml_tensor* pe,
-                             std::vector<ggml_tensor*> ref_latents = {},
-                             ggml_tensor* modulate_index           = nullptr) {
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                    struct ggml_tensor* x,
+                                    struct ggml_tensor* timestep,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* pe,
+                                    std::vector<ggml_tensor*> ref_latents = {}) {
             // Forward pass of DiT.
             // x: [N, C, H, W]
             // timestep: [N,]
@@ -448,17 +466,20 @@ namespace Qwen {
             int64_t C = x->ne[2];
             int64_t N = x->ne[3];
 
-            auto img           = DiT::pad_and_patchify(ctx, x, params.patch_size, params.patch_size);
-            int64_t img_tokens = img->ne[1];
+            auto img            = process_img(ctx->ggml_ctx, x);
+            uint64_t img_tokens = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = DiT::pad_and_patchify(ctx, ref, params.patch_size, params.patch_size);
+                    ref = process_img(ctx->ggml_ctx, ref);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
 
-            auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
+            int64_t h_len = ((H + (params.patch_size / 2)) / params.patch_size);
+            int64_t w_len = ((W + (params.patch_size / 2)) / params.patch_size);
+
+            auto out = forward_orig(ctx, img, timestep, context, pe);  // [N, h_len*w_len, ph*pw*C]
 
             if (out->ne[1] > img_tokens) {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [num_tokens, N, C * patch_size * patch_size]
@@ -466,7 +487,11 @@ namespace Qwen {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [N, h*w, C * patch_size * patch_size]
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, params.patch_size, params.patch_size);  // [N, C, H, W]
+            out = unpatchify(ctx->ggml_ctx, out, h_len, w_len);  // [N, C, H + pad_h, W + pad_w]
+
+            // slice
+            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N, C, H, W + pad_w]
+            out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N, C, H, W]
 
             return out;
         }
@@ -477,25 +502,19 @@ namespace Qwen {
         QwenImageParams qwen_image_params;
         QwenImageModel qwen_image;
         std::vector<float> pe_vec;
-        std::vector<float> modulate_index_vec;
         SDVersion version;
 
         QwenImageRunner(ggml_backend_t backend,
                         bool offload_params_to_cpu,
                         const String2TensorStorage& tensor_storage_map = {},
                         const std::string prefix                       = "",
-                        SDVersion version                              = VERSION_QWEN_IMAGE,
-                        bool zero_cond_t                               = false)
+                        SDVersion version                              = VERSION_QWEN_IMAGE)
             : GGMLRunner(backend, offload_params_to_cpu) {
-            qwen_image_params.num_layers  = 0;
-            qwen_image_params.zero_cond_t = zero_cond_t;
+            qwen_image_params.num_layers = 0;
             for (auto pair : tensor_storage_map) {
                 std::string tensor_name = pair.first;
                 if (tensor_name.find(prefix) == std::string::npos)
                     continue;
-                if (tensor_name.find("__index_timestep_zero__") != std::string::npos) {
-                    qwen_image_params.zero_cond_t = true;
-                }
                 size_t pos = tensor_name.find("transformer_blocks.");
                 if (pos != std::string::npos) {
                     tensor_name = tensor_name.substr(pos);  // remove prefix
@@ -510,9 +529,6 @@ namespace Qwen {
                 }
             }
             LOG_INFO("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers);
-            if (qwen_image_params.zero_cond_t) {
-                LOG_INFO("use zero_cond_t");
-            }
             qwen_image = QwenImageModel(qwen_image_params);
             qwen_image.init(params_ctx, tensor_storage_map, prefix);
         }
@@ -521,39 +537,36 @@ namespace Qwen {
             return "qwen_image";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
             qwen_image.get_param_tensors(tensors, prefix);
         }
 
-        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
-                                 const sd::Tensor<float>& timesteps_tensor,
-                                 const sd::Tensor<float>& context_tensor,
-                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false) {
-            ggml_cgraph* gf        = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
-            ggml_tensor* x         = make_input(x_tensor);
-            ggml_tensor* timesteps = make_input(timesteps_tensor);
+        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                        struct ggml_tensor* timesteps,
+                                        struct ggml_tensor* context,
+                                        std::vector<ggml_tensor*> ref_latents = {},
+                                        bool increase_ref_index               = false) {
             GGML_ASSERT(x->ne[3] == 1);
-            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));
+            struct ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
+
+            x         = to_backend(x);
+            context   = to_backend(context);
+            timesteps = to_backend(timesteps);
+
+            for (int i = 0; i < ref_latents.size(); i++) {
+                ref_latents[i] = to_backend(ref_latents[i]);
             }
 
-            pe_vec      = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
-                                                  static_cast<int>(x->ne[0]),
+            pe_vec      = Rope::gen_qwen_image_pe(x->ne[1],
+                                                  x->ne[0],
                                                   qwen_image_params.patch_size,
-                                                  static_cast<int>(x->ne[3]),
-                                                  static_cast<int>(context->ne[1]),
+                                                  x->ne[3],
+                                                  context->ne[1],
                                                   ref_latents,
                                                   increase_ref_index,
                                                   qwen_image_params.theta,
-                                                  circular_y_enabled,
-                                                  circular_x_enabled,
                                                   qwen_image_params.axes_dim);
-            int pos_len = static_cast<int>(pe_vec.size() / qwen_image_params.axes_dim_sum / 2);
+            int pos_len = pe_vec.size() / qwen_image_params.axes_dim_sum / 2;
             // LOG_DEBUG("pos_len %d", pos_len);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, qwen_image_params.axes_dim_sum / 2, pos_len);
             // pe->data = pe_vec.data();
@@ -561,100 +574,69 @@ namespace Qwen {
             // pe->data = nullptr;
             set_backend_tensor_data(pe, pe_vec.data());
 
-            ggml_tensor* modulate_index = nullptr;
-            if (qwen_image_params.zero_cond_t) {
-                modulate_index_vec.clear();
-
-                int64_t h_len          = ((x->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
-                int64_t w_len          = ((x->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
-                int64_t num_img_tokens = h_len * w_len;
-
-                modulate_index_vec.insert(modulate_index_vec.end(), num_img_tokens, 0.f);
-                int64_t num_ref_img_tokens = 0;
-                for (ggml_tensor* ref : ref_latents) {
-                    int64_t h_len = ((ref->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
-                    int64_t w_len = ((ref->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
-
-                    num_ref_img_tokens += h_len * w_len;
-                }
-
-                if (num_ref_img_tokens > 0) {
-                    modulate_index_vec.insert(modulate_index_vec.end(), num_ref_img_tokens, 1.f);
-                }
-
-                modulate_index = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, modulate_index_vec.size());
-                set_backend_tensor_data(modulate_index, modulate_index_vec.data());
-            }
-
             auto runner_ctx = get_context();
 
-            ggml_tensor* out = qwen_image.forward(&runner_ctx,
-                                                  x,
-                                                  timesteps,
-                                                  context,
-                                                  pe,
-                                                  ref_latents,
-                                                  modulate_index);
+            struct ggml_tensor* out = qwen_image.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 = {},
-                                  bool increase_ref_index                           = false) {
+        bool compute(int n_threads,
+                     struct ggml_tensor* x,
+                     struct ggml_tensor* timesteps,
+                     struct ggml_tensor* context,
+                     std::vector<ggml_tensor*> ref_latents = {},
+                     bool increase_ref_index               = false,
+                     struct ggml_tensor** output           = nullptr,
+                     struct ggml_context* output_ctx       = nullptr) {
             // x: [N, in_channels, h, w]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
-            auto get_graph = [&]() -> ggml_cgraph* {
+            auto get_graph = [&]() -> struct ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
         }
 
         void test() {
-            ggml_init_params params;
+            struct ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            ggml_context* ctx = ggml_init(params);
-            GGML_ASSERT(ctx != nullptr);
+            struct ggml_context* work_ctx = ggml_init(params);
+            GGML_ASSERT(work_ctx != nullptr);
 
             {
-                // auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 16, 16, 16, 1);
+                // auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 16, 16, 16, 1);
                 // ggml_set_f32(x, 0.01f);
-                auto x = sd::load_tensor_from_file_as_tensor<float>("./qwen_image_x.bin");
-                print_sd_tensor(x);
+                auto x = load_tensor_from_file(work_ctx, "./qwen_image_x.bin");
+                print_ggml_tensor(x);
 
                 std::vector<float> timesteps_vec(1, 1000.f);
-                auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
+                auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
 
-                // auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3584, 256, 1);
+                // auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 3584, 256, 1);
                 // ggml_set_f32(context, 0.01f);
-                auto context = sd::load_tensor_from_file_as_tensor<float>("./qwen_image_context.bin");
-                print_sd_tensor(context);
+                auto context = load_tensor_from_file(work_ctx, "./qwen_image_context.bin");
+                print_ggml_tensor(context);
 
-                sd::Tensor<float> out;
+                struct ggml_tensor* out = nullptr;
 
-                int64_t t0   = ggml_time_ms();
-                auto out_opt = compute(8,
-                                       x,
-                                       timesteps,
-                                       context,
-                                       {},
-                                       false);
-                int64_t t1   = ggml_time_ms();
+                int t0 = ggml_time_ms();
+                compute(8, x, timesteps, context, {}, false, &out, work_ctx);
+                int t1 = ggml_time_ms();
 
-                GGML_ASSERT(!out_opt.empty());
-                out = std::move(out_opt);
-                print_sd_tensor(out);
-                LOG_DEBUG("qwen_image test done in %lldms", t1 - t0);
+                print_ggml_tensor(out);
+                LOG_DEBUG("qwen_image test done in %dms", t1 - t0);
             }
         }
 

rng_mt19937.hpp

@@ -90,7 +90,7 @@ class MT19937RNG : public RNG {
             float u1    = 1.0f - data[j];
             float u2    = data[j + 8];
             float r     = std::sqrt(-2.0f * std::log(u1));
-            float theta = 2.0f * 3.14159265358979323846f * u2;
+            float theta = 2.0f * 3.14159265358979323846 * u2;
             data[j]     = r * std::cos(theta) * std + mean;
             data[j + 8] = r * std::sin(theta) * std + mean;
         }

rope.hpp

@@ -1,8 +1,6 @@
 #ifndef __ROPE_HPP__
 #define __ROPE_HPP__
 
-#include <algorithm>
-#include <cmath>
 #include <vector>
 #include "ggml_extend.hpp"
 
@@ -22,11 +20,11 @@ namespace Rope {
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> transpose(const std::vector<std::vector<float>>& mat) {
-        size_t rows = mat.size();
-        size_t cols = mat[0].size();
+        int rows = mat.size();
+        int cols = mat[0].size();
         std::vector<std::vector<float>> transposed(cols, std::vector<float>(rows));
-        for (size_t i = 0; i < rows; ++i) {
-            for (size_t j = 0; j < cols; ++j) {
+        for (int i = 0; i < rows; ++i) {
+            for (int j = 0; j < cols; ++j) {
                 transposed[j][i] = mat[i][j];
             }
         }
@@ -41,10 +39,7 @@ namespace Rope {
         return flat_vec;
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
-                                                           int dim,
-                                                           float theta,
-                                                           const std::vector<int>& axis_wrap_dims = {}) {
+    __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos, int dim, int theta) {
         assert(dim % 2 == 0);
         int half_dim = dim / 2;
 
@@ -52,31 +47,14 @@ namespace Rope {
 
         std::vector<float> omega(half_dim);
         for (int i = 0; i < half_dim; ++i) {
-            omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
+            omega[i] = 1.0 / std::pow(theta, scale[i]);
         }
 
-        size_t pos_size = pos.size();
+        int pos_size = pos.size();
         std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
-        for (size_t i = 0; i < pos_size; ++i) {
-            for (size_t j = 0; j < half_dim; ++j) {
-                float angle = pos[i] * omega[j];
-                if (!axis_wrap_dims.empty()) {
-                    size_t wrap_size = axis_wrap_dims.size();
-                    // mod batch size since we only store this for one item in the batch
-                    size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
-                    int wrap_dim    = axis_wrap_dims[wrap_idx];
-                    if (wrap_dim > 0) {
-                        constexpr float TWO_PI = 6.28318530717958647692f;
-                        float cycles           = omega[j] * wrap_dim / TWO_PI;
-                        // closest periodic harmonic, necessary to ensure things neatly tile
-                        // without this round, things don't tile at the boundaries and you end up
-                        // with the model knowing what is "center"
-                        float rounded = std::round(cycles);
-                        angle         = pos[i] * TWO_PI * rounded / wrap_dim;
-                    }
-                }
-
-                out[i][j] = angle;
+        for (int i = 0; i < pos_size; ++i) {
+            for (int j = 0; j < half_dim; ++j) {
+                out[i][j] = pos[i] * omega[j];
             }
         }
 
@@ -99,7 +77,7 @@ namespace Rope {
         for (int dim = 0; dim < axes_dim_num; dim++) {
             if (arange_dims.find(dim) != arange_dims.end()) {
                 for (int i = 0; i < bs * context_len; i++) {
-                    txt_ids[i][dim] = 1.f * (i % context_len);
+                    txt_ids[i][dim] = (i % context_len);
                 }
             }
         }
@@ -111,29 +89,20 @@ namespace Rope {
                                                                        int patch_size,
                                                                        int bs,
                                                                        int axes_dim_num,
-                                                                       int index       = 0,
-                                                                       int h_offset    = 0,
-                                                                       int w_offset    = 0,
-                                                                       bool scale_rope = false) {
+                                                                       int index    = 0,
+                                                                       int h_offset = 0,
+                                                                       int w_offset = 0) {
         int h_len = (h + (patch_size / 2)) / patch_size;
         int w_len = (w + (patch_size / 2)) / patch_size;
 
         std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
 
-        int h_start = h_offset;
-        int w_start = w_offset;
-
-        if (scale_rope) {
-            h_start -= h_len / 2;
-            w_start -= w_len / 2;
-        }
-
-        std::vector<float> row_ids = linspace<float>(1.f * h_start, 1.f * h_start + h_len - 1, h_len);
-        std::vector<float> col_ids = linspace<float>(1.f * w_start, 1.f * w_start + w_len - 1, w_len);
+        std::vector<float> row_ids = linspace<float>(h_offset, h_len - 1 + h_offset, h_len);
+        std::vector<float> col_ids = linspace<float>(w_offset, w_len - 1 + w_offset, w_len);
 
         for (int i = 0; i < h_len; ++i) {
             for (int j = 0; j < w_len; ++j) {
-                img_ids[i * w_len + j][0] = 1.f * index;
+                img_ids[i * w_len + j][0] = index;
                 img_ids[i * w_len + j][1] = row_ids[i];
                 img_ids[i * w_len + j][2] = col_ids[j];
             }
@@ -167,12 +136,11 @@ namespace Rope {
 
     __STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
                                                   int bs,
-                                                  const std::vector<float>& axis_thetas,
-                                                  const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {}) {
+                                                  int theta,
+                                                  const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> trans_ids = transpose(ids);
         size_t pos_len                            = ids.size() / bs;
-        size_t num_axes                           = axes_dim.size();
+        int num_axes                              = axes_dim.size();
         // for (int i = 0; i < pos_len; i++) {
         //     std::cout << trans_ids[0][i] << " " << trans_ids[1][i] << " " << trans_ids[2][i] << std::endl;
         // }
@@ -182,18 +150,9 @@ namespace Rope {
             emb_dim += d / 2;
 
         std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
-        size_t offset = 0;
-        for (size_t i = 0; i < num_axes; ++i) {
-            std::vector<int> axis_wrap_dims;
-            if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
-                axis_wrap_dims = wrap_dims[i];
-            }
-            float axis_theta = 10000.0f;
-            if (!axis_thetas.empty()) {
-                axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
-            }
-            std::vector<std::vector<float>> rope_emb =
-                rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims);  // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
+        int offset = 0;
+        for (int i = 0; i < num_axes; ++i) {
+            std::vector<std::vector<float>> rope_emb = rope(trans_ids[i], axes_dim[i], theta);  // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
             for (int b = 0; b < bs; ++b) {
                 for (int j = 0; j < pos_len; ++j) {
                     for (int k = 0; k < rope_emb[0].size(); ++k) {
@@ -207,55 +166,43 @@ namespace Rope {
         return flatten(emb);
     }
 
-    __STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
-                                                  int bs,
-                                                  float theta,
-                                                  const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {}) {
-        std::vector<float> axis_thetas(axes_dim.size(), theta);
-        return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims);
-    }
-
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
                                                                    int bs,
                                                                    int axes_dim_num,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
                                                                    bool increase_ref_index,
-                                                                   float ref_index_scale,
-                                                                   bool scale_rope) {
+                                                                   float ref_index_scale) {
         std::vector<std::vector<float>> ids;
-        int curr_h_offset = 0;
-        int curr_w_offset = 0;
-        int index         = 1;
+        uint64_t curr_h_offset = 0;
+        uint64_t curr_w_offset = 0;
+        int index              = 1;
         for (ggml_tensor* ref : ref_latents) {
-            int h_offset = 0;
-            int w_offset = 0;
+            uint64_t h_offset = 0;
+            uint64_t w_offset = 0;
             if (!increase_ref_index) {
                 if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
                     w_offset = curr_w_offset;
                 } else {
                     h_offset = curr_h_offset;
                 }
-                scale_rope = false;
             }
 
-            auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
-                                            static_cast<int>(ref->ne[0]),
+            auto ref_ids = gen_flux_img_ids(ref->ne[1],
+                                            ref->ne[0],
                                             patch_size,
                                             bs,
                                             axes_dim_num,
                                             static_cast<int>(index * ref_index_scale),
                                             h_offset,
-                                            w_offset,
-                                            scale_rope);
+                                            w_offset);
             ids          = concat_ids(ids, ref_ids, bs);
 
             if (increase_ref_index) {
                 index++;
             }
 
-            curr_h_offset = std::max(curr_h_offset, static_cast<int>(ref->ne[1]) + h_offset);
-            curr_w_offset = std::max(curr_w_offset, static_cast<int>(ref->ne[0]) + w_offset);
+            curr_h_offset = std::max(curr_h_offset, ref->ne[1] + h_offset);
+            curr_w_offset = std::max(curr_w_offset, ref->ne[0] + w_offset);
         }
         return ids;
     }
@@ -275,7 +222,7 @@ namespace Rope {
 
         auto ids = concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, ref_index_scale, false);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, ref_index_scale);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -292,8 +239,6 @@ namespace Rope {
                                                      bool increase_ref_index,
                                                      float ref_index_scale,
                                                      int theta,
-                                                     bool circular_h,
-                                                     bool circular_w,
                                                      const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_flux_ids(h,
                                                            w,
@@ -305,47 +250,7 @@ namespace Rope {
                                                            ref_latents,
                                                            increase_ref_index,
                                                            ref_index_scale);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int h_len = (h + (patch_size / 2)) / patch_size;
-            int w_len = (w + (patch_size / 2)) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                size_t pos_len = ids.size() / bs;
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-                size_t cursor           = context_len;  // text first
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-                cursor += img_tokens;
-                // reference latents
-                for (ggml_tensor* ref : ref_latents) {
-                    if (ref == nullptr) {
-                        continue;
-                    }
-                    int ref_h         = static_cast<int>(ref->ne[1]);
-                    int ref_w         = static_cast<int>(ref->ne[0]);
-                    int ref_h_l       = (ref_h + (patch_size / 2)) / patch_size;
-                    int ref_w_l       = (ref_w + (patch_size / 2)) / patch_size;
-                    size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
-                    for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
-                        if (circular_h) {
-                            wrap_dims[1][cursor + token_i] = ref_h_l;
-                        }
-                        if (circular_w) {
-                            wrap_dims[2][cursor + token_i] = ref_w_l;
-                        }
-                    }
-                    cursor += ref_tokens;
-                }
-            }
-        }
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+        return embed_nd(ids, bs, theta, axes_dim);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
@@ -358,7 +263,7 @@ namespace Rope {
         int h_len        = (h + (patch_size / 2)) / patch_size;
         int w_len        = (w + (patch_size / 2)) / patch_size;
         int txt_id_start = std::max(h_len, w_len);
-        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
+        auto txt_ids     = linspace<float>(txt_id_start, context_len + txt_id_start, context_len);
         std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
         for (int i = 0; i < bs; ++i) {
             for (int j = 0; j < txt_ids.size(); ++j) {
@@ -366,10 +271,10 @@ namespace Rope {
             }
         }
         int axes_dim_num = 3;
-        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
+        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num);
         auto ids         = concat_ids(txt_ids_repeated, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, 1.f, true);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, 1.f);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -384,57 +289,9 @@ namespace Rope {
                                                            const std::vector<ggml_tensor*>& ref_latents,
                                                            bool increase_ref_index,
                                                            int theta,
-                                                           bool circular_h,
-                                                           bool circular_w,
                                                            const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
-        std::vector<std::vector<int>> wrap_dims;
-        // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int pad_h = (patch_size - (h % patch_size)) % patch_size;
-            int pad_w = (patch_size - (w % patch_size)) % patch_size;
-            int h_len = (h + pad_h) / patch_size;
-            int w_len = (w + pad_w) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                const size_t total_tokens = ids.size();
-                // Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
-                size_t cursor           = context_len;  // ignore text tokens
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-                cursor += img_tokens;
-                // For each reference image, store wrap sizes as well
-                for (ggml_tensor* ref : ref_latents) {
-                    if (ref == nullptr) {
-                        continue;
-                    }
-                    int ref_h           = static_cast<int>(ref->ne[1]);
-                    int ref_w           = static_cast<int>(ref->ne[0]);
-                    int ref_pad_h       = (patch_size - (ref_h % patch_size)) % patch_size;
-                    int ref_pad_w       = (patch_size - (ref_w % patch_size)) % patch_size;
-                    int ref_h_len       = (ref_h + ref_pad_h) / patch_size;
-                    int ref_w_len       = (ref_w + ref_pad_w) / patch_size;
-                    size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
-                    for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
-                        if (circular_h) {
-                            wrap_dims[1][cursor + token_i] = ref_h_len;
-                        }
-                        if (circular_w) {
-                            wrap_dims[2][cursor + token_i] = ref_w_len;
-                        }
-                    }
-                    cursor += ref_n_tokens;
-                }
-            }
-        }
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+        return embed_nd(ids, bs, theta, axes_dim);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@@ -453,9 +310,9 @@ namespace Rope {
 
         std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
 
-        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
-        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
-        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
+        std::vector<float> t_ids = linspace<float>(t_offset, t_len - 1 + t_offset, t_len);
+        std::vector<float> h_ids = linspace<float>(h_offset, h_len - 1 + h_offset, h_len);
+        std::vector<float> w_ids = linspace<float>(w_offset, w_len - 1 + w_offset, w_len);
 
         for (int i = 0; i < t_len; ++i) {
             for (int j = 0; j < h_len; ++j) {
@@ -488,7 +345,7 @@ namespace Rope {
                                                     int theta,
                                                     const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs);
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+        return embed_nd(ids, bs, theta, axes_dim);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
@@ -506,8 +363,8 @@ namespace Rope {
 
                         GGML_ASSERT(i < grid_h * grid_w);
 
-                        ids[i][0] = static_cast<float>(ih + iy);
-                        ids[i][1] = static_cast<float>(iw + ix);
+                        ids[i][0] = ih + iy;
+                        ids[i][1] = iw + ix;
                         index++;
                     }
                 }
@@ -524,7 +381,7 @@ namespace Rope {
                                                         int theta,
                                                         const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_qwen2vl_ids(grid_h, grid_w, merge_size, window_index);
-        return embed_nd(ids, 1, static_cast<float>(theta), axes_dim);
+        return embed_nd(ids, 1, theta, axes_dim);
     }
 
     __STATIC_INLINE__ int bound_mod(int a, int m) {
@@ -571,39 +428,15 @@ namespace Rope {
                                                         const std::vector<ggml_tensor*>& ref_latents,
                                                         bool increase_ref_index,
                                                         int theta,
-                                                        bool circular_h,
-                                                        bool circular_w,
                                                         const std::vector<int>& axes_dim) {
         std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int pad_h = (patch_size - (h % patch_size)) % patch_size;
-            int pad_w = (patch_size - (w % patch_size)) % patch_size;
-            int h_len = (h + pad_h) / patch_size;
-            int w_len = (w + pad_w) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                size_t pos_len = ids.size() / bs;
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-                size_t cursor     = context_len + bound_mod(context_len, seq_multi_of);  // skip text (and its padding)
-                size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-            }
-        }
-
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+        return embed_nd(ids, bs, theta, axes_dim);
     }
 
-    __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
-                                              ggml_tensor* x,
-                                              ggml_tensor* pe,
-                                              bool rope_interleaved = true) {
+    __STATIC_INLINE__ struct ggml_tensor* apply_rope(struct ggml_context* ctx,
+                                                     struct ggml_tensor* x,
+                                                     struct ggml_tensor* pe,
+                                                     bool rope_interleaved = true) {
         // x: [N, L, n_head, d_head]
         // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
         int64_t d_head = x->ne[0];
@@ -641,21 +474,21 @@ namespace Rope {
         return x_out;
     }
 
-    __STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx,
-                                             ggml_tensor* q,
-                                             ggml_tensor* k,
-                                             ggml_tensor* v,
-                                             ggml_tensor* pe,
-                                             ggml_tensor* mask,
-                                             float kv_scale        = 1.0f,
-                                             bool rope_interleaved = true) {
+    __STATIC_INLINE__ struct ggml_tensor* attention(GGMLRunnerContext* ctx,
+                                                    struct ggml_tensor* q,
+                                                    struct ggml_tensor* k,
+                                                    struct ggml_tensor* v,
+                                                    struct ggml_tensor* pe,
+                                                    struct ggml_tensor* mask,
+                                                    float kv_scale        = 1.0f,
+                                                    bool rope_interleaved = true) {
         // q,k,v: [N, L, n_head, d_head]
         // pe: [L, d_head/2, 2, 2]
         // return: [N, L, n_head*d_head]
         q = apply_rope(ctx->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, v->ne[1], mask, false, true, ctx->flash_attn_enabled, kv_scale);  // [N, L, n_head*d_head]
         return x;
     }
 };  // namespace Rope

src/anima.hpp

@@ -1,683 +0,0 @@
-#ifndef __ANIMA_HPP__
-#define __ANIMA_HPP__
-
-#include <cmath>
-#include <memory>
-#include <utility>
-#include <vector>
-
-#include "common_block.hpp"
-#include "flux.hpp"
-#include "rope.hpp"
-
-namespace Anima {
-    constexpr int ANIMA_GRAPH_SIZE = 65536;
-
-    __STATIC_INLINE__ ggml_tensor* apply_gate(ggml_context* ctx,
-                                              ggml_tensor* x,
-                                              ggml_tensor* gate) {
-        gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]);  // [N, 1, C]
-        return ggml_mul(ctx, x, gate);
-    }
-
-    struct XEmbedder : public GGMLBlock {
-    public:
-        XEmbedder(int64_t in_dim, int64_t out_dim) {
-            blocks["proj.1"] = std::make_shared<Linear>(in_dim, out_dim, false);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj.1"]);
-            return proj->forward(ctx, x);
-        }
-    };
-
-    struct TimestepEmbedder : public GGMLBlock {
-    public:
-        TimestepEmbedder(int64_t in_dim, int64_t out_dim) {
-            blocks["1.linear_1"] = std::make_shared<Linear>(in_dim, in_dim, false);
-            blocks["1.linear_2"] = std::make_shared<Linear>(in_dim, out_dim, false);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_1"]);
-            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_2"]);
-
-            x = linear_1->forward(ctx, x);
-            x = ggml_silu_inplace(ctx->ggml_ctx, x);
-            x = linear_2->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct AdaLayerNormZero : public GGMLBlock {
-    protected:
-        int64_t in_features;
-
-    public:
-        AdaLayerNormZero(int64_t in_features, int64_t hidden_features = 256)
-            : in_features(in_features) {
-            blocks["norm"] = std::make_shared<LayerNorm>(in_features, 1e-6f, false, false);
-            blocks["1"]    = std::make_shared<Linear>(in_features, hidden_features, false);
-            blocks["2"]    = std::make_shared<Linear>(hidden_features, 3 * in_features, false);
-        }
-
-        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                      ggml_tensor* hidden_states,
-                                                      ggml_tensor* embedded_timestep,
-                                                      ggml_tensor* temb = nullptr) {
-            auto norm     = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
-            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
-            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
-
-            auto emb = ggml_silu(ctx->ggml_ctx, embedded_timestep);
-            emb      = linear_1->forward(ctx, emb);
-            emb      = linear_2->forward(ctx, emb);  // [N, 3*C]
-
-            if (temb != nullptr) {
-                emb = ggml_add(ctx->ggml_ctx, emb, temb);
-            }
-
-            auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 3, 0);
-            auto shift      = emb_chunks[0];
-            auto scale      = emb_chunks[1];
-            auto gate       = emb_chunks[2];
-
-            auto x = norm->forward(ctx, hidden_states);
-            x      = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
-
-            return {x, gate};
-        }
-    };
-
-    struct AdaLayerNorm : public GGMLBlock {
-    protected:
-        int64_t embedding_dim;
-
-    public:
-        AdaLayerNorm(int64_t in_features, int64_t hidden_features = 256)
-            : embedding_dim(in_features) {
-            blocks["norm"] = std::make_shared<LayerNorm>(in_features, 1e-6f, false, false);
-            blocks["1"]    = std::make_shared<Linear>(in_features, hidden_features, false);
-            blocks["2"]    = std::make_shared<Linear>(hidden_features, 2 * in_features, false);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* hidden_states,
-                             ggml_tensor* embedded_timestep,
-                             ggml_tensor* temb = nullptr) {
-            auto norm     = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
-            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
-            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
-
-            auto emb = ggml_silu(ctx->ggml_ctx, embedded_timestep);
-            emb      = linear_1->forward(ctx, emb);
-            emb      = linear_2->forward(ctx, emb);  // [N, 2*C]
-
-            if (temb != nullptr) {
-                auto temb_2c = ggml_view_2d(ctx->ggml_ctx, temb, 2 * embedding_dim, temb->ne[1], temb->nb[1], 0);
-                emb          = ggml_add(ctx->ggml_ctx, emb, temb_2c);
-            }
-
-            auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
-            auto shift      = emb_chunks[0];
-            auto scale      = emb_chunks[1];
-
-            auto x = norm->forward(ctx, hidden_states);
-            x      = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
-            return x;
-        }
-    };
-
-    struct AnimaAttention : public GGMLBlock {
-    protected:
-        int64_t num_heads;
-        int64_t head_dim;
-        std::string out_proj_name;
-
-    public:
-        AnimaAttention(int64_t query_dim,
-                       int64_t context_dim,
-                       int64_t num_heads,
-                       int64_t head_dim,
-                       const std::string& out_proj_name = "output_proj")
-            : num_heads(num_heads), head_dim(head_dim), out_proj_name(out_proj_name) {
-            int64_t inner_dim = num_heads * head_dim;
-
-            blocks["q_proj"]            = std::make_shared<Linear>(query_dim, inner_dim, false);
-            blocks["k_proj"]            = std::make_shared<Linear>(context_dim, inner_dim, false);
-            blocks["v_proj"]            = std::make_shared<Linear>(context_dim, inner_dim, false);
-            blocks["q_norm"]            = std::make_shared<RMSNorm>(head_dim, 1e-6f);
-            blocks["k_norm"]            = std::make_shared<RMSNorm>(head_dim, 1e-6f);
-            blocks[this->out_proj_name] = std::make_shared<Linear>(inner_dim, query_dim, false);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* hidden_states,
-                             ggml_tensor* encoder_hidden_states = nullptr,
-                             ggml_tensor* pe_q                  = nullptr,
-                             ggml_tensor* pe_k                  = nullptr) {
-            if (encoder_hidden_states == nullptr) {
-                encoder_hidden_states = hidden_states;
-            }
-
-            auto q_proj   = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
-            auto k_proj   = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
-            auto v_proj   = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
-            auto q_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
-            auto k_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
-            auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[out_proj_name]);
-
-            auto q = q_proj->forward(ctx, hidden_states);
-            auto k = k_proj->forward(ctx, encoder_hidden_states);
-            auto v = v_proj->forward(ctx, encoder_hidden_states);
-
-            int64_t N   = q->ne[2];
-            int64_t L_q = q->ne[1];
-            int64_t L_k = k->ne[1];
-
-            auto q4 = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, L_q, N);  // [N, L_q, H, D]
-            auto k4 = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, L_k, N);  // [N, L_k, H, D]
-            auto v4 = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, L_k, N);  // [N, L_k, H, D]
-
-            q4 = q_norm->forward(ctx, q4);
-            k4 = k_norm->forward(ctx, k4);
-
-            ggml_tensor* attn_out = nullptr;
-            if (pe_q != nullptr || pe_k != nullptr) {
-                if (pe_q == nullptr) {
-                    pe_q = pe_k;
-                }
-                if (pe_k == nullptr) {
-                    pe_k = pe_q;
-                }
-                auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
-                auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
-                attn_out    = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                                     ctx->backend,
-                                                     q_rope,
-                                                     k_rope,
-                                                     v4,
-                                                     num_heads,
-                                                     nullptr,
-                                                     true,
-                                                     ctx->flash_attn_enabled);
-            } else {
-                auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
-                auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
-                attn_out    = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                                     ctx->backend,
-                                                     q_flat,
-                                                     k_flat,
-                                                     v,
-                                                     num_heads,
-                                                     nullptr,
-                                                     false,
-                                                     ctx->flash_attn_enabled);
-            }
-
-            return out_proj->forward(ctx, attn_out);
-        }
-    };
-
-    struct AnimaMLP : public GGMLBlock {
-    public:
-        AnimaMLP(int64_t dim, int64_t hidden_dim) {
-            blocks["layer1"] = std::make_shared<Linear>(dim, hidden_dim, false);
-            blocks["layer2"] = std::make_shared<Linear>(hidden_dim, dim, false);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto layer1 = std::dynamic_pointer_cast<Linear>(blocks["layer1"]);
-            auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["layer2"]);
-
-            x = layer1->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
-            x = layer2->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct AdapterMLP : public GGMLBlock {
-    public:
-        AdapterMLP(int64_t dim, int64_t hidden_dim) {
-            blocks["0"] = std::make_shared<Linear>(dim, hidden_dim, true);
-            blocks["2"] = std::make_shared<Linear>(hidden_dim, dim, true);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto layer0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
-            auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
-
-            x = layer0->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
-            x = layer2->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct LLMAdapterBlock : public GGMLBlock {
-    public:
-        LLMAdapterBlock(int64_t model_dim = 1024, int64_t source_dim = 1024, int64_t num_heads = 16, int64_t head_dim = 64) {
-            blocks["norm_self_attn"]  = std::make_shared<RMSNorm>(model_dim, 1e-6f);
-            blocks["self_attn"]       = std::make_shared<AnimaAttention>(model_dim, model_dim, num_heads, head_dim, "o_proj");
-            blocks["norm_cross_attn"] = std::make_shared<RMSNorm>(model_dim, 1e-6f);
-            blocks["cross_attn"]      = std::make_shared<AnimaAttention>(model_dim, source_dim, num_heads, head_dim, "o_proj");
-            blocks["norm_mlp"]        = std::make_shared<RMSNorm>(model_dim, 1e-6f);
-            blocks["mlp"]             = std::make_shared<AdapterMLP>(model_dim, model_dim * 4);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* context,
-                             ggml_tensor* target_pe,
-                             ggml_tensor* context_pe) {
-            auto norm_self_attn  = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_self_attn"]);
-            auto self_attn       = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
-            auto norm_cross_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_cross_attn"]);
-            auto cross_attn      = std::dynamic_pointer_cast<AnimaAttention>(blocks["cross_attn"]);
-            auto norm_mlp        = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_mlp"]);
-            auto mlp             = std::dynamic_pointer_cast<AdapterMLP>(blocks["mlp"]);
-
-            auto h = norm_self_attn->forward(ctx, x);
-            h      = self_attn->forward(ctx, h, nullptr, target_pe, target_pe);
-            x      = ggml_add(ctx->ggml_ctx, x, h);
-
-            h = norm_cross_attn->forward(ctx, x);
-            h = cross_attn->forward(ctx, h, context, target_pe, context_pe);
-            x = ggml_add(ctx->ggml_ctx, x, h);
-
-            h = norm_mlp->forward(ctx, x);
-            h = mlp->forward(ctx, h);
-            x = ggml_add(ctx->ggml_ctx, x, h);
-
-            return x;
-        }
-    };
-
-    struct LLMAdapter : public GGMLBlock {
-    protected:
-        int num_layers;
-
-    public:
-        LLMAdapter(int64_t source_dim = 1024,
-                   int64_t target_dim = 1024,
-                   int64_t model_dim  = 1024,
-                   int num_layers     = 6,
-                   int num_heads      = 16)
-            : num_layers(num_layers) {
-            int64_t head_dim = model_dim / num_heads;
-
-            blocks["embed"] = std::make_shared<Embedding>(32128, target_dim);
-            for (int i = 0; i < num_layers; i++) {
-                blocks["blocks." + std::to_string(i)] =
-                    std::make_shared<LLMAdapterBlock>(model_dim, source_dim, num_heads, head_dim);
-            }
-            blocks["out_proj"] = std::make_shared<Linear>(model_dim, target_dim, true);
-            blocks["norm"]     = std::make_shared<RMSNorm>(target_dim, 1e-6f);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* source_hidden_states,
-                             ggml_tensor* target_input_ids,
-                             ggml_tensor* target_pe,
-                             ggml_tensor* source_pe) {
-            GGML_ASSERT(target_input_ids != nullptr);
-            if (ggml_n_dims(target_input_ids) == 1) {
-                target_input_ids = ggml_reshape_2d(ctx->ggml_ctx, target_input_ids, target_input_ids->ne[0], 1);
-            }
-
-            auto embed    = std::dynamic_pointer_cast<Embedding>(blocks["embed"]);
-            auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
-            auto norm     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
-
-            auto x = embed->forward(ctx, target_input_ids);  // [N, target_len, target_dim]
-
-            for (int i = 0; i < num_layers; i++) {
-                auto block = std::dynamic_pointer_cast<LLMAdapterBlock>(blocks["blocks." + std::to_string(i)]);
-                x          = block->forward(ctx, x, source_hidden_states, target_pe, source_pe);
-            }
-
-            x = out_proj->forward(ctx, x);
-            x = norm->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct TransformerBlock : public GGMLBlock {
-    public:
-        TransformerBlock(int64_t hidden_size,
-                         int64_t text_embed_dim,
-                         int64_t num_heads,
-                         int64_t head_dim,
-                         int64_t mlp_ratio      = 4,
-                         int64_t adaln_lora_dim = 256) {
-            blocks["adaln_modulation_self_attn"]  = std::make_shared<AdaLayerNormZero>(hidden_size, adaln_lora_dim);
-            blocks["self_attn"]                   = std::make_shared<AnimaAttention>(hidden_size, hidden_size, num_heads, head_dim);
-            blocks["adaln_modulation_cross_attn"] = std::make_shared<AdaLayerNormZero>(hidden_size, adaln_lora_dim);
-            blocks["cross_attn"]                  = std::make_shared<AnimaAttention>(hidden_size, text_embed_dim, num_heads, head_dim);
-            blocks["adaln_modulation_mlp"]        = std::make_shared<AdaLayerNormZero>(hidden_size, adaln_lora_dim);
-            blocks["mlp"]                         = std::make_shared<AnimaMLP>(hidden_size, hidden_size * mlp_ratio);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* hidden_states,
-                             ggml_tensor* encoder_hidden_states,
-                             ggml_tensor* embedded_timestep,
-                             ggml_tensor* temb,
-                             ggml_tensor* image_pe) {
-            auto norm1 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_self_attn"]);
-            auto attn1 = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
-            auto norm2 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_cross_attn"]);
-            auto attn2 = std::dynamic_pointer_cast<AnimaAttention>(blocks["cross_attn"]);
-            auto norm3 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_mlp"]);
-            auto mlp   = std::dynamic_pointer_cast<AnimaMLP>(blocks["mlp"]);
-
-            auto [normed1, gate1] = norm1->forward(ctx, hidden_states, embedded_timestep, temb);
-            auto h                = attn1->forward(ctx, normed1, nullptr, image_pe, image_pe);
-            hidden_states         = ggml_add(ctx->ggml_ctx, hidden_states, apply_gate(ctx->ggml_ctx, h, gate1));
-
-            auto [normed2, gate2] = norm2->forward(ctx, hidden_states, embedded_timestep, temb);
-            h                     = attn2->forward(ctx, normed2, encoder_hidden_states, nullptr, nullptr);
-            hidden_states         = ggml_add(ctx->ggml_ctx, hidden_states, apply_gate(ctx->ggml_ctx, h, gate2));
-
-            auto [normed3, gate3] = norm3->forward(ctx, hidden_states, embedded_timestep, temb);
-            h                     = mlp->forward(ctx, normed3);
-            hidden_states         = ggml_add(ctx->ggml_ctx, hidden_states, apply_gate(ctx->ggml_ctx, h, gate3));
-
-            return hidden_states;
-        }
-    };
-
-    struct FinalLayer : public GGMLBlock {
-    protected:
-        int64_t hidden_size;
-        int64_t patch_size;
-        int64_t out_channels;
-
-    public:
-        FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels)
-            : hidden_size(hidden_size), patch_size(patch_size), out_channels(out_channels) {
-            blocks["adaln_modulation"] = std::make_shared<AdaLayerNorm>(hidden_size, 256);
-            blocks["linear"]           = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, false);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* hidden_states,
-                             ggml_tensor* embedded_timestep,
-                             ggml_tensor* temb) {
-            auto adaln  = std::dynamic_pointer_cast<AdaLayerNorm>(blocks["adaln_modulation"]);
-            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
-
-            hidden_states = adaln->forward(ctx, hidden_states, embedded_timestep, temb);
-            hidden_states = linear->forward(ctx, hidden_states);
-            return hidden_states;
-        }
-    };
-
-    struct AnimaNet : public GGMLBlock {
-    public:
-        int64_t in_channels       = 16;
-        int64_t out_channels      = 16;
-        int64_t hidden_size       = 2048;
-        int64_t text_embed_dim    = 1024;
-        int64_t num_heads         = 16;
-        int64_t head_dim          = 128;
-        int patch_size            = 2;
-        int64_t num_layers        = 28;
-        std::vector<int> axes_dim = {44, 42, 42};
-        int theta                 = 10000;
-
-    public:
-        AnimaNet() = default;
-        explicit AnimaNet(int64_t num_layers)
-            : num_layers(num_layers) {
-            blocks["x_embedder"]       = std::make_shared<XEmbedder>((in_channels + 1) * patch_size * patch_size, hidden_size);
-            blocks["t_embedder"]       = std::make_shared<TimestepEmbedder>(hidden_size, hidden_size * 3);
-            blocks["t_embedding_norm"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
-            for (int i = 0; i < num_layers; i++) {
-                blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(hidden_size,
-                                                                                           text_embed_dim,
-                                                                                           num_heads,
-                                                                                           head_dim);
-            }
-            blocks["final_layer"] = std::make_shared<FinalLayer>(hidden_size, patch_size, out_channels);
-            blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* timestep,
-                             ggml_tensor* encoder_hidden_states,
-                             ggml_tensor* image_pe,
-                             ggml_tensor* t5_ids       = nullptr,
-                             ggml_tensor* t5_weights   = nullptr,
-                             ggml_tensor* adapter_q_pe = nullptr,
-                             ggml_tensor* adapter_k_pe = nullptr) {
-            GGML_ASSERT(x->ne[3] == 1);
-
-            auto x_embedder       = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
-            auto t_embedder       = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
-            auto t_embedding_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["t_embedding_norm"]);
-            auto final_layer      = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
-            auto llm_adapter      = std::dynamic_pointer_cast<LLMAdapter>(blocks["llm_adapter"]);
-
-            int64_t W = x->ne[0];
-            int64_t H = x->ne[1];
-
-            auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
-            x                 = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2);  // [N, C + 1, H, W]
-
-            x = DiT::pad_and_patchify(ctx, x, patch_size, patch_size);  // [N, h*w, (C+1)*ph*pw]
-
-            x = x_embedder->forward(ctx, x);
-
-            auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(hidden_size));
-            auto temb              = t_embedder->forward(ctx, timestep_proj);
-            auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj);
-
-            if (t5_ids != nullptr) {
-                auto adapted_context = llm_adapter->forward(ctx, encoder_hidden_states, t5_ids, adapter_q_pe, adapter_k_pe);
-                if (t5_weights != nullptr) {
-                    auto w = t5_weights;
-                    if (ggml_n_dims(w) == 1) {
-                        w = ggml_reshape_3d(ctx->ggml_ctx, w, 1, w->ne[0], 1);
-                    }
-                    w               = ggml_repeat_4d(ctx->ggml_ctx, w, adapted_context->ne[0], adapted_context->ne[1], adapted_context->ne[2], 1);
-                    adapted_context = ggml_mul(ctx->ggml_ctx, adapted_context, w);
-                }
-                if (adapted_context->ne[1] < 512) {
-                    auto pad_ctx    = ggml_ext_zeros(ctx->ggml_ctx,
-                                                     adapted_context->ne[0],
-                                                     512 - adapted_context->ne[1],
-                                                     adapted_context->ne[2],
-                                                     1);
-                    adapted_context = ggml_concat(ctx->ggml_ctx, adapted_context, pad_ctx, 1);
-                } else if (adapted_context->ne[1] > 512) {
-                    adapted_context = ggml_ext_slice(ctx->ggml_ctx, adapted_context, 1, 0, 512);
-                }
-                encoder_hidden_states = adapted_context;
-            }
-
-            for (int i = 0; i < num_layers; i++) {
-                auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
-                x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
-            }
-
-            x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
-
-            x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, false);  // [N, C, H, W]
-
-            return x;
-        }
-    };
-
-    struct AnimaRunner : public GGMLRunner {
-    public:
-        std::vector<float> image_pe_vec;
-        std::vector<float> adapter_q_pe_vec;
-        std::vector<float> adapter_k_pe_vec;
-        AnimaNet net;
-
-        AnimaRunner(ggml_backend_t backend,
-                    bool offload_params_to_cpu,
-                    const String2TensorStorage& tensor_storage_map = {},
-                    const std::string prefix                       = "model.diffusion_model")
-            : GGMLRunner(backend, offload_params_to_cpu) {
-            int64_t num_layers    = 0;
-            std::string layer_tag = prefix + ".net.blocks.";
-            for (const auto& kv : tensor_storage_map) {
-                const std::string& tensor_name = kv.first;
-                size_t pos                     = tensor_name.find(layer_tag);
-                if (pos == std::string::npos) {
-                    continue;
-                }
-                size_t start = pos + layer_tag.size();
-                size_t end   = tensor_name.find('.', start);
-                if (end == std::string::npos) {
-                    continue;
-                }
-                int64_t layer_id = atoll(tensor_name.substr(start, end - start).c_str());
-                num_layers       = std::max(num_layers, layer_id + 1);
-            }
-            if (num_layers <= 0) {
-                num_layers = 28;
-            }
-            LOG_INFO("anima net layers: %" PRId64, num_layers);
-
-            net = AnimaNet(num_layers);
-            net.init(params_ctx, tensor_storage_map, prefix + ".net");
-        }
-
-        std::string get_desc() override {
-            return "anima";
-        }
-
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-            net.get_param_tensors(tensors, prefix + ".net");
-        }
-
-        static std::vector<float> gen_1d_rope_pe_vec(int64_t seq_len, int dim, float theta = 10000.f) {
-            std::vector<float> pos(seq_len);
-            for (int64_t i = 0; i < seq_len; i++) {
-                pos[i] = static_cast<float>(i);
-            }
-            auto rope_emb = Rope::rope(pos, dim, theta);
-            return Rope::flatten(rope_emb);
-        }
-
-        static float calc_ntk_factor(float extrapolation_ratio, int axis_dim) {
-            if (extrapolation_ratio == 1.0f || axis_dim <= 2) {
-                return 1.0f;
-            }
-            return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
-        }
-
-        static std::vector<float> gen_anima_image_pe_vec(int bs,
-                                                         int h,
-                                                         int w,
-                                                         int patch_size,
-                                                         int theta,
-                                                         const std::vector<int>& axes_dim,
-                                                         float h_extrapolation_ratio,
-                                                         float w_extrapolation_ratio,
-                                                         float t_extrapolation_ratio) {
-            static const std::vector<ggml_tensor*> empty_ref_latents;
-            auto ids = Rope::gen_flux_ids(h,
-                                          w,
-                                          patch_size,
-                                          bs,
-                                          static_cast<int>(axes_dim.size()),
-                                          0,
-                                          {},
-                                          empty_ref_latents,
-                                          false,
-                                          1.0f);
-
-            std::vector<float> axis_thetas = {
-                static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
-                static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
-                static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
-            };
-            return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
-        }
-
-        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
-                                 const sd::Tensor<float>& timesteps_tensor,
-                                 const sd::Tensor<float>& context_tensor    = {},
-                                 const sd::Tensor<int32_t>& t5_ids_tensor   = {},
-                                 const sd::Tensor<float>& t5_weights_tensor = {}) {
-            ggml_tensor* x          = make_input(x_tensor);
-            ggml_tensor* timesteps  = make_input(timesteps_tensor);
-            ggml_tensor* context    = make_optional_input(context_tensor);
-            ggml_tensor* t5_ids     = make_optional_input(t5_ids_tensor);
-            ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
-            GGML_ASSERT(x->ne[3] == 1);
-            ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
-
-            int64_t pad_h = (net.patch_size - x->ne[1] % net.patch_size) % net.patch_size;
-            int64_t pad_w = (net.patch_size - x->ne[0] % net.patch_size) % net.patch_size;
-            int64_t h_pad = x->ne[1] + pad_h;
-            int64_t w_pad = x->ne[0] + pad_w;
-
-            image_pe_vec          = gen_anima_image_pe_vec(1,
-                                                           static_cast<int>(h_pad),
-                                                           static_cast<int>(w_pad),
-                                                           static_cast<int>(net.patch_size),
-                                                           net.theta,
-                                                           net.axes_dim,
-                                                           4.0f,
-                                                           4.0f,
-                                                           1.0f);
-            int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (net.head_dim / 2));
-            auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, net.head_dim / 2, image_pos_len);
-            set_backend_tensor_data(image_pe, image_pe_vec.data());
-
-            ggml_tensor* adapter_q_pe = nullptr;
-            ggml_tensor* adapter_k_pe = nullptr;
-            if (t5_ids != nullptr) {
-                int64_t target_len = t5_ids->ne[0];
-                int64_t source_len = context->ne[1];
-
-                adapter_q_pe_vec = gen_1d_rope_pe_vec(target_len, 64, 10000.f);
-                adapter_k_pe_vec = gen_1d_rope_pe_vec(source_len, 64, 10000.f);
-
-                int64_t target_pos_len = static_cast<int64_t>(adapter_q_pe_vec.size()) / (2 * 2 * 32);
-                int64_t source_pos_len = static_cast<int64_t>(adapter_k_pe_vec.size()) / (2 * 2 * 32);
-
-                adapter_q_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, target_pos_len);
-                adapter_k_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, source_pos_len);
-                set_backend_tensor_data(adapter_q_pe, adapter_q_pe_vec.data());
-                set_backend_tensor_data(adapter_k_pe, adapter_k_pe_vec.data());
-            }
-
-            auto runner_ctx = get_context();
-            auto out        = net.forward(&runner_ctx,
-                                          x,
-                                          timesteps,
-                                          context,
-                                          image_pe,
-                                          t5_ids,
-                                          t5_weights,
-                                          adapter_q_pe,
-                                          adapter_k_pe);
-
-            ggml_build_forward_expand(gf, out);
-            return gf;
-        }
-
-        sd::Tensor<float> compute(int n_threads,
-                                  const sd::Tensor<float>& x,
-                                  const sd::Tensor<float>& timesteps,
-                                  const sd::Tensor<float>& context    = {},
-                                  const sd::Tensor<int32_t>& t5_ids   = {},
-                                  const sd::Tensor<float>& t5_weights = {}) {
-            auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, t5_ids, t5_weights);
-            };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
-        }
-    };
-}  // namespace Anima
-
-#endif  // __ANIMA_HPP__

src/cache_dit.hpp

@@ -1,896 +0,0 @@
-#ifndef __CACHE_DIT_HPP__
-#define __CACHE_DIT_HPP__
-
-#include <algorithm>
-#include <cmath>
-#include <limits>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "condition_cache_utils.hpp"
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
-
-struct DBCacheConfig {
-    bool enabled                        = false;
-    int Fn_compute_blocks               = 8;
-    int Bn_compute_blocks               = 0;
-    float residual_diff_threshold       = 0.08f;
-    int max_warmup_steps                = 8;
-    int max_cached_steps                = -1;
-    int max_continuous_cached_steps     = -1;
-    float max_accumulated_residual_diff = -1.0f;
-    std::vector<int> steps_computation_mask;
-    bool scm_policy_dynamic = true;
-};
-
-struct TaylorSeerConfig {
-    bool enabled            = false;
-    int n_derivatives       = 1;
-    int max_warmup_steps    = 2;
-    int skip_interval_steps = 1;
-};
-
-struct CacheDitConfig {
-    DBCacheConfig dbcache;
-    TaylorSeerConfig taylorseer;
-    int double_Fn_blocks = -1;
-    int double_Bn_blocks = -1;
-    int single_Fn_blocks = -1;
-    int single_Bn_blocks = -1;
-};
-
-struct TaylorSeerState {
-    int n_derivatives      = 1;
-    int current_step       = -1;
-    int last_computed_step = -1;
-    std::vector<std::vector<float>> dY_prev;
-    std::vector<std::vector<float>> dY_current;
-
-    void init(int n_deriv, size_t hidden_size) {
-        n_derivatives = n_deriv;
-        int order     = n_derivatives + 1;
-        dY_prev.resize(order);
-        dY_current.resize(order);
-        for (int i = 0; i < order; i++) {
-            dY_prev[i].clear();
-            dY_current[i].clear();
-        }
-        current_step       = -1;
-        last_computed_step = -1;
-    }
-
-    void reset() {
-        for (auto& v : dY_prev)
-            v.clear();
-        for (auto& v : dY_current)
-            v.clear();
-        current_step       = -1;
-        last_computed_step = -1;
-    }
-
-    bool can_approximate() const {
-        return last_computed_step >= n_derivatives && !dY_prev.empty() && !dY_prev[0].empty();
-    }
-
-    void update_derivatives(const float* Y, size_t size, int step) {
-        int order = n_derivatives + 1;
-        dY_prev   = dY_current;
-        dY_current[0].resize(size);
-        for (size_t i = 0; i < size; i++) {
-            dY_current[0][i] = Y[i];
-        }
-
-        int window = step - last_computed_step;
-        if (window <= 0)
-            window = 1;
-
-        for (int d = 0; d < n_derivatives; d++) {
-            if (!dY_prev[d].empty() && dY_prev[d].size() == size) {
-                dY_current[d + 1].resize(size);
-                for (size_t i = 0; i < size; i++) {
-                    dY_current[d + 1][i] = (dY_current[d][i] - dY_prev[d][i]) / static_cast<float>(window);
-                }
-            } else {
-                dY_current[d + 1].clear();
-            }
-        }
-
-        current_step       = step;
-        last_computed_step = step;
-    }
-
-    void approximate(float* output, size_t size, int target_step) const {
-        if (!can_approximate() || dY_prev[0].size() != size) {
-            return;
-        }
-
-        int elapsed = target_step - last_computed_step;
-        if (elapsed <= 0)
-            elapsed = 1;
-
-        std::fill(output, output + size, 0.0f);
-        float factorial = 1.0f;
-        int order       = static_cast<int>(dY_prev.size());
-
-        for (int o = 0; o < order; o++) {
-            if (dY_prev[o].empty() || dY_prev[o].size() != size)
-                continue;
-            if (o > 0)
-                factorial *= static_cast<float>(o);
-            float coeff = ::powf(static_cast<float>(elapsed), static_cast<float>(o)) / factorial;
-            for (size_t i = 0; i < size; i++) {
-                output[i] += coeff * dY_prev[o][i];
-            }
-        }
-    }
-};
-
-struct BlockCacheEntry {
-    std::vector<float> residual_img;
-    std::vector<float> residual_txt;
-    std::vector<float> residual;
-    std::vector<float> prev_img;
-    std::vector<float> prev_txt;
-    std::vector<float> prev_output;
-    bool has_prev = false;
-};
-
-struct CacheDitState {
-    CacheDitConfig config;
-    bool initialized = false;
-
-    int total_double_blocks = 0;
-    int total_single_blocks = 0;
-    size_t hidden_size      = 0;
-
-    int current_step     = -1;
-    int total_steps      = 0;
-    int warmup_remaining = 0;
-    std::vector<int> cached_steps;
-    int continuous_cached_steps     = 0;
-    float accumulated_residual_diff = 0.0f;
-
-    std::vector<BlockCacheEntry> double_block_cache;
-    std::vector<BlockCacheEntry> single_block_cache;
-
-    std::vector<float> Fn_residual_img;
-    std::vector<float> Fn_residual_txt;
-    std::vector<float> prev_Fn_residual_img;
-    std::vector<float> prev_Fn_residual_txt;
-    bool has_prev_Fn_residual = false;
-
-    std::vector<float> Bn_buffer_img;
-    std::vector<float> Bn_buffer_txt;
-    std::vector<float> Bn_buffer;
-    bool has_Bn_buffer = false;
-
-    TaylorSeerState taylor_state;
-
-    bool can_cache_this_step  = false;
-    bool is_caching_this_step = false;
-
-    int total_blocks_computed = 0;
-    int total_blocks_cached   = 0;
-
-    void init(const CacheDitConfig& cfg, int num_double_blocks, int num_single_blocks, size_t h_size) {
-        config              = cfg;
-        total_double_blocks = num_double_blocks;
-        total_single_blocks = num_single_blocks;
-        hidden_size         = h_size;
-
-        initialized = cfg.dbcache.enabled || cfg.taylorseer.enabled;
-
-        if (!initialized)
-            return;
-
-        warmup_remaining = cfg.dbcache.max_warmup_steps;
-        double_block_cache.resize(total_double_blocks);
-        single_block_cache.resize(total_single_blocks);
-
-        if (cfg.taylorseer.enabled) {
-            taylor_state.init(cfg.taylorseer.n_derivatives, h_size);
-        }
-
-        reset_runtime();
-    }
-
-    void reset_runtime() {
-        current_step     = -1;
-        total_steps      = 0;
-        warmup_remaining = config.dbcache.max_warmup_steps;
-        cached_steps.clear();
-        continuous_cached_steps   = 0;
-        accumulated_residual_diff = 0.0f;
-
-        for (auto& entry : double_block_cache) {
-            entry.residual_img.clear();
-            entry.residual_txt.clear();
-            entry.prev_img.clear();
-            entry.prev_txt.clear();
-            entry.has_prev = false;
-        }
-
-        for (auto& entry : single_block_cache) {
-            entry.residual.clear();
-            entry.prev_output.clear();
-            entry.has_prev = false;
-        }
-
-        Fn_residual_img.clear();
-        Fn_residual_txt.clear();
-        prev_Fn_residual_img.clear();
-        prev_Fn_residual_txt.clear();
-        has_prev_Fn_residual = false;
-
-        Bn_buffer_img.clear();
-        Bn_buffer_txt.clear();
-        Bn_buffer.clear();
-        has_Bn_buffer = false;
-
-        taylor_state.reset();
-
-        can_cache_this_step  = false;
-        is_caching_this_step = false;
-
-        total_blocks_computed = 0;
-        total_blocks_cached   = 0;
-    }
-
-    bool enabled() const {
-        return initialized && (config.dbcache.enabled || config.taylorseer.enabled);
-    }
-
-    void begin_step(int step_index, float sigma = 0.0f) {
-        if (!enabled())
-            return;
-        if (step_index == current_step)
-            return;
-
-        current_step = step_index;
-        total_steps++;
-
-        bool in_warmup = warmup_remaining > 0;
-        if (in_warmup) {
-            warmup_remaining--;
-        }
-
-        bool scm_allows_cache = true;
-        if (!config.dbcache.steps_computation_mask.empty()) {
-            if (step_index < static_cast<int>(config.dbcache.steps_computation_mask.size())) {
-                scm_allows_cache = (config.dbcache.steps_computation_mask[step_index] == 0);
-                if (!config.dbcache.scm_policy_dynamic && scm_allows_cache) {
-                    can_cache_this_step  = true;
-                    is_caching_this_step = false;
-                    return;
-                }
-            }
-        }
-
-        bool max_cached_ok = (config.dbcache.max_cached_steps < 0) ||
-                             (static_cast<int>(cached_steps.size()) < config.dbcache.max_cached_steps);
-
-        bool max_cont_ok = (config.dbcache.max_continuous_cached_steps < 0) ||
-                           (continuous_cached_steps < config.dbcache.max_continuous_cached_steps);
-
-        bool accum_ok = (config.dbcache.max_accumulated_residual_diff < 0.0f) ||
-                        (accumulated_residual_diff < config.dbcache.max_accumulated_residual_diff);
-
-        can_cache_this_step  = !in_warmup && scm_allows_cache && max_cached_ok && max_cont_ok && accum_ok && has_prev_Fn_residual;
-        is_caching_this_step = false;
-    }
-
-    void end_step(bool was_cached) {
-        if (was_cached) {
-            cached_steps.push_back(current_step);
-            continuous_cached_steps++;
-        } else {
-            continuous_cached_steps = 0;
-        }
-    }
-
-    static float calculate_residual_diff(const float* prev, const float* curr, size_t size) {
-        if (size == 0)
-            return 0.0f;
-
-        float sum_diff = 0.0f;
-        float sum_abs  = 0.0f;
-
-        for (size_t i = 0; i < size; i++) {
-            sum_diff += std::fabs(prev[i] - curr[i]);
-            sum_abs += std::fabs(prev[i]);
-        }
-
-        return sum_diff / (sum_abs + 1e-6f);
-    }
-
-    static float calculate_residual_diff(const std::vector<float>& prev, const std::vector<float>& curr) {
-        if (prev.size() != curr.size() || prev.empty())
-            return 1.0f;
-        return calculate_residual_diff(prev.data(), curr.data(), prev.size());
-    }
-
-    int get_double_Fn_blocks() const {
-        return (config.double_Fn_blocks >= 0) ? config.double_Fn_blocks : config.dbcache.Fn_compute_blocks;
-    }
-
-    int get_double_Bn_blocks() const {
-        return (config.double_Bn_blocks >= 0) ? config.double_Bn_blocks : config.dbcache.Bn_compute_blocks;
-    }
-
-    int get_single_Fn_blocks() const {
-        return (config.single_Fn_blocks >= 0) ? config.single_Fn_blocks : config.dbcache.Fn_compute_blocks;
-    }
-
-    int get_single_Bn_blocks() const {
-        return (config.single_Bn_blocks >= 0) ? config.single_Bn_blocks : config.dbcache.Bn_compute_blocks;
-    }
-
-    bool is_Fn_double_block(int block_idx) const {
-        return block_idx < get_double_Fn_blocks();
-    }
-
-    bool is_Bn_double_block(int block_idx) const {
-        int Bn = get_double_Bn_blocks();
-        return Bn > 0 && block_idx >= (total_double_blocks - Bn);
-    }
-
-    bool is_Mn_double_block(int block_idx) const {
-        return !is_Fn_double_block(block_idx) && !is_Bn_double_block(block_idx);
-    }
-
-    bool is_Fn_single_block(int block_idx) const {
-        return block_idx < get_single_Fn_blocks();
-    }
-
-    bool is_Bn_single_block(int block_idx) const {
-        int Bn = get_single_Bn_blocks();
-        return Bn > 0 && block_idx >= (total_single_blocks - Bn);
-    }
-
-    bool is_Mn_single_block(int block_idx) const {
-        return !is_Fn_single_block(block_idx) && !is_Bn_single_block(block_idx);
-    }
-
-    void store_Fn_residual(const float* img, const float* txt, size_t img_size, size_t txt_size, const float* input_img, const float* input_txt) {
-        Fn_residual_img.resize(img_size);
-        Fn_residual_txt.resize(txt_size);
-
-        for (size_t i = 0; i < img_size; i++) {
-            Fn_residual_img[i] = img[i] - input_img[i];
-        }
-        for (size_t i = 0; i < txt_size; i++) {
-            Fn_residual_txt[i] = txt[i] - input_txt[i];
-        }
-    }
-
-    bool check_cache_decision() {
-        if (!can_cache_this_step) {
-            is_caching_this_step = false;
-            return false;
-        }
-
-        if (!has_prev_Fn_residual || prev_Fn_residual_img.empty()) {
-            is_caching_this_step = false;
-            return false;
-        }
-
-        float diff_img = calculate_residual_diff(prev_Fn_residual_img, Fn_residual_img);
-        float diff_txt = calculate_residual_diff(prev_Fn_residual_txt, Fn_residual_txt);
-        float diff     = (diff_img + diff_txt) / 2.0f;
-
-        if (diff < config.dbcache.residual_diff_threshold) {
-            is_caching_this_step = true;
-            accumulated_residual_diff += diff;
-            return true;
-        }
-
-        is_caching_this_step = false;
-        return false;
-    }
-
-    void update_prev_Fn_residual() {
-        prev_Fn_residual_img = Fn_residual_img;
-        prev_Fn_residual_txt = Fn_residual_txt;
-        has_prev_Fn_residual = !prev_Fn_residual_img.empty();
-    }
-
-    void store_double_block_residual(int block_idx, const float* img, const float* txt, size_t img_size, size_t txt_size, const float* prev_img, const float* prev_txt) {
-        if (block_idx < 0 || block_idx >= static_cast<int>(double_block_cache.size()))
-            return;
-
-        BlockCacheEntry& entry = double_block_cache[block_idx];
-
-        entry.residual_img.resize(img_size);
-        entry.residual_txt.resize(txt_size);
-        for (size_t i = 0; i < img_size; i++) {
-            entry.residual_img[i] = img[i] - prev_img[i];
-        }
-        for (size_t i = 0; i < txt_size; i++) {
-            entry.residual_txt[i] = txt[i] - prev_txt[i];
-        }
-
-        entry.prev_img.resize(img_size);
-        entry.prev_txt.resize(txt_size);
-        for (size_t i = 0; i < img_size; i++) {
-            entry.prev_img[i] = img[i];
-        }
-        for (size_t i = 0; i < txt_size; i++) {
-            entry.prev_txt[i] = txt[i];
-        }
-        entry.has_prev = true;
-    }
-
-    void apply_double_block_cache(int block_idx, float* img, float* txt, size_t img_size, size_t txt_size) {
-        if (block_idx < 0 || block_idx >= static_cast<int>(double_block_cache.size()))
-            return;
-
-        const BlockCacheEntry& entry = double_block_cache[block_idx];
-        if (entry.residual_img.size() != img_size || entry.residual_txt.size() != txt_size)
-            return;
-
-        for (size_t i = 0; i < img_size; i++) {
-            img[i] += entry.residual_img[i];
-        }
-        for (size_t i = 0; i < txt_size; i++) {
-            txt[i] += entry.residual_txt[i];
-        }
-
-        total_blocks_cached++;
-    }
-
-    void store_single_block_residual(int block_idx, const float* output, size_t size, const float* input) {
-        if (block_idx < 0 || block_idx >= static_cast<int>(single_block_cache.size()))
-            return;
-
-        BlockCacheEntry& entry = single_block_cache[block_idx];
-
-        entry.residual.resize(size);
-        for (size_t i = 0; i < size; i++) {
-            entry.residual[i] = output[i] - input[i];
-        }
-
-        entry.prev_output.resize(size);
-        for (size_t i = 0; i < size; i++) {
-            entry.prev_output[i] = output[i];
-        }
-        entry.has_prev = true;
-    }
-
-    void apply_single_block_cache(int block_idx, float* output, size_t size) {
-        if (block_idx < 0 || block_idx >= static_cast<int>(single_block_cache.size()))
-            return;
-
-        const BlockCacheEntry& entry = single_block_cache[block_idx];
-        if (entry.residual.size() != size)
-            return;
-
-        for (size_t i = 0; i < size; i++) {
-            output[i] += entry.residual[i];
-        }
-
-        total_blocks_cached++;
-    }
-
-    void store_Bn_buffer(const float* img, const float* txt, size_t img_size, size_t txt_size, const float* Bn_start_img, const float* Bn_start_txt) {
-        Bn_buffer_img.resize(img_size);
-        Bn_buffer_txt.resize(txt_size);
-
-        for (size_t i = 0; i < img_size; i++) {
-            Bn_buffer_img[i] = img[i] - Bn_start_img[i];
-        }
-        for (size_t i = 0; i < txt_size; i++) {
-            Bn_buffer_txt[i] = txt[i] - Bn_start_txt[i];
-        }
-        has_Bn_buffer = true;
-    }
-
-    void apply_Bn_buffer(float* img, float* txt, size_t img_size, size_t txt_size) {
-        if (!has_Bn_buffer)
-            return;
-        if (Bn_buffer_img.size() != img_size || Bn_buffer_txt.size() != txt_size)
-            return;
-
-        for (size_t i = 0; i < img_size; i++) {
-            img[i] += Bn_buffer_img[i];
-        }
-        for (size_t i = 0; i < txt_size; i++) {
-            txt[i] += Bn_buffer_txt[i];
-        }
-    }
-
-    void taylor_update(const float* hidden_state, size_t size) {
-        if (!config.taylorseer.enabled)
-            return;
-        taylor_state.update_derivatives(hidden_state, size, current_step);
-    }
-
-    bool taylor_can_approximate() const {
-        return config.taylorseer.enabled && taylor_state.can_approximate();
-    }
-
-    void taylor_approximate(float* output, size_t size) {
-        if (!config.taylorseer.enabled)
-            return;
-        taylor_state.approximate(output, size, current_step);
-    }
-
-    bool should_use_taylor_this_step() const {
-        if (!config.taylorseer.enabled)
-            return false;
-        if (current_step < config.taylorseer.max_warmup_steps)
-            return false;
-
-        int interval = config.taylorseer.skip_interval_steps;
-        if (interval <= 0)
-            interval = 1;
-
-        return (current_step % (interval + 1)) != 0;
-    }
-
-    void log_metrics() const {
-        if (!enabled())
-            return;
-
-        int total_blocks  = total_blocks_computed + total_blocks_cached;
-        float cache_ratio = (total_blocks > 0) ? (static_cast<float>(total_blocks_cached) / total_blocks * 100.0f) : 0.0f;
-
-        float step_cache_ratio = (total_steps > 0) ? (static_cast<float>(cached_steps.size()) / total_steps * 100.0f) : 0.0f;
-
-        LOG_INFO("CacheDIT: steps_cached=%zu/%d (%.1f%%), blocks_cached=%d/%d (%.1f%%), accum_diff=%.4f",
-                 cached_steps.size(), total_steps, step_cache_ratio,
-                 total_blocks_cached, total_blocks, cache_ratio,
-                 accumulated_residual_diff);
-    }
-
-    std::string get_summary() const {
-        char buf[256];
-        snprintf(buf, sizeof(buf),
-                 "CacheDIT[thresh=%.2f]: cached %zu/%d steps, %d/%d blocks",
-                 config.dbcache.residual_diff_threshold,
-                 cached_steps.size(), total_steps,
-                 total_blocks_cached, total_blocks_computed + total_blocks_cached);
-        return std::string(buf);
-    }
-};
-
-inline std::vector<int> parse_scm_mask(const std::string& mask_str) {
-    std::vector<int> mask;
-    if (mask_str.empty())
-        return mask;
-
-    size_t pos   = 0;
-    size_t start = 0;
-    while ((pos = mask_str.find(',', start)) != std::string::npos) {
-        std::string token = mask_str.substr(start, pos - start);
-        mask.push_back(std::stoi(token));
-        start = pos + 1;
-    }
-    if (start < mask_str.length()) {
-        mask.push_back(std::stoi(mask_str.substr(start)));
-    }
-
-    return mask;
-}
-
-inline std::vector<int> generate_scm_mask(
-    const std::vector<int>& compute_bins,
-    const std::vector<int>& cache_bins,
-    int total_steps) {
-    std::vector<int> mask;
-    size_t c_idx = 0, cache_idx = 0;
-
-    while (static_cast<int>(mask.size()) < total_steps) {
-        if (c_idx < compute_bins.size()) {
-            for (int i = 0; i < compute_bins[c_idx] && static_cast<int>(mask.size()) < total_steps; i++) {
-                mask.push_back(1);
-            }
-            c_idx++;
-        }
-        if (cache_idx < cache_bins.size()) {
-            for (int i = 0; i < cache_bins[cache_idx] && static_cast<int>(mask.size()) < total_steps; i++) {
-                mask.push_back(0);
-            }
-            cache_idx++;
-        }
-        if (c_idx >= compute_bins.size() && cache_idx >= cache_bins.size())
-            break;
-    }
-
-    if (!mask.empty()) {
-        mask.back() = 1;
-    }
-
-    return mask;
-}
-
-inline void parse_dbcache_options(const std::string& opts, DBCacheConfig& cfg) {
-    if (opts.empty())
-        return;
-
-    int Fn = 8, Bn = 0, warmup = 8, max_cached = -1, max_cont = -1;
-    float thresh = 0.08f;
-
-    sscanf(opts.c_str(), "%d,%d,%f,%d,%d,%d",
-           &Fn, &Bn, &thresh, &warmup, &max_cached, &max_cont);
-
-    cfg.Fn_compute_blocks           = Fn;
-    cfg.Bn_compute_blocks           = Bn;
-    cfg.residual_diff_threshold     = thresh;
-    cfg.max_warmup_steps            = warmup;
-    cfg.max_cached_steps            = max_cached;
-    cfg.max_continuous_cached_steps = max_cont;
-}
-
-inline void parse_taylorseer_options(const std::string& opts, TaylorSeerConfig& cfg) {
-    if (opts.empty())
-        return;
-
-    int n_deriv = 1, warmup = 2, interval = 1;
-    sscanf(opts.c_str(), "%d,%d,%d", &n_deriv, &warmup, &interval);
-
-    cfg.n_derivatives       = n_deriv;
-    cfg.max_warmup_steps    = warmup;
-    cfg.skip_interval_steps = interval;
-}
-
-struct CacheDitConditionState {
-    DBCacheConfig config;
-    TaylorSeerConfig taylor_config;
-    bool initialized = false;
-
-    int current_step_index = -1;
-    bool step_active       = false;
-    bool skip_current_step = false;
-    bool initial_step      = true;
-    int warmup_remaining   = 0;
-    std::vector<int> cached_steps;
-    int continuous_cached_steps     = 0;
-    float accumulated_residual_diff = 0.0f;
-    int total_steps_skipped         = 0;
-
-    const void* anchor_condition = nullptr;
-
-    struct CacheEntry {
-        std::vector<float> diff;
-        std::vector<float> prev_input;
-        std::vector<float> prev_output;
-        bool has_prev = false;
-    };
-    std::unordered_map<const void*, CacheEntry> cache_diffs;
-
-    TaylorSeerState taylor_state;
-
-    float start_sigma = std::numeric_limits<float>::max();
-    float end_sigma   = 0.0f;
-
-    void reset_runtime() {
-        current_step_index = -1;
-        step_active        = false;
-        skip_current_step  = false;
-        initial_step       = true;
-        warmup_remaining   = config.max_warmup_steps;
-        cached_steps.clear();
-        continuous_cached_steps   = 0;
-        accumulated_residual_diff = 0.0f;
-        total_steps_skipped       = 0;
-        anchor_condition          = nullptr;
-        cache_diffs.clear();
-        taylor_state.reset();
-    }
-
-    void init(const DBCacheConfig& dbcfg, const TaylorSeerConfig& tcfg) {
-        config        = dbcfg;
-        taylor_config = tcfg;
-        initialized   = dbcfg.enabled || tcfg.enabled;
-        reset_runtime();
-
-        if (taylor_config.enabled) {
-            taylor_state.init(taylor_config.n_derivatives, 0);
-        }
-    }
-
-    void set_sigmas(const std::vector<float>& sigmas) {
-        if (!initialized || sigmas.size() < 2)
-            return;
-
-        float start_percent = 0.15f;
-        float end_percent   = 0.95f;
-
-        size_t n_steps    = sigmas.size() - 1;
-        size_t start_step = static_cast<size_t>(start_percent * n_steps);
-        size_t end_step   = static_cast<size_t>(end_percent * n_steps);
-
-        if (start_step >= n_steps)
-            start_step = n_steps - 1;
-        if (end_step >= n_steps)
-            end_step = n_steps - 1;
-
-        start_sigma = sigmas[start_step];
-        end_sigma   = sigmas[end_step];
-
-        if (start_sigma < end_sigma) {
-            std::swap(start_sigma, end_sigma);
-        }
-    }
-
-    bool enabled() const {
-        return initialized && (config.enabled || taylor_config.enabled);
-    }
-
-    void begin_step(int step_index, float sigma) {
-        if (!enabled())
-            return;
-        if (step_index == current_step_index)
-            return;
-
-        current_step_index = step_index;
-        skip_current_step  = false;
-        step_active        = false;
-
-        if (sigma > start_sigma)
-            return;
-        if (!(sigma > end_sigma))
-            return;
-
-        step_active = true;
-
-        if (warmup_remaining > 0) {
-            warmup_remaining--;
-            return;
-        }
-
-        if (!config.steps_computation_mask.empty()) {
-            if (step_index < static_cast<int>(config.steps_computation_mask.size())) {
-                if (config.steps_computation_mask[step_index] == 1) {
-                    return;
-                }
-            }
-        }
-
-        if (config.max_cached_steps >= 0 &&
-            static_cast<int>(cached_steps.size()) >= config.max_cached_steps) {
-            return;
-        }
-
-        if (config.max_continuous_cached_steps >= 0 &&
-            continuous_cached_steps >= config.max_continuous_cached_steps) {
-            return;
-        }
-    }
-
-    bool step_is_active() const {
-        return enabled() && step_active;
-    }
-
-    bool is_step_skipped() const {
-        return enabled() && step_active && skip_current_step;
-    }
-
-    bool has_cache(const void* cond) const {
-        auto it = cache_diffs.find(cond);
-        return it != cache_diffs.end() && !it->second.diff.empty();
-    }
-
-    void update_cache(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
-        CacheEntry& entry = cache_diffs[cond];
-        if (!sd::store_condition_cache_diff(&entry.diff, input, output)) {
-            entry.prev_input.clear();
-            entry.prev_output.clear();
-            entry.has_prev = false;
-            return;
-        }
-
-        size_t size              = static_cast<size_t>(output.numel());
-        const float* input_data  = input.data();
-        const float* output_data = output.data();
-        entry.prev_input.resize(size);
-        entry.prev_output.resize(size);
-        for (size_t i = 0; i < size; i++) {
-            entry.prev_input[i]  = input_data[i];
-            entry.prev_output[i] = output_data[i];
-        }
-        entry.has_prev = true;
-    }
-
-    void apply_cache(const void* cond,
-                     const sd::Tensor<float>& input,
-                     sd::Tensor<float>* output) {
-        auto it = cache_diffs.find(cond);
-        if (it == cache_diffs.end() || it->second.diff.empty())
-            return;
-        sd::apply_condition_cache_diff(it->second.diff, input, output);
-    }
-
-    bool before_condition(const void* cond, const sd::Tensor<float>& input, sd::Tensor<float>* output, float sigma, int step_index) {
-        if (!enabled() || step_index < 0)
-            return false;
-
-        if (step_index != current_step_index) {
-            begin_step(step_index, sigma);
-        }
-
-        if (!step_active)
-            return false;
-
-        if (initial_step) {
-            anchor_condition = cond;
-            initial_step     = false;
-        }
-
-        bool is_anchor = (cond == anchor_condition);
-
-        if (skip_current_step) {
-            if (has_cache(cond)) {
-                apply_cache(cond, input, output);
-                return true;
-            }
-            return false;
-        }
-
-        if (!is_anchor)
-            return false;
-
-        auto it = cache_diffs.find(cond);
-        if (it == cache_diffs.end() || !it->second.has_prev)
-            return false;
-
-        size_t ne = static_cast<size_t>(input.numel());
-        if (it->second.prev_input.size() != ne)
-            return false;
-
-        const float* input_data = input.data();
-        float diff              = CacheDitState::calculate_residual_diff(
-                         it->second.prev_input.data(), input_data, ne);
-
-        float effective_threshold = config.residual_diff_threshold;
-        if (config.Fn_compute_blocks > 0) {
-            float fn_confidence = 1.0f + 0.02f * (config.Fn_compute_blocks - 8);
-            fn_confidence       = std::max(0.5f, std::min(2.0f, fn_confidence));
-            effective_threshold *= fn_confidence;
-        }
-        if (config.Bn_compute_blocks > 0) {
-            float bn_quality = 1.0f - 0.03f * config.Bn_compute_blocks;
-            bn_quality       = std::max(0.5f, std::min(1.0f, bn_quality));
-            effective_threshold *= bn_quality;
-        }
-
-        if (diff < effective_threshold) {
-            skip_current_step = true;
-            total_steps_skipped++;
-            cached_steps.push_back(current_step_index);
-            continuous_cached_steps++;
-            accumulated_residual_diff += diff;
-            apply_cache(cond, input, output);
-            return true;
-        }
-
-        continuous_cached_steps = 0;
-        return false;
-    }
-
-    void after_condition(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
-        if (!step_is_active())
-            return;
-
-        update_cache(cond, input, output);
-
-        if (cond == anchor_condition && taylor_config.enabled) {
-            taylor_state.update_derivatives(output.data(), static_cast<size_t>(output.numel()), current_step_index);
-        }
-    }
-
-    void log_metrics() const {
-        if (!enabled())
-            return;
-
-        LOG_INFO("CacheDIT: steps_skipped=%d/%d (%.1f%%), accum_residual_diff=%.4f",
-                 total_steps_skipped,
-                 current_step_index + 1,
-                 (current_step_index > 0) ? (100.0f * total_steps_skipped / (current_step_index + 1)) : 0.0f,
-                 accumulated_residual_diff);
-    }
-};
-
-#endif

src/common_dit.hpp

@@ -1,108 +0,0 @@
-#ifndef __COMMON_DIT_HPP__
-#define __COMMON_DIT_HPP__
-
-#include "ggml_extend.hpp"
-
-namespace DiT {
-    inline ggml_tensor* patchify(ggml_context* ctx,
-                                 ggml_tensor* x,
-                                 int pw,
-                                 int ph,
-                                 bool patch_last = true) {
-        // x: [N, C, H, W]
-        // return: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
-        int64_t N = x->ne[3];
-        int64_t C = x->ne[2];
-        int64_t H = x->ne[1];
-        int64_t W = x->ne[0];
-        int64_t h = H / ph;
-        int64_t w = W / pw;
-
-        GGML_ASSERT(h * ph == H && w * pw == W);
-
-        x = ggml_reshape_4d(ctx, x, pw, w, ph, h * C * N);     // [N*C*h, ph, w, pw]
-        x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*h, w, ph, pw]
-        x = ggml_reshape_4d(ctx, x, pw * ph, w * h, C, N);     // [N, C, h*w, ph*pw]
-        if (patch_last) {
-            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, h*w, C, ph*pw]
-            x = ggml_reshape_3d(ctx, x, pw * ph * C, w * h, N);    // [N, h*w, C*ph*pw]
-        } else {
-            x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, h*w, C, ph*pw]
-            x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N);              // [N, h*w, ph*pw*C]
-        }
-        return x;
-    }
-
-    inline ggml_tensor* unpatchify(ggml_context* ctx,
-                                   ggml_tensor* x,
-                                   int64_t h,
-                                   int64_t w,
-                                   int ph,
-                                   int pw,
-                                   bool patch_last = true) {
-        // x: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
-        // return: [N, C, H, W]
-        int64_t N = x->ne[2];
-        int64_t C = x->ne[0] / ph / pw;
-        int64_t H = h * ph;
-        int64_t W = w * pw;
-
-        GGML_ASSERT(C * ph * pw == x->ne[0]);
-
-        if (patch_last) {
-            x = ggml_reshape_4d(ctx, x, pw * ph, C, w * h, N);     // [N, h*w, C, ph*pw]
-            x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, C, h*w, ph*pw]
-        } else {
-            x = ggml_reshape_4d(ctx, x, C, pw * ph, w * h, N);     // [N, h*w, ph*pw, C]
-            x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3));  // [N, C, h*w, ph*pw]
-        }
-
-        x = ggml_reshape_4d(ctx, x, pw, ph, w, h * C * N);     // [N*C*h, w, ph, pw]
-        x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N*C*h, ph, w, pw]
-        x = ggml_reshape_4d(ctx, x, W, H, C, N);               // [N, C, h*ph, w*pw]
-
-        return x;
-    }
-
-    inline ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
-                                          ggml_tensor* x,
-                                          int ph,
-                                          int pw) {
-        int64_t W = x->ne[0];
-        int64_t H = x->ne[1];
-
-        int pad_h = (ph - H % ph) % ph;
-        int pad_w = (pw - W % pw) % pw;
-        x         = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-        return x;
-    }
-
-    inline ggml_tensor* pad_and_patchify(GGMLRunnerContext* ctx,
-                                         ggml_tensor* x,
-                                         int ph,
-                                         int pw,
-                                         bool patch_last = true) {
-        x = pad_to_patch_size(ctx, x, ph, pw);
-        x = patchify(ctx->ggml_ctx, x, ph, pw, patch_last);
-        return x;
-    }
-
-    inline ggml_tensor* unpatchify_and_crop(ggml_context* ctx,
-                                            ggml_tensor* x,
-                                            int64_t H,
-                                            int64_t W,
-                                            int ph,
-                                            int pw,
-                                            bool patch_last = true) {
-        int pad_h = (ph - H % ph) % ph;
-        int pad_w = (pw - W % pw) % pw;
-        int64_t h = ((H + pad_h) / ph);
-        int64_t w = ((W + pad_w) / pw);
-        x         = unpatchify(ctx, x, h, w, ph, pw, patch_last);  // [N, C, H + pad_h, W + pad_w]
-        x         = ggml_ext_slice(ctx, x, 1, 0, H);               // [N, C, H, W + pad_w]
-        x         = ggml_ext_slice(ctx, x, 0, 0, W);               // [N, C, H, W]
-        return x;
-    }
-}  // namespace DiT
-
-#endif  // __COMMON_DIT_HPP__

src/condition_cache_utils.hpp

@@ -1,64 +0,0 @@
-#ifndef __CONDITION_CACHE_UTILS_HPP__
-#define __CONDITION_CACHE_UTILS_HPP__
-
-#include <vector>
-
-#include "tensor.hpp"
-
-namespace sd {
-
-    inline bool store_condition_cache_diff(std::vector<float>* diff,
-                                           const sd::Tensor<float>& input,
-                                           const sd::Tensor<float>& output) {
-        if (diff == nullptr || input.empty() || output.empty()) {
-            return false;
-        }
-
-        size_t input_size  = static_cast<size_t>(input.numel());
-        size_t output_size = static_cast<size_t>(output.numel());
-        if (input_size == 0 || input_size != output_size) {
-            diff->clear();
-            return false;
-        }
-
-        const float* input_data  = input.data();
-        const float* output_data = output.data();
-        if (input_data == nullptr || output_data == nullptr) {
-            diff->clear();
-            return false;
-        }
-
-        diff->resize(output_size);
-        for (size_t i = 0; i < output_size; ++i) {
-            (*diff)[i] = output_data[i] - input_data[i];
-        }
-        return true;
-    }
-
-    inline bool apply_condition_cache_diff(const std::vector<float>& diff,
-                                           const sd::Tensor<float>& input,
-                                           sd::Tensor<float>* output) {
-        if (output == nullptr || input.empty() || diff.empty()) {
-            return false;
-        }
-
-        size_t input_size = static_cast<size_t>(input.numel());
-        if (input_size == 0 || diff.size() != input_size) {
-            return false;
-        }
-
-        *output            = input;
-        float* output_data = output->data();
-        if (output_data == nullptr) {
-            return false;
-        }
-
-        for (size_t i = 0; i < input_size; ++i) {
-            output_data[i] += diff[i];
-        }
-        return true;
-    }
-
-}  // namespace sd
-
-#endif  // __CONDITION_CACHE_UTILS_HPP__

src/diffusion_model.hpp

@@ -1,519 +0,0 @@
-#ifndef __DIFFUSION_MODEL_H__
-#define __DIFFUSION_MODEL_H__
-
-#include <optional>
-#include "anima.hpp"
-#include "flux.hpp"
-#include "mmdit.hpp"
-#include "qwen_image.hpp"
-#include "tensor_ggml.hpp"
-#include "unet.hpp"
-#include "wan.hpp"
-#include "z_image.hpp"
-
-struct DiffusionParams {
-    const sd::Tensor<float>* x                        = nullptr;
-    const sd::Tensor<float>* timesteps                = nullptr;
-    const sd::Tensor<float>* context                  = nullptr;
-    const sd::Tensor<float>* c_concat                 = nullptr;
-    const sd::Tensor<float>* y                        = nullptr;
-    const sd::Tensor<int32_t>* t5_ids                 = nullptr;
-    const sd::Tensor<float>* t5_weights               = nullptr;
-    const sd::Tensor<float>* guidance                 = nullptr;
-    const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
-    bool increase_ref_index                           = false;
-    int num_video_frames                              = -1;
-    const std::vector<sd::Tensor<float>>* controls    = nullptr;
-    float control_strength                            = 0.f;
-    const sd::Tensor<float>* vace_context             = nullptr;
-    float vace_strength                               = 1.f;
-    const std::vector<int>* skip_layers               = nullptr;
-};
-
-template <typename T>
-static inline const sd::Tensor<T>& tensor_or_empty(const sd::Tensor<T>* tensor) {
-    static const sd::Tensor<T> kEmpty;
-    return tensor != nullptr ? *tensor : kEmpty;
-}
-
-struct DiffusionModel {
-    virtual std::string get_desc()                                               = 0;
-    virtual sd::Tensor<float> compute(int n_threads,
-                                      const DiffusionParams& diffusion_params)   = 0;
-    virtual void alloc_params_buffer()                                           = 0;
-    virtual void free_params_buffer()                                            = 0;
-    virtual void free_compute_buffer()                                           = 0;
-    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
-    virtual size_t get_params_buffer_size()                                      = 0;
-    virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
-    virtual int64_t get_adm_in_channels()                            = 0;
-    virtual void set_flash_attention_enabled(bool enabled)           = 0;
-    virtual void set_circular_axes(bool circular_x, bool circular_y) = 0;
-};
-
-struct UNetModel : public DiffusionModel {
-    UNetModelRunner unet;
-
-    UNetModel(ggml_backend_t backend,
-              bool offload_params_to_cpu,
-              const String2TensorStorage& tensor_storage_map = {},
-              SDVersion version                              = VERSION_SD1)
-        : unet(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", version) {
-    }
-
-    std::string get_desc() override {
-        return unet.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        unet.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        unet.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        unet.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        unet.get_param_tensors(tensors, "model.diffusion_model");
-    }
-
-    size_t get_params_buffer_size() override {
-        return unet.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        unet.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return unet.unet.adm_in_channels;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        unet.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        unet.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_controls;
-        return unet.compute(n_threads,
-                            *diffusion_params.x,
-                            *diffusion_params.timesteps,
-                            tensor_or_empty(diffusion_params.context),
-                            tensor_or_empty(diffusion_params.c_concat),
-                            tensor_or_empty(diffusion_params.y),
-                            diffusion_params.num_video_frames,
-                            diffusion_params.controls ? *diffusion_params.controls : empty_controls,
-                            diffusion_params.control_strength);
-    }
-};
-
-struct MMDiTModel : public DiffusionModel {
-    MMDiTRunner mmdit;
-
-    MMDiTModel(ggml_backend_t backend,
-               bool offload_params_to_cpu,
-               const String2TensorStorage& tensor_storage_map = {})
-        : mmdit(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model") {
-    }
-
-    std::string get_desc() override {
-        return mmdit.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        mmdit.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        mmdit.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        mmdit.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        mmdit.get_param_tensors(tensors, "model.diffusion_model");
-    }
-
-    size_t get_params_buffer_size() override {
-        return mmdit.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        mmdit.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768 + 1280;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        mmdit.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        mmdit.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<int> empty_skip_layers;
-        return mmdit.compute(n_threads,
-                             *diffusion_params.x,
-                             *diffusion_params.timesteps,
-                             tensor_or_empty(diffusion_params.context),
-                             tensor_or_empty(diffusion_params.y),
-                             diffusion_params.skip_layers ? *diffusion_params.skip_layers : empty_skip_layers);
-    }
-};
-
-struct FluxModel : public DiffusionModel {
-    Flux::FluxRunner flux;
-
-    FluxModel(ggml_backend_t backend,
-              bool offload_params_to_cpu,
-              const String2TensorStorage& tensor_storage_map = {},
-              SDVersion version                              = VERSION_FLUX,
-              bool use_mask                                  = false)
-        : flux(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", version, use_mask) {
-    }
-
-    std::string get_desc() override {
-        return flux.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        flux.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        flux.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        flux.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        flux.get_param_tensors(tensors, "model.diffusion_model");
-    }
-
-    size_t get_params_buffer_size() override {
-        return flux.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        flux.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        flux.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        flux.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_ref_latents;
-        static const std::vector<int> empty_skip_layers;
-        return flux.compute(n_threads,
-                            *diffusion_params.x,
-                            *diffusion_params.timesteps,
-                            tensor_or_empty(diffusion_params.context),
-                            tensor_or_empty(diffusion_params.c_concat),
-                            tensor_or_empty(diffusion_params.y),
-                            tensor_or_empty(diffusion_params.guidance),
-                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                            diffusion_params.increase_ref_index,
-                            diffusion_params.skip_layers ? *diffusion_params.skip_layers : empty_skip_layers);
-    }
-};
-
-struct AnimaModel : public DiffusionModel {
-    std::string prefix;
-    Anima::AnimaRunner anima;
-
-    AnimaModel(ggml_backend_t backend,
-               bool offload_params_to_cpu,
-               const String2TensorStorage& tensor_storage_map = {},
-               const std::string prefix                       = "model.diffusion_model")
-        : prefix(prefix), anima(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
-    }
-
-    std::string get_desc() override {
-        return anima.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        anima.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        anima.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        anima.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        anima.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return anima.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        anima.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        anima.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        anima.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        return anima.compute(n_threads,
-                             *diffusion_params.x,
-                             *diffusion_params.timesteps,
-                             tensor_or_empty(diffusion_params.context),
-                             tensor_or_empty(diffusion_params.t5_ids),
-                             tensor_or_empty(diffusion_params.t5_weights));
-    }
-};
-
-struct WanModel : public DiffusionModel {
-    std::string prefix;
-    WAN::WanRunner wan;
-
-    WanModel(ggml_backend_t backend,
-             bool offload_params_to_cpu,
-             const String2TensorStorage& tensor_storage_map = {},
-             const std::string prefix                       = "model.diffusion_model",
-             SDVersion version                              = VERSION_WAN2)
-        : prefix(prefix), wan(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) {
-    }
-
-    std::string get_desc() override {
-        return wan.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        wan.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        wan.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        wan.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        wan.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return wan.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        wan.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        wan.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        wan.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        return wan.compute(n_threads,
-                           *diffusion_params.x,
-                           *diffusion_params.timesteps,
-                           tensor_or_empty(diffusion_params.context),
-                           tensor_or_empty(diffusion_params.y),
-                           tensor_or_empty(diffusion_params.c_concat),
-                           sd::Tensor<float>(),
-                           tensor_or_empty(diffusion_params.vace_context),
-                           diffusion_params.vace_strength);
-    }
-};
-
-struct QwenImageModel : public DiffusionModel {
-    std::string prefix;
-    Qwen::QwenImageRunner qwen_image;
-
-    QwenImageModel(ggml_backend_t backend,
-                   bool offload_params_to_cpu,
-                   const String2TensorStorage& tensor_storage_map = {},
-                   const std::string prefix                       = "model.diffusion_model",
-                   SDVersion version                              = VERSION_QWEN_IMAGE,
-                   bool zero_cond_t                               = false)
-        : prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version, zero_cond_t) {
-    }
-
-    std::string get_desc() override {
-        return qwen_image.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        qwen_image.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        qwen_image.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        qwen_image.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        qwen_image.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return qwen_image.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        qwen_image.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        qwen_image.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        qwen_image.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_ref_latents;
-        return qwen_image.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,
-                                  true);
-    }
-};
-
-struct ZImageModel : public DiffusionModel {
-    std::string prefix;
-    ZImage::ZImageRunner z_image;
-
-    ZImageModel(ggml_backend_t backend,
-                bool offload_params_to_cpu,
-                const String2TensorStorage& tensor_storage_map = {},
-                const std::string prefix                       = "model.diffusion_model",
-                SDVersion version                              = VERSION_Z_IMAGE)
-        : prefix(prefix), z_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) {
-    }
-
-    std::string get_desc() override {
-        return z_image.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        z_image.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        z_image.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        z_image.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        z_image.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return z_image.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        z_image.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        z_image.set_flash_attention_enabled(enabled);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        z_image.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_ref_latents;
-        return z_image.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,
-                               true);
-    }
-};
-
-#endif

src/preprocessing.hpp

@@ -1,278 +0,0 @@
-#ifndef __PREPROCESSING_HPP__
-#define __PREPROCESSING_HPP__
-
-#include <cmath>
-#include <limits>
-
-#include "ggml_extend.hpp"
-
-#define M_PI_ 3.14159265358979323846f
-
-static inline int64_t preprocessing_offset_4d(const sd::Tensor<float>& tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
-    const auto& shape = tensor.shape();
-    int64_t n0        = shape.size() > 0 ? shape[0] : 1;
-    int64_t n1        = shape.size() > 1 ? shape[1] : 1;
-    int64_t n2        = shape.size() > 2 ? shape[2] : 1;
-    return ((i3 * n2 + i2) * n1 + i1) * n0 + i0;
-}
-
-static inline float preprocessing_get_4d(const sd::Tensor<float>& tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
-    return tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))];
-}
-
-static inline void preprocessing_set_4d(sd::Tensor<float>& tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
-    tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))] = value;
-}
-
-static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t image) {
-    sd::Tensor<float> tensor({static_cast<int64_t>(image.width), static_cast<int64_t>(image.height), static_cast<int64_t>(image.channel), 1});
-    for (uint32_t y = 0; y < image.height; ++y) {
-        for (uint32_t x = 0; x < image.width; ++x) {
-            for (uint32_t c = 0; c < image.channel; ++c) {
-                preprocessing_set_4d(tensor, sd_image_get_f32(image, x, y, c), x, y, c, 0);
-            }
-        }
-    }
-    return tensor;
-}
-
-static inline void preprocessing_tensor_to_sd_image(const sd::Tensor<float>& tensor, uint8_t* image_data) {
-    GGML_ASSERT(tensor.dim() == 4);
-    GGML_ASSERT(tensor.shape()[3] == 1);
-    GGML_ASSERT(image_data != nullptr);
-
-    int width   = static_cast<int>(tensor.shape()[0]);
-    int height  = static_cast<int>(tensor.shape()[1]);
-    int channel = static_cast<int>(tensor.shape()[2]);
-    for (int y = 0; y < height; ++y) {
-        for (int x = 0; x < width; ++x) {
-            for (int c = 0; c < channel; ++c) {
-                float value                               = preprocessing_get_4d(tensor, x, y, c, 0);
-                value                                     = std::min(1.0f, std::max(0.0f, value));
-                image_data[(y * width + x) * channel + c] = static_cast<uint8_t>(std::round(value * 255.0f));
-            }
-        }
-    }
-}
-
-static inline sd::Tensor<float> gaussian_kernel_tensor(int kernel_size) {
-    sd::Tensor<float> kernel({kernel_size, kernel_size, 1, 1});
-    int ks_mid   = kernel_size / 2;
-    float sigma  = 1.4f;
-    float normal = 1.f / (2.0f * M_PI_ * std::pow(sigma, 2.0f));
-    for (int y = 0; y < kernel_size; ++y) {
-        float gx = static_cast<float>(-ks_mid + y);
-        for (int x = 0; x < kernel_size; ++x) {
-            float gy = static_cast<float>(-ks_mid + x);
-            float k  = std::exp(-((gx * gx + gy * gy) / (2.0f * std::pow(sigma, 2.0f)))) * normal;
-            preprocessing_set_4d(kernel, k, x, y, 0, 0);
-        }
-    }
-    return kernel;
-}
-
-static inline sd::Tensor<float> convolve_tensor(const sd::Tensor<float>& input, const sd::Tensor<float>& kernel, int padding) {
-    GGML_ASSERT(input.dim() == 4);
-    GGML_ASSERT(kernel.dim() == 4);
-    GGML_ASSERT(input.shape()[3] == 1);
-    GGML_ASSERT(kernel.shape()[2] == 1);
-    GGML_ASSERT(kernel.shape()[3] == 1);
-
-    sd::Tensor<float> output(input.shape());
-    int64_t width    = input.shape()[0];
-    int64_t height   = input.shape()[1];
-    int64_t channels = input.shape()[2];
-    int64_t kernel_w = kernel.shape()[0];
-    int64_t kernel_h = kernel.shape()[1];
-
-    for (int64_t c = 0; c < channels; ++c) {
-        for (int64_t y = 0; y < height; ++y) {
-            for (int64_t x = 0; x < width; ++x) {
-                float sum = 0.0f;
-                for (int64_t ky = 0; ky < kernel_h; ++ky) {
-                    int64_t iy = y + ky - padding;
-                    if (iy < 0 || iy >= height) {
-                        continue;
-                    }
-                    for (int64_t kx = 0; kx < kernel_w; ++kx) {
-                        int64_t ix = x + kx - padding;
-                        if (ix < 0 || ix >= width) {
-                            continue;
-                        }
-                        sum += preprocessing_get_4d(input, ix, iy, c, 0) * preprocessing_get_4d(kernel, kx, ky, 0, 0);
-                    }
-                }
-                preprocessing_set_4d(output, sum, x, y, c, 0);
-            }
-        }
-    }
-    return output;
-}
-
-static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& rgb_img) {
-    GGML_ASSERT(rgb_img.dim() == 4);
-    GGML_ASSERT(rgb_img.shape()[2] >= 3);
-    sd::Tensor<float> grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]});
-    for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) {
-        for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) {
-            float r    = preprocessing_get_4d(rgb_img, ix, iy, 0, 0);
-            float g    = preprocessing_get_4d(rgb_img, ix, iy, 1, 0);
-            float b    = preprocessing_get_4d(rgb_img, ix, iy, 2, 0);
-            float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b;
-            preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0);
-        }
-    }
-    return grayscale;
-}
-
-static inline sd::Tensor<float> tensor_hypot(const sd::Tensor<float>& x, const sd::Tensor<float>& y) {
-    sd::tensor_check_same_shape(x, y);
-    sd::Tensor<float> out(x.shape());
-    for (int64_t i = 0; i < out.numel(); ++i) {
-        out[i] = std::sqrt(x[i] * x[i] + y[i] * y[i]);
-    }
-    return out;
-}
-
-static inline sd::Tensor<float> tensor_arctan2(const sd::Tensor<float>& x, const sd::Tensor<float>& y) {
-    sd::tensor_check_same_shape(x, y);
-    sd::Tensor<float> out(x.shape());
-    for (int64_t i = 0; i < out.numel(); ++i) {
-        out[i] = std::atan2(y[i], x[i]);
-    }
-    return out;
-}
-
-static inline void normalize_tensor(sd::Tensor<float>* g) {
-    GGML_ASSERT(g != nullptr);
-    if (g->empty()) {
-        return;
-    }
-    float max_value = -std::numeric_limits<float>::infinity();
-    for (int64_t i = 0; i < g->numel(); ++i) {
-        max_value = std::max(max_value, (*g)[i]);
-    }
-    if (max_value == 0.0f || !std::isfinite(max_value)) {
-        return;
-    }
-    *g *= (1.0f / max_value);
-}
-
-static inline sd::Tensor<float> non_max_supression(const sd::Tensor<float>& G, const sd::Tensor<float>& D) {
-    GGML_ASSERT(G.shape() == D.shape());
-    sd::Tensor<float> result = sd::Tensor<float>::zeros(G.shape());
-    for (int64_t iy = 1; iy < result.shape()[1] - 1; ++iy) {
-        for (int64_t ix = 1; ix < result.shape()[0] - 1; ++ix) {
-            float angle = preprocessing_get_4d(D, ix, iy, 0, 0) * 180.0f / M_PI_;
-            angle       = angle < 0.0f ? angle + 180.0f : angle;
-            float q     = 1.0f;
-            float r     = 1.0f;
-
-            if ((0 >= angle && angle < 22.5f) || (157.5f >= angle && angle <= 180.0f)) {
-                q = preprocessing_get_4d(G, ix, iy + 1, 0, 0);
-                r = preprocessing_get_4d(G, ix, iy - 1, 0, 0);
-            } else if (22.5f >= angle && angle < 67.5f) {
-                q = preprocessing_get_4d(G, ix + 1, iy - 1, 0, 0);
-                r = preprocessing_get_4d(G, ix - 1, iy + 1, 0, 0);
-            } else if (67.5f >= angle && angle < 112.5f) {
-                q = preprocessing_get_4d(G, ix + 1, iy, 0, 0);
-                r = preprocessing_get_4d(G, ix - 1, iy, 0, 0);
-            } else if (112.5f >= angle && angle < 157.5f) {
-                q = preprocessing_get_4d(G, ix - 1, iy - 1, 0, 0);
-                r = preprocessing_get_4d(G, ix + 1, iy + 1, 0, 0);
-            }
-
-            float cur = preprocessing_get_4d(G, ix, iy, 0, 0);
-            preprocessing_set_4d(result, (cur >= q && cur >= r) ? cur : 0.0f, ix, iy, 0, 0);
-        }
-    }
-    return result;
-}
-
-static inline void threshold_hystersis(sd::Tensor<float>* img, float high_threshold, float low_threshold, float weak, float strong) {
-    GGML_ASSERT(img != nullptr);
-    if (img->empty()) {
-        return;
-    }
-    float max_value = -std::numeric_limits<float>::infinity();
-    for (int64_t i = 0; i < img->numel(); ++i) {
-        max_value = std::max(max_value, (*img)[i]);
-    }
-
-    float ht = max_value * high_threshold;
-    float lt = ht * low_threshold;
-    for (int64_t i = 0; i < img->numel(); ++i) {
-        float img_v = (*img)[i];
-        if (img_v >= ht) {
-            (*img)[i] = strong;
-        } else if (img_v <= ht && img_v >= lt) {
-            (*img)[i] = weak;
-        }
-    }
-
-    for (int64_t iy = 0; iy < img->shape()[1]; ++iy) {
-        for (int64_t ix = 0; ix < img->shape()[0]; ++ix) {
-            if (!(ix >= 3 && ix <= img->shape()[0] - 3 && iy >= 3 && iy <= img->shape()[1] - 3)) {
-                preprocessing_set_4d(*img, 0.0f, ix, iy, 0, 0);
-            }
-        }
-    }
-
-    for (int64_t iy = 1; iy < img->shape()[1] - 1; ++iy) {
-        for (int64_t ix = 1; ix < img->shape()[0] - 1; ++ix) {
-            float imd_v = preprocessing_get_4d(*img, ix, iy, 0, 0);
-            if (imd_v == weak) {
-                bool has_strong_neighbor =
-                    preprocessing_get_4d(*img, ix + 1, iy - 1, 0, 0) == strong ||
-                    preprocessing_get_4d(*img, ix + 1, iy, 0, 0) == strong ||
-                    preprocessing_get_4d(*img, ix, iy - 1, 0, 0) == strong ||
-                    preprocessing_get_4d(*img, ix, iy + 1, 0, 0) == strong ||
-                    preprocessing_get_4d(*img, ix - 1, iy - 1, 0, 0) == strong ||
-                    preprocessing_get_4d(*img, ix - 1, iy, 0, 0) == strong;
-                preprocessing_set_4d(*img, has_strong_neighbor ? strong : 0.0f, ix, iy, 0, 0);
-            }
-        }
-    }
-}
-
-bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
-    float kX[9] = {
-        -1, 0, 1,
-        -2, 0, 2,
-        -1, 0, 1};
-
-    float kY[9] = {
-        1, 2, 1,
-        0, 0, 0,
-        -1, -2, -1};
-
-    sd::Tensor<float> gkernel = gaussian_kernel_tensor(5);
-    sd::Tensor<float> sf_kx({3, 3, 1, 1}, std::vector<float>(kX, kX + 9));
-    sd::Tensor<float> sf_ky({3, 3, 1, 1}, std::vector<float>(kY, kY + 9));
-
-    sd::Tensor<float> image      = sd_image_to_preprocessing_tensor(img);
-    sd::Tensor<float> image_gray = grayscale_tensor(image);
-    image_gray                   = convolve_tensor(image_gray, gkernel, 2);
-    sd::Tensor<float> iX         = convolve_tensor(image_gray, sf_kx, 1);
-    sd::Tensor<float> iY         = convolve_tensor(image_gray, sf_ky, 1);
-    sd::Tensor<float> G          = tensor_hypot(iX, iY);
-    normalize_tensor(&G);
-    sd::Tensor<float> theta = tensor_arctan2(iX, iY);
-    image_gray              = non_max_supression(G, theta);
-    threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong);
-
-    for (uint32_t iy = 0; iy < img.height; ++iy) {
-        for (uint32_t ix = 0; ix < img.width; ++ix) {
-            float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0);
-            gray       = inverse ? 1.0f - gray : gray;
-            for (uint32_t c = 0; c < img.channel; ++c) {
-                preprocessing_set_4d(image, gray, ix, iy, c, 0);
-            }
-        }
-    }
-
-    preprocessing_tensor_to_sd_image(image, img.data);
-    return true;
-}
-
-#endif  // __PREPROCESSING_HPP__

src/sample-cache.cpp

@@ -1,361 +0,0 @@
-#include "sample-cache.h"
-
-namespace sd_sample {
-
-    static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
-        float reuse_threshold = params.reuse_threshold;
-        if (reuse_threshold == INFINITY) {
-            if (params.mode == SD_CACHE_EASYCACHE) {
-                reuse_threshold = 0.2f;
-            } else if (params.mode == SD_CACHE_UCACHE) {
-                reuse_threshold = 1.0f;
-            }
-        }
-        return std::max(0.0f, reuse_threshold);
-    }
-
-    bool SampleCacheRuntime::easycache_enabled() const {
-        return mode == SampleCacheMode::EASYCACHE;
-    }
-
-    bool SampleCacheRuntime::ucache_enabled() const {
-        return mode == SampleCacheMode::UCACHE;
-    }
-
-    bool SampleCacheRuntime::cachedit_enabled() const {
-        return mode == SampleCacheMode::CACHEDIT;
-    }
-
-    static bool has_valid_cache_percent_range(const sd_cache_params_t& cache_params) {
-        if (cache_params.mode != SD_CACHE_EASYCACHE && cache_params.mode != SD_CACHE_UCACHE) {
-            return true;
-        }
-
-        return cache_params.start_percent >= 0.0f &&
-               cache_params.start_percent < 1.0f &&
-               cache_params.end_percent > 0.0f &&
-               cache_params.end_percent <= 1.0f &&
-               cache_params.start_percent < cache_params.end_percent;
-    }
-
-    static void init_easycache_runtime(SampleCacheRuntime& runtime,
-                                       SDVersion version,
-                                       const sd_cache_params_t& cache_params,
-                                       Denoiser* denoiser) {
-        if (!sd_version_is_dit(version)) {
-            LOG_WARN("EasyCache requested but not supported for this model type");
-            return;
-        }
-
-        EasyCacheConfig config;
-        config.enabled         = true;
-        config.reuse_threshold = get_cache_reuse_threshold(cache_params);
-        config.start_percent   = cache_params.start_percent;
-        config.end_percent     = cache_params.end_percent;
-
-        runtime.easycache.init(config, denoiser);
-        if (!runtime.easycache.enabled()) {
-            LOG_WARN("EasyCache requested but could not be initialized for this run");
-            return;
-        }
-
-        runtime.mode = SampleCacheMode::EASYCACHE;
-        LOG_INFO("EasyCache enabled - threshold: %.3f, start: %.2f, end: %.2f",
-                 config.reuse_threshold,
-                 config.start_percent,
-                 config.end_percent);
-    }
-
-    static void init_ucache_runtime(SampleCacheRuntime& runtime,
-                                    SDVersion version,
-                                    const sd_cache_params_t& cache_params,
-                                    Denoiser* denoiser,
-                                    const std::vector<float>& sigmas) {
-        if (!sd_version_is_unet(version)) {
-            LOG_WARN("UCache requested but not supported for this model type (only UNET models)");
-            return;
-        }
-
-        UCacheConfig config;
-        config.enabled                = true;
-        config.reuse_threshold        = get_cache_reuse_threshold(cache_params);
-        config.start_percent          = cache_params.start_percent;
-        config.end_percent            = cache_params.end_percent;
-        config.error_decay_rate       = std::max(0.0f, std::min(1.0f, cache_params.error_decay_rate));
-        config.use_relative_threshold = cache_params.use_relative_threshold;
-        config.reset_error_on_compute = cache_params.reset_error_on_compute;
-
-        runtime.ucache.init(config, denoiser);
-        if (!runtime.ucache.enabled()) {
-            LOG_WARN("UCache requested but could not be initialized for this run");
-            return;
-        }
-
-        runtime.ucache.set_sigmas(sigmas);
-        runtime.mode = SampleCacheMode::UCACHE;
-        LOG_INFO("UCache enabled - threshold: %.3f, start: %.2f, end: %.2f, decay: %.2f, relative: %s, reset: %s",
-                 config.reuse_threshold,
-                 config.start_percent,
-                 config.end_percent,
-                 config.error_decay_rate,
-                 config.use_relative_threshold ? "true" : "false",
-                 config.reset_error_on_compute ? "true" : "false");
-    }
-
-    static void init_cachedit_runtime(SampleCacheRuntime& runtime,
-                                      SDVersion version,
-                                      const sd_cache_params_t& cache_params,
-                                      const std::vector<float>& sigmas) {
-        if (!sd_version_is_dit(version)) {
-            LOG_WARN("CacheDIT requested but not supported for this model type (only DiT models)");
-            return;
-        }
-
-        DBCacheConfig dbcfg;
-        dbcfg.enabled                     = (cache_params.mode == SD_CACHE_DBCACHE || cache_params.mode == SD_CACHE_CACHE_DIT);
-        dbcfg.Fn_compute_blocks           = cache_params.Fn_compute_blocks;
-        dbcfg.Bn_compute_blocks           = cache_params.Bn_compute_blocks;
-        dbcfg.residual_diff_threshold     = cache_params.residual_diff_threshold;
-        dbcfg.max_warmup_steps            = cache_params.max_warmup_steps;
-        dbcfg.max_cached_steps            = cache_params.max_cached_steps;
-        dbcfg.max_continuous_cached_steps = cache_params.max_continuous_cached_steps;
-        if (cache_params.scm_mask != nullptr && strlen(cache_params.scm_mask) > 0) {
-            dbcfg.steps_computation_mask = parse_scm_mask(cache_params.scm_mask);
-        }
-        dbcfg.scm_policy_dynamic = cache_params.scm_policy_dynamic;
-
-        TaylorSeerConfig tcfg;
-        tcfg.enabled             = (cache_params.mode == SD_CACHE_TAYLORSEER || cache_params.mode == SD_CACHE_CACHE_DIT);
-        tcfg.n_derivatives       = cache_params.taylorseer_n_derivatives;
-        tcfg.skip_interval_steps = cache_params.taylorseer_skip_interval;
-
-        runtime.cachedit.init(dbcfg, tcfg);
-        if (!runtime.cachedit.enabled()) {
-            LOG_WARN("CacheDIT requested but could not be initialized for this run");
-            return;
-        }
-
-        runtime.cachedit.set_sigmas(sigmas);
-        runtime.mode = SampleCacheMode::CACHEDIT;
-        LOG_INFO("CacheDIT enabled - mode: %s, Fn: %d, Bn: %d, threshold: %.3f, warmup: %d",
-                 cache_params.mode == SD_CACHE_CACHE_DIT ? "DBCache+TaylorSeer" : (cache_params.mode == SD_CACHE_DBCACHE ? "DBCache" : "TaylorSeer"),
-                 dbcfg.Fn_compute_blocks,
-                 dbcfg.Bn_compute_blocks,
-                 dbcfg.residual_diff_threshold,
-                 dbcfg.max_warmup_steps);
-    }
-
-    static void init_spectrum_runtime(SampleCacheRuntime& runtime,
-                                      SDVersion version,
-                                      const sd_cache_params_t& cache_params,
-                                      const std::vector<float>& sigmas) {
-        if (!sd_version_is_unet(version) && !sd_version_is_dit(version)) {
-            LOG_WARN("Spectrum requested but not supported for this model type (only UNET and DiT models)");
-            return;
-        }
-
-        SpectrumConfig config;
-        config.w            = cache_params.spectrum_w;
-        config.m            = cache_params.spectrum_m;
-        config.lam          = cache_params.spectrum_lam;
-        config.window_size  = cache_params.spectrum_window_size;
-        config.flex_window  = cache_params.spectrum_flex_window;
-        config.warmup_steps = cache_params.spectrum_warmup_steps;
-        config.stop_percent = cache_params.spectrum_stop_percent;
-
-        size_t total_steps = sigmas.size() > 0 ? sigmas.size() - 1 : 0;
-        runtime.spectrum.init(config, total_steps);
-        runtime.spectrum_enabled = true;
-
-        LOG_INFO("Spectrum enabled - w: %.2f, m: %d, lam: %.2f, window: %d, flex: %.2f, warmup: %d, stop: %.0f%%",
-                 config.w, config.m, config.lam,
-                 config.window_size, config.flex_window,
-                 config.warmup_steps, config.stop_percent * 100.0f);
-    }
-
-    SampleCacheRuntime init_sample_cache_runtime(SDVersion version,
-                                                 const sd_cache_params_t* cache_params,
-                                                 Denoiser* denoiser,
-                                                 const std::vector<float>& sigmas) {
-        SampleCacheRuntime runtime;
-        if (cache_params == nullptr || cache_params->mode == SD_CACHE_DISABLED) {
-            return runtime;
-        }
-
-        if (!has_valid_cache_percent_range(*cache_params)) {
-            LOG_WARN("Cache disabled due to invalid percent range (start=%.3f, end=%.3f)",
-                     cache_params->start_percent,
-                     cache_params->end_percent);
-            return runtime;
-        }
-
-        switch (cache_params->mode) {
-            case SD_CACHE_EASYCACHE:
-                init_easycache_runtime(runtime, version, *cache_params, denoiser);
-                break;
-            case SD_CACHE_UCACHE:
-                init_ucache_runtime(runtime, version, *cache_params, denoiser, sigmas);
-                break;
-            case SD_CACHE_DBCACHE:
-            case SD_CACHE_TAYLORSEER:
-            case SD_CACHE_CACHE_DIT:
-                init_cachedit_runtime(runtime, version, *cache_params, sigmas);
-                break;
-            case SD_CACHE_SPECTRUM:
-                init_spectrum_runtime(runtime, version, *cache_params, sigmas);
-                break;
-            default:
-                break;
-        }
-
-        return runtime;
-    }
-
-    SampleStepCacheDispatcher::SampleStepCacheDispatcher(SampleCacheRuntime& runtime, int step, float sigma)
-        : runtime(runtime), step(step), sigma(sigma), step_index(step > 0 ? (step - 1) : -1) {
-        if (step_index < 0) {
-            return;
-        }
-
-        switch (runtime.mode) {
-            case SampleCacheMode::EASYCACHE:
-                runtime.easycache.begin_step(step_index, sigma);
-                break;
-            case SampleCacheMode::UCACHE:
-                runtime.ucache.begin_step(step_index, sigma);
-                break;
-            case SampleCacheMode::CACHEDIT:
-                runtime.cachedit.begin_step(step_index, sigma);
-                break;
-            case SampleCacheMode::NONE:
-                break;
-        }
-    }
-
-    bool SampleStepCacheDispatcher::before_condition(const void* condition,
-                                                     const sd::Tensor<float>& input,
-                                                     sd::Tensor<float>* output) {
-        if (step_index < 0 || condition == nullptr || output == nullptr) {
-            return false;
-        }
-
-        switch (runtime.mode) {
-            case SampleCacheMode::EASYCACHE:
-                return runtime.easycache.before_condition(condition, input, output, sigma, step_index);
-            case SampleCacheMode::UCACHE:
-                return runtime.ucache.before_condition(condition, input, output, sigma, step_index);
-            case SampleCacheMode::CACHEDIT:
-                return runtime.cachedit.before_condition(condition, input, output, sigma, step_index);
-            case SampleCacheMode::NONE:
-                return false;
-        }
-
-        return false;
-    }
-
-    void SampleStepCacheDispatcher::after_condition(const void* condition,
-                                                    const sd::Tensor<float>& input,
-                                                    const sd::Tensor<float>& output) {
-        if (step_index < 0 || condition == nullptr) {
-            return;
-        }
-
-        switch (runtime.mode) {
-            case SampleCacheMode::EASYCACHE:
-                runtime.easycache.after_condition(condition, input, output);
-                break;
-            case SampleCacheMode::UCACHE:
-                runtime.ucache.after_condition(condition, input, output);
-                break;
-            case SampleCacheMode::CACHEDIT:
-                runtime.cachedit.after_condition(condition, input, output);
-                break;
-            case SampleCacheMode::NONE:
-                break;
-        }
-    }
-
-    bool SampleStepCacheDispatcher::is_step_skipped() const {
-        switch (runtime.mode) {
-            case SampleCacheMode::EASYCACHE:
-                return runtime.easycache.is_step_skipped();
-            case SampleCacheMode::UCACHE:
-                return runtime.ucache.is_step_skipped();
-            case SampleCacheMode::CACHEDIT:
-                return runtime.cachedit.is_step_skipped();
-            case SampleCacheMode::NONE:
-                return false;
-        }
-
-        return false;
-    }
-
-    void log_sample_cache_summary(const SampleCacheRuntime& runtime, size_t total_steps) {
-        if (runtime.easycache_enabled()) {
-            if (runtime.easycache.total_steps_skipped > 0 && total_steps > 0) {
-                if (runtime.easycache.total_steps_skipped < static_cast<int>(total_steps)) {
-                    double speedup = static_cast<double>(total_steps) /
-                                     static_cast<double>(total_steps - runtime.easycache.total_steps_skipped);
-                    LOG_INFO("EasyCache skipped %d/%zu steps (%.2fx estimated speedup)",
-                             runtime.easycache.total_steps_skipped,
-                             total_steps,
-                             speedup);
-                } else {
-                    LOG_INFO("EasyCache skipped %d/%zu steps",
-                             runtime.easycache.total_steps_skipped,
-                             total_steps);
-                }
-            } else if (total_steps > 0) {
-                LOG_INFO("EasyCache completed without skipping steps");
-            }
-        }
-
-        if (runtime.ucache_enabled()) {
-            if (runtime.ucache.total_steps_skipped > 0 && total_steps > 0) {
-                if (runtime.ucache.total_steps_skipped < static_cast<int>(total_steps)) {
-                    double speedup = static_cast<double>(total_steps) /
-                                     static_cast<double>(total_steps - runtime.ucache.total_steps_skipped);
-                    LOG_INFO("UCache skipped %d/%zu steps (%.2fx estimated speedup)",
-                             runtime.ucache.total_steps_skipped,
-                             total_steps,
-                             speedup);
-                } else {
-                    LOG_INFO("UCache skipped %d/%zu steps",
-                             runtime.ucache.total_steps_skipped,
-                             total_steps);
-                }
-            } else if (total_steps > 0) {
-                LOG_INFO("UCache completed without skipping steps");
-            }
-        }
-
-        if (runtime.cachedit_enabled()) {
-            if (runtime.cachedit.total_steps_skipped > 0 && total_steps > 0) {
-                if (runtime.cachedit.total_steps_skipped < static_cast<int>(total_steps)) {
-                    double speedup = static_cast<double>(total_steps) /
-                                     static_cast<double>(total_steps - runtime.cachedit.total_steps_skipped);
-                    LOG_INFO("CacheDIT skipped %d/%zu steps (%.2fx estimated speedup)",
-                             runtime.cachedit.total_steps_skipped,
-                             total_steps,
-                             speedup);
-                } else {
-                    LOG_INFO("CacheDIT skipped %d/%zu steps",
-                             runtime.cachedit.total_steps_skipped,
-                             total_steps);
-                }
-            } else if (total_steps > 0) {
-                LOG_INFO("CacheDIT completed without skipping steps");
-            }
-        }
-
-        if (runtime.spectrum_enabled && runtime.spectrum.total_steps_skipped > 0 && total_steps > 0) {
-            double speedup = static_cast<double>(total_steps) /
-                             static_cast<double>(total_steps - runtime.spectrum.total_steps_skipped);
-            LOG_INFO("Spectrum skipped %d/%zu steps (%.2fx estimated speedup)",
-                     runtime.spectrum.total_steps_skipped,
-                     total_steps,
-                     speedup);
-        }
-    }
-
-}  // namespace sd_sample

src/sample-cache.h

@@ -1,61 +0,0 @@
-#ifndef __SAMPLE_CACHE_H__
-#define __SAMPLE_CACHE_H__
-
-#include <vector>
-
-#include "cache_dit.hpp"
-#include "denoiser.hpp"
-#include "easycache.hpp"
-#include "model.h"
-#include "spectrum.hpp"
-#include "tensor.hpp"
-#include "ucache.hpp"
-#include "util.h"
-
-namespace sd_sample {
-
-    enum class SampleCacheMode {
-        NONE,
-        EASYCACHE,
-        UCACHE,
-        CACHEDIT,
-    };
-
-    struct SampleCacheRuntime {
-        SampleCacheMode mode = SampleCacheMode::NONE;
-
-        EasyCacheState easycache;
-        UCacheState ucache;
-        CacheDitConditionState cachedit;
-        SpectrumState spectrum;
-
-        bool spectrum_enabled = false;
-
-        bool easycache_enabled() const;
-        bool ucache_enabled() const;
-        bool cachedit_enabled() const;
-    };
-
-    struct SampleStepCacheDispatcher {
-        SampleCacheRuntime& runtime;
-        int step;
-        float sigma;
-        int step_index;
-
-        SampleStepCacheDispatcher(SampleCacheRuntime& runtime, int step, float sigma);
-
-        bool before_condition(const void* condition, const sd::Tensor<float>& input, sd::Tensor<float>* output);
-        void after_condition(const void* condition, const sd::Tensor<float>& input, const sd::Tensor<float>& output);
-        bool is_step_skipped() const;
-    };
-
-    SampleCacheRuntime init_sample_cache_runtime(SDVersion version,
-                                                 const sd_cache_params_t* cache_params,
-                                                 Denoiser* denoiser,
-                                                 const std::vector<float>& sigmas);
-
-    void log_sample_cache_summary(const SampleCacheRuntime& runtime, size_t total_steps);
-
-}  // namespace sd_sample
-
-#endif  // __SAMPLE_CACHE_H__

src/spectrum.hpp

@@ -1,187 +0,0 @@
-#ifndef __SPECTRUM_HPP__
-#define __SPECTRUM_HPP__
-
-#include <cmath>
-#include <cstring>
-#include <vector>
-
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
-
-struct SpectrumConfig {
-    float w            = 0.40f;
-    int m              = 3;
-    float lam          = 1.0f;
-    int window_size    = 2;
-    float flex_window  = 0.50f;
-    int warmup_steps   = 4;
-    float stop_percent = 0.9f;
-};
-
-struct SpectrumState {
-    SpectrumConfig config;
-    int cnt                 = 0;
-    int num_cached          = 0;
-    float curr_ws           = 2.0f;
-    int K                   = 6;
-    int stop_step           = 0;
-    int total_steps_skipped = 0;
-
-    std::vector<std::vector<float>> H_buf;
-    std::vector<float> T_buf;
-
-    void init(const SpectrumConfig& cfg, size_t total_steps) {
-        config              = cfg;
-        cnt                 = 0;
-        num_cached          = 0;
-        curr_ws             = (float)cfg.window_size;
-        K                   = std::max(cfg.m + 1, 6);
-        stop_step           = (int)(cfg.stop_percent * (float)total_steps);
-        total_steps_skipped = 0;
-        H_buf.clear();
-        T_buf.clear();
-    }
-
-    float taus(int step_cnt) const {
-        return (step_cnt / 50.0f) * 2.0f - 1.0f;
-    }
-
-    bool should_predict() {
-        if (cnt < config.warmup_steps)
-            return false;
-        if (stop_step > 0 && cnt >= stop_step)
-            return false;
-        if ((int)H_buf.size() < 2)
-            return false;
-
-        int ws = std::max(1, (int)std::floor(curr_ws));
-        return (num_cached + 1) % ws != 0;
-    }
-
-    void update(const sd::Tensor<float>& denoised) {
-        H_buf.emplace_back(denoised.data(), denoised.data() + denoised.numel());
-        T_buf.push_back(taus(cnt));
-
-        while ((int)H_buf.size() > K) {
-            H_buf.erase(H_buf.begin());
-            T_buf.erase(T_buf.begin());
-        }
-
-        if (cnt >= config.warmup_steps)
-            curr_ws += config.flex_window;
-
-        num_cached = 0;
-        cnt++;
-    }
-
-    void predict(sd::Tensor<float>* denoised) {
-        GGML_ASSERT(denoised != nullptr);
-        int64_t F    = (int64_t)H_buf[0].size();
-        int K_curr   = (int)H_buf.size();
-        int M1       = config.m + 1;
-        float tau_at = taus(cnt);
-
-        std::vector<float> X(K_curr * M1);
-        for (int i = 0; i < K_curr; i++) {
-            X[i * M1] = 1.0f;
-            if (M1 > 1)
-                X[i * M1 + 1] = T_buf[i];
-            for (int j = 2; j < M1; j++)
-                X[i * M1 + j] = 2.0f * T_buf[i] * X[i * M1 + j - 1] - X[i * M1 + j - 2];
-        }
-
-        std::vector<float> x_star(M1);
-        x_star[0] = 1.0f;
-        if (M1 > 1)
-            x_star[1] = tau_at;
-        for (int j = 2; j < M1; j++)
-            x_star[j] = 2.0f * tau_at * x_star[j - 1] - x_star[j - 2];
-
-        std::vector<float> XtX(M1 * M1, 0.0f);
-        for (int i = 0; i < M1; i++) {
-            for (int j = 0; j < M1; j++) {
-                float sum = 0.0f;
-                for (int k = 0; k < K_curr; k++)
-                    sum += X[k * M1 + i] * X[k * M1 + j];
-                XtX[i * M1 + j] = sum + (i == j ? config.lam : 0.0f);
-            }
-        }
-
-        std::vector<float> L(M1 * M1, 0.0f);
-        if (!cholesky_decompose(XtX.data(), L.data(), M1)) {
-            float trace = 0.0f;
-            for (int i = 0; i < M1; i++)
-                trace += XtX[i * M1 + i];
-            for (int i = 0; i < M1; i++)
-                XtX[i * M1 + i] += 1e-4f * trace / M1;
-            cholesky_decompose(XtX.data(), L.data(), M1);
-        }
-
-        std::vector<float> v(M1);
-        cholesky_solve(L.data(), x_star.data(), v.data(), M1);
-
-        std::vector<float> weights(K_curr, 0.0f);
-        for (int k = 0; k < K_curr; k++)
-            for (int j = 0; j < M1; j++)
-                weights[k] += X[k * M1 + j] * v[j];
-
-        float* out          = denoised->data();
-        float w_cheb        = config.w;
-        float w_taylor      = 1.0f - w_cheb;
-        const float* h_last = H_buf.back().data();
-        const float* h_prev = H_buf[H_buf.size() - 2].data();
-
-        for (int64_t f = 0; f < F; f++) {
-            float pred_cheb = 0.0f;
-            for (int k = 0; k < K_curr; k++)
-                pred_cheb += weights[k] * H_buf[k][f];
-
-            float pred_taylor = h_last[f] + 0.5f * (h_last[f] - h_prev[f]);
-
-            out[f] = w_taylor * pred_taylor + w_cheb * pred_cheb;
-        }
-
-        num_cached++;
-        total_steps_skipped++;
-        cnt++;
-    }
-
-private:
-    static bool cholesky_decompose(const float* A, float* L, int n) {
-        std::memset(L, 0, n * n * sizeof(float));
-        for (int i = 0; i < n; i++) {
-            for (int j = 0; j <= i; j++) {
-                float sum = 0.0f;
-                for (int k = 0; k < j; k++)
-                    sum += L[i * n + k] * L[j * n + k];
-                if (i == j) {
-                    float diag = A[i * n + i] - sum;
-                    if (diag <= 0.0f)
-                        return false;
-                    L[i * n + j] = std::sqrt(diag);
-                } else {
-                    L[i * n + j] = (A[i * n + j] - sum) / L[j * n + j];
-                }
-            }
-        }
-        return true;
-    }
-
-    static void cholesky_solve(const float* L, const float* b, float* x, int n) {
-        std::vector<float> y(n);
-        for (int i = 0; i < n; i++) {
-            float sum = 0.0f;
-            for (int j = 0; j < i; j++)
-                sum += L[i * n + j] * y[j];
-            y[i] = (b[i] - sum) / L[i * n + i];
-        }
-        for (int i = n - 1; i >= 0; i--) {
-            float sum = 0.0f;
-            for (int j = i + 1; j < n; j++)
-                sum += L[j * n + i] * x[j];
-            x[i] = (y[i] - sum) / L[i * n + i];
-        }
-    }
-};
-
-#endif  // __SPECTRUM_HPP__

src/tae.hpp

@@ -1,664 +0,0 @@
-#ifndef __TAE_HPP__
-#define __TAE_HPP__
-
-#include "ggml_extend.hpp"
-
-#include "model.h"
-
-/*
-    ===================================    TinyAutoEncoder  ===================================
-    References:
-    https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py
-    https://github.com/madebyollin/taesd/blob/main/taesd.py
-
-*/
-
-class TAEBlock : public UnaryBlock {
-protected:
-    int n_in;
-    int n_out;
-    bool use_midblock_gn;
-
-public:
-    TAEBlock(int n_in, int n_out, bool use_midblock_gn = false)
-        : n_in(n_in), n_out(n_out), use_midblock_gn(use_midblock_gn) {
-        blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
-        if (n_in != n_out) {
-            blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {1, 1}, {1, 1}, {1, 1}, {1, 1}, false));
-        }
-        if (use_midblock_gn) {
-            int n_gn         = n_in * 4;
-            blocks["pool.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_gn, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
-            blocks["pool.1"] = std::shared_ptr<GGMLBlock>(new GroupNorm(4, n_gn));
-            // pool.2 is ReLU, handled in forward
-            blocks["pool.3"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_gn, n_in, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
-        }
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-        // x: [n, n_in, h, w]
-        // return: [n, n_out, h, w]
-
-        if (use_midblock_gn) {
-            auto pool_0 = std::dynamic_pointer_cast<Conv2d>(blocks["pool.0"]);
-            auto pool_1 = std::dynamic_pointer_cast<GroupNorm>(blocks["pool.1"]);
-            auto pool_3 = std::dynamic_pointer_cast<Conv2d>(blocks["pool.3"]);
-
-            auto p = pool_0->forward(ctx, x);
-            p      = pool_1->forward(ctx, p);
-            p      = ggml_relu_inplace(ctx->ggml_ctx, p);
-            p      = pool_3->forward(ctx, p);
-
-            x = ggml_add(ctx->ggml_ctx, x, p);
-        }
-
-        auto conv_0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
-        auto conv_2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
-        auto conv_4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
-
-        auto h = conv_0->forward(ctx, x);
-        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
-        h      = conv_2->forward(ctx, h);
-        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
-        h      = conv_4->forward(ctx, h);
-
-        if (n_in != n_out) {
-            auto skip = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
-            LOG_DEBUG("skip");
-            x = skip->forward(ctx, x);
-        }
-
-        h = ggml_add(ctx->ggml_ctx, h, x);
-        h = ggml_relu_inplace(ctx->ggml_ctx, h);
-        return h;
-    }
-};
-
-class TinyEncoder : public UnaryBlock {
-    int in_channels = 3;
-    int channels    = 64;
-    int z_channels  = 4;
-    int num_blocks  = 3;
-
-public:
-    TinyEncoder(int z_channels = 4, bool use_midblock_gn = false)
-        : z_channels(z_channels) {
-        int index                       = 0;
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, channels, {3, 3}, {1, 1}, {1, 1}));
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
-
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
-        for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
-        }
-
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
-        for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
-        }
-
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
-        for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
-        }
-
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-        // x: [n, in_channels, h, w]
-        // return: [n, z_channels, h/8, w/8]
-
-        for (int i = 0; i < num_blocks * 3 + 6; i++) {
-            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
-
-            x = block->forward(ctx, x);
-        }
-
-        return x;
-    }
-};
-
-class TinyDecoder : public UnaryBlock {
-    int z_channels   = 4;
-    int channels     = 64;
-    int out_channels = 3;
-    int num_blocks   = 3;
-
-public:
-    TinyDecoder(int z_channels = 4, bool use_midblock_gn = false)
-        : z_channels(z_channels) {
-        int index = 0;
-
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels, {3, 3}, {1, 1}, {1, 1}));
-        index++;  // nn.ReLU()
-
-        for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
-        }
-        index++;  // nn.Upsample()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
-
-        for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
-        }
-        index++;  // nn.Upsample()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
-
-        for (int i = 0; i < num_blocks; i++) {
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
-        }
-        index++;  // nn.Upsample()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
-
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
-        // z: [n, z_channels, h, w]
-        // return: [n, out_channels, h*8, w*8]
-
-        auto h = ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f);
-        h      = ggml_tanh_inplace(ctx->ggml_ctx, h);
-        h      = ggml_ext_scale(ctx->ggml_ctx, h, 3.0f);
-
-        for (int i = 0; i < num_blocks * 3 + 10; i++) {
-            if (blocks.find(std::to_string(i)) == blocks.end()) {
-                if (i == 1) {
-                    h = ggml_relu_inplace(ctx->ggml_ctx, h);
-                } else {
-                    h = ggml_upscale(ctx->ggml_ctx, h, 2, GGML_SCALE_MODE_NEAREST);
-                }
-                continue;
-            }
-            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
-
-            h = block->forward(ctx, h);
-        }
-
-        return h;
-    }
-};
-
-class TPool : public UnaryBlock {
-    int stride;
-
-public:
-    TPool(int channels, int stride)
-        : stride(stride) {
-        blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * stride, channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-        auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
-        auto h    = x;
-        if (stride != 1) {
-            h = ggml_reshape_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2] * stride, h->ne[3] / stride);
-        }
-        h = conv->forward(ctx, h);
-        return h;
-    }
-};
-
-class TGrow : public UnaryBlock {
-    int stride;
-
-public:
-    TGrow(int channels, int stride)
-        : stride(stride) {
-        blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels * stride, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-        auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
-        auto h    = conv->forward(ctx, x);
-        if (stride != 1) {
-            h = ggml_reshape_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2] / stride, h->ne[3] * stride);
-        }
-        return h;
-    }
-};
-
-class MemBlock : public GGMLBlock {
-    bool has_skip_conv = false;
-
-public:
-    MemBlock(int channels, int out_channels)
-        : has_skip_conv(channels != out_channels) {
-        blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * 2, out_channels, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
-        if (has_skip_conv) {
-            blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
-        }
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* past) {
-        // x: [n, channels, h, w]
-        auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
-        auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
-        auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
-
-        auto h = ggml_concat(ctx->ggml_ctx, x, past, 2);
-        h      = conv0->forward(ctx, h);
-        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
-        h      = conv1->forward(ctx, h);
-        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
-        h      = conv2->forward(ctx, h);
-
-        auto skip = x;
-        if (has_skip_conv) {
-            auto skip_conv = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
-            skip           = skip_conv->forward(ctx, x);
-        }
-        h = ggml_add_inplace(ctx->ggml_ctx, h, skip);
-        h = ggml_relu_inplace(ctx->ggml_ctx, h);
-        return h;
-    }
-};
-
-ggml_tensor* patchify(ggml_context* ctx,
-                      ggml_tensor* x,
-                      int64_t patch_size,
-                      int64_t b = 1) {
-    // x: [f, b*c, h*q, w*r]
-    // return: [f, b*c*r*q, h, w]
-    if (patch_size == 1) {
-        return x;
-    }
-    int64_t r = patch_size;
-    int64_t q = patch_size;
-
-    int64_t W = x->ne[0];
-    int64_t H = x->ne[1];
-    int64_t C = x->ne[2];
-    int64_t f = x->ne[3];
-
-    int64_t w = W / r;
-    int64_t h = H / q;
-
-    x = ggml_reshape_4d(ctx, x, W, q, h, C * f);                         // [W, q, h, C*f]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [W, h, q, C*f]
-    x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f);                     // [r, w, h, q*C*f]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));  // [w, h, r, q*C*f]
-    x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f);                     // [f, b*c*r*q, h, w]
-
-    return x;
-}
-
-ggml_tensor* unpatchify(ggml_context* ctx,
-                        ggml_tensor* x,
-                        int64_t patch_size,
-                        int64_t b = 1) {
-    // x: [f, b*c*r*q, h, w]
-    // return: [f, b*c, h*q, w*r]
-    if (patch_size == 1) {
-        return x;
-    }
-    int64_t r = patch_size;
-    int64_t q = patch_size;
-    int64_t c = x->ne[2] / b / q / r;
-    int64_t f = x->ne[3];
-    int64_t h = x->ne[1];
-    int64_t w = x->ne[0];
-
-    x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f);                 // [q*c*b*f, r, h, w]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [r, w, h, q*c*b*f]
-    x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f);                 // [c*b*f, q, h, r*w]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [r*w, q, h, c*b*f]
-    x = ggml_reshape_4d(ctx, x, r * w, q * h, c * b, f);
-
-    return x;
-}
-
-class TinyVideoEncoder : public UnaryBlock {
-    int in_channels = 3;
-    int hidden      = 64;
-    int z_channels  = 4;
-    int num_blocks  = 3;
-    int num_layers  = 3;
-    int patch_size  = 1;
-
-public:
-    TinyVideoEncoder(int z_channels = 4, int patch_size = 1)
-        : z_channels(z_channels), patch_size(patch_size) {
-        int index                       = 0;
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels * patch_size * patch_size, hidden, {3, 3}, {1, 1}, {1, 1}));
-        index++;  // nn.ReLU()
-        for (int i = 0; i < num_layers; i++) {
-            int stride                      = i == num_layers - 1 ? 1 : 2;
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TPool(hidden, stride));
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, hidden, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
-            for (int j = 0; j < num_blocks; j++) {
-                blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(hidden, hidden));
-            }
-        }
-        blocks[std::to_string(index)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, z_channels, {3, 3}, {1, 1}, {1, 1}));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
-        auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
-
-        if (patch_size > 1) {
-            z = patchify(ctx->ggml_ctx, z, patch_size, 1);
-        }
-
-        auto h = first_conv->forward(ctx, z);
-        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
-
-        int index = 2;
-        for (int i = 0; i < num_layers; i++) {
-            auto pool = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
-            auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
-
-            h = pool->forward(ctx, h);
-            h = conv->forward(ctx, h);
-            for (int j = 0; j < num_blocks; j++) {
-                auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
-                auto mem   = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
-                mem        = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
-                h          = block->forward(ctx, h, mem);
-            }
-        }
-        auto last_conv = std::dynamic_pointer_cast<Conv2d>(blocks[std::to_string(index)]);
-        h              = last_conv->forward(ctx, h);
-        return h;
-    }
-};
-
-class TinyVideoDecoder : public UnaryBlock {
-    int z_channels               = 4;
-    int out_channels             = 3;
-    int num_blocks               = 3;
-    static const int num_layers  = 3;
-    int channels[num_layers + 1] = {256, 128, 64, 64};
-    int patch_size               = 1;
-
-public:
-    TinyVideoDecoder(int z_channels = 4, int patch_size = 1)
-        : z_channels(z_channels), patch_size(patch_size) {
-        int index                       = 1;  // Clamp()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels[0], {3, 3}, {1, 1}, {1, 1}));
-        index++;  // nn.ReLU()
-        for (int i = 0; i < num_layers; i++) {
-            int stride = i == 0 ? 1 : 2;
-            for (int j = 0; j < num_blocks; j++) {
-                blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(channels[i], channels[i]));
-            }
-            index++;  // nn.Upsample()
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TGrow(channels[i], stride));
-            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels[i], channels[i + 1], {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
-        }
-        index++;  // nn.ReLU()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels[num_layers], out_channels * patch_size * patch_size, {3, 3}, {1, 1}, {1, 1}));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
-        auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
-
-        // Clamp()
-        auto h = ggml_ext_scale(ctx->ggml_ctx,
-                                ggml_tanh_inplace(ctx->ggml_ctx,
-                                                  ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f)),
-                                3.0f,
-                                true);
-
-        h         = first_conv->forward(ctx, h);
-        h         = ggml_relu_inplace(ctx->ggml_ctx, h);
-        int index = 3;
-        for (int i = 0; i < num_layers; i++) {
-            for (int j = 0; j < num_blocks; j++) {
-                auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
-                auto mem   = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
-                mem        = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
-                h          = block->forward(ctx, h, mem);
-            }
-            // upsample
-            index++;
-            h          = ggml_upscale(ctx->ggml_ctx, h, 2, GGML_SCALE_MODE_NEAREST);
-            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
-            h          = block->forward(ctx, h);
-            block      = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
-            h          = block->forward(ctx, h);
-        }
-        h = ggml_relu_inplace(ctx->ggml_ctx, h);
-
-        auto last_conv = std::dynamic_pointer_cast<Conv2d>(blocks[std::to_string(++index)]);
-        h              = last_conv->forward(ctx, h);
-        if (patch_size > 1) {
-            h = unpatchify(ctx->ggml_ctx, h, patch_size, 1);
-        }
-        // shape(W, H, 3, 3 + T) => shape(W, H, 3, T)
-        h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - 3, h->nb[1], h->nb[2], h->nb[3], 3 * h->nb[3]);
-        return h;
-    }
-};
-
-class TAEHV : public GGMLBlock {
-protected:
-    bool decode_only;
-    SDVersion version;
-
-public:
-    int z_channels = 16;
-
-public:
-    TAEHV(bool decode_only = true, SDVersion version = VERSION_WAN2)
-        : decode_only(decode_only), version(version) {
-        int patch = 1;
-        if (version == VERSION_WAN2_2_TI2V) {
-            z_channels = 48;
-            patch      = 2;
-        }
-        blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch));
-        if (!decode_only) {
-            blocks["encoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoEncoder(z_channels, patch));
-        }
-    }
-
-    ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
-        auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
-        if (sd_version_is_wan(version)) {
-            // (W, H, C, T) -> (W, H, T, C)
-            z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
-        }
-        auto result = decoder->forward(ctx, z);
-        if (sd_version_is_wan(version)) {
-            // (W, H, C, T) -> (W, H, T, C)
-            result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
-        }
-        return result;
-    }
-
-    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));
-        int64_t num_frames = x->ne[3];
-        if (num_frames % 4) {
-            // pad to multiple of 4 at the end
-            auto last_frame = ggml_view_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], x->ne[2], 1, x->nb[1], x->nb[2], x->nb[3], (num_frames - 1) * x->nb[3]);
-            for (int i = 0; i < 4 - num_frames % 4; i++) {
-                x = ggml_concat(ctx->ggml_ctx, x, last_frame, 3);
-            }
-        }
-        x = encoder->forward(ctx, x);
-        x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
-        return x;
-    }
-};
-
-class TAESD : public GGMLBlock {
-protected:
-    bool decode_only;
-    bool taef2 = false;
-
-public:
-    int z_channels = 4;
-
-public:
-    TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
-        : decode_only(decode_only) {
-        bool use_midblock_gn = false;
-        taef2                = sd_version_is_flux2(version);
-
-        if (sd_version_is_dit(version)) {
-            z_channels = 16;
-        }
-        if (taef2) {
-            z_channels      = 32;
-            use_midblock_gn = true;
-        }
-        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels, use_midblock_gn));
-
-        if (!decode_only) {
-            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels, use_midblock_gn));
-        }
-    }
-
-    ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
-        auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
-        if (taef2) {
-            z = unpatchify(ctx->ggml_ctx, z, 2);
-        }
-        return decoder->forward(ctx, z);
-    }
-
-    ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
-        auto z       = encoder->forward(ctx, x);
-        if (taef2) {
-            z = patchify(ctx->ggml_ctx, z, 2);
-        }
-        return z;
-    }
-};
-
-struct TinyImageAutoEncoder : public VAE {
-    TAESD taesd;
-    bool decode_only = false;
-
-    TinyImageAutoEncoder(ggml_backend_t backend,
-                         bool offload_params_to_cpu,
-                         const String2TensorStorage& tensor_storage_map,
-                         const std::string prefix,
-                         bool decoder_only = true,
-                         SDVersion version = VERSION_SD1)
-        : decode_only(decoder_only),
-          taesd(decoder_only, version),
-          VAE(version, backend, offload_params_to_cpu) {
-        scale_input = false;
-        taesd.init(params_ctx, tensor_storage_map, prefix);
-    }
-
-    std::string get_desc() override {
-        return "taesd";
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-        taesd.get_param_tensors(tensors, prefix);
-    }
-
-    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
-        SD_UNUSED(rng);
-        return vae_output;
-    }
-
-    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
-        return latents;
-    }
-
-    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
-        return latents;
-    }
-
-    int get_encoder_output_channels(int input_channels) {
-        return taesd.z_channels;
-    }
-
-    ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
-        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
-        ggml_tensor* z   = make_input(z_tensor);
-        auto runner_ctx  = get_context();
-        ggml_tensor* out = decode_graph ? taesd.decode(&runner_ctx, z) : taesd.encode(&runner_ctx, z);
-        ggml_build_forward_expand(gf, out);
-        return gf;
-    }
-
-    sd::Tensor<float> _compute(const int n_threads,
-                               const sd::Tensor<float>& z_tensor,
-                               bool decode_graph) override {
-        auto get_graph = [&]() -> ggml_cgraph* {
-            return build_graph(z_tensor, decode_graph);
-        };
-
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
-    }
-};
-
-struct TinyVideoAutoEncoder : public VAE {
-    TAEHV taehv;
-    bool decode_only = false;
-
-    TinyVideoAutoEncoder(ggml_backend_t backend,
-                         bool offload_params_to_cpu,
-                         const String2TensorStorage& tensor_storage_map,
-                         const std::string prefix,
-                         bool decoder_only = true,
-                         SDVersion version = VERSION_WAN2)
-        : decode_only(decoder_only),
-          taehv(decoder_only, version),
-          VAE(version, backend, offload_params_to_cpu) {
-        scale_input = false;
-        taehv.init(params_ctx, tensor_storage_map, prefix);
-    }
-
-    std::string get_desc() override {
-        return "taehv";
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-        taehv.get_param_tensors(tensors, prefix);
-    }
-
-    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
-        SD_UNUSED(rng);
-        return vae_output;
-    }
-
-    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
-        return latents;
-    }
-
-    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
-        return latents;
-    }
-
-    int get_encoder_output_channels(int input_channels) {
-        return taehv.z_channels;
-    }
-
-    ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
-        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
-        ggml_tensor* z   = make_input(z_tensor);
-        auto runner_ctx  = get_context();
-        ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
-        ggml_build_forward_expand(gf, out);
-        return gf;
-    }
-
-    sd::Tensor<float> _compute(const int n_threads,
-                               const sd::Tensor<float>& z_tensor,
-                               bool decode_graph) override {
-        auto get_graph = [&]() -> ggml_cgraph* {
-            return build_graph(z_tensor, decode_graph);
-        };
-
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
-    }
-};
-
-#endif  // __TAE_HPP__

src/tensor_ggml.hpp

@@ -1,127 +0,0 @@
-#ifndef __SD_TENSOR_GGML_HPP__
-#define __SD_TENSOR_GGML_HPP__
-
-#include <array>
-#include <cstring>
-#include <fstream>
-#include <stdexcept>
-#include <string>
-#include <type_traits>
-
-#include "ggml.h"
-#include "tensor.hpp"
-
-namespace sd {
-
-    template <typename T>
-    struct GGMLTypeTraits;
-
-    template <>
-    struct GGMLTypeTraits<float> {
-        static constexpr ggml_type type = GGML_TYPE_F32;
-    };
-
-    template <>
-    struct GGMLTypeTraits<ggml_fp16_t> {
-        static constexpr ggml_type type = GGML_TYPE_F16;
-    };
-
-    template <>
-    struct GGMLTypeTraits<int32_t> {
-        static constexpr ggml_type type = GGML_TYPE_I32;
-    };
-
-    template <>
-    struct GGMLTypeTraits<int64_t> {
-        static constexpr ggml_type type = GGML_TYPE_I64;
-    };
-
-    inline std::vector<int64_t> shape_from_ggml(const ggml_tensor* tensor) {
-        std::vector<int64_t> shape;
-        shape.reserve(static_cast<size_t>(ggml_n_dims(tensor)));
-        for (int i = 0; i < ggml_n_dims(tensor); ++i) {
-            shape.push_back(tensor->ne[i]);
-        }
-        return shape;
-    }
-
-    template <typename T>
-    inline Tensor<T> make_sd_tensor_from_ggml(const ggml_tensor* tensor) {
-        if (tensor == nullptr) {
-            return {};
-        }
-        if (tensor->type != GGMLTypeTraits<T>::type) {
-            GGML_ABORT("ggml tensor type does not match sd::Tensor type");
-        }
-        Tensor<T> result(shape_from_ggml(tensor));
-        if (tensor->buffer != nullptr) {
-            ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
-        } else {
-            std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor));
-        }
-        return result;
-    }
-
-    template <typename T>
-    inline ggml_tensor* make_ggml_tensor(ggml_context* ctx, const Tensor<T>& tensor, bool copy_data = true) {
-        GGML_ASSERT(tensor.dim() > 0 && tensor.dim() <= 5);
-
-        int n_dims = std::min(static_cast<int>(tensor.dim()), GGML_MAX_DIMS);
-
-        std::array<int64_t, GGML_MAX_DIMS> ne = {1, 1, 1, 1};
-        for (int64_t i = 0; i < n_dims; ++i) {
-            ne[static_cast<size_t>(i)] = tensor.shape()[static_cast<size_t>(i)];
-        }
-
-        if (tensor.dim() == 5) {
-            ne[3] *= tensor.shape()[4];
-        }
-
-        ggml_tensor* result = ggml_new_tensor(ctx, GGMLTypeTraits<T>::type, n_dims, ne.data());
-        if (copy_data && tensor.numel() > 0) {
-            std::memcpy(result->data, tensor.data(), static_cast<size_t>(ggml_nbytes(result)));
-        }
-        return result;
-    }
-
-    template <typename T>
-    inline Tensor<T> load_tensor_from_file_as_tensor(const std::string& file_path) {
-        std::ifstream file(file_path, std::ios::binary);
-        if (!file.is_open()) {
-            throw std::runtime_error("failed to open tensor file: " + file_path);
-        }
-
-        int32_t n_dims = 0;
-        int32_t length = 0;
-        int32_t ttype  = 0;
-        file.read(reinterpret_cast<char*>(&n_dims), sizeof(n_dims));
-        file.read(reinterpret_cast<char*>(&length), sizeof(length));
-        file.read(reinterpret_cast<char*>(&ttype), sizeof(ttype));
-        if (!file.good()) {
-            throw std::runtime_error("incomplete tensor file header: " + file_path);
-        }
-        if (static_cast<ggml_type>(ttype) != GGMLTypeTraits<T>::type) {
-            throw std::invalid_argument("tensor file type does not match requested sd::Tensor type");
-        }
-
-        std::vector<int64_t> shape(4, 1);
-        for (int i = 0; i < n_dims; ++i) {
-            int32_t dim = 1;
-            file.read(reinterpret_cast<char*>(&dim), sizeof(dim));
-            shape[static_cast<size_t>(i)] = dim;
-        }
-        std::string name(static_cast<size_t>(length), '\0');
-        file.read(name.data(), length);
-
-        shape.resize(static_cast<size_t>(n_dims));
-        Tensor<T> tensor(shape);
-        file.read(reinterpret_cast<char*>(tensor.data()), static_cast<std::streamsize>(tensor.numel() * sizeof(T)));
-        if (!file.good()) {
-            throw std::runtime_error("incomplete tensor file data: " + file_path);
-        }
-        return tensor;
-    }
-
-}  // namespace sd
-
-#endif

src/ucache.hpp

@@ -1,423 +0,0 @@
-#ifndef __UCACHE_HPP__
-#define __UCACHE_HPP__
-
-#include <cmath>
-#include <limits>
-#include <unordered_map>
-#include <vector>
-
-#include "condition_cache_utils.hpp"
-#include "denoiser.hpp"
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
-
-struct UCacheConfig {
-    bool enabled                = false;
-    float reuse_threshold       = 1.0f;
-    float start_percent         = 0.15f;
-    float end_percent           = 0.95f;
-    float error_decay_rate      = 1.0f;
-    bool use_relative_threshold = true;
-    bool adaptive_threshold     = true;
-    float early_step_multiplier = 0.5f;
-    float late_step_multiplier  = 1.5f;
-    float relative_norm_gain    = 1.6f;
-    bool reset_error_on_compute = true;
-};
-
-struct UCacheCacheEntry {
-    std::vector<float> diff;
-};
-
-struct UCacheState {
-    UCacheConfig config;
-    Denoiser* denoiser           = nullptr;
-    float start_sigma            = std::numeric_limits<float>::max();
-    float end_sigma              = 0.0f;
-    bool initialized             = false;
-    bool initial_step            = true;
-    bool skip_current_step       = false;
-    bool step_active             = false;
-    const void* anchor_condition = nullptr;
-    std::unordered_map<const void*, UCacheCacheEntry> cache_diffs;
-    std::vector<float> prev_input;
-    std::vector<float> prev_output;
-    float output_prev_norm                = 0.0f;
-    bool has_prev_input                   = false;
-    bool has_prev_output                  = false;
-    bool has_output_prev_norm             = false;
-    bool has_relative_transformation_rate = false;
-    float relative_transformation_rate    = 0.0f;
-    float last_input_change               = 0.0f;
-    bool has_last_input_change            = false;
-    float output_change_ema               = 0.0f;
-    bool has_output_change_ema            = false;
-    int total_steps_skipped               = 0;
-    int current_step_index                = -1;
-    int steps_computed_since_active       = 0;
-    int expected_total_steps              = 0;
-    int consecutive_skipped_steps         = 0;
-    float accumulated_error               = 0.0f;
-
-    struct BlockMetrics {
-        float sum_transformation_rate = 0.0f;
-        float sum_output_norm         = 0.0f;
-        int sample_count              = 0;
-        float min_change_rate         = std::numeric_limits<float>::max();
-        float max_change_rate         = 0.0f;
-
-        void reset() {
-            sum_transformation_rate = 0.0f;
-            sum_output_norm         = 0.0f;
-            sample_count            = 0;
-            min_change_rate         = std::numeric_limits<float>::max();
-            max_change_rate         = 0.0f;
-        }
-
-        void record(float change_rate, float output_norm) {
-            if (std::isfinite(change_rate) && change_rate > 0.0f) {
-                sum_transformation_rate += change_rate;
-                sum_output_norm += output_norm;
-                sample_count++;
-                if (change_rate < min_change_rate)
-                    min_change_rate = change_rate;
-                if (change_rate > max_change_rate)
-                    max_change_rate = change_rate;
-            }
-        }
-
-        float avg_transformation_rate() const {
-            return (sample_count > 0) ? (sum_transformation_rate / sample_count) : 0.0f;
-        }
-
-        float avg_output_norm() const {
-            return (sample_count > 0) ? (sum_output_norm / sample_count) : 0.0f;
-        }
-    };
-    BlockMetrics block_metrics;
-    int total_active_steps = 0;
-
-    void reset_runtime() {
-        initial_step      = true;
-        skip_current_step = false;
-        step_active       = false;
-        anchor_condition  = nullptr;
-        cache_diffs.clear();
-        prev_input.clear();
-        prev_output.clear();
-        output_prev_norm                 = 0.0f;
-        has_prev_input                   = false;
-        has_prev_output                  = false;
-        has_output_prev_norm             = false;
-        has_relative_transformation_rate = false;
-        relative_transformation_rate     = 0.0f;
-        last_input_change                = 0.0f;
-        has_last_input_change            = false;
-        output_change_ema                = 0.0f;
-        has_output_change_ema            = false;
-        total_steps_skipped              = 0;
-        current_step_index               = -1;
-        steps_computed_since_active      = 0;
-        expected_total_steps             = 0;
-        consecutive_skipped_steps        = 0;
-        accumulated_error                = 0.0f;
-        block_metrics.reset();
-        total_active_steps = 0;
-    }
-
-    void init(const UCacheConfig& cfg, Denoiser* d) {
-        config      = cfg;
-        denoiser    = d;
-        initialized = cfg.enabled && d != nullptr;
-        reset_runtime();
-        if (initialized) {
-            start_sigma = percent_to_sigma(config.start_percent);
-            end_sigma   = percent_to_sigma(config.end_percent);
-        }
-    }
-
-    void set_sigmas(const std::vector<float>& sigmas) {
-        if (!initialized || sigmas.size() < 2) {
-            return;
-        }
-        size_t n_steps       = sigmas.size() - 1;
-        expected_total_steps = static_cast<int>(n_steps);
-
-        size_t start_step = static_cast<size_t>(config.start_percent * n_steps);
-        size_t end_step   = static_cast<size_t>(config.end_percent * n_steps);
-
-        if (start_step >= n_steps)
-            start_step = n_steps - 1;
-        if (end_step >= n_steps)
-            end_step = n_steps - 1;
-
-        start_sigma = sigmas[start_step];
-        end_sigma   = sigmas[end_step];
-
-        if (start_sigma < end_sigma) {
-            std::swap(start_sigma, end_sigma);
-        }
-    }
-
-    bool enabled() const {
-        return initialized && config.enabled;
-    }
-
-    float percent_to_sigma(float percent) const {
-        if (!denoiser) {
-            return 0.0f;
-        }
-        if (percent <= 0.0f) {
-            return std::numeric_limits<float>::max();
-        }
-        if (percent >= 1.0f) {
-            return 0.0f;
-        }
-        float t = (1.0f - percent) * (TIMESTEPS - 1);
-        return denoiser->t_to_sigma(t);
-    }
-
-    void begin_step(int step_index, float sigma) {
-        if (!enabled()) {
-            return;
-        }
-        if (step_index == current_step_index) {
-            return;
-        }
-        current_step_index    = step_index;
-        skip_current_step     = false;
-        has_last_input_change = false;
-        step_active           = false;
-
-        if (sigma > start_sigma) {
-            return;
-        }
-        if (!(sigma > end_sigma)) {
-            return;
-        }
-        step_active = true;
-        total_active_steps++;
-    }
-
-    bool step_is_active() const {
-        return enabled() && step_active;
-    }
-
-    bool is_step_skipped() const {
-        return enabled() && step_active && skip_current_step;
-    }
-
-    float get_adaptive_threshold(int estimated_total_steps = 0) const {
-        float base_threshold = config.reuse_threshold;
-
-        if (!config.adaptive_threshold) {
-            return base_threshold;
-        }
-
-        int effective_total = estimated_total_steps;
-        if (effective_total <= 0) {
-            effective_total = expected_total_steps;
-        }
-        if (effective_total <= 0) {
-            effective_total = std::max(20, steps_computed_since_active * 2);
-        }
-
-        float progress = (effective_total > 0) ? (static_cast<float>(steps_computed_since_active) / effective_total) : 0.0f;
-        progress       = std::max(0.0f, std::min(1.0f, progress));
-
-        float multiplier = 1.0f;
-        if (progress < 0.2f) {
-            multiplier = config.early_step_multiplier;
-        } else if (progress > 0.8f) {
-            multiplier = config.late_step_multiplier;
-        }
-
-        return base_threshold * multiplier;
-    }
-
-    bool has_cache(const void* cond) const {
-        auto it = cache_diffs.find(cond);
-        return it != cache_diffs.end() && !it->second.diff.empty();
-    }
-
-    void update_cache(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
-        UCacheCacheEntry& entry = cache_diffs[cond];
-        sd::store_condition_cache_diff(&entry.diff, input, output);
-    }
-
-    void apply_cache(const void* cond, const sd::Tensor<float>& input, sd::Tensor<float>* output) {
-        auto it = cache_diffs.find(cond);
-        if (it == cache_diffs.end() || it->second.diff.empty()) {
-            return;
-        }
-        sd::apply_condition_cache_diff(it->second.diff, input, output);
-    }
-
-    bool before_condition(const void* cond,
-                          const sd::Tensor<float>& input,
-                          sd::Tensor<float>* output,
-                          float sigma,
-                          int step_index) {
-        if (!enabled() || step_index < 0 || output == nullptr) {
-            return false;
-        }
-        if (step_index != current_step_index) {
-            begin_step(step_index, sigma);
-        }
-        if (!step_active) {
-            return false;
-        }
-
-        if (initial_step) {
-            anchor_condition = cond;
-            initial_step     = false;
-        }
-
-        bool is_anchor = (cond == anchor_condition);
-
-        if (skip_current_step) {
-            if (has_cache(cond)) {
-                apply_cache(cond, input, output);
-                return true;
-            }
-            return false;
-        }
-
-        if (!is_anchor) {
-            return false;
-        }
-
-        if (!has_prev_input || !has_prev_output || !has_cache(cond)) {
-            return false;
-        }
-
-        size_t ne = static_cast<size_t>(input.numel());
-        if (prev_input.size() != ne) {
-            return false;
-        }
-
-        const float* input_data = input.data();
-        last_input_change       = 0.0f;
-        for (size_t i = 0; i < ne; ++i) {
-            last_input_change += std::fabs(input_data[i] - prev_input[i]);
-        }
-        if (ne > 0) {
-            last_input_change /= static_cast<float>(ne);
-        }
-        has_last_input_change = true;
-
-        if (has_output_prev_norm && has_relative_transformation_rate &&
-            last_input_change > 0.0f && output_prev_norm > 0.0f) {
-            float approx_output_change = relative_transformation_rate * last_input_change;
-            float approx_output_change_rate;
-            if (config.use_relative_threshold) {
-                float base_scale          = std::max(output_prev_norm, 1e-6f);
-                float dyn_scale           = has_output_change_ema
-                                                ? std::max(output_change_ema * std::max(1.0f, config.relative_norm_gain), 1e-6f)
-                                                : base_scale;
-                float scale               = std::sqrt(base_scale * dyn_scale);
-                approx_output_change_rate = approx_output_change / scale;
-            } else {
-                approx_output_change_rate = approx_output_change;
-            }
-            // Increase estimated error with skip horizon to avoid long extrapolation streaks
-            approx_output_change_rate *= (1.0f + 0.50f * consecutive_skipped_steps);
-            accumulated_error = accumulated_error * config.error_decay_rate + approx_output_change_rate;
-
-            float effective_threshold = get_adaptive_threshold();
-            if (!config.use_relative_threshold && output_prev_norm > 0.0f) {
-                effective_threshold = effective_threshold * output_prev_norm;
-            }
-
-            if (accumulated_error < effective_threshold) {
-                skip_current_step = true;
-                total_steps_skipped++;
-                consecutive_skipped_steps++;
-                apply_cache(cond, input, output);
-                return true;
-            } else if (config.reset_error_on_compute) {
-                accumulated_error = 0.0f;
-            }
-        }
-
-        return false;
-    }
-
-    void after_condition(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
-        if (!step_is_active()) {
-            return;
-        }
-
-        update_cache(cond, input, output);
-
-        if (cond != anchor_condition) {
-            return;
-        }
-        steps_computed_since_active++;
-        consecutive_skipped_steps = 0;
-
-        size_t ne            = static_cast<size_t>(input.numel());
-        const float* in_data = input.data();
-        prev_input.resize(ne);
-        for (size_t i = 0; i < ne; ++i) {
-            prev_input[i] = in_data[i];
-        }
-        has_prev_input = true;
-
-        const float* out_data = output.data();
-        float output_change   = 0.0f;
-        if (has_prev_output && prev_output.size() == ne) {
-            for (size_t i = 0; i < ne; ++i) {
-                output_change += std::fabs(out_data[i] - prev_output[i]);
-            }
-            if (ne > 0) {
-                output_change /= static_cast<float>(ne);
-            }
-        }
-        if (std::isfinite(output_change) && output_change > 0.0f) {
-            if (!has_output_change_ema) {
-                output_change_ema     = output_change;
-                has_output_change_ema = true;
-            } else {
-                output_change_ema = 0.8f * output_change_ema + 0.2f * output_change;
-            }
-        }
-
-        prev_output.resize(ne);
-        for (size_t i = 0; i < ne; ++i) {
-            prev_output[i] = out_data[i];
-        }
-        has_prev_output = true;
-
-        float mean_abs = 0.0f;
-        for (size_t i = 0; i < ne; ++i) {
-            mean_abs += std::fabs(out_data[i]);
-        }
-        output_prev_norm     = (ne > 0) ? (mean_abs / static_cast<float>(ne)) : 0.0f;
-        has_output_prev_norm = output_prev_norm > 0.0f;
-
-        if (has_last_input_change && last_input_change > 0.0f && output_change > 0.0f) {
-            float rate = output_change / last_input_change;
-            if (std::isfinite(rate)) {
-                relative_transformation_rate     = rate;
-                has_relative_transformation_rate = true;
-                block_metrics.record(rate, output_prev_norm);
-            }
-        }
-
-        has_last_input_change = false;
-    }
-
-    void log_block_metrics() const {
-        if (block_metrics.sample_count > 0) {
-            LOG_INFO("UCacheBlockMetrics: samples=%d, avg_rate=%.4f, min=%.4f, max=%.4f, avg_norm=%.4f",
-                     block_metrics.sample_count,
-                     block_metrics.avg_transformation_rate(),
-                     block_metrics.min_change_rate,
-                     block_metrics.max_change_rate,
-                     block_metrics.avg_output_norm());
-        }
-    }
-};
-
-#endif  // __UCACHE_HPP__

src/vae.hpp

@@ -1,253 +0,0 @@
-#ifndef __VAE_HPP__
-#define __VAE_HPP__
-
-#include "common_block.hpp"
-#include "tensor_ggml.hpp"
-
-struct VAE : public GGMLRunner {
-protected:
-    SDVersion version;
-    bool scale_input                                      = true;
-    virtual sd::Tensor<float> _compute(const int n_threads,
-                                       const sd::Tensor<float>& z,
-                                       bool decode_graph) = 0;
-
-    static inline void scale_tensor_to_minus1_1(sd::Tensor<float>* tensor) {
-        GGML_ASSERT(tensor != nullptr);
-        for (int64_t i = 0; i < tensor->numel(); ++i) {
-            (*tensor)[i] = (*tensor)[i] * 2.0f - 1.0f;
-        }
-    }
-
-    static inline void scale_tensor_to_0_1(sd::Tensor<float>* tensor) {
-        GGML_ASSERT(tensor != nullptr);
-        for (int64_t i = 0; i < tensor->numel(); ++i) {
-            float value  = ((*tensor)[i] + 1.0f) * 0.5f;
-            (*tensor)[i] = std::max(0.0f, std::min(1.0f, value));
-        }
-    }
-
-    sd::Tensor<float> tiled_compute(const sd::Tensor<float>& input,
-                                    int n_threads,
-                                    int output_width,
-                                    int output_height,
-                                    int scale,
-                                    int p_tile_size_x,
-                                    int p_tile_size_y,
-                                    float tile_overlap_factor,
-                                    bool circular_x,
-                                    bool circular_y,
-                                    bool decode_graph,
-                                    const char* error_message,
-                                    bool silent = false) {
-        auto on_processing = [&](const sd::Tensor<float>& input_tile) {
-            auto output_tile = _compute(n_threads, input_tile, decode_graph);
-            if (output_tile.empty()) {
-                LOG_ERROR("%s", error_message);
-                return sd::Tensor<float>();
-            }
-            return output_tile;
-        };
-        return ::process_tiles_2d(input,
-                                  output_width,
-                                  output_height,
-                                  scale,
-                                  p_tile_size_x,
-                                  p_tile_size_y,
-                                  tile_overlap_factor,
-                                  circular_x,
-                                  circular_y,
-                                  on_processing,
-                                  silent);
-    }
-
-public:
-    VAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
-        : version(version), GGMLRunner(backend, offload_params_to_cpu) {}
-
-    int get_scale_factor() {
-        int scale_factor = 8;
-        if (version == VERSION_WAN2_2_TI2V) {
-            scale_factor = 16;
-        } else if (sd_version_is_flux2(version)) {
-            scale_factor = 16;
-        } else if (version == VERSION_CHROMA_RADIANCE) {
-            scale_factor = 1;
-        }
-        return scale_factor;
-    }
-
-    virtual int get_encoder_output_channels(int input_channels) = 0;
-
-    void get_tile_sizes(int& tile_size_x,
-                        int& tile_size_y,
-                        float& tile_overlap,
-                        const sd_tiling_params_t& params,
-                        int64_t latent_x,
-                        int64_t latent_y,
-                        float encoding_factor = 1.0f) {
-        tile_overlap       = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
-        auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
-            const int default_tile_size  = 32;
-            const int min_tile_dimension = 4;
-            int tile_size                = default_tile_size;
-            // factor <= 1 means simple fraction of the latent dimension
-            // factor > 1 means number of tiles across that dimension
-            if (factor > 0.f) {
-                if (factor > 1.0)
-                    factor = 1 / (factor - factor * tile_overlap + tile_overlap);
-                tile_size = static_cast<int>(std::round(latent_size * factor));
-            } else if (requested_size >= min_tile_dimension) {
-                tile_size = requested_size;
-            }
-            tile_size = static_cast<int>(tile_size * encoding_factor);
-            return std::max(std::min(tile_size, static_cast<int>(latent_size)), min_tile_dimension);
-        };
-
-        tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
-        tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
-    }
-
-    sd::Tensor<float> encode(int n_threads,
-                             const sd::Tensor<float>& x,
-                             sd_tiling_params_t tiling_params,
-                             bool circular_x = false,
-                             bool circular_y = false) {
-        int64_t t0              = ggml_time_ms();
-        sd::Tensor<float> input = x;
-        sd::Tensor<float> output;
-        if (scale_input) {
-            scale_tensor_to_minus1_1(&input);
-        }
-
-        if (tiling_params.enabled) {
-            const int scale_factor = get_scale_factor();
-            int64_t W              = input.shape()[0] / scale_factor;
-            int64_t H              = input.shape()[1] / scale_factor;
-            float tile_overlap;
-            int tile_size_x, tile_size_y;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, 1.30539f);
-            LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
-            output = tiled_compute(input,
-                                   n_threads,
-                                   static_cast<int>(W),
-                                   static_cast<int>(H),
-                                   scale_factor,
-                                   tile_size_x,
-                                   tile_size_y,
-                                   tile_overlap,
-                                   circular_x,
-                                   circular_y,
-                                   false,
-                                   "vae encode compute failed while processing a tile");
-        } else {
-            output = _compute(n_threads, input, false);
-            free_compute_buffer();
-        }
-
-        if (output.empty()) {
-            LOG_ERROR("vae encode compute failed");
-            return {};
-        }
-        int64_t t1 = ggml_time_ms();
-        LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
-        return std::move(output);
-    }
-
-    sd::Tensor<float> decode(int n_threads,
-                             const sd::Tensor<float>& x,
-                             sd_tiling_params_t tiling_params,
-                             bool decode_video = false,
-                             bool circular_x   = false,
-                             bool circular_y   = false,
-                             bool silent       = false) {
-        int64_t t0              = ggml_time_ms();
-        sd::Tensor<float> input = x;
-        sd::Tensor<float> output;
-
-        if (tiling_params.enabled) {
-            const int scale_factor = get_scale_factor();
-            int64_t W              = input.shape()[0] * scale_factor;
-            int64_t H              = input.shape()[1] * scale_factor;
-            float tile_overlap;
-            int tile_size_x, tile_size_y;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
-            if (!silent) {
-                LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
-            }
-            output = tiled_compute(
-                input,
-                n_threads,
-                static_cast<int>(W),
-                static_cast<int>(H),
-                scale_factor,
-                tile_size_x,
-                tile_size_y,
-                tile_overlap,
-                circular_x,
-                circular_y,
-                true,
-                "vae decode compute failed while processing a tile",
-                silent);
-        } else {
-            output = _compute(n_threads, input, true);
-        }
-
-        free_compute_buffer();
-
-        if (output.empty()) {
-            LOG_ERROR("vae decode compute failed");
-            return {};
-        }
-        if (scale_input) {
-            scale_tensor_to_0_1(&output);
-        }
-        int64_t t1 = ggml_time_ms();
-        LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
-        return std::move(output);
-    }
-
-    virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
-    virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents)                           = 0;
-    virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents)                           = 0;
-    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix)         = 0;
-    virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
-};
-
-struct FakeVAE : public VAE {
-    FakeVAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
-        : VAE(version, backend, offload_params_to_cpu) {}
-
-    int get_encoder_output_channels(int input_channels) {
-        return input_channels;
-    }
-
-    sd::Tensor<float> _compute(const int n_threads,
-                               const sd::Tensor<float>& z,
-                               bool decode_graph) override {
-        SD_UNUSED(n_threads);
-        SD_UNUSED(decode_graph);
-        return z;
-    }
-
-    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
-        SD_UNUSED(rng);
-        return vae_output;
-    }
-
-    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
-        return latents;
-    }
-
-    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
-        return latents;
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {}
-
-    std::string get_desc() override {
-        return "fake_vae";
-    }
-};
-
-#endif  // __VAE_HPP__

src/vocab/vocab.cpp

@@ -1,35 +0,0 @@
-#include "vocab.h"
-#include "clip_t5.hpp"
-#include "mistral.hpp"
-#include "qwen.hpp"
-#include "umt5.hpp"
-
-std::string load_clip_merges() {
-    std::string merges_utf8_str(reinterpret_cast<const char*>(clip_merges_utf8_c_str), sizeof(clip_merges_utf8_c_str));
-    return merges_utf8_str;
-}
-
-std::string load_qwen2_merges() {
-    std::string merges_utf8_str(reinterpret_cast<const char*>(qwen2_merges_utf8_c_str), sizeof(qwen2_merges_utf8_c_str));
-    return merges_utf8_str;
-}
-
-std::string load_mistral_merges() {
-    std::string merges_utf8_str(reinterpret_cast<const char*>(mistral_merges_utf8_c_str), sizeof(mistral_merges_utf8_c_str));
-    return merges_utf8_str;
-}
-
-std::string load_mistral_vocab_json() {
-    std::string json_str(reinterpret_cast<const char*>(mistral_vocab_json_utf8_c_str), sizeof(mistral_vocab_json_utf8_c_str));
-    return json_str;
-}
-
-std::string load_t5_tokenizer_json() {
-    std::string json_str(reinterpret_cast<const char*>(t5_tokenizer_json_str), sizeof(t5_tokenizer_json_str));
-    return json_str;
-}
-
-std::string load_umt5_tokenizer_json() {
-    std::string json_str(reinterpret_cast<const char*>(umt5_tokenizer_json_str), sizeof(umt5_tokenizer_json_str));
-    return json_str;
-}

src/vocab/vocab.h

@@ -1,13 +0,0 @@
-#ifndef __VOCAB_H__
-#define __VOCAB_H__
-
-#include <string>
-
-std::string load_clip_merges();
-std::string load_qwen2_merges();
-std::string load_mistral_merges();
-std::string load_mistral_vocab_json();
-std::string load_t5_tokenizer_json();
-std::string load_umt5_tokenizer_json();
-
-#endif  // __VOCAB_H__

stable-diffusion.h

@@ -48,8 +48,6 @@ enum sample_method_t {
     LCM_SAMPLE_METHOD,
     DDIM_TRAILING_SAMPLE_METHOD,
     TCD_SAMPLE_METHOD,
-    RES_MULTISTEP_SAMPLE_METHOD,
-    RES_2S_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 
@@ -62,9 +60,7 @@ enum scheduler_t {
     SGM_UNIFORM_SCHEDULER,
     SIMPLE_SCHEDULER,
     SMOOTHSTEP_SCHEDULER,
-    KL_OPTIMAL_SCHEDULER,
     LCM_SCHEDULER,
-    BONG_TANGENT_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -120,8 +116,7 @@ enum sd_type_t {
     // SD_TYPE_IQ4_NL_4_8 = 37,
     // SD_TYPE_IQ4_NL_8_8 = 38,
     SD_TYPE_MXFP4 = 39,  // MXFP4 (1 block)
-    SD_TYPE_NVFP4 = 40,  // NVFP4 (4 blocks, E4M3 scale)
-    SD_TYPE_COUNT = 41,
+    SD_TYPE_COUNT = 40,
 };
 
 enum sd_log_level_t {
@@ -173,6 +168,7 @@ typedef struct {
     const char* vae_path;
     const char* taesd_path;
     const char* control_net_path;
+    const char* lora_model_dir;
     const sd_embedding_t* embeddings;
     uint32_t embedding_count;
     const char* photo_maker_path;
@@ -186,22 +182,18 @@ typedef struct {
     enum prediction_t prediction;
     enum lora_apply_mode_t lora_apply_mode;
     bool offload_params_to_cpu;
-    bool enable_mmap;
     bool keep_clip_on_cpu;
     bool keep_control_net_on_cpu;
     bool keep_vae_on_cpu;
-    bool flash_attn;
     bool diffusion_flash_attn;
     bool tae_preview_only;
     bool diffusion_conv_direct;
     bool vae_conv_direct;
-    bool circular_x;
-    bool circular_y;
     bool force_sdxl_vae_conv_scale;
     bool chroma_use_dit_mask;
     bool chroma_use_t5_mask;
     int chroma_t5_mask_pad;
-    bool qwen_image_zero_cond_t;
+    float flow_shift;
 } sd_ctx_params_t;
 
 typedef struct {
@@ -233,9 +225,6 @@ typedef struct {
     int sample_steps;
     float eta;
     int shifted_timestep;
-    float* custom_sigmas;
-    int custom_sigmas_count;
-    float flow_shift;
 } sd_sample_params_t;
 
 typedef struct {
@@ -245,42 +234,12 @@ typedef struct {
     float style_strength;
 } sd_pm_params_t;  // photo maker
 
-enum sd_cache_mode_t {
-    SD_CACHE_DISABLED = 0,
-    SD_CACHE_EASYCACHE,
-    SD_CACHE_UCACHE,
-    SD_CACHE_DBCACHE,
-    SD_CACHE_TAYLORSEER,
-    SD_CACHE_CACHE_DIT,
-    SD_CACHE_SPECTRUM,
-};
-
 typedef struct {
-    enum sd_cache_mode_t mode;
+    bool enabled;
     float reuse_threshold;
     float start_percent;
     float end_percent;
-    float error_decay_rate;
-    bool use_relative_threshold;
-    bool reset_error_on_compute;
-    int Fn_compute_blocks;
-    int Bn_compute_blocks;
-    float residual_diff_threshold;
-    int max_warmup_steps;
-    int max_cached_steps;
-    int max_continuous_cached_steps;
-    int taylorseer_n_derivatives;
-    int taylorseer_skip_interval;
-    const char* scm_mask;
-    bool scm_policy_dynamic;
-    float spectrum_w;
-    int spectrum_m;
-    float spectrum_lam;
-    int spectrum_window_size;
-    float spectrum_flex_window;
-    int spectrum_warmup_steps;
-    float spectrum_stop_percent;
-} sd_cache_params_t;
+} sd_easycache_params_t;
 
 typedef struct {
     bool is_high_noise;
@@ -310,7 +269,7 @@ typedef struct {
     float control_strength;
     sd_pm_params_t pm_params;
     sd_tiling_params_t vae_tiling_params;
-    sd_cache_params_t cache;
+    sd_easycache_params_t easycache;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -332,8 +291,7 @@ typedef struct {
     int64_t seed;
     int video_frames;
     float vace_strength;
-    sd_tiling_params_t vae_tiling_params;
-    sd_cache_params_t cache;
+    sd_easycache_params_t easycache;
 } sd_vid_gen_params_t;
 
 typedef struct sd_ctx_t sd_ctx_t;
@@ -363,7 +321,7 @@ SD_API enum preview_t str_to_preview(const char* str);
 SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
 SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
 
-SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
+SD_API void sd_easycache_params_init(sd_easycache_params_t* easycache_params);
 
 SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params);
 SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
@@ -375,7 +333,7 @@ SD_API void sd_sample_params_init(sd_sample_params_t* sample_params);
 SD_API char* sd_sample_params_to_str(const sd_sample_params_t* sample_params);
 
 SD_API enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx);
-SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method);
+SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx);
 
 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);
@@ -403,8 +361,7 @@ SD_API bool convert(const char* input_path,
                     const char* vae_path,
                     const char* output_path,
                     enum sd_type_t output_type,
-                    const char* tensor_type_rules,
-                    bool convert_name);
+                    const char* tensor_type_rules);
 
 SD_API bool preprocess_canny(sd_image_t image,
                              float high_threshold,

t5.hpp

@@ -1,4 +1,4 @@
-#ifndef __T5_HPP__
+#ifndef __T5_HPP__
 #define __T5_HPP__
 
 #include <cfloat>
@@ -14,7 +14,6 @@
 #include "ggml_extend.hpp"
 #include "json.hpp"
 #include "model.h"
-#include "vocab/vocab.h"
 
 // Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
 // and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
@@ -97,7 +96,7 @@ protected:
 
         try {
             data = nlohmann::json::parse(json_str);
-        } catch (const nlohmann::json::parse_error&) {
+        } catch (const nlohmann::json::parse_error& e) {
             status_ = INVLIAD_JSON;
             return;
         }
@@ -169,9 +168,9 @@ protected:
             kMaxTrieResultsSize);
         trie_results_size_ = 0;
         for (const auto& p : *pieces) {
-            const size_t num_nodes = trie_->commonPrefixSearch(
+            const int num_nodes = trie_->commonPrefixSearch(
                 p.first.data(), results.data(), results.size(), p.first.size());
-            trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
+            trie_results_size_ = std::max(trie_results_size_, num_nodes);
         }
 
         if (trie_results_size_ == 0)
@@ -211,9 +210,9 @@ protected:
         // implementation. It's based on the following three ideas:
         //
         // 1. Because it uses the *unigram* model:
-        //     best_score(x1, x2, ... xt) = best_score(x1, x2, ... x{t-1}) + score(xt)
+        //     best_score(x1, x2, …, xt) = best_score(x1, x2, …, x{t-1}) + score(xt)
         // Deciding the best path (and score) can be decoupled into two isolated
-        // terms: (a) the best path ended before the last token `best_score(x1, x2, ...)`
+        // terms: (a) the best path ended before the last token `best_score(x1, x2, …,
         // x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
         // not related to each other at all.
         //
@@ -227,7 +226,7 @@ protected:
         // position, where n is the input length and k is the maximum number of tokens
         // that can be recognized starting at each position.
         //
-        // 2. Again, because it uses the *unigram* model, we don't need to actually
+        // 2. Again, because it uses the *unigram* model, we don’t need to actually
         // store the lattice nodes. We still recognize all the tokens and lattice
         // nodes from the input, but along identifying them, we use and discard them
         // on the fly. There is no need to actually store them for best path Viterbi
@@ -269,7 +268,7 @@ protected:
                 -1;  // The starting position (in utf-8) of this node. The entire best
                      // path can be constructed by backtracking along this link.
         };
-        const int size        = static_cast<int>(normalized.size());
+        const int size        = normalized.size();
         const float unk_score = min_score() - kUnkPenalty;
         // The ends are exclusive.
         std::vector<BestPathNode> best_path_ends_at(size + 1);
@@ -282,7 +281,7 @@ protected:
                 best_path_ends_at[starts_at].best_path_score;
             bool has_single_node = false;
             const int mblen =
-                std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)),
+                std::min<int>(OneCharLen(normalized.data() + starts_at),
                               size - starts_at);
             while (key_pos < size) {
                 const int ret =
@@ -303,7 +302,7 @@ protected:
                         score + best_path_score_till_here;
                     if (target_node.starts_at == -1 ||
                         candidate_best_path_score > target_node.best_path_score) {
-                        target_node.best_path_score = static_cast<float>(candidate_best_path_score);
+                        target_node.best_path_score = candidate_best_path_score;
                         target_node.starts_at       = starts_at;
                         target_node.id              = ret;
                     }
@@ -342,9 +341,9 @@ protected:
 public:
     explicit T5UniGramTokenizer(bool is_umt5 = false) {
         if (is_umt5) {
-            InitializePieces(load_umt5_tokenizer_json());
+            InitializePieces(ModelLoader::load_umt5_tokenizer_json());
         } else {
-            InitializePieces(load_t5_tokenizer_json());
+            InitializePieces(ModelLoader::load_t5_tokenizer_json());
         }
 
         min_score_ = FLT_MAX;
@@ -395,7 +394,7 @@ public:
                     bool padding      = false) {
         if (max_length > 0 && padding) {
             size_t orig_token_num = tokens.size() - 1;
-            size_t n              = static_cast<size_t>(std::ceil(orig_token_num * 1.0 / (max_length - 1)));
+            size_t n              = std::ceil(orig_token_num * 1.0 / (max_length - 1));
             if (n == 0) {
                 n = 1;
             }
@@ -462,7 +461,7 @@ protected:
     int64_t hidden_size;
     float eps;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         enum ggml_type wtype = GGML_TYPE_F32;
         params["weight"]     = ggml_new_tensor_1d(ctx, wtype, hidden_size);
     }
@@ -473,10 +472,10 @@ public:
         : hidden_size(hidden_size),
           eps(eps) {}
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-        ggml_tensor* w = params["weight"];
-        x              = ggml_rms_norm(ctx->ggml_ctx, x, eps);
-        x              = ggml_mul(ctx->ggml_ctx, x, w);
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+        struct ggml_tensor* w = params["weight"];
+        x                     = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+        x                     = ggml_mul(ctx->ggml_ctx, x, w);
         return x;
     }
 };
@@ -488,7 +487,7 @@ public:
         blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [N, n_token, model_dim]
         auto wi = std::dynamic_pointer_cast<Linear>(blocks["wi"]);
         auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
@@ -510,13 +509,13 @@ public:
         blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false, false, false, scale));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [N, n_token, model_dim]
         auto wi_0 = std::dynamic_pointer_cast<Linear>(blocks["wi_0"]);
         auto wi_1 = std::dynamic_pointer_cast<Linear>(blocks["wi_1"]);
         auto wo   = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
 
-        auto hidden_gelu   = ggml_ext_gelu(ctx->ggml_ctx, wi_0->forward(ctx, x), true);
+        auto hidden_gelu   = ggml_gelu_inplace(ctx->ggml_ctx, wi_0->forward(ctx, x));
         auto hidden_linear = wi_1->forward(ctx, x);
         x                  = ggml_mul_inplace(ctx->ggml_ctx, hidden_gelu, hidden_linear);
         x                  = wo->forward(ctx, x);
@@ -531,7 +530,7 @@ public:
         blocks["layer_norm"]     = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [N, n_token, model_dim]
         auto DenseReluDense = std::dynamic_pointer_cast<T5DenseGatedActDense>(blocks["DenseReluDense"]);
         auto layer_norm     = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
@@ -570,8 +569,8 @@ public:
         }
     }
 
-    ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
-                              ggml_tensor* relative_position_bucket) {
+    struct ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
+                                     struct ggml_tensor* relative_position_bucket) {
         auto relative_attention_bias = std::dynamic_pointer_cast<Embedding>(blocks["relative_attention_bias"]);
 
         auto values = relative_attention_bias->forward(ctx, relative_position_bucket);            // shape (query_length, key_length, num_heads)
@@ -580,11 +579,11 @@ public:
     }
 
     // x: [N, n_token, model_dim]
-    std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                  ggml_tensor* x,
-                                                  ggml_tensor* past_bias                = nullptr,
-                                                  ggml_tensor* mask                     = nullptr,
-                                                  ggml_tensor* relative_position_bucket = nullptr) {
+    std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                                struct ggml_tensor* x,
+                                                                struct ggml_tensor* past_bias                = nullptr,
+                                                                struct ggml_tensor* mask                     = nullptr,
+                                                                struct ggml_tensor* relative_position_bucket = nullptr) {
         auto q_proj   = std::dynamic_pointer_cast<Linear>(blocks["q"]);
         auto k_proj   = std::dynamic_pointer_cast<Linear>(blocks["k"]);
         auto v_proj   = std::dynamic_pointer_cast<Linear>(blocks["v"]);
@@ -609,7 +608,7 @@ public:
             }
         }
 
-        k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
+        k = ggml_scale_inplace(ctx->ggml_ctx, k, sqrt(d_head));
 
         x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
 
@@ -629,11 +628,11 @@ public:
         blocks["layer_norm"]    = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
     }
 
-    std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                  ggml_tensor* x,
-                                                  ggml_tensor* past_bias                = nullptr,
-                                                  ggml_tensor* mask                     = nullptr,
-                                                  ggml_tensor* relative_position_bucket = nullptr) {
+    std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                                struct ggml_tensor* x,
+                                                                struct ggml_tensor* past_bias                = nullptr,
+                                                                struct ggml_tensor* mask                     = nullptr,
+                                                                struct ggml_tensor* relative_position_bucket = nullptr) {
         // x: [N, n_token, model_dim]
         auto SelfAttention = std::dynamic_pointer_cast<T5Attention>(blocks["SelfAttention"]);
         auto layer_norm    = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
@@ -655,11 +654,11 @@ public:
         blocks["layer.1"] = std::shared_ptr<GGMLBlock>(new T5LayerFF(model_dim, ff_dim));
     }
 
-    std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                  ggml_tensor* x,
-                                                  ggml_tensor* past_bias                = nullptr,
-                                                  ggml_tensor* mask                     = nullptr,
-                                                  ggml_tensor* relative_position_bucket = nullptr) {
+    std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                                struct ggml_tensor* x,
+                                                                struct ggml_tensor* past_bias                = nullptr,
+                                                                struct ggml_tensor* mask                     = nullptr,
+                                                                struct ggml_tensor* relative_position_bucket = nullptr) {
         // x: [N, n_token, model_dim]
         auto layer_0 = std::dynamic_pointer_cast<T5LayerSelfAttention>(blocks["layer.0"]);
         auto layer_1 = std::dynamic_pointer_cast<T5LayerFF>(blocks["layer.1"]);
@@ -690,11 +689,11 @@ public:
         blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* past_bias                = nullptr,
-                         ggml_tensor* attention_mask           = nullptr,
-                         ggml_tensor* relative_position_bucket = nullptr) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* past_bias                = nullptr,
+                                struct ggml_tensor* attention_mask           = nullptr,
+                                struct ggml_tensor* relative_position_bucket = nullptr) {
         // x: [N, n_token, model_dim]
         for (int i = 0; i < num_layers; i++) {
             auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
@@ -737,11 +736,11 @@ public:
                                                                      params.model_dim));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids,
-                         ggml_tensor* past_bias                = nullptr,
-                         ggml_tensor* attention_mask           = nullptr,
-                         ggml_tensor* relative_position_bucket = nullptr) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* input_ids,
+                                struct ggml_tensor* past_bias                = nullptr,
+                                struct ggml_tensor* attention_mask           = nullptr,
+                                struct ggml_tensor* relative_position_bucket = nullptr) {
         // input_ids: [N, n_token]
 
         auto shared  = std::dynamic_pointer_cast<Embedding>(blocks["shared"]);
@@ -776,14 +775,14 @@ struct T5Runner : public GGMLRunner {
         return "t5";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids,
-                         ggml_tensor* relative_position_bucket,
-                         ggml_tensor* attention_mask = nullptr) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* input_ids,
+                                struct ggml_tensor* relative_position_bucket,
+                                struct ggml_tensor* attention_mask = nullptr) {
         size_t N       = input_ids->ne[1];
         size_t n_token = input_ids->ne[0];
 
@@ -791,13 +790,14 @@ struct T5Runner : public GGMLRunner {
         return hidden_states;
     }
 
-    ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
-                             const sd::Tensor<float>& attention_mask_tensor = {}) {
-        ggml_cgraph* gf             = ggml_new_graph(compute_ctx);
-        ggml_tensor* input_ids      = make_input(input_ids_tensor);
-        ggml_tensor* attention_mask = attention_mask_tensor.empty() ? nullptr : make_input(attention_mask_tensor);
+    struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
+                                    struct ggml_tensor* attention_mask = nullptr) {
+        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
-        relative_position_bucket_vec = compute_relative_position_bucket(static_cast<int>(input_ids->ne[0]), static_cast<int>(input_ids->ne[0]));
+        input_ids      = to_backend(input_ids);
+        attention_mask = to_backend(attention_mask);
+
+        relative_position_bucket_vec = compute_relative_position_bucket(input_ids->ne[0], input_ids->ne[0]);
 
         // for (int i = 0; i < relative_position_bucket_vec.size(); i++) {
         //     if (i % 77 == 0) {
@@ -812,21 +812,23 @@ struct T5Runner : public GGMLRunner {
                                                            input_ids->ne[0]);
         set_backend_tensor_data(relative_position_bucket, relative_position_bucket_vec.data());
 
-        auto runner_ctx            = get_context();
-        ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
+        auto runner_ctx                   = get_context();
+        struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
 
         ggml_build_forward_expand(gf, hidden_states);
 
         return gf;
     }
 
-    sd::Tensor<float> compute(const int n_threads,
-                              const sd::Tensor<int32_t>& input_ids,
-                              const sd::Tensor<float>& attention_mask) {
-        auto get_graph = [&]() -> ggml_cgraph* {
+    bool compute(const int n_threads,
+                 struct ggml_tensor* input_ids,
+                 struct ggml_tensor* attention_mask,
+                 ggml_tensor** output,
+                 ggml_context* output_ctx = nullptr) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true), 3);
+        return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
     }
 
     static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
@@ -909,7 +911,7 @@ struct T5Embedder {
         : model(backend, offload_params_to_cpu, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
@@ -959,16 +961,17 @@ struct T5Embedder {
     }
 
     void test() {
-        ggml_init_params params;
+        struct ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        ggml_context* ctx = ggml_init(params);
-        GGML_ASSERT(ctx != nullptr);
+        struct ggml_context* work_ctx = ggml_init(params);
+        GGML_ASSERT(work_ctx != nullptr);
 
         {
             std::string text("a lovely cat");
+            // std::string text("一只可爱的猫"); // umt5 chinease test
             auto tokens_and_weights     = tokenize(text, 512, true);
             std::vector<int>& tokens    = std::get<0>(tokens_and_weights);
             std::vector<float>& weights = std::get<1>(tokens_and_weights);
@@ -977,18 +980,16 @@ struct T5Embedder {
                 printf("%d ", token);
             }
             printf("\n");
-            auto input_ids      = sd::Tensor<int32_t>::from_vector(tokens);
-            auto attention_mask = sd::Tensor<float>::from_vector(masks);
-            sd::Tensor<float> out;
+            auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
+            auto attention_mask     = vector_to_ggml_tensor(work_ctx, masks);
+            struct ggml_tensor* out = nullptr;
 
-            int64_t t0   = ggml_time_ms();
-            auto out_opt = model.compute(8, input_ids, attention_mask);
-            int64_t t1   = ggml_time_ms();
+            int t0 = ggml_time_ms();
+            model.compute(8, input_ids, attention_mask, &out, work_ctx);
+            int t1 = ggml_time_ms();
 
-            GGML_ASSERT(!out_opt.empty());
-            out = std::move(out_opt);
-            print_sd_tensor(out);
-            LOG_DEBUG("t5 test done in %lldms", t1 - t0);
+            print_ggml_tensor(out);
+            LOG_DEBUG("t5 test done in %dms", t1 - t0);
         }
     }
 

tae.hpp

@@ -0,0 +1,263 @@
+#ifndef __TAE_HPP__
+#define __TAE_HPP__
+
+#include "ggml_extend.hpp"
+
+#include "model.h"
+
+/*
+    ===================================    TinyAutoEncoder  ===================================
+    References:
+    https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py
+    https://github.com/madebyollin/taesd/blob/main/taesd.py
+
+*/
+
+class TAEBlock : public UnaryBlock {
+protected:
+    int n_in;
+    int n_out;
+
+public:
+    TAEBlock(int n_in, int n_out)
+        : n_in(n_in), n_out(n_out) {
+        blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
+        if (n_in != n_out) {
+            blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {1, 1}, {1, 1}, {1, 1}, {1, 1}, false));
+        }
+    }
+
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+        // x: [n, n_in, h, w]
+        // return: [n, n_out, h, w]
+
+        auto conv_0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
+        auto conv_2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
+        auto conv_4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
+
+        auto h = conv_0->forward(ctx, x);
+        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
+        h      = conv_2->forward(ctx, h);
+        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
+        h      = conv_4->forward(ctx, h);
+
+        if (n_in != n_out) {
+            auto skip = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
+            LOG_DEBUG("skip");
+            x = skip->forward(ctx, x);
+        }
+
+        h = ggml_add(ctx->ggml_ctx, h, x);
+        h = ggml_relu_inplace(ctx->ggml_ctx, h);
+        return h;
+    }
+};
+
+class TinyEncoder : public UnaryBlock {
+    int in_channels = 3;
+    int channels    = 64;
+    int z_channels  = 4;
+    int num_blocks  = 3;
+
+public:
+    TinyEncoder(int z_channels = 4)
+        : z_channels(z_channels) {
+        int index                       = 0;
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, channels, {3, 3}, {1, 1}, {1, 1}));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
+    }
+
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+        // x: [n, in_channels, h, w]
+        // return: [n, z_channels, h/8, w/8]
+
+        for (int i = 0; i < num_blocks * 3 + 6; i++) {
+            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
+
+            x = block->forward(ctx, x);
+        }
+
+        return x;
+    }
+};
+
+class TinyDecoder : public UnaryBlock {
+    int z_channels   = 4;
+    int channels     = 64;
+    int out_channels = 3;
+    int num_blocks   = 3;
+
+public:
+    TinyDecoder(int z_channels = 4)
+        : z_channels(z_channels) {
+        int index = 0;
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels, {3, 3}, {1, 1}, {1, 1}));
+        index++;  // nn.ReLU()
+
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+        index++;  // nn.Upsample()
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+        index++;  // nn.Upsample()
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+
+        for (int i = 0; i < num_blocks; i++) {
+            blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        }
+        index++;  // nn.Upsample()
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
+    }
+
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) override {
+        // z: [n, z_channels, h, w]
+        // return: [n, out_channels, h*8, w*8]
+
+        auto h = ggml_scale(ctx->ggml_ctx, z, 1.0f / 3.0f);
+        h      = ggml_tanh_inplace(ctx->ggml_ctx, h);
+        h      = ggml_scale(ctx->ggml_ctx, h, 3.0f);
+
+        for (int i = 0; i < num_blocks * 3 + 10; i++) {
+            if (blocks.find(std::to_string(i)) == blocks.end()) {
+                if (i == 1) {
+                    h = ggml_relu_inplace(ctx->ggml_ctx, h);
+                } else {
+                    h = ggml_upscale(ctx->ggml_ctx, h, 2, GGML_SCALE_MODE_NEAREST);
+                }
+                continue;
+            }
+            auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
+
+            h = block->forward(ctx, h);
+        }
+
+        return h;
+    }
+};
+
+class TAESD : public GGMLBlock {
+protected:
+    bool decode_only;
+
+public:
+    TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
+        : decode_only(decode_only) {
+        int z_channels = 4;
+        if (sd_version_is_dit(version)) {
+            z_channels = 16;
+        }
+        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels));
+
+        if (!decode_only) {
+            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels));
+        }
+    }
+
+    struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
+        auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
+        return decoder->forward(ctx, z);
+    }
+
+    struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
+        return encoder->forward(ctx, x);
+    }
+};
+
+struct TinyAutoEncoder : public GGMLRunner {
+    TAESD taesd;
+    bool decode_only = false;
+
+    TinyAutoEncoder(ggml_backend_t backend,
+                    bool offload_params_to_cpu,
+                    const String2TensorStorage& tensor_storage_map,
+                    const std::string prefix,
+                    bool decoder_only = true,
+                    SDVersion version = VERSION_SD1)
+        : decode_only(decoder_only),
+          taesd(decoder_only, version),
+          GGMLRunner(backend, offload_params_to_cpu) {
+        taesd.init(params_ctx, tensor_storage_map, prefix);
+    }
+
+    std::string get_desc() override {
+        return "taesd";
+    }
+
+    bool load_from_file(const std::string& file_path, int n_threads) {
+        LOG_INFO("loading taesd from '%s', decode_only = %s", file_path.c_str(), decode_only ? "true" : "false");
+        alloc_params_buffer();
+        std::map<std::string, ggml_tensor*> taesd_tensors;
+        taesd.get_param_tensors(taesd_tensors);
+        std::set<std::string> ignore_tensors;
+        if (decode_only) {
+            ignore_tensors.insert("encoder.");
+        }
+
+        ModelLoader model_loader;
+        if (!model_loader.init_from_file_and_convert_name(file_path)) {
+            LOG_ERROR("init taesd model loader from file failed: '%s'", file_path.c_str());
+            return false;
+        }
+
+        bool success = model_loader.load_tensors(taesd_tensors, ignore_tensors, n_threads);
+
+        if (!success) {
+            LOG_ERROR("load tae tensors from model loader failed");
+            return false;
+        }
+
+        LOG_INFO("taesd model loaded");
+        return success;
+    }
+
+    struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
+        struct ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        z                       = to_backend(z);
+        auto runner_ctx         = get_context();
+        struct ggml_tensor* out = decode_graph ? taesd.decode(&runner_ctx, z) : taesd.encode(&runner_ctx, z);
+        ggml_build_forward_expand(gf, out);
+        return gf;
+    }
+
+    bool compute(const int n_threads,
+                 struct ggml_tensor* z,
+                 bool decode_graph,
+                 struct ggml_tensor** output,
+                 struct ggml_context* output_ctx = nullptr) {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
+            return build_graph(z, decode_graph);
+        };
+
+        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+    }
+};
+
+#endif  // __TAE_HPP__

thirdparty/darts.h

@@ -845,7 +845,7 @@ inline void BitVector::build() {
 
   num_ones_ = 0;
   for (std::size_t i = 0; i < units_.size(); ++i) {
-    ranks_[i] = static_cast<id_type>(num_ones_);
+    ranks_[i] = num_ones_;
     num_ones_ += pop_count(units_[i]);
   }
 }
@@ -1769,7 +1769,7 @@ id_type DoubleArrayBuilder::arrange_from_keyset(const Keyset<T> &keyset,
 
 inline id_type DoubleArrayBuilder::find_valid_offset(id_type id) const {
   if (extras_head_ >= units_.size()) {
-    return static_cast<id_type>(units_.size()) | (id & LOWER_MASK);
+    return units_.size() | (id & LOWER_MASK);
   }
 
   id_type unfixed_id = extras_head_;
@@ -1781,7 +1781,7 @@ inline id_type DoubleArrayBuilder::find_valid_offset(id_type id) const {
     unfixed_id = extras(unfixed_id).next();
   } while (unfixed_id != extras_head_);
 
-  return static_cast<id_type>(units_.size()) | (id & LOWER_MASK);
+  return units_.size() | (id & LOWER_MASK);
 }
 
 inline bool DoubleArrayBuilder::is_valid_offset(id_type id,
@@ -1812,7 +1812,7 @@ inline void DoubleArrayBuilder::reserve_id(id_type id) {
   if (id == extras_head_) {
     extras_head_ = extras(id).next();
     if (extras_head_ == id) {
-      extras_head_ = static_cast<id_type>(units_.size());
+      extras_head_ = units_.size();
     }
   }
   extras(extras(id).prev()).set_next(extras(id).next());
@@ -1821,8 +1821,8 @@ inline void DoubleArrayBuilder::reserve_id(id_type id) {
 }
 
 inline void DoubleArrayBuilder::expand_units() {
-  id_type src_num_units = static_cast<id_type>(units_.size());
-  id_type src_num_blocks = static_cast<id_type>(num_blocks());
+  id_type src_num_units = units_.size();
+  id_type src_num_blocks = num_blocks();
 
   id_type dest_num_units = src_num_units + BLOCK_SIZE;
   id_type dest_num_blocks = src_num_blocks + 1;
@@ -1834,7 +1834,7 @@ inline void DoubleArrayBuilder::expand_units() {
   units_.resize(dest_num_units);
 
   if (dest_num_blocks > NUM_EXTRA_BLOCKS) {
-    for (id_type id = src_num_units; id < dest_num_units; ++id) {
+    for (std::size_t id = src_num_units; id < dest_num_units; ++id) {
       extras(id).set_is_used(false);
       extras(id).set_is_fixed(false);
     }
@@ -1858,9 +1858,9 @@ inline void DoubleArrayBuilder::expand_units() {
 inline void DoubleArrayBuilder::fix_all_blocks() {
   id_type begin = 0;
   if (num_blocks() > NUM_EXTRA_BLOCKS) {
-    begin = static_cast<id_type>(num_blocks() - NUM_EXTRA_BLOCKS);
+    begin = num_blocks() - NUM_EXTRA_BLOCKS;
   }
-  id_type end = static_cast<id_type>(num_blocks());
+  id_type end = num_blocks();
 
   for (id_type block_id = begin; block_id != end; ++block_id) {
     fix_block(block_id);

thirdparty/stb_image_write.h

@@ -257,10 +257,6 @@ int stbi_write_tga_with_rle = 1;
 int stbi_write_force_png_filter = -1;
 #endif
 
-#ifndef STBMIN
-#define STBMIN(a, b) ((a) < (b) ? (a) : (b))
-#endif // STBMIN
-
 static int stbi__flip_vertically_on_write = 0;
 
 STBIWDEF void stbi_flip_vertically_on_write(int flag)
@@ -1183,8 +1179,8 @@ STBIWDEF unsigned char *stbi_write_png_to_mem(const unsigned char *pixels, int s
    if (!zlib) return 0;
 
    if(parameters != NULL) {
-      param_length = (int)strlen(parameters);
-      param_length += (int)strlen("parameters") + 1; // For the name and the null-byte
+      param_length = strlen(parameters);
+      param_length += strlen("parameters") + 1; // For the name and the null-byte
    }
 
    // each tag requires 12 bytes of overhead
@@ -1530,11 +1526,11 @@ static int stbi_write_jpg_core(stbi__write_context *s, int width, int height, in
       if(parameters != NULL) {
          stbiw__putc(s, 0xFF /* comnent */ );
          stbiw__putc(s, 0xFE /* marker  */ );
-         int param_length = STBMIN(2 + (int)strlen("parameters") + 1 + (int)strlen(parameters) + 1, 0xFFFF);
+         size_t param_length = std::min(2 + strlen("parameters") + 1 + strlen(parameters) + 1, (size_t) 0xFFFF);
          stbiw__putc(s, param_length >> 8); // no need to mask, length < 65536
          stbiw__putc(s, param_length & 0xFF);
-         s->func(s->context, (void*)"parameters", (int)strlen("parameters") + 1); // std::string is zero-terminated
-         s->func(s->context, (void*)parameters, STBMIN(param_length, 65534) - 2 - (int)strlen("parameters") - 1);
+         s->func(s->context, (void*)"parameters", strlen("parameters") + 1); // std::string is zero-terminated
+         s->func(s->context, (void*)parameters, std::min(param_length, (size_t) 65534) - 2 - strlen("parameters") - 1);
          if(param_length > 65534) stbiw__putc(s, 0); // always zero-terminate for safety
          if(param_length & 1) stbiw__putc(s, 0xFF); // pad to even length
       }

tokenize_util.cpp

@@ -1,4 +1,4 @@
-#include <algorithm>
+#include <algorithm>
 #include <iostream>
 #include <string>
 #include <vector>
@@ -919,21 +919,15 @@ std::vector<std::string> token_split(const std::string& text) {
 
         // `\s*[\r\n]+|\s+(?!\S)|\s+`
         if (is_space(cp)) {
-            std::string token;
-            bool saw_new_line = false;
+            std::string token = codepoint_to_utf8(cp);
+            ++i;
 
             while (i < cps.size() && is_space(cps[i])) {
                 token += codepoint_to_utf8(cps[i]);
-
-                if (cps[i] == U'\r' || cps[i] == U'\n') {
-                    saw_new_line = true;
-                } else {
-                    if (saw_new_line) {
-                        break;
-                    }
-                }
-
                 ++i;
+                if (cps[i] == U'\r' || cps[i] == U'\n') {
+                    break;
+                }
             }
 
             tokens.push_back(token);
@@ -982,7 +976,7 @@ std::vector<std::string> split_with_special_tokens(
 }
 
 // int main() {
-//     std::string text = "I'm testing C++ token_split function. Hello world 123";
+//     std::string text = "I'm testing C++ token_split function. 你好，世界! 123";
 //     auto tokens = token_split(text);
 
 //     for (const auto& t : tokens) {

unet.hpp

@@ -1,7 +1,8 @@
 #ifndef __UNET_HPP__
 #define __UNET_HPP__
 
-#include "common_block.hpp"
+#include "common.hpp"
+#include "ggml_extend.hpp"
 #include "model.h"
 
 /*==================================================== UnetModel =====================================================*/
@@ -11,7 +12,7 @@
 class SpatialVideoTransformer : public SpatialTransformer {
 protected:
     int64_t time_depth;
-    int max_time_embed_period;
+    int64_t max_time_embed_period;
 
 public:
     SpatialVideoTransformer(int64_t in_channels,
@@ -20,8 +21,8 @@ public:
                             int64_t depth,
                             int64_t context_dim,
                             bool use_linear,
-                            int64_t time_depth        = 1,
-                            int max_time_embed_period = 10000)
+                            int64_t time_depth            = 1,
+                            int64_t max_time_embed_period = 10000)
         : SpatialTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear),
           max_time_embed_period(max_time_embed_period) {
         // We will convert unet transformer linear to conv2d 1x1 when loading the weights, so use_linear is always False
@@ -60,10 +61,10 @@ public:
         blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* context,
-                         int timesteps) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* context,
+                                int timesteps) {
         // x: [N, in_channels, h, w] aka [b*t, in_channels, h, w], t == timesteps
         // context: [N, max_position(aka n_context), hidden_size(aka context_dim)] aka [b*t, n_context, context_dim], t == timesteps
         // t_emb: [N, in_channels] aka [b*t, in_channels]
@@ -111,9 +112,9 @@ public:
         x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));  // [N, h, w, inner_dim]
         x = ggml_reshape_3d(ctx->ggml_ctx, x, inner_dim, w * h, n);                // [N, h * w, inner_dim]
 
-        auto num_frames = ggml_arange(ctx->ggml_ctx, 0.f, static_cast<float>(timesteps), 1.f);
+        auto num_frames = ggml_arange(ctx->ggml_ctx, 0, timesteps, 1);
         // since b is 1, no need to do repeat
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, num_frames, static_cast<int>(in_channels), max_time_embed_period);  // [N, in_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, num_frames, in_channels, max_time_embed_period);  // [N, in_channels]
 
         auto emb = time_pos_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);
@@ -200,9 +201,6 @@ public:
             num_head_channels     = 64;
             num_heads             = -1;
             use_linear_projection = true;
-            if (version == VERSION_SDXL_VEGA) {
-                transformer_depth = {1, 1, 2};
-            }
         } else if (version == VERSION_SVD) {
             in_channels           = 8;
             out_channels          = 4;
@@ -217,13 +215,10 @@ public:
         } else if (sd_version_is_unet_edit(version)) {
             in_channels = 8;
         }
-        if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS) {
+        if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET) {
             num_res_blocks = 1;
             channel_mult   = {1, 2, 4};
             tiny_unet      = true;
-            if (version == VERSION_SDXS) {
-                attention_resolutions = {4, 2};  // here just like SDXL
-            }
         }
 
         // dims is always 2
@@ -321,7 +316,7 @@ public:
         }
         if (!tiny_unet) {
             blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
-            if (version != VERSION_SDXL_SSD1B && version != VERSION_SDXL_VEGA) {
+            if (version != VERSION_SDXL_SSD1B) {
                 blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
                                                                                           n_head,
                                                                                           d_head,
@@ -388,11 +383,11 @@ public:
         blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    ggml_tensor* resblock_forward(std::string name,
-                                  GGMLRunnerContext* ctx,
-                                  ggml_tensor* x,
-                                  ggml_tensor* emb,
-                                  int num_video_frames) {
+    struct ggml_tensor* resblock_forward(std::string name,
+                                         GGMLRunnerContext* ctx,
+                                         struct ggml_tensor* x,
+                                         struct ggml_tensor* emb,
+                                         int num_video_frames) {
         if (version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
 
@@ -404,11 +399,11 @@ public:
         }
     }
 
-    ggml_tensor* attention_layer_forward(std::string name,
-                                         GGMLRunnerContext* ctx,
-                                         ggml_tensor* x,
-                                         ggml_tensor* context,
-                                         int timesteps) {
+    struct ggml_tensor* attention_layer_forward(std::string name,
+                                                GGMLRunnerContext* ctx,
+                                                struct ggml_tensor* x,
+                                                struct ggml_tensor* context,
+                                                int timesteps) {
         if (version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
 
@@ -420,15 +415,15 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x,
-                         ggml_tensor* timesteps,
-                         ggml_tensor* context,
-                         ggml_tensor* c_concat              = nullptr,
-                         ggml_tensor* y                     = nullptr,
-                         int num_video_frames               = -1,
-                         std::vector<ggml_tensor*> controls = {},
-                         float control_strength             = 0.f) {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x,
+                                struct ggml_tensor* timesteps,
+                                struct ggml_tensor* context,
+                                struct ggml_tensor* c_concat              = nullptr,
+                                struct ggml_tensor* y                     = nullptr,
+                                int num_video_frames                      = -1,
+                                std::vector<struct ggml_tensor*> controls = {},
+                                float control_strength                    = 0.f) {
         // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
         // timesteps: [N,]
         // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
@@ -480,7 +475,7 @@ public:
         }
 
         // input_blocks
-        std::vector<ggml_tensor*> hs;
+        std::vector<struct ggml_tensor*> hs;
 
         // input block 0
         auto h = input_blocks_0_0->forward(ctx, x);
@@ -522,16 +517,16 @@ public:
         // middle_block
         if (!tiny_unet) {
             h = resblock_forward("middle_block.0", ctx, h, emb, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
-            if (version != VERSION_SDXL_SSD1B && version != VERSION_SDXL_VEGA) {
+            if (version != VERSION_SDXL_SSD1B) {
                 h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
                 h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
             }
         }
         if (controls.size() > 0) {
-            auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true);
+            auto cs = ggml_scale_inplace(ctx->ggml_ctx, controls[controls.size() - 1], control_strength);
             h       = ggml_add(ctx->ggml_ctx, h, cs);  // middle control
         }
-        int control_offset = static_cast<int>(controls.size() - 2);
+        int control_offset = controls.size() - 2;
 
         // output_blocks
         int output_block_idx = 0;
@@ -541,7 +536,7 @@ public:
                 hs.pop_back();
 
                 if (controls.size() > 0) {
-                    auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, true);
+                    auto cs = ggml_scale_inplace(ctx->ggml_ctx, controls[control_offset], control_strength);
                     h_skip  = ggml_add(ctx->ggml_ctx, h_skip, cs);  // control net condition
                     control_offset--;
                 }
@@ -605,81 +600,82 @@ struct UNetModelRunner : public GGMLRunner {
         return "unet";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         unet.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 sd::Tensor<float>& c_concat_tensor              = {},
-                             const sd::Tensor<float>& y_tensor                     = {},
-                             int num_video_frames                                  = -1,
-                             const std::vector<sd::Tensor<float>>& controls_tensor = {},
-                             float control_strength                                = 0.f) {
-        ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
-
-        ggml_tensor* x         = make_input(x_tensor);
-        ggml_tensor* timesteps = make_input(timesteps_tensor);
-        ggml_tensor* context   = make_optional_input(context_tensor);
-        ggml_tensor* c_concat  = make_optional_input(c_concat_tensor);
-        ggml_tensor* y         = make_optional_input(y_tensor);
-        std::vector<ggml_tensor*> controls;
-        controls.reserve(controls_tensor.size());
-        for (const auto& control_tensor : controls_tensor) {
-            controls.push_back(make_input(control_tensor));
-        }
+    struct ggml_cgraph* build_graph(struct ggml_tensor* x,
+                                    struct ggml_tensor* timesteps,
+                                    struct ggml_tensor* context,
+                                    struct ggml_tensor* c_concat              = nullptr,
+                                    struct ggml_tensor* y                     = nullptr,
+                                    int num_video_frames                      = -1,
+                                    std::vector<struct ggml_tensor*> controls = {},
+                                    float control_strength                    = 0.f) {
+        struct ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
 
         if (num_video_frames == -1) {
-            num_video_frames = static_cast<int>(x->ne[3]);
+            num_video_frames = x->ne[3];
+        }
+
+        x         = to_backend(x);
+        context   = to_backend(context);
+        y         = to_backend(y);
+        timesteps = to_backend(timesteps);
+        c_concat  = to_backend(c_concat);
+
+        for (int i = 0; i < controls.size(); i++) {
+            controls[i] = to_backend(controls[i]);
         }
 
         auto runner_ctx = get_context();
 
-        ggml_tensor* out = unet.forward(&runner_ctx,
-                                        x,
-                                        timesteps,
-                                        context,
-                                        c_concat,
-                                        y,
-                                        num_video_frames,
-                                        controls,
-                                        control_strength);
+        struct ggml_tensor* out = unet.forward(&runner_ctx,
+                                               x,
+                                               timesteps,
+                                               context,
+                                               c_concat,
+                                               y,
+                                               num_video_frames,
+                                               controls,
+                                               control_strength);
 
         ggml_build_forward_expand(gf, out);
 
         return gf;
     }
 
-    sd::Tensor<float> compute(int n_threads,
-                              const sd::Tensor<float>& x,
-                              const sd::Tensor<float>& timesteps,
-                              const sd::Tensor<float>& context               = {},
-                              const sd::Tensor<float>& c_concat              = {},
-                              const sd::Tensor<float>& y                     = {},
-                              int num_video_frames                           = -1,
-                              const std::vector<sd::Tensor<float>>& controls = {},
-                              float control_strength                         = 0.f) {
+    bool compute(int n_threads,
+                 struct ggml_tensor* x,
+                 struct ggml_tensor* timesteps,
+                 struct ggml_tensor* context,
+                 struct ggml_tensor* c_concat,
+                 struct ggml_tensor* y,
+                 int num_video_frames                      = -1,
+                 std::vector<struct ggml_tensor*> controls = {},
+                 float control_strength                    = 0.f,
+                 struct ggml_tensor** output               = nullptr,
+                 struct ggml_context* output_ctx           = nullptr) {
         // x: [N, in_channels, h, w]
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
         // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
-        auto get_graph = [&]() -> ggml_cgraph* {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
     }
 
     void test() {
-        ggml_init_params params;
+        struct ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        ggml_context* ctx = ggml_init(params);
-        GGML_ASSERT(ctx != nullptr);
+        struct ggml_context* work_ctx = ggml_init(params);
+        GGML_ASSERT(work_ctx != nullptr);
 
         {
             // CPU, num_video_frames = 1, x{num_video_frames, 8, 8, 8}: Pass
@@ -688,38 +684,28 @@ struct UNetModelRunner : public GGMLRunner {
             // CUDA, num_video_frames = 3, x{num_video_frames, 8, 8, 8}: nan
             int num_video_frames = 3;
 
-            sd::Tensor<float> x({8, 8, 8, num_video_frames});
+            auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 8, num_video_frames);
             std::vector<float> timesteps_vec(num_video_frames, 999.f);
-            auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
-            x.fill_(0.5f);
+            auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
+            ggml_set_f32(x, 0.5f);
             // print_ggml_tensor(x);
 
-            sd::Tensor<float> context({1024, 1, num_video_frames});
-            context.fill_(0.5f);
+            auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 1024, 1, num_video_frames);
+            ggml_set_f32(context, 0.5f);
             // print_ggml_tensor(context);
 
-            sd::Tensor<float> y({768, num_video_frames});
-            y.fill_(0.5f);
+            auto y = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 768, num_video_frames);
+            ggml_set_f32(y, 0.5f);
             // print_ggml_tensor(y);
 
-            sd::Tensor<float> out;
+            struct ggml_tensor* out = nullptr;
 
-            int64_t t0   = ggml_time_ms();
-            auto out_opt = compute(8,
-                                   x,
-                                   timesteps,
-                                   context,
-                                   {},
-                                   y,
-                                   num_video_frames,
-                                   {},
-                                   0.f);
-            int64_t t1   = ggml_time_ms();
+            int t0 = ggml_time_ms();
+            compute(8, x, timesteps, context, nullptr, y, num_video_frames, {}, 0.f, &out, work_ctx);
+            int t1 = ggml_time_ms();
 
-            GGML_ASSERT(!out_opt.empty());
-            out = std::move(out_opt);
-            print_sd_tensor(out);
-            LOG_DEBUG("unet test done in %lldms", t1 - t0);
+            print_ggml_tensor(out);
+            LOG_DEBUG("unet test done in %dms", t1 - t0);
         }
     }
 };

upscaler.cpp

@@ -2,7 +2,6 @@
 #include "ggml_extend.hpp"
 #include "model.h"
 #include "stable-diffusion.h"
-#include "util.h"
 
 struct UpscalerGGML {
     ggml_backend_t backend    = nullptr;  // general backend
@@ -65,39 +64,6 @@ struct UpscalerGGML {
         return true;
     }
 
-    sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor) {
-        sd::Tensor<float> upscaled;
-        if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
-            upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
-        } else {
-            auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
-                auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
-                if (output_tile.empty()) {
-                    LOG_ERROR("esrgan compute failed while processing a tile");
-                    return {};
-                }
-                return output_tile;
-            };
-
-            upscaled = process_tiles_2d(input_tensor,
-                                        static_cast<int>(input_tensor.shape()[0] * esrgan_upscaler->scale),
-                                        static_cast<int>(input_tensor.shape()[1] * esrgan_upscaler->scale),
-                                        esrgan_upscaler->scale,
-                                        tile_size,
-                                        tile_size,
-                                        0.25f,
-                                        false,
-                                        false,
-                                        on_processing);
-        }
-        esrgan_upscaler->free_compute_buffer();
-        if (upscaled.empty()) {
-            LOG_ERROR("esrgan compute failed");
-            return {};
-        }
-        return upscaled;
-    }
-
     sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor) {
         // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
         sd_image_t upscaled_image = {0, 0, 0, nullptr};
@@ -106,17 +72,39 @@ struct UpscalerGGML {
         LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
                  input_image.width, input_image.height, output_width, output_height);
 
-        sd::Tensor<float> input_tensor = sd_image_to_tensor(input_image);
-        sd::Tensor<float> upscaled;
-        int64_t t0 = ggml_time_ms();
-        upscaled   = upscale_tensor(input_tensor);
-        if (upscaled.empty()) {
+        struct ggml_init_params params;
+        params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
+        params.mem_buffer = nullptr;
+        params.no_alloc   = false;
+
+        // draft context
+        struct ggml_context* upscale_ctx = ggml_init(params);
+        if (!upscale_ctx) {
+            LOG_ERROR("ggml_init() failed");
             return upscaled_image;
         }
-        sd_image_t upscaled_data = tensor_to_sd_image(upscaled);
-        int64_t t3               = ggml_time_ms();
+        // LOG_DEBUG("upscale work buffer size: %.2f MB", params.mem_size / 1024.f / 1024.f);
+        ggml_tensor* input_image_tensor = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, input_image.width, input_image.height, 3, 1);
+        sd_image_to_ggml_tensor(input_image, input_image_tensor);
+
+        ggml_tensor* upscaled = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, output_width, output_height, 3, 1);
+        auto on_tiling        = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
+            esrgan_upscaler->compute(n_threads, in, &out);
+        };
+        int64_t t0 = ggml_time_ms();
+        sd_tiling(input_image_tensor, upscaled, esrgan_upscaler->scale, esrgan_upscaler->tile_size, 0.25f, on_tiling);
+        esrgan_upscaler->free_compute_buffer();
+        ggml_ext_tensor_clamp_inplace(upscaled, 0.f, 1.f);
+        uint8_t* upscaled_data = ggml_tensor_to_sd_image(upscaled);
+        ggml_free(upscale_ctx);
+        int64_t t3 = ggml_time_ms();
         LOG_INFO("input_image_tensor upscaled, taking %.2fs", (t3 - t0) / 1000.0f);
-        upscaled_image = upscaled_data;
+        upscaled_image = {
+            (uint32_t)output_width,
+            (uint32_t)output_height,
+            3,
+            upscaled_data,
+        };
         return upscaled_image;
     }
 };

util.cpp

@@ -95,71 +95,9 @@ bool is_directory(const std::string& path) {
     return (attributes != INVALID_FILE_ATTRIBUTES && (attributes & FILE_ATTRIBUTE_DIRECTORY));
 }
 
-class MmapWrapperImpl : public MmapWrapper {
-public:
-    MmapWrapperImpl(void* data, size_t size, HANDLE hfile, HANDLE hmapping)
-        : MmapWrapper(data, size), hfile_(hfile), hmapping_(hmapping) {}
-
-    ~MmapWrapperImpl() override {
-        UnmapViewOfFile(data_);
-        CloseHandle(hmapping_);
-        CloseHandle(hfile_);
-    }
-
-private:
-    HANDLE hfile_;
-    HANDLE hmapping_;
-};
-
-std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
-    void* mapped_data = nullptr;
-    size_t file_size  = 0;
-
-    HANDLE file_handle = CreateFileA(
-        filename.c_str(),
-        GENERIC_READ,
-        FILE_SHARE_READ,
-        NULL,
-        OPEN_EXISTING,
-        FILE_ATTRIBUTE_NORMAL,
-        NULL);
-
-    if (file_handle == INVALID_HANDLE_VALUE) {
-        return nullptr;
-    }
-
-    LARGE_INTEGER size;
-    if (!GetFileSizeEx(file_handle, &size)) {
-        CloseHandle(file_handle);
-        return nullptr;
-    }
-
-    file_size = static_cast<size_t>(size.QuadPart);
-
-    HANDLE mapping_handle = CreateFileMapping(file_handle, NULL, PAGE_READONLY, 0, 0, NULL);
-
-    if (mapping_handle == NULL) {
-        CloseHandle(file_handle);
-        return nullptr;
-    }
-
-    mapped_data = MapViewOfFile(mapping_handle, FILE_MAP_READ, 0, 0, file_size);
-
-    if (mapped_data == NULL) {
-        CloseHandle(mapping_handle);
-        CloseHandle(file_handle);
-        return nullptr;
-    }
-
-    return std::make_unique<MmapWrapperImpl>(mapped_data, file_size, file_handle, mapping_handle);
-}
-
 #else  // Unix
 #include <dirent.h>
-#include <fcntl.h>
-#include <sys/mman.h>
 #include <sys/stat.h>
-#include <unistd.h>
 
 bool file_exists(const std::string& filename) {
     struct stat buffer;
@@ -171,64 +109,8 @@ bool is_directory(const std::string& path) {
     return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-class MmapWrapperImpl : public MmapWrapper {
-public:
-    MmapWrapperImpl(void* data, size_t size)
-        : MmapWrapper(data, size) {}
-
-    ~MmapWrapperImpl() override {
-        munmap(data_, size_);
-    }
-};
-
-std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
-    int file_descriptor = open(filename.c_str(), O_RDONLY);
-    if (file_descriptor == -1) {
-        return nullptr;
-    }
-
-    int mmap_flags = MAP_PRIVATE;
-
-#ifdef __linux__
-    // performance flags used by llama.cpp
-    // posix_fadvise(file_descriptor, 0, 0, POSIX_FADV_SEQUENTIAL);
-    // mmap_flags |= MAP_POPULATE;
 #endif
 
-    struct stat sb;
-    if (fstat(file_descriptor, &sb) == -1) {
-        close(file_descriptor);
-        return nullptr;
-    }
-
-    size_t file_size = sb.st_size;
-
-    void* mapped_data = mmap(NULL, file_size, PROT_READ, mmap_flags, file_descriptor, 0);
-
-    close(file_descriptor);
-
-    if (mapped_data == MAP_FAILED) {
-        return nullptr;
-    }
-
-#ifdef __linux__
-    // performance flags used by llama.cpp
-    // posix_madvise(mapped_data, file_size, POSIX_MADV_WILLNEED);
-#endif
-
-    return std::make_unique<MmapWrapperImpl>(mapped_data, file_size);
-}
-
-#endif
-
-bool MmapWrapper::copy_data(void* buf, size_t n, size_t offset) const {
-    if (offset >= size_ || n > (size_ - offset)) {
-        return false;
-    }
-    std::memcpy(buf, data() + offset, n);
-    return true;
-}
-
 // get_num_physical_cores is copy from
 // https://github.com/ggerganov/llama.cpp/blob/master/examples/common.cpp
 // LICENSE: https://github.com/ggerganov/llama.cpp/blob/master/LICENSE
@@ -479,96 +361,158 @@ const char* sd_get_system_info() {
     return buffer;
 }
 
-sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index) {
-    const auto& shape = tensor.shape();
-    GGML_ASSERT(shape.size() == 4 || shape.size() == 5);
-    int width     = static_cast<int>(shape[0]);
-    int height    = static_cast<int>(shape[1]);
-    int channel   = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
-    uint8_t* data = (uint8_t*)malloc(static_cast<size_t>(width * height * channel));
-    GGML_ASSERT(data != nullptr);
+sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image) {
+    sd_image_f32_t converted_image;
+    converted_image.width   = image.width;
+    converted_image.height  = image.height;
+    converted_image.channel = image.channel;
 
-    for (int iw = 0; iw < width; ++iw) {
-        for (int ih = 0; ih < height; ++ih) {
-            for (int ic = 0; ic < channel; ++ic) {
-                float value                            = shape.size() == 5 ? tensor.index(iw, ih, frame_index, ic, 0)
-                                                                           : tensor.index(iw, ih, ic, frame_index);
-                value                                  = std::clamp(value, 0.0f, 1.0f);
-                data[(ih * width + iw) * channel + ic] = static_cast<uint8_t>(std::round(value * 255.0f));
-            }
-        }
+    // Allocate memory for float data
+    converted_image.data = (float*)malloc(image.width * image.height * image.channel * sizeof(float));
+
+    for (int i = 0; i < image.width * image.height * image.channel; i++) {
+        // Convert uint8_t to float
+        converted_image.data[i] = (float)image.data[i];
     }
-    return {
-        static_cast<uint32_t>(width),
-        static_cast<uint32_t>(height),
-        static_cast<uint32_t>(channel),
-        data,
-    };
+
+    return converted_image;
 }
 
-sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
-                                     int target_width,
-                                     int target_height,
-                                     bool scale) {
-    sd::Tensor<float> tensor = sd::zeros<float>({static_cast<int64_t>(image.width),
-                                                 static_cast<int64_t>(image.height),
-                                                 static_cast<int64_t>(image.channel),
-                                                 1});
-    for (uint32_t iw = 0; iw < image.width; ++iw) {
-        for (uint32_t ih = 0; ih < image.height; ++ih) {
-            for (uint32_t ic = 0; ic < image.channel; ++ic) {
-                tensor.index(iw, ih, ic, 0) = sd_image_get_f32(image, iw, ih, ic, scale);
+// Function to perform double linear interpolation
+float interpolate(float v1, float v2, float v3, float v4, float x_ratio, float y_ratio) {
+    return v1 * (1 - x_ratio) * (1 - y_ratio) + v2 * x_ratio * (1 - y_ratio) + v3 * (1 - x_ratio) * y_ratio + v4 * x_ratio * y_ratio;
+}
+
+sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height) {
+    sd_image_f32_t resized_image;
+    resized_image.width   = target_width;
+    resized_image.height  = target_height;
+    resized_image.channel = image.channel;
+
+    // Allocate memory for resized float data
+    resized_image.data = (float*)malloc(target_width * target_height * image.channel * sizeof(float));
+
+    for (int y = 0; y < target_height; y++) {
+        for (int x = 0; x < target_width; x++) {
+            float original_x = (float)x * image.width / target_width;
+            float original_y = (float)y * image.height / target_height;
+
+            uint32_t x1 = (uint32_t)original_x;
+            uint32_t y1 = (uint32_t)original_y;
+            uint32_t x2 = std::min(x1 + 1, image.width - 1);
+            uint32_t y2 = std::min(y1 + 1, image.height - 1);
+
+            for (int k = 0; k < image.channel; k++) {
+                float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
+                float v2 = *(image.data + y1 * image.width * image.channel + x2 * image.channel + k);
+                float v3 = *(image.data + y2 * image.width * image.channel + x1 * image.channel + k);
+                float v4 = *(image.data + y2 * image.width * image.channel + x2 * image.channel + k);
+
+                float x_ratio = original_x - x1;
+                float y_ratio = original_y - y1;
+
+                float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
+
+                *(resized_image.data + y * target_width * image.channel + x * image.channel + k) = value;
             }
         }
     }
-    if (target_width >= 0 && target_height >= 0 &&
-        (tensor.shape()[0] != target_width || tensor.shape()[1] != target_height)) {
-        tensor = sd::ops::interpolate(tensor,
-                                      {target_width,
-                                       target_height,
-                                       tensor.shape()[2],
-                                       tensor.shape()[3]});
+
+    return resized_image;
+}
+
+void normalize_sd_image_f32_t(sd_image_f32_t image, float means[3], float stds[3]) {
+    for (int y = 0; y < image.height; y++) {
+        for (int x = 0; x < image.width; x++) {
+            for (int k = 0; k < image.channel; k++) {
+                int index         = (y * image.width + x) * image.channel + k;
+                image.data[index] = (image.data[index] - means[k]) / stds[k];
+            }
+        }
     }
-    return tensor;
 }
 
 // Constants for means and std
-float means[3] = {0.48145466f, 0.4578275f, 0.40821073f};
-float stds[3]  = {0.26862954f, 0.26130258f, 0.27577711f};
+float means[3] = {0.48145466, 0.4578275, 0.40821073};
+float stds[3]  = {0.26862954, 0.26130258, 0.27577711};
 
-sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_width, int target_height) {
-    GGML_ASSERT(image.dim() == 4);
-    GGML_ASSERT(image.shape()[2] == 3);
-    GGML_ASSERT(image.shape()[3] == 1);
-    GGML_ASSERT(target_width > 0 && target_height > 0);
+// Function to clip and preprocess sd_image_f32_t
+sd_image_f32_t clip_preprocess(sd_image_f32_t image, int target_width, int target_height) {
+    float width_scale  = (float)target_width / image.width;
+    float height_scale = (float)target_height / image.height;
 
-    float width_scale  = static_cast<float>(target_width) / static_cast<float>(image.shape()[0]);
-    float height_scale = static_cast<float>(target_height) / static_cast<float>(image.shape()[1]);
-    float scale        = std::fmax(width_scale, height_scale);
+    float scale = std::fmax(width_scale, height_scale);
 
-    int64_t resized_width  = static_cast<int64_t>(scale * static_cast<float>(image.shape()[0]));
-    int64_t resized_height = static_cast<int64_t>(scale * static_cast<float>(image.shape()[1]));
+    // Interpolation
+    int resized_width   = (int)(scale * image.width);
+    int resized_height  = (int)(scale * image.height);
+    float* resized_data = (float*)malloc(resized_width * resized_height * image.channel * sizeof(float));
 
-    sd::Tensor<float> resized = sd::ops::interpolate(
-        image,
-        {resized_width, resized_height, image.shape()[2], image.shape()[3]});
+    for (int y = 0; y < resized_height; y++) {
+        for (int x = 0; x < resized_width; x++) {
+            float original_x = (float)x * image.width / resized_width;
+            float original_y = (float)y * image.height / resized_height;
 
-    int64_t h_offset = std::max<int64_t>((resized_height - target_height) / 2, 0);
-    int64_t w_offset = std::max<int64_t>((resized_width - target_width) / 2, 0);
+            uint32_t x1 = (uint32_t)original_x;
+            uint32_t y1 = (uint32_t)original_y;
+            uint32_t x2 = std::min(x1 + 1, image.width - 1);
+            uint32_t y2 = std::min(y1 + 1, image.height - 1);
 
-    sd::Tensor<float> cropped({target_width, target_height, image.shape()[2], image.shape()[3]});
-    for (int64_t y = 0; y < target_height; ++y) {
-        for (int64_t x = 0; x < target_width; ++x) {
-            for (int64_t c = 0; c < image.shape()[2]; ++c) {
-                cropped.index(x, y, c, 0) = resized.index(x + w_offset, y + h_offset, c, 0);
+            for (int k = 0; k < image.channel; k++) {
+                float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
+                float v2 = *(image.data + y1 * image.width * image.channel + x2 * image.channel + k);
+                float v3 = *(image.data + y2 * image.width * image.channel + x1 * image.channel + k);
+                float v4 = *(image.data + y2 * image.width * image.channel + x2 * image.channel + k);
+
+                float x_ratio = original_x - x1;
+                float y_ratio = original_y - y1;
+
+                float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
+
+                *(resized_data + y * resized_width * image.channel + x * image.channel + k) = value;
             }
         }
     }
 
-    sd::Tensor<float> normalized = sd::ops::clamp(cropped, 0.0f, 1.0f);
-    sd::Tensor<float> mean({1, 1, 3, 1}, {means[0], means[1], means[2]});
-    sd::Tensor<float> std({1, 1, 3, 1}, {stds[0], stds[1], stds[2]});
-    return (normalized - mean) / std;
+    // Clip and preprocess
+    int h_offset = std::max((int)(resized_height - target_height) / 2, 0);
+    int w_offset = std::max((int)(resized_width - target_width) / 2, 0);
+
+    sd_image_f32_t result;
+    result.width   = target_width;
+    result.height  = target_height;
+    result.channel = image.channel;
+    result.data    = (float*)malloc(target_height * target_width * image.channel * sizeof(float));
+
+    for (int k = 0; k < image.channel; k++) {
+        for (int i = 0; i < result.height; i++) {
+            for (int j = 0; j < result.width; j++) {
+                int src_y = std::min(i + h_offset, resized_height - 1);
+                int src_x = std::min(j + w_offset, resized_width - 1);
+                *(result.data + i * result.width * image.channel + j * image.channel + k) =
+                    fmin(fmax(*(resized_data + src_y * resized_width * image.channel + src_x * image.channel + k), 0.0f), 255.0f) / 255.0f;
+            }
+        }
+    }
+
+    // Free allocated memory
+    free(resized_data);
+
+    // Normalize
+    for (int k = 0; k < image.channel; k++) {
+        for (int i = 0; i < result.height; i++) {
+            for (int j = 0; j < result.width; j++) {
+                // *(result.data + i * size * image.channel + j * image.channel + k) = 0.5f;
+                int offset  = i * result.width * image.channel + j * image.channel + k;
+                float value = *(result.data + offset);
+                value       = (value - means[k]) / stds[k];
+                // value = 0.5f;
+                *(result.data + offset) = value;
+            }
+        }
+    }
+
+    return result;
 }
 
 // Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/prompt_parser.py#L345

util.h

@@ -2,12 +2,10 @@
 #define __UTIL_H__
 
 #include <cstdint>
-#include <memory>
 #include <string>
 #include <vector>
 
 #include "stable-diffusion.h"
-#include "tensor.hpp"
 
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
@@ -30,36 +28,20 @@ std::string utf32_to_utf8(const std::u32string& utf32_str);
 std::u32string unicode_value_to_utf32(int unicode_value);
 // std::string sd_basename(const std::string& path);
 
-sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index = 0);
+typedef struct {
+    uint32_t width;
+    uint32_t height;
+    uint32_t channel;
+    float* data;
+} sd_image_f32_t;
 
-sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
-                                     int target_width  = -1,
-                                     int target_height = -1,
-                                     bool scale        = true);
+void normalize_sd_image_f32_t(sd_image_f32_t image, float means[3], float stds[3]);
 
-sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_width, int target_height);
+sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image);
 
-class MmapWrapper {
-public:
-    static std::unique_ptr<MmapWrapper> create(const std::string& filename);
+sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height);
 
-    virtual ~MmapWrapper() = default;
-
-    MmapWrapper(const MmapWrapper&)            = delete;
-    MmapWrapper& operator=(const MmapWrapper&) = delete;
-    MmapWrapper(MmapWrapper&&)                 = delete;
-    MmapWrapper& operator=(MmapWrapper&&)      = delete;
-
-    const uint8_t* data() const { return static_cast<uint8_t*>(data_); }
-    size_t size() const { return size_; }
-    bool copy_data(void* buf, size_t n, size_t offset) const;
-
-protected:
-    MmapWrapper(void* data, size_t size)
-        : data_(data), size_(size) {}
-    void* data_  = nullptr;
-    size_t size_ = 0;
-};
+sd_image_f32_t clip_preprocess(sd_image_f32_t image, int target_width, int target_height);
 
 std::string path_join(const std::string& p1, const std::string& p2);
 std::vector<std::string> split_string(const std::string& str, char delimiter);

vae.hpp

@@ -1,7 +1,8 @@
-#ifndef __AUTO_ENCODER_KL_HPP__
-#define __AUTO_ENCODER_KL_HPP__
+#ifndef __VAE_HPP__
+#define __VAE_HPP__
 
-#include "vae.hpp"
+#include "common.hpp"
+#include "ggml_extend.hpp"
 
 /*================================================== AutoEncoderKL ===================================================*/
 
@@ -29,7 +30,7 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [N, in_channels, h, w]
         // t_emb is always None
         auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
@@ -65,7 +66,7 @@ protected:
     int64_t in_channels;
     bool use_linear;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
         auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
         if (iter != tensor_storage_map.end()) {
             if (iter->second.n_dims == 4 && use_linear) {
@@ -101,7 +102,7 @@ public:
         }
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [N, in_channels, h, w]
         auto norm     = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
         auto q_proj   = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
@@ -126,6 +127,8 @@ public:
             q = q_proj->forward(ctx, h_);  // [N, h * w, in_channels]
             k = k_proj->forward(ctx, h_);  // [N, h * w, in_channels]
             v = v_proj->forward(ctx, h_);  // [N, h * w, in_channels]
+
+            v = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));  // [N, in_channels, h * w]
         } else {
             q = q_proj->forward(ctx, h_);                                              // [N, in_channels, h, w]
             q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 1, 2, 0, 3));  // [N, h, w, in_channels]
@@ -135,12 +138,11 @@ public:
             k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 1, 2, 0, 3));  // [N, h, w, in_channels]
             k = ggml_reshape_3d(ctx->ggml_ctx, k, c, h * w, n);                        // [N, h * w, in_channels]
 
-            v = v_proj->forward(ctx, h_);                                              // [N, in_channels, h, w]
-            v = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3));  // [N, h, w, in_channels]
-            v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n);                        // [N, h * w, in_channels]
+            v = v_proj->forward(ctx, h_);                        // [N, in_channels, h, w]
+            v = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n);  // [N, in_channels, h * w]
         }
 
-        h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
+        h_ = ggml_ext_attention(ctx->ggml_ctx, q, k, v, false);  // [N, h * w, in_channels]
 
         if (use_linear) {
             h_ = proj_out->forward(ctx, h_);  // [N, h * w, in_channels]
@@ -164,27 +166,27 @@ public:
     AE3DConv(int64_t in_channels,
              int64_t out_channels,
              std::pair<int, int> kernel_size,
-             int video_kernel_size        = 3,
+             int64_t video_kernel_size    = 3,
              std::pair<int, int> stride   = {1, 1},
              std::pair<int, int> padding  = {0, 0},
              std::pair<int, int> dilation = {1, 1},
              bool bias                    = true)
         : Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias) {
-        int kernel_padding      = video_kernel_size / 2;
-        blocks["time_mix_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(out_channels,
-                                                                        out_channels,
-                                                                        {video_kernel_size, 1, 1},
-                                                                        {1, 1, 1},
-                                                                        {kernel_padding, 0, 0}));
+        int64_t kernel_padding  = video_kernel_size / 2;
+        blocks["time_mix_conv"] = std::shared_ptr<GGMLBlock>(new Conv3dnx1x1(out_channels,
+                                                                             out_channels,
+                                                                             video_kernel_size,
+                                                                             1,
+                                                                             kernel_padding));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                struct ggml_tensor* x) override {
         // timesteps always None
         // skip_video always False
         // x: [N, IC, IH, IW]
         // result: [N, OC, OH, OW]
-        auto time_mix_conv = std::dynamic_pointer_cast<Conv3d>(blocks["time_mix_conv"]);
+        auto time_mix_conv = std::dynamic_pointer_cast<Conv3dnx1x1>(blocks["time_mix_conv"]);
 
         x = Conv2d::forward(ctx, x);
         // timesteps = x.shape[0]
@@ -208,7 +210,7 @@ public:
 
 class VideoResnetBlock : public ResnetBlock {
 protected:
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         enum ggml_type wtype = get_type(prefix + "mix_factor", tensor_storage_map, GGML_TYPE_F32);
         params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
     }
@@ -227,7 +229,7 @@ public:
         blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, 0, out_channels, {video_kernel_size, 1}, 3, false, true));
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
         // x: [N, in_channels, h, w] aka [b*t, in_channels, h, w]
         // return: [N, out_channels, h, w] aka [b*t, out_channels, h, w]
         // t_emb is always None
@@ -252,8 +254,8 @@ public:
 
         float alpha = get_alpha();
         x           = ggml_add(ctx->ggml_ctx,
-                               ggml_ext_scale(ctx->ggml_ctx, x, alpha),
-                               ggml_ext_scale(ctx->ggml_ctx, x_mix, 1.0f - alpha));
+                               ggml_scale(ctx->ggml_ctx, x, alpha),
+                               ggml_scale(ctx->ggml_ctx, x_mix, 1.0f - alpha));
 
         x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // b c t (h w) -> b t c (h w)
         x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, C, T * B);                     // b t c (h w) -> (b t) c h w
@@ -317,7 +319,7 @@ public:
         blocks["conv_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(block_in, double_z ? z_channels * 2 : z_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, in_channels, h, w]
 
         auto conv_in     = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"]);
@@ -407,8 +409,8 @@ public:
           z_channels(z_channels),
           video_decoder(video_decoder),
           video_kernel_size(video_kernel_size) {
-        int num_resolutions = static_cast<int>(ch_mult.size());
-        int block_in        = ch * ch_mult[num_resolutions - 1];
+        size_t num_resolutions = ch_mult.size();
+        int block_in           = ch * ch_mult[num_resolutions - 1];
 
         blocks["conv_in"] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, block_in, {3, 3}, {1, 1}, {1, 1}));
 
@@ -435,7 +437,7 @@ public:
         blocks["conv_out"] = get_conv_out(block_in, out_ch, {3, 3}, {1, 1}, {1, 1});
     }
 
-    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) {
+    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
         // z: [N, z_channels, h, w]
         // alpha is always 0
         // merge_strategy is always learned
@@ -459,7 +461,7 @@ public:
         h = mid_block_2->forward(ctx, h);  // [N, block_in, h, w]
 
         // upsampling
-        int num_resolutions = static_cast<int>(ch_mult.size());
+        size_t num_resolutions = ch_mult.size();
         for (int i = num_resolutions - 1; i >= 0; i--) {
             for (int j = 0; j < num_res_blocks + 1; j++) {
                 std::string name = "up." + std::to_string(i) + ".block." + std::to_string(j);
@@ -483,7 +485,7 @@ public:
 };
 
 // ldm.models.autoencoder.AutoencoderKL
-class AutoEncoderKLModel : public GGMLBlock {
+class AutoencodingEngine : public GGMLBlock {
 protected:
     SDVersion version;
     bool decode_only       = true;
@@ -502,7 +504,7 @@ protected:
     } dd_config;
 
 public:
-    AutoEncoderKLModel(SDVersion version          = VERSION_SD1,
+    AutoencodingEngine(SDVersion version          = VERSION_SD1,
                        bool decode_only           = true,
                        bool use_linear_projection = false,
                        bool use_video_decoder     = false)
@@ -549,7 +551,7 @@ public:
         }
     }
 
-    ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
+    struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
         // z: [N, z_channels, h, w]
         if (sd_version_is_flux2(version)) {
             // [N, C*p*p, h, w] -> [N, C, h*p, w*p]
@@ -581,7 +583,7 @@ public:
         return h;
     }
 
-    ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
         // x: [N, in_channels, h, w]
         auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
 
@@ -610,21 +612,48 @@ public:
         }
         return z;
     }
+};
 
-    int get_encoder_output_channels() {
-        int factor = dd_config.double_z ? 2 : 1;
-        if (sd_version_is_flux2(version)) {
-            return dd_config.z_channels * 4;
+struct VAE : public GGMLRunner {
+    VAE(ggml_backend_t backend, bool offload_params_to_cpu)
+        : GGMLRunner(backend, offload_params_to_cpu) {}
+    virtual bool compute(const int n_threads,
+                         struct ggml_tensor* z,
+                         bool decode_graph,
+                         struct ggml_tensor** output,
+                         struct ggml_context* output_ctx)                                                         = 0;
+    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) = 0;
+    virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
+};
+
+struct FakeVAE : public VAE {
+    FakeVAE(ggml_backend_t backend, bool offload_params_to_cpu)
+        : VAE(backend, offload_params_to_cpu) {}
+    bool compute(const int n_threads,
+                 struct ggml_tensor* z,
+                 bool decode_graph,
+                 struct ggml_tensor** output,
+                 struct ggml_context* output_ctx) override {
+        if (*output == nullptr && output_ctx != nullptr) {
+            *output = ggml_dup_tensor(output_ctx, z);
         }
-        return dd_config.z_channels * factor;
+        ggml_ext_tensor_iter(z, [&](ggml_tensor* z, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+            float value = ggml_ext_tensor_get_f32(z, i0, i1, i2, i3);
+            ggml_ext_tensor_set_f32(*output, value, i0, i1, i2, i3);
+        });
+        return true;
+    }
+
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {}
+
+    std::string get_desc() override {
+        return "fake_vae";
     }
 };
 
 struct AutoEncoderKL : public VAE {
-    float scale_factor = 1.f;
-    float shift_factor = 0.f;
-    bool decode_only   = true;
-    AutoEncoderKLModel ae;
+    bool decode_only = true;
+    AutoencodingEngine ae;
 
     AutoEncoderKL(ggml_backend_t backend,
                   bool offload_params_to_cpu,
@@ -633,23 +662,7 @@ struct AutoEncoderKL : public VAE {
                   bool decode_only       = false,
                   bool use_video_decoder = false,
                   SDVersion version      = VERSION_SD1)
-        : decode_only(decode_only), VAE(version, backend, offload_params_to_cpu) {
-        if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
-            scale_factor = 0.18215f;
-            shift_factor = 0.f;
-        } else if (sd_version_is_sdxl(version)) {
-            scale_factor = 0.13025f;
-            shift_factor = 0.f;
-        } 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)) {
-            scale_factor = 0.3611f;
-            shift_factor = 0.1159f;
-        } else if (sd_version_is_flux2(version)) {
-            scale_factor = 1.0f;
-            shift_factor = 0.f;
-        }
+        : decode_only(decode_only), VAE(backend, offload_params_to_cpu) {
         bool use_linear_projection = false;
         for (const auto& [name, tensor_storage] : tensor_storage_map) {
             if (!starts_with(name, prefix)) {
@@ -662,7 +675,7 @@ struct AutoEncoderKL : public VAE {
                 break;
             }
         }
-        ae = AutoEncoderKLModel(version, decode_only, use_linear_projection, use_video_decoder);
+        ae = AutoencodingEngine(version, decode_only, use_linear_projection, use_video_decoder);
         ae.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -681,150 +694,63 @@ struct AutoEncoderKL : public VAE {
         return "vae";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
+    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {
         ae.get_param_tensors(tensors, prefix);
     }
 
-    ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
-        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
-        ggml_tensor* z  = make_input(z_tensor);
+    struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
+        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+
+        z = to_backend(z);
 
         auto runner_ctx = get_context();
 
-        ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
+        struct ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
 
         ggml_build_forward_expand(gf, out);
 
         return gf;
     }
 
-    sd::Tensor<float> _compute(const int n_threads,
-                               const sd::Tensor<float>& z,
-                               bool decode_graph) override {
+    bool compute(const int n_threads,
+                 struct ggml_tensor* z,
+                 bool decode_graph,
+                 struct ggml_tensor** output,
+                 struct ggml_context* output_ctx = nullptr) override {
         GGML_ASSERT(!decode_only || decode_graph);
-        auto get_graph = [&]() -> ggml_cgraph* {
+        auto get_graph = [&]() -> struct ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
-    }
-
-    sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {
-        // ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample
-        auto chunks               = sd::ops::chunk(moments, 2, 2);
-        const auto& mean          = chunks[0];
-        const auto& logvar        = chunks[1];
-        sd::Tensor<float> stddev  = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -30.0f, 20.0f));
-        sd::Tensor<float> noise   = sd::Tensor<float>::randn_like(mean, rng);
-        sd::Tensor<float> latents = mean + stddev * noise;
-        return latents;
-    }
-
-    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
-        if (sd_version_is_flux2(version)) {
-            return vae_output;
-        } else if (version == VERSION_SD1_PIX2PIX) {
-            return sd::ops::chunk(vae_output, 2, 2)[0];
-        } else {
-            return gaussian_latent_sample(vae_output, rng);
-        }
-    }
-
-    std::pair<sd::Tensor<float>, sd::Tensor<float>> get_latents_mean_std(const sd::Tensor<float>& latents, int channel_dim) {
-        GGML_ASSERT(channel_dim >= 0 && static_cast<size_t>(channel_dim) < static_cast<size_t>(latents.dim()));
-        if (sd_version_is_flux2(version)) {
-            GGML_ASSERT(latents.shape()[channel_dim] == 128);
-            std::vector<int64_t> stats_shape(static_cast<size_t>(latents.dim()), 1);
-            stats_shape[static_cast<size_t>(channel_dim)] = latents.shape()[channel_dim];
-
-            auto mean_tensor = sd::Tensor<float>::from_vector({-0.0676f, -0.0715f, -0.0753f, -0.0745f, 0.0223f, 0.0180f, 0.0142f, 0.0184f,
-                                                               -0.0001f, -0.0063f, -0.0002f, -0.0031f, -0.0272f, -0.0281f, -0.0276f, -0.0290f,
-                                                               -0.0769f, -0.0672f, -0.0902f, -0.0892f, 0.0168f, 0.0152f, 0.0079f, 0.0086f,
-                                                               0.0083f, 0.0015f, 0.0003f, -0.0043f, -0.0439f, -0.0419f, -0.0438f, -0.0431f,
-                                                               -0.0102f, -0.0132f, -0.0066f, -0.0048f, -0.0311f, -0.0306f, -0.0279f, -0.0180f,
-                                                               0.0030f, 0.0015f, 0.0126f, 0.0145f, 0.0347f, 0.0338f, 0.0337f, 0.0283f,
-                                                               0.0020f, 0.0047f, 0.0047f, 0.0050f, 0.0123f, 0.0081f, 0.0081f, 0.0146f,
-                                                               0.0681f, 0.0679f, 0.0767f, 0.0732f, -0.0462f, -0.0474f, -0.0392f, -0.0511f,
-                                                               -0.0528f, -0.0477f, -0.0470f, -0.0517f, -0.0317f, -0.0316f, -0.0345f, -0.0283f,
-                                                               0.0510f, 0.0445f, 0.0578f, 0.0458f, -0.0412f, -0.0458f, -0.0487f, -0.0467f,
-                                                               -0.0088f, -0.0106f, -0.0088f, -0.0046f, -0.0376f, -0.0432f, -0.0436f, -0.0499f,
-                                                               0.0118f, 0.0166f, 0.0203f, 0.0279f, 0.0113f, 0.0129f, 0.0016f, 0.0072f,
-                                                               -0.0118f, -0.0018f, -0.0141f, -0.0054f, -0.0091f, -0.0138f, -0.0145f, -0.0187f,
-                                                               0.0323f, 0.0305f, 0.0259f, 0.0300f, 0.0540f, 0.0614f, 0.0495f, 0.0590f,
-                                                               -0.0511f, -0.0603f, -0.0478f, -0.0524f, -0.0227f, -0.0274f, -0.0154f, -0.0255f,
-                                                               -0.0572f, -0.0565f, -0.0518f, -0.0496f, 0.0116f, 0.0054f, 0.0163f, 0.0104f});
-            mean_tensor.reshape_(stats_shape);
-            auto std_tensor = sd::Tensor<float>::from_vector({1.8029f, 1.7786f, 1.7868f, 1.7837f, 1.7717f, 1.7590f, 1.7610f, 1.7479f,
-                                                              1.7336f, 1.7373f, 1.7340f, 1.7343f, 1.8626f, 1.8527f, 1.8629f, 1.8589f,
-                                                              1.7593f, 1.7526f, 1.7556f, 1.7583f, 1.7363f, 1.7400f, 1.7355f, 1.7394f,
-                                                              1.7342f, 1.7246f, 1.7392f, 1.7304f, 1.7551f, 1.7513f, 1.7559f, 1.7488f,
-                                                              1.8449f, 1.8454f, 1.8550f, 1.8535f, 1.8240f, 1.7813f, 1.7854f, 1.7945f,
-                                                              1.8047f, 1.7876f, 1.7695f, 1.7676f, 1.7782f, 1.7667f, 1.7925f, 1.7848f,
-                                                              1.7579f, 1.7407f, 1.7483f, 1.7368f, 1.7961f, 1.7998f, 1.7920f, 1.7925f,
-                                                              1.7780f, 1.7747f, 1.7727f, 1.7749f, 1.7526f, 1.7447f, 1.7657f, 1.7495f,
-                                                              1.7775f, 1.7720f, 1.7813f, 1.7813f, 1.8162f, 1.8013f, 1.8023f, 1.8033f,
-                                                              1.7527f, 1.7331f, 1.7563f, 1.7482f, 1.7610f, 1.7507f, 1.7681f, 1.7613f,
-                                                              1.7665f, 1.7545f, 1.7828f, 1.7726f, 1.7896f, 1.7999f, 1.7864f, 1.7760f,
-                                                              1.7613f, 1.7625f, 1.7560f, 1.7577f, 1.7783f, 1.7671f, 1.7810f, 1.7799f,
-                                                              1.7201f, 1.7068f, 1.7265f, 1.7091f, 1.7793f, 1.7578f, 1.7502f, 1.7455f,
-                                                              1.7587f, 1.7500f, 1.7525f, 1.7362f, 1.7616f, 1.7572f, 1.7444f, 1.7430f,
-                                                              1.7509f, 1.7610f, 1.7634f, 1.7612f, 1.7254f, 1.7135f, 1.7321f, 1.7226f,
-                                                              1.7664f, 1.7624f, 1.7718f, 1.7664f, 1.7457f, 1.7441f, 1.7569f, 1.7530f});
-            std_tensor.reshape_(stats_shape);
-            return {std::move(mean_tensor), std::move(std_tensor)};
-        } else {
-            GGML_ABORT("unknown version %d", version);
-        }
-    }
-
-    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
-        if (sd_version_is_flux2(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;
-        }
-        return (latents / scale_factor) + shift_factor;
-    }
-
-    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
-        if (sd_version_is_flux2(version)) {
-            int channel_dim                = 2;
-            auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
-            return ((latents - mean_tensor) * scale_factor) / std_tensor;
-        }
-        return (latents - shift_factor) * scale_factor;
-    }
-
-    int get_encoder_output_channels(int input_channels) {
-        return ae.get_encoder_output_channels();
+        // ggml_set_f32(z, 0.5f);
+        // print_ggml_tensor(z);
+        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
     }
 
     void test() {
-        ggml_init_params params;
+        struct ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        ggml_context* ctx = ggml_init(params);
-        GGML_ASSERT(ctx != nullptr);
+        struct ggml_context* work_ctx = ggml_init(params);
+        GGML_ASSERT(work_ctx != nullptr);
 
         {
             // CPU, x{1, 3, 64, 64}: Pass
             // CUDA, x{1, 3, 64, 64}: Pass, but sill get wrong result for some image, may be due to interlnal nan
             // CPU, x{2, 3, 64, 64}: Wrong result
             // CUDA, x{2, 3, 64, 64}: Wrong result, and different from CPU result
-            sd::Tensor<float> x({64, 64, 3, 2});
-            x.fill_(0.5f);
-            print_sd_tensor(x);
-            sd::Tensor<float> out;
+            auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 64, 64, 3, 2);
+            ggml_set_f32(x, 0.5f);
+            print_ggml_tensor(x);
+            struct ggml_tensor* out = nullptr;
 
-            int64_t t0   = ggml_time_ms();
-            auto out_opt = _compute(8, x, false);
-            int64_t t1   = ggml_time_ms();
+            int t0 = ggml_time_ms();
+            compute(8, x, false, &out, work_ctx);
+            int t1 = ggml_time_ms();
 
-            GGML_ASSERT(!out_opt.empty());
-            out = std::move(out_opt);
-            print_sd_tensor(out);
-            LOG_DEBUG("encode test done in %lldms", t1 - t0);
+            print_ggml_tensor(out);
+            LOG_DEBUG("encode test done in %dms", t1 - t0);
         }
 
         if (false) {
@@ -832,21 +758,19 @@ struct AutoEncoderKL : public VAE {
             // CUDA, z{1, 4, 8, 8}: Pass
             // CPU, z{3, 4, 8, 8}: Wrong result
             // CUDA, z{3, 4, 8, 8}: Wrong result, and different from CPU result
-            sd::Tensor<float> z({8, 8, 4, 1});
-            z.fill_(0.5f);
-            print_sd_tensor(z);
-            sd::Tensor<float> out;
+            auto z = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 4, 1);
+            ggml_set_f32(z, 0.5f);
+            print_ggml_tensor(z);
+            struct ggml_tensor* out = nullptr;
 
-            int64_t t0   = ggml_time_ms();
-            auto out_opt = _compute(8, z, true);
-            int64_t t1   = ggml_time_ms();
+            int t0 = ggml_time_ms();
+            compute(8, z, true, &out, work_ctx);
+            int t1 = ggml_time_ms();
 
-            GGML_ASSERT(!out_opt.empty());
-            out = std::move(out_opt);
-            print_sd_tensor(out);
-            LOG_DEBUG("decode test done in %lldms", t1 - t0);
+            print_ggml_tensor(out);
+            LOG_DEBUG("decode test done in %dms", t1 - t0);
         }
     };
 };
 
-#endif  // __AUTO_ENCODER_KL_HPP__
+#endif