fix: correct double increment on flow denoisers sigma calculations (#1372)

* inpaint: get max pixel max instead of single sample

* inpaint: masked diffusion for inpainting models with inflated mask

* refactor tensor interpolate nearest-like reduction paths and generalize max_pool_2d

---------

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

.github/workflows/build.yml

@@ -64,7 +64,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Dependencies
         id: depends
@@ -127,7 +127,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Dependencies
         id: depends
@@ -205,7 +205,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Get commit hash
         id: commit
@@ -239,6 +239,7 @@ jobs:
         id: build-push
         uses: docker/build-push-action@v6
         with:
+          context: .
           platforms: linux/amd64
           push: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
           file: Dockerfile.${{ matrix.variant }}
@@ -264,7 +265,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Dependencies
         id: depends
@@ -345,7 +346,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Install cuda-toolkit
         id: cuda-toolkit
@@ -460,7 +461,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Cache ROCm Installation
         id: cache-rocm
@@ -573,7 +574,7 @@ jobs:
       - name: Setup pnpm
         uses: pnpm/action-setup@v4
         with:
-          version: 9
+          version: 10.15.1
 
       - name: Free disk space
         run: |

.gitmodules

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

CMakeLists.txt

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

docs/build.md

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

docs/caching.md

@@ -11,7 +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 |
+| `spectrum` | UNET and DiT models | Chebyshev + Taylor output forecasting |
 
 ### UCache (UNET Models)
 
@@ -111,9 +111,9 @@ Mask values: `1` = compute, `0` = can cache.
 --scm-policy dynamic
 ```
 
-### Spectrum (UNET Models)
+### Spectrum (UNET and DiT 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).
+Spectrum uses Chebyshev polynomial fitting blended with Taylor extrapolation to predict denoised outputs, skipping entire forward passes. Based on the paper [Spectrum: Adaptive Spectral Feature Forecasting for Efficient Diffusion Sampling](https://github.com/tingyu215/Spectrum).
 
 ```bash
 sd-cli -m model.safetensors -p "a cat" --cache-mode spectrum

examples/cli/CMakeLists.txt

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

examples/cli/README.md

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

examples/cli/avi_writer.h

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

examples/cli/image_metadata.h

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

examples/cli/main.cpp

@@ -15,9 +15,12 @@
 // #include "preprocessing.hpp"
 #include "stable-diffusion.h"
 
-#include "common/common.hpp"
+#include "common/common.h"
+#include "common/media_io.h"
+#include "common/resource_owners.hpp"
+#include "image_metadata.h"
 
-#include "avi_writer.h"
+namespace fs = std::filesystem;
 
 const char* previews_str[] = {
     "none",
@@ -32,6 +35,8 @@ struct SDCliParams {
     SDMode mode             = IMG_GEN;
     std::string output_path = "output.png";
     int output_begin_idx    = -1;
+    std::string image_path;
+    std::string metadata_format = "text";
 
     bool verbose          = false;
     bool canny_preprocess = false;
@@ -44,6 +49,9 @@ struct SDCliParams {
     bool taesd_preview       = false;
     bool preview_noisy       = false;
     bool color               = false;
+    bool metadata_raw        = false;
+    bool metadata_brief      = false;
+    bool metadata_all        = false;
 
     bool normal_exit = false;
 
@@ -53,11 +61,19 @@ struct SDCliParams {
         options.string_options = {
             {"-o",
              "--output",
-             "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png)",
+             "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp",
              &output_path},
+            {"",
+             "--image",
+             "path to the image to inspect (for metadata mode)",
+             &image_path},
+            {"",
+             "--metadata-format",
+             "metadata output format, one of [text, json] (default: text)",
+             &metadata_format},
             {"",
              "--preview-path",
-             "path to write preview image to (default: ./preview.png)",
+             "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
              &preview_path},
         };
 
@@ -97,6 +113,18 @@ struct SDCliParams {
              "--preview-noisy",
              "enables previewing noisy inputs of the models rather than the denoised outputs",
              true, &preview_noisy},
+            {"",
+             "--metadata-raw",
+             "include raw hex previews for unparsed metadata payloads",
+             true, &metadata_raw},
+            {"",
+             "--metadata-brief",
+             "truncate long metadata text values in text output",
+             true, &metadata_brief},
+            {"",
+             "--metadata-all",
+             "include structural/container entries such as IHDR, IDAT, and non-metadata JPEG segments",
+             true, &metadata_all},
 
         };
 
@@ -149,7 +177,7 @@ struct SDCliParams {
         options.manual_options = {
             {"-M",
              "--mode",
-             "run mode, one of [img_gen, vid_gen, upscale, convert], default: img_gen",
+             "run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen",
              on_mode_arg},
             {"",
              "--preview",
@@ -165,7 +193,7 @@ struct SDCliParams {
     };
 
     bool process_and_check() {
-        if (output_path.length() == 0) {
+        if (mode != METADATA && output_path.length() == 0) {
             LOG_ERROR("error: the following arguments are required: output_path");
             return false;
         }
@@ -174,6 +202,16 @@ struct SDCliParams {
             if (output_path == "output.png") {
                 output_path = "output.gguf";
             }
+        } else if (mode == METADATA) {
+            if (image_path.empty()) {
+                LOG_ERROR("error: metadata mode needs an image path (--image)");
+                return false;
+            }
+            if (metadata_format != "text" && metadata_format != "json") {
+                LOG_ERROR("error: invalid metadata format %s, must be one of [text, json]",
+                          metadata_format.c_str());
+                return false;
+            }
         }
         return true;
     }
@@ -183,6 +221,8 @@ struct SDCliParams {
         oss << "SDCliParams {\n"
             << "  mode: " << modes_str[mode] << ",\n"
             << "  output_path: \"" << output_path << "\",\n"
+            << "  image_path: \"" << image_path << "\",\n"
+            << "  metadata_format: \"" << metadata_format << "\",\n"
             << "  verbose: " << (verbose ? "true" : "false") << ",\n"
             << "  color: " << (color ? "true" : "false") << ",\n"
             << "  canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
@@ -192,7 +232,10 @@ struct SDCliParams {
             << "  preview_path: \"" << preview_path << "\",\n"
             << "  preview_fps: " << preview_fps << ",\n"
             << "  taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n"
-            << "  preview_noisy: " << (preview_noisy ? "true" : "false") << "\n"
+            << "  preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n"
+            << "  metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n"
+            << "  metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n"
+            << "  metadata_all: " << (metadata_all ? "true" : "false") << "\n"
             << "}";
         return oss.str();
     }
@@ -217,78 +260,25 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
         exit(cli_params.normal_exit ? 0 : 1);
     }
 
-    if (!cli_params.process_and_check() ||
-        !ctx_params.process_and_check(cli_params.mode) ||
-        !gen_params.process_and_check(cli_params.mode, ctx_params.lora_model_dir)) {
+    bool valid = cli_params.process_and_check();
+    if (valid && cli_params.mode != METADATA) {
+        valid = ctx_params.process_and_check(cli_params.mode) &&
+                gen_params.process_and_check(cli_params.mode, ctx_params.lora_model_dir);
+    }
+
+    if (!valid) {
         print_usage(argc, argv, options_vec);
         exit(1);
     }
 }
 
-std::string get_image_params(const SDCliParams& cli_params, const SDContextParams& ctx_params, const SDGenerationParams& gen_params, int64_t seed) {
-    std::string parameter_string = gen_params.prompt_with_lora + "\n";
-    if (gen_params.negative_prompt.size() != 0) {
-        parameter_string += "Negative prompt: " + gen_params.negative_prompt + "\n";
-    }
-    parameter_string += "Steps: " + std::to_string(gen_params.sample_params.sample_steps) + ", ";
-    parameter_string += "CFG scale: " + std::to_string(gen_params.sample_params.guidance.txt_cfg) + ", ";
-    if (gen_params.sample_params.guidance.slg.scale != 0 && gen_params.skip_layers.size() != 0) {
-        parameter_string += "SLG scale: " + std::to_string(gen_params.sample_params.guidance.txt_cfg) + ", ";
-        parameter_string += "Skip layers: [";
-        for (const auto& layer : gen_params.skip_layers) {
-            parameter_string += std::to_string(layer) + ", ";
-        }
-        parameter_string += "], ";
-        parameter_string += "Skip layer start: " + std::to_string(gen_params.sample_params.guidance.slg.layer_start) + ", ";
-        parameter_string += "Skip layer end: " + std::to_string(gen_params.sample_params.guidance.slg.layer_end) + ", ";
-    }
-    parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", ";
-    parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", ";
-    parameter_string += "Seed: " + std::to_string(seed) + ", ";
-    parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
-    parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";
-    parameter_string += "RNG: " + std::string(sd_rng_type_name(ctx_params.rng_type)) + ", ";
-    if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) {
-        parameter_string += "Sampler RNG: " + std::string(sd_rng_type_name(ctx_params.sampler_rng_type)) + ", ";
-    }
-    parameter_string += "Sampler: " + std::string(sd_sample_method_name(gen_params.sample_params.sample_method));
-    if (!gen_params.custom_sigmas.empty()) {
-        parameter_string += ", Custom Sigmas: [";
-        for (size_t i = 0; i < gen_params.custom_sigmas.size(); ++i) {
-            std::ostringstream oss;
-            oss << std::fixed << std::setprecision(4) << gen_params.custom_sigmas[i];
-            parameter_string += oss.str() + (i == gen_params.custom_sigmas.size() - 1 ? "" : ", ");
-        }
-        parameter_string += "]";
-    } else if (gen_params.sample_params.scheduler != SCHEDULER_COUNT) {  // Only show schedule if not using custom sigmas
-        parameter_string += " " + std::string(sd_scheduler_name(gen_params.sample_params.scheduler));
-    }
-    parameter_string += ", ";
-    for (const auto& te : {ctx_params.clip_l_path, ctx_params.clip_g_path, ctx_params.t5xxl_path, ctx_params.llm_path, ctx_params.llm_vision_path}) {
-        if (!te.empty()) {
-            parameter_string += "TE: " + sd_basename(te) + ", ";
-        }
-    }
-    if (!ctx_params.diffusion_model_path.empty()) {
-        parameter_string += "Unet: " + sd_basename(ctx_params.diffusion_model_path) + ", ";
-    }
-    if (!ctx_params.vae_path.empty()) {
-        parameter_string += "VAE: " + sd_basename(ctx_params.vae_path) + ", ";
-    }
-    if (gen_params.clip_skip != -1) {
-        parameter_string += "Clip skip: " + std::to_string(gen_params.clip_skip) + ", ";
-    }
-    parameter_string += "Version: stable-diffusion.cpp";
-    return parameter_string;
-}
-
 void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
     SDCliParams* cli_params = (SDCliParams*)data;
     log_print(level, log, cli_params->verbose, cli_params->color);
 }
 
 bool load_images_from_dir(const std::string dir,
-                          std::vector<sd_image_t>& images,
+                          SDImageVec& images,
                           int expected_width  = 0,
                           int expected_height = 0,
                           int max_image_num   = 0,
@@ -315,7 +305,7 @@ bool load_images_from_dir(const std::string dir,
         std::string ext  = entry.path().extension().string();
         std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
 
-        if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp") {
+        if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp" || ext == ".webp") {
             LOG_DEBUG("load image %zu from '%s'", images.size(), path.c_str());
             int width             = 0;
             int height            = 0;
@@ -330,7 +320,7 @@ bool load_images_from_dir(const std::string dir,
                               3,
                               image_buffer});
 
-            if (max_image_num > 0 && images.size() >= max_image_num) {
+            if (max_image_num > 0 && static_cast<int>(images.size()) >= max_image_num) {
                 break;
             }
         }
@@ -345,9 +335,17 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy,
     // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
     // unused in this app, it will either be always noisy or always denoised here
     if (frame_count == 1) {
-        stbi_write_png(cli_params->preview_path.c_str(), image->width, image->height, image->channel, image->data, 0);
+        if (!write_image_to_file(cli_params->preview_path,
+                                 image->data,
+                                 image->width,
+                                 image->height,
+                                 image->channel)) {
+            LOG_ERROR("save preview image to '%s' failed", cli_params->preview_path.c_str());
+        }
     } else {
-        create_mjpg_avi_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps);
+        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps) != 0) {
+            LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
+        }
     }
 }
 
@@ -397,9 +395,13 @@ bool save_results(const SDCliParams& cli_params,
 
     std::string ext_lower = ext.string();
     std::transform(ext_lower.begin(), ext_lower.end(), ext_lower.begin(), ::tolower);
-    bool is_jpg = (ext_lower == ".jpg" || ext_lower == ".jpeg" || ext_lower == ".jpe");
+    const EncodedImageFormat output_format = encoded_image_format_from_path(out_path.string());
     if (!ext.empty()) {
-        if (is_jpg || ext_lower == ".png") {
+        if (output_format == EncodedImageFormat::JPEG ||
+            output_format == EncodedImageFormat::PNG ||
+            output_format == EncodedImageFormat::WEBP ||
+            ext_lower == ".avi" ||
+            ext_lower == ".webm") {
             base_path.replace_extension();
         }
     }
@@ -414,21 +416,18 @@ bool save_results(const SDCliParams& cli_params,
         if (!img.data)
             return false;
 
-        std::string params = get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + idx);
-        int ok             = 0;
-        if (is_jpg) {
-            ok = stbi_write_jpg(path.string().c_str(), img.width, img.height, img.channel, img.data, 90, params.c_str());
-        } else {
-            ok = stbi_write_png(path.string().c_str(), img.width, img.height, img.channel, img.data, 0, params.c_str());
-        }
+        std::string params = gen_params.embed_image_metadata
+                                 ? get_image_params(ctx_params, gen_params, gen_params.seed + idx)
+                                 : "";
+        const bool ok      = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, 90);
         LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
-        return ok != 0;
+        return ok;
     };
 
     int sucessful_reults = 0;
 
     if (std::regex_search(cli_params.output_path, format_specifier_regex)) {
-        if (!is_jpg && ext_lower != ".png")
+        if (output_format == EncodedImageFormat::UNKNOWN)
             ext = ".png";
         fs::path pattern = base_path;
         pattern += ext;
@@ -444,20 +443,20 @@ bool save_results(const SDCliParams& cli_params,
     }
 
     if (cli_params.mode == VID_GEN && num_results > 1) {
-        if (ext_lower != ".avi")
+        if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
             ext = ".avi";
         fs::path video_path = base_path;
         video_path += ext;
-        if (create_mjpg_avi_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
-            LOG_INFO("save result MJPG AVI video to '%s'", video_path.string().c_str());
+        if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
+            LOG_INFO("save result video to '%s'", video_path.string().c_str());
             return true;
         } else {
-            LOG_ERROR("Failed to save result MPG AVI video to '%s'", video_path.string().c_str());
+            LOG_ERROR("Failed to save result video to '%s'", video_path.string().c_str());
             return false;
         }
     }
 
-    if (!is_jpg && ext_lower != ".png")
+    if (output_format == EncodedImageFormat::UNKNOWN)
         ext = ".png";
 
     for (int i = 0; i < num_results; ++i) {
@@ -485,6 +484,27 @@ int main(int argc, const char* argv[]) {
     SDGenerationParams gen_params;
 
     parse_args(argc, argv, cli_params, ctx_params, gen_params);
+    sd_set_log_callback(sd_log_cb, (void*)&cli_params);
+    log_verbose = cli_params.verbose;
+    log_color   = cli_params.color;
+
+    if (cli_params.mode == METADATA) {
+        MetadataReadOptions options;
+        options.output_format      = cli_params.metadata_format == "json"
+                                         ? MetadataOutputFormat::JSON
+                                         : MetadataOutputFormat::TEXT;
+        options.include_raw        = cli_params.metadata_raw;
+        options.brief              = cli_params.metadata_brief;
+        options.include_structural = cli_params.metadata_all;
+
+        std::string error;
+        if (!print_image_metadata(cli_params.image_path, options, std::cout, error)) {
+            LOG_ERROR("%s", error.c_str());
+            return 1;
+        }
+        return 0;
+    }
+
     if (gen_params.video_frames > 4) {
         size_t last_dot_pos   = cli_params.preview_path.find_last_of(".");
         std::string base_path = cli_params.preview_path;
@@ -502,9 +522,6 @@ int main(int argc, const char* argv[]) {
     if (cli_params.preview_method == PREVIEW_PROJ)
         cli_params.preview_fps /= 4;
 
-    sd_set_log_callback(sd_log_cb, (void*)&cli_params);
-    log_verbose = cli_params.verbose;
-    log_color   = cli_params.color;
     sd_set_preview_callback(step_callback,
                             cli_params.preview_method,
                             cli_params.preview_interval,
@@ -540,39 +557,17 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    bool vae_decode_only     = true;
-    sd_image_t init_image    = {0, 0, 3, nullptr};
-    sd_image_t end_image     = {0, 0, 3, nullptr};
-    sd_image_t control_image = {0, 0, 3, nullptr};
-    sd_image_t mask_image    = {0, 0, 1, nullptr};
-    std::vector<sd_image_t> ref_images;
-    std::vector<sd_image_t> pmid_images;
-    std::vector<sd_image_t> control_frames;
-
-    auto release_all_resources = [&]() {
-        free(init_image.data);
-        free(end_image.data);
-        free(control_image.data);
-        free(mask_image.data);
-        for (auto image : ref_images) {
-            free(image.data);
-            image.data = nullptr;
-        }
-        ref_images.clear();
-        for (auto image : pmid_images) {
-            free(image.data);
-            image.data = nullptr;
-        }
-        pmid_images.clear();
-        for (auto image : control_frames) {
-            free(image.data);
-            image.data = nullptr;
-        }
-        control_frames.clear();
-    };
+    bool vae_decode_only = true;
+    SDImageOwner init_image({0, 0, 3, nullptr});
+    SDImageOwner end_image({0, 0, 3, nullptr});
+    SDImageOwner control_image({0, 0, 3, nullptr});
+    SDImageOwner mask_image({0, 0, 1, nullptr});
+    SDImageVec ref_images;
+    SDImageVec pmid_images;
+    SDImageVec control_frames;
 
     auto load_image_and_update_size = [&](const std::string& path,
-                                          sd_image_t& image,
+                                          SDImageOwner& image,
                                           bool resize_image    = true,
                                           int expected_channel = 3) -> bool {
         int expected_width  = 0;
@@ -582,13 +577,12 @@ int main(int argc, const char* argv[]) {
             expected_height = gen_params.height;
         }
 
-        if (!load_sd_image_from_file(&image, path.c_str(), expected_width, expected_height, expected_channel)) {
+        if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) {
             LOG_ERROR("load image from '%s' failed", path.c_str());
-            release_all_resources();
             return false;
         }
 
-        gen_params.set_width_and_height_if_unset(image.width, image.height);
+        gen_params.set_width_and_height_if_unset(image.get().width, image.get().height);
         return true;
     };
 
@@ -601,7 +595,7 @@ int main(int argc, const char* argv[]) {
 
     if (gen_params.end_image_path.size() > 0) {
         vae_decode_only = false;
-        if (!load_image_and_update_size(gen_params.init_image_path, end_image)) {
+        if (!load_image_and_update_size(gen_params.end_image_path, end_image)) {
             return 1;
         }
     }
@@ -609,47 +603,46 @@ int main(int argc, const char* argv[]) {
     if (gen_params.ref_image_paths.size() > 0) {
         vae_decode_only = false;
         for (auto& path : gen_params.ref_image_paths) {
-            sd_image_t ref_image = {0, 0, 3, nullptr};
+            SDImageOwner ref_image({0, 0, 3, nullptr});
             if (!load_image_and_update_size(path, ref_image, false)) {
                 return 1;
             }
-            ref_images.push_back(ref_image);
+            ref_images.push_back(std::move(ref_image));
         }
     }
 
     if (gen_params.mask_image_path.size() > 0) {
-        if (!load_sd_image_from_file(&mask_image,
+        if (!load_sd_image_from_file(mask_image.put(),
                                      gen_params.mask_image_path.c_str(),
                                      gen_params.get_resolved_width(),
                                      gen_params.get_resolved_height(),
                                      1)) {
             LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
-            release_all_resources();
             return 1;
         }
     } else {
-        mask_image.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
-        if (mask_image.data == nullptr) {
+        sd_image_t generated_mask = {0, 0, 1, nullptr};
+        generated_mask.data       = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        if (generated_mask.data == nullptr) {
             LOG_ERROR("malloc mask image failed");
-            release_all_resources();
             return 1;
         }
-        mask_image.width  = gen_params.get_resolved_width();
-        mask_image.height = gen_params.get_resolved_height();
-        memset(mask_image.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        generated_mask.width  = gen_params.get_resolved_width();
+        generated_mask.height = gen_params.get_resolved_height();
+        memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
+        mask_image.reset(generated_mask);
     }
 
     if (gen_params.control_image_path.size() > 0) {
-        if (!load_sd_image_from_file(&control_image,
+        if (!load_sd_image_from_file(control_image.put(),
                                      gen_params.control_image_path.c_str(),
                                      gen_params.get_resolved_width(),
                                      gen_params.get_resolved_height())) {
             LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
-            release_all_resources();
             return 1;
         }
         if (cli_params.canny_preprocess) {  // apply preprocessor
-            preprocess_canny(control_image,
+            preprocess_canny(control_image.get(),
                              0.08f,
                              0.08f,
                              0.8f,
@@ -665,7 +658,6 @@ int main(int argc, const char* argv[]) {
                                   gen_params.get_resolved_height(),
                                   gen_params.video_frames,
                                   cli_params.verbose)) {
-            release_all_resources();
             return 1;
         }
     }
@@ -677,7 +669,6 @@ int main(int argc, const char* argv[]) {
                                   0,
                                   0,
                                   cli_params.verbose)) {
-            release_all_resources();
             return 1;
         }
     }
@@ -688,39 +679,30 @@ int main(int argc, const char* argv[]) {
 
     sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
 
-    sd_image_t* results = nullptr;
-    int num_results     = 0;
+    SDImageVec results;
+    int num_results = 0;
 
     if (cli_params.mode == UPSCALE) {
         num_results = 1;
-        results     = (sd_image_t*)calloc(num_results, sizeof(sd_image_t));
-        if (results == nullptr) {
-            LOG_INFO("failed to allocate results array");
-            release_all_resources();
-            return 1;
-        }
-
-        results[0]      = init_image;
-        init_image.data = nullptr;
+        results.push_back(init_image.release());
     } else {
-        sd_ctx_t* sd_ctx = new_sd_ctx(&sd_ctx_params);
+        SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
 
         if (sd_ctx == nullptr) {
             LOG_INFO("new_sd_ctx_t failed");
-            release_all_resources();
             return 1;
         }
 
         if (gen_params.sample_params.sample_method == SAMPLE_METHOD_COUNT) {
-            gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
+            gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get());
         }
 
         if (gen_params.high_noise_sample_params.sample_method == SAMPLE_METHOD_COUNT) {
-            gen_params.high_noise_sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
+            gen_params.high_noise_sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get());
         }
 
         if (gen_params.sample_params.scheduler == SCHEDULER_COUNT) {
-            gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx, gen_params.sample_params.sample_method);
+            gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method);
         }
 
         if (cli_params.mode == IMG_GEN) {
@@ -730,19 +712,19 @@ int main(int argc, const char* argv[]) {
                 gen_params.prompt.c_str(),
                 gen_params.negative_prompt.c_str(),
                 gen_params.clip_skip,
-                init_image,
+                init_image.get(),
                 ref_images.data(),
                 (int)ref_images.size(),
                 gen_params.auto_resize_ref_image,
                 gen_params.increase_ref_index,
-                mask_image,
+                mask_image.get(),
                 gen_params.get_resolved_width(),
                 gen_params.get_resolved_height(),
                 gen_params.sample_params,
                 gen_params.strength,
                 gen_params.seed,
                 gen_params.batch_count,
-                control_image,
+                control_image.get(),
                 gen_params.control_strength,
                 {
                     pmid_images.data(),
@@ -750,12 +732,12 @@ int main(int argc, const char* argv[]) {
                     gen_params.pm_id_embed_path.c_str(),
                     gen_params.pm_style_strength,
                 },  // pm_params
-                ctx_params.vae_tiling_params,
+                gen_params.vae_tiling_params,
                 gen_params.cache_params,
             };
 
-            results     = generate_image(sd_ctx, &img_gen_params);
             num_results = gen_params.batch_count;
+            results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
         } else if (cli_params.mode == VID_GEN) {
             sd_vid_gen_params_t vid_gen_params = {
                 gen_params.lora_vec.data(),
@@ -763,8 +745,8 @@ int main(int argc, const char* argv[]) {
                 gen_params.prompt.c_str(),
                 gen_params.negative_prompt.c_str(),
                 gen_params.clip_skip,
-                init_image,
-                end_image,
+                init_image.get(),
+                end_image.get(),
                 control_frames.data(),
                 (int)control_frames.size(),
                 gen_params.get_resolved_width(),
@@ -776,29 +758,27 @@ int main(int argc, const char* argv[]) {
                 gen_params.seed,
                 gen_params.video_frames,
                 gen_params.vace_strength,
-                ctx_params.vae_tiling_params,
+                gen_params.vae_tiling_params,
                 gen_params.cache_params,
             };
 
-            results = generate_video(sd_ctx, &vid_gen_params, &num_results);
+            sd_image_t* generated_video = generate_video(sd_ctx.get(), &vid_gen_params, &num_results);
+            results.adopt(generated_video, num_results);
         }
 
-        if (results == nullptr) {
+        if (!results) {
             LOG_ERROR("generate failed");
-            free_sd_ctx(sd_ctx);
             return 1;
         }
-
-        free_sd_ctx(sd_ctx);
     }
 
     int upscale_factor = 4;  // unused for RealESRGAN_x4plus_anime_6B.pth
     if (ctx_params.esrgan_path.size() > 0 && gen_params.upscale_repeats > 0) {
-        upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(ctx_params.esrgan_path.c_str(),
-                                                        ctx_params.offload_params_to_cpu,
-                                                        ctx_params.diffusion_conv_direct,
-                                                        ctx_params.n_threads,
-                                                        gen_params.upscale_tile_size);
+        UpscalerCtxPtr upscaler_ctx(new_upscaler_ctx(ctx_params.esrgan_path.c_str(),
+                                                     ctx_params.offload_params_to_cpu,
+                                                     ctx_params.diffusion_conv_direct,
+                                                     ctx_params.n_threads,
+                                                     gen_params.upscale_tile_size));
 
         if (upscaler_ctx == nullptr) {
             LOG_ERROR("new_upscaler_ctx failed");
@@ -807,32 +787,24 @@ int main(int argc, const char* argv[]) {
                 if (results[i].data == nullptr) {
                     continue;
                 }
-                sd_image_t current_image = results[i];
+                SDImageOwner current_image(results[i]);
+                results[i] = {0, 0, 0, nullptr};
                 for (int u = 0; u < gen_params.upscale_repeats; ++u) {
-                    sd_image_t upscaled_image = upscale(upscaler_ctx, current_image, upscale_factor);
-                    if (upscaled_image.data == nullptr) {
+                    SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor));
+                    if (upscaled_image.get().data == nullptr) {
                         LOG_ERROR("upscale failed");
                         break;
                     }
-                    free(current_image.data);
-                    current_image = upscaled_image;
+                    current_image = std::move(upscaled_image);
                 }
-                results[i] = current_image;  // Set the final upscaled image as the result
+                results[i] = current_image.release();  // Set the final upscaled image as the result
             }
         }
     }
 
-    if (!save_results(cli_params, ctx_params, gen_params, results, num_results)) {
+    if (!save_results(cli_params, ctx_params, gen_params, results.data(), num_results)) {
         return 1;
     }
 
-    for (int i = 0; i < num_results; i++) {
-        free(results[i].data);
-        results[i].data = nullptr;
-    }
-    free(results);
-
-    release_all_resources();
-
     return 0;
 }

examples/common/common.h

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

examples/common/log.cpp

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

examples/common/log.h

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

examples/common/media_io.h

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

examples/common/resource_owners.hpp

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

examples/server/CMakeLists.txt

@@ -56,7 +56,12 @@ else()
     message(STATUS "Frontend disabled or directory not found: ${FRONTEND_DIR}")
 endif()
 
-add_executable(${TARGET} main.cpp)
+add_executable(${TARGET}
+    ../common/common.cpp
+    ../common/log.cpp
+    ../common/media_io.cpp
+    main.cpp
+)
 
 if(HAVE_FRONTEND_BUILD)
     add_dependencies(${TARGET} ${TARGET}_frontend)
@@ -70,4 +75,18 @@ 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)
+if(SD_WEBP)
+    target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBP)
+    target_link_libraries(${TARGET} PRIVATE webp libwebpmux)
+endif()
+if(SD_WEBM)
+    target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBM)
+    target_link_libraries(${TARGET} PRIVATE webm)
+endif()
+
+# due to httplib; it contains a pragma for MSVC, but other things need explicit flags
+if(WIN32 AND NOT MSVC)
+    target_link_libraries(${TARGET} PRIVATE ws2_32)
+endif()
+
+target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17)

examples/server/README.md

@@ -189,7 +189,7 @@ Default Generation Options:
                                            medium
   --skip-layer-start <float>               SLG enabling point (default: 0.01)
   --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            eta in DDIM, only for DDIM and TCD (default: 0)
+  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
   --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
   --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)
@@ -197,7 +197,7 @@ Default Generation Options:
   --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: 0)
   --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
   --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) eta in DDIM, only for DDIM and TCD (default: 0)
+  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)
   --strength <float>                       strength for noising/unnoising (default: 0.75)
   --pm-style-strength <float>
   --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image
@@ -205,6 +205,7 @@ Default Generation Options:
   --vace-strength <float>                  wan vace strength
   --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
+  --disable-image-metadata                 do not embed generation metadata on image files
   -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

format-code.sh

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

include/stable-diffusion.h

@@ -120,7 +120,8 @@ enum sd_type_t {
     // SD_TYPE_IQ4_NL_4_8 = 37,
     // SD_TYPE_IQ4_NL_8_8 = 38,
     SD_TYPE_MXFP4 = 39,  // MXFP4 (1 block)
-    SD_TYPE_COUNT = 40,
+    SD_TYPE_NVFP4 = 40,  // NVFP4 (4 blocks, E4M3 scale)
+    SD_TYPE_COUNT = 41,
 };
 
 enum sd_log_level_t {

src/anima.hpp

@@ -13,9 +13,9 @@
 namespace Anima {
     constexpr int ANIMA_GRAPH_SIZE = 65536;
 
-    __STATIC_INLINE__ struct ggml_tensor* apply_gate(struct ggml_context* ctx,
-                                                     struct ggml_tensor* x,
-                                                     struct ggml_tensor* gate) {
+    __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);
     }
@@ -26,7 +26,7 @@ namespace Anima {
             blocks["proj.1"] = std::make_shared<Linear>(in_dim, out_dim, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj.1"]);
             return proj->forward(ctx, x);
         }
@@ -39,7 +39,7 @@ namespace Anima {
             blocks["1.linear_2"] = std::make_shared<Linear>(in_dim, out_dim, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        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"]);
 
@@ -62,10 +62,10 @@ namespace Anima {
             blocks["2"]    = std::make_shared<Linear>(hidden_features, 3 * in_features, false);
         }
 
-        std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
-                                                                    struct ggml_tensor* hidden_states,
-                                                                    struct ggml_tensor* embedded_timestep,
-                                                                    struct ggml_tensor* temb = nullptr) {
+        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"]);
@@ -102,10 +102,10 @@ namespace Anima {
             blocks["2"]    = std::make_shared<Linear>(hidden_features, 2 * in_features, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* hidden_states,
-                                    struct ggml_tensor* embedded_timestep,
-                                    struct ggml_tensor* temb = nullptr) {
+        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"]);
@@ -152,11 +152,11 @@ namespace Anima {
             blocks[this->out_proj_name] = std::make_shared<Linear>(inner_dim, query_dim, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* hidden_states,
-                                    struct ggml_tensor* encoder_hidden_states = nullptr,
-                                    struct ggml_tensor* pe_q                  = nullptr,
-                                    struct ggml_tensor* pe_k                  = nullptr) {
+        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;
             }
@@ -183,7 +183,7 @@ namespace Anima {
             q4 = q_norm->forward(ctx, q4);
             k4 = k_norm->forward(ctx, k4);
 
-            struct ggml_tensor* attn_out = nullptr;
+            ggml_tensor* attn_out = nullptr;
             if (pe_q != nullptr || pe_k != nullptr) {
                 if (pe_q == nullptr) {
                     pe_q = pe_k;
@@ -227,7 +227,7 @@ namespace Anima {
             blocks["layer2"] = std::make_shared<Linear>(hidden_dim, dim, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        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"]);
 
@@ -245,7 +245,7 @@ namespace Anima {
             blocks["2"] = std::make_shared<Linear>(hidden_dim, dim, true);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
+        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"]);
 
@@ -267,11 +267,11 @@ namespace Anima {
             blocks["mlp"]             = std::make_shared<AdapterMLP>(model_dim, model_dim * 4);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* context,
-                                    struct ggml_tensor* target_pe,
-                                    struct ggml_tensor* context_pe) {
+        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"]);
@@ -317,11 +317,11 @@ namespace Anima {
             blocks["norm"]     = std::make_shared<RMSNorm>(target_dim, 1e-6f);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* source_hidden_states,
-                                    struct ggml_tensor* target_input_ids,
-                                    struct ggml_tensor* target_pe,
-                                    struct ggml_tensor* source_pe) {
+        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);
@@ -360,12 +360,12 @@ namespace Anima {
             blocks["mlp"]                         = std::make_shared<AnimaMLP>(hidden_size, hidden_size * mlp_ratio);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* hidden_states,
-                                    struct ggml_tensor* encoder_hidden_states,
-                                    struct ggml_tensor* embedded_timestep,
-                                    struct ggml_tensor* temb,
-                                    struct ggml_tensor* image_pe) {
+        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"]);
@@ -402,10 +402,10 @@ namespace Anima {
             blocks["linear"]           = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, false);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* hidden_states,
-                                    struct ggml_tensor* embedded_timestep,
-                                    struct ggml_tensor* temb) {
+        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"]);
 
@@ -445,15 +445,15 @@ namespace Anima {
             blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
         }
 
-        struct ggml_tensor* forward(GGMLRunnerContext* ctx,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* timestep,
-                                    struct ggml_tensor* encoder_hidden_states,
-                                    struct ggml_tensor* image_pe,
-                                    struct ggml_tensor* t5_ids       = nullptr,
-                                    struct ggml_tensor* t5_weights   = nullptr,
-                                    struct ggml_tensor* adapter_q_pe = nullptr,
-                                    struct ggml_tensor* adapter_k_pe = nullptr) {
+        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"]);
@@ -553,7 +553,7 @@ namespace Anima {
             return "anima";
         }
 
-        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) {
             net.get_param_tensors(tensors, prefix + ".net");
         }
 
@@ -602,19 +602,18 @@ namespace Anima {
             return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
         }
 
-        struct ggml_cgraph* build_graph(struct ggml_tensor* x,
-                                        struct ggml_tensor* timesteps,
-                                        struct ggml_tensor* context,
-                                        struct ggml_tensor* t5_ids     = nullptr,
-                                        struct ggml_tensor* t5_weights = nullptr) {
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor    = {},
+                                 const sd::Tensor<int32_t>& t5_ids_tensor   = {},
+                                 const sd::Tensor<float>& t5_weights_tensor = {}) {
+            ggml_tensor* x          = make_input(x_tensor);
+            ggml_tensor* timesteps  = make_input(timesteps_tensor);
+            ggml_tensor* context    = make_optional_input(context_tensor);
+            ggml_tensor* t5_ids     = make_optional_input(t5_ids_tensor);
+            ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
             GGML_ASSERT(x->ne[3] == 1);
-            struct 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);
+            ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
 
             int64_t pad_h = (net.patch_size - x->ne[1] % net.patch_size) % net.patch_size;
             int64_t pad_w = (net.patch_size - x->ne[0] % net.patch_size) % net.patch_size;
@@ -667,18 +666,16 @@ namespace Anima {
             return gf;
         }
 
-        bool compute(int n_threads,
-                     struct ggml_tensor* x,
-                     struct ggml_tensor* timesteps,
-                     struct ggml_tensor* context,
-                     struct ggml_tensor* t5_ids      = nullptr,
-                     struct ggml_tensor* t5_weights  = nullptr,
-                     struct ggml_tensor** output     = nullptr,
-                     struct ggml_context* output_ctx = nullptr) {
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context    = {},
+                                  const sd::Tensor<int32_t>& t5_ids   = {},
+                                  const sd::Tensor<float>& t5_weights = {}) {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, t5_ids, t5_weights);
             };
-            return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
     };
 }  // namespace Anima

src/auto_encoder_kl.hpp

@@ -1,4 +1,4 @@
-#ifndef __AUTO_ENCODER_KL_HPP__
+#ifndef __AUTO_ENCODER_KL_HPP__
 #define __AUTO_ENCODER_KL_HPP__
 
 #include "vae.hpp"
@@ -29,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"]);
@@ -65,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) {
@@ -101,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"]);
@@ -178,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]
@@ -208,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);
     }
@@ -227,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
@@ -317,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"]);
@@ -435,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
@@ -549,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]
@@ -581,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"]);
 
@@ -613,6 +613,9 @@ public:
 
     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;
+        }
         return dd_config.z_channels * factor;
     }
 };
@@ -678,202 +681,117 @@ 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);
-
-        z = to_backend(z);
+    ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
+        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+        ggml_tensor* z  = make_input(z_tensor);
 
         auto runner_ctx = get_context();
 
-        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 {
+    sd::Tensor<float> _compute(const int n_threads,
+                               const sd::Tensor<float>& z,
+                               bool decode_graph) 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);
-        // print_ggml_tensor(z);
-        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
     }
 
-    ggml_tensor* gaussian_latent_sample(ggml_context* work_ctx, ggml_tensor* moments, std::shared_ptr<RNG> rng) {
+    sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& 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]);
-        struct 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);
-                        }
-                    }
-                }
-            }
-        }
+        auto chunks               = sd::ops::chunk(moments, 2, 2);
+        const auto& mean          = chunks[0];
+        const auto& logvar        = chunks[1];
+        sd::Tensor<float> stddev  = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -30.0f, 20.0f));
+        sd::Tensor<float> noise   = sd::Tensor<float>::randn_like(mean, rng);
+        sd::Tensor<float> latents = mean + stddev * noise;
         return latents;
     }
 
-    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
         if (sd_version_is_flux2(version)) {
             return vae_output;
         } else if (version == VERSION_SD1_PIX2PIX) {
-            return 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);
+            return sd::ops::chunk(vae_output, 2, 2)[0];
         } else {
-            return gaussian_latent_sample(work_ctx, vae_output, rng);
+            return gaussian_latent_sample(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
+    std::pair<sd::Tensor<float>, sd::Tensor<float>> get_latents_mean_std(const sd::Tensor<float>& latents, int channel_dim) {
+        GGML_ASSERT(channel_dim >= 0 && static_cast<size_t>(channel_dim) < static_cast<size_t>(latents.dim()));
         if (sd_version_is_flux2(version)) {
-            GGML_ASSERT(latents->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};
+            GGML_ASSERT(latents.shape()[channel_dim] == 128);
+            std::vector<int64_t> stats_shape(static_cast<size_t>(latents.dim()), 1);
+            stats_shape[static_cast<size_t>(channel_dim)] = latents.shape()[channel_dim];
+
+            auto mean_tensor = sd::Tensor<float>::from_vector({-0.0676f, -0.0715f, -0.0753f, -0.0745f, 0.0223f, 0.0180f, 0.0142f, 0.0184f,
+                                                               -0.0001f, -0.0063f, -0.0002f, -0.0031f, -0.0272f, -0.0281f, -0.0276f, -0.0290f,
+                                                               -0.0769f, -0.0672f, -0.0902f, -0.0892f, 0.0168f, 0.0152f, 0.0079f, 0.0086f,
+                                                               0.0083f, 0.0015f, 0.0003f, -0.0043f, -0.0439f, -0.0419f, -0.0438f, -0.0431f,
+                                                               -0.0102f, -0.0132f, -0.0066f, -0.0048f, -0.0311f, -0.0306f, -0.0279f, -0.0180f,
+                                                               0.0030f, 0.0015f, 0.0126f, 0.0145f, 0.0347f, 0.0338f, 0.0337f, 0.0283f,
+                                                               0.0020f, 0.0047f, 0.0047f, 0.0050f, 0.0123f, 0.0081f, 0.0081f, 0.0146f,
+                                                               0.0681f, 0.0679f, 0.0767f, 0.0732f, -0.0462f, -0.0474f, -0.0392f, -0.0511f,
+                                                               -0.0528f, -0.0477f, -0.0470f, -0.0517f, -0.0317f, -0.0316f, -0.0345f, -0.0283f,
+                                                               0.0510f, 0.0445f, 0.0578f, 0.0458f, -0.0412f, -0.0458f, -0.0487f, -0.0467f,
+                                                               -0.0088f, -0.0106f, -0.0088f, -0.0046f, -0.0376f, -0.0432f, -0.0436f, -0.0499f,
+                                                               0.0118f, 0.0166f, 0.0203f, 0.0279f, 0.0113f, 0.0129f, 0.0016f, 0.0072f,
+                                                               -0.0118f, -0.0018f, -0.0141f, -0.0054f, -0.0091f, -0.0138f, -0.0145f, -0.0187f,
+                                                               0.0323f, 0.0305f, 0.0259f, 0.0300f, 0.0540f, 0.0614f, 0.0495f, 0.0590f,
+                                                               -0.0511f, -0.0603f, -0.0478f, -0.0524f, -0.0227f, -0.0274f, -0.0154f, -0.0255f,
+                                                               -0.0572f, -0.0565f, -0.0518f, -0.0496f, 0.0116f, 0.0054f, 0.0163f, 0.0104f});
+            mean_tensor.reshape_(stats_shape);
+            auto std_tensor = sd::Tensor<float>::from_vector({1.8029f, 1.7786f, 1.7868f, 1.7837f, 1.7717f, 1.7590f, 1.7610f, 1.7479f,
+                                                              1.7336f, 1.7373f, 1.7340f, 1.7343f, 1.8626f, 1.8527f, 1.8629f, 1.8589f,
+                                                              1.7593f, 1.7526f, 1.7556f, 1.7583f, 1.7363f, 1.7400f, 1.7355f, 1.7394f,
+                                                              1.7342f, 1.7246f, 1.7392f, 1.7304f, 1.7551f, 1.7513f, 1.7559f, 1.7488f,
+                                                              1.8449f, 1.8454f, 1.8550f, 1.8535f, 1.8240f, 1.7813f, 1.7854f, 1.7945f,
+                                                              1.8047f, 1.7876f, 1.7695f, 1.7676f, 1.7782f, 1.7667f, 1.7925f, 1.7848f,
+                                                              1.7579f, 1.7407f, 1.7483f, 1.7368f, 1.7961f, 1.7998f, 1.7920f, 1.7925f,
+                                                              1.7780f, 1.7747f, 1.7727f, 1.7749f, 1.7526f, 1.7447f, 1.7657f, 1.7495f,
+                                                              1.7775f, 1.7720f, 1.7813f, 1.7813f, 1.8162f, 1.8013f, 1.8023f, 1.8033f,
+                                                              1.7527f, 1.7331f, 1.7563f, 1.7482f, 1.7610f, 1.7507f, 1.7681f, 1.7613f,
+                                                              1.7665f, 1.7545f, 1.7828f, 1.7726f, 1.7896f, 1.7999f, 1.7864f, 1.7760f,
+                                                              1.7613f, 1.7625f, 1.7560f, 1.7577f, 1.7783f, 1.7671f, 1.7810f, 1.7799f,
+                                                              1.7201f, 1.7068f, 1.7265f, 1.7091f, 1.7793f, 1.7578f, 1.7502f, 1.7455f,
+                                                              1.7587f, 1.7500f, 1.7525f, 1.7362f, 1.7616f, 1.7572f, 1.7444f, 1.7430f,
+                                                              1.7509f, 1.7610f, 1.7634f, 1.7612f, 1.7254f, 1.7135f, 1.7321f, 1.7226f,
+                                                              1.7664f, 1.7624f, 1.7718f, 1.7664f, 1.7457f, 1.7441f, 1.7569f, 1.7530f});
+            std_tensor.reshape_(stats_shape);
+            return {std::move(mean_tensor), std::move(std_tensor)};
         } else {
             GGML_ABORT("unknown version %d", version);
         }
     }
 
-    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        ggml_tensor* vae_latents = ggml_dup(work_ctx, latents);
+    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
         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);
-            });
+            int channel_dim                = 2;
+            auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
+            return (latents * std_tensor) / scale_factor + mean_tensor;
         }
