refactor: simplify sample cache flow (#1350)

* feat: turn flow_shift into a generation parameter

* format code

* simplify set_shift/set_parameters

* fix sd_sample_params_to_str

* remove unused variable

* update docs

---------

Co-authored-by: leejet <leejet714@gmail.com>

.github/workflows/build.yml

@@ -21,11 +21,13 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "examples/server/frontend/**",
       ]
   pull_request:
     types: [opened, synchronize, reopened]
     paths:
       [
+        ".github/workflows/**",
         "**/CMakeLists.txt",
         "**/Makefile",
         "**/*.h",
@@ -33,6 +35,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "examples/server/frontend/**",
       ]
 
 env:
@@ -53,6 +56,16 @@ 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: |
@@ -70,7 +83,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: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Fetch system info
         id: system-info
@@ -106,6 +119,16 @@ 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: |
@@ -123,7 +146,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: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Fetch system info
         id: system-info
@@ -162,7 +185,7 @@ jobs:
 
     strategy:
       matrix:
-        variant: [musa, sycl, vulkan]
+        variant: [musa, sycl, vulkan, cuda]
 
     env:
       REGISTRY: ghcr.io
@@ -174,10 +197,20 @@ 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: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -223,6 +256,16 @@ 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: |
@@ -240,7 +283,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: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Fetch system info
         id: system-info
@@ -294,6 +337,16 @@ 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' }}
@@ -340,7 +393,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: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Pack artifacts
         id: pack_artifacts
@@ -399,6 +452,16 @@ 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
@@ -463,7 +526,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: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Pack artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -485,6 +548,156 @@ 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' }}
 
@@ -493,6 +706,7 @@ jobs:
     needs:
       - ubuntu-latest-cmake
       - ubuntu-latest-cmake-vulkan
+      - ubuntu-latest-rocm
       - build-and-push-docker-images
       - macOS-latest-cmake
       - windows-latest-cmake
@@ -519,7 +733,7 @@ jobs:
 
       - name: Get commit hash
         id: commit
-        uses: pr-mpt/actions-commit-hash@v2
+        uses: prompt/actions-commit-hash@v2
 
       - name: Create release
         id: create_release

.gitmodules

@@ -1,3 +1,6 @@
 [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

@@ -36,7 +36,6 @@ 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)
@@ -70,26 +69,22 @@ 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
-    "*.h"
-    "*.cpp"
-    "*.hpp"
+    "src/*.h"
+    "src/*.cpp"
+    "src/*.hpp"
+    "src/vocab/*.h"
+    "src/vocab/*.cpp"
 )
 
 find_program(GIT_EXE NAMES git git.exe NO_CMAKE_FIND_ROOT_PATH)
@@ -119,7 +114,7 @@ endif()
 message(STATUS "stable-diffusion.cpp commit ${SDCPP_BUILD_COMMIT}")
 
 set_property(
-  SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp
+  SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/src/version.cpp
   APPEND PROPERTY COMPILE_DEFINITIONS
   SDCPP_BUILD_COMMIT=${SDCPP_BUILD_COMMIT} SDCPP_BUILD_VERSION=${SDCPP_BUILD_VERSION}
 )
@@ -182,6 +177,7 @@ endif()
 add_subdirectory(thirdparty)
 
 target_link_libraries(${SD_LIB} PUBLIC ggml zip)
+target_include_directories(${SD_LIB} PUBLIC . include)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
 target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17)
 
@@ -190,7 +186,7 @@ if (SD_BUILD_EXAMPLES)
     add_subdirectory(examples)
 endif()
 
-set(SD_PUBLIC_HEADERS stable-diffusion.h)
+set(SD_PUBLIC_HEADERS include/stable-diffusion.h)
 set_target_properties(${SD_LIB} PROPERTIES PUBLIC_HEADER "${SD_PUBLIC_HEADERS}")
 
 install(TARGETS ${SD_LIB} LIBRARY PUBLIC_HEADER)

Dockerfile.cuda

@@ -0,0 +1,25 @@
+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" ]

README.md

@@ -53,6 +53,7 @@ API and command-line option may change frequently.***
     - [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)
@@ -139,6 +140,7 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 - [🔥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)

docs/anima.md

@@ -0,0 +1,21 @@
+# 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

@@ -11,6 +11,7 @@ Caching methods accelerate diffusion inference by reusing intermediate computati
 | `dbcache` | DiT models | Block-level L1 residual threshold |
 | `taylorseer` | DiT models | Taylor series approximation |
 | `cache-dit` | DiT models | Combined DBCache + TaylorSeer |
+| `spectrum` | UNET models | Chebyshev + Taylor output forecasting |
 
 ### UCache (UNET Models)
 
@@ -79,7 +80,7 @@ Uses Taylor series approximation to predict block outputs:
 Combines DBCache and TaylorSeer:
 
 ```bash
---cache-mode cache-dit --cache-preset fast
+--cache-mode cache-dit
 ```
 
 #### Parameters
@@ -91,14 +92,6 @@ Combines DBCache and TaylorSeer:
 | `threshold` | L1 residual difference threshold | 0.08 |
 | `warmup` | Steps before caching starts | 8 |
 
-#### Presets
-
-Available presets: `slow`, `medium`, `fast`, `ultra` (or `s`, `m`, `f`, `u`).
-
-```bash
---cache-mode cache-dit --cache-preset fast
-```
-
 #### SCM Options
 
 Steps Computation Mask controls which steps can be cached:
@@ -118,6 +111,28 @@ Mask values: `1` = compute, `0` = can cache.
 --scm-policy dynamic
 ```
 
+### Spectrum (UNET Models)
+
+Spectrum uses Chebyshev polynomial fitting blended with Taylor extrapolation to predict denoised outputs, skipping entire UNet 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

examples/cli/README.md

@@ -4,11 +4,12 @@
 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)
-  --output-begin-idx <int>    starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)
+  -o, --output <string>       path to write result image to. you can use printf-style %d format specifiers for image sequences (default:
+                              ./output.png) (eg. output_%03d.png)
   --preview-path <string>     path to write preview image to (default: ./preview.png)
   --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
@@ -44,7 +45,6 @@ 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
@@ -60,6 +60,7 @@ Context Options:
   --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
@@ -108,14 +109,15 @@ Generation Options:
                                            medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            eta in DDIM, only for DDIM/TCD/res_multistep/res_2s (default: 0)
+  --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)
   --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM/TCD/res_multistep/res_2s (default: 0)
+  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM and TCD (default: 0)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>
   --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image
@@ -124,21 +126,24 @@ Generation Options:
   --increase-ref-index                     automatically increase the indices of references images based on the order they are listed (starting with 1).
   --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)
-  --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
+  --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)
+  --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").
   --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)
+  --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"
-  --cache-preset                           cache-dit preset: 'slow'/'s', 'medium'/'m', 'fast'/'f', 'ultra'/'u'
+                                           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'
 ```

examples/cli/main.cpp