-        return vae_latents;
+        return (latents / scale_factor) + shift_factor;
     }
 
-    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        ggml_tensor* diffusion_latents = ggml_dup(work_ctx, latents);
+    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
         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);
-            });
+            int channel_dim                = 2;
+            auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
+            return ((latents - mean_tensor) * scale_factor) / std_tensor;
         }
-        return diffusion_latents;
+        return (latents - shift_factor) * scale_factor;
     }
 
     int get_encoder_output_channels(int input_channels) {
@@ -881,29 +799,31 @@ struct AutoEncoderKL : public VAE {
     }
 
     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_ASSERT(work_ctx != nullptr);
+        ggml_context* ctx = ggml_init(params);
+        GGML_ASSERT(ctx != nullptr);
 
         {
             // CPU, x{1, 3, 64, 64}: Pass
             // CUDA, x{1, 3, 64, 64}: Pass, but sill get wrong result for some image, may be due to interlnal nan
             // CPU, x{2, 3, 64, 64}: Wrong result
             // CUDA, x{2, 3, 64, 64}: Wrong result, and different from CPU result
-            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;
+            sd::Tensor<float> x({64, 64, 3, 2});
+            x.fill_(0.5f);
+            print_sd_tensor(x);
+            sd::Tensor<float> out;
 
-            int64_t t0 = ggml_time_ms();
-            _compute(8, x, false, &out, work_ctx);
-            int64_t t1 = ggml_time_ms();
+            int64_t t0   = ggml_time_ms();
+            auto out_opt = _compute(8, x, false);
+            int64_t t1   = ggml_time_ms();
 
-            print_ggml_tensor(out);
+            GGML_ASSERT(!out_opt.empty());
+            out = std::move(out_opt);
+            print_sd_tensor(out);
             LOG_DEBUG("encode test done in %lldms", t1 - t0);
         }
 
@@ -912,19 +832,21 @@ struct AutoEncoderKL : public VAE {
             // CUDA, z{1, 4, 8, 8}: Pass
             // CPU, z{3, 4, 8, 8}: Wrong result
             // CUDA, z{3, 4, 8, 8}: Wrong result, and different from CPU result
-            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;
+            sd::Tensor<float> z({8, 8, 4, 1});
+            z.fill_(0.5f);
+            print_sd_tensor(z);
+            sd::Tensor<float> out;
 
-            int64_t t0 = ggml_time_ms();
-            _compute(8, z, true, &out, work_ctx);
-            int64_t t1 = ggml_time_ms();
+            int64_t t0   = ggml_time_ms();
+            auto out_opt = _compute(8, z, true);
+            int64_t t1   = ggml_time_ms();
 
-            print_ggml_tensor(out);
+            GGML_ASSERT(!out_opt.empty());
+            out = std::move(out_opt);
+            print_sd_tensor(out);
             LOG_DEBUG("decode test done in %lldms", t1 - t0);
         }
     };
 };
 
-#endif  // __AUTO_ENCODER_KL_HPP__
+#endif  // __AUTO_ENCODER_KL_HPP__

src/cache_dit.hpp

@@ -8,7 +8,9 @@
 #include <unordered_map>
 #include <vector>
 
+#include "condition_cache_utils.hpp"
 #include "ggml_extend.hpp"
+#include "tensor.hpp"
 
 struct DBCacheConfig {
     bool enabled                        = false;
@@ -771,35 +773,37 @@ struct CacheDitConditionState {
         return it != cache_diffs.end() && !it->second.diff.empty();
     }
 
-    void update_cache(const void* cond, const float* input, const float* output, size_t size) {
+    void update_cache(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
         CacheEntry& entry = cache_diffs[cond];
-        entry.diff.resize(size);
-        for (size_t i = 0; i < size; i++) {
-            entry.diff[i] = output[i] - input[i];
+        if (!sd::store_condition_cache_diff(&entry.diff, input, output)) {
+            entry.prev_input.clear();
+            entry.prev_output.clear();
+            entry.has_prev = false;
+            return;
         }
 
+        size_t size              = static_cast<size_t>(output.numel());
+        const float* input_data  = input.data();
+        const float* output_data = output.data();
         entry.prev_input.resize(size);
         entry.prev_output.resize(size);
         for (size_t i = 0; i < size; i++) {
-            entry.prev_input[i]  = input[i];
-            entry.prev_output[i] = output[i];
+            entry.prev_input[i]  = input_data[i];
+            entry.prev_output[i] = output_data[i];
         }
         entry.has_prev = true;
     }
 
-    void apply_cache(const void* cond, const float* input, float* output, size_t size) {
+    void apply_cache(const void* cond,
+                     const sd::Tensor<float>& input,
+                     sd::Tensor<float>* output) {
         auto it = cache_diffs.find(cond);
         if (it == cache_diffs.end() || it->second.diff.empty())
             return;
-        if (it->second.diff.size() != size)
-            return;
-
-        for (size_t i = 0; i < size; i++) {
-            output[i] = input[i] + it->second.diff[i];
-        }
+        sd::apply_condition_cache_diff(it->second.diff, input, output);
     }
 
-    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, const sd::Tensor<float>& input, sd::Tensor<float>* output, float sigma, int step_index) {
         if (!enabled() || step_index < 0)
             return false;
 
@@ -819,8 +823,7 @@ struct CacheDitConditionState {
 
         if (skip_current_step) {
             if (has_cache(cond)) {
-                apply_cache(cond, (float*)input->data, (float*)output->data,
-                            static_cast<size_t>(ggml_nelements(output)));
+                apply_cache(cond, input, output);
                 return true;
             }
             return false;
@@ -833,13 +836,13 @@ struct CacheDitConditionState {
         if (it == cache_diffs.end() || !it->second.has_prev)
             return false;
 
-        size_t ne = static_cast<size_t>(ggml_nelements(input));
+        size_t ne = static_cast<size_t>(input.numel());
         if (it->second.prev_input.size() != ne)
             return false;
 
-        float* input_data = (float*)input->data;
-        float diff        = CacheDitState::calculate_residual_diff(
-                   it->second.prev_input.data(), input_data, ne);
+        const float* input_data = input.data();
+        float diff              = CacheDitState::calculate_residual_diff(
+                         it->second.prev_input.data(), input_data, ne);
 
         float effective_threshold = config.residual_diff_threshold;
         if (config.Fn_compute_blocks > 0) {
@@ -859,7 +862,7 @@ struct CacheDitConditionState {
             cached_steps.push_back(current_step_index);
             continuous_cached_steps++;
             accumulated_residual_diff += diff;
-            apply_cache(cond, input_data, (float*)output->data, ne);
+            apply_cache(cond, input, output);
             return true;
         }
 
@@ -867,15 +870,14 @@ struct CacheDitConditionState {
         return false;
     }
 
-    void after_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output) {
+    void after_condition(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
         if (!step_is_active())
             return;
 
-        size_t ne = static_cast<size_t>(ggml_nelements(output));
-        update_cache(cond, (float*)input->data, (float*)output->data, ne);
+        update_cache(cond, input, output);
 
         if (cond == anchor_condition && taylor_config.enabled) {
-            taylor_state.update_derivatives((float*)output->data, ne, current_step_index);
+            taylor_state.update_derivatives(output.data(), static_cast<size_t>(output.numel()), current_step_index);
         }
     }
 

src/clip.hpp

@@ -473,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"]);
@@ -511,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"]);
@@ -541,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);
@@ -573,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);
@@ -597,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"];
@@ -630,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;
@@ -653,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);
@@ -663,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]
     }
 };
@@ -693,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);
@@ -734,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"]);
@@ -804,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"]);
@@ -839,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);
@@ -856,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));
         }
@@ -886,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"]);
@@ -936,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) {
@@ -957,17 +957,16 @@ 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(const sd::Tensor<int32_t>& input_ids_tensor,
+                             int num_custom_embeddings    = 0,
+                             void* custom_embeddings_data = nullptr,
+                             size_t max_token_idx         = 0,
+                             bool return_pooled           = false,
+                             int clip_skip                = -1) {
+        ggml_cgraph* gf        = new_graph_custom(2048);
+        ggml_tensor* input_ids = make_input(input_ids_tensor);
 
-        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();
@@ -997,26 +996,28 @@ 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);
 
         return gf;
     }
 
-    bool compute(const int n_threads,
-                 struct ggml_tensor* input_ids,
-                 int num_custom_embeddings,
-                 void* custom_embeddings_data,
-                 size_t max_token_idx,
-                 bool return_pooled,
-                 int clip_skip,
-                 ggml_tensor** output,
-                 ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
+    sd::Tensor<float> compute(const int n_threads,
+                              const sd::Tensor<int32_t>& input_ids,
+                              int num_custom_embeddings,
+                              void* custom_embeddings_data,
+                              size_t max_token_idx,
+                              bool return_pooled,
+                              int clip_skip) {
+        auto get_graph = [&]() -> ggml_cgraph* {
             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);
+        auto result = GGMLRunner::compute<float>(get_graph, n_threads, true);
+        if (return_pooled) {
+            return take_or_empty(std::move(result));
+        }
+        return restore_trailing_singleton_dims(std::move(result), 3);
     }
 };
 

src/common_block.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.])

src/common_dit.hpp

@@ -4,11 +4,11 @@
 #include "ggml_extend.hpp"
 
 namespace DiT {
-    ggml_tensor* patchify(ggml_context* ctx,
-                          ggml_tensor* x,
-                          int pw,
-                          int ph,
-                          bool patch_last = true) {
+    inline ggml_tensor* patchify(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 int pw,
+                                 int ph,
+                                 bool patch_last = true) {
         // x: [N, C, H, W]
         // return: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
         int64_t N = x->ne[3];
@@ -33,13 +33,13 @@ namespace DiT {
         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) {
+    inline ggml_tensor* unpatchify(ggml_context* ctx,
+                                   ggml_tensor* x,
+                                   int64_t h,
+                                   int64_t w,
+                                   int ph,
+                                   int pw,
+                                   bool patch_last = true) {
         // x: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
         // return: [N, C, H, W]
         int64_t N = x->ne[2];
@@ -64,10 +64,10 @@ namespace DiT {
         return x;
     }
 
-    ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
-                                   ggml_tensor* x,
-                                   int ph,
-                                   int pw) {
+    inline ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
+                                          ggml_tensor* x,
+                                          int ph,
+                                          int pw) {
         int64_t W = x->ne[0];
         int64_t H = x->ne[1];
 
@@ -77,23 +77,23 @@ namespace DiT {
         return x;
     }
 
-    ggml_tensor* pad_and_patchify(GGMLRunnerContext* ctx,
-                                  ggml_tensor* x,
-                                  int ph,
-                                  int pw,
-                                  bool patch_last = true) {
+    inline ggml_tensor* pad_and_patchify(GGMLRunnerContext* ctx,
+                                         ggml_tensor* x,
+                                         int ph,
+                                         int pw,
+                                         bool patch_last = true) {
         x = pad_to_patch_size(ctx, x, ph, pw);
         x = patchify(ctx->ggml_ctx, x, ph, pw, patch_last);
         return x;
     }
 
-    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) {
+    inline ggml_tensor* unpatchify_and_crop(ggml_context* ctx,
+                                            ggml_tensor* x,
+                                            int64_t H,
+                                            int64_t W,
+                                            int ph,
+                                            int pw,
+                                            bool patch_last = true) {
         int pad_h = (ph - H % ph) % ph;
         int pad_w = (pw - W % pw) % pw;
         int64_t h = ((H + pad_h) / ph);
@@ -105,4 +105,4 @@ namespace DiT {
     }
 }  // namespace DiT
 
-#endif  // __COMMON_DIT_HPP__
+#endif  // __COMMON_DIT_HPP__

src/condition_cache_utils.hpp

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

src/control.hpp

@@ -164,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++) {
@@ -198,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]
@@ -246,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);
@@ -310,11 +310,13 @@ struct ControlNet : public GGMLRunner {
     SDVersion version = VERSION_SD1;
     ControlNetBlock control_net;
 
-    ggml_backend_buffer_t control_buffer = nullptr;  // keep control output tensors in backend memory
+    ggml_backend_buffer_t control_buffer = nullptr;
     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*> control_outputs_ggml;
+    ggml_tensor* guided_hint_output_ggml = nullptr;
+    std::vector<sd::Tensor<float>> controls;
+    sd::Tensor<float> guided_hint;
+    bool guided_hint_cached = false;
 
     ControlNet(ggml_backend_t backend,
                bool offload_params_to_cpu,
@@ -328,23 +330,23 @@ 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;
         control_ctx       = ggml_init(params);
 
-        controls.resize(outs.size() - 1);
+        control_outputs_ggml.resize(outs.size() - 1);
 
         size_t control_buffer_size = 0;
 
-        guided_hint = ggml_dup_tensor(control_ctx, outs[0]);
-        control_buffer_size += ggml_nbytes(guided_hint);
+        guided_hint_output_ggml = ggml_dup_tensor(control_ctx, outs[0]);
+        control_buffer_size += ggml_nbytes(guided_hint_output_ggml);
 
         for (int i = 0; i < outs.size() - 1; i++) {
-            controls[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
-            control_buffer_size += ggml_nbytes(controls[i]);
+            control_outputs_ggml[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
+            control_buffer_size += ggml_nbytes(control_outputs_ggml[i]);
         }
 
         control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, runtime_backend);
@@ -361,8 +363,10 @@ struct ControlNet : public GGMLRunner {
             ggml_free(control_ctx);
             control_ctx = nullptr;
         }
-        guided_hint        = nullptr;
-        guided_hint_cached = false;
+        guided_hint_output_ggml = nullptr;
+        guided_hint_cached      = false;
+        guided_hint             = {};
+        control_outputs_ggml.clear();
         controls.clear();
     }
 
@@ -370,33 +374,37 @@ 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(const sd::Tensor<float>& x_tensor,
+                             const sd::Tensor<float>& hint_tensor,
+                             const sd::Tensor<float>& timesteps_tensor,
+                             const sd::Tensor<float>& context_tensor = {},
+                             const sd::Tensor<float>& y_tensor       = {}) {
+        ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
 
-        x = to_backend(x);
-        if (guided_hint_cached) {
-            hint = nullptr;
+        ggml_tensor* x         = make_input(x_tensor);
+        ggml_tensor* hint      = nullptr;
+        ggml_tensor* timesteps = make_input(timesteps_tensor);
+        ggml_tensor* context   = make_optional_input(context_tensor);
+        ggml_tensor* y         = make_optional_input(y_tensor);
+
+        ggml_tensor* guided_hint_input = nullptr;
+        if (guided_hint_cached && !guided_hint.empty()) {
+            guided_hint_input = make_input(guided_hint);
+            hint              = nullptr;
         } else {
-            hint = to_backend(hint);
+            hint = make_input(hint_tensor);
         }
-        context   = to_backend(context);
-        y         = to_backend(y);
-        timesteps = to_backend(timesteps);
 
         auto runner_ctx = get_context();
 
         auto outs = control_net.forward(&runner_ctx,
                                         x,
                                         hint,
-                                        guided_hint_cached ? guided_hint : nullptr,
+                                        guided_hint_input,
                                         timesteps,
                                         context,
                                         y);
@@ -405,36 +413,46 @@ struct ControlNet : public GGMLRunner {
             alloc_control_ctx(outs);
         }
 
-        ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint));
+        ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint_output_ggml));
         for (int i = 0; i < outs.size() - 1; i++) {
-            ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], controls[i]));
+            ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], control_outputs_ggml[i]));
         }
 
         return gf;
     }
 