@@ -394,12 +394,15 @@ bool save_results(const SDCliParams& cli_params,
 
     fs::path base_path = out_path;
     fs::path ext       = out_path.has_extension() ? out_path.extension() : fs::path{};
-    if (!ext.empty())
-        base_path.replace_extension();
 
     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) {
@@ -409,7 +412,7 @@ bool save_results(const SDCliParams& cli_params,
     auto write_image = [&](const fs::path& path, int idx) {
         const sd_image_t& img = results[idx];
         if (!img.data)
-            return;
+            return false;
 
         std::string params = get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + idx);
         int ok             = 0;
@@ -419,8 +422,11 @@ bool save_results(const SDCliParams& cli_params,
             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";
@@ -429,9 +435,12 @@ bool save_results(const SDCliParams& cli_params,
 
         for (int i = 0; i < num_results; ++i) {
             fs::path img_path = format_frame_idx(pattern.string(), output_begin_idx + i);
-            write_image(img_path, i);
+            if (write_image(img_path, i)) {
+                sucessful_reults++;
+            }
         }
-        return true;
+        LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
+        return sucessful_reults != 0;
     }
 
     if (cli_params.mode == VID_GEN && num_results > 1) {
@@ -439,9 +448,13 @@ bool save_results(const SDCliParams& cli_params,
             ext = ".avi";
         fs::path video_path = base_path;
         video_path += ext;
-        create_mjpg_avi_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps);
-        LOG_INFO("save result MJPG AVI video to '%s'", video_path.string().c_str());
-        return true;
+        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")
@@ -453,10 +466,12 @@ bool save_results(const SDCliParams& cli_params,
             img_path += "_" + std::to_string(output_begin_idx + i);
         }
         img_path += ext;
-        write_image(img_path, i);
+        if (write_image(img_path, i)) {
+            sucessful_reults++;
+        }
     }
-
-    return true;
+    LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
+    return sucessful_reults != 0;
 }
 
 int main(int argc, const char* argv[]) {

examples/common/common.hpp

@@ -581,10 +581,6 @@ struct SDContextParams {
              "--vae-tile-overlap",
              "tile overlap for vae tiling, in fraction of tile size (default: 0.5)",
              &vae_tiling_params.target_overlap},
-            {"",
-             "--flow-shift",
-             "shift value for Flow models like SD3.x or WAN (default: auto)",
-             &flow_shift},
         };
 
         options.bool_options = {
@@ -903,7 +899,6 @@ struct SDContextParams {
             << "  photo_maker_path: \"" << photo_maker_path << "\",\n"
             << "  rng_type: " << sd_rng_type_name(rng_type) << ",\n"
             << "  sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
-            << "  flow_shift: " << (std::isinf(flow_shift) ? "INF" : std::to_string(flow_shift)) << "\n"
             << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
             << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
             << "  control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
@@ -986,7 +981,6 @@ struct SDContextParams {
             chroma_use_t5_mask,
             chroma_t5_mask_pad,
             qwen_image_zero_cond_t,
-            flow_shift,
         };
         return sd_ctx_params;
     }
@@ -1053,7 +1047,6 @@ struct SDGenerationParams {
 
     std::string cache_mode;
     std::string cache_option;
-    std::string cache_preset;
     std::string scm_mask;
     bool scm_policy_dynamic = true;
     sd_cache_params_t cache_params{};
@@ -1206,6 +1199,10 @@ struct SDGenerationParams {
              "--eta",
              "eta in DDIM, only for DDIM and TCD (default: 0)",
              &sample_params.eta},
+            {"",
+             "--flow-shift",
+             "shift value for Flow models like SD3.x or WAN (default: auto)",
+             &sample_params.flow_shift},
             {"",
              "--high-noise-cfg-scale",
              "(high noise) unconditional guidance scale: (default: 7.0)",
@@ -1424,8 +1421,8 @@ struct SDGenerationParams {
             }
             cache_mode = argv_to_utf8(index, argv);
             if (cache_mode != "easycache" && cache_mode != "ucache" &&
-                cache_mode != "dbcache" && cache_mode != "taylorseer" && cache_mode != "cache-dit") {
-                fprintf(stderr, "error: invalid cache mode '%s', must be 'easycache', 'ucache', 'dbcache', 'taylorseer', or 'cache-dit'\n", cache_mode.c_str());
+                cache_mode != "dbcache" && cache_mode != "taylorseer" && cache_mode != "cache-dit" && cache_mode != "spectrum") {
+                fprintf(stderr, "error: invalid cache mode '%s', must be 'easycache', 'ucache', 'dbcache', 'taylorseer', 'cache-dit', or 'spectrum'\n", cache_mode.c_str());
                 return -1;
             }
             return 1;
@@ -1463,21 +1460,6 @@ struct SDGenerationParams {
             return 1;
         };
 
-        auto on_cache_preset_arg = [&](int argc, const char** argv, int index) {
-            if (++index >= argc) {
-                return -1;
-            }
-            cache_preset = argv_to_utf8(index, argv);
-            if (cache_preset != "slow" && cache_preset != "s" && cache_preset != "S" &&
-                cache_preset != "medium" && cache_preset != "m" && cache_preset != "M" &&
-                cache_preset != "fast" && cache_preset != "f" && cache_preset != "F" &&
-                cache_preset != "ultra" && cache_preset != "u" && cache_preset != "U") {
-                fprintf(stderr, "error: invalid cache preset '%s', must be 'slow'/'s', 'medium'/'m', 'fast'/'f', or 'ultra'/'u'\n", cache_preset.c_str());
-                return -1;
-            }
-            return 1;
-        };
-
         options.manual_options = {
             {"-s",
              "--seed",
@@ -1515,16 +1497,12 @@ struct SDGenerationParams {
              on_ref_image_arg},
             {"",
              "--cache-mode",
-             "caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level)",
+             "caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",
              on_cache_mode_arg},
             {"",
              "--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\"",
+             "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\"",
              on_cache_option_arg},
-            {"",
-             "--cache-preset",
-             "cache-dit preset: 'slow'/'s', 'medium'/'m', 'fast'/'f', 'ultra'/'u'",
-             on_cache_preset_arg},
             {"",
              "--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",
@@ -1577,7 +1555,6 @@ struct SDGenerationParams {
         load_if_exists("negative_prompt", negative_prompt);
         load_if_exists("cache_mode", cache_mode);
         load_if_exists("cache_option", cache_option);
-        load_if_exists("cache_preset", cache_preset);
         load_if_exists("scm_mask", scm_mask);
 
         load_if_exists("clip_skip", clip_skip);
@@ -1606,6 +1583,7 @@ struct SDGenerationParams {
         load_if_exists("cfg_scale", sample_params.guidance.txt_cfg);
         load_if_exists("img_cfg_scale", sample_params.guidance.img_cfg);
         load_if_exists("guidance", sample_params.guidance.distilled_guidance);
+        load_if_exists("flow_shift", sample_params.flow_shift);
 
         auto load_sampler_if_exists = [&](const char* key, enum sample_method_t& out) {
             if (j.contains(key) && j[key].is_string()) {
@@ -1780,7 +1758,23 @@ struct SDGenerationParams {
                     } else if (key == "Bn" || key == "bn") {
                         cache_params.Bn_compute_blocks = std::stoi(val);
                     } else if (key == "warmup") {
-                        cache_params.max_warmup_steps = std::stoi(val);
+                        if (cache_mode == "spectrum") {
+                            cache_params.spectrum_warmup_steps = std::stoi(val);
+                        } else {
+                            cache_params.max_warmup_steps = std::stoi(val);
+                        }
+                    } else if (key == "w") {
+                        cache_params.spectrum_w = std::stof(val);
+                    } else if (key == "m") {
+                        cache_params.spectrum_m = std::stoi(val);
+                    } else if (key == "lam") {
+                        cache_params.spectrum_lam = std::stof(val);
+                    } else if (key == "window") {
+                        cache_params.spectrum_window_size = std::stoi(val);
+                    } else if (key == "flex") {
+                        cache_params.spectrum_flex_window = std::stof(val);
+                    } else if (key == "stop") {
+                        cache_params.spectrum_stop_percent = std::stof(val);
                     } else {
                         LOG_ERROR("error: unknown cache parameter '%s'", key.c_str());
                         return false;
@@ -1795,39 +1789,17 @@ struct SDGenerationParams {
 
         if (!cache_mode.empty()) {
             if (cache_mode == "easycache") {
-                cache_params.mode                   = SD_CACHE_EASYCACHE;
-                cache_params.reuse_threshold        = 0.2f;
-                cache_params.start_percent          = 0.15f;
-                cache_params.end_percent            = 0.95f;
-                cache_params.error_decay_rate       = 1.0f;
-                cache_params.use_relative_threshold = true;
-                cache_params.reset_error_on_compute = true;
+                cache_params.mode = SD_CACHE_EASYCACHE;
             } else if (cache_mode == "ucache") {
-                cache_params.mode                   = SD_CACHE_UCACHE;
-                cache_params.reuse_threshold        = 1.0f;
-                cache_params.start_percent          = 0.15f;
-                cache_params.end_percent            = 0.95f;
-                cache_params.error_decay_rate       = 1.0f;
-                cache_params.use_relative_threshold = true;
-                cache_params.reset_error_on_compute = true;
+                cache_params.mode = SD_CACHE_UCACHE;
             } else if (cache_mode == "dbcache") {
-                cache_params.mode                    = SD_CACHE_DBCACHE;
-                cache_params.Fn_compute_blocks       = 8;
-                cache_params.Bn_compute_blocks       = 0;
-                cache_params.residual_diff_threshold = 0.08f;
-                cache_params.max_warmup_steps        = 8;
+                cache_params.mode = SD_CACHE_DBCACHE;
             } else if (cache_mode == "taylorseer") {
-                cache_params.mode                    = SD_CACHE_TAYLORSEER;
-                cache_params.Fn_compute_blocks       = 8;
-                cache_params.Bn_compute_blocks       = 0;
-                cache_params.residual_diff_threshold = 0.08f;
-                cache_params.max_warmup_steps        = 8;
+                cache_params.mode = SD_CACHE_TAYLORSEER;
             } else if (cache_mode == "cache-dit") {
-                cache_params.mode                    = SD_CACHE_CACHE_DIT;
-                cache_params.Fn_compute_blocks       = 8;
-                cache_params.Bn_compute_blocks       = 0;
-                cache_params.residual_diff_threshold = 0.08f;
-                cache_params.max_warmup_steps        = 8;
+                cache_params.mode = SD_CACHE_CACHE_DIT;
+            } else if (cache_mode == "spectrum") {
+                cache_params.mode = SD_CACHE_SPECTRUM;
             }
 
             if (!cache_option.empty()) {

examples/server/CMakeLists.txt

@@ -1,6 +1,73 @@
 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,15 +1,104 @@
+# 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)
-  --listen-port <int>         server listen port (default: 1234)
-  --serve-html-path <string>  path to HTML file to serve at root (optional)
-  -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)        
+  --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
 
 Context Options:
   -m, --model <string>                     path to full model
@@ -36,14 +125,13 @@ 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)
-  --mmap                                   whether to memory-map model
   --fa                                     use flash attention
   --diffusion-fa                           use flash attention in the diffusion model only
   --diffusion-conv-direct                  use ggml_conv2d_direct in the diffusion model
@@ -52,6 +140,7 @@ Context Options:
   --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
@@ -100,14 +189,15 @@ Default Generation Options:
                                            medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            eta in DDIM, only for DDIM/TCD/res_multistep/res_2s (default: 0)
+  --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)
   --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM/TCD/res_multistep/res_2s (default: 0)
+  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM and TCD (default: 0)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>
   --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image
@@ -116,21 +206,22 @@ Default Generation Options:
   --increase-ref-index                     automatically increase the indices of references images based on the order they are listed (starting with 1).
   --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)
-  --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
+  --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)
+  --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").
   --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)
+  --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"
-  --cache-preset                           cache-dit preset: 'slow'/'s', 'medium'/'m', 'fast'/'f', 'ultra'/'u'
   --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'
 ```

examples/server/frontend

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

examples/server/main.cpp

@@ -13,6 +13,10 @@
 
 #include "common/common.hpp"
 
+#ifdef HAVE_INDEX_HTML
+#include "frontend/dist/gen_index_html.h"
+#endif
+
 namespace fs = std::filesystem;
 
 // ----------------------- helpers -----------------------
@@ -266,8 +270,21 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
 struct LoraEntry {
     std::string name;
     std::string path;
+    std::string fullpath;
 };
 
+void free_results(sd_image_t* result_images, int num_results) {
+    if (result_images) {
+        for (int i = 0; i < num_results; ++i) {
+            if (result_images[i].data) {
+                stbi_image_free(result_images[i].data);
+                result_images[i].data = nullptr;
+            }
+        }
+    }
+    free(result_images);
+}
+
 int main(int argc, const char** argv) {
     if (argc > 1 && std::string(argv[1]) == "--version") {
         std::cout << version_string() << "\n";
@@ -321,7 +338,8 @@ int main(int argc, const char** argv) {
 
                 LoraEntry e;
                 e.name          = p.stem().u8string();
-                std::string rel = fs::relative(p, lora_dir).u8string();
+                e.fullpath      = p.u8string();
+                std::string rel = p.lexically_relative(lora_dir).u8string();
                 std::replace(rel.begin(), rel.end(), '\\', '/');
                 e.path = rel;
 
@@ -340,10 +358,11 @@ int main(int argc, const char** argv) {
         }
     };
 
-    auto is_valid_lora_path = [&](const std::string& path) -> bool {
+    auto get_lora_full_path = [&](const std::string& path) -> std::string {
         std::lock_guard<std::mutex> lock(lora_mutex);
-        return std::any_of(lora_cache.begin(), lora_cache.end(),
-                           [&](const LoraEntry& e) { return e.path == path; });
+        auto it = std::find_if(lora_cache.begin(), lora_cache.end(),
+                               [&](const LoraEntry& e) { return e.path == path; });
+        return (it != lora_cache.end()) ? it->fullpath : "";
     };
 
     httplib::Server svr;
@@ -365,7 +384,13 @@ int main(int argc, const char** argv) {
         return httplib::Server::HandlerResponse::Unhandled;
     });
 
-    // root
+    // index html
+    std::string index_html;
+#ifdef HAVE_INDEX_HTML
+    index_html.assign(reinterpret_cast<const char*>(index_html_bytes), index_html_size);
+#else
+    index_html = "Stable Diffusion Server is running";
+#endif
     svr.Get("/", [&](const httplib::Request&, httplib::Response& res) {
         if (!svr_params.serve_html_path.empty()) {
             std::ifstream file(svr_params.serve_html_path);
@@ -377,7 +402,7 @@ int main(int argc, const char** argv) {
                 res.set_content("Error: Unable to read HTML file", "text/plain");
             }
         } else {
-            res.set_content("Stable Diffusion Server is running", "text/plain");
+            res.set_content(index_html, "text/html");
         }
     });
 
@@ -404,8 +429,8 @@ int main(int argc, const char** argv) {
             std::string size          = j.value("size", "");
             std::string output_format = j.value("output_format", "png");
             int output_compression    = j.value("output_compression", 100);
-            int width                 = 512;
-            int height                = 512;
+            int width                 = default_gen_params.width > 0 ? default_gen_params.width : 512;
+            int height                = default_gen_params.width > 0 ? default_gen_params.height : 512;
             if (!size.empty()) {
                 auto pos = size.find('x');
                 if (pos != std::string::npos) {
@@ -534,6 +559,7 @@ int main(int argc, const char** argv) {
                 item["b64_json"] = b64;
                 out["data"].push_back(item);
             }
+            free_results(results, num_results);
 
             res.set_content(out.dump(), "application/json");
             res.status = 200;
@@ -564,8 +590,9 @@ int main(int argc, const char** argv) {
 
             std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(prompt);
 
-            size_t image_count = req.form.get_file_count("image[]");
-            if (image_count == 0) {
+            size_t image_count    = req.form.get_file_count("image[]");
+            bool has_legacy_image = req.form.has_file("image");
+            if (image_count == 0 && !has_legacy_image) {
                 res.status = 400;
                 res.set_content(R"({"error":"at least one image[] required"})", "application/json");
                 return;
@@ -576,6 +603,10 @@ int main(int argc, const char** argv) {
                 auto file = req.form.get_file("image[]", i);
                 images_bytes.emplace_back(file.content.begin(), file.content.end());
             }
+            if (image_count == 0 && has_legacy_image) {
+                auto file = req.form.get_file("image");
+                images_bytes.emplace_back(file.content.begin(), file.content.end());
+            }
 
             std::vector<uint8_t> mask_bytes;
             if (req.form.has_file("mask")) {
@@ -593,7 +624,7 @@ int main(int argc, const char** argv) {
             n = std::clamp(n, 1, 8);
 
             std::string size = req.form.get_field("size");
-            int width = 512, height = 512;
+            int width = -1, height = -1;
             if (!size.empty()) {
                 auto pos = size.find('x');
                 if (pos != std::string::npos) {
@@ -650,15 +681,31 @@ int main(int argc, const char** argv) {
 
             LOG_DEBUG("%s\n", gen_params.to_string().c_str());
 
-            sd_image_t init_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 init_image    = {0, 0, 3, nullptr};
+            sd_image_t control_image = {0, 0, 3, nullptr};
             std::vector<sd_image_t> pmid_images;
 
+            auto get_resolved_width = [&gen_params, &default_gen_params]() -> int {
+                if (gen_params.width > 0)
+                    return gen_params.width;
+                if (default_gen_params.width > 0)
+                    return default_gen_params.width;
+                return 512;
+            };
+            auto get_resolved_height = [&gen_params, &default_gen_params]() -> int {
+                if (gen_params.height > 0)
+                    return gen_params.height;
+                if (default_gen_params.height > 0)
+                    return default_gen_params.height;
+                return 512;
+            };
+
             std::vector<sd_image_t> ref_images;
             ref_images.reserve(images_bytes.size());
             for (auto& bytes : images_bytes) {
-                int img_w           = width;
-                int img_h           = height;
+                int img_w;
+                int img_h;
+
                 uint8_t* raw_pixels = load_image_from_memory(
                     reinterpret_cast<const char*>(bytes.data()),
                     static_cast<int>(bytes.size()),
@@ -670,22 +717,31 @@ int main(int argc, const char** argv) {
                 }
 
                 sd_image_t img{(uint32_t)img_w, (uint32_t)img_h, 3, raw_pixels};
+                gen_params.set_width_and_height_if_unset(img.width, img.height);
                 ref_images.push_back(img);
             }
 
             sd_image_t mask_image = {0};
             if (!mask_bytes.empty()) {
-                int mask_w        = width;
-                int mask_h        = height;
+                int expected_width  = 0;
+                int expected_height = 0;
+                if (gen_params.width_and_height_are_set()) {
+                    expected_width  = gen_params.width;
+                    expected_height = gen_params.height;
+                }
+                int mask_w;
+                int mask_h;
+
                 uint8_t* mask_raw = load_image_from_memory(
                     reinterpret_cast<const char*>(mask_bytes.data()),
                     static_cast<int>(mask_bytes.size()),
                     mask_w, mask_h,
-                    width, height, 1);
+                    expected_width, expected_height, 1);
                 mask_image = {(uint32_t)mask_w, (uint32_t)mask_h, 1, mask_raw};
+                gen_params.set_width_and_height_if_unset(mask_image.width, mask_image.height);
             } else {
-                mask_image.width   = width;
-                mask_image.height  = height;
+                mask_image.width   = get_resolved_width();
+                mask_image.height  = get_resolved_height();
                 mask_image.channel = 1;
                 mask_image.data    = nullptr;
             }
@@ -702,8 +758,8 @@ int main(int argc, const char** argv) {
                 gen_params.auto_resize_ref_image,
                 gen_params.increase_ref_index,
                 mask_image,
-                gen_params.width,
-                gen_params.height,
+                get_resolved_width(),
+                get_resolved_height(),
                 gen_params.sample_params,
                 gen_params.strength,
                 gen_params.seed,
@@ -748,6 +804,7 @@ int main(int argc, const char** argv) {
                 item["b64_json"] = b64;
                 out["data"].push_back(item);
             }
+            free_results(results, num_results);
 
             res.set_content(out.dump(), "application/json");
             res.status = 200;
@@ -786,8 +843,8 @@ int main(int argc, const char** argv) {
             std::string negative_prompt = j.value("negative_prompt", "");
             int width                   = j.value("width", 512);
             int height                  = j.value("height", 512);
-            int steps                   = j.value("steps", -1);
-            float cfg_scale             = j.value("cfg_scale", 7.f);
+            int steps                   = j.value("steps", default_gen_params.sample_params.sample_steps);
+            float cfg_scale             = j.value("cfg_scale", default_gen_params.sample_params.guidance.txt_cfg);
             int64_t seed                = j.value("seed", -1);
             int batch_size              = j.value("batch_size", 1);
             int clip_skip               = j.value("clip_skip", -1);
@@ -837,11 +894,12 @@ int main(int argc, const char** argv) {
                         return bad("lora.path required");
                     }
 
-                    if (!is_valid_lora_path(path)) {
+                    std::string fullpath = get_lora_full_path(path);
+                    if (fullpath.empty()) {
                         return bad("invalid lora path: " + path);
                     }
 
-                    lora_path_storage.push_back(path);
+                    lora_path_storage.push_back(fullpath);
                     sd_lora_t l;
                     l.is_high_noise = is_high_noise;
                     l.multiplier    = multiplier;
@@ -883,16 +941,13 @@ int main(int argc, const char** argv) {
 
             enum scheduler_t scheduler = str_to_scheduler(scheduler_name.c_str());
 
-            // avoid excessive resource usage
-
-            SDGenerationParams gen_params         = default_gen_params;
-            gen_params.prompt                     = prompt;
-            gen_params.negative_prompt            = negative_prompt;
-            gen_params.width                      = width;
-            gen_params.height                     = height;
-            gen_params.seed                       = seed;
-            gen_params.sample_params.sample_steps = steps;
-            gen_params.batch_count                = batch_size;
+            SDGenerationParams gen_params             = default_gen_params;
+            gen_params.prompt                         = prompt;
+            gen_params.negative_prompt                = negative_prompt;
+            gen_params.seed                           = seed;
+            gen_params.sample_params.sample_steps     = steps;
+            gen_params.batch_count                    = batch_size;
+            gen_params.sample_params.guidance.txt_cfg = cfg_scale;
 
             if (clip_skip > 0) {
                 gen_params.clip_skip = clip_skip;
@@ -906,38 +961,66 @@ int main(int argc, const char** argv) {
                 gen_params.sample_params.scheduler = scheduler;
             }
 
+            // re-read to avoid applying 512 as default before the provided
+            // images and/or server command-line
+            gen_params.width  = j.value("width", -1);
+            gen_params.height = j.value("height", -1);
+
             LOG_DEBUG("%s\n", gen_params.to_string().c_str());
 
-            sd_image_t init_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};
+            sd_image_t init_image    = {0, 0, 3, nullptr};
+            sd_image_t control_image = {0, 0, 3, nullptr};
+            sd_image_t mask_image    = {0, 0, 1, nullptr};
             std::vector<uint8_t> mask_data;
             std::vector<sd_image_t> pmid_images;
             std::vector<sd_image_t> ref_images;
 
-            if (img2img) {
-                auto decode_image = [](sd_image_t& image, std::string encoded) -> bool {
-                    // remove data URI prefix if present ("data:image/png;base64,")
-                    auto comma_pos = encoded.find(',');
-                    if (comma_pos != std::string::npos) {
-                        encoded = encoded.substr(comma_pos + 1);
-                    }
-                    std::vector<uint8_t> img_data = base64_decode(encoded);
-                    if (!img_data.empty()) {
-                        int img_w         = image.width;
-                        int img_h         = image.height;
-                        uint8_t* raw_data = load_image_from_memory(
-                            (const char*)img_data.data(), (int)img_data.size(),
-                            img_w, img_h,
-                            image.width, image.height, image.channel);
-                        if (raw_data) {
-                            image = {(uint32_t)img_w, (uint32_t)img_h, image.channel, raw_data};
-                            return true;
-                        }
-                    }
-                    return false;
-                };
+            auto get_resolved_width = [&gen_params, &default_gen_params]() -> int {
+                if (gen_params.width > 0)
+                    return gen_params.width;
+                if (default_gen_params.width > 0)
+                    return default_gen_params.width;
+                return 512;
+            };
+            auto get_resolved_height = [&gen_params, &default_gen_params]() -> int {
+                if (gen_params.height > 0)
+                    return gen_params.height;
+                if (default_gen_params.height > 0)
+                    return default_gen_params.height;
+                return 512;
+            };
 
+            auto decode_image = [&gen_params](sd_image_t& image, std::string encoded) -> bool {
+                // remove data URI prefix if present ("data:image/png;base64,")
+                auto comma_pos = encoded.find(',');
+                if (comma_pos != std::string::npos) {
+                    encoded = encoded.substr(comma_pos + 1);
+                }
+                std::vector<uint8_t> img_data = base64_decode(encoded);
+                if (!img_data.empty()) {
+                    int expected_width  = 0;
+                    int expected_height = 0;
+                    if (gen_params.width_and_height_are_set()) {
+                        expected_width  = gen_params.width;
+                        expected_height = gen_params.height;
+                    }
+                    int img_w;
+                    int img_h;
+
+                    uint8_t* raw_data = load_image_from_memory(
+                        (const char*)img_data.data(), (int)img_data.size(),
+                        img_w, img_h,
+                        expected_width, expected_height, image.channel);
+                    if (raw_data) {
+                        image = {(uint32_t)img_w, (uint32_t)img_h, image.channel, raw_data};
+                        gen_params.set_width_and_height_if_unset(image.width, image.height);
+                        return true;
+                    }
+                }
+                return false;
+            };
+
+            if (img2img) {
                 if (j.contains("init_images") && j["init_images"].is_array() && !j["init_images"].empty()) {
                     std::string encoded = j["init_images"][0].get<std::string>();
                     decode_image(init_image, encoded);
@@ -953,23 +1036,15 @@ int main(int argc, const char** argv) {
                         }
                     }
                 } else {
-                    mask_data          = std::vector<uint8_t>(width * height, 255);
-                    mask_image.width   = width;
-                    mask_image.height  = height;
+                    int m_width        = get_resolved_width();
+                    int m_height       = get_resolved_height();
+                    mask_data          = std::vector<uint8_t>(m_width * m_height, 255);
+                    mask_image.width   = m_width;
+                    mask_image.height  = m_height;
                     mask_image.channel = 1;
                     mask_image.data    = mask_data.data();
                 }
 
-                if (j.contains("extra_images") && j["extra_images"].is_array()) {
-                    for (auto extra_image : j["extra_images"]) {
-                        std::string encoded  = extra_image.get<std::string>();
-                        sd_image_t tmp_image = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
-                        if (decode_image(tmp_image, encoded)) {
-                            ref_images.push_back(tmp_image);
-                        }
-                    }
-                }
-
                 float denoising_strength = j.value("denoising_strength", -1.f);
                 if (denoising_strength >= 0.f) {
                     denoising_strength  = std::min(denoising_strength, 1.0f);
@@ -977,6 +1052,16 @@ int main(int argc, const char** argv) {
                 }
             }
 
+            if (j.contains("extra_images") && j["extra_images"].is_array()) {
+                for (auto extra_image : j["extra_images"]) {
+                    std::string encoded  = extra_image.get<std::string>();
+                    sd_image_t tmp_image = {(uint32_t)gen_params.width, (uint32_t)gen_params.height, 3, nullptr};
+                    if (decode_image(tmp_image, encoded)) {
+                        ref_images.push_back(tmp_image);
+                    }
+                }
+            }
+
             sd_img_gen_params_t img_gen_params = {
                 sd_loras.data(),
                 static_cast<uint32_t>(sd_loras.size()),
@@ -989,8 +1074,8 @@ int main(int argc, const char** argv) {
                 gen_params.auto_resize_ref_image,
                 gen_params.increase_ref_index,
                 mask_image,
-                gen_params.width,
-                gen_params.height,
+                get_resolved_width(),
+                get_resolved_height(),
                 gen_params.sample_params,
                 gen_params.strength,
                 gen_params.seed,
@@ -1040,6 +1125,7 @@ int main(int argc, const char** argv) {
                 std::string b64 = base64_encode(image_bytes);
                 out["images"].push_back(b64);
             }
+            free_results(results, num_results);
 
             res.set_content(out.dump(), "application/json");
             res.status = 200;

format-code.sh

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

include/stable-diffusion.h

@@ -201,7 +201,6 @@ typedef struct {
     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 {
@@ -235,6 +234,7 @@ typedef struct {
     int shifted_timestep;
     float* custom_sigmas;
     int custom_sigmas_count;
+    float flow_shift;
 } sd_sample_params_t;
 
 typedef struct {
@@ -251,6 +251,7 @@ enum sd_cache_mode_t {
     SD_CACHE_DBCACHE,
     SD_CACHE_TAYLORSEER,
     SD_CACHE_CACHE_DIT,
+    SD_CACHE_SPECTRUM,
 };
 
 typedef struct {
@@ -271,6 +272,13 @@ typedef struct {
     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;
 
 typedef struct {

script/face_detect.py

@@ -1,88 +1,88 @@
-import os
-import sys
-
-import numpy as np
-import torch
-from diffusers.utils import load_image
-# pip install insightface==0.7.3
-from insightface.app import FaceAnalysis
-from insightface.data import get_image as ins_get_image
-from safetensors.torch import save_file
-
-### 
-# https://github.com/cubiq/ComfyUI_IPAdapter_plus/issues/165#issue-2055829543
-###
-class FaceAnalysis2(FaceAnalysis):
-    # NOTE: allows setting det_size for each detection call.
-    # the model allows it but the wrapping code from insightface
-    # doesn't show it, and people end up loading duplicate models
-    # for different sizes where there is absolutely no need to
-    def get(self, img, max_num=0, det_size=(640, 640)):
-        if det_size is not None:
-            self.det_model.input_size = det_size
-
-        return super().get(img, max_num)
-
-def analyze_faces(face_analysis: FaceAnalysis, img_data: np.ndarray, det_size=(640, 640)):
-    # NOTE: try detect faces, if no faces detected, lower det_size until it does
-    detection_sizes = [None] + [(size, size) for size in range(640, 256, -64)] + [(256, 256)]
-
-    for size in detection_sizes:
-        faces = face_analysis.get(img_data, det_size=size)
-        if len(faces) > 0:
-            return faces
-
-    return []
-
-if __name__ == "__main__":
-    #face_detector = FaceAnalysis2(providers=['CUDAExecutionProvider'], allowed_modules=['detection', 'recognition'])
-    face_detector = FaceAnalysis2(providers=['CPUExecutionProvider'], allowed_modules=['detection', 'recognition'])
-    face_detector.prepare(ctx_id=0, det_size=(640, 640))
-    #input_folder_name = './scarletthead_woman'
-    input_folder_name = sys.argv[1]
-    image_basename_list = os.listdir(input_folder_name)
-    image_path_list = sorted([os.path.join(input_folder_name, basename) for basename in image_basename_list])
-
-    input_id_images = []
-    for image_path in image_path_list:
-        input_id_images.append(load_image(image_path))
-    
-    id_embed_list = []
-    
-    for img in input_id_images:
-        img = np.array(img)
-        img = img[:, :, ::-1]
-        faces = analyze_faces(face_detector, img)
-        if len(faces) > 0:
-            id_embed_list.append(torch.from_numpy((faces[0]['embedding'])))
-    
-    if len(id_embed_list) == 0:
-        raise ValueError(f"No face detected in input image pool")
-    
-    id_embeds = torch.stack(id_embed_list)    
-    
-    # for r in id_embeds:
-    #     print(r)
-    # #torch.save(id_embeds, input_folder_name+'/id_embeds.pt');
-    # weights = dict()
-    # weights["id_embeds"] = id_embeds
-    # save_file(weights, input_folder_name+'/id_embeds.safetensors')
-
-    binary_data = id_embeds.numpy().tobytes()
-    two = 4
-    zero = 0
-    one = 1
-    tensor_name = "id_embeds"
-# Write binary data to a file
-    with open(input_folder_name+'/id_embeds.bin', "wb") as f:
-        f.write(two.to_bytes(4, byteorder='little'))
-        f.write((len(tensor_name)).to_bytes(4, byteorder='little'))
-        f.write(zero.to_bytes(4, byteorder='little'))
-        f.write((id_embeds.shape[1]).to_bytes(4, byteorder='little'))
-        f.write((id_embeds.shape[0]).to_bytes(4, byteorder='little'))
-        f.write(one.to_bytes(4, byteorder='little'))
-        f.write(one.to_bytes(4, byteorder='little'))
-        f.write(tensor_name.encode('ascii'))
-        f.write(binary_data)
-
+import os
+import sys
+
+import numpy as np
+import torch
+from diffusers.utils import load_image
+# pip install insightface==0.7.3
+from insightface.app import FaceAnalysis
+from insightface.data import get_image as ins_get_image
+from safetensors.torch import save_file
+
+### 
+# https://github.com/cubiq/ComfyUI_IPAdapter_plus/issues/165#issue-2055829543
+###
+class FaceAnalysis2(FaceAnalysis):
+    # NOTE: allows setting det_size for each detection call.
+    # the model allows it but the wrapping code from insightface
+    # doesn't show it, and people end up loading duplicate models
+    # for different sizes where there is absolutely no need to
+    def get(self, img, max_num=0, det_size=(640, 640)):
+        if det_size is not None:
+            self.det_model.input_size = det_size
+
+        return super().get(img, max_num)
+
+def analyze_faces(face_analysis: FaceAnalysis, img_data: np.ndarray, det_size=(640, 640)):
+    # NOTE: try detect faces, if no faces detected, lower det_size until it does
+    detection_sizes = [None] + [(size, size) for size in range(640, 256, -64)] + [(256, 256)]
+
+    for size in detection_sizes:
+        faces = face_analysis.get(img_data, det_size=size)
+        if len(faces) > 0:
+            return faces
+
+    return []
+
+if __name__ == "__main__":
+    #face_detector = FaceAnalysis2(providers=['CUDAExecutionProvider'], allowed_modules=['detection', 'recognition'])
+    face_detector = FaceAnalysis2(providers=['CPUExecutionProvider'], allowed_modules=['detection', 'recognition'])
+    face_detector.prepare(ctx_id=0, det_size=(640, 640))
+    #input_folder_name = './scarletthead_woman'
+    input_folder_name = sys.argv[1]
+    image_basename_list = os.listdir(input_folder_name)
+    image_path_list = sorted([os.path.join(input_folder_name, basename) for basename in image_basename_list])
+
+    input_id_images = []
+    for image_path in image_path_list:
+        input_id_images.append(load_image(image_path))
+    
+    id_embed_list = []
+    
+    for img in input_id_images:
+        img = np.array(img)
+        img = img[:, :, ::-1]
+        faces = analyze_faces(face_detector, img)
+        if len(faces) > 0:
+            id_embed_list.append(torch.from_numpy((faces[0]['embedding'])))
+    
+    if len(id_embed_list) == 0:
+        raise ValueError(f"No face detected in input image pool")
+    
+    id_embeds = torch.stack(id_embed_list)    
+    
+    # for r in id_embeds:
+    #     print(r)
+    # #torch.save(id_embeds, input_folder_name+'/id_embeds.pt');
+    # weights = dict()
+    # weights["id_embeds"] = id_embeds
+    # save_file(weights, input_folder_name+'/id_embeds.safetensors')
+
+    binary_data = id_embeds.numpy().tobytes()
+    two = 4
+    zero = 0
+    one = 1
+    tensor_name = "id_embeds"
+# Write binary data to a file
+    with open(input_folder_name+'/id_embeds.bin', "wb") as f:
+        f.write(two.to_bytes(4, byteorder='little'))
+        f.write((len(tensor_name)).to_bytes(4, byteorder='little'))
+        f.write(zero.to_bytes(4, byteorder='little'))
+        f.write((id_embeds.shape[1]).to_bytes(4, byteorder='little'))
+        f.write((id_embeds.shape[0]).to_bytes(4, byteorder='little'))
+        f.write(one.to_bytes(4, byteorder='little'))
+        f.write(one.to_bytes(4, byteorder='little'))
+        f.write(tensor_name.encode('ascii'))
+        f.write(binary_data)
+
     

src/anima.hpp

@@ -0,0 +1,686 @@
+#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(ggml_tensor* x,
+                                 ggml_tensor* timesteps,
+                                 ggml_tensor* context,
+                                 ggml_tensor* t5_ids     = nullptr,
+                                 ggml_tensor* t5_weights = nullptr) {
+            GGML_ASSERT(x->ne[3] == 1);
+            ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
+
+            x          = to_backend(x);
+            timesteps  = to_backend(timesteps);
+            context    = to_backend(context);
+            t5_ids     = to_backend(t5_ids);
+            t5_weights = to_backend(t5_weights);
+
+            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;
+        }
+
+        bool compute(int n_threads,
+                     ggml_tensor* x,
+                     ggml_tensor* timesteps,
+                     ggml_tensor* context,
+                     ggml_tensor* t5_ids      = nullptr,
+                     ggml_tensor* t5_weights  = nullptr,
+                     ggml_tensor** output     = nullptr,
+                     ggml_context* output_ctx = nullptr) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, t5_ids, t5_weights);
+            };
+            return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        }
+    };
+}  // namespace Anima
+
+#endif  // __ANIMA_HPP__

src/auto_encoder_kl.hpp

@@ -1,8 +1,7 @@
-#ifndef __VAE_HPP__
-#define __VAE_HPP__
+#ifndef __AUTO_ENCODER_KL_HPP__
+#define __AUTO_ENCODER_KL_HPP__
 
-#include "common.hpp"
-#include "ggml_extend.hpp"
+#include "vae.hpp"
 
 /*================================================== AutoEncoderKL ===================================================*/
 
@@ -30,7 +29,7 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         // x: [N, in_channels, h, w]
         // t_emb is always None
         auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
@@ -66,7 +65,7 @@ protected:
     int64_t in_channels;
     bool use_linear;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+    void init_params(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) {
@@ -102,7 +101,7 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -141,7 +140,7 @@ public:
             v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n);                        // [N, h * w, in_channels]
         }
 
-        h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, true, ctx->flash_attn_enabled);
+        h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
 
         if (use_linear) {
             h_ = proj_out->forward(ctx, h_);  // [N, h * w, in_channels]
@@ -179,8 +178,8 @@ public:
                                                                         {kernel_padding, 0, 0}));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x) override {
         // timesteps always None
         // skip_video always False
         // x: [N, IC, IH, IW]
@@ -209,7 +208,7 @@ public:
 
 class VideoResnetBlock : public ResnetBlock {
 protected:
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(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);
     }
@@ -228,7 +227,7 @@ public:
         blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, 0, out_channels, {video_kernel_size, 1}, 3, false, true));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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
@@ -318,7 +317,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 struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [N, in_channels, h, w]
 
         auto conv_in     = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"]);
@@ -436,7 +435,7 @@ public:
         blocks["conv_out"] = get_conv_out(block_in, out_ch, {3, 3}, {1, 1}, {1, 1});
     }
 
-    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) {
         // z: [N, z_channels, h, w]
         // alpha is always 0
         // merge_strategy is always learned
@@ -484,7 +483,7 @@ public:
 };
 
 // ldm.models.autoencoder.AutoencoderKL
-class AutoencodingEngine : public GGMLBlock {
+class AutoEncoderKLModel : public GGMLBlock {
 protected:
     SDVersion version;
     bool decode_only       = true;
@@ -503,7 +502,7 @@ protected:
     } dd_config;
 
 public:
-    AutoencodingEngine(SDVersion version          = VERSION_SD1,
+    AutoEncoderKLModel(SDVersion version          = VERSION_SD1,
                        bool decode_only           = true,
                        bool use_linear_projection = false,
                        bool use_video_decoder     = false)
@@ -550,7 +549,7 @@ public:
         }
     }
 
-    struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
+    ggml_tensor* decode(GGMLRunnerContext* ctx, 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]
@@ -582,7 +581,7 @@ public:
         return h;
     }
 
-    struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [N, in_channels, h, w]
         auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
 
@@ -611,48 +610,21 @@ public:
         }
         return z;
     }
-};
 
-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);
+    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;
         }
-        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";
+        return dd_config.z_channels * factor;
     }
 };
 
 struct AutoEncoderKL : public VAE {
-    bool decode_only = true;
-    AutoencodingEngine ae;
+    float scale_factor = 1.f;
+    float shift_factor = 0.f;
+    bool decode_only   = true;
+    AutoEncoderKLModel ae;
 
     AutoEncoderKL(ggml_backend_t backend,
                   bool offload_params_to_cpu,
@@ -661,7 +633,23 @@ struct AutoEncoderKL : public VAE {
                   bool decode_only       = false,
                   bool use_video_decoder = false,
                   SDVersion version      = VERSION_SD1)
-        : decode_only(decode_only), VAE(backend, offload_params_to_cpu) {
+        : 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;
+        }
         bool use_linear_projection = false;
         for (const auto& [name, tensor_storage] : tensor_storage_map) {
             if (!starts_with(name, prefix)) {
@@ -674,7 +662,7 @@ struct AutoEncoderKL : public VAE {
                 break;
             }
         }
-        ae = AutoencodingEngine(version, decode_only, use_linear_projection, use_video_decoder);
+        ae = AutoEncoderKLModel(version, decode_only, use_linear_projection, use_video_decoder);
         ae.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -693,31 +681,31 @@ struct AutoEncoderKL : public VAE {
         return "vae";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
         ae.get_param_tensors(tensors, prefix);
     }
 
-    struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
-        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+    ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
+        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
         z = to_backend(z);
 
         auto runner_ctx = get_context();
 
-        struct ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
+        ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.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) override {
+    bool _compute(const int n_threads,
+                  ggml_tensor* z,
+                  bool decode_graph,
+                  ggml_tensor** output,
+                  ggml_context* output_ctx = nullptr) override {
         GGML_ASSERT(!decode_only || decode_graph);
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
         // ggml_set_f32(z, 0.5f);
@@ -725,13 +713,183 @@ struct AutoEncoderKL : public VAE {
         return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
     }
 
+    ggml_tensor* gaussian_latent_sample(ggml_context* work_ctx, ggml_tensor* moments, std::shared_ptr<RNG> rng) {
+        // ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample
+        ggml_tensor* latents = ggml_new_tensor_4d(work_ctx, moments->type, moments->ne[0], moments->ne[1], moments->ne[2] / 2, moments->ne[3]);
+        ggml_tensor* noise   = ggml_dup_tensor(work_ctx, latents);
+        ggml_ext_im_set_randn_f32(noise, rng);
+        {
+            float mean   = 0;
+            float logvar = 0;
+            float value  = 0;
+            float std_   = 0;
+            for (int i = 0; i < latents->ne[3]; i++) {
+                for (int j = 0; j < latents->ne[2]; j++) {
+                    for (int k = 0; k < latents->ne[1]; k++) {
+                        for (int l = 0; l < latents->ne[0]; l++) {
+                            mean   = ggml_ext_tensor_get_f32(moments, l, k, j, i);
+                            logvar = ggml_ext_tensor_get_f32(moments, l, k, j + (int)latents->ne[2], i);
+                            logvar = std::max(-30.0f, std::min(logvar, 20.0f));
+                            std_   = std::exp(0.5f * logvar);
+                            value  = mean + std_ * ggml_ext_tensor_get_f32(noise, l, k, j, i);
+                            // printf("%d %d %d %d -> %f\n", i, j, k, l, value);
+                            ggml_ext_tensor_set_f32(latents, value, l, k, j, i);
+                        }
+                    }
+                }
+            }
+        }
+        return latents;
+    }
+
+    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+        if (sd_version_is_flux2(version)) {
+            return vae_output;
+        } else if (version == VERSION_SD1_PIX2PIX) {
+            return ggml_view_3d(work_ctx,
+                                vae_output,
+                                vae_output->ne[0],
+                                vae_output->ne[1],
+                                vae_output->ne[2] / 2,
+                                vae_output->nb[1],
+                                vae_output->nb[2],
+                                0);
+        } else {
+            return gaussian_latent_sample(work_ctx, vae_output, rng);
+        }
+    }
+
+    void get_latents_mean_std_vec(ggml_tensor* latents, int channel_dim, std::vector<float>& latents_mean_vec, std::vector<float>& latents_std_vec) {
+        // flux2
+        if (sd_version_is_flux2(version)) {
+            GGML_ASSERT(latents->ne[channel_dim] == 128);
+            latents_mean_vec = {-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};
+            latents_std_vec  = {
+                 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};
+        } else {
+            GGML_ABORT("unknown version %d", version);
+        }
+    }
+
+    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        ggml_tensor* vae_latents = ggml_dup(work_ctx, latents);
+        if (sd_version_is_flux2(version)) {
+            int channel_dim = 2;
+            std::vector<float> latents_mean_vec;
+            std::vector<float> latents_std_vec;
+            get_latents_mean_std_vec(latents, channel_dim, latents_mean_vec, latents_std_vec);
+
+            float mean;
+            float std_;
+            for (int i = 0; i < latents->ne[3]; i++) {
+                if (channel_dim == 3) {
+                    mean = latents_mean_vec[i];
+                    std_ = latents_std_vec[i];
+                }
+                for (int j = 0; j < latents->ne[2]; j++) {
+                    if (channel_dim == 2) {
+                        mean = latents_mean_vec[j];
+                        std_ = latents_std_vec[j];
+                    }
+                    for (int k = 0; k < latents->ne[1]; k++) {
+                        for (int l = 0; l < latents->ne[0]; l++) {
+                            float value = ggml_ext_tensor_get_f32(latents, l, k, j, i);
+                            value       = value * std_ / scale_factor + mean;
+                            ggml_ext_tensor_set_f32(vae_latents, value, l, k, j, i);
+                        }
+                    }
+                }
+            }
+        } else {
+            ggml_ext_tensor_iter(latents, [&](ggml_tensor* latents, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+                float value = ggml_ext_tensor_get_f32(latents, i0, i1, i2, i3);
+                value       = (value / scale_factor) + shift_factor;
+                ggml_ext_tensor_set_f32(vae_latents, value, i0, i1, i2, i3);
+            });
+        }
+        return vae_latents;
+    }
+
+    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        ggml_tensor* diffusion_latents = ggml_dup(work_ctx, latents);
+        if (sd_version_is_flux2(version)) {
+            int channel_dim = 2;
+            std::vector<float> latents_mean_vec;
+            std::vector<float> latents_std_vec;
+            get_latents_mean_std_vec(latents, channel_dim, latents_mean_vec, latents_std_vec);
+
+            float mean;
+            float std_;
+            for (int i = 0; i < latents->ne[3]; i++) {
+                if (channel_dim == 3) {
+                    mean = latents_mean_vec[i];
+                    std_ = latents_std_vec[i];
+                }
+                for (int j = 0; j < latents->ne[2]; j++) {
+                    if (channel_dim == 2) {
+                        mean = latents_mean_vec[j];
+                        std_ = latents_std_vec[j];
+                    }
+                    for (int k = 0; k < latents->ne[1]; k++) {
+                        for (int l = 0; l < latents->ne[0]; l++) {
+                            float value = ggml_ext_tensor_get_f32(latents, l, k, j, i);
+                            value       = (value - mean) * scale_factor / std_;
+                            ggml_ext_tensor_set_f32(diffusion_latents, value, l, k, j, i);
+                        }
+                    }
+                }
+            }
+        } else {
+            ggml_ext_tensor_iter(latents, [&](ggml_tensor* latents, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+                float value = ggml_ext_tensor_get_f32(latents, i0, i1, i2, i3);
+                value       = (value - shift_factor) * scale_factor;
+                ggml_ext_tensor_set_f32(diffusion_latents, value, i0, i1, i2, i3);
+            });
+        }
+        return diffusion_latents;
+    }
+
+    int get_encoder_output_channels(int input_channels) {
+        return ae.get_encoder_output_channels();
+    }
+
     void test() {
-        struct ggml_init_params params;
+        ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        struct ggml_context* work_ctx = ggml_init(params);
+        ggml_context* work_ctx = ggml_init(params);
         GGML_ASSERT(work_ctx != nullptr);
 
         {
@@ -742,10 +900,10 @@ struct AutoEncoderKL : public VAE {
             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;
+            ggml_tensor* out = nullptr;
 
             int64_t t0 = ggml_time_ms();
-            compute(8, x, false, &out, work_ctx);
+            _compute(8, x, false, &out, work_ctx);
             int64_t t1 = ggml_time_ms();
 
             print_ggml_tensor(out);
@@ -760,10 +918,10 @@ struct AutoEncoderKL : public VAE {
             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;
+            ggml_tensor* out = nullptr;
 
             int64_t t0 = ggml_time_ms();
-            compute(8, z, true, &out, work_ctx);
+            _compute(8, z, true, &out, work_ctx);
             int64_t t1 = ggml_time_ms();
 
             print_ggml_tensor(out);
@@ -772,4 +930,4 @@ struct AutoEncoderKL : public VAE {
     };
 };
 
-#endif
+#endif  // __AUTO_ENCODER_KL_HPP__

src/cache_dit.hpp

@@ -603,87 +603,6 @@ inline std::vector<int> generate_scm_mask(
     return mask;
 }
 
-inline std::vector<int> get_scm_preset(const std::string& preset, int total_steps) {
-    struct Preset {
-        std::vector<int> compute_bins;
-        std::vector<int> cache_bins;
-    };
-
-    Preset slow   = {{8, 3, 3, 2, 1, 1}, {1, 2, 2, 2, 3}};
-    Preset medium = {{6, 2, 2, 2, 2, 1}, {1, 3, 3, 3, 3}};
-    Preset fast   = {{6, 1, 1, 1, 1, 1}, {1, 3, 4, 5, 4}};
-    Preset ultra  = {{4, 1, 1, 1, 1}, {2, 5, 6, 7}};
-
-    Preset* p = nullptr;
-    if (preset == "slow" || preset == "s" || preset == "S")
-        p = &slow;
-    else if (preset == "medium" || preset == "m" || preset == "M")
-        p = &medium;
-    else if (preset == "fast" || preset == "f" || preset == "F")
-        p = &fast;
-    else if (preset == "ultra" || preset == "u" || preset == "U")
-        p = &ultra;
-    else
-        return {};
-
-    if (total_steps != 28 && total_steps > 0) {
-        float scale = static_cast<float>(total_steps) / 28.0f;
-        std::vector<int> scaled_compute, scaled_cache;
-
-        for (int v : p->compute_bins) {
-            scaled_compute.push_back(std::max(1, static_cast<int>(v * scale + 0.5f)));
-        }
-        for (int v : p->cache_bins) {
-            scaled_cache.push_back(std::max(1, static_cast<int>(v * scale + 0.5f)));
-        }
-
-        return generate_scm_mask(scaled_compute, scaled_cache, total_steps);
-    }
-
-    return generate_scm_mask(p->compute_bins, p->cache_bins, total_steps);
-}
-
-inline float get_preset_threshold(const std::string& preset) {
-    if (preset == "slow" || preset == "s" || preset == "S")
-        return 0.20f;
-    if (preset == "medium" || preset == "m" || preset == "M")
-        return 0.25f;
-    if (preset == "fast" || preset == "f" || preset == "F")
-        return 0.30f;
-    if (preset == "ultra" || preset == "u" || preset == "U")
-        return 0.34f;
-    return 0.08f;
-}
-
-inline int get_preset_warmup(const std::string& preset) {
-    if (preset == "slow" || preset == "s" || preset == "S")
-        return 8;
-    if (preset == "medium" || preset == "m" || preset == "M")
-        return 6;
-    if (preset == "fast" || preset == "f" || preset == "F")
-        return 6;
-    if (preset == "ultra" || preset == "u" || preset == "U")
-        return 4;
-    return 8;
-}
-
-inline int get_preset_Fn(const std::string& preset) {
-    if (preset == "slow" || preset == "s" || preset == "S")
-        return 8;
-    if (preset == "medium" || preset == "m" || preset == "M")
-        return 8;
-    if (preset == "fast" || preset == "f" || preset == "F")
-        return 6;
-    if (preset == "ultra" || preset == "u" || preset == "U")
-        return 4;
-    return 8;
-}
-
-inline int get_preset_Bn(const std::string& preset) {
-    (void)preset;
-    return 0;
-}
-
 inline void parse_dbcache_options(const std::string& opts, DBCacheConfig& cfg) {
     if (opts.empty())
         return;
@@ -880,7 +799,7 @@ struct CacheDitConditionState {
         }
     }
 
-    bool before_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output, float sigma, int step_index) {
+    bool before_condition(const void* cond, ggml_tensor* input, ggml_tensor* output, float sigma, int step_index) {
         if (!enabled() || step_index < 0)
             return false;
 
@@ -948,7 +867,7 @@ struct CacheDitConditionState {
         return false;
     }
 
-    void after_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output) {
+    void after_condition(const void* cond, ggml_tensor* input, ggml_tensor* output) {
         if (!step_is_active())
             return;
 

src/clip.hpp

@@ -4,6 +4,7 @@
 #include "ggml_extend.hpp"
 #include "model.h"
 #include "tokenize_util.h"
+#include "vocab/vocab.h"
 
 /*================================================== CLIPTokenizer ===================================================*/
 
@@ -110,7 +111,7 @@ public:
         if (merges_utf8_str.size() > 0) {
             load_from_merges(merges_utf8_str);
         } else {
-            load_from_merges(ModelLoader::load_merges());
+            load_from_merges(load_clip_merges());
         }
         add_special_token("<|startoftext|>");
         add_special_token("<|endoftext|>");
@@ -472,7 +473,7 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -510,7 +511,7 @@ public:
         blocks["mlp"] = std::shared_ptr<GGMLBlock>(new CLIPMLP(d_model, intermediate_size));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* mask = nullptr) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask = nullptr) {
         // 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"]);
@@ -540,10 +541,10 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* mask = nullptr,
-                                int clip_skip            = -1) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* mask = nullptr,
+                         int clip_skip     = -1) {
         // x: [N, n_token, d_model]
         int layer_idx = n_layer - 1;
         // LOG_DEBUG("clip_skip %d", clip_skip);
@@ -572,7 +573,7 @@ protected:
     int64_t num_positions;
     bool force_clip_f32;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(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);
@@ -596,13 +597,13 @@ public:
           force_clip_f32(force_clip_f32) {
     }
 
-    struct ggml_tensor* get_token_embed_weight() {
+    ggml_tensor* get_token_embed_weight() {
         return params["token_embedding.weight"];
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* input_ids,
-                                struct ggml_tensor* custom_embed_weight) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* input_ids,
+                         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"];
@@ -629,7 +630,7 @@ protected:
     int num_patches;
     int64_t num_positions;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(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;
@@ -652,7 +653,7 @@ public:
         num_positions = num_patches + 1;
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* pixel_values) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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);
@@ -662,20 +663,20 @@ public:
         auto position_embed_weight = params["position_embedding.weight"];
 
         // concat(patch_embedding, class_embedding) + position_embedding
-        struct ggml_tensor* patch_embedding;
+        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]
 
-        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* 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* 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);
+        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]
     }
 };
@@ -692,7 +693,7 @@ enum CLIPVersion {
 
 class CLIPTextModel : public GGMLBlock {
 protected:
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(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);
@@ -733,18 +734,18 @@ public:
         blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
     }
 
-    struct ggml_tensor* get_token_embed_weight() {
+    ggml_tensor* get_token_embed_weight() {
         auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
         return embeddings->get_token_embed_weight();
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* input_ids,
-                                struct ggml_tensor* tkn_embeddings,
-                                struct ggml_tensor* mask = nullptr,
-                                size_t max_token_idx     = 0,
-                                bool return_pooled       = false,
-                                int clip_skip            = -1) {
+    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) {
         // input_ids: [N, n_token]
         auto embeddings       = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
         auto encoder          = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
@@ -803,10 +804,10 @@ public:
         blocks["post_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* pixel_values,
-                                bool return_pooled = true,
-                                int clip_skip      = -1) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         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"]);
@@ -838,7 +839,7 @@ protected:
     int64_t out_features;
     bool transpose_weight;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(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);
@@ -855,8 +856,8 @@ public:
           out_features(out_features),
           transpose_weight(transpose_weight) {}
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
-        struct ggml_tensor* w = params["weight"];
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
         if (transpose_weight) {
             w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
         }
@@ -885,10 +886,10 @@ public:
         blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new CLIPProjection(hidden_size, projection_dim, transpose_proj_w));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* pixel_values,
-                                bool return_pooled = true,
-                                int clip_skip      = -1) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         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"]);
@@ -935,17 +936,17 @@ struct CLIPTextModelRunner : public GGMLRunner {
         return "clip";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* input_ids,
-                                struct ggml_tensor* embeddings,
-                                struct ggml_tensor* mask,
-                                size_t max_token_idx = 0,
-                                bool return_pooled   = false,
-                                int clip_skip        = -1) {
+    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) {
         size_t N       = input_ids->ne[1];
         size_t n_token = input_ids->ne[0];
         if (input_ids->ne[0] > model.n_token) {
@@ -956,17 +957,17 @@ struct CLIPTextModelRunner : public GGMLRunner {
         return model.forward(ctx, input_ids, embeddings, mask, max_token_idx, return_pooled, clip_skip);
     }
 
-    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_cgraph* build_graph(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) {
+        ggml_cgraph* gf = new_graph_custom(2048);
 
         input_ids = to_backend(input_ids);
 
-        struct ggml_tensor* embeddings = nullptr;
+        ggml_tensor* embeddings = nullptr;
 
         if (num_custom_embeddings > 0 && custom_embeddings_data != nullptr) {
             auto token_embed_weight = model.get_token_embed_weight();
@@ -996,7 +997,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
 
         auto runner_ctx = get_context();
 
-        struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, attention_mask, max_token_idx, return_pooled, clip_skip);
+        ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, attention_mask, max_token_idx, return_pooled, clip_skip);
 
         ggml_build_forward_expand(gf, hidden_states);
 
@@ -1004,7 +1005,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
     }
 
     bool compute(const int n_threads,
-                 struct ggml_tensor* input_ids,
+                 ggml_tensor* input_ids,
                  int num_custom_embeddings,
                  void* custom_embeddings_data,
                  size_t max_token_idx,
@@ -1012,7 +1013,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
                  int clip_skip,
                  ggml_tensor** output,
                  ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
         return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);

src/common_block.hpp

@@ -1,5 +1,5 @@
-#ifndef __COMMON_HPP__
-#define __COMMON_HPP__
+#ifndef __COMMON_BLOCK_HPP__
+#define __COMMON_BLOCK_HPP__
 
 #include "ggml_extend.hpp"
 
@@ -23,7 +23,7 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [N, channels, h, w]
         if (vae_downsample) {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
@@ -52,7 +52,7 @@ public:
         blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [N, channels, h, w]
         auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
@@ -121,7 +121,7 @@ public:
         }
     }
 
-    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* emb = nullptr) {
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, 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,7 +188,7 @@ public:
         blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out * 2));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -214,7 +214,7 @@ public:
         blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out, bias));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -258,7 +258,7 @@ public:
         blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [ne3, ne2, ne1, dim]
         // return: [ne3, ne2, ne1, dim_out]
 
@@ -297,9 +297,9 @@ public:
         // to_out_1 is nn.Dropout(), skip for inference
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* context) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* context) {
         // x: [N, n_token, query_dim]
         // context: [N, n_context, context_dim]
         // return: [N, n_token, query_dim]
@@ -355,9 +355,9 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* context) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* context) {
         // x: [N, n_token, query_dim]
         // context: [N, n_context, context_dim]
         // return: [N, n_token, query_dim]
@@ -406,7 +406,7 @@ protected:
     int64_t context_dim = 768;  // hidden_size, 1024 for VERSION_SD2
     bool use_linear     = false;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+    void init_params(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 +456,9 @@ public:
         }
     }
 
-    virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                        struct ggml_tensor* x,
-                                        struct ggml_tensor* context) {
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                 ggml_tensor* x,
+                                 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 +510,7 @@ public:
 
 class AlphaBlender : public GGMLBlock {
 protected:
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+    void init_params(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,9 +530,9 @@ public:
         // since mix_factor.shape is [1,], we don't need rearrange using  rearrange_pattern
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x_spatial,
-                                struct ggml_tensor* x_temporal) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x_spatial,
+                         ggml_tensor* x_temporal) {
         // image_only_indicator is always tensor([0.])
         float alpha = get_alpha();
         auto x      = ggml_add(ctx->ggml_ctx,
@@ -555,10 +555,10 @@ public:
         blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* emb,
-                                int num_video_frames) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         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 +590,4 @@ public:
     }
 };
 
-#endif  // __COMMON_HPP__
+#endif  // __COMMON_BLOCK_HPP__

src/common_dit.hpp

@@ -0,0 +1,108 @@
+#ifndef __COMMON_DIT_HPP__
+#define __COMMON_DIT_HPP__
+
+#include "ggml_extend.hpp"
+
+namespace DiT {
+    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;
+    }
+
+    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;
+    }
+
+    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;
+    }
+
+    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;
+    }
+
+    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/conditioner.hpp

@@ -6,13 +6,18 @@
 #include "t5.hpp"
 
 struct SDCondition {
-    struct ggml_tensor* c_crossattn = nullptr;  // aka context
-    struct ggml_tensor* c_vector    = nullptr;  // aka y
-    struct ggml_tensor* c_concat    = nullptr;
+    ggml_tensor* c_crossattn = nullptr;  // aka context
+    ggml_tensor* c_vector    = nullptr;  // aka y
+    ggml_tensor* c_concat    = nullptr;
+
+    std::vector<ggml_tensor*> extra_c_crossattns;
 
     SDCondition() = default;
-    SDCondition(struct ggml_tensor* c_crossattn, struct ggml_tensor* c_vector, struct ggml_tensor* c_concat)
-        : c_crossattn(c_crossattn), c_vector(c_vector), c_concat(c_concat) {}
+    SDCondition(ggml_tensor* c_crossattn,
+                ggml_tensor* c_vector,
+                ggml_tensor* c_concat,
+                const std::vector<ggml_tensor*>& extra_c_crossattns = {})
+        : c_crossattn(c_crossattn), c_vector(c_vector), c_concat(c_concat), extra_c_crossattns(extra_c_crossattns) {}
 };
 
 struct ConditionerParams {
@@ -32,7 +37,7 @@ struct Conditioner {
                                               const ConditionerParams& conditioner_params) = 0;
     virtual void alloc_params_buffer()                                                     = 0;
     virtual void free_params_buffer()                                                      = 0;
-    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors)    = 0;
+    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
     virtual size_t get_params_buffer_size()                                                = 0;
     virtual void set_flash_attention_enabled(bool enabled)                                 = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
@@ -87,7 +92,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         text_model->get_param_tensors(tensors, "cond_stage_model.transformer.text_model");
         if (sd_version_is_sdxl(version)) {
             text_model2->get_param_tensors(tensors, "cond_stage_model.1.transformer.text_model");
@@ -144,14 +149,14 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             }
             return true;
         }
-        struct ggml_init_params params;
-        params.mem_size               = 100 * 1024 * 1024;  // max for custom embeddings 100 MB
-        params.mem_buffer             = nullptr;
-        params.no_alloc               = false;
-        struct ggml_context* embd_ctx = ggml_init(params);
-        struct ggml_tensor* embd      = nullptr;
-        struct ggml_tensor* embd2     = nullptr;
-        auto on_load                  = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
+        ggml_init_params params;
+        params.mem_size        = 100 * 1024 * 1024;  // max for custom embeddings 100 MB
+        params.mem_buffer      = nullptr;
+        params.no_alloc        = false;
+        ggml_context* embd_ctx = ggml_init(params);
+        ggml_tensor* embd      = nullptr;
+        ggml_tensor* embd2     = nullptr;
+        auto on_load           = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
             if (tensor_storage.ne[0] != text_model->model.hidden_size) {
                 if (text_model2) {
                     if (tensor_storage.ne[0] == text_model2->model.hidden_size) {
@@ -430,12 +435,12 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                              int height,
                                              int adm_in_channels  = -1,
                                              bool zero_out_masked = false) {
-        int64_t t0                               = ggml_time_ms();
-        struct ggml_tensor* hidden_states        = nullptr;  // [N, n_token, hidden_size]
-        struct ggml_tensor* chunk_hidden_states  = nullptr;  // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2]
-        struct ggml_tensor* chunk_hidden_states1 = nullptr;  // [n_token, hidden_size]
-        struct ggml_tensor* chunk_hidden_states2 = nullptr;  // [n_token, hidden_size2]
-        struct ggml_tensor* pooled               = nullptr;
+        int64_t t0                        = ggml_time_ms();
+        ggml_tensor* hidden_states        = nullptr;  // [N, n_token, hidden_size]
+        ggml_tensor* chunk_hidden_states  = nullptr;  // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2]
+        ggml_tensor* chunk_hidden_states1 = nullptr;  // [n_token, hidden_size]
+        ggml_tensor* chunk_hidden_states2 = nullptr;  // [n_token, hidden_size2]
+        ggml_tensor* pooled               = nullptr;
         std::vector<float> hidden_states_vec;
 
         if (clip_skip <= 0) {
@@ -450,9 +455,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             std::vector<float> chunk_weights(weights.begin() + chunk_idx * chunk_len,
                                              weights.begin() + (chunk_idx + 1) * chunk_len);
 
-            auto input_ids                 = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
-            struct ggml_tensor* input_ids2 = nullptr;
-            size_t max_token_idx           = 0;
+            auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
+            ggml_tensor* input_ids2 = nullptr;
+            size_t max_token_idx    = 0;
             if (sd_version_is_sdxl(version)) {
                 auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer.EOS_TOKEN_ID);
                 if (it != chunk_tokens.end()) {
@@ -671,18 +676,18 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
         return "clip_vision";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
         vision_model.get_param_tensors(tensors, "cond_stage_model.transformer");
     }
 
-    struct ggml_cgraph* build_graph(struct ggml_tensor* pixel_values, bool return_pooled, int clip_skip) {
-        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+    ggml_cgraph* build_graph(ggml_tensor* pixel_values, bool return_pooled, int clip_skip) {
+        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
         pixel_values = to_backend(pixel_values);
 
         auto runner_ctx = get_context();
 
-        struct ggml_tensor* hidden_states = vision_model.forward(&runner_ctx, pixel_values, return_pooled, clip_skip);
+        ggml_tensor* hidden_states = vision_model.forward(&runner_ctx, pixel_values, return_pooled, clip_skip);
 
         ggml_build_forward_expand(gf, hidden_states);
 
@@ -695,7 +700,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
                  int clip_skip,
                  ggml_tensor** output,
                  ggml_context* output_ctx) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(pixel_values, return_pooled, clip_skip);
         };
         return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
@@ -741,7 +746,7 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (clip_l) {
             clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
         }
@@ -904,15 +909,15 @@ struct SD3CLIPEmbedder : public Conditioner {
             clip_skip = 2;
         }
 
-        int64_t t0                                 = ggml_time_ms();
-        struct ggml_tensor* hidden_states          = nullptr;  // [N, n_token*2, 4096]
-        struct ggml_tensor* chunk_hidden_states    = nullptr;  // [n_token*2, 4096]
-        struct ggml_tensor* chunk_hidden_states_l  = nullptr;  // [n_token, hidden_size_l]
-        struct ggml_tensor* chunk_hidden_states_g  = nullptr;  // [n_token, hidden_size_g]
-        struct ggml_tensor* chunk_hidden_states_t5 = nullptr;  // [n_token, hidden_size_t5]
-        struct ggml_tensor* pooled                 = nullptr;
-        struct ggml_tensor* pooled_l               = nullptr;  // [768,]
-        struct ggml_tensor* pooled_g               = nullptr;  // [1280,]
+        int64_t t0                          = ggml_time_ms();
+        ggml_tensor* hidden_states          = nullptr;  // [N, n_token*2, 4096]
+        ggml_tensor* chunk_hidden_states    = nullptr;  // [n_token*2, 4096]
+        ggml_tensor* chunk_hidden_states_l  = nullptr;  // [n_token, hidden_size_l]
+        ggml_tensor* chunk_hidden_states_g  = nullptr;  // [n_token, hidden_size_g]
+        ggml_tensor* chunk_hidden_states_t5 = nullptr;  // [n_token, hidden_size_t5]
+        ggml_tensor* pooled                 = nullptr;
+        ggml_tensor* pooled_l               = nullptr;  // [768,]
+        ggml_tensor* pooled_g               = nullptr;  // [1280,]
         std::vector<float> hidden_states_vec;
 
         size_t chunk_len   = 77;
@@ -1173,7 +1178,7 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (clip_l) {
             clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
         }
@@ -1301,10 +1306,10 @@ struct FluxCLIPEmbedder : public Conditioner {
             clip_skip = 2;
         }
 
-        int64_t t0                              = ggml_time_ms();
-        struct ggml_tensor* hidden_states       = nullptr;  // [N, n_token, 4096]
-        struct ggml_tensor* chunk_hidden_states = nullptr;  // [n_token, 4096]
-        struct ggml_tensor* pooled              = nullptr;  // [768,]
+        int64_t t0                       = ggml_time_ms();
+        ggml_tensor* hidden_states       = nullptr;  // [N, n_token, 4096]
+        ggml_tensor* chunk_hidden_states = nullptr;  // [n_token, 4096]
+        ggml_tensor* pooled              = nullptr;  // [768,]
         std::vector<float> hidden_states_vec;
 
         size_t chunk_count = std::max(clip_l_tokens.size() > 0 ? chunk_len : 0, t5_tokens.size()) / chunk_len;
@@ -1443,7 +1448,7 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (t5) {
             t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
         }
@@ -1518,7 +1523,7 @@ struct T5CLIPEmbedder : public Conditioner {
         return {t5_tokens, t5_weights, t5_mask};
     }
 
-    void modify_mask_to_attend_padding(struct ggml_tensor* mask, int max_seq_length, int num_extra_padding = 8) {
+    void modify_mask_to_attend_padding(ggml_tensor* mask, int max_seq_length, int num_extra_padding = 8) {
         float* mask_data = (float*)mask->data;
         int num_pad      = 0;
         for (int64_t i = 0; i < max_seq_length; i++) {
@@ -1549,11 +1554,11 @@ struct T5CLIPEmbedder : public Conditioner {
         auto& t5_weights       = std::get<1>(token_and_weights);
         auto& t5_attn_mask_vec = std::get<2>(token_and_weights);
 
-        int64_t t0                              = ggml_time_ms();
-        struct ggml_tensor* hidden_states       = nullptr;  // [N, n_token, 4096]
-        struct ggml_tensor* chunk_hidden_states = nullptr;  // [n_token, 4096]
-        struct ggml_tensor* pooled              = nullptr;
-        struct ggml_tensor* t5_attn_mask        = vector_to_ggml_tensor(work_ctx, t5_attn_mask_vec);  // [n_token]
+        int64_t t0                       = ggml_time_ms();
+        ggml_tensor* hidden_states       = nullptr;  // [N, n_token, 4096]
+        ggml_tensor* chunk_hidden_states = nullptr;  // [n_token, 4096]
+        ggml_tensor* pooled              = nullptr;
+        ggml_tensor* t5_attn_mask        = vector_to_ggml_tensor(work_ctx, t5_attn_mask_vec);  // [n_token]
 
         std::vector<float> hidden_states_vec;
 
@@ -1636,6 +1641,142 @@ struct T5CLIPEmbedder : public Conditioner {
     }
 };
 
+struct AnimaConditioner : public Conditioner {
+    std::shared_ptr<LLM::BPETokenizer> qwen_tokenizer;
+    T5UniGramTokenizer t5_tokenizer;
+    std::shared_ptr<LLM::LLMRunner> llm;
+
+    AnimaConditioner(ggml_backend_t backend,
+                     bool offload_params_to_cpu,
+                     const String2TensorStorage& tensor_storage_map = {}) {
+        qwen_tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
+        llm            = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
+                                               backend,
+                                               offload_params_to_cpu,
+                                               tensor_storage_map,
+                                               "text_encoders.llm",
+                                               false);
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
+    void alloc_params_buffer() override {
+        llm->alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        llm->free_params_buffer();
+    }
+
+    size_t get_params_buffer_size() override {
+        return llm->get_params_buffer_size();
+    }
+
+    void set_flash_attention_enabled(bool enabled) override {
+        llm->set_flash_attention_enabled(enabled);
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        llm->set_weight_adapter(adapter);
+    }
+
+    std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
+        auto parsed_attention = parse_prompt_attention(text);
+
+        {
+            std::stringstream ss;
+            ss << "[";
+            for (const auto& item : parsed_attention) {
+                ss << "['" << item.first << "', " << item.second << "], ";
+            }
+            ss << "]";
+            LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
+        }
+
+        std::vector<int> qwen_tokens;
+        std::vector<float> qwen_weights;
+        std::vector<int> t5_tokens;
+        std::vector<float> t5_weights;
+
+        for (const auto& item : parsed_attention) {
+            const std::string& curr_text = item.first;
+            std::vector<int> curr_tokens = qwen_tokenizer->tokenize(curr_text, nullptr);
+            qwen_tokens.insert(qwen_tokens.end(), curr_tokens.begin(), curr_tokens.end());
+            // Anima uses uniform Qwen token weights.
+            qwen_weights.insert(qwen_weights.end(), curr_tokens.size(), 1.f);
+        }
+        if (qwen_tokens.empty()) {
+            qwen_tokens.push_back(151643);  // qwen3 pad token
+            qwen_weights.push_back(1.f);
+        }
+
+        for (const auto& item : parsed_attention) {
+            const std::string& curr_text = item.first;
+            float curr_weight            = item.second;
+            std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
+            t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
+            t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
+        }
+
+        return {qwen_tokens, qwen_weights, t5_tokens, t5_weights};
+    }
+
+    SDCondition get_learned_condition(ggml_context* work_ctx,
+                                      int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        int64_t t0 = ggml_time_ms();
+
+        auto tokenized     = tokenize(conditioner_params.text);
+        auto& qwen_tokens  = std::get<0>(tokenized);
+        auto& qwen_weights = std::get<1>(tokenized);
+        auto& t5_tokens    = std::get<2>(tokenized);
+        auto& t5_weights   = std::get<3>(tokenized);
+
+        auto input_ids = vector_to_ggml_tensor_i32(work_ctx, qwen_tokens);
+
+        ggml_tensor* hidden_states = nullptr;  // [N, n_token, 1024]
+        llm->compute(n_threads,
+                     input_ids,
+                     nullptr,
+                     {},
+                     {},
+                     &hidden_states,
+                     work_ctx);
+
+        {
+            auto tensor         = hidden_states;
+            float original_mean = ggml_ext_tensor_mean(tensor);
+            for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+                for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                    for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                        float value = ggml_ext_tensor_get_f32(tensor, i0, i1, i2);
+                        value *= qwen_weights[i1];
+                        ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2);
+                    }
+                }
+            }
+            float new_mean = ggml_ext_tensor_mean(tensor);
+            if (new_mean != 0.f) {
+                ggml_ext_tensor_scale_inplace(tensor, (original_mean / new_mean));
+            }
+        }
+
+        ggml_tensor* t5_ids_tensor    = nullptr;
+        ggml_tensor* t5_weight_tensor = nullptr;
+        if (!t5_tokens.empty()) {
+            t5_ids_tensor    = vector_to_ggml_tensor_i32(work_ctx, t5_tokens);
+            t5_weight_tensor = vector_to_ggml_tensor(work_ctx, t5_weights);
+        }
+
+        int64_t t1 = ggml_time_ms();
+        LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+
+        return {hidden_states, t5_weight_tensor, t5_ids_tensor};
+    }
+};
+
 struct LLMEmbedder : public Conditioner {
     SDVersion version;
     std::shared_ptr<LLM::BPETokenizer> tokenizer;
@@ -1667,7 +1808,7 @@ struct LLMEmbedder : public Conditioner {
                                                enable_vision);
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
 
@@ -1696,18 +1837,23 @@ struct LLMEmbedder : public Conditioner {
     }
 
     std::tuple<std::vector<int>, std::vector<float>> tokenize(std::string text,
-                                                              std::pair<int, int> attn_range,
+                                                              const std::pair<int, int>& attn_range,
                                                               size_t max_length = 0,
                                                               bool padding      = false) {
         std::vector<std::pair<std::string, float>> parsed_attention;
-        parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
-        if (attn_range.second - attn_range.first > 0) {
-            auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
-            parsed_attention.insert(parsed_attention.end(),
-                                    new_parsed_attention.begin(),
-                                    new_parsed_attention.end());
+        if (attn_range.first >= 0 && attn_range.second > 0) {
+            parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
+            if (attn_range.second - attn_range.first > 0) {
+                auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
+                parsed_attention.insert(parsed_attention.end(),
+                                        new_parsed_attention.begin(),
+                                        new_parsed_attention.end());
+            }
+            parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
+        } else {
+            parsed_attention.emplace_back(text, 1.f);
         }
-        parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
+
         {
             std::stringstream ss;
             ss << "[";
@@ -1738,156 +1884,27 @@ struct LLMEmbedder : public Conditioner {
         return {tokens, weights};
     }
 
-    SDCondition get_learned_condition(ggml_context* work_ctx,
-                                      int n_threads,
-                                      const ConditionerParams& conditioner_params) override {
-        std::string prompt;
-        std::vector<std::pair<int, ggml_tensor*>> image_embeds;
-        std::pair<int, int> prompt_attn_range;
-        int prompt_template_encode_start_idx = 34;
-        int max_length                       = 0;
-        std::set<int> out_layers;
-        std::vector<int> tokens;
-        std::vector<float> weights;
+    ggml_tensor* encode_prompt(ggml_context* work_ctx,
+                               int n_threads,
+                               const std::string prompt,
+                               const std::pair<int, int>& prompt_attn_range,
+                               int max_length,
+                               int min_length,
+                               std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                               const std::set<int>& out_layers,
+                               int prompt_template_encode_start_idx) {
+        auto tokens_and_weights = tokenize(prompt, prompt_attn_range);
+        auto& tokens            = std::get<0>(tokens_and_weights);
+        auto& weights           = std::get<1>(tokens_and_weights);
         std::vector<float> mask;
-        if (llm->enable_vision && conditioner_params.ref_images.size() > 0) {
-            LOG_INFO("QwenImageEditPlusPipeline");
-            prompt_template_encode_start_idx = 64;
-            int image_embed_idx              = 64 + 6;
-
-            int min_pixels          = 384 * 384;
-            int max_pixels          = 560 * 560;
-            std::string placeholder = "<|image_pad|>";
-            std::string img_prompt;
-
-            for (int i = 0; i < conditioner_params.ref_images.size(); i++) {
-                sd_image_f32_t image = sd_image_t_to_sd_image_f32_t(*conditioner_params.ref_images[i]);
-                double factor        = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size;
-                int height           = image.height;
-                int width            = image.width;
-                int h_bar            = static_cast<int>(std::round(height / factor) * factor);
-                int w_bar            = static_cast<int>(std::round(width / factor) * factor);
-
-                if (static_cast<double>(h_bar) * w_bar > max_pixels) {
-                    double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
-                    h_bar       = std::max(static_cast<int>(factor),
-                                           static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
-                    w_bar       = std::max(static_cast<int>(factor),
-                                           static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
-                } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
-                    double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
-                    h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
-                    w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
-                }
-
-                LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, image.height, image.width, h_bar, w_bar);
-
-                sd_image_f32_t resized_image = clip_preprocess(image, w_bar, h_bar);
-                free(image.data);
-                image.data = nullptr;
-
-                ggml_tensor* image_tensor = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, resized_image.width, resized_image.height, 3, 1);
-                sd_image_f32_to_ggml_tensor(resized_image, image_tensor, false);
-                free(resized_image.data);
-                resized_image.data = nullptr;
-
-                ggml_tensor* image_embed = nullptr;
-                llm->encode_image(n_threads, image_tensor, &image_embed, work_ctx);
-                image_embeds.emplace_back(image_embed_idx, image_embed);
-                image_embed_idx += 1 + static_cast<int>(image_embed->ne[1]) + 6;
-
-                img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";  // [24669, 220, index, 25, 220, 151652]
-                int64_t num_image_tokens = image_embed->ne[1];
-                img_prompt.reserve(num_image_tokens * placeholder.size());
-                for (int j = 0; j < num_image_tokens; j++) {
-                    img_prompt += placeholder;
-                }
-                img_prompt += "<|vision_end|>";
-            }
-
-            prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
-            prompt += img_prompt;
-
-            prompt_attn_range.first = static_cast<int>(prompt.size());
-            prompt += conditioner_params.text;
-            prompt_attn_range.second = static_cast<int>(prompt.size());
-
-            prompt += "<|im_end|>\n<|im_start|>assistant\n";
-        } else if (version == VERSION_FLUX2) {
-            prompt_template_encode_start_idx = 0;
-            out_layers                       = {10, 20, 30};
-
-            prompt = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
-
-            prompt_attn_range.first = static_cast<int>(prompt.size());
-            prompt += conditioner_params.text;
-            prompt_attn_range.second = static_cast<int>(prompt.size());
-
-            prompt += "[/INST]";
-        } else if (sd_version_is_z_image(version)) {
-            prompt_template_encode_start_idx = 0;
-            out_layers                       = {35};  // -2
-
-            prompt = "<|im_start|>user\n";
-
-            prompt_attn_range.first = static_cast<int>(prompt.size());
-            prompt += conditioner_params.text;
-            prompt_attn_range.second = static_cast<int>(prompt.size());
-
-            prompt += "<|im_end|>\n<|im_start|>assistant\n";
-        } else if (version == VERSION_FLUX2_KLEIN) {
-            prompt_template_encode_start_idx = 0;
-            max_length                       = 512;
-            out_layers                       = {9, 18, 27};
-
-            prompt = "<|im_start|>user\n";
-
-            prompt_attn_range.first = static_cast<int>(prompt.size());
-            prompt += conditioner_params.text;
-            prompt_attn_range.second = static_cast<int>(prompt.size());
-
-            prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
-
-            auto tokens_and_weights = tokenize(prompt, prompt_attn_range, 0, false);
-            tokens                  = std::get<0>(tokens_and_weights);
-            weights                 = std::get<1>(tokens_and_weights);
 
+        if (max_length > 0 && tokens.size() < max_length) {
             mask.insert(mask.end(), tokens.size(), 1.f);
-            if (tokens.size() < max_length) {
-                mask.insert(mask.end(), max_length - tokens.size(), 0.f);
-                tokenizer->pad_tokens(tokens, weights, max_length, true);
-            }
-        } else if (version == VERSION_OVIS_IMAGE) {
-            prompt_template_encode_start_idx = 28;
-            max_length                       = prompt_template_encode_start_idx + 256;
-
-            prompt = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:";
-
-            prompt_attn_range.first = static_cast<int>(prompt.size());
-            prompt += " " + conditioner_params.text;
-            prompt_attn_range.second = static_cast<int>(prompt.size());
-
-            prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
-        } else {
-            prompt_template_encode_start_idx = 34;
-
-            prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
-
-            prompt_attn_range.first = static_cast<int>(prompt.size());
-            prompt += conditioner_params.text;
-            prompt_attn_range.second = static_cast<int>(prompt.size());
-
-            prompt += "<|im_end|>\n<|im_start|>assistant\n";
+            mask.insert(mask.end(), max_length - tokens.size(), 0.f);
+            tokenizer->pad_tokens(tokens, weights, max_length, true);
         }
 
-        if (tokens.empty()) {
-            auto tokens_and_weights = tokenize(prompt, prompt_attn_range, max_length, max_length > 0);
-            tokens                  = std::get<0>(tokens_and_weights);
-            weights                 = std::get<1>(tokens_and_weights);
-        }
-
-        int64_t t0                        = ggml_time_ms();
-        struct ggml_tensor* hidden_states = nullptr;  // [N, n_token, 3584]
+        ggml_tensor* hidden_states = nullptr;  // [N, n_token, hidden_size]
 
         auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
 
@@ -1930,11 +1947,6 @@ struct LLMEmbedder : public Conditioner {
 
         GGML_ASSERT(hidden_states->ne[1] > prompt_template_encode_start_idx);
 
-        int64_t min_length = 0;
-        if (version == VERSION_FLUX2) {
-            min_length = 512;
-        }
-
         int64_t zero_pad_len = 0;
         if (min_length > 0) {
             if (hidden_states->ne[1] - prompt_template_encode_start_idx < min_length) {
@@ -1956,11 +1968,186 @@ struct LLMEmbedder : public Conditioner {
             ggml_ext_tensor_set_f32(new_hidden_states, value, i0, i1, i2, i3);
         });
 
-        // print_ggml_tensor(new_hidden_states);
+        return new_hidden_states;
+    }
+
+    SDCondition get_learned_condition(ggml_context* work_ctx,
+                                      int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        std::string prompt;
+        std::pair<int, int> prompt_attn_range;
+        std::vector<std::string> extra_prompts;
+        std::vector<std::pair<int, int>> extra_prompts_attn_range;
+        std::vector<std::pair<int, ggml_tensor*>> image_embeds;
+        int prompt_template_encode_start_idx = 34;
+        int max_length                       = 0;  // pad tokens
+        int min_length                       = 0;  // zero pad hidden_states
+        std::set<int> out_layers;
+
+        int64_t t0 = ggml_time_ms();
+
+        if (sd_version_is_qwen_image(version)) {
+            if (llm->enable_vision && !conditioner_params.ref_images.empty()) {
+                LOG_INFO("QwenImageEditPlusPipeline");
+                prompt_template_encode_start_idx = 64;
+                int image_embed_idx              = 64 + 6;
+
+                int min_pixels          = 384 * 384;
+                int max_pixels          = 560 * 560;
+                std::string placeholder = "<|image_pad|>";
+                std::string img_prompt;
+
+                for (int i = 0; i < conditioner_params.ref_images.size(); i++) {
+                    sd_image_f32_t image = sd_image_t_to_sd_image_f32_t(*conditioner_params.ref_images[i]);
+                    double factor        = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size;
+                    int height           = image.height;
+                    int width            = image.width;
+                    int h_bar            = static_cast<int>(std::round(height / factor) * factor);
+                    int w_bar            = static_cast<int>(std::round(width / factor) * factor);
+
+                    if (static_cast<double>(h_bar) * w_bar > max_pixels) {
+                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
+                        h_bar       = std::max(static_cast<int>(factor),
+                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
+                        w_bar       = std::max(static_cast<int>(factor),
+                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
+                    } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
+                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
+                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+                    }
+
+                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, image.height, image.width, h_bar, w_bar);
+
+                    sd_image_f32_t resized_image = clip_preprocess(image, w_bar, h_bar);
+                    free(image.data);
+                    image.data = nullptr;
+
+                    ggml_tensor* image_tensor = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, resized_image.width, resized_image.height, 3, 1);
+                    sd_image_f32_to_ggml_tensor(resized_image, image_tensor, false);
+                    free(resized_image.data);
+                    resized_image.data = nullptr;
+
+                    ggml_tensor* image_embed = nullptr;
+                    llm->encode_image(n_threads, image_tensor, &image_embed, work_ctx);
+                    image_embeds.emplace_back(image_embed_idx, image_embed);
+                    image_embed_idx += 1 + static_cast<int>(image_embed->ne[1]) + 6;
+
+                    img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";  // [24669, 220, index, 25, 220, 151652]
+                    int64_t num_image_tokens = image_embed->ne[1];
+                    img_prompt.reserve(num_image_tokens * placeholder.size());
+                    for (int j = 0; j < num_image_tokens; j++) {
+                        img_prompt += placeholder;
+                    }
+                    img_prompt += "<|vision_end|>";
+                }
+
+                prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
+                prompt += img_prompt;
+
+                prompt_attn_range.first = static_cast<int>(prompt.size());
+                prompt += conditioner_params.text;
+                prompt_attn_range.second = static_cast<int>(prompt.size());
+
+                prompt += "<|im_end|>\n<|im_start|>assistant\n";
+            } else {
+                prompt_template_encode_start_idx = 34;
+
+                prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
+
+                prompt_attn_range.first = static_cast<int>(prompt.size());
+                prompt += conditioner_params.text;
+                prompt_attn_range.second = static_cast<int>(prompt.size());
+
+                prompt += "<|im_end|>\n<|im_start|>assistant\n";
+            }
+        } else if (version == VERSION_FLUX2) {
+            prompt_template_encode_start_idx = 0;
+            min_length                       = 512;
+            out_layers                       = {10, 20, 30};
+
+            prompt = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "[/INST]";
+        } else if (sd_version_is_z_image(version)) {
+            prompt_template_encode_start_idx = 0;
+            out_layers                       = {35};  // -2
+
+            if (!conditioner_params.ref_images.empty()) {
+                LOG_INFO("ZImageOmniPipeline");
+                prompt = "<|im_start|>user\n<|vision_start|>";
+                for (int i = 0; i < conditioner_params.ref_images.size() - 1; i++) {
+                    extra_prompts.push_back("<|vision_end|><|vision_start|>");
+                }
+                extra_prompts.push_back("<|vision_end|>" + conditioner_params.text + "<|im_end|>\n<|im_start|>assistant\n<|vision_start|>");
+                extra_prompts.push_back("<|vision_end|><|im_end|>");
+            } else {
+                prompt = "<|im_start|>user\n";
+
+                prompt_attn_range.first = static_cast<int>(prompt.size());
+                prompt += conditioner_params.text;
+                prompt_attn_range.second = static_cast<int>(prompt.size());
+
+                prompt += "<|im_end|>\n<|im_start|>assistant\n";
+            }
+        } else if (version == VERSION_FLUX2_KLEIN) {
+            prompt_template_encode_start_idx = 0;
+            max_length                       = 512;
+            out_layers                       = {9, 18, 27};
+
+            prompt = "<|im_start|>user\n";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
+        } else if (version == VERSION_OVIS_IMAGE) {
+            prompt_template_encode_start_idx = 28;
+            max_length                       = prompt_template_encode_start_idx + 256;
+
+            prompt = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += " " + conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
+        } else {
+            GGML_ABORT("unknown version %d", version);
+        }
+
+        auto hidden_states = encode_prompt(work_ctx,
+                                           n_threads,
+                                           prompt,
+                                           prompt_attn_range,
+                                           max_length,
+                                           min_length,
+                                           image_embeds,
+                                           out_layers,
+                                           prompt_template_encode_start_idx);
+
+        std::vector<ggml_tensor*> extra_hidden_states_vec;
+        for (int i = 0; i < extra_prompts.size(); i++) {
+            auto extra_hidden_states = encode_prompt(work_ctx,
+                                                     n_threads,
+                                                     extra_prompts[i],
+                                                     extra_prompts_attn_range[i],
+                                                     max_length,
+                                                     min_length,
+                                                     image_embeds,
+                                                     out_layers,
+                                                     prompt_template_encode_start_idx);
+            extra_hidden_states_vec.push_back(extra_hidden_states);
+        }
 
         int64_t t1 = ggml_time_ms();
         LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
-        return {new_hidden_states, nullptr, nullptr};
+        return {hidden_states, nullptr, nullptr, extra_hidden_states_vec};
     }
 };
 

src/control.hpp

@@ -1,8 +1,7 @@
 #ifndef __CONTROL_HPP__
 #define __CONTROL_HPP__
 
-#include "common.hpp"
-#include "ggml_extend.hpp"
+#include "common_block.hpp"
 #include "model.h"
 
 #define CONTROL_NET_GRAPH_SIZE 1536
@@ -165,26 +164,26 @@ public:
         blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
     }
 
-    struct ggml_tensor* resblock_forward(std::string name,
-                                         GGMLRunnerContext* ctx,
-                                         struct ggml_tensor* x,
-                                         struct ggml_tensor* emb) {
+    ggml_tensor* resblock_forward(std::string name,
+                                  GGMLRunnerContext* ctx,
+                                  ggml_tensor* x,
+                                  ggml_tensor* emb) {
         auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
         return block->forward(ctx, x, emb);
     }
 
-    struct ggml_tensor* attention_layer_forward(std::string name,
-                                                GGMLRunnerContext* ctx,
-                                                struct ggml_tensor* x,
-                                                struct ggml_tensor* context) {
+    ggml_tensor* attention_layer_forward(std::string name,
+                                         GGMLRunnerContext* ctx,
+                                         ggml_tensor* x,
+                                         ggml_tensor* context) {
         auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
         return block->forward(ctx, x, context);
     }
 
-    struct ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
-                                                 struct ggml_tensor* hint,
-                                                 struct ggml_tensor* emb,
-                                                 struct ggml_tensor* context) {
+    ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
+                                          ggml_tensor* hint,
+                                          ggml_tensor* emb,
+                                          ggml_tensor* context) {
         int num_input_blocks = 15;
         auto h               = hint;
         for (int i = 0; i < num_input_blocks; i++) {
@@ -199,13 +198,13 @@ public:
         return h;
     }
 
-    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) {
+    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) {
         // 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]
@@ -247,7 +246,7 @@ public:
             emb = ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
         }
 
-        std::vector<struct ggml_tensor*> outs;
+        std::vector<ggml_tensor*> outs;
 
         if (guided_hint == nullptr) {
             guided_hint = input_hint_block_forward(ctx, hint, emb, context);
@@ -313,9 +312,9 @@ struct ControlNet : public GGMLRunner {
 
     ggml_backend_buffer_t control_buffer = nullptr;  // keep control output tensors in backend memory
     ggml_context* control_ctx            = nullptr;
-    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;
+    std::vector<ggml_tensor*> controls;  // (12 input block outputs, 1 middle block output) SD 1.5
+    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,
@@ -329,8 +328,8 @@ struct ControlNet : public GGMLRunner {
         free_control_ctx();
     }
 
-    void alloc_control_ctx(std::vector<struct ggml_tensor*> outs) {
-        struct ggml_init_params params;
+    void alloc_control_ctx(std::vector<ggml_tensor*> outs) {
+        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;
@@ -371,16 +370,16 @@ struct ControlNet : public GGMLRunner {
         return "control_net";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         control_net.get_param_tensors(tensors, prefix);
     }
 
-    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_cgraph* build_graph(ggml_tensor* x,
+                             ggml_tensor* hint,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* context,
+                             ggml_tensor* y = nullptr) {
+        ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
 
         x = to_backend(x);
         if (guided_hint_cached) {
@@ -415,18 +414,18 @@ struct ControlNet : public GGMLRunner {
     }
 
     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) {
+                 ggml_tensor* x,
+                 ggml_tensor* hint,
+                 ggml_tensor* timesteps,
+                 ggml_tensor* context,
+                 ggml_tensor* y,
+                 ggml_tensor** output     = nullptr,
+                 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 = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x, hint, timesteps, context, y);
         };
 

src/denoiser.hpp

@@ -657,9 +657,8 @@ struct DiscreteFlowDenoiser : public Denoiser {
 
     float sigma_data = 1.0f;
 
-    DiscreteFlowDenoiser(float shift = 3.0f)
-        : shift(shift) {
-        set_parameters();
+    DiscreteFlowDenoiser(float shift = 3.0f) {
+        set_shift(shift);
     }
 
     void set_parameters() {
@@ -668,6 +667,11 @@ struct DiscreteFlowDenoiser : public Denoiser {
         }
     }
 
+    void set_shift(float shift) {
+        this->shift = shift;
+        set_parameters();
+    }
+
     float sigma_min() override {
         return sigmas[0];
     }
@@ -710,34 +714,8 @@ float flux_time_shift(float mu, float sigma, float t) {
     return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
 }
 
-struct FluxFlowDenoiser : public Denoiser {
-    float sigmas[TIMESTEPS];
-    float shift = 1.15f;
-
-    float sigma_data = 1.0f;
-
-    FluxFlowDenoiser(float shift = 1.15f) {
-        set_parameters(shift);
-    }
-
-    void set_shift(float shift) {
-        this->shift = shift;
-    }
-
-    void set_parameters(float shift) {
-        set_shift(shift);
-        for (int i = 0; i < TIMESTEPS; i++) {
-            sigmas[i] = t_to_sigma(static_cast<float>(i));
-        }
-    }
-
-    float sigma_min() override {
-        return sigmas[0];
-    }
-
-    float sigma_max() override {
-        return sigmas[TIMESTEPS - 1];
-    }
+struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
+    FluxFlowDenoiser() = default;
 
     float sigma_to_t(float sigma) override {
         return sigma;
@@ -747,26 +725,6 @@ struct FluxFlowDenoiser : public Denoiser {
         t = t + 1;
         return flux_time_shift(shift, 1.0f, t / TIMESTEPS);
     }
-
-    std::vector<float> get_scalings(float sigma) override {
-        float c_skip = 1.0f;
-        float c_out  = -sigma;
-        float c_in   = 1.0f;
-        return {c_skip, c_out, c_in};
-    }
-
-    // this function will modify noise/latent
-    ggml_tensor* noise_scaling(float sigma, ggml_tensor* noise, ggml_tensor* latent) override {
-        ggml_ext_tensor_scale_inplace(noise, sigma);
-        ggml_ext_tensor_scale_inplace(latent, 1.0f - sigma);
-        ggml_ext_tensor_add_inplace(latent, noise);
-        return latent;
-    }
-
-    ggml_tensor* inverse_noise_scaling(float sigma, ggml_tensor* latent) override {
-        ggml_ext_tensor_scale_inplace(latent, 1.0f / (1.0f - sigma));
-        return latent;
-    }
 };
 
 struct Flux2FlowDenoiser : public FluxFlowDenoiser {
@@ -815,8 +773,8 @@ static bool sample_k_diffusion(sample_method_t method,
     // sample_euler_ancestral
     switch (method) {
         case EULER_A_SAMPLE_METHOD: {
-            struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* d     = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* d     = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
                 float sigma = sigmas[i];
@@ -872,7 +830,7 @@ static bool sample_k_diffusion(sample_method_t method,
         } break;
         case EULER_SAMPLE_METHOD:  // Implemented without any sigma churn
         {
-            struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
                 float sigma = sigmas[i];
@@ -907,8 +865,8 @@ static bool sample_k_diffusion(sample_method_t method,
             }
         } break;
         case HEUN_SAMPLE_METHOD: {
-            struct ggml_tensor* d  = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* d  = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
                 // denoise
@@ -963,8 +921,8 @@ static bool sample_k_diffusion(sample_method_t method,
             }
         } break;
         case DPM2_SAMPLE_METHOD: {
-            struct ggml_tensor* d  = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* d  = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
                 // denoise
@@ -1021,8 +979,8 @@ static bool sample_k_diffusion(sample_method_t method,
 
         } break;
         case DPMPP2S_A_SAMPLE_METHOD: {
-            struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* x2    = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* x2    = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
                 // denoise
@@ -1092,7 +1050,7 @@ static bool sample_k_diffusion(sample_method_t method,
         } break;
         case DPMPP2M_SAMPLE_METHOD:  // DPM++ (2M) from Karras et al (2022)
         {
-            struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
 
             auto t_fn = [](float sigma) -> float { return -log(sigma); };
 
@@ -1134,7 +1092,7 @@ static bool sample_k_diffusion(sample_method_t method,
         } break;
         case DPMPP2Mv2_SAMPLE_METHOD:  // Modified DPM++ (2M) from https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457
         {
-            struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
 
             auto t_fn = [](float sigma) -> float { return -log(sigma); };
 
@@ -1199,8 +1157,8 @@ static bool sample_k_diffusion(sample_method_t method,
                 }
                 float* vec_denoised = (float*)denoised->data;
                 // d_cur = (x_cur - denoised) / sigma
-                struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x_cur);
-                float* vec_d_cur          = (float*)d_cur->data;
+                ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x_cur);
+                float* vec_d_cur   = (float*)d_cur->data;
 
                 for (int j = 0; j < ggml_nelements(d_cur); j++) {
                     vec_d_cur[j] = (vec_x_cur[j] - vec_denoised[j]) / sigma;
@@ -1267,11 +1225,11 @@ static bool sample_k_diffusion(sample_method_t method,
                 float t_next = sigmas[i + 1];
 
                 // Denoising step
-                ggml_tensor* denoised     = model(x, sigma, i + 1);
-                float* vec_denoised       = (float*)denoised->data;
-                struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x);
-                float* vec_d_cur          = (float*)d_cur->data;
-                float* vec_x              = (float*)x->data;
+                ggml_tensor* denoised = model(x, sigma, i + 1);
+                float* vec_denoised   = (float*)denoised->data;
+                ggml_tensor* d_cur    = ggml_dup_tensor(work_ctx, x);
+                float* vec_d_cur      = (float*)d_cur->data;
+                float* vec_x          = (float*)x->data;
 
                 // d_cur = (x - denoised) / sigma
                 for (int j = 0; j < ggml_nelements(d_cur); j++) {
@@ -1332,8 +1290,8 @@ static bool sample_k_diffusion(sample_method_t method,
         } break;
         case LCM_SAMPLE_METHOD:  // Latent Consistency Models
         {
-            struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* d     = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* d     = ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
                 float sigma = sigmas[i];
@@ -1400,9 +1358,9 @@ static bool sample_k_diffusion(sample_method_t method,
                               alphas_cumprod[i]);
             }
 
-            struct ggml_tensor* pred_original_sample =
+            ggml_tensor* pred_original_sample =
                 ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* variance_noise =
+            ggml_tensor* variance_noise =
                 ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
@@ -1464,7 +1422,7 @@ static bool sample_k_diffusion(sample_method_t method,
                 // model_output = model() is the D(x, sigma) as
                 // defined in Karras et al. (2022), p. 3, Table 1 and
                 // p. 8 (7), compare also p. 38 (226) therein.
-                struct ggml_tensor* model_output =
+                ggml_tensor* model_output =
                     model(x, sigma, i + 1);
                 // Here model_output is still the k-diffusion denoiser
                 // output, not the U-net output F_theta(c_in(sigma) x;
@@ -1587,9 +1545,9 @@ static bool sample_k_diffusion(sample_method_t method,
             }
             int original_steps = 50;
 
-            struct ggml_tensor* pred_original_sample =
+            ggml_tensor* pred_original_sample =
                 ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* noise =
+            ggml_tensor* noise =
                 ggml_dup_tensor(work_ctx, x);
 
             for (int i = 0; i < steps; i++) {
@@ -1623,7 +1581,7 @@ static bool sample_k_diffusion(sample_method_t method,
                         vec_x[j] *= std::sqrt(sigma * sigma + 1);
                     }
                 }
-                struct ggml_tensor* model_output =
+                ggml_tensor* model_output =
                     model(x, sigma, i + 1);
                 {
                     float* vec_x = (float*)x->data;
@@ -1731,8 +1689,8 @@ static bool sample_k_diffusion(sample_method_t method,
         } break;
         case RES_MULTISTEP_SAMPLE_METHOD:  // Res Multistep sampler
         {
-            struct ggml_tensor* noise        = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* noise        = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
 
             bool have_old_sigma  = false;
             float old_sigma_down = 0.0f;
@@ -1839,9 +1797,9 @@ static bool sample_k_diffusion(sample_method_t method,
         } break;
         case RES_2S_SAMPLE_METHOD:  // Res 2s sampler
         {
-            struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* x0    = ggml_dup_tensor(work_ctx, x);
-            struct ggml_tensor* x2    = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* x0    = ggml_dup_tensor(work_ctx, x);
+            ggml_tensor* x2    = ggml_dup_tensor(work_ctx, x);
 
             const float c2 = 0.5f;
             auto t_fn      = [](float sigma) -> float { return -logf(sigma); };

src/diffusion_model.hpp

@@ -1,6 +1,7 @@
 #ifndef __DIFFUSION_MODEL_H__
 #define __DIFFUSION_MODEL_H__
 
+#include "anima.hpp"
 #include "flux.hpp"
 #include "mmdit.hpp"
 #include "qwen_image.hpp"
@@ -9,33 +10,33 @@
 #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              = {};
+    ggml_tensor* x                        = nullptr;
+    ggml_tensor* timesteps                = nullptr;
+    ggml_tensor* context                  = nullptr;
+    ggml_tensor* c_concat                 = nullptr;
+    ggml_tensor* y                        = nullptr;
+    ggml_tensor* guidance                 = nullptr;
+    std::vector<ggml_tensor*> ref_latents = {};
+    bool increase_ref_index               = false;
+    int num_video_frames                  = -1;
+    std::vector<ggml_tensor*> controls    = {};
+    float control_strength                = 0.f;
+    ggml_tensor* vace_context             = nullptr;
+    float vace_strength                   = 1.f;
+    std::vector<int> skip_layers          = {};
 };
 
 struct DiffusionModel {
-    virtual std::string get_desc()                                                      = 0;
+    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;
+                         ggml_tensor** output     = nullptr,
+                         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, 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;
@@ -68,7 +69,7 @@ struct UNetModel : public DiffusionModel {
         unet.free_compute_buffer();
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         unet.get_param_tensors(tensors, "model.diffusion_model");
     }
 
@@ -94,8 +95,8 @@ struct UNetModel : public DiffusionModel {
 
     bool compute(int n_threads,
                  DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
         return unet.compute(n_threads,
                             diffusion_params.x,
                             diffusion_params.timesteps,
@@ -133,7 +134,7 @@ struct MMDiTModel : public DiffusionModel {
         mmdit.free_compute_buffer();
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         mmdit.get_param_tensors(tensors, "model.diffusion_model");
     }
 
@@ -159,8 +160,8 @@ struct MMDiTModel : public DiffusionModel {
 
     bool compute(int n_threads,
                  DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
         return mmdit.compute(n_threads,
                              diffusion_params.x,
                              diffusion_params.timesteps,
@@ -199,7 +200,7 @@ struct FluxModel : public DiffusionModel {
         flux.free_compute_buffer();
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         flux.get_param_tensors(tensors, "model.diffusion_model");
     }
 
@@ -225,8 +226,8 @@ struct FluxModel : public DiffusionModel {
 
     bool compute(int n_threads,
                  DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
         return flux.compute(n_threads,
                             diffusion_params.x,
                             diffusion_params.timesteps,
@@ -242,6 +243,72 @@ struct FluxModel : public DiffusionModel {
     }
 };
 
+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);
+    }
+
+    bool compute(int n_threads,
+                 DiffusionParams diffusion_params,
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
+        return anima.compute(n_threads,
+                             diffusion_params.x,
+                             diffusion_params.timesteps,
+                             diffusion_params.context,
+                             diffusion_params.c_concat,
+                             diffusion_params.y,
+                             output,
+                             output_ctx);
+    }
+};
+
 struct WanModel : public DiffusionModel {
     std::string prefix;
     WAN::WanRunner wan;
@@ -270,7 +337,7 @@ struct WanModel : public DiffusionModel {
         wan.free_compute_buffer();
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         wan.get_param_tensors(tensors, prefix);
     }
 
@@ -296,8 +363,8 @@ struct WanModel : public DiffusionModel {
 
     bool compute(int n_threads,
                  DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
         return wan.compute(n_threads,
                            diffusion_params.x,
                            diffusion_params.timesteps,
@@ -341,7 +408,7 @@ struct QwenImageModel : public DiffusionModel {
         qwen_image.free_compute_buffer();
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         qwen_image.get_param_tensors(tensors, prefix);
     }
 
@@ -367,8 +434,8 @@ struct QwenImageModel : public DiffusionModel {
 
     bool compute(int n_threads,
                  DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
         return qwen_image.compute(n_threads,
                                   diffusion_params.x,
                                   diffusion_params.timesteps,
@@ -408,7 +475,7 @@ struct ZImageModel : public DiffusionModel {
         z_image.free_compute_buffer();
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         z_image.get_param_tensors(tensors, prefix);
     }
 
@@ -434,8 +501,8 @@ struct ZImageModel : public DiffusionModel {
 
     bool compute(int n_threads,
                  DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+                 ggml_tensor** output     = nullptr,
+                 ggml_context* output_ctx = nullptr) override {
         return z_image.compute(n_threads,
                                diffusion_params.x,
                                diffusion_params.timesteps,

src/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}));
     }
 
-    struct ggml_tensor* lrelu(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
         return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [n, num_feat, h, w]
         // return: [n, num_feat, h, w]
 
@@ -64,7 +64,7 @@ public:
         blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [n, num_feat, h, w]
         // return: [n, num_feat, h, w]
 
@@ -112,11 +112,11 @@ public:
     int get_scale() { return scale; }
     int get_num_block() { return num_block; }
 
-    struct ggml_tensor* lrelu(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
         return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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,24 @@ struct ESRGAN : public GGMLRunner {
         return success;
     }
 
-    struct ggml_cgraph* build_graph(struct ggml_tensor* x) {
+    ggml_cgraph* build_graph(ggml_tensor* x) {
         if (!rrdb_net)
             return nullptr;
         constexpr int kGraphNodes = 1 << 16;  // 65k
-        struct ggml_cgraph* gf    = new_graph_custom(kGraphNodes);
+        ggml_cgraph* gf           = new_graph_custom(kGraphNodes);
         x                         = to_backend(x);
 
-        auto runner_ctx         = get_context();
-        struct ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
+        auto runner_ctx  = get_context();
+        ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
         ggml_build_forward_expand(gf, out);
         return gf;
     }
 
     bool compute(const int n_threads,
-                 struct ggml_tensor* x,
+                 ggml_tensor* x,
                  ggml_tensor** output,
                  ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x);
         };
         return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);

src/flux.hpp

@@ -4,7 +4,7 @@
 #include <memory>
 #include <vector>
 
-#include "ggml_extend.hpp"
+#include "common_dit.hpp"
 #include "model.h"
 #include "rope.hpp"
 
@@ -19,7 +19,7 @@ namespace Flux {
             blocks["out_layer"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_dim, hidden_dim, bias));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             // x: [..., in_dim]
             // return: [..., hidden_dim]
             auto in_layer  = std::dynamic_pointer_cast<Linear>(blocks["in_layer"]);
@@ -37,7 +37,7 @@ namespace Flux {
         int64_t hidden_size;
         float eps;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             ggml_type wtype = GGML_TYPE_F32;
             params["scale"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
         }
@@ -48,10 +48,10 @@ namespace Flux {
             : hidden_size(hidden_size),
               eps(eps) {}
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
-            struct ggml_tensor* w = params["scale"];
-            x                     = ggml_rms_norm(ctx->ggml_ctx, x, eps);
-            x                     = ggml_mul(ctx->ggml_ctx, x, w);
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            ggml_tensor* w = params["scale"];
+            x              = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+            x              = ggml_mul(ctx->ggml_ctx, x, w);
             return x;
         }
     };
@@ -63,7 +63,7 @@ namespace Flux {
             blocks["key_norm"]   = std::shared_ptr<GGMLBlock>(new RMSNorm(dim));
         }
 
-        struct ggml_tensor* query_norm(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* query_norm(GGMLRunnerContext* ctx, ggml_tensor* x) {
             // x: [..., dim]
             // return: [..., dim]
             auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["query_norm"]);
@@ -72,7 +72,7 @@ namespace Flux {
             return x;
         }
 
-        struct ggml_tensor* key_norm(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* key_norm(GGMLRunnerContext* ctx, ggml_tensor* x) {
             // x: [..., dim]
             // return: [..., dim]
             auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["key_norm"]);
@@ -98,32 +98,34 @@ namespace Flux {
             blocks["proj"]   = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, proj_bias));
         }
 
-        std::vector<struct ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
             auto norm     = std::dynamic_pointer_cast<QKNorm>(blocks["norm"]);
 
             auto qkv         = qkv_proj->forward(ctx, x);
-            auto qkv_vec     = ggml_ext_chunk(ctx->ggml_ctx, qkv, 3, 0, true);
-            int64_t head_dim = qkv_vec[0]->ne[0] / num_heads;
-            auto q           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]);
-            auto k           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]);
-            auto v           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, qkv_vec[2]->ne[1], qkv_vec[2]->ne[2]);
+            int64_t head_dim = qkv->ne[0] / 3 / num_heads;
+            auto q           = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
+                                            qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
+            auto k           = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
+                                            qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], (qkv->nb[0]) * qkv->ne[0] / 3);
+            auto v           = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
+                                            qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], (qkv->nb[0]) * 2 * qkv->ne[0] / 3);
             q                = norm->query_norm(ctx, q);
             k                = norm->key_norm(ctx, k);
             return {q, k, v};
         }
 
-        struct ggml_tensor* post_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* post_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
             x = proj->forward(ctx, x);  // [N, n_token, dim]
             return x;
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mask) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* mask) {
             // x: [N, n_token, dim]
             // pe: [n_token, d_head/2, 2, 2]
             // return [N, n_token, dim]
@@ -145,7 +147,7 @@ namespace Flux {
             blocks["2"]             = std::make_shared<Linear>(intermediate_size, hidden_size, bias);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
 
@@ -168,7 +170,7 @@ namespace Flux {
             blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, bias);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
             auto up_proj   = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
             auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
@@ -210,7 +212,7 @@ namespace Flux {
             blocks["lin"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * multiplier, bias));
         }
 
-        std::vector<ModulationOut> forward(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
+        std::vector<ModulationOut> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             // x: [N, dim]
             // return: [ModulationOut, ModulationOut]
             auto lin = std::dynamic_pointer_cast<Linear>(blocks["lin"]);
@@ -230,11 +232,11 @@ namespace Flux {
         }
     };
 
-    __STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
-                                                   struct ggml_tensor* x,
-                                                   struct ggml_tensor* shift,
-                                                   struct ggml_tensor* scale,
-                                                   bool skip_reshape = false) {
+    __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+                                            ggml_tensor* x,
+                                            ggml_tensor* shift,
+                                            ggml_tensor* scale,
+                                            bool skip_reshape = false) {
         // x: [N, L, C]
         // scale: [N, C]
         // shift: [N, C]
@@ -292,7 +294,7 @@ namespace Flux {
             }
         }
 
-        std::vector<ModulationOut> get_distil_img_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
+        std::vector<ModulationOut> get_distil_img_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             // TODO: not hardcoded?
             const int single_blocks_count = 38;
             const int double_blocks_count = 19;
@@ -301,7 +303,7 @@ namespace Flux {
             return {ModulationOut(ctx, vec, offset), ModulationOut(ctx, vec, offset + 3)};
         }
 
-        std::vector<ModulationOut> get_distil_txt_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
+        std::vector<ModulationOut> get_distil_txt_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             // TODO: not hardcoded?
             const int single_blocks_count = 38;
             const int double_blocks_count = 19;
@@ -310,14 +312,14 @@ namespace Flux {
             return {ModulationOut(ctx, vec, offset), ModulationOut(ctx, vec, offset + 3)};
         }
 
-        std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                                    struct ggml_tensor* img,
-                                                                    struct ggml_tensor* txt,
-                                                                    struct ggml_tensor* vec,
-                                                                    struct ggml_tensor* pe,
-                                                                    struct ggml_tensor* mask            = nullptr,
-                                                                    std::vector<ModulationOut> img_mods = {},
-                                                                    std::vector<ModulationOut> txt_mods = {}) {
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* img,
+                                                      ggml_tensor* txt,
+                                                      ggml_tensor* vec,
+                                                      ggml_tensor* pe,
+                                                      ggml_tensor* mask                   = nullptr,
+                                                      std::vector<ModulationOut> img_mods = {},
+                                                      std::vector<ModulationOut> txt_mods = {}) {
             // 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]
@@ -455,17 +457,17 @@ namespace Flux {
             }
         }
 
-        ModulationOut get_distil_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
+        ModulationOut get_distil_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             int64_t offset = 3 * idx;
             return ModulationOut(ctx, vec, offset);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* vec,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mask        = nullptr,
-                                    std::vector<ModulationOut> mods = {}) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* vec,
+                             ggml_tensor* pe,
+                             ggml_tensor* mask               = nullptr,
+                             std::vector<ModulationOut> mods = {}) {
             // x: [N, n_token, hidden_size]
             // pe: [n_token, d_head/2, 2, 2]
             // return: [N, n_token, hidden_size]
@@ -491,15 +493,14 @@ namespace Flux {
             auto x_mod   = Flux::modulate(ctx->ggml_ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale);
             auto qkv_mlp = linear1->forward(ctx, x_mod);  // [N, n_token, hidden_size * 3 + mlp_hidden_dim*mlp_mult_factor]
 
-            auto q = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, hidden_size, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], 0);
-            auto k = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, hidden_size, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], hidden_size * qkv_mlp->nb[0]);
-            auto v = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, hidden_size, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], hidden_size * 2 * qkv_mlp->nb[0]);
-
             int64_t head_dim = hidden_size / num_heads;
 
-            q = ggml_reshape_4d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim, num_heads, q->ne[1], q->ne[2]);  // [N, n_token, n_head, d_head]
-            k = ggml_reshape_4d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim, num_heads, k->ne[1], k->ne[2]);  // [N, n_token, n_head, d_head]
-            v = ggml_reshape_4d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim, num_heads, v->ne[1], v->ne[2]);  // [N, n_token, n_head, d_head]
+            auto q = ggml_view_4d(ctx->ggml_ctx, qkv_mlp, head_dim, num_heads, qkv_mlp->ne[1], qkv_mlp->ne[2],
+                                  qkv_mlp->nb[0] * head_dim, qkv_mlp->nb[1], qkv_mlp->nb[2], 0);
+            auto k = ggml_view_4d(ctx->ggml_ctx, qkv_mlp, head_dim, num_heads, qkv_mlp->ne[1], qkv_mlp->ne[2],
+                                  qkv_mlp->nb[0] * head_dim, qkv_mlp->nb[1], qkv_mlp->nb[2], (qkv_mlp->nb[0]) * hidden_size);
+            auto v = ggml_view_4d(ctx->ggml_ctx, qkv_mlp, head_dim, num_heads, qkv_mlp->ne[1], qkv_mlp->ne[2],
+                                  qkv_mlp->nb[0] * head_dim, qkv_mlp->nb[1], qkv_mlp->nb[2], (qkv_mlp->nb[0]) * 2 * hidden_size);
 
             q         = norm->query_norm(ctx, q);
             k         = norm->key_norm(ctx, k);
@@ -538,7 +539,7 @@ namespace Flux {
             }
         }
 
-        ModulationOut get_distil_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
+        ModulationOut get_distil_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             int64_t offset = vec->ne[2] - 2;
             int64_t stride = vec->nb[1] * vec->ne[1];
             auto shift     = ggml_view_2d(ctx->ggml_ctx, vec, vec->ne[0], vec->ne[1], vec->nb[1], stride * (offset + 0));  // [N, dim]
@@ -547,15 +548,15 @@ namespace Flux {
             return {shift, scale, nullptr};
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* c) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* c) {
             // x: [N, n_token, hidden_size]
             // c: [N, hidden_size]
             // return: [N, n_token, patch_size * patch_size * out_channels]
             auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
             auto linear     = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
-            struct ggml_tensor *shift, *scale;
+            ggml_tensor *shift, *scale;
             if (prune_mod) {
                 auto mod = get_distil_mod(ctx, c);
                 shift    = mod.shift;
@@ -588,7 +589,7 @@ namespace Flux {
             blocks["out_proj"] = std::shared_ptr<GGMLBlock>(new Linear(inner_size, hidden_size, true));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto in_proj  = std::dynamic_pointer_cast<Linear>(blocks["in_proj"]);
             auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
 
@@ -611,9 +612,9 @@ namespace Flux {
             blocks["embedder.0"] = std::make_shared<Linear>(in_channels + max_freqs * max_freqs, hidden_size_input);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* dct) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* dct) {
             // x: (B, P^2, C)
             // dct: (1, P^2, max_freqs^2)
             // return: (B, P^2, hidden_size_input)
@@ -638,9 +639,9 @@ namespace Flux {
             blocks["norm"]            = std::make_shared<RMSNorm>(hidden_size_x);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* s) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* s) {
             // x: (batch_size, n_token, hidden_size_x)
             // s: (batch_size, hidden_size_s)
             // return: (batch_size, n_token, hidden_size_x)
@@ -688,8 +689,8 @@ namespace Flux {
             blocks["linear"] = std::make_shared<Linear>(hidden_size, out_channels);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x) {
             auto norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
 
@@ -707,8 +708,8 @@ namespace Flux {
             blocks["conv"] = std::make_shared<Conv2d>(hidden_size, out_channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x) {
             // x: [N, C, H, W]
             auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
@@ -846,79 +847,15 @@ namespace Flux {
             }
         }
 
-        struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* 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_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-            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(GGMLRunnerContext* ctx,
-                                        struct ggml_tensor* x) {
-            // img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
-            x = pad_to_patch_size(ctx, x);
-            x = patchify(ctx->ggml_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* img,
-                                         struct ggml_tensor* txt,
-                                         struct ggml_tensor* timesteps,
-                                         struct ggml_tensor* y,
-                                         struct ggml_tensor* guidance,
-                                         struct ggml_tensor* pe,
-                                         struct ggml_tensor* mod_index_arange = nullptr,
-                                         std::vector<int> skip_layers         = {}) {
+        ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
+                                  ggml_tensor* img,
+                                  ggml_tensor* txt,
+                                  ggml_tensor* timesteps,
+                                  ggml_tensor* y,
+                                  ggml_tensor* guidance,
+                                  ggml_tensor* pe,
+                                  ggml_tensor* mod_index_arange = nullptr,
+                                  std::vector<int> skip_layers  = {}) {
             auto img_in      = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
             auto txt_in      = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
             auto final_layer = std::dynamic_pointer_cast<LastLayer>(blocks["final_layer"]);
@@ -927,8 +864,8 @@ namespace Flux {
                 img = img_in->forward(ctx, img);
             }
 
-            struct ggml_tensor* vec;
-            struct ggml_tensor* txt_img_mask = nullptr;
+            ggml_tensor* vec;
+            ggml_tensor* txt_img_mask = nullptr;
             if (params.is_chroma) {
                 int64_t mod_index_length = 344;
                 auto approx              = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
@@ -1030,27 +967,27 @@ namespace Flux {
             return img;
         }
 
-        struct ggml_tensor* _apply_x0_residual(GGMLRunnerContext* ctx,
-                                               struct ggml_tensor* predicted,
-                                               struct ggml_tensor* noisy,
-                                               struct ggml_tensor* timesteps) {
+        ggml_tensor* _apply_x0_residual(GGMLRunnerContext* ctx,
+                                        ggml_tensor* predicted,
+                                        ggml_tensor* noisy,
+                                        ggml_tensor* timesteps) {
             auto x = ggml_sub(ctx->ggml_ctx, noisy, predicted);
             x      = ggml_div(ctx->ggml_ctx, x, timesteps);
             return x;
         }
 
-        struct ggml_tensor* forward_chroma_radiance(GGMLRunnerContext* ctx,
-                                                    struct ggml_tensor* x,
-                                                    struct ggml_tensor* timestep,
-                                                    struct ggml_tensor* context,
-                                                    struct ggml_tensor* c_concat,
-                                                    struct ggml_tensor* y,
-                                                    struct ggml_tensor* guidance,
-                                                    struct ggml_tensor* pe,
-                                                    struct ggml_tensor* mod_index_arange  = nullptr,
-                                                    struct ggml_tensor* dct               = nullptr,
-                                                    std::vector<ggml_tensor*> ref_latents = {},
-                                                    std::vector<int> skip_layers          = {}) {
+        ggml_tensor* forward_chroma_radiance(GGMLRunnerContext* ctx,
+                                             ggml_tensor* x,
+                                             ggml_tensor* timestep,
+                                             ggml_tensor* context,
+                                             ggml_tensor* c_concat,
+                                             ggml_tensor* y,
+                                             ggml_tensor* guidance,
+                                             ggml_tensor* pe,
+                                             ggml_tensor* mod_index_arange         = nullptr,
+                                             ggml_tensor* dct                      = nullptr,
+                                             std::vector<ggml_tensor*> ref_latents = {},
+                                             std::vector<int> skip_layers          = {}) {
             GGML_ASSERT(x->ne[3] == 1);
 
             int64_t W      = x->ne[0];
@@ -1060,7 +997,7 @@ namespace Flux {
             int pad_h      = (patch_size - H % patch_size) % patch_size;
             int pad_w      = (patch_size - W % patch_size) % patch_size;
 
-            auto img      = pad_to_patch_size(ctx, x);
+            auto img      = DiT::pad_to_patch_size(ctx, x, params.patch_size, params.patch_size);
             auto orig_img = img;
 
             if (params.chroma_radiance_params.fake_patch_size_x2) {
@@ -1082,7 +1019,7 @@ namespace Flux {
             auto nerf_image_embedder   = std::dynamic_pointer_cast<NerfEmbedder>(blocks["nerf_image_embedder"]);
             auto nerf_final_layer_conv = std::dynamic_pointer_cast<NerfFinalLayerConv>(blocks["nerf_final_layer_conv"]);
 
-            auto nerf_pixels    = patchify(ctx->ggml_ctx, orig_img);  // [N, num_patches, C * patch_size * patch_size]
+            auto nerf_pixels    = DiT::patchify(ctx->ggml_ctx, orig_img, patch_size, patch_size);  // [N, num_patches, C * patch_size * patch_size]
             int64_t num_patches = nerf_pixels->ne[1];
             nerf_pixels         = ggml_reshape_3d(ctx->ggml_ctx,
                                                   nerf_pixels,
@@ -1102,7 +1039,7 @@ namespace Flux {
 
             img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img_dct, 1, 0, 2, 3));                                 // [N*num_patches, nerf_hidden_size, patch_size*patch_size]
             img_dct = ggml_reshape_3d(ctx->ggml_ctx, img_dct, img_dct->ne[0] * img_dct->ne[1], num_patches, img_dct->ne[2] / num_patches);  // [N, num_patches, nerf_hidden_size*patch_size*patch_size]
-            img_dct = unpatchify(ctx->ggml_ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size);                               // [N, nerf_hidden_size, H, W]
+            img_dct = DiT::unpatchify(ctx->ggml_ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size);  // [N, nerf_hidden_size, H, W]
 
             out = nerf_final_layer_conv->forward(ctx, img_dct);  // [N, C, H, W]
 
@@ -1113,18 +1050,18 @@ namespace Flux {
             return out;
         }
 
-        struct ggml_tensor* forward_flux_chroma(GGMLRunnerContext* ctx,
-                                                struct ggml_tensor* x,
-                                                struct ggml_tensor* timestep,
-                                                struct ggml_tensor* context,
-                                                struct ggml_tensor* c_concat,
-                                                struct ggml_tensor* y,
-                                                struct ggml_tensor* guidance,
-                                                struct ggml_tensor* pe,
-                                                struct ggml_tensor* mod_index_arange  = nullptr,
-                                                struct ggml_tensor* dct               = nullptr,
-                                                std::vector<ggml_tensor*> ref_latents = {},
-                                                std::vector<int> skip_layers          = {}) {
+        ggml_tensor* forward_flux_chroma(GGMLRunnerContext* ctx,
+                                         ggml_tensor* x,
+                                         ggml_tensor* timestep,
+                                         ggml_tensor* context,
+                                         ggml_tensor* c_concat,
+                                         ggml_tensor* y,
+                                         ggml_tensor* guidance,
+                                         ggml_tensor* pe,
+                                         ggml_tensor* mod_index_arange         = nullptr,
+                                         ggml_tensor* dct                      = nullptr,
+                                         std::vector<ggml_tensor*> ref_latents = {},
+                                         std::vector<int> skip_layers          = {}) {
             GGML_ASSERT(x->ne[3] == 1);
 
             int64_t W      = x->ne[0];
@@ -1134,7 +1071,7 @@ namespace Flux {
             int pad_h      = (patch_size - H % patch_size) % patch_size;
             int pad_w      = (patch_size - W % patch_size) % patch_size;
 
-            auto img           = process_img(ctx, x);
+            auto img           = DiT::pad_and_patchify(ctx, x, patch_size, patch_size);
             int64_t img_tokens = img->ne[1];
 
             if (params.version == VERSION_FLUX_FILL) {
@@ -1142,8 +1079,8 @@ namespace Flux {
                 ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
                 ggml_tensor* mask   = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
 
-                masked = process_img(ctx, masked);
-                mask   = process_img(ctx, mask);
+                masked = DiT::pad_and_patchify(ctx, masked, patch_size, patch_size);
+                mask   = DiT::pad_and_patchify(ctx, mask, patch_size, patch_size);
 
                 img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, masked, mask, 0), 0);
             } else if (params.version == VERSION_FLEX_2) {
@@ -1152,21 +1089,21 @@ namespace Flux {
                 ggml_tensor* mask    = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
                 ggml_tensor* control = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * (C + 1));
 
-                masked  = process_img(ctx, masked);
-                mask    = process_img(ctx, mask);
-                control = process_img(ctx, control);
+                masked  = DiT::pad_and_patchify(ctx, masked, patch_size, patch_size);
+                mask    = DiT::pad_and_patchify(ctx, mask, patch_size, patch_size);
+                control = DiT::pad_and_patchify(ctx, control, patch_size, patch_size);
 
                 img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, masked, mask, 0), control, 0), 0);
             } else if (params.version == VERSION_FLUX_CONTROLS) {
                 GGML_ASSERT(c_concat != nullptr);
 
-                auto control = process_img(ctx, c_concat);
+                auto control = DiT::pad_and_patchify(ctx, c_concat, patch_size, patch_size);
                 img          = ggml_concat(ctx->ggml_ctx, img, control, 0);
             }
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = process_img(ctx, ref);
+                    ref = DiT::pad_and_patchify(ctx, ref, patch_size, patch_size);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
@@ -1178,23 +1115,22 @@ namespace Flux {
                 out = ggml_cont(ctx->ggml_ctx, out);
             }
 
-            // rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)
-            out = unpatchify(ctx->ggml_ctx, out, (H + pad_h) / patch_size, (W + pad_w) / patch_size);  // [N, C, H + pad_h, W + pad_w]
+            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size);  // [N, C, H, W]
             return out;
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* timestep,
-                                    struct ggml_tensor* context,
-                                    struct ggml_tensor* c_concat,
-                                    struct ggml_tensor* y,
-                                    struct ggml_tensor* guidance,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mod_index_arange  = nullptr,
-                                    struct ggml_tensor* dct               = nullptr,
-                                    std::vector<ggml_tensor*> ref_latents = {},
-                                    std::vector<int> skip_layers          = {}) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* c_concat,
+                             ggml_tensor* y,
+                             ggml_tensor* guidance,
+                             ggml_tensor* pe,
+                             ggml_tensor* mod_index_arange         = nullptr,
+                             ggml_tensor* dct                      = nullptr,
+                             std::vector<ggml_tensor*> ref_latents = {},
+                             std::vector<int> skip_layers          = {}) {
             // Forward pass of DiT.
             // x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
             // timestep: (N,) tensor of diffusion timesteps
@@ -1363,7 +1299,7 @@ namespace Flux {
             return "flux";
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
             flux.get_param_tensors(tensors, prefix);
         }
 
@@ -1417,20 +1353,20 @@ namespace Flux {
             return dct;
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
-                                        struct ggml_tensor* timesteps,
-                                        struct ggml_tensor* context,
-                                        struct ggml_tensor* c_concat,
-                                        struct ggml_tensor* y,
-                                        struct ggml_tensor* guidance,
-                                        std::vector<ggml_tensor*> ref_latents = {},
-                                        bool increase_ref_index               = false,
-                                        std::vector<int> skip_layers          = {}) {
+        ggml_cgraph* build_graph(ggml_tensor* x,
+                                 ggml_tensor* timesteps,
+                                 ggml_tensor* context,
+                                 ggml_tensor* c_concat,
+                                 ggml_tensor* y,
+                                 ggml_tensor* guidance,
+                                 std::vector<ggml_tensor*> ref_latents = {},
+                                 bool increase_ref_index               = false,
+                                 std::vector<int> skip_layers          = {}) {
             GGML_ASSERT(x->ne[3] == 1);
-            struct ggml_cgraph* gf = new_graph_custom(FLUX_GRAPH_SIZE);
+            ggml_cgraph* gf = new_graph_custom(FLUX_GRAPH_SIZE);
 
-            struct ggml_tensor* mod_index_arange = nullptr;
-            struct ggml_tensor* dct              = nullptr;  // for chroma radiance
+            ggml_tensor* mod_index_arange = nullptr;
+            ggml_tensor* dct              = nullptr;  // for chroma radiance
 
             x       = to_backend(x);
             context = to_backend(context);
@@ -1501,18 +1437,18 @@ namespace Flux {
 
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* out = flux.forward(&runner_ctx,
-                                                   x,
-                                                   timesteps,
-                                                   context,
-                                                   c_concat,
-                                                   y,
-                                                   guidance,
-                                                   pe,
-                                                   mod_index_arange,
-                                                   dct,
-                                                   ref_latents,
-                                                   skip_layers);
+            ggml_tensor* out = flux.forward(&runner_ctx,
+                                            x,
+                                            timesteps,
+                                            context,
+                                            c_concat,
+                                            y,
+                                            guidance,
+                                            pe,
+                                            mod_index_arange,
+                                            dct,
+                                            ref_latents,
+                                            skip_layers);
 
             ggml_build_forward_expand(gf, out);
 
@@ -1520,23 +1456,23 @@ namespace Flux {
         }
 
         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,
-                     struct ggml_tensor* guidance,
+                     ggml_tensor* x,
+                     ggml_tensor* timesteps,
+                     ggml_tensor* context,
+                     ggml_tensor* c_concat,
+                     ggml_tensor* y,
+                     ggml_tensor* guidance,
                      std::vector<ggml_tensor*> ref_latents = {},
                      bool increase_ref_index               = false,
-                     struct ggml_tensor** output           = nullptr,
-                     struct ggml_context* output_ctx       = nullptr,
+                     ggml_tensor** output                  = nullptr,
+                     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]
             // y: [N, adm_in_channels] or [1, adm_in_channels]
             // guidance: [N, ]
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
             };
 
@@ -1544,12 +1480,12 @@ namespace Flux {
         }
 
         void test() {
-            struct ggml_init_params params;
+            ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            struct ggml_context* work_ctx = ggml_init(params);
+            ggml_context* work_ctx = ggml_init(params);
             GGML_ASSERT(work_ctx != nullptr);
 
             {
@@ -1577,7 +1513,7 @@ namespace Flux {
                 auto y = nullptr;
                 // print_ggml_tensor(y);
 
-                struct ggml_tensor* out = nullptr;
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 compute(8, x, timesteps, context, nullptr, y, guidance, {}, false, &out, work_ctx);

src/latent-preview.h

@@ -1,234 +1,234 @@
-#include <cstddef>
-#include <cstdint>
-#include "ggml.h"
-
-const float wan_21_latent_rgb_proj[16][3] = {
-    {0.015123f, -0.148418f, 0.479828f},
-    {0.003652f, -0.010680f, -0.037142f},
-    {0.212264f, 0.063033f, 0.016779f},
-    {0.232999f, 0.406476f, 0.220125f},
-    {-0.051864f, -0.082384f, -0.069396f},
-    {0.085005f, -0.161492f, 0.010689f},
-    {-0.245369f, -0.506846f, -0.117010f},
-    {-0.151145f, 0.017721f, 0.007207f},
-    {-0.293239f, -0.207936f, -0.421135f},
-    {-0.187721f, 0.050783f, 0.177649f},
-    {-0.013067f, 0.265964f, 0.166578f},
-    {0.028327f, 0.109329f, 0.108642f},
-    {-0.205343f, 0.043991f, 0.148914f},
-    {0.014307f, -0.048647f, -0.007219f},
-    {0.217150f, 0.053074f, 0.319923f},
-    {0.155357f, 0.083156f, 0.064780f}};
-float wan_21_latent_rgb_bias[3] = {-0.270270f, -0.234976f, -0.456853f};
-
-const float wan_22_latent_rgb_proj[48][3] = {
-    {0.017126f, -0.027230f, -0.019257f},
-    {-0.113739f, -0.028715f, -0.022885f},
-    {-0.000106f, 0.021494f, 0.004629f},
-    {-0.013273f, -0.107137f, -0.033638f},
-    {-0.000381f, 0.000279f, 0.025877f},
-    {-0.014216f, -0.003975f, 0.040528f},
-    {0.001638f, -0.000748f, 0.011022f},
-    {0.029238f, -0.006697f, 0.035933f},
-    {0.021641f, -0.015874f, 0.040531f},
-    {-0.101984f, -0.070160f, -0.028855f},
-    {0.033207f, -0.021068f, 0.002663f},
-    {-0.104711f, 0.121673f, 0.102981f},
-    {0.082647f, -0.004991f, 0.057237f},
-    {-0.027375f, 0.031581f, 0.006868f},
-    {-0.045434f, 0.029444f, 0.019287f},
-    {-0.046572f, -0.012537f, 0.006675f},
-    {0.074709f, 0.033690f, 0.025289f},
-    {-0.008251f, -0.002745f, -0.006999f},
-    {0.012685f, -0.061856f, -0.048658f},
-    {0.042304f, -0.007039f, 0.000295f},
-    {-0.007644f, -0.060843f, -0.033142f},
-    {0.159909f, 0.045628f, 0.367541f},
-    {0.095171f, 0.086438f, 0.010271f},
-    {0.006812f, 0.019643f, 0.029637f},
-    {0.003467f, -0.010705f, 0.014252f},
-    {-0.099681f, -0.066272f, -0.006243f},
-    {0.047357f, 0.037040f, 0.000185f},
-    {-0.041797f, -0.089225f, -0.032257f},
-    {0.008928f, 0.017028f, 0.018684f},
-    {-0.042255f, 0.016045f, 0.006849f},
-    {0.011268f, 0.036462f, 0.037387f},
-    {0.011553f, -0.016375f, -0.048589f},
-    {0.046266f, -0.027189f, 0.056979f},
-    {0.009640f, -0.017576f, 0.030324f},
-    {-0.045794f, -0.036083f, -0.010616f},
-    {0.022418f, 0.039783f, -0.032939f},
-    {-0.052714f, -0.015525f, 0.007438f},
-    {0.193004f, 0.223541f, 0.264175f},
-    {-0.059406f, -0.008188f, 0.022867f},
-    {-0.156742f, -0.263791f, -0.007385f},
-    {-0.015717f, 0.016570f, 0.033969f},
-    {0.037969f, 0.109835f, 0.200449f},
-    {-0.000782f, -0.009566f, -0.008058f},
-    {0.010709f, 0.052960f, -0.044195f},
-    {0.017271f, 0.045839f, 0.034569f},
-    {0.009424f, 0.013088f, -0.001714f},
-    {-0.024805f, -0.059378f, -0.033756f},
-    {-0.078293f, 0.029070f, 0.026129f}};
-float wan_22_latent_rgb_bias[3] = {0.013160f, -0.096492f, -0.071323f};
-
-const float flux_latent_rgb_proj[16][3] = {
-    {-0.041168f, 0.019917f, 0.097253f},
-    {0.028096f, 0.026730f, 0.129576f},
-    {0.065618f, -0.067950f, -0.014651f},
-    {-0.012998f, -0.014762f, 0.081251f},
-    {0.078567f, 0.059296f, -0.024687f},
-    {-0.015987f, -0.003697f, 0.005012f},
-    {0.033605f, 0.138999f, 0.068517f},
-    {-0.024450f, -0.063567f, -0.030101f},
-    {-0.040194f, -0.016710f, 0.127185f},
-    {0.112681f, 0.088764f, -0.041940f},
-    {-0.023498f, 0.093664f, 0.025543f},
-    {0.082899f, 0.048320f, 0.007491f},
-    {0.075712f, 0.074139f, 0.081965f},
-    {-0.143501f, 0.018263f, -0.136138f},
-    {-0.025767f, -0.082035f, -0.040023f},
-    {-0.111849f, -0.055589f, -0.032361f}};
-float flux_latent_rgb_bias[3] = {0.024600f, -0.006937f, -0.008089f};
-
-const float flux2_latent_rgb_proj[32][3] = {
-    {0.000736f, -0.008385f, -0.019710f},
-    {-0.001352f, -0.016392f, 0.020693f},
-    {-0.006376f, 0.002428f, 0.036736f},
-    {0.039384f, 0.074167f, 0.119789f},
-    {0.007464f, -0.005705f, -0.004734f},
-    {-0.004086f, 0.005287f, -0.000409f},
-    {-0.032835f, 0.050802f, -0.028120f},
-    {-0.003158f, -0.000835f, 0.000406f},
-    {-0.112840f, -0.084337f, -0.023083f},
-    {0.001462f, -0.006656f, 0.000549f},
-    {-0.009980f, -0.007480f, 0.009702f},
-    {0.032540f, 0.000214f, -0.061388f},
-    {0.011023f, 0.000694f, 0.007143f},
-    {-0.001468f, -0.006723f, -0.001678f},
-    {-0.005921f, -0.010320f, -0.003907f},
-    {-0.028434f, 0.027584f, 0.018457f},
-    {0.014349f, 0.011523f, 0.000441f},
-    {0.009874f, 0.003081f, 0.001507f},
-    {0.002218f, 0.005712f, 0.001563f},
-    {0.053010f, -0.019844f, 0.008683f},
-    {-0.002507f, 0.005384f, 0.000938f},
-    {-0.002177f, -0.011366f, 0.003559f},
-    {-0.000261f, 0.015121f, -0.003240f},
-    {-0.003944f, -0.002083f, 0.005043f},
-    {-0.009138f, 0.011336f, 0.003781f},
-    {0.011429f, 0.003985f, -0.003855f},
-    {0.010518f, -0.005586f, 0.010131f},
-    {0.007883f, 0.002912f, -0.001473f},
-    {-0.003318f, -0.003160f, 0.003684f},
-    {-0.034560f, -0.008740f, 0.012996f},
-    {0.000166f, 0.001079f, -0.012153f},
-    {0.017772f, 0.000937f, -0.011953f}};
-float flux2_latent_rgb_bias[3] = {-0.028738f, -0.098463f, -0.107619f};
-
-// This one was taken straight from
-// https://github.com/Stability-AI/sd3.5/blob/8565799a3b41eb0c7ba976d18375f0f753f56402/sd3_impls.py#L288-L303
-// (MiT Licence)
-const float sd3_latent_rgb_proj[16][3] = {
-    {-0.0645f, 0.0177f, 0.1052f},
-    {0.0028f, 0.0312f, 0.0650f},
-    {0.1848f, 0.0762f, 0.0360f},
-    {0.0944f, 0.0360f, 0.0889f},
-    {0.0897f, 0.0506f, -0.0364f},
-    {-0.0020f, 0.1203f, 0.0284f},
-    {0.0855f, 0.0118f, 0.0283f},
-    {-0.0539f, 0.0658f, 0.1047f},
-    {-0.0057f, 0.0116f, 0.0700f},
-    {-0.0412f, 0.0281f, -0.0039f},
-    {0.1106f, 0.1171f, 0.1220f},
-    {-0.0248f, 0.0682f, -0.0481f},
-    {0.0815f, 0.0846f, 0.1207f},
-    {-0.0120f, -0.0055f, -0.0867f},
-    {-0.0749f, -0.0634f, -0.0456f},
-    {-0.1418f, -0.1457f, -0.1259f},
-};
-float sd3_latent_rgb_bias[3] = {0, 0, 0};
-
-const float sdxl_latent_rgb_proj[4][3] = {
-    {0.258303f, 0.277640f, 0.329699f},
-    {-0.299701f, 0.105446f, 0.014194f},
-    {0.050522f, 0.186163f, -0.143257f},
-    {-0.211938f, -0.149892f, -0.080036f}};
-float sdxl_latent_rgb_bias[3] = {0.144381f, -0.033313f, 0.007061f};
-
-const float sd_latent_rgb_proj[4][3] = {
-    {0.337366f, 0.216344f, 0.257386f},
-    {0.165636f, 0.386828f, 0.046994f},
-    {-0.267803f, 0.237036f, 0.223517f},
-    {-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, 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 frames        = 1;
-    if (ggml_n_dims(latents) == 4) {
-        frames = static_cast<uint32_t>(latents->ne[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++) {
-                int latent_x = rgb_x / patch_size;
-                int latent_y = rgb_y / patch_size;
-
-                int channel_offset = 0;
-                if (patch_size > 1) {
-                    channel_offset = ((rgb_y % patch_size) * patch_size + (rgb_x % patch_size));
-                }
-
-                size_t latent_id = (latent_x * latents->nb[0] + latent_y * latents->nb[1] + k * latents->nb[2]);
-
-                // should be incremented by 1 for each pixel
-                size_t pixel_id = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x;
-
-                float r = 0, g = 0, b = 0;
-                if (latent_rgb_proj != nullptr) {
-                    for (uint32_t 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];
-                        b += value * latent_rgb_proj[d][2];
-                    }
-                } else {
-                    // interpret first 3 channels as RGB
-                    r = *(float*)((char*)latents->data + latent_id + 0 * latents->nb[ggml_n_dims(latents) - 1]);
-                    g = *(float*)((char*)latents->data + latent_id + 1 * latents->nb[ggml_n_dims(latents) - 1]);
-                    b = *(float*)((char*)latents->data + latent_id + 2 * latents->nb[ggml_n_dims(latents) - 1]);
-                }
-                if (latent_rgb_bias != nullptr) {
-                    // bias
-                    r += latent_rgb_bias[0];
-                    g += latent_rgb_bias[1];
-                    b += latent_rgb_bias[2];
-                }
-                // change range
-                r = r * .5f + .5f;
-                g = g * .5f + .5f;
-                b = b * .5f + .5f;
-
-                // clamp rgb values to [0,1] range
-                r = r >= 0 ? r <= 1 ? r : 1 : 0;
-                g = g >= 0 ? g <= 1 ? g : 1 : 0;
-                b = b >= 0 ? b <= 1 ? b : 1 : 0;
-
-                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);
-            }
-        }
-    }
-}
+#include <cstddef>
+#include <cstdint>
+#include "ggml.h"
+
+const float wan_21_latent_rgb_proj[16][3] = {
+    {0.015123f, -0.148418f, 0.479828f},
+    {0.003652f, -0.010680f, -0.037142f},
+    {0.212264f, 0.063033f, 0.016779f},
+    {0.232999f, 0.406476f, 0.220125f},
+    {-0.051864f, -0.082384f, -0.069396f},
+    {0.085005f, -0.161492f, 0.010689f},
+    {-0.245369f, -0.506846f, -0.117010f},
+    {-0.151145f, 0.017721f, 0.007207f},
+    {-0.293239f, -0.207936f, -0.421135f},
+    {-0.187721f, 0.050783f, 0.177649f},
+    {-0.013067f, 0.265964f, 0.166578f},
+    {0.028327f, 0.109329f, 0.108642f},
+    {-0.205343f, 0.043991f, 0.148914f},
+    {0.014307f, -0.048647f, -0.007219f},
+    {0.217150f, 0.053074f, 0.319923f},
+    {0.155357f, 0.083156f, 0.064780f}};
+float wan_21_latent_rgb_bias[3] = {-0.270270f, -0.234976f, -0.456853f};
+
+const float wan_22_latent_rgb_proj[48][3] = {
+    {0.017126f, -0.027230f, -0.019257f},
+    {-0.113739f, -0.028715f, -0.022885f},
+    {-0.000106f, 0.021494f, 0.004629f},
+    {-0.013273f, -0.107137f, -0.033638f},
+    {-0.000381f, 0.000279f, 0.025877f},
+    {-0.014216f, -0.003975f, 0.040528f},
+    {0.001638f, -0.000748f, 0.011022f},
+    {0.029238f, -0.006697f, 0.035933f},
+    {0.021641f, -0.015874f, 0.040531f},
+    {-0.101984f, -0.070160f, -0.028855f},
+    {0.033207f, -0.021068f, 0.002663f},
+    {-0.104711f, 0.121673f, 0.102981f},
+    {0.082647f, -0.004991f, 0.057237f},
+    {-0.027375f, 0.031581f, 0.006868f},
+    {-0.045434f, 0.029444f, 0.019287f},
+    {-0.046572f, -0.012537f, 0.006675f},
+    {0.074709f, 0.033690f, 0.025289f},
+    {-0.008251f, -0.002745f, -0.006999f},
+    {0.012685f, -0.061856f, -0.048658f},
+    {0.042304f, -0.007039f, 0.000295f},
+    {-0.007644f, -0.060843f, -0.033142f},
+    {0.159909f, 0.045628f, 0.367541f},
+    {0.095171f, 0.086438f, 0.010271f},
+    {0.006812f, 0.019643f, 0.029637f},
+    {0.003467f, -0.010705f, 0.014252f},
+    {-0.099681f, -0.066272f, -0.006243f},
+    {0.047357f, 0.037040f, 0.000185f},
+    {-0.041797f, -0.089225f, -0.032257f},
+    {0.008928f, 0.017028f, 0.018684f},
+    {-0.042255f, 0.016045f, 0.006849f},
+    {0.011268f, 0.036462f, 0.037387f},
+    {0.011553f, -0.016375f, -0.048589f},
+    {0.046266f, -0.027189f, 0.056979f},
+    {0.009640f, -0.017576f, 0.030324f},
+    {-0.045794f, -0.036083f, -0.010616f},
+    {0.022418f, 0.039783f, -0.032939f},
+    {-0.052714f, -0.015525f, 0.007438f},
+    {0.193004f, 0.223541f, 0.264175f},
+    {-0.059406f, -0.008188f, 0.022867f},
+    {-0.156742f, -0.263791f, -0.007385f},
+    {-0.015717f, 0.016570f, 0.033969f},
+    {0.037969f, 0.109835f, 0.200449f},
+    {-0.000782f, -0.009566f, -0.008058f},
+    {0.010709f, 0.052960f, -0.044195f},
+    {0.017271f, 0.045839f, 0.034569f},
+    {0.009424f, 0.013088f, -0.001714f},
+    {-0.024805f, -0.059378f, -0.033756f},
+    {-0.078293f, 0.029070f, 0.026129f}};
+float wan_22_latent_rgb_bias[3] = {0.013160f, -0.096492f, -0.071323f};
+
+const float flux_latent_rgb_proj[16][3] = {
+    {-0.041168f, 0.019917f, 0.097253f},
+    {0.028096f, 0.026730f, 0.129576f},
+    {0.065618f, -0.067950f, -0.014651f},
+    {-0.012998f, -0.014762f, 0.081251f},
+    {0.078567f, 0.059296f, -0.024687f},
+    {-0.015987f, -0.003697f, 0.005012f},
+    {0.033605f, 0.138999f, 0.068517f},
+    {-0.024450f, -0.063567f, -0.030101f},
+    {-0.040194f, -0.016710f, 0.127185f},
+    {0.112681f, 0.088764f, -0.041940f},
+    {-0.023498f, 0.093664f, 0.025543f},
+    {0.082899f, 0.048320f, 0.007491f},
+    {0.075712f, 0.074139f, 0.081965f},
+    {-0.143501f, 0.018263f, -0.136138f},
+    {-0.025767f, -0.082035f, -0.040023f},
+    {-0.111849f, -0.055589f, -0.032361f}};
+float flux_latent_rgb_bias[3] = {0.024600f, -0.006937f, -0.008089f};
+
+const float flux2_latent_rgb_proj[32][3] = {
+    {0.000736f, -0.008385f, -0.019710f},
+    {-0.001352f, -0.016392f, 0.020693f},
+    {-0.006376f, 0.002428f, 0.036736f},
+    {0.039384f, 0.074167f, 0.119789f},
+    {0.007464f, -0.005705f, -0.004734f},
+    {-0.004086f, 0.005287f, -0.000409f},
+    {-0.032835f, 0.050802f, -0.028120f},
+    {-0.003158f, -0.000835f, 0.000406f},
+    {-0.112840f, -0.084337f, -0.023083f},
+    {0.001462f, -0.006656f, 0.000549f},
+    {-0.009980f, -0.007480f, 0.009702f},
+    {0.032540f, 0.000214f, -0.061388f},
+    {0.011023f, 0.000694f, 0.007143f},
+    {-0.001468f, -0.006723f, -0.001678f},
+    {-0.005921f, -0.010320f, -0.003907f},
+    {-0.028434f, 0.027584f, 0.018457f},
+    {0.014349f, 0.011523f, 0.000441f},
+    {0.009874f, 0.003081f, 0.001507f},
+    {0.002218f, 0.005712f, 0.001563f},
+    {0.053010f, -0.019844f, 0.008683f},
+    {-0.002507f, 0.005384f, 0.000938f},
+    {-0.002177f, -0.011366f, 0.003559f},
+    {-0.000261f, 0.015121f, -0.003240f},
+    {-0.003944f, -0.002083f, 0.005043f},
+    {-0.009138f, 0.011336f, 0.003781f},
+    {0.011429f, 0.003985f, -0.003855f},
+    {0.010518f, -0.005586f, 0.010131f},
+    {0.007883f, 0.002912f, -0.001473f},
+    {-0.003318f, -0.003160f, 0.003684f},
+    {-0.034560f, -0.008740f, 0.012996f},
+    {0.000166f, 0.001079f, -0.012153f},
+    {0.017772f, 0.000937f, -0.011953f}};
+float flux2_latent_rgb_bias[3] = {-0.028738f, -0.098463f, -0.107619f};
+
+// This one was taken straight from
+// https://github.com/Stability-AI/sd3.5/blob/8565799a3b41eb0c7ba976d18375f0f753f56402/sd3_impls.py#L288-L303
+// (MiT Licence)
+const float sd3_latent_rgb_proj[16][3] = {
+    {-0.0645f, 0.0177f, 0.1052f},
+    {0.0028f, 0.0312f, 0.0650f},
+    {0.1848f, 0.0762f, 0.0360f},
+    {0.0944f, 0.0360f, 0.0889f},
+    {0.0897f, 0.0506f, -0.0364f},
+    {-0.0020f, 0.1203f, 0.0284f},
+    {0.0855f, 0.0118f, 0.0283f},
+    {-0.0539f, 0.0658f, 0.1047f},
+    {-0.0057f, 0.0116f, 0.0700f},
+    {-0.0412f, 0.0281f, -0.0039f},
+    {0.1106f, 0.1171f, 0.1220f},
+    {-0.0248f, 0.0682f, -0.0481f},
+    {0.0815f, 0.0846f, 0.1207f},
+    {-0.0120f, -0.0055f, -0.0867f},
+    {-0.0749f, -0.0634f, -0.0456f},
+    {-0.1418f, -0.1457f, -0.1259f},
+};
+float sd3_latent_rgb_bias[3] = {0, 0, 0};
+
+const float sdxl_latent_rgb_proj[4][3] = {
+    {0.258303f, 0.277640f, 0.329699f},
+    {-0.299701f, 0.105446f, 0.014194f},
+    {0.050522f, 0.186163f, -0.143257f},
+    {-0.211938f, -0.149892f, -0.080036f}};
+float sdxl_latent_rgb_bias[3] = {0.144381f, -0.033313f, 0.007061f};
+
+const float sd_latent_rgb_proj[4][3] = {
+    {0.337366f, 0.216344f, 0.257386f},
+    {0.165636f, 0.386828f, 0.046994f},
+    {-0.267803f, 0.237036f, 0.223517f},
+    {-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) {
+    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 frames        = 1;
+    if (ggml_n_dims(latents) == 4) {
+        frames = static_cast<uint32_t>(latents->ne[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++) {
+                int latent_x = rgb_x / patch_size;
+                int latent_y = rgb_y / patch_size;
+
+                int channel_offset = 0;
+                if (patch_size > 1) {
+                    channel_offset = ((rgb_y % patch_size) * patch_size + (rgb_x % patch_size));
+                }
+
+                size_t latent_id = (latent_x * latents->nb[0] + latent_y * latents->nb[1] + k * latents->nb[2]);
+
+                // should be incremented by 1 for each pixel
+                size_t pixel_id = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x;
+
+                float r = 0, g = 0, b = 0;
+                if (latent_rgb_proj != nullptr) {
+                    for (uint32_t 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];
+                        b += value * latent_rgb_proj[d][2];
+                    }
+                } else {
+                    // interpret first 3 channels as RGB
+                    r = *(float*)((char*)latents->data + latent_id + 0 * latents->nb[ggml_n_dims(latents) - 1]);
+                    g = *(float*)((char*)latents->data + latent_id + 1 * latents->nb[ggml_n_dims(latents) - 1]);
+                    b = *(float*)((char*)latents->data + latent_id + 2 * latents->nb[ggml_n_dims(latents) - 1]);
+                }
+                if (latent_rgb_bias != nullptr) {
+                    // bias
+                    r += latent_rgb_bias[0];
+                    g += latent_rgb_bias[1];
+                    b += latent_rgb_bias[2];
+                }
+                // change range
+                r = r * .5f + .5f;
+                g = g * .5f + .5f;
+                b = b * .5f + .5f;
+
+                // clamp rgb values to [0,1] range
+                r = r >= 0 ? r <= 1 ? r : 1 : 0;
+                g = g >= 0 ? g <= 1 ? g : 1 : 0;
+                b = b >= 0 ? b <= 1 ? b : 1 : 0;
+
+                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);
+            }
+        }
+    }
+}

src/llm.hpp

@@ -19,6 +19,7 @@
 #include "json.hpp"
 #include "rope.hpp"
 #include "tokenize_util.h"
+#include "vocab/vocab.h"
 
 namespace LLM {
     constexpr int LLM_GRAPH_SIZE = 10240;
@@ -365,7 +366,7 @@ namespace LLM {
             if (merges_utf8_str.size() > 0) {
                 load_from_merges(merges_utf8_str);
             } else {
-                load_from_merges(ModelLoader::load_qwen2_merges());
+                load_from_merges(load_qwen2_merges());
             }
         }
     };
@@ -466,7 +467,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(ModelLoader::load_mistral_merges(), ModelLoader::load_mistral_vocab_json());
+                load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
             }
         }
     };
@@ -521,7 +522,7 @@ namespace LLM {
             blocks["down_proj"] = std::shared_ptr<GGMLBlock>(new Linear(intermediate_size, hidden_size, bias));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -581,7 +582,7 @@ namespace LLM {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, 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,
@@ -630,7 +631,7 @@ namespace LLM {
             blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, dim));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -667,10 +668,10 @@ namespace LLM {
             blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mask = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* mask = nullptr) {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
@@ -717,10 +718,10 @@ namespace LLM {
             blocks["norm2"] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_size, eps));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mask = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             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"]);
@@ -777,12 +778,12 @@ namespace LLM {
             blocks["merger"] = std::shared_ptr<GGMLBlock>(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size));
         }
 
-        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) {
+        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) {
             // 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)]
@@ -835,10 +836,10 @@ namespace LLM {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* input_pos,
-                                    struct ggml_tensor* attention_mask = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* input_pos,
+                             ggml_tensor* attention_mask = nullptr) {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
@@ -897,10 +898,10 @@ namespace LLM {
             blocks["post_attention_layernorm"] = std::make_shared<RMSNorm>(params.hidden_size, params.rms_norm_eps);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* input_pos,
-                                    struct ggml_tensor* attention_mask = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* input_pos,
+                             ggml_tensor* attention_mask = nullptr) {
             // 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"]);
@@ -935,12 +936,12 @@ namespace LLM {
             blocks["norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.hidden_size, params.rms_norm_eps));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* input_ids,
-                                    struct ggml_tensor* input_pos,
-                                    struct ggml_tensor* attention_mask,
-                                    std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                                    std::set<int> out_layers) {
+        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) {
             // input_ids: [N, n_token]
             // return: [N, n_token, hidden_size]
 
@@ -1036,12 +1037,12 @@ namespace LLM {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* input_ids,
-                                    struct ggml_tensor* input_pos,
-                                    struct ggml_tensor* attention_mask,
-                                    std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                                    std::set<int> out_layers) {
+        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) {
             // input_ids: [N, n_token]
             auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
 
@@ -1049,12 +1050,12 @@ namespace LLM {
             return x;
         }
 
-        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_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) {
             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);
@@ -1155,35 +1156,35 @@ namespace LLM {
             return llm_arch_to_str[static_cast<int>(params.arch)];
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
             model.get_param_tensors(tensors, prefix);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* input_ids,
-                                    struct ggml_tensor* input_pos,
-                                    struct ggml_tensor* attention_mask,
-                                    std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                                    std::set<int> out_layers) {
+        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]
             return hidden_states;
         }
 
-        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) {
+        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) {
             auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask);
             return hidden_states;
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
-                                        struct ggml_tensor* attention_mask,
-                                        std::vector<std::pair<int, ggml_tensor*>> image_embeds,
-                                        std::set<int> out_layers) {
-            struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+        ggml_cgraph* build_graph(ggml_tensor* input_ids,
+                                 ggml_tensor* attention_mask,
+                                 std::vector<std::pair<int, ggml_tensor*>> image_embeds,
+                                 std::set<int> out_layers) {
+            ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
             input_ids = to_backend(input_ids);
 
@@ -1231,7 +1232,7 @@ namespace LLM {
 
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);
+            ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);
 
             ggml_build_forward_expand(gf, hidden_states);
 
@@ -1239,13 +1240,13 @@ namespace LLM {
         }
 
         bool compute(const int n_threads,
-                     struct ggml_tensor* input_ids,
-                     struct ggml_tensor* attention_mask,
+                     ggml_tensor* input_ids,
+                     ggml_tensor* attention_mask,
                      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* {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input_ids, attention_mask, image_embeds, out_layers);
             };
             return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
@@ -1260,7 +1261,7 @@ namespace LLM {
             return grid_t * grid_h * grid_w;
         }
 
-        struct ggml_tensor* process_image(struct ggml_context* ctx, struct ggml_tensor* image) {
+        ggml_tensor* process_image(ggml_context* ctx, 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];
@@ -1287,8 +1288,8 @@ namespace LLM {
             return image;
         }
 
-        struct ggml_cgraph* build_encode_image_graph(struct ggml_tensor* image) {
-            struct ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
+        ggml_cgraph* build_encode_image_graph(ggml_tensor* image) {
+            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);
@@ -1398,23 +1399,23 @@ namespace LLM {
             // pe->data = nullptr;
             set_backend_tensor_data(pe, pe_vec.data());
 
-            auto runnter_ctx                  = get_context();
-            struct ggml_tensor* hidden_states = vision_forward(&runnter_ctx,
-                                                               pixel_values,
-                                                               pe,
-                                                               window_index,
-                                                               window_inverse_index,
-                                                               window_mask);
+            auto runnter_ctx           = get_context();
+            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;
         }
 
         void encode_image(const int n_threads,
-                          struct ggml_tensor* image,
+                          ggml_tensor* image,
                           ggml_tensor** output,
                           ggml_context* output_ctx = nullptr) {
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
             GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
@@ -1439,7 +1440,7 @@ namespace LLM {
             }
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
             model.get_param_tensors(tensors, prefix);
         }
 
@@ -1491,12 +1492,12 @@ namespace LLM {
         }
 
         void test() {
-            struct ggml_init_params params;
+            ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            struct ggml_context* work_ctx = ggml_init(params);
+            ggml_context* work_ctx = ggml_init(params);
             GGML_ASSERT(work_ctx != nullptr);
             bool test_mistral          = false;
             bool test_qwen3            = true;
@@ -1508,7 +1509,7 @@ namespace LLM {
                 {
                     auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
                     print_ggml_tensor(image, false, "image");
-                    struct ggml_tensor* out = nullptr;
+                    ggml_tensor* out = nullptr;
 
                     int64_t t0 = ggml_time_ms();
                     model.encode_image(8, image, &out, work_ctx);
@@ -1546,8 +1547,8 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
-                struct ggml_tensor* out = nullptr;
+                auto input_ids   = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 model.compute(8, input_ids, nullptr, image_embeds, {}, &out, work_ctx);
@@ -1560,7 +1561,7 @@ namespace LLM {
                 // ggml_set_f32(image, 0.f);
                 auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
                 print_ggml_tensor(image, false, "image");
-                struct ggml_tensor* out = nullptr;
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 model.encode_image(8, image, &out, work_ctx);
@@ -1586,8 +1587,8 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
-                struct ggml_tensor* out = nullptr;
+                auto input_ids   = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 model.compute(8, input_ids, nullptr, {}, {10, 20, 30}, &out, work_ctx);
@@ -1609,8 +1610,8 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
-                struct ggml_tensor* out = nullptr;
+                auto input_ids   = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 model.compute(8, input_ids, nullptr, {}, {35}, &out, work_ctx);
@@ -1632,8 +1633,8 @@ namespace LLM {
                     printf("%d ", token);
                 }
                 printf("\n");
-                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, tokens);
-                struct ggml_tensor* out = nullptr;
+                auto input_ids   = vector_to_ggml_tensor_i32(work_ctx, tokens);
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 model.compute(8, input_ids, nullptr, {}, {}, &out, work_ctx);

src/lora.hpp

@@ -9,7 +9,7 @@
 struct LoraModel : public GGMLRunner {
     std::string lora_id;
     float multiplier = 1.0f;
-    std::unordered_map<std::string, struct ggml_tensor*> lora_tensors;
+    std::unordered_map<std::string, 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;
-            struct 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;
+            ggml_tensor* real  = ggml_new_tensor(params_ctx,
+                                                 ts.type,
+                                                 ts.n_dims,
+                                                 ts.ne);
+            lora_tensors[name] = real;
         }
 
         alloc_params_buffer();
@@ -337,10 +337,10 @@ struct LoraModel : public GGMLRunner {
             }
             scale_value *= multiplier;
 
-            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_ext_scale(ctx, curr_updown, scale_value, true);
+            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);
             if (updown == nullptr) {
                 updown = curr_updown;
             } else {
@@ -747,9 +747,9 @@ struct LoraModel : public GGMLRunner {
         return out_diff;
     }
 
-    struct ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
+    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;
-        struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
+        ggml_cgraph* gf        = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
 
         preprocess_lora_tensors(model_tensors);
 
@@ -788,8 +788,8 @@ struct LoraModel : public GGMLRunner {
         return gf;
     }
 
-    void apply(std::map<std::string, struct ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+    void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_lora_graph(model_tensors, version);
         };
         GGMLRunner::compute(get_graph, n_threads, false);

src/ltxv.hpp

@@ -1,8 +1,7 @@
 #ifndef __LTXV_HPP__
 #define __LTXV_HPP__
 
-#include "common.hpp"
-#include "ggml_extend.hpp"
+#include "common_block.hpp"
 
 namespace LTXV {
 
@@ -27,9 +26,9 @@ namespace LTXV {
                                                                      bias));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    bool causal = true) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             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"]);

src/mmdit.hpp

@@ -27,7 +27,7 @@ public:
         blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -72,7 +72,7 @@ public:
                                                                bias));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [N, C, H, W]
         // return: [N, H*W, embed_dim]
         auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
@@ -111,7 +111,7 @@ public:
         blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_channels, true, true));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* t) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
         // x: [N, input_dim]
         // return: [N, hidden_size]
         auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
@@ -175,7 +175,7 @@ public:
         }
     }
 
-    std::vector<struct ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, 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};
     }
 
-    struct ggml_tensor* post_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* post_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
         GGML_ASSERT(!pre_only);
 
         auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
@@ -208,8 +208,8 @@ public:
     }
 
     // x: [N, n_token, dim]
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         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]
@@ -217,10 +217,10 @@ public:
     }
 };
 
-__STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
-                                               struct ggml_tensor* x,
-                                               struct ggml_tensor* shift,
-                                               struct ggml_tensor* scale) {
+__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+                                        ggml_tensor* x,
+                                        ggml_tensor* shift,
+                                        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,
-                                                                                                                struct ggml_tensor* x,
-                                                                                                                struct ggml_tensor* c) {
+                                                                                                                ggml_tensor* x,
+                                                                                                                ggml_tensor* c) {
         GGML_ASSERT(self_attn);
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
@@ -309,9 +309,9 @@ public:
         return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}};
     }
 
-    std::pair<std::vector<struct ggml_tensor*>, std::vector<struct ggml_tensor*>> pre_attention(GGMLRunnerContext* ctx,
-                                                                                                struct ggml_tensor* x,
-                                                                                                struct ggml_tensor* c) {
+    std::pair<std::vector<ggml_tensor*>, std::vector<ggml_tensor*>> pre_attention(GGMLRunnerContext* ctx,
+                                                                                  ggml_tensor* x,
+                                                                                  ggml_tensor* c) {
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
         auto norm1              = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
@@ -346,15 +346,15 @@ public:
         }
     }
 
-    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) {
+    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) {
         // attn_out: [N, n_token, hidden_size]
         // x: [N, n_token, hidden_size]
         // gate_msa: [N, hidden_size]
@@ -384,13 +384,13 @@ public:
         return x;
     }
 
-    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) {
+    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) {
         // attn_out: [N, n_token, hidden_size]
         // x: [N, n_token, hidden_size]
         // gate_msa: [N, hidden_size]
@@ -416,9 +416,9 @@ public:
         return x;
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* c) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* c) {
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
         // return: [N, n_token, hidden_size]
@@ -463,11 +463,11 @@ public:
     }
 };
 
-__STATIC_INLINE__ std::pair<struct ggml_tensor*, struct ggml_tensor*>
+__STATIC_INLINE__ std::pair<ggml_tensor*, ggml_tensor*>
 block_mixing(GGMLRunnerContext* ctx,
-             struct ggml_tensor* context,
-             struct ggml_tensor* x,
-             struct ggml_tensor* c,
+             ggml_tensor* context,
+             ggml_tensor* x,
+             ggml_tensor* c,
              std::shared_ptr<DismantledBlock> context_block,
              std::shared_ptr<DismantledBlock> x_block) {
     // context: [N, n_context, hidden_size]
@@ -489,7 +489,7 @@ block_mixing(GGMLRunnerContext* ctx,
         x_qkv                    = x_qkv_intermediates.first;
         x_intermediates          = x_qkv_intermediates.second;
     }
-    std::vector<struct ggml_tensor*> qkv;
+    std::vector<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));
     }
@@ -563,10 +563,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<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                                struct ggml_tensor* context,
-                                                                struct ggml_tensor* x,
-                                                                struct ggml_tensor* c) {
+    std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                  ggml_tensor* context,
+                                                  ggml_tensor* x,
+                                                  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 +586,9 @@ public:
         blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* c) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* c) {
         // x: [N, n_token, hidden_size]
         // c: [N, hidden_size]
         // return: [N, n_token, patch_size * patch_size * out_channels]
@@ -626,7 +626,7 @@ protected:
     int64_t hidden_size;
     std::string qk_norm;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+    void init_params(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 +705,8 @@ public:
         blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(hidden_size, patch_size, out_channels));
     }
 
-    struct ggml_tensor*
-    cropped_pos_embed(struct ggml_context* ctx,
+    ggml_tensor*
+    cropped_pos_embed(ggml_context* ctx,
                       int64_t h,
                       int64_t w) {
         auto pos_embed = params["pos_embed"];
@@ -745,33 +745,11 @@ public:
         return spatial_pos_embed;
     }
 
-    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>()) {
+    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>()) {
         // x: [N, H*W, hidden_size]
         // context: [N, n_context, d_context]
         // c: [N, hidden_size]
@@ -796,12 +774,12 @@ public:
         return x;
     }
 
-    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>()) {
+    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>()) {
         // 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
@@ -811,11 +789,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]
@@ -834,7 +812,7 @@ public:
 
         x = forward_core_with_concat(ctx, x, c, context, skip_layers);  // (N, H*W, patch_size ** 2 * out_channels)
 
-        x = unpatchify(ctx->ggml_ctx, x, h, w);  // [N, C, H, W]
+        x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, /*patch_last*/ false);  // [N, C, H, W]
 
         return x;
     }
@@ -854,29 +832,29 @@ struct MMDiTRunner : public GGMLRunner {
         return "mmdit";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         mmdit.get_param_tensors(tensors, prefix);
     }
 
-    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_cgraph* build_graph(ggml_tensor* x,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* context,
+                             ggml_tensor* y,
+                             std::vector<int> skip_layers = std::vector<int>()) {
+        ggml_cgraph* gf = new_graph_custom(MMDIT_GRAPH_SIZE);
 
         x         = to_backend(x);
         context   = to_backend(context);
         y         = to_backend(y);
         timesteps = to_backend(timesteps);
 
-        auto runner_ctx         = get_context();
-        struct ggml_tensor* out = mmdit.forward(&runner_ctx,
-                                                x,
-                                                timesteps,
-                                                y,
-                                                context,
-                                                skip_layers);
+        auto runner_ctx  = get_context();
+        ggml_tensor* out = mmdit.forward(&runner_ctx,
+                                         x,
+                                         timesteps,
+                                         y,
+                                         context,
+                                         skip_layers);
 
         ggml_build_forward_expand(gf, out);
 
@@ -884,18 +862,18 @@ struct MMDiTRunner : public GGMLRunner {
     }
 
     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>()) {
+                 ggml_tensor* x,
+                 ggml_tensor* timesteps,
+                 ggml_tensor* context,
+                 ggml_tensor* y,
+                 ggml_tensor** output         = nullptr,
+                 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 = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
@@ -903,12 +881,12 @@ struct MMDiTRunner : public GGMLRunner {
     }
 
     void test() {
-        struct ggml_init_params params;
+        ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        struct ggml_context* work_ctx = ggml_init(params);
+        ggml_context* work_ctx = ggml_init(params);
         GGML_ASSERT(work_ctx != nullptr);
 
         {
@@ -930,7 +908,7 @@ struct MMDiTRunner : public GGMLRunner {
             ggml_set_f32(y, 0.01f);
             // print_ggml_tensor(y);
 
-            struct ggml_tensor* out = nullptr;
+            ggml_tensor* out = nullptr;
 
             int64_t t0 = ggml_time_ms();
             compute(8, x, timesteps, context, y, &out, work_ctx);

src/model.cpp

@@ -16,10 +16,6 @@
 #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"
@@ -291,7 +287,7 @@ void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
 }
 
 bool is_zip_file(const std::string& file_path) {
-    struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
     if (zip == nullptr) {
         return false;
     }
@@ -457,9 +453,9 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
     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);
-        struct 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);
+        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());
 
@@ -816,7 +812,7 @@ struct PickleTensorReader {
         }
     }
 
-    void read_string(const std::string& str, struct zip_t* zip, std::string dir) {
+    void read_string(const std::string& str, zip_t* zip, std::string dir) {
         if (str == "storage") {
             read_global_type = true;
         } else if (str != "state_dict") {
@@ -999,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;
 
-    struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    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;
@@ -1061,6 +1057,9 @@ 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;
             }
@@ -1105,10 +1104,12 @@ 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) {
+        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) {
+        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" ||
@@ -1340,36 +1341,6 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
     }
 }
 
-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, bool enable_mmap) {
     int64_t process_time_ms = 0;
     std::atomic<int64_t> read_time_ms(0);
@@ -1442,7 +1413,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;
-                struct zip_t* zip = nullptr;
+                zip_t* zip = nullptr;
                 if (is_zip) {
                     zip = zip_open(file_path.c_str(), 0, 'r');
                     if (zip == nullptr) {
@@ -1630,7 +1601,7 @@ 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, struct ggml_tensor*>& tensors,
+bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                std::set<std::string> ignore_tensors,
                                int n_threads,
                                bool enable_mmap) {
@@ -1644,7 +1615,7 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
             tensor_names_in_file.insert(name);
         }
 
-        struct ggml_tensor* real;
+        ggml_tensor* real;
         if (tensors.find(name) != tensors.end()) {
             real = tensors[name];
         } else {

src/model.h

@@ -45,6 +45,7 @@ enum SDVersion {
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
     VERSION_QWEN_IMAGE,
+    VERSION_ANIMA,
     VERSION_FLUX2,
     VERSION_FLUX2_KLEIN,
     VERSION_Z_IMAGE,
@@ -122,6 +123,13 @@ 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;
@@ -146,6 +154,7 @@ 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;
     }
@@ -314,7 +323,7 @@ public:
     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, struct ggml_tensor*>& tensors,
+    bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                       std::set<std::string> ignore_tensors = {},
                       int n_threads                        = 0,
                       bool use_mmap                        = false);
@@ -331,13 +340,6 @@ 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__

src/name_conversion.cpp

@@ -653,6 +653,14 @@ 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);
@@ -664,6 +672,8 @@ 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;
 }
@@ -1110,7 +1120,11 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
         for (const auto& prefix : first_stage_model_prefix_vec) {
             if (starts_with(name, prefix)) {
                 name = convert_first_stage_model_name(name.substr(prefix.size()), prefix);
-                name = prefix + name;
+                if (version == VERSION_SDXS) {
+                    name = "tae." + name;
+                } else {
+                    name = prefix + name;
+                }
                 break;
             }
         }

src/pmid.hpp

@@ -21,14 +21,14 @@ public:
         blocks["layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(in_dim));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
 
-        struct ggml_tensor* r = x;
+        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);
@@ -54,8 +54,8 @@ public:
         blocks["1"]   = std::shared_ptr<GGMLBlock>(new Mlp(dim, inner_dim, dim, false));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x) {
         auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["0"]);
         auto ff   = std::dynamic_pointer_cast<Mlp>(blocks["1"]);
 
@@ -81,9 +81,9 @@ public:
         blocks["to_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim, false));
     }
 
-    struct ggml_tensor* reshape_tensor(struct ggml_context* ctx,
-                                       struct ggml_tensor* x,
-                                       int heads) {
+    ggml_tensor* reshape_tensor(ggml_context* ctx,
+                                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<struct ggml_tensor*> chunk_half(struct ggml_context* ctx,
-                                                struct ggml_tensor* x) {
+    std::vector<ggml_tensor*> chunk_half(ggml_context* ctx,
+                                         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)};
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* latents) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* latents) {
         // x (torch.Tensor): image features
         //     shape (b, n1, D)
         // latent (torch.Tensor): latent features
@@ -176,9 +176,9 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* latents,
-                                struct ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* latents,
+                         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));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* last_hidden_state) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         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);
-        struct 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);
+        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));
     }
 
-    struct ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* prompt_embeds,
-                                struct ggml_tensor* id_embeds) {
+    ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
+                         ggml_tensor* prompt_embeds,
+                         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;
     }
 
-    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) {
+    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) {
         // x: [N, channels, h, w]
 
-        struct ggml_tensor* valid_id_embeds = id_embeds;
+        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");
-        struct ggml_tensor* image_token_embeds = ggml_get_rows(ctx->ggml_ctx, prompt_embeds, class_tokens_mask_pos);
+        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]);
-        struct 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]);
+        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);
-        struct 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);
+        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));
     }
 
-    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) {
+    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) {
         // 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"]);
 
-        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]
+        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]
 
         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));
 
-        struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
-                                                                         prompt_embeds,
-                                                                         id_embeds,
-                                                                         class_tokens_mask,
-                                                                         class_tokens_mask_pos,
-                                                                         left, right);
+        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));
     }
 
-    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) {
+    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) {
         // 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"]);
 
-        // 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* 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* updated_prompt_embeds = fuse_module->forward(ctx,
-                                                                         prompt_embeds,
-                                                                         id_embeds,
-                                                                         class_tokens_mask,
-                                                                         class_tokens_mask_pos,
-                                                                         left, right);
+        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,18 +436,18 @@ public:
         return pm_version;
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, 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);
     }
 
-    struct ggml_cgraph* build_graph(  // struct ggml_allocr* allocr,
-        struct ggml_tensor* id_pixel_values,
-        struct ggml_tensor* prompt_embeds,
+    ggml_cgraph* build_graph(  // ggml_allocr* allocr,
+        ggml_tensor* id_pixel_values,
+        ggml_tensor* prompt_embeds,
         std::vector<bool>& class_tokens_mask,
-        struct ggml_tensor* id_embeds) {
+        ggml_tensor* id_embeds) {
         ctm.clear();
         ctmf16.clear();
         ctmpos.clear();
@@ -458,20 +458,20 @@ public:
 
         auto runner_ctx = get_context();
 
-        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+        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;
 
-        struct ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
+        ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
 
-        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);
+        ggml_tensor* id_pixel_values_d = to_backend(id_pixel_values);
+        ggml_tensor* prompt_embeds_d   = to_backend(prompt_embeds);
+        ggml_tensor* id_embeds_d       = to_backend(id_embeds);
 
-        struct ggml_tensor* left  = nullptr;
-        struct ggml_tensor* right = nullptr;
+        ggml_tensor* left  = nullptr;
+        ggml_tensor* right = nullptr;
         for (int i = 0; i < class_tokens_mask.size(); i++) {
             if (class_tokens_mask[i]) {
                 // printf(" 1,");
@@ -495,7 +495,7 @@ public:
             right = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type,
                                        hidden_size, seq_length - ctmpos[ctmpos.size() - 1] - 1, 1);
         }
-        struct ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
+        ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
 
         {
             if (type == GGML_TYPE_F16)
@@ -526,7 +526,7 @@ public:
                 }
             }
         }
-        struct ggml_tensor* updated_prompt_embeds = nullptr;
+        ggml_tensor* updated_prompt_embeds = nullptr;
         if (pm_version == PM_VERSION_1)
             updated_prompt_embeds = id_encoder.forward(&runner_ctx,
                                                        id_pixel_values_d,
@@ -549,13 +549,13 @@ public:
     }
 
     bool compute(const int n_threads,
-                 struct ggml_tensor* id_pixel_values,
-                 struct ggml_tensor* prompt_embeds,
-                 struct ggml_tensor* id_embeds,
+                 ggml_tensor* id_pixel_values,
+                 ggml_tensor* prompt_embeds,
+                 ggml_tensor* id_embeds,
                  std::vector<bool>& class_tokens_mask,
-                 struct ggml_tensor** updated_prompt_embeds,
+                 ggml_tensor** updated_prompt_embeds,
                  ggml_context* output_ctx) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> 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);
         };
@@ -566,7 +566,7 @@ public:
 };
 
 struct PhotoMakerIDEmbed : public GGMLRunner {
-    std::map<std::string, struct ggml_tensor*> tensors;
+    std::map<std::string, ggml_tensor*> tensors;
     std::string file_path;
     ModelLoader* model_loader;
     bool load_failed = false;
@@ -606,11 +606,11 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             }
             if (dry_run) {
                 std::lock_guard<std::mutex> lock(tensor_mutex);
-                struct ggml_tensor* real = ggml_new_tensor(params_ctx,
-                                                           tensor_storage.type,
-                                                           tensor_storage.n_dims,
-                                                           tensor_storage.ne);
-                tensors[name]            = real;
+                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;
@@ -629,8 +629,8 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
         return true;
     }
 
-    struct ggml_tensor* get() {
-        std::map<std::string, struct ggml_tensor*>::iterator pos;
+    ggml_tensor* get() {
+        std::map<std::string, ggml_tensor*>::iterator pos;
         pos = tensors.find("pmid.id_embeds");
         if (pos != tensors.end())
             return pos->second;

src/preprocessing.hpp

@@ -4,13 +4,13 @@
 #include "ggml_extend.hpp"
 #define M_PI_ 3.14159265358979323846f
 
-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);
+void convolve(ggml_tensor* input, ggml_tensor* output, ggml_tensor* kernel, int padding) {
+    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;
+    ggml_context* ctx0       = ggml_init(params);
+    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);
@@ -19,7 +19,7 @@ void convolve(struct ggml_tensor* input, struct ggml_tensor* output, struct ggml
     ggml_free(ctx0);
 }
 
-void gaussian_kernel(struct ggml_tensor* kernel) {
+void gaussian_kernel(ggml_tensor* kernel) {
     int ks_mid   = static_cast<int>(kernel->ne[0] / 2);
     float sigma  = 1.4f;
     float normal = 1.f / (2.0f * M_PI_ * powf(sigma, 2.0f));
@@ -33,7 +33,7 @@ void gaussian_kernel(struct ggml_tensor* kernel) {
     }
 }
 
-void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
+void grayscale(ggml_tensor* rgb_img, 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);
@@ -45,7 +45,7 @@ void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
     }
 }
 
-void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
+void prop_hypot(ggml_tensor* x, ggml_tensor* y, ggml_tensor* h) {
     int n_elements = static_cast<int>(ggml_nelements(h));
     float* dx      = (float*)x->data;
     float* dy      = (float*)y->data;
@@ -55,7 +55,7 @@ void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor
     }
 }
 
-void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
+void prop_arctan2(ggml_tensor* x, ggml_tensor* y, ggml_tensor* h) {
     int n_elements = static_cast<int>(ggml_nelements(h));
     float* dx      = (float*)x->data;
     float* dy      = (float*)y->data;
@@ -65,7 +65,7 @@ void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tens
     }
 }
 
-void normalize_tensor(struct ggml_tensor* g) {
+void normalize_tensor(ggml_tensor* g) {
     int n_elements = static_cast<int>(ggml_nelements(g));
     float* dg      = (float*)g->data;
     float max      = -INFINITY;
@@ -78,7 +78,7 @@ void normalize_tensor(struct ggml_tensor* g) {
     }
 }
 
-void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struct ggml_tensor* D) {
+void non_max_supression(ggml_tensor* result, ggml_tensor* G, 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_;
@@ -117,7 +117,7 @@ void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struc
     }
 }
 
-void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float low_threshold, float weak, float strong) {
+void threshold_hystersis(ggml_tensor* img, float high_threshold, float low_threshold, float weak, float strong) {
     int n_elements = static_cast<int>(ggml_nelements(img));
     float* imd     = (float*)img->data;
     float max      = -INFINITY;
@@ -163,11 +163,11 @@ void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float lo
 }
 
 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);
+    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;
+    ggml_context* work_ctx = ggml_init(params);
 
     if (!work_ctx) {
         LOG_ERROR("ggml_init() failed");
@@ -185,19 +185,19 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
         -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);
+    int kernel_size      = 5;
+    ggml_tensor* gkernel = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, 1);
+    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);
+    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);
+    ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
+    ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
+    ggml_tensor* iX         = ggml_dup_tensor(work_ctx, image_gray);
+    ggml_tensor* iY         = ggml_dup_tensor(work_ctx, image_gray);
+    ggml_tensor* G          = ggml_dup_tensor(work_ctx, image_gray);
+    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);

src/qwen_image.hpp

@@ -3,9 +3,8 @@
 
 #include <memory>
 
-#include "common.hpp"
+#include "common_block.hpp"
 #include "flux.hpp"
-#include "ggml_extend.hpp"
 
 namespace Qwen {
     constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480;
@@ -27,9 +26,9 @@ namespace Qwen {
             blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, out_dim, sample_proj_bias));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* sample,
-                                    struct ggml_tensor* condition = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* sample,
+                             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));
@@ -50,8 +49,8 @@ namespace Qwen {
             blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* timesteps) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* timesteps) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
@@ -108,10 +107,10 @@ namespace Qwen {
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                      struct ggml_tensor* img,
-                                                      struct ggml_tensor* txt,
-                                                      struct ggml_tensor* pe,
-                                                      struct ggml_tensor* mask = nullptr) {
+                                                      ggml_tensor* img,
+                                                      ggml_tensor* txt,
+                                                      ggml_tensor* pe,
+                                                      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]
@@ -250,11 +249,11 @@ namespace Qwen {
         }
 
         virtual std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                              struct ggml_tensor* img,
-                                                              struct ggml_tensor* txt,
-                                                              struct ggml_tensor* t_emb,
-                                                              struct ggml_tensor* pe,
-                                                              struct ggml_tensor* modulate_index = nullptr) {
+                                                              ggml_tensor* img,
+                                                              ggml_tensor* txt,
+                                                              ggml_tensor* t_emb,
+                                                              ggml_tensor* pe,
+                                                              ggml_tensor* modulate_index = 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]
@@ -326,9 +325,9 @@ namespace Qwen {
             blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(conditioning_embedding_dim, embedding_dim * 2, bias));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* c) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* c) {
             // x: [N, n_token, hidden_size]
             // c: [N, hidden_size]
             // return: [N, n_token, patch_size * patch_size * out_channels]
@@ -390,75 +389,12 @@ namespace Qwen {
             blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
         }
 
-        struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* 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_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-            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(GGMLRunnerContext* ctx,
-                                        struct ggml_tensor* x) {
-            x = pad_to_patch_size(ctx, x);
-            x = patchify(ctx->ggml_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,
-                                         struct ggml_tensor* modulate_index = nullptr) {
+        ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
+                                  ggml_tensor* x,
+                                  ggml_tensor* timestep,
+                                  ggml_tensor* context,
+                                  ggml_tensor* pe,
+                                  ggml_tensor* modulate_index = nullptr) {
             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"]);
@@ -468,7 +404,7 @@ namespace Qwen {
 
             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(ctx->ggml_ctx, timestep->ne[0], timestep->ne[1], timestep->ne[2], timestep->ne[3]));
+                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);
@@ -493,13 +429,13 @@ namespace Qwen {
             return img;
         }
 
-        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 = {},
-                                    struct ggml_tensor* modulate_index    = nullptr) {
+        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) {
             // Forward pass of DiT.
             // x: [N, C, H, W]
             // timestep: [N,]
@@ -512,19 +448,16 @@ namespace Qwen {
             int64_t C = x->ne[2];
             int64_t N = x->ne[3];
 
-            auto img           = process_img(ctx, x);
+            auto img           = DiT::pad_and_patchify(ctx, x, params.patch_size, params.patch_size);
             int64_t img_tokens = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = process_img(ctx, ref);
+                    ref = DiT::pad_and_patchify(ctx, ref, params.patch_size, params.patch_size);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
 
-            int64_t h_len = ((H + (params.patch_size / 2)) / params.patch_size);
-            int64_t w_len = ((W + (params.patch_size / 2)) / params.patch_size);
-
             auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
 
             if (out->ne[1] > img_tokens) {
@@ -533,11 +466,7 @@ 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 = 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]
+            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, params.patch_size, params.patch_size);  // [N, C, H, W]
 
             return out;
         }
@@ -592,17 +521,17 @@ namespace Qwen {
             return "qwen_image";
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
             qwen_image.get_param_tensors(tensors, prefix);
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
-                                        struct ggml_tensor* timesteps,
-                                        struct ggml_tensor* context,
-                                        std::vector<ggml_tensor*> ref_latents = {},
-                                        bool increase_ref_index               = false) {
+        ggml_cgraph* build_graph(ggml_tensor* x,
+                                 ggml_tensor* timesteps,
+                                 ggml_tensor* context,
+                                 std::vector<ggml_tensor*> ref_latents = {},
+                                 bool increase_ref_index               = false) {
             GGML_ASSERT(x->ne[3] == 1);
-            struct ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
+            ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
 
             x         = to_backend(x);
             context   = to_backend(context);
@@ -658,13 +587,13 @@ namespace Qwen {
 
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* out = qwen_image.forward(&runner_ctx,
-                                                         x,
-                                                         timesteps,
-                                                         context,
-                                                         pe,
-                                                         ref_latents,
-                                                         modulate_index);
+            ggml_tensor* out = qwen_image.forward(&runner_ctx,
+                                                  x,
+                                                  timesteps,
+                                                  context,
+                                                  pe,
+                                                  ref_latents,
+                                                  modulate_index);
 
             ggml_build_forward_expand(gf, out);
 
@@ -672,17 +601,17 @@ namespace Qwen {
         }
 
         bool compute(int n_threads,
-                     struct ggml_tensor* x,
-                     struct ggml_tensor* timesteps,
-                     struct ggml_tensor* context,
+                     ggml_tensor* x,
+                     ggml_tensor* timesteps,
+                     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) {
+                     ggml_tensor** output                  = nullptr,
+                     ggml_context* output_ctx              = nullptr) {
             // x: [N, in_channels, h, w]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
@@ -690,12 +619,12 @@ namespace Qwen {
         }
 
         void test() {
-            struct ggml_init_params params;
+            ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            struct ggml_context* work_ctx = ggml_init(params);
+            ggml_context* work_ctx = ggml_init(params);
             GGML_ASSERT(work_ctx != nullptr);
 
             {
@@ -712,7 +641,7 @@ namespace Qwen {
                 auto context = load_tensor_from_file(work_ctx, "./qwen_image_context.bin");
                 print_ggml_tensor(context);
 
-                struct ggml_tensor* out = nullptr;
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 compute(8, x, timesteps, context, {}, false, &out, work_ctx);

src/rope.hpp

@@ -43,7 +43,7 @@ namespace Rope {
 
     __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
                                                            int dim,
-                                                           int theta,
+                                                           float theta,
                                                            const std::vector<int>& axis_wrap_dims = {}) {
         assert(dim % 2 == 0);
         int half_dim = dim / 2;
@@ -167,7 +167,7 @@ namespace Rope {
 
     __STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
                                                   int bs,
-                                                  int theta,
+                                                  const std::vector<float>& axis_thetas,
                                                   const std::vector<int>& axes_dim,
                                                   const std::vector<std::vector<int>>& wrap_dims = {}) {
         std::vector<std::vector<float>> trans_ids = transpose(ids);
@@ -188,8 +188,12 @@ namespace Rope {
             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], theta, axis_wrap_dims);  // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
+                rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims);  // [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) {
@@ -203,6 +207,15 @@ 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,
@@ -332,7 +345,7 @@ namespace Rope {
                 }
             }
         }
-        return embed_nd(ids, bs, theta, axes_dim, wrap_dims);
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
@@ -421,7 +434,7 @@ namespace Rope {
                 }
             }
         }
-        return embed_nd(ids, bs, theta, axes_dim, wrap_dims);
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@@ -475,7 +488,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, theta, axes_dim);
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
@@ -511,7 +524,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, theta, axes_dim);
+        return embed_nd(ids, 1, static_cast<float>(theta), axes_dim);
     }
 
     __STATIC_INLINE__ int bound_mod(int a, int m) {
@@ -584,13 +597,13 @@ namespace Rope {
             }
         }
 
-        return embed_nd(ids, bs, theta, axes_dim, wrap_dims);
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
-    __STATIC_INLINE__ struct ggml_tensor* apply_rope(struct ggml_context* ctx,
-                                                     struct ggml_tensor* x,
-                                                     struct ggml_tensor* pe,
-                                                     bool rope_interleaved = true) {
+    __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
+                                              ggml_tensor* x,
+                                              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];
@@ -628,14 +641,14 @@ namespace Rope {
         return x_out;
     }
 
-    __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) {
+    __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) {
         // q,k,v: [N, L, n_head, d_head]
         // pe: [L, d_head/2, 2, 2]
         // return: [N, L, n_head*d_head]

src/spectrum.hpp

@@ -0,0 +1,195 @@
+#ifndef __SPECTRUM_HPP__
+#define __SPECTRUM_HPP__
+
+#include <cmath>
+#include <cstring>
+#include <vector>
+
+#include "ggml_extend.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 ggml_tensor* denoised) {
+        int64_t ne        = ggml_nelements(denoised);
+        const float* data = (const float*)denoised->data;
+
+        H_buf.emplace_back(data, data + ne);
+        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(ggml_tensor* denoised) {
+        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);
+
+        // Design matrix X: K_curr x M1 (Chebyshev basis)
+        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];
+        }
+
+        // x_star: Chebyshev basis at current tau
+        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];
+
+        // XtX = X^T X + lambda I
+        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);
+            }
+        }
+
+        // Cholesky decomposition
+        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);
+        }
+
+        // Solve XtX v = x_star
+        std::vector<float> v(M1);
+        cholesky_solve(L.data(), x_star.data(), v.data(), M1);
+
+        // Prediction weights per history entry
+        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];
+
+        // Blend Chebyshev and Taylor predictions
+        float* out          = (float*)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/t5.hpp

@@ -14,6 +14,7 @@
 #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.
@@ -210,9 +211,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, <EFBFBD>? xt) = best_score(x1, x2, <20>? 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, <EFBFBD>?
         // x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
         // not related to each other at all.
         //
@@ -341,9 +342,9 @@ protected:
 public:
     explicit T5UniGramTokenizer(bool is_umt5 = false) {
         if (is_umt5) {
-            InitializePieces(ModelLoader::load_umt5_tokenizer_json());
+            InitializePieces(load_umt5_tokenizer_json());
         } else {
-            InitializePieces(ModelLoader::load_t5_tokenizer_json());
+            InitializePieces(load_t5_tokenizer_json());
         }
 
         min_score_ = FLT_MAX;
@@ -461,7 +462,7 @@ protected:
     int64_t hidden_size;
     float eps;
 
-    void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+    void init_params(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);
     }
@@ -472,10 +473,10 @@ public:
         : hidden_size(hidden_size),
           eps(eps) {}
 
-    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);
+    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);
         return x;
     }
 };
@@ -487,7 +488,7 @@ public:
         blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -509,7 +510,7 @@ public:
         blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false, false, false, scale));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -530,7 +531,7 @@ public:
         blocks["layer_norm"]     = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, 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"]);
@@ -569,8 +570,8 @@ public:
         }
     }
 
-    struct ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
-                                     struct ggml_tensor* relative_position_bucket) {
+    ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
+                              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)
@@ -579,11 +580,11 @@ public:
     }
 
     // x: [N, n_token, model_dim]
-    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) {
+    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) {
         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"]);
@@ -628,11 +629,11 @@ public:
         blocks["layer_norm"]    = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
     }
 
-    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) {
+    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) {
         // 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"]);
@@ -654,11 +655,11 @@ public:
         blocks["layer.1"] = std::shared_ptr<GGMLBlock>(new T5LayerFF(model_dim, ff_dim));
     }
 
-    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) {
+    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) {
         // 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"]);
@@ -689,11 +690,11 @@ public:
         blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
     }
 
-    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) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* past_bias                = nullptr,
+                         ggml_tensor* attention_mask           = nullptr,
+                         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)]);
@@ -736,11 +737,11 @@ public:
                                                                      params.model_dim));
     }
 
-    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) {
+    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) {
         // input_ids: [N, n_token]
 
         auto shared  = std::dynamic_pointer_cast<Embedding>(blocks["shared"]);
@@ -775,14 +776,14 @@ struct T5Runner : public GGMLRunner {
         return "t5";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* input_ids,
-                                struct ggml_tensor* relative_position_bucket,
-                                struct ggml_tensor* attention_mask = nullptr) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* input_ids,
+                         ggml_tensor* relative_position_bucket,
+                         ggml_tensor* attention_mask = nullptr) {
         size_t N       = input_ids->ne[1];
         size_t n_token = input_ids->ne[0];
 
@@ -790,9 +791,9 @@ struct T5Runner : public GGMLRunner {
         return hidden_states;
     }
 
-    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);
+    ggml_cgraph* build_graph(ggml_tensor* input_ids,
+                             ggml_tensor* attention_mask = nullptr) {
+        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
 
         input_ids      = to_backend(input_ids);
         attention_mask = to_backend(attention_mask);
@@ -812,8 +813,8 @@ 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();
-        struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
+        auto runner_ctx            = get_context();
+        ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
 
         ggml_build_forward_expand(gf, hidden_states);
 
@@ -821,11 +822,11 @@ struct T5Runner : public GGMLRunner {
     }
 
     bool compute(const int n_threads,
-                 struct ggml_tensor* input_ids,
-                 struct ggml_tensor* attention_mask,
+                 ggml_tensor* input_ids,
+                 ggml_tensor* attention_mask,
                  ggml_tensor** output,
                  ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
         return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
@@ -911,7 +912,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, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
@@ -961,17 +962,17 @@ struct T5Embedder {
     }
 
     void test() {
-        struct ggml_init_params params;
+        ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        struct ggml_context* work_ctx = ggml_init(params);
+        ggml_context* work_ctx = ggml_init(params);
         GGML_ASSERT(work_ctx != nullptr);
 
         {
             std::string text("a lovely cat");
-            // std::string text("一只可爱的猫"); // umt5 chinease test
+            // std::string text("一只可爱的<EFBFBD>?); // 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);
@@ -980,9 +981,9 @@ struct T5Embedder {
                 printf("%d ", token);
             }
             printf("\n");
-            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;
+            auto input_ids      = vector_to_ggml_tensor_i32(work_ctx, tokens);
+            auto attention_mask = vector_to_ggml_tensor(work_ctx, masks);
+            ggml_tensor* out    = nullptr;
 
             int64_t t0 = ggml_time_ms();
             model.compute(8, input_ids, attention_mask, &out, work_ctx);

src/tae.hpp

@@ -37,7 +37,7 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         // x: [n, n_in, h, w]
         // return: [n, n_out, h, w]
 
@@ -107,7 +107,7 @@ public:
         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 {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         // x: [n, in_channels, h, w]
         // return: [n, z_channels, h/8, w/8]
 
@@ -157,7 +157,7 @@ public:
         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 {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
         // z: [n, z_channels, h, w]
         // return: [n, out_channels, h*8, w*8]
 
@@ -192,7 +192,7 @@ public:
         blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * stride, channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
         auto h    = x;
         if (stride != 1) {
@@ -212,7 +212,7 @@ public:
         blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels * stride, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+    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) {
@@ -236,7 +236,7 @@ public:
         }
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* past) {
+    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"]);
@@ -260,10 +260,10 @@ public:
     }
 };
 
-struct ggml_tensor* patchify(struct ggml_context* ctx,
-                             struct ggml_tensor* x,
-                             int64_t patch_size,
-                             int64_t b = 1) {
+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) {
@@ -289,10 +289,10 @@ struct ggml_tensor* patchify(struct ggml_context* ctx,
     return x;
 }
 
-struct ggml_tensor* unpatchify(struct ggml_context* ctx,
-                               struct ggml_tensor* x,
-                               int64_t patch_size,
-                               int64_t b = 1) {
+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) {
@@ -339,7 +339,7 @@ public:
         blocks[std::to_string(index)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, z_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
         auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
 
         if (patch_size > 1) {
@@ -396,7 +396,7 @@ public:
         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}));
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) override {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
         auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
 
         // Clamp()
@@ -442,11 +442,13 @@ 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 z_channels = 16;
-        int patch      = 1;
+        int patch = 1;
         if (version == VERSION_WAN2_2_TI2V) {
             z_channels = 48;
             patch      = 2;
@@ -457,7 +459,7 @@ public:
         }
     }
 
-    struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
+    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)
@@ -471,7 +473,7 @@ public:
         return result;
     }
 
-    struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    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));
@@ -494,10 +496,12 @@ 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) {
-        int z_channels       = 4;
         bool use_midblock_gn = false;
         taef2                = sd_version_is_flux2(version);
 
@@ -515,7 +519,7 @@ public:
         }
     }
 
-    struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
+    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);
@@ -523,7 +527,7 @@ public:
         return decoder->forward(ctx, z);
     }
 
-    struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+    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) {
@@ -533,20 +537,7 @@ public:
     }
 };
 
-struct TinyAutoEncoder : public GGMLRunner {
-    TinyAutoEncoder(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 = nullptr) = 0;
-
-    virtual bool load_from_file(const std::string& file_path, int n_threads)                                      = 0;
-    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) = 0;
-};
-
-struct TinyImageAutoEncoder : public TinyAutoEncoder {
+struct TinyImageAutoEncoder : public VAE {
     TAESD taesd;
     bool decode_only = false;
 
@@ -558,7 +549,8 @@ struct TinyImageAutoEncoder : public TinyAutoEncoder {
                          SDVersion version = VERSION_SD1)
         : decode_only(decoder_only),
           taesd(decoder_only, version),
-          TinyAutoEncoder(backend, offload_params_to_cpu) {
+          VAE(version, backend, offload_params_to_cpu) {
+        scale_input = false;
         taesd.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -566,52 +558,41 @@ struct TinyImageAutoEncoder : public TinyAutoEncoder {
         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;
-    }
-
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         taesd.get_param_tensors(tensors, prefix);
     }
 
-    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_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+        return vae_output;
+    }
+
+    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    }
+
+    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    }
+
+    int get_encoder_output_channels(int input_channels) {
+        return taesd.z_channels;
+    }
+
+    ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
+        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        z                = to_backend(z);
+        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;
     }
 
-    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* {
+    bool _compute(const int n_threads,
+                  ggml_tensor* z,
+                  bool decode_graph,
+                  ggml_tensor** output,
+                  ggml_context* output_ctx = nullptr) {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
 
@@ -619,7 +600,7 @@ struct TinyImageAutoEncoder : public TinyAutoEncoder {
     }
 };
 
-struct TinyVideoAutoEncoder : public TinyAutoEncoder {
+struct TinyVideoAutoEncoder : public VAE {
     TAEHV taehv;
     bool decode_only = false;
 
@@ -631,7 +612,8 @@ struct TinyVideoAutoEncoder : public TinyAutoEncoder {
                          SDVersion version = VERSION_WAN2)
         : decode_only(decoder_only),
           taehv(decoder_only, version),
-          TinyAutoEncoder(backend, offload_params_to_cpu) {
+          VAE(version, backend, offload_params_to_cpu) {
+        scale_input = false;
         taehv.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -639,52 +621,41 @@ struct TinyVideoAutoEncoder : public TinyAutoEncoder {
         return "taehv";
     }
 
-    bool load_from_file(const std::string& file_path, int n_threads) {
-        LOG_INFO("loading taehv from '%s', decode_only = %s", file_path.c_str(), decode_only ? "true" : "false");
-        alloc_params_buffer();
-        std::map<std::string, ggml_tensor*> taehv_tensors;
-        taehv.get_param_tensors(taehv_tensors);
-        std::set<std::string> ignore_tensors;
-        if (decode_only) {
-            ignore_tensors.insert("encoder.");
-        }
-
-        ModelLoader model_loader;
-        if (!model_loader.init_from_file(file_path)) {
-            LOG_ERROR("init taehv model loader from file failed: '%s'", file_path.c_str());
-            return false;
-        }
-
-        bool success = model_loader.load_tensors(taehv_tensors, ignore_tensors, n_threads);
-
-        if (!success) {
-            LOG_ERROR("load tae tensors from model loader failed");
-            return false;
-        }
-
-        LOG_INFO("taehv model loaded");
-        return success;
-    }
-
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         taehv.get_param_tensors(tensors, prefix);
     }
 
-    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 ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
+    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+        return vae_output;
+    }
+
+    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    }
+
+    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    }
+
+    int get_encoder_output_channels(int input_channels) {
+        return taehv.z_channels;
+    }
+
+    ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
+        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        z                = to_backend(z);
+        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;
     }
 
-    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* {
+    bool _compute(const int n_threads,
+                  ggml_tensor* z,
+                  bool decode_graph,
+                  ggml_tensor** output,
+                  ggml_context* output_ctx = nullptr) {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
 

src/tokenize_util.cpp

@@ -919,15 +919,21 @@ std::vector<std::string> token_split(const std::string& text) {
 
         // `\s*[\r\n]+|\s+(?!\S)|\s+`
         if (is_space(cp)) {
-            std::string token = codepoint_to_utf8(cp);
-            ++i;
+            std::string token;
+            bool saw_new_line = false;
 
             while (i < cps.size() && is_space(cps[i])) {
                 token += codepoint_to_utf8(cps[i]);
-                ++i;
+
                 if (cps[i] == U'\r' || cps[i] == U'\n') {
-                    break;
+                    saw_new_line = true;
+                } else {
+                    if (saw_new_line) {
+                        break;
+                    }
                 }
+
+                ++i;
             }
 
             tokens.push_back(token);

src/ucache.hpp

@@ -19,6 +19,7 @@ struct UCacheConfig {
     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;
 };
 
@@ -45,14 +46,16 @@ struct UCacheState {
     bool has_output_prev_norm             = false;
     bool has_relative_transformation_rate = false;
     float relative_transformation_rate    = 0.0f;
-    float cumulative_change_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;
-    float reference_output_norm           = 0.0f;
 
     struct BlockMetrics {
         float sum_transformation_rate = 0.0f;
@@ -106,14 +109,16 @@ struct UCacheState {
         has_output_prev_norm             = false;
         has_relative_transformation_rate = false;
         relative_transformation_rate     = 0.0f;
-        cumulative_change_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;
-        reference_output_norm            = 0.0f;
         block_metrics.reset();
         total_active_steps = 0;
     }
@@ -133,7 +138,8 @@ struct UCacheState {
         if (!initialized || sigmas.size() < 2) {
             return;
         }
-        size_t n_steps = sigmas.size() - 1;
+        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);
@@ -207,11 +213,15 @@ struct UCacheState {
         }
 
         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) {
@@ -309,17 +319,31 @@ struct UCacheState {
 
         if (has_output_prev_norm && has_relative_transformation_rate &&
             last_input_change > 0.0f && output_prev_norm > 0.0f) {
-            float approx_output_change_rate = (relative_transformation_rate * last_input_change) / output_prev_norm;
-            accumulated_error               = accumulated_error * config.error_decay_rate + approx_output_change_rate;
+            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 && reference_output_norm > 0.0f) {
-                effective_threshold = effective_threshold * reference_output_norm;
+            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) {
@@ -340,6 +364,8 @@ struct UCacheState {
         if (cond != anchor_condition) {
             return;
         }
+        steps_computed_since_active++;
+        consecutive_skipped_steps = 0;
 
         size_t ne      = static_cast<size_t>(ggml_nelements(input));
         float* in_data = (float*)input->data;
@@ -359,6 +385,14 @@ struct UCacheState {
                 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) {
@@ -373,10 +407,6 @@ struct UCacheState {
         output_prev_norm     = (ne > 0) ? (mean_abs / static_cast<float>(ne)) : 0.0f;
         has_output_prev_norm = output_prev_norm > 0.0f;
 
-        if (reference_output_norm == 0.0f) {
-            reference_output_norm = output_prev_norm;
-        }
-
         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)) {

src/unet.hpp

@@ -1,8 +1,7 @@
 #ifndef __UNET_HPP__
 #define __UNET_HPP__
 
-#include "common.hpp"
-#include "ggml_extend.hpp"
+#include "common_block.hpp"
 #include "model.h"
 
 /*==================================================== UnetModel =====================================================*/
@@ -61,10 +60,10 @@ public:
         blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
     }
 
-    struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                struct ggml_tensor* x,
-                                struct ggml_tensor* context,
-                                int timesteps) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         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]
@@ -389,11 +388,11 @@ public:
         blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
     }
 
-    struct ggml_tensor* resblock_forward(std::string name,
-                                         GGMLRunnerContext* ctx,
-                                         struct ggml_tensor* x,
-                                         struct ggml_tensor* emb,
-                                         int num_video_frames) {
+    ggml_tensor* resblock_forward(std::string name,
+                                  GGMLRunnerContext* ctx,
+                                  ggml_tensor* x,
+                                  ggml_tensor* emb,
+                                  int num_video_frames) {
         if (version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
 
@@ -405,11 +404,11 @@ public:
         }
     }
 
-    struct ggml_tensor* attention_layer_forward(std::string name,
-                                                GGMLRunnerContext* ctx,
-                                                struct ggml_tensor* x,
-                                                struct ggml_tensor* context,
-                                                int timesteps) {
+    ggml_tensor* attention_layer_forward(std::string name,
+                                         GGMLRunnerContext* ctx,
+                                         ggml_tensor* x,
+                                         ggml_tensor* context,
+                                         int timesteps) {
         if (version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
 
@@ -421,15 +420,15 @@ public:
         }
     }
 
-    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) {
+    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) {
         // 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]
@@ -481,7 +480,7 @@ public:
         }
 
         // input_blocks
-        std::vector<struct ggml_tensor*> hs;
+        std::vector<ggml_tensor*> hs;
 
         // input block 0
         auto h = input_blocks_0_0->forward(ctx, x);
@@ -606,19 +605,19 @@ struct UNetModelRunner : public GGMLRunner {
         return "unet";
     }
 
-    void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         unet.get_param_tensors(tensors, prefix);
     }
 
-    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);
+    ggml_cgraph* build_graph(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) {
+        ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
 
         if (num_video_frames == -1) {
             num_video_frames = static_cast<int>(x->ne[3]);
@@ -636,15 +635,15 @@ struct UNetModelRunner : public GGMLRunner {
 
         auto runner_ctx = get_context();
 
-        struct ggml_tensor* out = unet.forward(&runner_ctx,
-                                               x,
-                                               timesteps,
-                                               context,
-                                               c_concat,
-                                               y,
-                                               num_video_frames,
-                                               controls,
-                                               control_strength);
+        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);
 
@@ -652,22 +651,22 @@ struct UNetModelRunner : public GGMLRunner {
     }
 
     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) {
+                 ggml_tensor* x,
+                 ggml_tensor* timesteps,
+                 ggml_tensor* context,
+                 ggml_tensor* c_concat,
+                 ggml_tensor* y,
+                 int num_video_frames               = -1,
+                 std::vector<ggml_tensor*> controls = {},
+                 float control_strength             = 0.f,
+                 ggml_tensor** output               = nullptr,
+                 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 = [&]() -> struct ggml_cgraph* {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
         };
 
@@ -675,12 +674,12 @@ struct UNetModelRunner : public GGMLRunner {
     }
 
     void test() {
-        struct ggml_init_params params;
+        ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
         params.mem_buffer = nullptr;
         params.no_alloc   = false;
 
-        struct ggml_context* work_ctx = ggml_init(params);
+        ggml_context* work_ctx = ggml_init(params);
         GGML_ASSERT(work_ctx != nullptr);
 
         {
@@ -704,7 +703,7 @@ struct UNetModelRunner : public GGMLRunner {
             ggml_set_f32(y, 0.5f);
             // print_ggml_tensor(y);
 
-            struct ggml_tensor* out = nullptr;
+            ggml_tensor* out = nullptr;
 
             int64_t t0 = ggml_time_ms();
             compute(8, x, timesteps, context, nullptr, y, num_video_frames, {}, 0.f, &out, work_ctx);

src/upscaler.cpp

@@ -72,13 +72,13 @@ struct UpscalerGGML {
         LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
                  input_image.width, input_image.height, output_width, output_height);
 
-        struct ggml_init_params params;
+        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);
+        ggml_context* upscale_ctx = ggml_init(params);
         if (!upscale_ctx) {
             LOG_ERROR("ggml_init() failed");
             return upscaled_image;
@@ -89,10 +89,11 @@ struct UpscalerGGML {
 
         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);
+            return 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);
+        // TODO: circular upscaling?
+        sd_tiling(input_image_tensor, upscaled, esrgan_upscaler->scale, esrgan_upscaler->tile_size, 0.25f, false, false, 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);

src/vae.hpp

@@ -0,0 +1,235 @@
+#ifndef __VAE_HPP__
+#define __VAE_HPP__
+
+#include "common_block.hpp"
+
+struct VAE : public GGMLRunner {
+protected:
+    SDVersion version;
+    bool scale_input                                = true;
+    virtual bool _compute(const int n_threads,
+                          ggml_tensor* z,
+                          bool decode_graph,
+                          ggml_tensor** output,
+                          ggml_context* output_ctx) = 0;
+
+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);
+    }
+
+    ggml_tensor* encode(int n_threads,
+                        ggml_context* work_ctx,
+                        ggml_tensor* x,
+                        sd_tiling_params_t tiling_params,
+                        bool circular_x = false,
+                        bool circular_y = false) {
+        int64_t t0             = ggml_time_ms();
+        ggml_tensor* result    = nullptr;
+        const int scale_factor = get_scale_factor();
+        int64_t W              = x->ne[0] / scale_factor;
+        int64_t H              = x->ne[1] / scale_factor;
+        int channel_dim        = sd_version_is_wan(version) ? 3 : 2;
+        int64_t C              = get_encoder_output_channels(static_cast<int>(x->ne[channel_dim]));
+        int64_t ne2;
+        int64_t ne3;
+        if (sd_version_is_wan(version)) {
+            int64_t T = x->ne[2];
+            ne2       = (T - 1) / 4 + 1;
+            ne3       = C;
+        } else {
+            ne2 = C;
+            ne3 = x->ne[3];
+        }
+        result = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, W, H, ne2, ne3);
+
+        if (scale_input) {
+            scale_to_minus1_1(x);
+        }
+
+        if (sd_version_is_qwen_image(version) || sd_version_is_anima(version)) {
+            x = ggml_reshape_4d(work_ctx, x, x->ne[0], x->ne[1], 1, x->ne[2] * x->ne[3]);
+        }
+
+        if (tiling_params.enabled) {
+            float tile_overlap;
+            int tile_size_x, tile_size_y;
+            // multiply tile size for encode to keep the compute buffer size consistent
+            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);
+
+            auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
+                return _compute(n_threads, in, false, &out, work_ctx);
+            };
+            sd_tiling_non_square(x, result, scale_factor, tile_size_x, tile_size_y, tile_overlap, circular_x, circular_y, on_tiling);
+        } else {
+            _compute(n_threads, x, false, &result, work_ctx);
+        }
+        free_compute_buffer();
+
+        int64_t t1 = ggml_time_ms();
+        LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
+        return result;
+    }
+
+    ggml_tensor* decode(int n_threads,
+                        ggml_context* work_ctx,
+                        ggml_tensor* x,
+                        sd_tiling_params_t tiling_params,
+                        bool decode_video   = false,
+                        bool circular_x     = false,
+                        bool circular_y     = false,
+                        ggml_tensor* result = nullptr,
+                        bool silent         = false) {
+        const int scale_factor = get_scale_factor();
+        int64_t W              = x->ne[0] * scale_factor;
+        int64_t H              = x->ne[1] * scale_factor;
+        int64_t C              = 3;
+        if (result == nullptr) {
+            if (decode_video) {
+                int64_t T = x->ne[2];
+                if (sd_version_is_wan(version)) {
+                    T = ((T - 1) * 4) + 1;
+                }
+                result = ggml_new_tensor_4d(work_ctx,
+                                            GGML_TYPE_F32,
+                                            W,
+                                            H,
+                                            T,
+                                            3);
+            } else {
+                result = ggml_new_tensor_4d(work_ctx,
+                                            GGML_TYPE_F32,
+                                            W,
+                                            H,
+                                            C,
+                                            x->ne[3]);
+            }
+        }
+        int64_t t0 = ggml_time_ms();
+        if (sd_version_is_qwen_image(version) || sd_version_is_anima(version)) {
+            x = ggml_reshape_4d(work_ctx, x, x->ne[0], x->ne[1], 1, x->ne[2] * x->ne[3]);
+        }
+        if (tiling_params.enabled) {
+            float tile_overlap;
+            int tile_size_x, tile_size_y;
+            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, x->ne[0], x->ne[1]);
+
+            if (!silent) {
+                LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
+            }
+
+            auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
+                return _compute(n_threads, in, true, &out, nullptr);
+            };
+            sd_tiling_non_square(x, result, scale_factor, tile_size_x, tile_size_y, tile_overlap, circular_x, circular_y, on_tiling, silent);
+        } else {
+            if (!_compute(n_threads, x, true, &result, work_ctx)) {
+                LOG_ERROR("Failed to decode latetnts");
+                free_compute_buffer();
+                return nullptr;
+            }
+        }
+        free_compute_buffer();
+        if (scale_input) {
+            scale_to_0_1(result);
+        }
+        int64_t t1 = ggml_time_ms();
+        LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
+        ggml_ext_tensor_clamp_inplace(result, 0.0f, 1.0f);
+        return result;
+    }
+
+    virtual ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) = 0;
+    virtual ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents)                           = 0;
+    virtual ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* 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;
+    }
+
+    bool _compute(const int n_threads,
+                  ggml_tensor* z,
+                  bool decode_graph,
+                  ggml_tensor** output,
+                  ggml_context* output_ctx) override {
+        if (*output == nullptr && output_ctx != nullptr) {
+            *output = ggml_dup_tensor(output_ctx, z);
+        }
+        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;
+    }
+
+    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+        return vae_output;
+    }
+
+    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    }
+
+    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+        return ggml_ext_dup_and_cpy_tensor(work_ctx, 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/clip_t5.hpp

@@ -1,4 +1,4 @@
-static unsigned char merges_utf8_c_str[] = {
+static const unsigned char clip_merges_utf8_c_str[] = {
     0x23,
     0x76,
     0x65,
@@ -524620,7 +524620,7 @@ static unsigned char merges_utf8_c_str[] = {
     0x0a,
 };
 
-static unsigned char t5_tokenizer_json_str[] = {
+static const unsigned char t5_tokenizer_json_str[] = {
     0x7b,
     0x0a,
     0x20,

src/vocab/mistral.hpp

@@ -1,4 +1,4 @@
-unsigned char mistral_merges_utf8_c_str[] = {
+static const unsigned char mistral_merges_utf8_c_str[] = {
   0xc4, 0xa0, 0x20, 0xc4, 0xa0, 0x0a, 0xc4, 0xa0, 0x20, 0x74, 0x0a, 0x65,
   0x20, 0x72, 0x0a, 0x69, 0x20, 0x6e, 0x0a, 0xc4, 0xa0, 0x20, 0xc4, 0xa0,
   0xc4, 0xa0, 0xc4, 0xa0, 0x0a, 0xc4, 0xa0, 0xc4, 0xa0, 0x20, 0xc4, 0xa0,
@@ -260614,7 +260614,7 @@ unsigned char mistral_merges_utf8_c_str[] = {
   0xc3, 0xa5, 0xc4, 0xb2, 0xc4, 0xb0, 0x20, 0xc3, 0xa6, 0xc2, 0xb1, 0xc4,
   0xab, 0xc3, 0xa4, 0xc2, 0xb9, 0xc2, 0xa6, 0x0a,
 };
-unsigned char mistral_vocab_json_utf8_c_str[] = {
+static const unsigned char mistral_vocab_json_utf8_c_str[] = {
   0x7b, 0x22, 0x3c, 0x75, 0x6e, 0x6b, 0x3e, 0x22, 0x3a, 0x20, 0x30, 0x2c,
   0x20, 0x22, 0x3c, 0x73, 0x3e, 0x22, 0x3a, 0x20, 0x31, 0x2c, 0x20, 0x22,
   0x3c, 0x2f, 0x73, 0x3e, 0x22, 0x3a, 0x20, 0x32, 0x2c, 0x20, 0x22, 0x5b,

src/vocab/qwen.hpp

@@ -1,4 +1,4 @@
-unsigned char qwen2_merges_utf8_c_str[] = {
+static const unsigned char qwen2_merges_utf8_c_str[] = {
   0xc4, 0xa0, 0x20, 0xc4, 0xa0, 0x0a, 0xc4, 0xa0, 0xc4, 0xa0, 0x20, 0xc4,
   0xa0, 0xc4, 0xa0, 0x0a, 0x69, 0x20, 0x6e, 0x0a, 0xc4, 0xa0, 0x20, 0x74,
   0x0a, 0xc4, 0xa0, 0xc4, 0xa0, 0xc4, 0xa0, 0xc4, 0xa0, 0x20, 0xc4, 0xa0,

src/vocab/umt5.hpp

@@ -1,4 +1,4 @@
-unsigned char umt5_tokenizer_json_str[] = {
+static const unsigned char umt5_tokenizer_json_str[] = {
   0x7b, 0x22, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x22, 0x3a, 0x20,
   0x22, 0x31, 0x2e, 0x30, 0x22, 0x2c, 0x20, 0x22, 0x74, 0x72, 0x75, 0x6e,
   0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x22, 0x3a, 0x20, 0x6e, 0x75, 0x6c,

src/vocab/vocab.cpp

@@ -0,0 +1,35 @@
+#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

@@ -0,0 +1,13 @@
+#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__

src/wan.hpp

@@ -5,9 +5,8 @@
 #include <memory>
 #include <utility>
 
-#include "common.hpp"
+#include "common_block.hpp"
 #include "flux.hpp"
-#include "ggml_extend.hpp"
 #include "rope.hpp"
 #include "vae.hpp"
 
@@ -26,7 +25,7 @@ namespace WAN {
         std::tuple<int, int, int> dilation;
         bool bias;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             params["weight"] = ggml_new_tensor_4d(ctx,
                                                   GGML_TYPE_F16,
                                                   std::get<2>(kernel_size),
@@ -54,11 +53,11 @@ namespace WAN {
               dilation(std::move(dilation)),
               bias(bias) {}
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* cache_x = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) {
             // x: [N*IC, ID, IH, IW]
             // result: x: [N*OC, ID, IH, IW]
-            struct ggml_tensor* w = params["weight"];
-            struct ggml_tensor* b = nullptr;
+            ggml_tensor* w = params["weight"];
+            ggml_tensor* b = nullptr;
             if (bias) {
                 b = params["bias"];
             }
@@ -87,7 +86,7 @@ namespace WAN {
     protected:
         int64_t dim;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             ggml_type wtype = GGML_TYPE_F32;
             auto iter       = tensor_storage_map.find(prefix + "gamma");
             if (iter != tensor_storage_map.end()) {
@@ -101,16 +100,16 @@ namespace WAN {
         RMS_norm(int64_t dim)
             : dim(dim) {}
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             // x: [N*IC, ID, IH, IW], IC == dim
             // assert N == 1
 
-            struct ggml_tensor* w = params["gamma"];
-            w                     = ggml_reshape_1d(ctx->ggml_ctx, w, ggml_nelements(w));
-            auto h                = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2));  // [ID, IH, IW, N*IC]
-            h                     = ggml_rms_norm(ctx->ggml_ctx, h, 1e-12f);
-            h                     = ggml_mul(ctx->ggml_ctx, h, w);
-            h                     = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
+            ggml_tensor* w = params["gamma"];
+            w              = ggml_reshape_1d(ctx->ggml_ctx, w, ggml_nelements(w));
+            auto h         = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2));  // [ID, IH, IW, N*IC]
+            h              = ggml_rms_norm(ctx->ggml_ctx, h, 1e-12f);
+            h              = ggml_mul(ctx->ggml_ctx, h, w);
+            h              = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
 
             return h;
         }
@@ -149,12 +148,12 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b,
-                                    std::vector<struct ggml_tensor*>& feat_cache,
-                                    int& feat_idx,
-                                    int chunk_idx) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b,
+                             std::vector<ggml_tensor*>& feat_cache,
+                             int& feat_idx,
+                             int chunk_idx) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
             int64_t c = x->ne[3] / b;
@@ -255,9 +254,9 @@ namespace WAN {
             GGML_ASSERT(in_channels * factor % out_channels == 0);
             group_size = in_channels * factor / out_channels;
         }
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t B = 1) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t B = 1) {
             // x: [B*IC, T, H, W]
             // return: [B*OC, T/factor_t, H/factor_s, W/factor_s]
             GGML_ASSERT(B == 1);
@@ -302,10 +301,10 @@ namespace WAN {
             GGML_ASSERT(out_channels * factor % in_channels == 0);
             repeats = out_channels * factor / in_channels;
         }
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    bool first_chunk = false,
-                                    int64_t B        = 1) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             bool first_chunk = false,
+                             int64_t B        = 1) {
             // x: [B*IC, T, H, W]
             // return: [B*OC, T/factor_t, H/factor_s, W/factor_s]
             GGML_ASSERT(B == 1);
@@ -357,14 +356,14 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b,
-                                    std::vector<struct ggml_tensor*>& feat_cache,
-                                    int& feat_idx) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b,
+                             std::vector<ggml_tensor*>& feat_cache,
+                             int& feat_idx) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
-            struct ggml_tensor* h = x;
+            ggml_tensor* h = x;
             if (in_dim != out_dim) {
                 auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
 
@@ -431,15 +430,15 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b,
-                                    std::vector<struct ggml_tensor*>& feat_cache,
-                                    int& feat_idx,
-                                    int chunk_idx) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b,
+                             std::vector<ggml_tensor*>& feat_cache,
+                             int& feat_idx,
+                             int chunk_idx) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
-            struct ggml_tensor* x_copy = x;
+            ggml_tensor* x_copy = x;
 
             auto avg_shortcut = std::dynamic_pointer_cast<AvgDown3D>(blocks["avg_shortcut"]);
 
@@ -493,15 +492,15 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b,
-                                    std::vector<struct ggml_tensor*>& feat_cache,
-                                    int& feat_idx,
-                                    int chunk_idx) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b,
+                             std::vector<ggml_tensor*>& feat_cache,
+                             int& feat_idx,
+                             int chunk_idx) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
-            struct ggml_tensor* x_copy = x;
+            ggml_tensor* x_copy = x;
 
             int i = 0;
             for (; i < mult; i++) {
@@ -538,9 +537,9 @@ namespace WAN {
             blocks["proj"]   = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {1, 1}));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
             auto norm   = std::dynamic_pointer_cast<RMS_norm>(blocks["norm"]);
@@ -572,8 +571,8 @@ namespace WAN {
             auto v = qkv_vec[2];
             v      = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n);  // [t, c, h * w]
 
-            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));                           // [t, h * w, c]
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, true, ctx->flash_attn_enabled);  // [t, h * w, c]
+            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));                            // [t, h * w, c]
+            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [t, h * w, c]
 
             x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [t, c, h * w]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n);                             // [t, c, h, w]
@@ -660,12 +659,12 @@ namespace WAN {
             blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b,
-                                    std::vector<struct ggml_tensor*>& feat_cache,
-                                    int& feat_idx,
-                                    int chunk_idx) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b,
+                             std::vector<ggml_tensor*>& feat_cache,
+                             int& feat_idx,
+                             int chunk_idx) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
             auto conv1    = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
@@ -831,12 +830,12 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    int64_t b,
-                                    std::vector<struct ggml_tensor*>& feat_cache,
-                                    int& feat_idx,
-                                    int chunk_idx) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t b,
+                             std::vector<ggml_tensor*>& feat_cache,
+                             int& feat_idx,
+                             int chunk_idx) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
             auto conv1    = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
@@ -935,16 +934,16 @@ namespace WAN {
 
         int _conv_num = 33;
         int _conv_idx = 0;
-        std::vector<struct ggml_tensor*> _feat_map;
+        std::vector<ggml_tensor*> _feat_map;
         int _enc_conv_num = 28;
         int _enc_conv_idx = 0;
-        std::vector<struct ggml_tensor*> _enc_feat_map;
+        std::vector<ggml_tensor*> _enc_feat_map;
 
         void clear_cache() {
             _conv_idx     = 0;
-            _feat_map     = std::vector<struct ggml_tensor*>(_conv_num, nullptr);
+            _feat_map     = std::vector<ggml_tensor*>(_conv_num, nullptr);
             _enc_conv_idx = 0;
-            _enc_feat_map = std::vector<struct ggml_tensor*>(_enc_conv_num, nullptr);
+            _enc_feat_map = std::vector<ggml_tensor*>(_enc_conv_num, nullptr);
         }
 
     public:
@@ -967,10 +966,10 @@ namespace WAN {
             blocks["conv2"]   = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
         }
 
-        struct ggml_tensor* patchify(struct ggml_context* ctx,
-                                     struct ggml_tensor* x,
-                                     int64_t patch_size,
-                                     int64_t b = 1) {
+        ggml_tensor* patchify(ggml_context* ctx,
+                              ggml_tensor* x,
+                              int64_t patch_size,
+                              int64_t b = 1) {
             // x: [b*c, f, h*q, w*r]
             // return: [b*c*r*q, f, h, w]
             if (patch_size == 1) {
@@ -994,10 +993,10 @@ namespace WAN {
             return x;
         }
 
-        struct ggml_tensor* unpatchify(struct ggml_context* ctx,
-                                       struct ggml_tensor* x,
-                                       int64_t patch_size,
-                                       int64_t b = 1) {
+        ggml_tensor* unpatchify(ggml_context* ctx,
+                                ggml_tensor* x,
+                                int64_t patch_size,
+                                int64_t b = 1) {
             // x: [b*c*r*q, f, h, w]
             // return: [b*c, f, h*q, w*r]
             if (patch_size == 1) {
@@ -1020,9 +1019,9 @@ namespace WAN {
             return x;
         }
 
-        struct ggml_tensor* encode(GGMLRunnerContext* ctx,
-                                   struct ggml_tensor* x,
-                                   int64_t b = 1) {
+        ggml_tensor* encode(GGMLRunnerContext* ctx,
+                            ggml_tensor* x,
+                            int64_t b = 1) {
             // x: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
             GGML_ASSERT(decode_only == false);
@@ -1038,7 +1037,7 @@ namespace WAN {
 
             int64_t t     = x->ne[2];
             int64_t iter_ = 1 + (t - 1) / 4;
-            struct ggml_tensor* out;
+            ggml_tensor* out;
             for (int i = 0; i < iter_; i++) {
                 _enc_conv_idx = 0;
                 if (i == 0) {
@@ -1056,9 +1055,9 @@ namespace WAN {
             return mu;
         }
 
-        struct ggml_tensor* decode(GGMLRunnerContext* ctx,
-                                   struct ggml_tensor* z,
-                                   int64_t b = 1) {
+        ggml_tensor* decode(GGMLRunnerContext* ctx,
+                            ggml_tensor* z,
+                            int64_t b = 1) {
             // z: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
 
@@ -1069,7 +1068,7 @@ namespace WAN {
 
             int64_t iter_ = z->ne[2];
             auto x        = conv2->forward(ctx, z);
-            struct ggml_tensor* out;
+            ggml_tensor* out;
             for (int i = 0; i < iter_; i++) {
                 _conv_idx = 0;
                 if (i == 0) {
@@ -1088,10 +1087,10 @@ namespace WAN {
             return out;
         }
 
-        struct ggml_tensor* decode_partial(GGMLRunnerContext* ctx,
-                                           struct ggml_tensor* z,
-                                           int i,
-                                           int64_t b = 1) {
+        ggml_tensor* decode_partial(GGMLRunnerContext* ctx,
+                                    ggml_tensor* z,
+                                    int i,
+                                    int64_t b = 1) {
             // z: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
 
@@ -1110,7 +1109,8 @@ namespace WAN {
     };
 
     struct WanVAERunner : public VAE {
-        bool decode_only = true;
+        float scale_factor = 1.0f;
+        bool decode_only   = true;
         WanVAE ae;
 
         WanVAERunner(ggml_backend_t backend,
@@ -1119,7 +1119,7 @@ namespace WAN {
                      const std::string prefix                       = "",
                      bool decode_only                               = false,
                      SDVersion version                              = VERSION_WAN2)
-            : decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V), VAE(backend, offload_params_to_cpu) {
+            : decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V), VAE(version, backend, offload_params_to_cpu) {
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -1127,26 +1127,121 @@ namespace WAN {
             return "wan_vae";
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
             ae.get_param_tensors(tensors, prefix);
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
-            struct ggml_cgraph* gf = new_graph_custom(10240 * z->ne[2]);
+        ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+            return vae_output;
+        }
+
+        void get_latents_mean_std_vec(ggml_tensor* latents, int channel_dim, std::vector<float>& latents_mean_vec, std::vector<float>& latents_std_vec) {
+            GGML_ASSERT(latents->ne[channel_dim] == 16 || latents->ne[channel_dim] == 48);
+            if (latents->ne[channel_dim] == 16) {  // Wan2.1 VAE
+                latents_mean_vec = {-0.7571f, -0.7089f, -0.9113f, 0.1075f, -0.1745f, 0.9653f, -0.1517f, 1.5508f,
+                                    0.4134f, -0.0715f, 0.5517f, -0.3632f, -0.1922f, -0.9497f, 0.2503f, -0.2921f};
+                latents_std_vec  = {2.8184f, 1.4541f, 2.3275f, 2.6558f, 1.2196f, 1.7708f, 2.6052f, 2.0743f,
+                                    3.2687f, 2.1526f, 2.8652f, 1.5579f, 1.6382f, 1.1253f, 2.8251f, 1.9160f};
+            } else if (latents->ne[channel_dim] == 48) {  // Wan2.2 VAE
+                latents_mean_vec = {-0.2289f, -0.0052f, -0.1323f, -0.2339f, -0.2799f, 0.0174f, 0.1838f, 0.1557f,
+                                    -0.1382f, 0.0542f, 0.2813f, 0.0891f, 0.1570f, -0.0098f, 0.0375f, -0.1825f,
+                                    -0.2246f, -0.1207f, -0.0698f, 0.5109f, 0.2665f, -0.2108f, -0.2158f, 0.2502f,
+                                    -0.2055f, -0.0322f, 0.1109f, 0.1567f, -0.0729f, 0.0899f, -0.2799f, -0.1230f,
+                                    -0.0313f, -0.1649f, 0.0117f, 0.0723f, -0.2839f, -0.2083f, -0.0520f, 0.3748f,
+                                    0.0152f, 0.1957f, 0.1433f, -0.2944f, 0.3573f, -0.0548f, -0.1681f, -0.0667f};
+                latents_std_vec  = {
+                     0.4765f, 1.0364f, 0.4514f, 1.1677f, 0.5313f, 0.4990f, 0.4818f, 0.5013f,
+                     0.8158f, 1.0344f, 0.5894f, 1.0901f, 0.6885f, 0.6165f, 0.8454f, 0.4978f,
+                     0.5759f, 0.3523f, 0.7135f, 0.6804f, 0.5833f, 1.4146f, 0.8986f, 0.5659f,
+                     0.7069f, 0.5338f, 0.4889f, 0.4917f, 0.4069f, 0.4999f, 0.6866f, 0.4093f,
+                     0.5709f, 0.6065f, 0.6415f, 0.4944f, 0.5726f, 1.2042f, 0.5458f, 1.6887f,
+                     0.3971f, 1.0600f, 0.3943f, 0.5537f, 0.5444f, 0.4089f, 0.7468f, 0.7744f};
+            }
+        }
+
+        ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+            ggml_tensor* vae_latents = ggml_dup(work_ctx, latents);
+            int channel_dim          = sd_version_is_wan(version) ? 3 : 2;
+            std::vector<float> latents_mean_vec;
+            std::vector<float> latents_std_vec;
+            get_latents_mean_std_vec(latents, channel_dim, latents_mean_vec, latents_std_vec);
+
+            float mean;
+            float std_;
+            for (int i = 0; i < latents->ne[3]; i++) {
+                if (channel_dim == 3) {
+                    mean = latents_mean_vec[i];
+                    std_ = latents_std_vec[i];
+                }
+                for (int j = 0; j < latents->ne[2]; j++) {
+                    if (channel_dim == 2) {
+                        mean = latents_mean_vec[j];
+                        std_ = latents_std_vec[j];
+                    }
+                    for (int k = 0; k < latents->ne[1]; k++) {
+                        for (int l = 0; l < latents->ne[0]; l++) {
+                            float value = ggml_ext_tensor_get_f32(latents, l, k, j, i);
+                            value       = value * std_ / scale_factor + mean;
+                            ggml_ext_tensor_set_f32(vae_latents, value, l, k, j, i);
+                        }
+                    }
+                }
+            }
+
+            return vae_latents;
+        }
+
+        ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
+            ggml_tensor* diffusion_latents = ggml_dup(work_ctx, latents);
+            int channel_dim                = sd_version_is_wan(version) ? 3 : 2;
+            std::vector<float> latents_mean_vec;
+            std::vector<float> latents_std_vec;
+            get_latents_mean_std_vec(latents, channel_dim, latents_mean_vec, latents_std_vec);
+
+            float mean;
+            float std_;
+            for (int i = 0; i < latents->ne[3]; i++) {
+                if (channel_dim == 3) {
+                    mean = latents_mean_vec[i];
+                    std_ = latents_std_vec[i];
+                }
+                for (int j = 0; j < latents->ne[2]; j++) {
+                    if (channel_dim == 2) {
+                        mean = latents_mean_vec[j];
+                        std_ = latents_std_vec[j];
+                    }
+                    for (int k = 0; k < latents->ne[1]; k++) {
+                        for (int l = 0; l < latents->ne[0]; l++) {
+                            float value = ggml_ext_tensor_get_f32(latents, l, k, j, i);
+                            value       = (value - mean) * scale_factor / std_;
+                            ggml_ext_tensor_set_f32(diffusion_latents, value, l, k, j, i);
+                        }
+                    }
+                }
+            }
+            return diffusion_latents;
+        }
+
+        int get_encoder_output_channels(int input_channels) {
+            return static_cast<int>(ae.z_dim);
+        }
+
+        ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
+            ggml_cgraph* gf = new_graph_custom(10240 * z->ne[2]);
 
             z = to_backend(z);
 
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
+            ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
 
             ggml_build_forward_expand(gf, out);
 
             return gf;
         }
 
-        struct ggml_cgraph* build_graph_partial(struct ggml_tensor* z, bool decode_graph, int i) {
-            struct ggml_cgraph* gf = new_graph_custom(20480);
+        ggml_cgraph* build_graph_partial(ggml_tensor* z, bool decode_graph, int i) {
+            ggml_cgraph* gf = new_graph_custom(20480);
 
             ae.clear_cache();
 
@@ -1159,7 +1254,7 @@ namespace WAN {
 
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* out = decode_graph ? ae.decode_partial(&runner_ctx, z, i) : ae.encode(&runner_ctx, z);
+            ggml_tensor* out = decode_graph ? ae.decode_partial(&runner_ctx, z, i) : ae.encode(&runner_ctx, z);
 
             for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
                 ggml_tensor* feat_cache = ae._feat_map[feat_idx];
@@ -1174,13 +1269,13 @@ namespace WAN {
             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) override {
+        bool _compute(const int n_threads,
+                      ggml_tensor* z,
+                      bool decode_graph,
+                      ggml_tensor** output,
+                      ggml_context* output_ctx = nullptr) override {
             if (true) {
-                auto get_graph = [&]() -> struct ggml_cgraph* {
+                auto get_graph = [&]() -> ggml_cgraph* {
                     return build_graph(z, decode_graph);
                 };
                 return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
@@ -1188,11 +1283,11 @@ namespace WAN {
                 ae.clear_cache();
                 int64_t t      = z->ne[2];
                 int i          = 0;
-                auto get_graph = [&]() -> struct ggml_cgraph* {
+                auto get_graph = [&]() -> ggml_cgraph* {
                     return build_graph_partial(z, decode_graph, i);
                 };
-                struct ggml_tensor* out = nullptr;
-                bool res                = GGMLRunner::compute(get_graph, n_threads, true, &out, output_ctx);
+                ggml_tensor* out = nullptr;
+                bool res         = GGMLRunner::compute(get_graph, n_threads, true, &out, output_ctx);
                 ae.clear_cache();
                 if (t == 1) {
                     *output = out;
@@ -1230,12 +1325,12 @@ namespace WAN {
         }
 
         void test() {
-            struct ggml_init_params params;
+            ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            struct ggml_context* work_ctx = ggml_init(params);
+            ggml_context* work_ctx = ggml_init(params);
             GGML_ASSERT(work_ctx != nullptr);
 
             if (true) {
@@ -1247,10 +1342,10 @@ namespace WAN {
                 ggml_set_f32(z, 0.5f);
                 z = load_tensor_from_file(work_ctx, "wan_vae_z.bin");
                 print_ggml_tensor(z);
-                struct ggml_tensor* out = nullptr;
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
-                compute(8, z, true, &out, work_ctx);
+                _compute(8, z, true, &out, work_ctx);
                 int64_t t1 = ggml_time_ms();
 
                 print_ggml_tensor(out);
@@ -1315,10 +1410,10 @@ namespace WAN {
             }
         }
 
-        virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                            struct ggml_tensor* x,
-                                            struct ggml_tensor* pe,
-                                            struct ggml_tensor* mask = nullptr) {
+        virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                     ggml_tensor* x,
+                                     ggml_tensor* pe,
+                                     ggml_tensor* mask = nullptr) {
             // x: [N, n_token, dim]
             // pe: [n_token, d_head/2, 2, 2]
             // return [N, n_token, dim]
@@ -1356,10 +1451,10 @@ namespace WAN {
                           bool qk_norm = true,
                           float eps    = 1e-6)
             : WanSelfAttention(dim, num_heads, qk_norm, eps) {}
-        virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                            struct ggml_tensor* x,
-                                            struct ggml_tensor* context,
-                                            int64_t context_img_len) = 0;
+        virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                     ggml_tensor* x,
+                                     ggml_tensor* context,
+                                     int64_t context_img_len) = 0;
     };
 
     class WanT2VCrossAttention : public WanCrossAttention {
@@ -1369,10 +1464,10 @@ namespace WAN {
                              bool qk_norm = true,
                              float eps    = 1e-6)
             : WanCrossAttention(dim, num_heads, qk_norm, eps) {}
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* context,
-                                    int64_t context_img_len) override {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* context,
+                             int64_t context_img_len) override {
             // x: [N, n_token, dim]
             // context: [N, n_context, dim]
             // context_img_len: unused
@@ -1417,10 +1512,10 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* context,
-                                    int64_t context_img_len) override {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* context,
+                             int64_t context_img_len) override {
             // x: [N, n_token, dim]
             // context: [N, context_img_len + context_txt_len, dim]
             // return [N, n_token, dim]
@@ -1465,7 +1560,7 @@ namespace WAN {
         }
     };
 
-    static struct ggml_tensor* modulate_add(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* e) {
+    static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
         // x: [N, n_token, dim]
         // e: [N, 1, dim] or [N, T, 1, dim]
         if (ggml_n_dims(e) == 3) {
@@ -1479,7 +1574,7 @@ namespace WAN {
         return x;
     }
 
-    static struct ggml_tensor* modulate_mul(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* e) {
+    static ggml_tensor* modulate_mul(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
         // x: [N, n_token, dim]
         // e: [N, 1, dim] or [N, T, 1, dim]
         if (ggml_n_dims(e) == 3) {
@@ -1497,7 +1592,7 @@ namespace WAN {
     protected:
         int64_t dim;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(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);
             params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
         }
@@ -1531,12 +1626,12 @@ namespace WAN {
             blocks["ffn.2"] = std::shared_ptr<GGMLBlock>(new Linear(ffn_dim, dim));
         }
 
-        virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                            struct ggml_tensor* x,
-                                            struct ggml_tensor* e,
-                                            struct ggml_tensor* pe,
-                                            struct ggml_tensor* context,
-                                            int64_t context_img_len = 257) {
+        virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                     ggml_tensor* x,
+                                     ggml_tensor* e,
+                                     ggml_tensor* pe,
+                                     ggml_tensor* context,
+                                     int64_t context_img_len = 257) {
             // x: [N, n_token, dim]
             // e: [N, 6, dim] or [N, T, 6, dim]
             // context: [N, context_img_len + context_txt_len, dim]
@@ -1585,7 +1680,7 @@ namespace WAN {
     class VaceWanAttentionBlock : public WanAttentionBlock {
     protected:
         int block_id;
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(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);
             params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
         }
@@ -1607,11 +1702,11 @@ namespace WAN {
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                      struct ggml_tensor* c,
-                                                      struct ggml_tensor* x,
-                                                      struct ggml_tensor* e,
-                                                      struct ggml_tensor* pe,
-                                                      struct ggml_tensor* context,
+                                                      ggml_tensor* c,
+                                                      ggml_tensor* x,
+                                                      ggml_tensor* e,
+                                                      ggml_tensor* pe,
+                                                      ggml_tensor* context,
                                                       int64_t context_img_len = 257) {
             // x: [N, n_token, dim]
             // e: [N, 6, dim] or [N, T, 6, dim]
@@ -1637,7 +1732,7 @@ namespace WAN {
     protected:
         int64_t dim;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(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);
             params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 2, 1);
         }
@@ -1654,9 +1749,9 @@ namespace WAN {
             blocks["head"] = std::shared_ptr<GGMLBlock>(new Linear(dim, out_dim));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* e) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* e) {
             // x: [N, n_token, dim]
             // e: [N, dim] or [N, T, dim]
             // return [N, n_token, out_dim]
@@ -1684,7 +1779,7 @@ namespace WAN {
         int64_t in_dim;
         int64_t flf_pos_embed_token_number;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             if (flf_pos_embed_token_number > 0) {
                 params["emb_pos"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, in_dim, flf_pos_embed_token_number, 1);
             }
@@ -1702,8 +1797,8 @@ namespace WAN {
             blocks["proj.4"] = std::shared_ptr<GGMLBlock>(new LayerNorm(out_dim));
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* image_embeds) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* image_embeds) {
             if (flf_pos_embed_token_number > 0) {
                 auto emb_pos = params["emb_pos"];
 
@@ -1822,8 +1917,8 @@ namespace WAN {
             }
         }
 
-        struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
-                                              struct ggml_tensor* x) {
+        ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
+                                       ggml_tensor* x) {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
             int64_t T = x->ne[2];
@@ -1835,11 +1930,11 @@ namespace WAN {
             return x;
         }
 
-        struct ggml_tensor* unpatchify(struct ggml_context* ctx,
-                                       struct ggml_tensor* x,
-                                       int64_t t_len,
-                                       int64_t h_len,
-                                       int64_t w_len) {
+        ggml_tensor* unpatchify(ggml_context* ctx,
+                                ggml_tensor* x,
+                                int64_t t_len,
+                                int64_t h_len,
+                                int64_t w_len) {
             // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
             // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
             int64_t N  = x->ne[3];
@@ -1862,15 +1957,15 @@ namespace WAN {
             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,
-                                         struct ggml_tensor* clip_fea     = nullptr,
-                                         struct ggml_tensor* vace_context = nullptr,
-                                         float vace_strength              = 1.f,
-                                         int64_t N                        = 1) {
+        ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
+                                  ggml_tensor* x,
+                                  ggml_tensor* timestep,
+                                  ggml_tensor* context,
+                                  ggml_tensor* pe,
+                                  ggml_tensor* clip_fea     = nullptr,
+                                  ggml_tensor* vace_context = nullptr,
+                                  float vace_strength       = 1.f,
+                                  int64_t N                 = 1) {
             // x: [N*C, T, H, W], C => in_dim
             // vace_context: [N*vace_in_dim, T, H, W]
             // timestep: [N,] or [T]
@@ -1956,16 +2051,16 @@ namespace WAN {
             return x;
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* timestep,
-                                    struct ggml_tensor* context,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* clip_fea        = nullptr,
-                                    struct ggml_tensor* time_dim_concat = nullptr,
-                                    struct ggml_tensor* vace_context    = nullptr,
-                                    float vace_strength                 = 1.f,
-                                    int64_t N                           = 1) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe,
+                             ggml_tensor* clip_fea        = nullptr,
+                             ggml_tensor* time_dim_concat = nullptr,
+                             ggml_tensor* vace_context    = nullptr,
+                             float vace_strength          = 1.f,
+                             int64_t N                    = 1) {
             // Forward pass of DiT.
             // x: [N*C, T, H, W]
             // timestep: [N,]
@@ -2130,19 +2225,19 @@ namespace WAN {
             return desc;
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
             wan.get_param_tensors(tensors, prefix);
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
-                                        struct ggml_tensor* timesteps,
-                                        struct ggml_tensor* context,
-                                        struct ggml_tensor* clip_fea        = nullptr,
-                                        struct ggml_tensor* c_concat        = nullptr,
-                                        struct ggml_tensor* time_dim_concat = nullptr,
-                                        struct ggml_tensor* vace_context    = nullptr,
-                                        float vace_strength                 = 1.f) {
-            struct ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
+        ggml_cgraph* build_graph(ggml_tensor* x,
+                                 ggml_tensor* timesteps,
+                                 ggml_tensor* context,
+                                 ggml_tensor* clip_fea        = nullptr,
+                                 ggml_tensor* c_concat        = nullptr,
+                                 ggml_tensor* time_dim_concat = nullptr,
+                                 ggml_tensor* vace_context    = nullptr,
+                                 float vace_strength          = 1.f) {
+            ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
 
             x               = to_backend(x);
             timesteps       = to_backend(timesteps);
@@ -2175,15 +2270,15 @@ namespace WAN {
 
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* out = wan.forward(&runner_ctx,
-                                                  x,
-                                                  timesteps,
-                                                  context,
-                                                  pe,
-                                                  clip_fea,
-                                                  time_dim_concat,
-                                                  vace_context,
-                                                  vace_strength);
+            ggml_tensor* out = wan.forward(&runner_ctx,
+                                           x,
+                                           timesteps,
+                                           context,
+                                           pe,
+                                           clip_fea,
+                                           time_dim_concat,
+                                           vace_context,
+                                           vace_strength);
 
             ggml_build_forward_expand(gf, out);
 
@@ -2191,17 +2286,17 @@ namespace WAN {
         }
 
         bool compute(int n_threads,
-                     struct ggml_tensor* x,
-                     struct ggml_tensor* timesteps,
-                     struct ggml_tensor* context,
-                     struct ggml_tensor* clip_fea        = nullptr,
-                     struct ggml_tensor* c_concat        = nullptr,
-                     struct ggml_tensor* time_dim_concat = nullptr,
-                     struct ggml_tensor* vace_context    = nullptr,
-                     float vace_strength                 = 1.f,
-                     struct ggml_tensor** output         = nullptr,
-                     struct ggml_context* output_ctx     = nullptr) {
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+                     ggml_tensor* x,
+                     ggml_tensor* timesteps,
+                     ggml_tensor* context,
+                     ggml_tensor* clip_fea        = nullptr,
+                     ggml_tensor* c_concat        = nullptr,
+                     ggml_tensor* time_dim_concat = nullptr,
+                     ggml_tensor* vace_context    = nullptr,
+                     float vace_strength          = 1.f,
+                     ggml_tensor** output         = nullptr,
+                     ggml_context* output_ctx     = nullptr) {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
             };
 
@@ -2209,12 +2304,12 @@ namespace WAN {
         }
 
         void test() {
-            struct ggml_init_params params;
+            ggml_init_params params;
             params.mem_size   = static_cast<size_t>(200 * 1024 * 1024);  // 200 MB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            struct ggml_context* work_ctx = ggml_init(params);
+            ggml_context* work_ctx = ggml_init(params);
             GGML_ASSERT(work_ctx != nullptr);
 
             {
@@ -2237,7 +2332,7 @@ namespace WAN {
                 // auto clip_fea = load_tensor_from_file(work_ctx, "wan_dit_clip_fea.bin");
                 // print_ggml_tensor(clip_fea);
 
-                struct ggml_tensor* out = nullptr;
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 compute(8, x, timesteps, context, nullptr, nullptr, nullptr, nullptr, 1.f, &out, work_ctx);

src/z_image.hpp

@@ -42,10 +42,10 @@ namespace ZImage {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mask = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* mask = nullptr) {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
@@ -124,23 +124,23 @@ namespace ZImage {
             blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
             auto x1 = w1->forward(ctx, x);
             auto x3 = w3->forward(ctx, x);
-            x       = ggml_mul(ctx->ggml_ctx, ggml_silu(ctx->ggml_ctx, x1), x3);
+            x       = ggml_swiglu_split(ctx->ggml_ctx, x1, x3);
             x       = w2->forward(ctx, x);
 
             return x;
         }
     };
 
-    __STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
-                                                   struct ggml_tensor* x,
-                                                   struct ggml_tensor* scale) {
+    __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+                                            ggml_tensor* x,
+                                            ggml_tensor* scale) {
         // x: [N, L, C]
         // scale: [N, C]
         scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);  // [N, 1, C]
@@ -175,11 +175,11 @@ namespace ZImage {
             }
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* mask        = nullptr,
-                                    struct ggml_tensor* adaln_input = nullptr) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* mask        = nullptr,
+                             ggml_tensor* adaln_input = nullptr) {
             auto attention       = std::dynamic_pointer_cast<JointAttention>(blocks["attention"]);
             auto feed_forward    = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"]);
             auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
@@ -241,9 +241,9 @@ namespace ZImage {
             blocks["adaLN_modulation.1"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), hidden_size);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* c) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* c) {
             // x: [N, n_token, hidden_size]
             // c: [N, hidden_size]
             // return: [N, n_token, patch_size * patch_size * out_channels]
@@ -284,7 +284,7 @@ namespace ZImage {
     protected:
         ZImageParams z_image_params;
 
-        void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
             params["x_pad_token"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
         }
@@ -346,74 +346,11 @@ namespace ZImage {
             blocks["final_layer"] = std::make_shared<FinalLayer>(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels);
         }
 
-        struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
-                                              struct ggml_tensor* x) {
-            int64_t W = x->ne[0];
-            int64_t H = x->ne[1];
-
-            int pad_h = (z_image_params.patch_size - H % z_image_params.patch_size) % z_image_params.patch_size;
-            int pad_w = (z_image_params.patch_size - W % z_image_params.patch_size) % z_image_params.patch_size;
-            x         = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-            return x;
-        }
-
-        struct ggml_tensor* patchify(struct ggml_context* ctx,
-                                     struct ggml_tensor* x) {
-            // x: [N, C, H, W]
-            // return: [N, h*w, patch_size*patch_size*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 p = z_image_params.patch_size;
-            int64_t h = H / z_image_params.patch_size;
-            int64_t w = W / z_image_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_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, h*w, C, p*p]
-            x = ggml_reshape_3d(ctx, x, C * p * p, w * h, N);                // [N, h*w, p*p*C]
-            return x;
-        }
-
-        struct ggml_tensor* process_img(GGMLRunnerContext* ctx,
-                                        struct ggml_tensor* x) {
-            x = pad_to_patch_size(ctx, x);
-            x = patchify(ctx->ggml_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, patch_size*patch_size*C]
-            // return: [N, C, H, W]
-            int64_t N = x->ne[2];
-            int64_t C = x->ne[0] / z_image_params.patch_size / z_image_params.patch_size;
-            int64_t H = h * z_image_params.patch_size;
-            int64_t W = w * z_image_params.patch_size;
-            int64_t p = z_image_params.patch_size;
-
-            GGML_ASSERT(C * p * p == x->ne[0]);
-
-            x = ggml_reshape_4d(ctx, x, C, p * p, w * h, N);                 // [N, h*w, p*p, C]
-            x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 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_core(GGMLRunnerContext* ctx,
-                                         struct ggml_tensor* x,
-                                         struct ggml_tensor* timestep,
-                                         struct ggml_tensor* context,
-                                         struct ggml_tensor* pe) {
+        ggml_tensor* forward_core(GGMLRunnerContext* ctx,
+                                  ggml_tensor* x,
+                                  ggml_tensor* timestep,
+                                  ggml_tensor* context,
+                                  ggml_tensor* pe) {
             auto x_embedder     = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
             auto t_embedder     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
             auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
@@ -477,12 +414,12 @@ namespace ZImage {
             return img;
         }
 
-        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 = {}) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe,
+                             std::vector<ggml_tensor*> ref_latents = {}) {
             // Forward pass of DiT.
             // x: [N, C, H, W]
             // timestep: [N,]
@@ -495,27 +432,22 @@ namespace ZImage {
             int64_t C = x->ne[2];
             int64_t N = x->ne[3];
 
-            auto img             = process_img(ctx, x);
+            int patch_size = z_image_params.patch_size;
+
+            auto img             = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
             uint64_t n_img_token = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = process_img(ctx, ref);
+                    ref = DiT::pad_and_patchify(ctx, ref, patch_size, patch_size, false);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
 
-            int64_t h_len = ((H + (z_image_params.patch_size / 2)) / z_image_params.patch_size);
-            int64_t w_len = ((W + (z_image_params.patch_size / 2)) / z_image_params.patch_size);
-
             auto out = forward_core(ctx, img, timestep, context, pe);
 
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token);  // [N, n_img_token, ph*pw*C]
-            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]
+            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token);                              // [N, n_img_token, ph*pw*C]
+            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);  // [N, C, H, W]
 
             out = ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
 
@@ -545,17 +477,17 @@ namespace ZImage {
             return "z_image";
         }
 
-        void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
             z_image.get_param_tensors(tensors, prefix);
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
-                                        struct ggml_tensor* timesteps,
-                                        struct ggml_tensor* context,
-                                        std::vector<ggml_tensor*> ref_latents = {},
-                                        bool increase_ref_index               = false) {
+        ggml_cgraph* build_graph(ggml_tensor* x,
+                                 ggml_tensor* timesteps,
+                                 ggml_tensor* context,
+                                 std::vector<ggml_tensor*> ref_latents = {},
+                                 bool increase_ref_index               = false) {
             GGML_ASSERT(x->ne[3] == 1);
-            struct ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
+            ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
 
             x         = to_backend(x);
             context   = to_backend(context);
@@ -586,12 +518,12 @@ namespace ZImage {
             set_backend_tensor_data(pe, pe_vec.data());
             auto runner_ctx = get_context();
 
-            struct ggml_tensor* out = z_image.forward(&runner_ctx,
-                                                      x,
-                                                      timesteps,
-                                                      context,
-                                                      pe,
-                                                      ref_latents);
+            ggml_tensor* out = z_image.forward(&runner_ctx,
+                                               x,
+                                               timesteps,
+                                               context,
+                                               pe,
+                                               ref_latents);
 
             ggml_build_forward_expand(gf, out);
 
@@ -599,17 +531,17 @@ namespace ZImage {
         }
 
         bool compute(int n_threads,
-                     struct ggml_tensor* x,
-                     struct ggml_tensor* timesteps,
-                     struct ggml_tensor* context,
+                     ggml_tensor* x,
+                     ggml_tensor* timesteps,
+                     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) {
+                     ggml_tensor** output                  = nullptr,
+                     ggml_context* output_ctx              = nullptr) {
             // x: [N, in_channels, h, w]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
@@ -617,12 +549,12 @@ namespace ZImage {
         }
 
         void test() {
-            struct ggml_init_params params;
+            ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
             params.mem_buffer = nullptr;
             params.no_alloc   = false;
 
-            struct ggml_context* work_ctx = ggml_init(params);
+            ggml_context* work_ctx = ggml_init(params);
             GGML_ASSERT(work_ctx != nullptr);
 
             {
@@ -639,7 +571,7 @@ namespace ZImage {
                 auto context = load_tensor_from_file(work_ctx, "./z_image_context.bin");
                 print_ggml_tensor(context);
 
-                struct ggml_tensor* out = nullptr;
+                ggml_tensor* out = nullptr;
 
                 int64_t t0 = ggml_time_ms();
                 compute(8, x, timesteps, context, {}, false, &out, work_ctx);