-    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) {
+    std::optional<std::vector<sd::Tensor<float>>> compute(int n_threads,
+                                                          const sd::Tensor<float>& x,
+                                                          const sd::Tensor<float>& hint,
+                                                          const sd::Tensor<float>& timesteps,
+                                                          const sd::Tensor<float>& context = {},
+                                                          const sd::Tensor<float>& y       = {}) {
         // 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);
         };
 
-        bool res = GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
-        if (res) {
-            // cache guided_hint
-            guided_hint_cached = true;
+        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false);
+        if (!compute_result.has_value()) {
+            return std::nullopt;
         }
-        return res;
+
+        if (guided_hint_output_ggml != nullptr) {
+            guided_hint = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(guided_hint_output_ggml),
+                                                          4);
+        }
+        controls.clear();
+        controls.reserve(control_outputs_ggml.size());
+        for (ggml_tensor* control : control_outputs_ggml) {
+            auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
+            GGML_ASSERT(!control_host.empty());
+            controls.push_back(std::move(control_host));
+        }
+        guided_hint_cached = true;
+        return controls;
     }
 
     bool load_from_file(const std::string& file_path, int n_threads) {
@@ -462,4 +480,4 @@ struct ControlNet : public GGMLRunner {
     }
 };
 
-#endif  // __CONTROL_HPP__
+#endif  // __CONTROL_HPP__

src/diffusion_model.hpp

@@ -1,42 +1,50 @@
 #ifndef __DIFFUSION_MODEL_H__
 #define __DIFFUSION_MODEL_H__
 
+#include <optional>
 #include "anima.hpp"
 #include "flux.hpp"
 #include "mmdit.hpp"
 #include "qwen_image.hpp"
+#include "tensor_ggml.hpp"
 #include "unet.hpp"
 #include "wan.hpp"
 #include "z_image.hpp"
 
 struct DiffusionParams {
-    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              = {};
+    const sd::Tensor<float>* x                        = nullptr;
+    const sd::Tensor<float>* timesteps                = nullptr;
+    const sd::Tensor<float>* context                  = nullptr;
+    const sd::Tensor<float>* c_concat                 = nullptr;
+    const sd::Tensor<float>* y                        = nullptr;
+    const sd::Tensor<int32_t>* t5_ids                 = nullptr;
+    const sd::Tensor<float>* t5_weights               = nullptr;
+    const sd::Tensor<float>* guidance                 = nullptr;
+    const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
+    bool increase_ref_index                           = false;
+    int num_video_frames                              = -1;
+    const std::vector<sd::Tensor<float>>* controls    = nullptr;
+    float control_strength                            = 0.f;
+    const sd::Tensor<float>* vace_context             = nullptr;
+    float vace_strength                               = 1.f;
+    const std::vector<int>* skip_layers               = nullptr;
 };
 
+template <typename T>
+static inline const sd::Tensor<T>& tensor_or_empty(const sd::Tensor<T>* tensor) {
+    static const sd::Tensor<T> kEmpty;
+    return tensor != nullptr ? *tensor : kEmpty;
+}
+
 struct DiffusionModel {
-    virtual std::string get_desc()                                                      = 0;
-    virtual bool compute(int n_threads,
-                         DiffusionParams diffusion_params,
-                         struct ggml_tensor** output     = nullptr,
-                         struct ggml_context* output_ctx = nullptr)                     = 0;
-    virtual void alloc_params_buffer()                                                  = 0;
-    virtual void free_params_buffer()                                                   = 0;
-    virtual void free_compute_buffer()                                                  = 0;
-    virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) = 0;
-    virtual size_t get_params_buffer_size()                                             = 0;
+    virtual std::string get_desc()                                               = 0;
+    virtual sd::Tensor<float> compute(int n_threads,
+                                      const DiffusionParams& diffusion_params)   = 0;
+    virtual void alloc_params_buffer()                                           = 0;
+    virtual void free_params_buffer()                                            = 0;
+    virtual void free_compute_buffer()                                           = 0;
+    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
+    virtual size_t get_params_buffer_size()                                      = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
     virtual int64_t get_adm_in_channels()                            = 0;
     virtual void set_flash_attention_enabled(bool enabled)           = 0;
@@ -69,7 +77,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");
     }
 
@@ -93,19 +101,20 @@ struct UNetModel : public DiffusionModel {
         unet.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        static const std::vector<sd::Tensor<float>> empty_controls;
         return unet.compute(n_threads,
-                            diffusion_params.x,
-                            diffusion_params.timesteps,
-                            diffusion_params.context,
-                            diffusion_params.c_concat,
-                            diffusion_params.y,
+                            *diffusion_params.x,
+                            *diffusion_params.timesteps,
+                            tensor_or_empty(diffusion_params.context),
+                            tensor_or_empty(diffusion_params.c_concat),
+                            tensor_or_empty(diffusion_params.y),
                             diffusion_params.num_video_frames,
-                            diffusion_params.controls,
-                            diffusion_params.control_strength, output, output_ctx);
+                            diffusion_params.controls ? *diffusion_params.controls : empty_controls,
+                            diffusion_params.control_strength);
     }
 };
 
@@ -134,7 +143,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");
     }
 
@@ -158,18 +167,17 @@ struct MMDiTModel : public DiffusionModel {
         mmdit.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        static const std::vector<int> empty_skip_layers;
         return mmdit.compute(n_threads,
-                             diffusion_params.x,
-                             diffusion_params.timesteps,
-                             diffusion_params.context,
-                             diffusion_params.y,
-                             output,
-                             output_ctx,
-                             diffusion_params.skip_layers);
+                             *diffusion_params.x,
+                             *diffusion_params.timesteps,
+                             tensor_or_empty(diffusion_params.context),
+                             tensor_or_empty(diffusion_params.y),
+                             diffusion_params.skip_layers ? *diffusion_params.skip_layers : empty_skip_layers);
     }
 };
 
@@ -200,7 +208,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");
     }
 
@@ -224,22 +232,22 @@ struct FluxModel : public DiffusionModel {
         flux.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        static const std::vector<sd::Tensor<float>> empty_ref_latents;
+        static const std::vector<int> empty_skip_layers;
         return flux.compute(n_threads,
-                            diffusion_params.x,
-                            diffusion_params.timesteps,
-                            diffusion_params.context,
-                            diffusion_params.c_concat,
-                            diffusion_params.y,
-                            diffusion_params.guidance,
-                            diffusion_params.ref_latents,
+                            *diffusion_params.x,
+                            *diffusion_params.timesteps,
+                            tensor_or_empty(diffusion_params.context),
+                            tensor_or_empty(diffusion_params.c_concat),
+                            tensor_or_empty(diffusion_params.y),
+                            tensor_or_empty(diffusion_params.guidance),
+                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
                             diffusion_params.increase_ref_index,
-                            output,
-                            output_ctx,
-                            diffusion_params.skip_layers);
+                            diffusion_params.skip_layers ? *diffusion_params.skip_layers : empty_skip_layers);
     }
 };
 
@@ -270,7 +278,7 @@ struct AnimaModel : public DiffusionModel {
         anima.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 {
         anima.get_param_tensors(tensors, prefix);
     }
 
@@ -294,18 +302,16 @@ struct AnimaModel : public DiffusionModel {
         anima.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
         return anima.compute(n_threads,
-                             diffusion_params.x,
-                             diffusion_params.timesteps,
-                             diffusion_params.context,
-                             diffusion_params.c_concat,
-                             diffusion_params.y,
-                             output,
-                             output_ctx);
+                             *diffusion_params.x,
+                             *diffusion_params.timesteps,
+                             tensor_or_empty(diffusion_params.context),
+                             tensor_or_empty(diffusion_params.t5_ids),
+                             tensor_or_empty(diffusion_params.t5_weights));
     }
 };
 
@@ -337,7 +343,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);
     }
 
@@ -361,21 +367,19 @@ struct WanModel : public DiffusionModel {
         wan.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
         return wan.compute(n_threads,
-                           diffusion_params.x,
-                           diffusion_params.timesteps,
-                           diffusion_params.context,
-                           diffusion_params.y,
-                           diffusion_params.c_concat,
-                           nullptr,
-                           diffusion_params.vace_context,
-                           diffusion_params.vace_strength,
-                           output,
-                           output_ctx);
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           tensor_or_empty(diffusion_params.y),
+                           tensor_or_empty(diffusion_params.c_concat),
+                           sd::Tensor<float>(),
+                           tensor_or_empty(diffusion_params.vace_context),
+                           diffusion_params.vace_strength);
     }
 };
 
@@ -408,7 +412,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);
     }
 
@@ -432,18 +436,17 @@ struct QwenImageModel : public DiffusionModel {
         qwen_image.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        static const std::vector<sd::Tensor<float>> empty_ref_latents;
         return qwen_image.compute(n_threads,
-                                  diffusion_params.x,
-                                  diffusion_params.timesteps,
-                                  diffusion_params.context,
-                                  diffusion_params.ref_latents,
-                                  true,  // increase_ref_index
-                                  output,
-                                  output_ctx);
+                                  *diffusion_params.x,
+                                  *diffusion_params.timesteps,
+                                  tensor_or_empty(diffusion_params.context),
+                                  diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                                  true);
     }
 };
 
@@ -475,7 +478,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);
     }
 
@@ -499,18 +502,17 @@ struct ZImageModel : public DiffusionModel {
         z_image.set_circular_axes(circular_x, circular_y);
     }
 
-    bool compute(int n_threads,
-                 DiffusionParams diffusion_params,
-                 struct ggml_tensor** output     = nullptr,
-                 struct ggml_context* output_ctx = nullptr) override {
+    sd::Tensor<float> compute(int n_threads,
+                              const DiffusionParams& diffusion_params) override {
+        GGML_ASSERT(diffusion_params.x != nullptr);
+        GGML_ASSERT(diffusion_params.timesteps != nullptr);
+        static const std::vector<sd::Tensor<float>> empty_ref_latents;
         return z_image.compute(n_threads,
-                               diffusion_params.x,
-                               diffusion_params.timesteps,
-                               diffusion_params.context,
-                               diffusion_params.ref_latents,
-                               true,  // increase_ref_index
-                               output,
-                               output_ctx);
+                               *diffusion_params.x,
+                               *diffusion_params.timesteps,
+                               tensor_or_empty(diffusion_params.context),
+                               diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                               true);
     }
 };
 

src/easycache.hpp

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

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,28 +341,25 @@ struct ESRGAN : public GGMLRunner {
         return success;
     }
 
-    struct ggml_cgraph* build_graph(struct ggml_tensor* x) {
+    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
         if (!rrdb_net)
             return nullptr;
         constexpr int kGraphNodes = 1 << 16;  // 65k
-        struct ggml_cgraph* gf    = new_graph_custom(kGraphNodes);
-        x                         = to_backend(x);
+        ggml_cgraph* gf           = new_graph_custom(kGraphNodes);
+        ggml_tensor* x            = make_input(x_tensor);
 
-        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** output,
-                 ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(x);
-        };
-        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+    sd::Tensor<float> compute(const int n_threads,
+                              const sd::Tensor<float>& x) {
+        auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
+        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return result;
     }
 };
 
-#endif  // __ESRGAN_HPP__
+#endif  // __ESRGAN_HPP__

src/flux.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,7 +98,7 @@ 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"]);
 
@@ -115,17 +115,17 @@ namespace Flux {
             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]
@@ -147,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"]);
 
@@ -170,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"]);
@@ -212,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"]);
@@ -232,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]
@@ -294,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;
@@ -303,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;
@@ -312,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]
@@ -457,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]
@@ -539,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]
@@ -548,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;
@@ -589,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"]);
 
@@ -612,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)
@@ -639,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)
@@ -689,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"]);
 
@@ -708,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"]);
@@ -847,15 +847,15 @@ namespace Flux {
             }
         }
 
-        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"]);
@@ -864,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"]);
@@ -967,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];
@@ -1050,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];
@@ -1119,18 +1119,18 @@ namespace Flux {
             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
@@ -1178,6 +1178,7 @@ namespace Flux {
         std::vector<float> pe_vec;
         std::vector<float> mod_index_arange_vec;
         std::vector<float> dct_vec;
+        sd::Tensor<float> guidance_tensor;
         SDVersion version;
         bool use_mask = false;
 
@@ -1299,7 +1300,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);
         }
 
@@ -1353,29 +1354,42 @@ 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_ASSERT(x->ne[3] == 1);
-            struct ggml_cgraph* gf = new_graph_custom(FLUX_GRAPH_SIZE);
-
-            struct ggml_tensor* mod_index_arange = nullptr;
-            struct ggml_tensor* dct              = nullptr;  // for chroma radiance
-
-            x       = to_backend(x);
-            context = to_backend(context);
-            if (c_concat != nullptr) {
-                c_concat = to_backend(c_concat);
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor                  = {},
+                                 const sd::Tensor<float>& c_concat_tensor                 = {},
+                                 const sd::Tensor<float>& y_tensor                        = {},
+                                 const sd::Tensor<float>& guidance_tensor                 = {},
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
+                                 bool increase_ref_index                                  = false,
+                                 std::vector<int> skip_layers                             = {}) {
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            ggml_tensor* context   = make_optional_input(context_tensor);
+            ggml_tensor* c_concat  = make_optional_input(c_concat_tensor);
+            ggml_tensor* y         = make_optional_input(y_tensor);
+            if (flux_params.guidance_embed || flux_params.is_chroma) {
+                if (!guidance_tensor.empty()) {
+                    this->guidance_tensor = guidance_tensor;
+                    if (flux_params.is_chroma) {
+                        this->guidance_tensor.fill_(0.f);
+                    }
+                }
+            }
+            ggml_tensor* guidance = make_optional_input(this->guidance_tensor);
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
             }
-            if (flux_params.is_chroma) {
-                guidance = ggml_set_f32(guidance, 0);
 
+            GGML_ASSERT(x->ne[3] == 1);
+            ggml_cgraph* gf = new_graph_custom(FLUX_GRAPH_SIZE);
+
+            ggml_tensor* mod_index_arange = nullptr;
+            ggml_tensor* dct              = nullptr;  // for chroma radiance
+
+            if (flux_params.is_chroma) {
                 if (!use_mask) {
                     y = nullptr;
                 }
@@ -1385,16 +1399,6 @@ namespace Flux {
                 mod_index_arange     = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, mod_index_arange_vec.size());
                 set_backend_tensor_data(mod_index_arange, mod_index_arange_vec.data());
             }
-            y = to_backend(y);
-
-            timesteps = to_backend(timesteps);
-            if (flux_params.guidance_embed || flux_params.is_chroma) {
-                guidance = to_backend(guidance);
-            }
-            for (int i = 0; i < ref_latents.size(); i++) {
-                ref_latents[i] = to_backend(ref_latents[i]);
-            }
-
             std::set<int> txt_arange_dims;
             if (sd_version_is_flux2(version)) {
                 txt_arange_dims    = {3};
@@ -1437,89 +1441,98 @@ 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);
 
             return gf;
         }
 
-        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,
-                     std::vector<ggml_tensor*> ref_latents = {},
-                     bool increase_ref_index               = false,
-                     struct ggml_tensor** output           = nullptr,
-                     struct ggml_context* output_ctx       = nullptr,
-                     std::vector<int> skip_layers          = std::vector<int>()) {
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context                  = {},
+                                  const sd::Tensor<float>& c_concat                 = {},
+                                  const sd::Tensor<float>& y                        = {},
+                                  const sd::Tensor<float>& guidance                 = {},
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {},
+                                  bool increase_ref_index                           = false,
+                                  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);
             };
 
-            return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return result;
         }
 
         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_ASSERT(work_ctx != nullptr);
+            ggml_context* ctx = ggml_init(params);
+            GGML_ASSERT(ctx != nullptr);
 
             {
                 // cpu f16:
                 // cuda f16: nan
                 // cuda q8_0: pass
-                auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 16, 16, 128, 1);
+                sd::Tensor<float> x({16, 16, 128, 1});
                 // ggml_set_f32(x, 0.01f);
-                // auto x = load_tensor_from_file(work_ctx, "chroma_x.bin");
+                // auto x = load_tensor_from_file(ctx, "chroma_x.bin");
                 // print_ggml_tensor(x);
 
                 std::vector<float> timesteps_vec(1, 1.f);
-                auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
+                auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
 
                 std::vector<float> guidance_vec(1, 0.f);
-                auto guidance = vector_to_ggml_tensor(work_ctx, guidance_vec);
+                auto guidance = sd::Tensor<float>::from_vector(guidance_vec);
 
-                auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 15360, 256, 1);
+                sd::Tensor<float> context({15360, 256, 1});
                 // ggml_set_f32(context, 0.01f);
-                // auto context = load_tensor_from_file(work_ctx, "chroma_context.bin");
+                // auto context = load_tensor_from_file(ctx, "chroma_context.bin");
                 // print_ggml_tensor(context);
 
-                // auto y = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 768, 1);
+                // auto y = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 768, 1);
                 // ggml_set_f32(y, 0.01f);
                 auto y = nullptr;
                 // print_ggml_tensor(y);
 
-                struct ggml_tensor* out = nullptr;
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                compute(8, x, timesteps, context, nullptr, y, guidance, {}, false, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = compute(8,
+                                       x,
+                                       timesteps,
+                                       context,
+                                       {},
+                                       {},
+                                       guidance,
+                                       {},
+                                       false);
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("flux test done in %lldms", t1 - t0);
             }
         }

src/latent-preview.h

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

src/llm.hpp

@@ -194,6 +194,7 @@ namespace LLM {
                         bool padding      = false) {
             if (add_bos_token) {
                 tokens.insert(tokens.begin(), BOS_TOKEN_ID);
+                weights.insert(weights.begin(), 1.f);
             }
             if (max_length > 0 && padding) {
                 size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.f / max_length));
@@ -522,7 +523,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"]);
@@ -582,7 +583,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,
@@ -631,7 +632,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"]);
@@ -668,10 +669,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];
@@ -718,10 +719,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"]);
@@ -778,12 +779,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)]
@@ -836,10 +837,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];
@@ -898,10 +899,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"]);
@@ -936,12 +937,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]
 
@@ -1037,12 +1038,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"]);
 
@@ -1050,12 +1051,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);
@@ -1156,40 +1157,41 @@ 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);
-
-            input_ids = to_backend(input_ids);
-
-            for (auto& image_embed : image_embeds) {
-                image_embed.second = to_backend(image_embed.second);
+        ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
+                                 const sd::Tensor<float>& attention_mask_tensor,
+                                 const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
+                                 std::set<int> out_layers) {
+            ggml_cgraph* gf        = ggml_new_graph(compute_ctx);
+            ggml_tensor* input_ids = make_input(input_ids_tensor);
+            std::vector<std::pair<int, ggml_tensor*>> image_embeds;
+            image_embeds.reserve(image_embeds_tensor.size());
+            for (const auto& [idx, embed_tensor] : image_embeds_tensor) {
+                ggml_tensor* embed = make_input(embed_tensor);
+                image_embeds.emplace_back(idx, embed);
             }
 
             int64_t n_tokens = input_ids->ne[0];
@@ -1213,8 +1215,9 @@ namespace LLM {
                                                 input_pos_vec.size());
             set_backend_tensor_data(input_pos, input_pos_vec.data());
 
-            if (attention_mask != nullptr) {
-                attention_mask = to_backend(attention_mask);
+            ggml_tensor* attention_mask = nullptr;
+            if (!attention_mask_tensor.empty()) {
+                attention_mask = make_input(attention_mask_tensor);
             } else {
                 attention_mask_vec.resize(n_tokens * n_tokens);
                 for (int i0 = 0; i0 < n_tokens; i0++) {
@@ -1232,24 +1235,22 @@ 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);
 
             return gf;
         }
 
-        bool compute(const int n_threads,
-                     struct ggml_tensor* input_ids,
-                     struct 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* {
+        sd::Tensor<float> compute(const int n_threads,
+                                  const sd::Tensor<int32_t>& input_ids,
+                                  const sd::Tensor<float>& attention_mask,
+                                  const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
+                                  std::set<int> out_layers) {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input_ids, attention_mask, image_embeds, out_layers);
             };
-            return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
         }
 
         int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
@@ -1261,7 +1262,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];
@@ -1288,8 +1289,9 @@ 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(const sd::Tensor<float>& image_tensor) {
+            ggml_cgraph* gf    = new_graph_custom(LLM_GRAPH_SIZE);
+            ggml_tensor* image = make_input(image_tensor);
 
             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);
@@ -1301,8 +1303,6 @@ namespace LLM {
             int llm_grid_w             = grid_w / params.vision.spatial_merge_size;
             int vit_merger_window_size = params.vision.window_size / params.vision.patch_size / params.vision.spatial_merge_size;
 
-            image = to_backend(image);
-
             auto pixel_values = process_image(compute_ctx, image);
 
             // window index
@@ -1399,26 +1399,24 @@ 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** output,
-                          ggml_context* output_ctx = nullptr) {
-            auto get_graph = [&]() -> struct ggml_cgraph* {
+        sd::Tensor<float> encode_image(const int n_threads,
+                                       const sd::Tensor<float>& image) {
+            auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
         }
     };
 
@@ -1440,7 +1438,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);
         }
 
@@ -1492,44 +1490,46 @@ 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_ASSERT(work_ctx != nullptr);
+            ggml_context* ctx = ggml_init(params);
+            GGML_ASSERT(ctx != nullptr);
             bool test_mistral          = false;
             bool test_qwen3            = true;
             bool test_vit              = false;
             bool test_decoder_with_vit = false;
 
             if (test_decoder_with_vit) {
-                ggml_tensor* image_embed = nullptr;
+                sd::Tensor<float> image_embed;
                 {
-                    auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
-                    print_ggml_tensor(image, false, "image");
-                    struct ggml_tensor* out = nullptr;
+                    auto image = sd::load_tensor_from_file_as_tensor<float>("qwen2vl_normalized.bin");
+                    print_sd_tensor(image, false, "image");
+                    sd::Tensor<float> out;
 
-                    int64_t t0 = ggml_time_ms();
-                    model.encode_image(8, image, &out, work_ctx);
-                    int64_t t1 = ggml_time_ms();
+                    int64_t t0   = ggml_time_ms();
+                    auto out_opt = model.encode_image(8, image);
+                    int64_t t1   = ggml_time_ms();
 
-                    print_ggml_tensor(out, false, "image_embed");
+                    GGML_ASSERT(!out_opt.empty());
+                    out = std::move(out_opt);
+                    print_sd_tensor(out, false, "image_embed");
                     image_embed = out;
                     LOG_DEBUG("llm encode_image test done in %lldms", t1 - t0);
                 }
 
                 std::string placeholder  = "<|image_pad|>";
                 std::string img_prompt   = "Picture 1: <|vision_start|>";  // [24669, 220, 16, 25, 220, 151652]
-                int64_t num_image_tokens = image_embed->ne[1];
+                int64_t num_image_tokens = image_embed.shape()[1];
                 img_prompt.reserve(num_image_tokens * placeholder.size());
                 for (int i = 0; i < num_image_tokens; i++) {
                     img_prompt += placeholder;
                 }
                 img_prompt += "<|vision_end|>";
 
-                std::vector<std::pair<int, ggml_tensor*>> image_embeds;
+                std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
                 image_embeds.emplace_back(64, image_embed);
 
                 std::pair<int, int> prompt_attn_range;
@@ -1547,29 +1547,33 @@ 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 = sd::Tensor<int32_t>::from_vector(tokens);
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                model.compute(8, input_ids, nullptr, image_embeds, {}, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), image_embeds, {});
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("llm test done in %lldms", t1 - t0);
             } else if (test_vit) {
-                // auto image = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 280, 280, 3);
+                // auto image = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 280, 280, 3);
                 // 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;
+                auto image = sd::load_tensor_from_file_as_tensor<float>("qwen2vl_normalized.bin");
+                print_sd_tensor(image, false, "image");
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                model.encode_image(8, image, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = model.encode_image(8, image);
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out, false, "out");
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out, false, "out");
 
-                // auto ref_out = load_tensor_from_file(work_ctx, "qwen2vl.bin");
+                // auto ref_out = load_tensor_from_file(ctx, "qwen2vl.bin");
                 // ggml_ext_tensor_diff(ref_out, out, 0.01f);
 
                 LOG_DEBUG("llm test done in %lldms", t1 - t0);
@@ -1587,14 +1591,16 @@ 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 = sd::Tensor<int32_t>::from_vector(tokens);
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                model.compute(8, input_ids, nullptr, {}, {10, 20, 30}, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), {}, {10, 20, 30});
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("llm test done in %lldms", t1 - t0);
             } else if (test_qwen3) {
                 std::pair<int, int> prompt_attn_range;
@@ -1610,14 +1616,16 @@ 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 = sd::Tensor<int32_t>::from_vector(tokens);
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                model.compute(8, input_ids, nullptr, {}, {35}, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), {}, {35});
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("llm test done in %lldms", t1 - t0);
             } else {
                 std::pair<int, int> prompt_attn_range;
@@ -1633,14 +1641,16 @@ 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 = sd::Tensor<int32_t>::from_vector(tokens);
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                model.compute(8, input_ids, nullptr, {}, {}, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = model.compute(8, input_ids, sd::Tensor<float>(), {}, {});
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("llm test done in %lldms", t1 - t0);
             }
         }

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,11 +788,11 @@ 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);
+        GGMLRunner::compute<float>(get_graph, n_threads, false, true);
         stat();
         for (auto item : original_tensor_to_final_tensor) {
             ggml_tensor* original_tensor = item.first;

src/ltxv.hpp

@@ -26,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,11 +745,11 @@ public:
         return spatial_pos_embed;
     }
 
-    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]
@@ -774,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
@@ -832,89 +832,93 @@ 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(const sd::Tensor<float>& x_tensor,
+                             const sd::Tensor<float>& timesteps_tensor,
+                             const sd::Tensor<float>& context_tensor = {},
+                             const sd::Tensor<float>& y_tensor       = {},
+                             std::vector<int> skip_layers            = std::vector<int>()) {
+        ggml_cgraph* gf = new_graph_custom(MMDIT_GRAPH_SIZE);
 
-        x         = to_backend(x);
-        context   = to_backend(context);
-        y         = to_backend(y);
-        timesteps = to_backend(timesteps);
+        ggml_tensor* x         = make_input(x_tensor);
+        ggml_tensor* timesteps = make_input(timesteps_tensor);
+        ggml_tensor* context   = make_optional_input(context_tensor);
+        ggml_tensor* y         = make_optional_input(y_tensor);
 
-        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);
 
         return gf;
     }
 
-    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>()) {
+    sd::Tensor<float> compute(int n_threads,
+                              const sd::Tensor<float>& x,
+                              const sd::Tensor<float>& timesteps,
+                              const sd::Tensor<float>& context = {},
+                              const sd::Tensor<float>& y       = {},
+                              std::vector<int> skip_layers     = std::vector<int>()) {
         // 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);
         };
 
-        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
     }
 
     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_ASSERT(work_ctx != nullptr);
+        ggml_context* ctx = ggml_init(params);
+        GGML_ASSERT(ctx != nullptr);
 
         {
             // cpu f16: pass
             // cpu f32: pass
             // cuda f16: pass
             // cuda f32: pass
-            auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 128, 128, 16, 1);
+            sd::Tensor<float> x({128, 128, 16, 1});
             std::vector<float> timesteps_vec(1, 999.f);
-            auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
-            ggml_set_f32(x, 0.01f);
+            auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
+            x.fill_(0.01f);
             // print_ggml_tensor(x);
 
-            auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 4096, 154, 1);
-            ggml_set_f32(context, 0.01f);
+            sd::Tensor<float> context({4096, 154, 1});
+            context.fill_(0.01f);
             // print_ggml_tensor(context);
 
-            auto y = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 2048, 1);
-            ggml_set_f32(y, 0.01f);
+            sd::Tensor<float> y({2048, 1});
+            y.fill_(0.01f);
             // print_ggml_tensor(y);
 
-            struct ggml_tensor* out = nullptr;
+            sd::Tensor<float> out;
 
-            int64_t t0 = ggml_time_ms();
-            compute(8, x, timesteps, context, y, &out, work_ctx);
-            int64_t t1 = ggml_time_ms();
+            int64_t t0   = ggml_time_ms();
+            auto out_opt = compute(8,
+                                   x,
+                                   timesteps,
+                                   context,
+                                   y);
+            int64_t t1   = ggml_time_ms();
 
-            print_ggml_tensor(out);
+            GGML_ASSERT(!out_opt.empty());
+            out = std::move(out_opt);
+            print_sd_tensor(out);
             LOG_DEBUG("mmdit test done in %lldms", t1 - t0);
         }
     }

src/model.cpp

@@ -162,43 +162,7 @@ uint16_t f8_e4m3_to_f16(uint8_t f8) {
 }
 
 uint16_t f8_e5m2_to_f16(uint8_t fp8) {
-    uint8_t sign     = (fp8 >> 7) & 0x1;
-    uint8_t exponent = (fp8 >> 2) & 0x1F;
-    uint8_t mantissa = fp8 & 0x3;
-
-    uint16_t fp16_sign = sign << 15;
-    uint16_t fp16_exponent;
-    uint16_t fp16_mantissa;
-
-    if (exponent == 0 && mantissa == 0) {  // zero
-        return fp16_sign;
-    }
-
-    if (exponent == 0x1F) {  // NAN and INF
-        fp16_exponent = 0x1F;
-        fp16_mantissa = mantissa ? (mantissa << 8) : 0;
-        return fp16_sign | (fp16_exponent << 10) | fp16_mantissa;
-    }
-
-    if (exponent == 0) {  // subnormal numbers
-        fp16_mantissa = (mantissa << 8);
-        return fp16_sign | fp16_mantissa;
-    }
-
-    // normal numbers
-    int16_t true_exponent = (int16_t)exponent - 15 + 15;
-    if (true_exponent <= 0) {
-        fp16_exponent = 0;
-        fp16_mantissa = (mantissa << 8);
-    } else if (true_exponent >= 0x1F) {
-        fp16_exponent = 0x1F;
-        fp16_mantissa = 0;
-    } else {
-        fp16_exponent = (uint16_t)true_exponent;
-        fp16_mantissa = mantissa << 8;
-    }
-
-    return fp16_sign | (fp16_exponent << 10) | fp16_mantissa;
+    return static_cast<uint16_t>(fp8) << 8;
 }
 
 void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
@@ -287,7 +251,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;
     }
@@ -453,9 +417,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());
 
@@ -812,7 +776,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") {
@@ -995,7 +959,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;
@@ -1347,6 +1311,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     std::atomic<int64_t> memcpy_time_ms(0);
     std::atomic<int64_t> copy_to_backend_time_ms(0);
     std::atomic<int64_t> convert_time_ms(0);
+    std::atomic<uint64_t> bytes_processed(0);
 
     int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
     LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
@@ -1413,7 +1378,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) {
@@ -1558,6 +1523,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         t1 = ggml_time_ms();
                         copy_to_backend_time_ms.fetch_add(t1 - t0);
                     }
+
+                    bytes_processed.fetch_add((uint64_t)nbytes_to_read);
                 }
                 if (zip != nullptr) {
                     zip_close(zip);
@@ -1570,8 +1537,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
             if (current_idx >= file_tensors.size() || failed) {
                 break;
             }
-            size_t curr_num = total_tensors_processed + current_idx;
-            pretty_progress(static_cast<int>(curr_num), static_cast<int>(total_tensors_to_process), (ggml_time_ms() - t_start) / 1000.0f / (curr_num + 1e-6f));
+            size_t curr_num       = total_tensors_processed + current_idx;
+            float elapsed_seconds = (ggml_time_ms() - t_start) / 1000.0f;
+            pretty_bytes_progress(static_cast<int>(curr_num),
+                                  static_cast<int>(total_tensors_to_process),
+                                  bytes_processed.load(),
+                                  elapsed_seconds);
             std::this_thread::sleep_for(std::chrono::milliseconds(200));
         }
 
@@ -1584,7 +1555,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
             break;
         }
         total_tensors_processed += file_tensors.size();
-        pretty_progress(static_cast<int>(total_tensors_processed), static_cast<int>(total_tensors_to_process), (ggml_time_ms() - t_start) / 1000.0f / (total_tensors_processed + 1e-6f));
+        pretty_bytes_progress(static_cast<int>(total_tensors_processed),
+                              static_cast<int>(total_tensors_to_process),
+                              bytes_processed.load(),
+                              (ggml_time_ms() - t_start) / 1000.0f);
         if (total_tensors_processed < total_tensors_to_process) {
             printf("\n");
         }
@@ -1601,7 +1575,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) {
@@ -1615,7 +1589,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

@@ -323,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);

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,17 @@ 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,
-        std::vector<bool>& class_tokens_mask,
-        struct ggml_tensor* id_embeds) {
+    ggml_cgraph* build_graph(const sd::Tensor<float>& id_pixel_values_tensor,
+                             const sd::Tensor<float>& prompt_embeds_tensor,
+                             std::vector<bool>& class_tokens_mask,
+                             const sd::Tensor<float>& id_embeds_tensor = {}) {
         ctm.clear();
         ctmf16.clear();
         ctmpos.clear();
@@ -458,20 +457,20 @@ public:
 
         auto runner_ctx = get_context();
 
-        struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+
+        ggml_tensor* id_pixel_values = make_input(id_pixel_values_tensor);
+        ggml_tensor* prompt_embeds   = make_input(prompt_embeds_tensor);
+        ggml_tensor* id_embeds       = make_optional_input(id_embeds_tensor);
 
         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);
-
-        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 +494,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,21 +525,21 @@ 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,
-                                                       prompt_embeds_d,
+                                                       id_pixel_values,
+                                                       prompt_embeds,
                                                        class_tokens_mask_d,
                                                        class_tokens_mask_pos,
                                                        left, right);
         else if (pm_version == PM_VERSION_2)
             updated_prompt_embeds = id_encoder2.forward(&runner_ctx,
-                                                        id_pixel_values_d,
-                                                        prompt_embeds_d,
+                                                        id_pixel_values,
+                                                        prompt_embeds,
                                                         class_tokens_mask_d,
                                                         class_tokens_mask_pos,
-                                                        id_embeds_d,
+                                                        id_embeds,
                                                         left, right);
 
         ggml_build_forward_expand(gf, updated_prompt_embeds);
@@ -548,25 +547,21 @@ public:
         return gf;
     }
 
-    bool compute(const int n_threads,
-                 struct ggml_tensor* id_pixel_values,
-                 struct ggml_tensor* prompt_embeds,
-                 struct ggml_tensor* id_embeds,
-                 std::vector<bool>& class_tokens_mask,
-                 struct ggml_tensor** updated_prompt_embeds,
-                 ggml_context* output_ctx) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            // return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask);
+    sd::Tensor<float> compute(const int n_threads,
+                              const sd::Tensor<float>& id_pixel_values,
+                              const sd::Tensor<float>& prompt_embeds,
+                              const sd::Tensor<float>& id_embeds,
+                              std::vector<bool>& class_tokens_mask) {
+        auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        // GGMLRunner::compute(get_graph, n_threads, updated_prompt_embeds);
-        return GGMLRunner::compute(get_graph, n_threads, true, updated_prompt_embeds, output_ctx);
+        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
     }
 };
 
 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 +601,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 +624,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

@@ -1,179 +1,241 @@
 #ifndef __PREPROCESSING_HPP__
 #define __PREPROCESSING_HPP__
 
+#include <cmath>
+#include <limits>
+
 #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);
-    ggml_fp32_to_fp16_row((float*)kernel->data, (ggml_fp16_t*)kernel_fp16->data, ggml_nelements(kernel));
-    ggml_tensor* h  = ggml_conv_2d(ctx0, kernel_fp16, input, 1, 1, padding, padding, 1, 1);
-    ggml_cgraph* gf = ggml_new_graph(ctx0);
-    ggml_build_forward_expand(gf, ggml_cpy(ctx0, h, output));
-    ggml_graph_compute_with_ctx(ctx0, gf, 1);
-    ggml_free(ctx0);
+static inline int64_t preprocessing_offset_4d(const sd::Tensor<float>& tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    const auto& shape = tensor.shape();
+    int64_t n0        = shape.size() > 0 ? shape[0] : 1;
+    int64_t n1        = shape.size() > 1 ? shape[1] : 1;
+    int64_t n2        = shape.size() > 2 ? shape[2] : 1;
+    return ((i3 * n2 + i2) * n1 + i1) * n0 + i0;
 }
 
-void gaussian_kernel(struct ggml_tensor* kernel) {
-    int ks_mid   = static_cast<int>(kernel->ne[0] / 2);
+static inline float preprocessing_get_4d(const sd::Tensor<float>& tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    return tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))];
+}
+
+static inline void preprocessing_set_4d(sd::Tensor<float>& tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))] = value;
+}
+
+static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t image) {
+    sd::Tensor<float> tensor({static_cast<int64_t>(image.width), static_cast<int64_t>(image.height), static_cast<int64_t>(image.channel), 1});
+    for (uint32_t y = 0; y < image.height; ++y) {
+        for (uint32_t x = 0; x < image.width; ++x) {
+            for (uint32_t c = 0; c < image.channel; ++c) {
+                preprocessing_set_4d(tensor, sd_image_get_f32(image, x, y, c), x, y, c, 0);
+            }
+        }
+    }
+    return tensor;
+}
+
+static inline void preprocessing_tensor_to_sd_image(const sd::Tensor<float>& tensor, uint8_t* image_data) {
+    GGML_ASSERT(tensor.dim() == 4);
+    GGML_ASSERT(tensor.shape()[3] == 1);
+    GGML_ASSERT(image_data != nullptr);
+
+    int width   = static_cast<int>(tensor.shape()[0]);
+    int height  = static_cast<int>(tensor.shape()[1]);
+    int channel = static_cast<int>(tensor.shape()[2]);
+    for (int y = 0; y < height; ++y) {
+        for (int x = 0; x < width; ++x) {
+            for (int c = 0; c < channel; ++c) {
+                float value                               = preprocessing_get_4d(tensor, x, y, c, 0);
+                value                                     = std::min(1.0f, std::max(0.0f, value));
+                image_data[(y * width + x) * channel + c] = static_cast<uint8_t>(std::round(value * 255.0f));
+            }
+        }
+    }
+}
+
+static inline sd::Tensor<float> gaussian_kernel_tensor(int kernel_size) {
+    sd::Tensor<float> kernel({kernel_size, kernel_size, 1, 1});
+    int ks_mid   = kernel_size / 2;
     float sigma  = 1.4f;
-    float normal = 1.f / (2.0f * M_PI_ * powf(sigma, 2.0f));
-    for (int y = 0; y < kernel->ne[0]; y++) {
+    float normal = 1.f / (2.0f * M_PI_ * std::pow(sigma, 2.0f));
+    for (int y = 0; y < kernel_size; ++y) {
         float gx = static_cast<float>(-ks_mid + y);
-        for (int x = 0; x < kernel->ne[1]; x++) {
+        for (int x = 0; x < kernel_size; ++x) {
             float gy = static_cast<float>(-ks_mid + x);
-            float k_ = expf(-((gx * gx + gy * gy) / (2.0f * powf(sigma, 2.0f)))) * normal;
-            ggml_ext_tensor_set_f32(kernel, k_, x, y);
+            float k  = std::exp(-((gx * gx + gy * gy) / (2.0f * std::pow(sigma, 2.0f)))) * normal;
+            preprocessing_set_4d(kernel, k, x, y, 0, 0);
         }
     }
+    return kernel;
 }
 
-void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
-    for (int iy = 0; iy < rgb_img->ne[1]; iy++) {
-        for (int ix = 0; ix < rgb_img->ne[0]; ix++) {
-            float r    = ggml_ext_tensor_get_f32(rgb_img, ix, iy);
-            float g    = ggml_ext_tensor_get_f32(rgb_img, ix, iy, 1);
-            float b    = ggml_ext_tensor_get_f32(rgb_img, ix, iy, 2);
+static inline sd::Tensor<float> convolve_tensor(const sd::Tensor<float>& input, const sd::Tensor<float>& kernel, int padding) {
+    GGML_ASSERT(input.dim() == 4);
+    GGML_ASSERT(kernel.dim() == 4);
+    GGML_ASSERT(input.shape()[3] == 1);
+    GGML_ASSERT(kernel.shape()[2] == 1);
+    GGML_ASSERT(kernel.shape()[3] == 1);
+
+    sd::Tensor<float> output(input.shape());
+    int64_t width    = input.shape()[0];
+    int64_t height   = input.shape()[1];
+    int64_t channels = input.shape()[2];
+    int64_t kernel_w = kernel.shape()[0];
+    int64_t kernel_h = kernel.shape()[1];
+
+    for (int64_t c = 0; c < channels; ++c) {
+        for (int64_t y = 0; y < height; ++y) {
+            for (int64_t x = 0; x < width; ++x) {
+                float sum = 0.0f;
+                for (int64_t ky = 0; ky < kernel_h; ++ky) {
+                    int64_t iy = y + ky - padding;
+                    if (iy < 0 || iy >= height) {
+                        continue;
+                    }
+                    for (int64_t kx = 0; kx < kernel_w; ++kx) {
+                        int64_t ix = x + kx - padding;
+                        if (ix < 0 || ix >= width) {
+                            continue;
+                        }
+                        sum += preprocessing_get_4d(input, ix, iy, c, 0) * preprocessing_get_4d(kernel, kx, ky, 0, 0);
+                    }
+                }
+                preprocessing_set_4d(output, sum, x, y, c, 0);
+            }
+        }
+    }
+    return output;
+}
+
+static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& rgb_img) {
+    GGML_ASSERT(rgb_img.dim() == 4);
+    GGML_ASSERT(rgb_img.shape()[2] >= 3);
+    sd::Tensor<float> grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]});
+    for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) {
+        for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) {
+            float r    = preprocessing_get_4d(rgb_img, ix, iy, 0, 0);
+            float g    = preprocessing_get_4d(rgb_img, ix, iy, 1, 0);
+            float b    = preprocessing_get_4d(rgb_img, ix, iy, 2, 0);
             float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b;
-            ggml_ext_tensor_set_f32(grayscale, gray, ix, iy);
+            preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0);
         }
     }
+    return grayscale;
 }
 
-void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
-    int n_elements = static_cast<int>(ggml_nelements(h));
-    float* dx      = (float*)x->data;
-    float* dy      = (float*)y->data;
-    float* dh      = (float*)h->data;
-    for (int i = 0; i < n_elements; i++) {
-        dh[i] = sqrtf(dx[i] * dx[i] + dy[i] * dy[i]);
+static inline sd::Tensor<float> tensor_hypot(const sd::Tensor<float>& x, const sd::Tensor<float>& y) {
+    sd::tensor_check_same_shape(x, y);
+    sd::Tensor<float> out(x.shape());
+    for (int64_t i = 0; i < out.numel(); ++i) {
+        out[i] = std::sqrt(x[i] * x[i] + y[i] * y[i]);
     }
+    return out;
 }
 
-void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
-    int n_elements = static_cast<int>(ggml_nelements(h));
-    float* dx      = (float*)x->data;
-    float* dy      = (float*)y->data;
-    float* dh      = (float*)h->data;
-    for (int i = 0; i < n_elements; i++) {
-        dh[i] = atan2f(dy[i], dx[i]);
+static inline sd::Tensor<float> tensor_arctan2(const sd::Tensor<float>& x, const sd::Tensor<float>& y) {
+    sd::tensor_check_same_shape(x, y);
+    sd::Tensor<float> out(x.shape());
+    for (int64_t i = 0; i < out.numel(); ++i) {
+        out[i] = std::atan2(y[i], x[i]);
     }
+    return out;
 }
 
-void normalize_tensor(struct ggml_tensor* g) {
-    int n_elements = static_cast<int>(ggml_nelements(g));
-    float* dg      = (float*)g->data;
-    float max      = -INFINITY;
-    for (int i = 0; i < n_elements; i++) {
-        max = dg[i] > max ? dg[i] : max;
+static inline void normalize_tensor(sd::Tensor<float>* g) {
+    GGML_ASSERT(g != nullptr);
+    if (g->empty()) {
+        return;
     }
-    max = 1.0f / max;
-    for (int i = 0; i < n_elements; i++) {
-        dg[i] *= max;
+    float max_value = -std::numeric_limits<float>::infinity();
+    for (int64_t i = 0; i < g->numel(); ++i) {
+        max_value = std::max(max_value, (*g)[i]);
     }
+    if (max_value == 0.0f || !std::isfinite(max_value)) {
+        return;
+    }
+    *g *= (1.0f / max_value);
 }
 
-void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struct ggml_tensor* D) {
-    for (int iy = 1; iy < result->ne[1] - 1; iy++) {
-        for (int ix = 1; ix < result->ne[0] - 1; ix++) {
-            float angle = ggml_ext_tensor_get_f32(D, ix, iy) * 180.0f / M_PI_;
-            angle       = angle < 0.0f ? angle += 180.0f : angle;
+static inline sd::Tensor<float> non_max_supression(const sd::Tensor<float>& G, const sd::Tensor<float>& D) {
+    GGML_ASSERT(G.shape() == D.shape());
+    sd::Tensor<float> result = sd::Tensor<float>::zeros(G.shape());
+    for (int64_t iy = 1; iy < result.shape()[1] - 1; ++iy) {
+        for (int64_t ix = 1; ix < result.shape()[0] - 1; ++ix) {
+            float angle = preprocessing_get_4d(D, ix, iy, 0, 0) * 180.0f / M_PI_;
+            angle       = angle < 0.0f ? angle + 180.0f : angle;
             float q     = 1.0f;
             float r     = 1.0f;
 
-            // angle 0
-            if ((0 >= angle && angle < 22.5f) || (157.5f >= angle && angle <= 180)) {
-                q = ggml_ext_tensor_get_f32(G, ix, iy + 1);
-                r = ggml_ext_tensor_get_f32(G, ix, iy - 1);
-            }
-            // angle 45
-            else if (22.5f >= angle && angle < 67.5f) {
-                q = ggml_ext_tensor_get_f32(G, ix + 1, iy - 1);
-                r = ggml_ext_tensor_get_f32(G, ix - 1, iy + 1);
-            }
-            // angle 90
-            else if (67.5f >= angle && angle < 112.5) {
-                q = ggml_ext_tensor_get_f32(G, ix + 1, iy);
-                r = ggml_ext_tensor_get_f32(G, ix - 1, iy);
-            }
-            // angle 135
-            else if (112.5 >= angle && angle < 157.5f) {
-                q = ggml_ext_tensor_get_f32(G, ix - 1, iy - 1);
-                r = ggml_ext_tensor_get_f32(G, ix + 1, iy + 1);
+            if ((0 >= angle && angle < 22.5f) || (157.5f >= angle && angle <= 180.0f)) {
+                q = preprocessing_get_4d(G, ix, iy + 1, 0, 0);
+                r = preprocessing_get_4d(G, ix, iy - 1, 0, 0);
+            } else if (22.5f >= angle && angle < 67.5f) {
+                q = preprocessing_get_4d(G, ix + 1, iy - 1, 0, 0);
+                r = preprocessing_get_4d(G, ix - 1, iy + 1, 0, 0);
+            } else if (67.5f >= angle && angle < 112.5f) {
+                q = preprocessing_get_4d(G, ix + 1, iy, 0, 0);
+                r = preprocessing_get_4d(G, ix - 1, iy, 0, 0);
+            } else if (112.5f >= angle && angle < 157.5f) {
+                q = preprocessing_get_4d(G, ix - 1, iy - 1, 0, 0);
+                r = preprocessing_get_4d(G, ix + 1, iy + 1, 0, 0);
             }
 
-            float cur = ggml_ext_tensor_get_f32(G, ix, iy);
-            if ((cur >= q) && (cur >= r)) {
-                ggml_ext_tensor_set_f32(result, cur, ix, iy);
-            } else {
-                ggml_ext_tensor_set_f32(result, 0.0f, ix, iy);
-            }
+            float cur = preprocessing_get_4d(G, ix, iy, 0, 0);
+            preprocessing_set_4d(result, (cur >= q && cur >= r) ? cur : 0.0f, ix, iy, 0, 0);
         }
     }
+    return result;
 }
 
-void threshold_hystersis(struct 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;
-    for (int i = 0; i < n_elements; i++) {
-        max = imd[i] > max ? imd[i] : max;
+static inline void threshold_hystersis(sd::Tensor<float>* img, float high_threshold, float low_threshold, float weak, float strong) {
+    GGML_ASSERT(img != nullptr);
+    if (img->empty()) {
+        return;
     }
-    float ht = max * high_threshold;
+    float max_value = -std::numeric_limits<float>::infinity();
+    for (int64_t i = 0; i < img->numel(); ++i) {
+        max_value = std::max(max_value, (*img)[i]);
+    }
+
+    float ht = max_value * high_threshold;
     float lt = ht * low_threshold;
-    for (int i = 0; i < n_elements; i++) {
-        float img_v = imd[i];
-        if (img_v >= ht) {  // strong pixel
-            imd[i] = strong;
-        } else if (img_v <= ht && img_v >= lt) {  // strong pixel
-            imd[i] = weak;
+    for (int64_t i = 0; i < img->numel(); ++i) {
+        float img_v = (*img)[i];
+        if (img_v >= ht) {
+            (*img)[i] = strong;
+        } else if (img_v <= ht && img_v >= lt) {
+            (*img)[i] = weak;
         }
     }
 
-    for (int iy = 0; iy < img->ne[1]; iy++) {
-        for (int ix = 0; ix < img->ne[0]; ix++) {
-            if (ix >= 3 && ix <= img->ne[0] - 3 && iy >= 3 && iy <= img->ne[1] - 3) {
-                ggml_ext_tensor_set_f32(img, ggml_ext_tensor_get_f32(img, ix, iy), ix, iy);
-            } else {
-                ggml_ext_tensor_set_f32(img, 0.0f, ix, iy);
+    for (int64_t iy = 0; iy < img->shape()[1]; ++iy) {
+        for (int64_t ix = 0; ix < img->shape()[0]; ++ix) {
+            if (!(ix >= 3 && ix <= img->shape()[0] - 3 && iy >= 3 && iy <= img->shape()[1] - 3)) {
+                preprocessing_set_4d(*img, 0.0f, ix, iy, 0, 0);
             }
         }
     }
 
-    // hysteresis
-    for (int iy = 1; iy < img->ne[1] - 1; iy++) {
-        for (int ix = 1; ix < img->ne[0] - 1; ix++) {
-            float imd_v = ggml_ext_tensor_get_f32(img, ix, iy);
+    for (int64_t iy = 1; iy < img->shape()[1] - 1; ++iy) {
+        for (int64_t ix = 1; ix < img->shape()[0] - 1; ++ix) {
+            float imd_v = preprocessing_get_4d(*img, ix, iy, 0, 0);
             if (imd_v == weak) {
-                if (ggml_ext_tensor_get_f32(img, ix + 1, iy - 1) == strong || ggml_ext_tensor_get_f32(img, ix + 1, iy) == strong ||
-                    ggml_ext_tensor_get_f32(img, ix, iy - 1) == strong || ggml_ext_tensor_get_f32(img, ix, iy + 1) == strong ||
-                    ggml_ext_tensor_get_f32(img, ix - 1, iy - 1) == strong || ggml_ext_tensor_get_f32(img, ix - 1, iy) == strong) {
-                    ggml_ext_tensor_set_f32(img, strong, ix, iy);
-                } else {
-                    ggml_ext_tensor_set_f32(img, 0.0f, ix, iy);
-                }
+                bool has_strong_neighbor =
+                    preprocessing_get_4d(*img, ix + 1, iy - 1, 0, 0) == strong ||
+                    preprocessing_get_4d(*img, ix + 1, iy, 0, 0) == strong ||
+                    preprocessing_get_4d(*img, ix, iy - 1, 0, 0) == strong ||
+                    preprocessing_get_4d(*img, ix, iy + 1, 0, 0) == strong ||
+                    preprocessing_get_4d(*img, ix - 1, iy - 1, 0, 0) == strong ||
+                    preprocessing_get_4d(*img, ix - 1, iy, 0, 0) == strong;
+                preprocessing_set_4d(*img, has_strong_neighbor ? strong : 0.0f, ix, iy, 0, 0);
             }
         }
     }
 }
 
 bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
-    struct ggml_init_params params;
-    params.mem_size               = static_cast<size_t>(40 * img.width * img.height);  // 10MB for 512x512
-    params.mem_buffer             = nullptr;
-    params.no_alloc               = false;
-    struct ggml_context* work_ctx = ggml_init(params);
-
-    if (!work_ctx) {
-        LOG_ERROR("ggml_init() failed");
-        return false;
-    }
-
     float kX[9] = {
         -1, 0, 1,
         -2, 0, 2,
@@ -184,43 +246,33 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
         0, 0, 0,
         -1, -2, -1};
 
-    // generate kernel
-    int kernel_size             = 5;
-    struct ggml_tensor* gkernel = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, 1);
-    struct ggml_tensor* sf_kx   = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
-    memcpy(sf_kx->data, kX, ggml_nbytes(sf_kx));
-    struct ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
-    memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
-    gaussian_kernel(gkernel);
-    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
-    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
-    struct ggml_tensor* iX         = ggml_dup_tensor(work_ctx, image_gray);
-    struct ggml_tensor* iY         = ggml_dup_tensor(work_ctx, image_gray);
-    struct ggml_tensor* G          = ggml_dup_tensor(work_ctx, image_gray);
-    struct ggml_tensor* tetha      = ggml_dup_tensor(work_ctx, image_gray);
-    sd_image_to_ggml_tensor(img, image);
-    grayscale(image, image_gray);
-    convolve(image_gray, image_gray, gkernel, 2);
-    convolve(image_gray, iX, sf_kx, 1);
-    convolve(image_gray, iY, sf_ky, 1);
-    prop_hypot(iX, iY, G);
-    normalize_tensor(G);
-    prop_arctan2(iX, iY, tetha);
-    non_max_supression(image_gray, G, tetha);
-    threshold_hystersis(image_gray, high_threshold, low_threshold, weak, strong);
-    // to RGB channels
-    for (uint32_t iy = 0; iy < img.height; iy++) {
-        for (uint32_t ix = 0; ix < img.width; ix++) {
-            float gray = ggml_ext_tensor_get_f32(image_gray, ix, iy);
+    sd::Tensor<float> gkernel = gaussian_kernel_tensor(5);
+    sd::Tensor<float> sf_kx({3, 3, 1, 1}, std::vector<float>(kX, kX + 9));
+    sd::Tensor<float> sf_ky({3, 3, 1, 1}, std::vector<float>(kY, kY + 9));
+
+    sd::Tensor<float> image      = sd_image_to_preprocessing_tensor(img);
+    sd::Tensor<float> image_gray = grayscale_tensor(image);
+    image_gray                   = convolve_tensor(image_gray, gkernel, 2);
+    sd::Tensor<float> iX         = convolve_tensor(image_gray, sf_kx, 1);
+    sd::Tensor<float> iY         = convolve_tensor(image_gray, sf_ky, 1);
+    sd::Tensor<float> G          = tensor_hypot(iX, iY);
+    normalize_tensor(&G);
+    sd::Tensor<float> theta = tensor_arctan2(iX, iY);
+    image_gray              = non_max_supression(G, theta);
+    threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong);
+
+    for (uint32_t iy = 0; iy < img.height; ++iy) {
+        for (uint32_t ix = 0; ix < img.width; ++ix) {
+            float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0);
             gray       = inverse ? 1.0f - gray : gray;
-            ggml_ext_tensor_set_f32(image, gray, ix, iy);
-            ggml_ext_tensor_set_f32(image, gray, ix, iy, 1);
-            ggml_ext_tensor_set_f32(image, gray, ix, iy, 2);
+            for (uint32_t c = 0; c < img.channel; ++c) {
+                preprocessing_set_4d(image, gray, ix, iy, c, 0);
+            }
         }
     }
-    ggml_tensor_to_sd_image(image, img.data);
-    ggml_free(work_ctx);
+
+    preprocessing_tensor_to_sd_image(image, img.data);
     return true;
 }
 
-#endif  // __PREPROCESSING_HPP__
+#endif  // __PREPROCESSING_HPP__

src/qwen_image.hpp

@@ -26,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));
@@ -49,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"]);
@@ -107,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]
@@ -249,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]
@@ -325,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]
@@ -389,12 +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* 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"]);
@@ -429,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,]
@@ -521,24 +521,25 @@ 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(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
+                                 bool increase_ref_index                                  = false) {
+            ggml_cgraph* gf        = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
             GGML_ASSERT(x->ne[3] == 1);
-            struct ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
-
-            x         = to_backend(x);
-            context   = to_backend(context);
-            timesteps = to_backend(timesteps);
-
-            for (int i = 0; i < ref_latents.size(); i++) {
-                ref_latents[i] = to_backend(ref_latents[i]);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
             pe_vec      = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
@@ -587,67 +588,72 @@ 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);
 
             return gf;
         }
 
-        bool compute(int n_threads,
-                     struct ggml_tensor* x,
-                     struct ggml_tensor* timesteps,
-                     struct ggml_tensor* context,
-                     std::vector<ggml_tensor*> ref_latents = {},
-                     bool increase_ref_index               = false,
-                     struct ggml_tensor** output           = nullptr,
-                     struct ggml_context* output_ctx       = nullptr) {
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {},
+                                  bool increase_ref_index                           = false) {
             // 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);
             };
 
-            return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         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_ASSERT(work_ctx != nullptr);
+            ggml_context* ctx = ggml_init(params);
+            GGML_ASSERT(ctx != nullptr);
 
             {
-                // auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 16, 16, 16, 1);
+                // auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 16, 16, 16, 1);
                 // ggml_set_f32(x, 0.01f);
-                auto x = load_tensor_from_file(work_ctx, "./qwen_image_x.bin");
-                print_ggml_tensor(x);
+                auto x = sd::load_tensor_from_file_as_tensor<float>("./qwen_image_x.bin");
+                print_sd_tensor(x);
 
                 std::vector<float> timesteps_vec(1, 1000.f);
-                auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
+                auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
 
-                // auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 3584, 256, 1);
+                // auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3584, 256, 1);
                 // ggml_set_f32(context, 0.01f);
-                auto context = load_tensor_from_file(work_ctx, "./qwen_image_context.bin");
-                print_ggml_tensor(context);
+                auto context = sd::load_tensor_from_file_as_tensor<float>("./qwen_image_context.bin");
+                print_sd_tensor(context);
 
-                struct ggml_tensor* out = nullptr;
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                compute(8, x, timesteps, context, {}, false, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = compute(8,
+                                       x,
+                                       timesteps,
+                                       context,
+                                       {},
+                                       false);
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("qwen_image test done in %lldms", t1 - t0);
             }
         }

src/rope.hpp

@@ -600,10 +600,10 @@ namespace Rope {
         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];
@@ -641,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/sample-cache.cpp

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

src/sample-cache.h

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

src/spectrum.hpp

@@ -6,6 +6,7 @@
 #include <vector>
 
 #include "ggml_extend.hpp"
+#include "tensor.hpp"
 
 struct SpectrumConfig {
     float w            = 0.40f;
@@ -57,11 +58,8 @@ struct SpectrumState {
         return (num_cached + 1) % ws != 0;
     }
 
-    void update(const struct ggml_tensor* denoised) {
-        int64_t ne        = ggml_nelements(denoised);
-        const float* data = (const float*)denoised->data;
-
-        H_buf.emplace_back(data, data + ne);
+    void update(const sd::Tensor<float>& denoised) {
+        H_buf.emplace_back(denoised.data(), denoised.data() + denoised.numel());
         T_buf.push_back(taus(cnt));
 
         while ((int)H_buf.size() > K) {
@@ -76,13 +74,13 @@ struct SpectrumState {
         cnt++;
     }
 
-    void predict(struct ggml_tensor* denoised) {
+    void predict(sd::Tensor<float>* denoised) {
+        GGML_ASSERT(denoised != nullptr);
         int64_t F    = (int64_t)H_buf[0].size();
         int K_curr   = (int)H_buf.size();
         int M1       = config.m + 1;
         float tau_at = taus(cnt);
 
-        // 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;
@@ -92,7 +90,6 @@ struct SpectrumState {
                 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)
@@ -100,7 +97,6 @@ struct SpectrumState {
         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++) {
@@ -111,7 +107,6 @@ struct SpectrumState {
             }
         }
 
-        // Cholesky decomposition
         std::vector<float> L(M1 * M1, 0.0f);
         if (!cholesky_decompose(XtX.data(), L.data(), M1)) {
             float trace = 0.0f;
@@ -122,18 +117,15 @@ struct SpectrumState {
             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* out          = denoised->data();
         float w_cheb        = config.w;
         float w_taylor      = 1.0f - w_cheb;
         const float* h_last = H_buf.back().data();

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()
@@ -459,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)
@@ -473,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));
@@ -519,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);
@@ -527,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) {
@@ -558,45 +558,44 @@ struct TinyImageAutoEncoder : public VAE {
         return "taesd";
     }
 
-    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);
     }
 
-    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
+        SD_UNUSED(rng);
         return vae_output;
     }
 
-    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+        return latents;
     }
 
-    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+        return latents;
     }
 
     int get_encoder_output_channels(int input_channels) {
         return taesd.z_channels;
     }
 
-    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_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
+        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        ggml_tensor* z   = make_input(z_tensor);
+        auto runner_ctx  = get_context();
+        ggml_tensor* out = decode_graph ? taesd.decode(&runner_ctx, z) : taesd.encode(&runner_ctx, z);
         ggml_build_forward_expand(gf, out);
         return gf;
     }
 
-    bool _compute(const int n_threads,
-                  struct ggml_tensor* z,
-                  bool decode_graph,
-                  struct ggml_tensor** output,
-                  struct ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(z, decode_graph);
+    sd::Tensor<float> _compute(const int n_threads,
+                               const sd::Tensor<float>& z_tensor,
+                               bool decode_graph) override {
+        auto get_graph = [&]() -> ggml_cgraph* {
+            return build_graph(z_tensor, decode_graph);
         };
 
-        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 
@@ -621,46 +620,45 @@ struct TinyVideoAutoEncoder : public VAE {
         return "taehv";
     }
 
-    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);
     }
 
-    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
+        SD_UNUSED(rng);
         return vae_output;
     }
 
-    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+        return latents;
     }
 
-    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+        return latents;
     }
 
     int get_encoder_output_channels(int input_channels) {
         return taehv.z_channels;
     }
 
-    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_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
+        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        ggml_tensor* z   = make_input(z_tensor);
+        auto runner_ctx  = get_context();
+        ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
         ggml_build_forward_expand(gf, out);
         return gf;
     }
 
-    bool _compute(const int n_threads,
-                  struct ggml_tensor* z,
-                  bool decode_graph,
-                  struct ggml_tensor** output,
-                  struct ggml_context* output_ctx = nullptr) {
-        auto get_graph = [&]() -> struct ggml_cgraph* {
-            return build_graph(z, decode_graph);
+    sd::Tensor<float> _compute(const int n_threads,
+                               const sd::Tensor<float>& z_tensor,
+                               bool decode_graph) override {
+        auto get_graph = [&]() -> ggml_cgraph* {
+            return build_graph(z_tensor, decode_graph);
         };
 
-        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 
-#endif  // __TAE_HPP__
+#endif  // __TAE_HPP__

src/tensor_ggml.hpp

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

src/ucache.hpp

@@ -6,8 +6,10 @@
 #include <unordered_map>
 #include <vector>
 
+#include "condition_cache_utils.hpp"
 #include "denoiser.hpp"
 #include "ggml_extend.hpp"
+#include "tensor.hpp"
 
 struct UCacheConfig {
     bool enabled                = false;
@@ -29,15 +31,15 @@ struct UCacheCacheEntry {
 
 struct UCacheState {
     UCacheConfig config;
-    Denoiser* denoiser                  = nullptr;
-    float start_sigma                   = std::numeric_limits<float>::max();
-    float end_sigma                     = 0.0f;
-    bool initialized                    = false;
-    bool initial_step                   = true;
-    bool skip_current_step              = false;
-    bool step_active                    = false;
-    const SDCondition* anchor_condition = nullptr;
-    std::unordered_map<const SDCondition*, UCacheCacheEntry> cache_diffs;
+    Denoiser* denoiser           = nullptr;
+    float start_sigma            = std::numeric_limits<float>::max();
+    float end_sigma              = 0.0f;
+    bool initialized             = false;
+    bool initial_step            = true;
+    bool skip_current_step       = false;
+    bool step_active             = false;
+    const void* anchor_condition = nullptr;
+    std::unordered_map<const void*, UCacheCacheEntry> cache_diffs;
     std::vector<float> prev_input;
     std::vector<float> prev_output;
     float output_prev_norm                = 0.0f;
@@ -233,43 +235,30 @@ struct UCacheState {
         return base_threshold * multiplier;
     }
 
-    bool has_cache(const SDCondition* cond) const {
+    bool has_cache(const void* cond) const {
         auto it = cache_diffs.find(cond);
         return it != cache_diffs.end() && !it->second.diff.empty();
     }
 
-    void update_cache(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) {
+    void update_cache(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
         UCacheCacheEntry& entry = cache_diffs[cond];
-        size_t ne               = static_cast<size_t>(ggml_nelements(output));
-        entry.diff.resize(ne);
-        float* out_data = (float*)output->data;
-        float* in_data  = (float*)input->data;
-
-        for (size_t i = 0; i < ne; ++i) {
-            entry.diff[i] = out_data[i] - in_data[i];
-        }
+        sd::store_condition_cache_diff(&entry.diff, input, output);
     }
 
-    void apply_cache(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) {
+    void apply_cache(const void* cond, const sd::Tensor<float>& input, sd::Tensor<float>* output) {
         auto it = cache_diffs.find(cond);
         if (it == cache_diffs.end() || it->second.diff.empty()) {
             return;
         }
-
-        copy_ggml_tensor(output, input);
-        float* out_data                = (float*)output->data;
-        const std::vector<float>& diff = it->second.diff;
-        for (size_t i = 0; i < diff.size(); ++i) {
-            out_data[i] += diff[i];
-        }
+        sd::apply_condition_cache_diff(it->second.diff, input, output);
     }
 
-    bool before_condition(const SDCondition* cond,
-                          ggml_tensor* input,
-                          ggml_tensor* output,
+    bool before_condition(const void* cond,
+                          const sd::Tensor<float>& input,
+                          sd::Tensor<float>* output,
                           float sigma,
                           int step_index) {
-        if (!enabled() || step_index < 0) {
+        if (!enabled() || step_index < 0 || output == nullptr) {
             return false;
         }
         if (step_index != current_step_index) {
@@ -302,13 +291,13 @@ struct UCacheState {
             return false;
         }
 
-        size_t ne = static_cast<size_t>(ggml_nelements(input));
+        size_t ne = static_cast<size_t>(input.numel());
         if (prev_input.size() != ne) {
             return false;
         }
 
-        float* input_data = (float*)input->data;
-        last_input_change = 0.0f;
+        const float* input_data = input.data();
+        last_input_change       = 0.0f;
         for (size_t i = 0; i < ne; ++i) {
             last_input_change += std::fabs(input_data[i] - prev_input[i]);
         }
@@ -354,7 +343,7 @@ struct UCacheState {
         return false;
     }
 
-    void after_condition(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) {
+    void after_condition(const void* cond, const sd::Tensor<float>& input, const sd::Tensor<float>& output) {
         if (!step_is_active()) {
             return;
         }
@@ -367,16 +356,16 @@ struct UCacheState {
         steps_computed_since_active++;
         consecutive_skipped_steps = 0;
 
-        size_t ne      = static_cast<size_t>(ggml_nelements(input));
-        float* in_data = (float*)input->data;
+        size_t ne            = static_cast<size_t>(input.numel());
+        const float* in_data = input.data();
         prev_input.resize(ne);
         for (size_t i = 0; i < ne; ++i) {
             prev_input[i] = in_data[i];
         }
         has_prev_input = true;
 
-        float* out_data     = (float*)output->data;
-        float output_change = 0.0f;
+        const float* out_data = output.data();
+        float output_change   = 0.0f;
         if (has_prev_output && prev_output.size() == ne) {
             for (size_t i = 0; i < ne; ++i) {
                 output_change += std::fabs(out_data[i] - prev_output[i]);

src/unet.hpp

@@ -60,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]
@@ -388,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]);
 
@@ -404,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]);
 
@@ -420,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]
@@ -480,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);
@@ -605,82 +605,81 @@ 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(const sd::Tensor<float>& x_tensor,
+                             const sd::Tensor<float>& timesteps_tensor,
+                             const sd::Tensor<float>& context_tensor               = {},
+                             const sd::Tensor<float>& c_concat_tensor              = {},
+                             const sd::Tensor<float>& y_tensor                     = {},
+                             int num_video_frames                                  = -1,
+                             const std::vector<sd::Tensor<float>>& controls_tensor = {},
+                             float control_strength                                = 0.f) {
+        ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
+
+        ggml_tensor* x         = make_input(x_tensor);
+        ggml_tensor* timesteps = make_input(timesteps_tensor);
+        ggml_tensor* context   = make_optional_input(context_tensor);
+        ggml_tensor* c_concat  = make_optional_input(c_concat_tensor);
+        ggml_tensor* y         = make_optional_input(y_tensor);
+        std::vector<ggml_tensor*> controls;
+        controls.reserve(controls_tensor.size());
+        for (const auto& control_tensor : controls_tensor) {
+            controls.push_back(make_input(control_tensor));
+        }
 
         if (num_video_frames == -1) {
             num_video_frames = static_cast<int>(x->ne[3]);
         }
 
-        x         = to_backend(x);
-        context   = to_backend(context);
-        y         = to_backend(y);
-        timesteps = to_backend(timesteps);
-        c_concat  = to_backend(c_concat);
-
-        for (int i = 0; i < controls.size(); i++) {
-            controls[i] = to_backend(controls[i]);
-        }
-
         auto runner_ctx = get_context();
 
-        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);
 
         return gf;
     }
 
-    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) {
+    sd::Tensor<float> compute(int n_threads,
+                              const sd::Tensor<float>& x,
+                              const sd::Tensor<float>& timesteps,
+                              const sd::Tensor<float>& context               = {},
+                              const sd::Tensor<float>& c_concat              = {},
+                              const sd::Tensor<float>& y                     = {},
+                              int num_video_frames                           = -1,
+                              const std::vector<sd::Tensor<float>>& controls = {},
+                              float control_strength                         = 0.f) {
         // 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);
         };
 
-        return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
     }
 
     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_ASSERT(work_ctx != nullptr);
+        ggml_context* ctx = ggml_init(params);
+        GGML_ASSERT(ctx != nullptr);
 
         {
             // CPU, num_video_frames = 1, x{num_video_frames, 8, 8, 8}: Pass
@@ -689,27 +688,37 @@ struct UNetModelRunner : public GGMLRunner {
             // CUDA, num_video_frames = 3, x{num_video_frames, 8, 8, 8}: nan
             int num_video_frames = 3;
 
-            auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 8, num_video_frames);
+            sd::Tensor<float> x({8, 8, 8, num_video_frames});
             std::vector<float> timesteps_vec(num_video_frames, 999.f);
-            auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
-            ggml_set_f32(x, 0.5f);
+            auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
+            x.fill_(0.5f);
             // print_ggml_tensor(x);
 
-            auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 1024, 1, num_video_frames);
-            ggml_set_f32(context, 0.5f);
+            sd::Tensor<float> context({1024, 1, num_video_frames});
+            context.fill_(0.5f);
             // print_ggml_tensor(context);
 
-            auto y = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 768, num_video_frames);
-            ggml_set_f32(y, 0.5f);
+            sd::Tensor<float> y({768, num_video_frames});
+            y.fill_(0.5f);
             // print_ggml_tensor(y);
 
-            struct ggml_tensor* out = nullptr;
+            sd::Tensor<float> out;
 
-            int64_t t0 = ggml_time_ms();
-            compute(8, x, timesteps, context, nullptr, y, num_video_frames, {}, 0.f, &out, work_ctx);
-            int64_t t1 = ggml_time_ms();
+            int64_t t0   = ggml_time_ms();
+            auto out_opt = compute(8,
+                                   x,
+                                   timesteps,
+                                   context,
+                                   {},
+                                   y,
+                                   num_video_frames,
+                                   {},
+                                   0.f);
+            int64_t t1   = ggml_time_ms();
 
-            print_ggml_tensor(out);
+            GGML_ASSERT(!out_opt.empty());
+            out = std::move(out_opt);
+            print_sd_tensor(out);
             LOG_DEBUG("unet test done in %lldms", t1 - t0);
         }
     }

src/upscaler.cpp

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

src/util.cpp

@@ -337,17 +337,13 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
     return result;
 }
 
-void pretty_progress(int step, int steps, float time) {
-    if (sd_progress_cb) {
-        sd_progress_cb(step, steps, time, sd_progress_cb_data);
-        return;
-    }
-    if (step == 0) {
-        return;
-    }
+static std::string build_progress_bar(int step, int steps) {
     std::string progress = "  |";
     int max_progress     = 50;
-    int32_t current      = (int32_t)(step * 1.f * max_progress / steps);
+    int32_t current      = 0;
+    if (steps > 0) {
+        current = (int32_t)(step * 1.f * max_progress / steps);
+    }
     for (int i = 0; i < 50; i++) {
         if (i > current) {
             progress += " ";
@@ -358,16 +354,57 @@ void pretty_progress(int step, int steps, float time) {
         }
     }
     progress += "|";
+    return progress;
+}
 
-    const char* lf   = (step == steps ? "\n" : "");
+static void print_progress_line(int step, int steps, const std::string& speed_text) {
+    if (step == 0) {
+        return;
+    }
+    std::string progress = build_progress_bar(step, steps);
+    const char* lf       = (step == steps ? "\n" : "");
+    printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf);
+    fflush(stdout);  // for linux
+}
+
+void pretty_progress(int step, int steps, float time) {
+    if (sd_progress_cb) {
+        sd_progress_cb(step, steps, time, sd_progress_cb_data);
+        return;
+    }
+    if (step == 0) {
+        return;
+    }
     const char* unit = "s/it";
     float speed      = time;
     if (speed < 1.0f && speed > 0.f) {
         speed = 1.0f / speed;
         unit  = "it/s";
     }
-    printf("\r%s %i/%i - %.2f%s\033[K%s", progress.c_str(), step, steps, speed, unit, lf);
-    fflush(stdout);  // for linux
+    print_progress_line(step, steps, sd_format("%.2f%s", speed, unit));
+}
+
+void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds) {
+    if (sd_progress_cb) {
+        float time = elapsed_seconds / (step + 1e-6f);
+        sd_progress_cb(step, steps, time, sd_progress_cb_data);
+        return;
+    }
+    if (step == 0) {
+        return;
+    }
+
+    double bytes_per_second = 0.0;
+    if (elapsed_seconds > 0.0f) {
+        bytes_per_second = bytes_processed / (double)elapsed_seconds;
+    }
+
+    double speed_mb = bytes_per_second / (1024.0 * 1024.0);
+    if (speed_mb >= 1024.0) {
+        print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0));
+    } else {
+        print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb));
+    }
 }
 
 std::string ltrim(const std::string& s) {
@@ -479,158 +516,96 @@ const char* sd_get_system_info() {
     return buffer;
 }
 
-sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image) {
-    sd_image_f32_t converted_image;
-    converted_image.width   = image.width;
-    converted_image.height  = image.height;
-    converted_image.channel = image.channel;
+sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index) {
+    const auto& shape = tensor.shape();
+    GGML_ASSERT(shape.size() == 4 || shape.size() == 5);
+    int width     = static_cast<int>(shape[0]);
+    int height    = static_cast<int>(shape[1]);
+    int channel   = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
+    uint8_t* data = (uint8_t*)malloc(static_cast<size_t>(width * height * channel));
+    GGML_ASSERT(data != nullptr);
 
-    // Allocate memory for float data
-    converted_image.data = (float*)malloc(image.width * image.height * image.channel * sizeof(float));
-
-    for (uint32_t i = 0; i < image.width * image.height * image.channel; i++) {
-        // Convert uint8_t to float
-        converted_image.data[i] = (float)image.data[i];
-    }
-
-    return converted_image;
-}
-
-// Function to perform double linear interpolation
-float interpolate(float v1, float v2, float v3, float v4, float x_ratio, float y_ratio) {
-    return v1 * (1 - x_ratio) * (1 - y_ratio) + v2 * x_ratio * (1 - y_ratio) + v3 * (1 - x_ratio) * y_ratio + v4 * x_ratio * y_ratio;
-}
-
-sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height) {
-    sd_image_f32_t resized_image;
-    resized_image.width   = target_width;
-    resized_image.height  = target_height;
-    resized_image.channel = image.channel;
-
-    // Allocate memory for resized float data
-    resized_image.data = (float*)malloc(target_width * target_height * image.channel * sizeof(float));
-
-    for (int y = 0; y < target_height; y++) {
-        for (int x = 0; x < target_width; x++) {
-            float original_x = (float)x * image.width / target_width;
-            float original_y = (float)y * image.height / target_height;
-
-            uint32_t x1 = (uint32_t)original_x;
-            uint32_t y1 = (uint32_t)original_y;
-            uint32_t x2 = std::min(x1 + 1, image.width - 1);
-            uint32_t y2 = std::min(y1 + 1, image.height - 1);
-
-            for (uint32_t k = 0; k < image.channel; k++) {
-                float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
-                float v2 = *(image.data + y1 * image.width * image.channel + x2 * image.channel + k);
-                float v3 = *(image.data + y2 * image.width * image.channel + x1 * image.channel + k);
-                float v4 = *(image.data + y2 * image.width * image.channel + x2 * image.channel + k);
-
-                float x_ratio = original_x - x1;
-                float y_ratio = original_y - y1;
-
-                float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
-
-                *(resized_image.data + y * target_width * image.channel + x * image.channel + k) = value;
+    for (int iw = 0; iw < width; ++iw) {
+        for (int ih = 0; ih < height; ++ih) {
+            for (int ic = 0; ic < channel; ++ic) {
+                float value                            = shape.size() == 5 ? tensor.index(iw, ih, frame_index, ic, 0)
+                                                                           : tensor.index(iw, ih, ic, frame_index);
+                value                                  = std::clamp(value, 0.0f, 1.0f);
+                data[(ih * width + iw) * channel + ic] = static_cast<uint8_t>(std::round(value * 255.0f));
             }
         }
     }
-
-    return resized_image;
+    return {
+        static_cast<uint32_t>(width),
+        static_cast<uint32_t>(height),
+        static_cast<uint32_t>(channel),
+        data,
+    };
 }
 
-void normalize_sd_image_f32_t(sd_image_f32_t image, float means[3], float stds[3]) {
-    for (uint32_t y = 0; y < image.height; y++) {
-        for (uint32_t x = 0; x < image.width; x++) {
-            for (uint32_t k = 0; k < image.channel; k++) {
-                int index         = (y * image.width + x) * image.channel + k;
-                image.data[index] = (image.data[index] - means[k]) / stds[k];
+sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
+                                     int target_width,
+                                     int target_height,
+                                     bool scale) {
+    sd::Tensor<float> tensor = sd::zeros<float>({static_cast<int64_t>(image.width),
+                                                 static_cast<int64_t>(image.height),
+                                                 static_cast<int64_t>(image.channel),
+                                                 1});
+    for (uint32_t iw = 0; iw < image.width; ++iw) {
+        for (uint32_t ih = 0; ih < image.height; ++ih) {
+            for (uint32_t ic = 0; ic < image.channel; ++ic) {
+                tensor.index(iw, ih, ic, 0) = sd_image_get_f32(image, iw, ih, ic, scale);
             }
         }
     }
+    if (target_width >= 0 && target_height >= 0 &&
+        (tensor.shape()[0] != target_width || tensor.shape()[1] != target_height)) {
+        tensor = sd::ops::interpolate(tensor,
+                                      {target_width,
+                                       target_height,
+                                       tensor.shape()[2],
+                                       tensor.shape()[3]});
+    }
+    return tensor;
 }
 
 // Constants for means and std
 float means[3] = {0.48145466f, 0.4578275f, 0.40821073f};
 float stds[3]  = {0.26862954f, 0.26130258f, 0.27577711f};
 
-// Function to clip and preprocess sd_image_f32_t
-sd_image_f32_t clip_preprocess(sd_image_f32_t image, int target_width, int target_height) {
-    float width_scale  = (float)target_width / image.width;
-    float height_scale = (float)target_height / image.height;
+sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_width, int target_height) {
+    GGML_ASSERT(image.dim() == 4);
+    GGML_ASSERT(image.shape()[2] == 3);
+    GGML_ASSERT(image.shape()[3] == 1);
+    GGML_ASSERT(target_width > 0 && target_height > 0);
 
-    float scale = std::fmax(width_scale, height_scale);
+    float width_scale  = static_cast<float>(target_width) / static_cast<float>(image.shape()[0]);
+    float height_scale = static_cast<float>(target_height) / static_cast<float>(image.shape()[1]);
+    float scale        = std::fmax(width_scale, height_scale);
 
-    // Interpolation
-    int resized_width   = (int)(scale * image.width);
-    int resized_height  = (int)(scale * image.height);
-    float* resized_data = (float*)malloc(resized_width * resized_height * image.channel * sizeof(float));
+    int64_t resized_width  = static_cast<int64_t>(scale * static_cast<float>(image.shape()[0]));
+    int64_t resized_height = static_cast<int64_t>(scale * static_cast<float>(image.shape()[1]));
 
-    for (int y = 0; y < resized_height; y++) {
-        for (int x = 0; x < resized_width; x++) {
-            float original_x = (float)x * image.width / resized_width;
-            float original_y = (float)y * image.height / resized_height;
+    sd::Tensor<float> resized = sd::ops::interpolate(
+        image,
+        {resized_width, resized_height, image.shape()[2], image.shape()[3]});
 
-            uint32_t x1 = (uint32_t)original_x;
-            uint32_t y1 = (uint32_t)original_y;
-            uint32_t x2 = std::min(x1 + 1, image.width - 1);
-            uint32_t y2 = std::min(y1 + 1, image.height - 1);
+    int64_t h_offset = std::max<int64_t>((resized_height - target_height) / 2, 0);
+    int64_t w_offset = std::max<int64_t>((resized_width - target_width) / 2, 0);
 
-            for (uint32_t k = 0; k < image.channel; k++) {
-                float v1 = *(image.data + y1 * image.width * image.channel + x1 * image.channel + k);
-                float v2 = *(image.data + y1 * image.width * image.channel + x2 * image.channel + k);
-                float v3 = *(image.data + y2 * image.width * image.channel + x1 * image.channel + k);
-                float v4 = *(image.data + y2 * image.width * image.channel + x2 * image.channel + k);
-
-                float x_ratio = original_x - x1;
-                float y_ratio = original_y - y1;
-
-                float value = interpolate(v1, v2, v3, v4, x_ratio, y_ratio);
-
-                *(resized_data + y * resized_width * image.channel + x * image.channel + k) = value;
+    sd::Tensor<float> cropped({target_width, target_height, image.shape()[2], image.shape()[3]});
+    for (int64_t y = 0; y < target_height; ++y) {
+        for (int64_t x = 0; x < target_width; ++x) {
+            for (int64_t c = 0; c < image.shape()[2]; ++c) {
+                cropped.index(x, y, c, 0) = resized.index(x + w_offset, y + h_offset, c, 0);
             }
         }
     }
 
-    // Clip and preprocess
-    int h_offset = std::max((int)(resized_height - target_height) / 2, 0);
-    int w_offset = std::max((int)(resized_width - target_width) / 2, 0);
-
-    sd_image_f32_t result;
-    result.width   = target_width;
-    result.height  = target_height;
-    result.channel = image.channel;
-    result.data    = (float*)malloc(target_height * target_width * image.channel * sizeof(float));
-
-    for (uint32_t k = 0; k < image.channel; k++) {
-        for (uint32_t i = 0; i < result.height; i++) {
-            for (uint32_t j = 0; j < result.width; j++) {
-                int src_y = std::min(static_cast<int>(i + h_offset), resized_height - 1);
-                int src_x = std::min(static_cast<int>(j + w_offset), resized_width - 1);
-                *(result.data + i * result.width * image.channel + j * image.channel + k) =
-                    fmin(fmax(*(resized_data + src_y * resized_width * image.channel + src_x * image.channel + k), 0.0f), 255.0f) / 255.0f;
-            }
-        }
-    }
-
-    // Free allocated memory
-    free(resized_data);
-
-    // Normalize
-    for (uint32_t k = 0; k < image.channel; k++) {
-        for (uint32_t i = 0; i < result.height; i++) {
-            for (uint32_t j = 0; j < result.width; j++) {
-                // *(result.data + i * size * image.channel + j * image.channel + k) = 0.5f;
-                int offset  = i * result.width * image.channel + j * image.channel + k;
-                float value = *(result.data + offset);
-                value       = (value - means[k]) / stds[k];
-                // value = 0.5f;
-                *(result.data + offset) = value;
-            }
-        }
-    }
-
-    return result;
+    sd::Tensor<float> normalized = sd::ops::clamp(cropped, 0.0f, 1.0f);
+    sd::Tensor<float> mean({1, 1, 3, 1}, {means[0], means[1], means[2]});
+    sd::Tensor<float> std({1, 1, 3, 1}, {stds[0], stds[1], stds[2]});
+    return (normalized - mean) / std;
 }
 
 // Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/prompt_parser.py#L345

src/util.h

@@ -7,6 +7,7 @@
 #include <vector>
 
 #include "stable-diffusion.h"
+#include "tensor.hpp"
 
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
@@ -29,20 +30,14 @@ std::string utf32_to_utf8(const std::u32string& utf32_str);
 std::u32string unicode_value_to_utf32(int unicode_value);
 // std::string sd_basename(const std::string& path);
 
-typedef struct {
-    uint32_t width;
-    uint32_t height;
-    uint32_t channel;
-    float* data;
-} sd_image_f32_t;
+sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index = 0);
 
-void normalize_sd_image_f32_t(sd_image_f32_t image, float means[3], float stds[3]);
+sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
+                                     int target_width  = -1,
+                                     int target_height = -1,
+                                     bool scale        = true);
 
-sd_image_f32_t sd_image_t_to_sd_image_f32_t(sd_image_t image);
-
-sd_image_f32_t resize_sd_image_f32_t(sd_image_f32_t image, int target_width, int target_height);
-
-sd_image_f32_t clip_preprocess(sd_image_f32_t image, int target_width, int target_height);
+sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_width, int target_height);
 
 class MmapWrapper {
 public:
@@ -69,6 +64,7 @@ protected:
 std::string path_join(const std::string& p1, const std::string& p2);
 std::vector<std::string> split_string(const std::string& str, char delimiter);
 void pretty_progress(int step, int steps, float time);
+void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
 

src/vae.hpp

@@ -2,16 +2,64 @@
 #define __VAE_HPP__
 
 #include "common_block.hpp"
+#include "tensor_ggml.hpp"
 
 struct VAE : public GGMLRunner {
 protected:
     SDVersion version;
-    bool scale_input                                       = true;
-    virtual bool _compute(const int n_threads,
-                          struct ggml_tensor* z,
-                          bool decode_graph,
-                          struct ggml_tensor** output,
-                          struct ggml_context* output_ctx) = 0;
+    bool scale_input                                      = true;
+    virtual sd::Tensor<float> _compute(const int n_threads,
+                                       const sd::Tensor<float>& z,
+                                       bool decode_graph) = 0;
+
+    static inline void scale_tensor_to_minus1_1(sd::Tensor<float>* tensor) {
+        GGML_ASSERT(tensor != nullptr);
+        for (int64_t i = 0; i < tensor->numel(); ++i) {
+            (*tensor)[i] = (*tensor)[i] * 2.0f - 1.0f;
+        }
+    }
+
+    static inline void scale_tensor_to_0_1(sd::Tensor<float>* tensor) {
+        GGML_ASSERT(tensor != nullptr);
+        for (int64_t i = 0; i < tensor->numel(); ++i) {
+            float value  = ((*tensor)[i] + 1.0f) * 0.5f;
+            (*tensor)[i] = std::max(0.0f, std::min(1.0f, value));
+        }
+    }
+
+    sd::Tensor<float> tiled_compute(const sd::Tensor<float>& input,
+                                    int n_threads,
+                                    int output_width,
+                                    int output_height,
+                                    int scale,
+                                    int p_tile_size_x,
+                                    int p_tile_size_y,
+                                    float tile_overlap_factor,
+                                    bool circular_x,
+                                    bool circular_y,
+                                    bool decode_graph,
+                                    const char* error_message,
+                                    bool silent = false) {
+        auto on_processing = [&](const sd::Tensor<float>& input_tile) {
+            auto output_tile = _compute(n_threads, input_tile, decode_graph);
+            if (output_tile.empty()) {
+                LOG_ERROR("%s", error_message);
+                return sd::Tensor<float>();
+            }
+            return output_tile;
+        };
+        return ::process_tiles_2d(input,
+                                  output_width,
+                                  output_height,
+                                  scale,
+                                  p_tile_size_x,
+                                  p_tile_size_y,
+                                  tile_overlap_factor,
+                                  circular_x,
+                                  circular_y,
+                                  on_processing,
+                                  silent);
+    }
 
 public:
     VAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
@@ -60,133 +108,109 @@ public:
         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);
-
+    sd::Tensor<float> encode(int n_threads,
+                             const sd::Tensor<float>& x,
+                             sd_tiling_params_t tiling_params,
+                             bool circular_x = false,
+                             bool circular_y = false) {
+        int64_t t0              = ggml_time_ms();
+        sd::Tensor<float> input = x;
+        sd::Tensor<float> output;
         if (scale_input) {
-            scale_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]);
+            scale_tensor_to_minus1_1(&input);
         }
 
         if (tiling_params.enabled) {
+            const int scale_factor = get_scale_factor();
+            int64_t W              = input.shape()[0] / scale_factor;
+            int64_t H              = input.shape()[1] / scale_factor;
             float tile_overlap;
             int tile_size_x, tile_size_y;
-            // 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);
+            output = tiled_compute(input,
+                                   n_threads,
+                                   static_cast<int>(W),
+                                   static_cast<int>(H),
+                                   scale_factor,
+                                   tile_size_x,
+                                   tile_size_y,
+                                   tile_overlap,
+                                   circular_x,
+                                   circular_y,
+                                   false,
+                                   "vae encode compute failed while processing a tile");
         } else {
-            _compute(n_threads, x, false, &result, work_ctx);
+            output = _compute(n_threads, input, false);
+            free_compute_buffer();
         }
-        free_compute_buffer();
 
+        if (output.empty()) {
+            LOG_ERROR("vae encode compute failed");
+            return {};
+        }
         int64_t t1 = ggml_time_ms();
         LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
-        return result;
+        return std::move(output);
     }
 
-    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]);
-        }
+    sd::Tensor<float> decode(int n_threads,
+                             const sd::Tensor<float>& x,
+                             sd_tiling_params_t tiling_params,
+                             bool decode_video = false,
+                             bool circular_x   = false,
+                             bool circular_y   = false,
+                             bool silent       = false) {
+        int64_t t0              = ggml_time_ms();
+        sd::Tensor<float> input = x;
+        sd::Tensor<float> output;
+
         if (tiling_params.enabled) {
+            const int scale_factor = get_scale_factor();
+            int64_t W              = input.shape()[0] * scale_factor;
+            int64_t H              = input.shape()[1] * scale_factor;
             float tile_overlap;
             int tile_size_x, tile_size_y;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, x->ne[0], x->ne[1]);
-
+            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
             if (!silent) {
                 LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
             }
-
-            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);
+            output = tiled_compute(
+                input,
+                n_threads,
+                static_cast<int>(W),
+                static_cast<int>(H),
+                scale_factor,
+                tile_size_x,
+                tile_size_y,
+                tile_overlap,
+                circular_x,
+                circular_y,
+                true,
+                "vae decode compute failed while processing a tile",
+                silent);
         } else {
-            if (!_compute(n_threads, x, true, &result, work_ctx)) {
-                LOG_ERROR("Failed to decode latetnts");
-                free_compute_buffer();
-                return nullptr;
-            }
+            output = _compute(n_threads, input, true);
         }
+
         free_compute_buffer();
+
+        if (output.empty()) {
+            LOG_ERROR("vae decode compute failed");
+            return {};
+        }
         if (scale_input) {
-            scale_to_0_1(result);
+            scale_tensor_to_0_1(&output);
         }
         int64_t t1 = ggml_time_ms();
         LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
-        ggml_ext_tensor_clamp_inplace(result, 0.0f, 1.0f);
-        return result;
+        return std::move(output);
     }
 
-    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, struct ggml_tensor*>& tensors, const std::string prefix)         = 0;
+    virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
+    virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents)                           = 0;
+    virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents)                           = 0;
+    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix)         = 0;
     virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
 };
 
@@ -198,34 +222,28 @@ struct FakeVAE : public VAE {
         return input_channels;
     }
 
-    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);
-        }
-        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;
+    sd::Tensor<float> _compute(const int n_threads,
+                               const sd::Tensor<float>& z,
+                               bool decode_graph) override {
+        SD_UNUSED(n_threads);
+        SD_UNUSED(decode_graph);
+        return z;
     }
 
-    ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) {
+    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
+        SD_UNUSED(rng);
         return vae_output;
     }
 
-    ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+        return latents;
     }
 
-    ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) {
-        return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
+    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+        return latents;
     }
 
-    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 {}
 
     std::string get_desc() override {
         return "fake_vae";

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,11 +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* 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"]);
@@ -414,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,]
@@ -477,24 +477,25 @@ 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(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
+                                 bool increase_ref_index                                  = false) {
+            ggml_cgraph* gf        = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
             GGML_ASSERT(x->ne[3] == 1);
-            struct ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
-
-            x         = to_backend(x);
-            context   = to_backend(context);
-            timesteps = to_backend(timesteps);
-
-            for (int i = 0; i < ref_latents.size(); i++) {
-                ref_latents[i] = to_backend(ref_latents[i]);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
             pe_vec      = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
@@ -518,66 +519,71 @@ 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);
 
             return gf;
         }
 
-        bool compute(int n_threads,
-                     struct ggml_tensor* x,
-                     struct ggml_tensor* timesteps,
-                     struct ggml_tensor* context,
-                     std::vector<ggml_tensor*> ref_latents = {},
-                     bool increase_ref_index               = false,
-                     struct ggml_tensor** output           = nullptr,
-                     struct ggml_context* output_ctx       = nullptr) {
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {},
+                                  bool increase_ref_index                           = false) {
             // 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);
             };
 
-            return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         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_ASSERT(work_ctx != nullptr);
+            ggml_context* ctx = ggml_init(params);
+            GGML_ASSERT(ctx != nullptr);
 
             {
-                // auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 16, 16, 16, 1);
+                // auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 16, 16, 16, 1);
                 // ggml_set_f32(x, 0.01f);
-                auto x = load_tensor_from_file(work_ctx, "./z_image_x.bin");
-                print_ggml_tensor(x);
+                auto x = sd::load_tensor_from_file_as_tensor<float>("./z_image_x.bin");
+                print_sd_tensor(x);
 
                 std::vector<float> timesteps_vec(1, 0.f);
-                auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
+                auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
 
-                // auto context = ggml_new_tensor_3d(work_ctx, GGML_TYPE_F32, 2560, 256, 1);
+                // auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2560, 256, 1);
                 // ggml_set_f32(context, 0.01f);
-                auto context = load_tensor_from_file(work_ctx, "./z_image_context.bin");
-                print_ggml_tensor(context);
+                auto context = sd::load_tensor_from_file_as_tensor<float>("./z_image_context.bin");
+                print_sd_tensor(context);
 
-                struct ggml_tensor* out = nullptr;
+                sd::Tensor<float> out;
 
-                int64_t t0 = ggml_time_ms();
-                compute(8, x, timesteps, context, {}, false, &out, work_ctx);
-                int64_t t1 = ggml_time_ms();
+                int64_t t0   = ggml_time_ms();
+                auto out_opt = compute(8,
+                                       x,
+                                       timesteps,
+                                       context,
+                                       {},
+                                       false);
+                int64_t t1   = ggml_time_ms();
 
-                print_ggml_tensor(out);
+                GGML_ASSERT(!out_opt.empty());
+                out = std::move(out_opt);
+                print_sd_tensor(out);
                 LOG_DEBUG("z_image test done in %lldms", t1 - t0);
             }
         }

thirdparty/CMakeLists.txt

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