fix: preserve BF16 embedding weights for get_rows (#1959)

CMakeLists.txt

@@ -229,6 +229,8 @@ file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/model/*/*.h"
     "src/model/*/*.cpp"
     "src/model/*/*.hpp"
+    "src/pipeline/*.h"
+    "src/pipeline/*.cpp"
     "src/runtime/*.h"
     "src/runtime/*.cpp"
     "src/runtime/*.hpp"

docs/backend.md

@@ -188,8 +188,9 @@ weights, compute buffers and caches must
 still fit the runner's capacity checks. Offloading weights does not guarantee
 that every resolution or frame count will fit, and auto-fit does not change a
 component to CPU computation solely because its full weights exceed VRAM.
-If a VAE decode fails, auto-fit retries with spatial tiling; supported video
-decoders try temporal tiling first and can then add spatial tiling.
+If a VAE decode fails, decoding retries with spatial tiling even when `--auto-fit`
+is off; supported video decoders try temporal tiling first and can then add
+spatial tiling. Spatial retries use half-size tiles along each latent dimension.
 
 ## Modules
 

docs/rpc.md

@@ -57,7 +57,7 @@ The RPC server acts as the worker. You must explicitly enable the **backend** (t
 
 To find the correct flags for your system, refer to the official documentation for the [`llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) repository.
 
-> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=128`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled.
+> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=160`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled.
 >
 > I recommend disabling the `LLAMA_CURL` flag to avoid unnecessary dependencies, and disabling shared library builds to avoid potential conflicts.
 
@@ -72,8 +72,8 @@ cmake .. -DGGML_RPC=ON \
     -DGGML_VULKAN=ON \        # Ensure backend is enabled
     -DGGML_BUILD_SHARED_LIBS=OFF \
     -DLLAMA_CURL=OFF \
-    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \
-    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 \
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
 cmake --build . --config Release --target rpc-server -j $(nproc)
 ```
 
@@ -86,8 +86,8 @@ cmake .. -DGGML_RPC=ON \
     -DGGML_METAL=ON \
     -DGGML_BUILD_SHARED_LIBS=OFF \
     -DLLAMA_CURL=OFF \
-    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \
-    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 \
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
 cmake --build . --config Release --target rpc-server
 ```
 
@@ -101,8 +101,8 @@ cmake .. -G "Visual Studio 17 2022" -A x64 `
     -DGGML_VULKAN=ON `
     -DGGML_BUILD_SHARED_LIBS=OFF `
     -DLLAMA_CURL=OFF `
-    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 `
-    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 `
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
 cmake --build . --config Release --target rpc-server
 ```
 

scripts/format-code.ps1

@@ -11,6 +11,8 @@ $patterns = @(
     "src/extensions/*.cpp"
     "src/extensions/*.h"
     "src/extensions/*.hpp"
+    "src/pipeline/*.cpp"
+    "src/pipeline/*.h"
     "src/runtime/*.cpp"
     "src/runtime/*.h"
     "src/runtime/*.hpp"

scripts/format-code.sh

@@ -9,6 +9,7 @@ for f in src/*.cpp src/*.h src/*.hpp \
          src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
          src/core/*.cpp src/core/*.h src/core/*.hpp \
          src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \
+         src/pipeline/*.cpp src/pipeline/*.h \
          src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \
          src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \
          src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \

src/conditioning/conditioner.hpp

@@ -1,11 +1,13 @@
 #ifndef __SD_CONDITIONING_CONDITIONER_HPP__
 #define __SD_CONDITIONING_CONDITIONER_HPP__
 
+#include <cinttypes>
 #include <cmath>
 #include <iomanip>
 #include <limits>
 #include <optional>
 #include <sstream>
+#include "core/ggml_tensor_utils.h"
 
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
@@ -160,9 +162,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::shared_ptr<CLIPTextModelRunner> text_model2;
 
     std::map<std::string, std::string> embedding_map;
-    int32_t num_custom_embeddings   = 0;
-    int32_t num_custom_embeddings_2 = 0;
+    int32_t num_custom_embeddings = 0;
     std::vector<uint8_t> token_embed_custom;
+    std::vector<uint8_t> token_embed_custom2;
     std::map<std::string, std::pair<int, int>> embedding_pos_map;
 
     FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend,
@@ -249,74 +251,93 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             LOG_ERROR("embedding '%s' failed", embd_name.c_str());
             return false;
         }
+        auto push_ids = [&](int pos_start, int pos_end, bool cached) {
+            for (int i = pos_start; i < pos_end; i++) {
+                bpe_tokens.push_back(text_model->model.vocab_size + i);
+            }
+            if (!cached) {
+                LOG_VERBOSE("embedding '%s' applied: %i token(s), custom embeddings: %i", embd_name.c_str(), pos_end - pos_start, num_custom_embeddings);
+            }
+        };
         auto iter = embedding_pos_map.find(embd_name);
         if (iter != embedding_pos_map.end()) {
             LOG_VERBOSE("embedding already read in: %s", embd_name.c_str());
-            for (int i = iter->second.first; i < iter->second.second; i++) {
-                bpe_tokens.push_back(text_model->model.vocab_size + i);
-            }
+            push_ids(iter->second.first, iter->second.second, true);
             return true;
         }
         ggml_init_params params;
-        params.mem_size        = 100 * 1024 * 1024;  // max for custom embeddings 100 MB
-        params.mem_buffer      = nullptr;
-        params.no_alloc        = false;
-        ggml_context* embd_ctx = ggml_init(params);
-        ggml_tensor* embd      = nullptr;
-        ggml_tensor* embd2     = nullptr;
-        auto on_load           = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
-            if (tensor_storage.ne[0] != text_model->model.hidden_size) {
-                if (text_model2) {
-                    if (tensor_storage.ne[0] == text_model2->model.hidden_size) {
-                        embd2       = ggml_new_tensor_2d(embd_ctx, tensor_storage.type, text_model2->model.hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
-                        *dst_tensor = embd2;
-                    } else {
-                        LOG_VERBOSE("embedding wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], text_model->model.hidden_size, text_model2->model.hidden_size);
-                        return false;
-                    }
-                } else {
-                    LOG_VERBOSE("embedding wrong hidden size, got %i, expected %i", tensor_storage.ne[0], text_model->model.hidden_size);
+        params.mem_size       = 100 * 1024 * 1024;  // max for custom embeddings 100 MB
+        params.mem_buffer     = nullptr;
+        params.no_alloc       = false;
+        auto ggml_ctx_deleter = [](ggml_context* ctx) { ggml_free(ctx); };
+        auto embd_ctx         = std::unique_ptr<ggml_context, decltype(ggml_ctx_deleter)>(ggml_init(params), ggml_ctx_deleter);
+        if (!embd_ctx.get()) {
+            LOG_ERROR("ggml_init failed when loading embeddings file");
+            return false;
+        }
+        ggml_tensor* embd    = nullptr;
+        ggml_tensor* embd2   = nullptr;
+        ggml_type embd_type  = text_model->model.get_token_embed_weight()->type;
+        ggml_type embd2_type = text_model2 ? text_model2->model.get_token_embed_weight()->type : embd_type;
+        int64_t hidden_size  = text_model->model.hidden_size;
+        int64_t hidden_size2 = text_model2 ? text_model2->model.hidden_size : 0;
+        auto on_load         = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
+            if (tensor_storage.ne[0] == hidden_size) {
+                embd = ggml_new_tensor_2d(embd_ctx.get(), embd_type, hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
+                if (embd == nullptr) {
                     return false;
                 }
-            } else {
-                embd        = ggml_new_tensor_2d(embd_ctx, tensor_storage.type, text_model->model.hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
                 *dst_tensor = embd;
+            } else if (text_model2) {
+                if (tensor_storage.ne[0] == hidden_size2) {
+                    embd2 = ggml_new_tensor_2d(embd_ctx.get(), embd2_type, hidden_size2, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
+                    if (embd2 == nullptr) {
+                        return false;
+                    }
+                    *dst_tensor = embd2;
+                } else {
+                    LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], hidden_size, hidden_size2);
+                }
+            } else {
+                LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i", tensor_storage.ne[0], hidden_size);
             }
             return true;
         };
         model_loader.set_n_threads(1);
-        model_loader.load_tensors(on_load);
-        int pos_start = num_custom_embeddings;
-        if (embd) {
-            int64_t hidden_size = text_model->model.hidden_size;
-            token_embed_custom.resize(token_embed_custom.size() + ggml_nbytes(embd));
-            memcpy((void*)(token_embed_custom.data() + num_custom_embeddings * hidden_size * ggml_type_size(embd->type)),
-                   embd->data,
-                   ggml_nbytes(embd));
-            for (int i = 0; i < embd->ne[1]; i++) {
-                bpe_tokens.push_back(text_model->model.vocab_size + num_custom_embeddings);
-                // LOG_VERBOSE("new custom token: %i", text_model.vocab_size + num_custom_embeddings);
-                num_custom_embeddings++;
-            }
-            LOG_VERBOSE("embedding '%s' applied, custom embeddings: %i", embd_name.c_str(), num_custom_embeddings);
-        }
-        if (embd2) {
-            int64_t hidden_size = text_model2->model.hidden_size;
-            token_embed_custom.resize(token_embed_custom.size() + ggml_nbytes(embd2));
-            memcpy((void*)(token_embed_custom.data() + num_custom_embeddings_2 * hidden_size * ggml_type_size(embd2->type)),
-                   embd2->data,
-                   ggml_nbytes(embd2));
-            for (int i = 0; i < embd2->ne[1]; i++) {
-                bpe_tokens.push_back(text_model2->model.vocab_size + num_custom_embeddings_2);
-                // LOG_VERBOSE("new custom token: %i", text_model.vocab_size + num_custom_embeddings);
-                num_custom_embeddings_2++;
-            }
-            LOG_VERBOSE("embedding '%s' applied, custom embeddings: %i (text model 2)", embd_name.c_str(), num_custom_embeddings_2);
-        }
-        int pos_end = num_custom_embeddings;
-        if (pos_end == pos_start) {
+        if (!model_loader.load_tensors(on_load)) {
+            LOG_ERROR("embedding '%s' failed", embd_name.c_str());
             return false;
         }
+        if (!embd && !embd2) {
+            LOG_WARN("embedding '%s' has no usable tensor", embd_name.c_str());
+            return false;
+        }
+        int pos_start      = num_custom_embeddings;
+        int64_t embd_rows  = embd ? embd->ne[1] : 0;
+        int64_t embd2_rows = embd2 ? embd2->ne[1] : 0;
+        if (embd_rows < embd2_rows) {
+            LOG_WARN("embedding '%s' has fewer rows for text model 1, zero-padding", embd_name.c_str());
+        } else if (text_model2 && embd2_rows < embd_rows) {
+            LOG_WARN("embedding '%s' has fewer rows for text model 2, zero-padding", embd_name.c_str());
+        }
+        int64_t rows      = std::max(embd_rows, embd2_rows);
+        size_t embd_bytes = hidden_size * ggml_type_size(embd_type);
+        token_embed_custom.resize(token_embed_custom.size() + embd_bytes * rows);
+        if (embd) {
+            memcpy((void*)(token_embed_custom.data() + embd_bytes * num_custom_embeddings),
+                   embd->data, embd_bytes * embd_rows);
+        }
+        if (text_model2) {
+            size_t embd2_bytes = hidden_size2 * ggml_type_size(embd2_type);
+            token_embed_custom2.resize(token_embed_custom2.size() + embd2_bytes * rows);
+            if (embd2) {
+                memcpy((void*)(token_embed_custom2.data() + embd2_bytes * num_custom_embeddings),
+                       embd2->data, embd2_bytes * embd2_rows);
+            }
+        }
+        num_custom_embeddings += (int)rows;
+        int pos_end = num_custom_embeddings;
+        push_ids(pos_start, pos_end, false);
         embedding_pos_map[embd_name] = std::pair{pos_start, pos_end};
         return true;
     }
@@ -459,7 +480,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                     auto chunk_hidden_states2 = text_model2->compute(n_threads,
                                                                      input_ids2,
                                                                      num_custom_embeddings,
-                                                                     token_embed_custom.data(),
+                                                                     token_embed_custom2.data(),
                                                                      max_token_idx,
                                                                      false,
                                                                      clip_skip,
@@ -471,7 +492,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                         pooled = text_model2->compute(n_threads,
                                                       input_ids2,
                                                       num_custom_embeddings,
-                                                      token_embed_custom.data(),
+                                                      token_embed_custom2.data(),
                                                       max_token_idx,
                                                       true,
                                                       clip_skip,
@@ -594,7 +615,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(pixel_values, return_pooled, clip_skip);
         };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
     }
 };
 
@@ -2876,7 +2897,7 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x);
         };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
     }
 };
 

src/core/backend_fit.cpp

@@ -390,6 +390,8 @@ namespace sd::backend_fit {
             retry_mode                    = tiling_params.enabled ? "spatial+temporal" : "temporal";
         } else if (!tiling_params.enabled) {
             tiling_params.enabled = true;
+            tiling_params.rel_size_x = 0.5f;
+            tiling_params.rel_size_y = 0.5f;
             if (tiling_params.tile_size_x <= 0) {
                 tiling_params.tile_size_x = 256;
             }
@@ -401,7 +403,7 @@ namespace sd::backend_fit {
             return false;
         }
 
-        LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
+        LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
                  retry_mode);
         return true;
     }

src/core/ggml_extend.cpp

@@ -0,0 +1,727 @@
+#include "core/ggml_extend.h"
+
+#include <cmath>
+#include <utility>
+
+#include "core/ggml_extend_backend.h"
+
+ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode) {
+    // reshape A
+    // swap 0th and nth axis
+    a           = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0));
+    int64_t ne1 = a->ne[1];
+    int64_t ne2 = a->ne[2];
+    int64_t ne3 = a->ne[3];
+    // make 2D
+    a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1)));
+
+    ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b)));
+
+    // reshape output (same shape as a after permutation except first dim)
+    result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3);
+    // swap back 0th and nth axis
+    result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0);
+    return result;
+}
+
+ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b) {
+    return ggml_mul(ctx,
+                    ggml_interpolate(ctx,
+                                     a,
+                                     a->ne[0] * b->ne[0],
+                                     a->ne[1] * b->ne[1],
+                                     a->ne[2] * b->ne[2],
+                                     a->ne[3] * b->ne[3],
+                                     GGML_SCALE_MODE_NEAREST),
+                    b);
+}
+
+ggml_tensor* ggml_ext_cont(ggml_context* ctx,
+                           ggml_tensor* x) {
+    if (ggml_is_contiguous(x)) {
+        return x;
+    }
+    return ggml_cont(ctx, x);
+}
+
+ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int axis0,
+                                    int axis1,
+                                    int axis2,
+                                    int axis3) {
+    int torch_axes[4] = {axis0, axis1, axis2, axis3};
+
+    int ggml_axes[4] = {0};
+    for (int i = 0; i < 4; ++i) {
+        int found = 0;
+        for (int j = 0; j < 4; ++j) {
+            if (torch_axes[j] == i) {
+                ggml_axes[i] = j;
+                found        = 1;
+                break;
+            }
+        }
+        GGML_ASSERT(found && "Invalid permute input: must be a permutation of 0-3");
+    }
+
+    return ggml_permute(ctx, x, ggml_axes[0], ggml_axes[1], ggml_axes[2], ggml_axes[3]);
+}
+
+ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+                            ggml_tensor* x,
+                            int dim,
+                            int64_t start,
+                            int64_t end,
+                            bool cont) {
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    if (x->ne[dim] == 1) {
+        return x;
+    }
+    while (start < 0) {
+        start = x->ne[dim] + start;
+    }
+    while (end < 0) {
+        end = x->ne[dim] + end;
+    }
+    GGML_ASSERT(end > start);
+    GGML_ASSERT(start >= 0 && start < x->ne[dim]);
+    GGML_ASSERT(end > start && end <= x->ne[dim]);
+
+    int64_t slice_size  = end - start;
+    int64_t slice_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
+    slice_ne[dim]       = slice_size;
+
+    x = ggml_view_4d(ctx, x,
+                     slice_ne[0], slice_ne[1], slice_ne[2], slice_ne[3],
+                     x->nb[1], x->nb[2], x->nb[3], start * x->nb[dim]);
+
+    if (cont) {
+        x = ggml_cont(ctx, x);
+    }
+
+    return x;
+}
+
+std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+                                         ggml_tensor* x,
+                                         int num,
+                                         int64_t dim,
+                                         bool cont) {
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(x->ne[dim] % num == 0);
+
+    std::vector<ggml_tensor*> chunks;
+    int64_t chunk_size  = x->ne[dim] / num;
+    int64_t stride      = chunk_size * x->nb[dim];
+    int64_t chunk_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
+    chunk_ne[dim]       = chunk_size;
+    for (int i = 0; i < num; i++) {
+        auto chunk = ggml_view_4d(
+            ctx, x,
+            chunk_ne[0], chunk_ne[1], chunk_ne[2], chunk_ne[3],
+            x->nb[1], x->nb[2], x->nb[3], stride * i);
+        if (cont) {
+            chunk = ggml_cont(ctx, chunk);
+        }
+        chunks.push_back(chunk);
+    }
+
+    return chunks;
+}
+
+ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first) {
+    // x: [ne3, ne2, ne1, ne0]
+    // return: [ne3, ne2, ne1, ne0/2]
+
+    auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false);
+    ggml_tensor* gate;
+    if (gate_first) {
+        gate = x_vec[0];
+        x    = x_vec[1];
+    } else {
+        x    = x_vec[0];
+        gate = x_vec[1];
+    }
+    gate = ggml_cont(ctx, gate);
+    gate = ggml_silu_inplace(ctx, gate);
+
+    x = ggml_mul(ctx, x, gate);  // [ne3, ne2, ne1, ne0/2]
+
+    return x;
+}
+
+ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
+                                    ggml_tensor* a) {
+    const float eps = 1e-6f;  // default eps parameter
+    return ggml_group_norm(ctx, a, 32, eps);
+}
+
+static bool ggml_ext_is_padded_1d(const ggml_tensor* x) {
+    return x->nb[0] == ggml_type_size(x->type) &&
+           x->nb[2] == x->nb[1] * x->ne[1] &&
+           x->nb[3] == x->nb[2] * x->ne[2];
+}
+
+ggml_tensor* ggml_ext_scale(ggml_context* ctx,
+                            ggml_tensor* x,
+                            float factor,
+                            bool inplace) {
+    if (!ggml_ext_is_padded_1d(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_scale_inplace(ctx, x, factor);
+    } else {
+        x = ggml_scale(ctx, x, factor);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
+                           ggml_tensor* x,
+                           bool inplace) {
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_gelu_inplace(ctx, x);
+    } else {
+        x = ggml_gelu(ctx, x);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 bool inplace) {
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_gelu_quick_inplace(ctx, x);
+    } else {
+        x = ggml_gelu_quick(ctx, x);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* w,
+                             ggml_tensor* b,
+                             bool force_prec_f32,
+                             float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+    if (x->ne[2] * x->ne[3] > 1024) {
+        // workaround: avoid ggml cuda error
+        int64_t ne2 = x->ne[2];
+        int64_t ne3 = x->ne[3];
+        x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
+        x           = ggml_mul_mat(ctx, w, x);
+        if (force_prec_f32) {
+            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        }
+        x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
+    } else {
+        x = ggml_mul_mat(ctx, w, x);
+        if (force_prec_f32) {
+            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        }
+    }
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
+    if (b != nullptr) {
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* weight_scale,
+                                           ggml_tensor* b,
+                                           int convrot_group_size,
+                                           float scale) {
+    GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f));
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+
+    ggml_tensor* fused_bias = scale == 1.f ? b : nullptr;
+    if (x->ne[2] * x->ne[3] > 1024) {
+        int64_t ne2 = x->ne[2];
+        int64_t ne3 = x->ne[3];
+        x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
+        x           = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+        x           = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
+    } else {
+        x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+    }
+
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+        if (b != nullptr) {
+            x = ggml_add_inplace(ctx, x, b);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              int lp0,
+                              int rp0,
+                              int lp1,
+                              int rp1,
+                              int lp2,
+                              int rp2,
+                              int lp3,
+                              int rp3,
+                              bool circular_x,
+                              bool circular_y) {
+    if (circular_x && circular_y) {
+        return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+    }
+
+    if (circular_x && (lp0 != 0 || rp0 != 0)) {
+        x   = ggml_pad_ext_circular(ctx, x, lp0, rp0, 0, 0, 0, 0, 0, 0);
+        lp0 = rp0 = 0;
+    }
+    if (circular_y && (lp1 != 0 || rp1 != 0)) {
+        x   = ggml_pad_ext_circular(ctx, x, 0, 0, lp1, rp1, 0, 0, 0, 0);
+        lp1 = rp1 = 0;
+    }
+
+    if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
+        ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
+            x = padded;
+        } else {
+            // Some backends (e.g. Metal) only implement right-padding for
+            // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
+            // the padding around to the left. shift < ne always holds because
+            // ne grew by lp+rp.
+            x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
+            x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+                          ggml_tensor* x,
+                          int p0,
+                          int p1,
+                          int p2,
+                          int p3,
+                          bool circular_x,
+                          bool circular_y) {
+    return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
+}
+
+ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0,
+                              int s1,
+                              int p0,
+                              int p1,
+                              int d0,
+                              int d1,
+                              bool direct,
+                              bool circular_x,
+                              bool circular_y,
+                              float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+    if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]);
+    }
+
+    if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
+        x  = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
+        p0 = 0;
+        p1 = 0;
+    }
+
+    if (direct) {
+        x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1);
+    } else {
+        x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
+    }
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int64_t IC,
+                              int s0,
+                              int s1,
+                              int s2,
+                              int p0,
+                              int p1,
+                              int p2,
+                              int d0,
+                              int d1,
+                              int d2,
+                              bool force_prec_f32) {
+    if (force_prec_f32) {
+        ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+
+        int64_t OC = w->ne[3] / IC;
+        int64_t N  = x->ne[3] / IC;
+        x          = ggml_mul_mat(ctx,
+                                  ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]),
+                                  ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC));
+        ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+
+        int64_t OD = im2col->ne[3] / N;
+        x          = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC);
+        x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
+        x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
+    } else {
+        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
+        // (e.g. Metal, see #850) do not implement. Fall back to
+        // GGML_OP_CONV_3D on those backends.
+        bool im2col_3d_supported = true;
+        if (backend != nullptr) {
+            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
+        }
+        if (im2col_3d_supported) {
+            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        } else {
+            int64_t OC = w->ne[3] / IC;
+            int64_t N  = x->ne[3] / IC;
+            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+        }
+    }
+
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]);  // [OC, 1, 1, 1]
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* w,
+                                    ggml_tensor* b,
+                                    int s2,
+                                    int p2,
+                                    int d2) {
+    x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2);  // [N, OC, T, OH * OW]
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add(ctx, x, b);
+    }
+    return x;  // [N, OC, T, OH * OW]
+}
+
+std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
+                                    ggml_tensor* qkv) {
+    qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]);  // [N, L, 3, C]
+    qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2));                    // [3, N, L, C]
+
+    int64_t offset = qkv->nb[2] * qkv->ne[2];
+    auto q         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0);  // [N, L, C]
+    auto k         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1);  // [N, L, C]
+    auto v         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2);  // [N, L, C]
+    return {q, k, v};
+}
+
+std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
+                                          ggml_tensor* qkv) {
+    int64_t W   = qkv->ne[0];
+    int64_t H   = qkv->ne[1];
+    int64_t C   = qkv->ne[2] / 3;
+    int64_t N   = qkv->ne[3];
+    int64_t nb1 = qkv->nb[1];
+    int64_t nb2 = qkv->nb[2];
+    qkv         = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N);                     // [N, 3, C, H*W]
+    qkv         = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2));  // [3, N, C, H*W]
+
+    int64_t offset = qkv->nb[2] * qkv->ne[2];
+    auto q         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0);  // [N, C, H, W]
+    auto k         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1);  // [N, C, H, W]
+    auto v         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2);  // [N, C, H, W]
+    return {q, k, v};
+}
+
+ggml_tensor* ggml_ext_full(ggml_context* ctx,
+                           float value,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3) {
+    auto one = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:one");
+    auto t   = ggml_ext_scale(ctx, one, value);             // [1,]
+    t        = ggml_repeat_4d(ctx, t, ne0, ne1, ne2, ne3);  // [ne0, ne1, ne2, ne3]
+    return t;
+}
+
+ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
+                            int64_t ne0,
+                            int64_t ne1,
+                            int64_t ne2,
+                            int64_t ne3) {
+    return ggml_ext_full(ctx, 0.f, ne0, ne1, ne2, ne3);
+}
+
+ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
+                                 ggml_tensor* x) {
+    return ggml_ext_zeros(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]);
+}
+
+ggml_tensor* ggml_ext_ones(ggml_context* ctx,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3) {
+    return ggml_ext_full(ctx, 1.f, ne0, ne1, ne2, ne3);
+}
+
+ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
+                                ggml_tensor* x) {
+    return ggml_ext_ones(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]);
+}
+
+ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a) {
+    if (sd_backend_is(backend, "Vulkan")) {
+        auto zero_index = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:zero_int");
+        auto out        = ggml_reshape_1d(ctx, a, ggml_nelements(a));
+        out             = ggml_get_rows(ctx, out, zero_index);
+        out             = ggml_reshape(ctx, out, a);
+        // auto out = ggml_cast(ctx, a, GGML_TYPE_F32);
+        return out;
+    } else {
+        auto out         = ggml_reshape_2d(ctx, a, 1, ggml_nelements(a));
+        ggml_tensor* one = ggml_ext_ones(ctx, 1, 1, 1, 1);  // [1,]
+        if (ggml_is_transposed(out)) {
+            out = ggml_mul_mat(ctx, one, out);
+        } else {
+            out = ggml_mul_mat(ctx, out, one);
+        }
+        out = ggml_reshape(ctx, out, a);
+        return out;
+    }
+}
+
+ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask,
+                                    bool skip_reshape,
+                                    bool flash_attn,
+                                    float kv_scale) {  // avoid overflow
+    int64_t L_q;
+    int64_t L_k;
+    int64_t C;
+    int64_t N;
+    int64_t d_head;
+    int64_t n_kv_head;
+    if (!skip_reshape) {
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        C         = q->ne[0];
+        N         = q->ne[2];
+        d_head    = C / n_head;
+        n_kv_head = k->ne[0] / d_head;
+
+        q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N);       // [N, L_q, n_head, d_head]
+        q = ggml_ext_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));  // [N, n_head, L_q, d_head]
+        q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N);      // [N * n_head, L_q, d_head]
+
+        k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N);    // [N, L_k, n_kv_head, d_head]
+        k = ggml_ext_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));  // [N, n_kv_head, L_k, d_head]
+        k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N);   // [N * n_kv_head, L_k, d_head]
+
+        v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N);  // [N, L_k, n_kv_head, d_head]
+    } else {
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        d_head    = v->ne[0];
+        N         = v->ne[3];
+        n_kv_head = k->ne[2] / N;
+        C         = d_head * n_head;
+    }
+
+    float scale = (1.0f / sqrt((float)d_head));
+
+    ggml_tensor* kqv = nullptr;
+
+    auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* {
+        if (kv_scale != 1.0f) {
+            k_in = ggml_ext_scale(ctx, k_in, kv_scale);
+        }
+        k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16);
+
+        v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
+        v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
+        if (kv_scale != 1.0f) {
+            v_in = ggml_ext_scale(ctx, v_in, kv_scale);
+        }
+        v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
+
+        if (mask_in != nullptr) {
+            // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped
+            // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and,
+            // unlike the manual-attention path, does not broadcast the query dimension.
+            // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over
+            // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the
+            // kernel indexes it correctly. (A bare ggml_transpose here produced a
+            // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank
+            // output for masked flash attention.)
+            if (mask_in->ne[1] != L_q) {
+                mask_in = ggml_repeat(ctx, mask_in,
+                                      ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3]));
+            }
+            mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
+        }
+
+        auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
+        if (!ggml_backend_supports_op(backend, out)) {
+            return nullptr;
+        }
+        ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
+        if (kv_scale != 1.0f) {
+            out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
+        }
+        return out;
+    };
+
+    if (flash_attn) {
+        // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
+        bool can_use_flash_attn = true;
+        if (mask != nullptr) {
+            // TODO: figure out if we can bend t5 to work too
+            can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1;
+        }
+
+        if (can_use_flash_attn) {
+            kqv = build_kqv(q, k, v, mask);
+            if (kqv != nullptr) {
+                kqv = ggml_view_4d(ctx,
+                                   kqv,
+                                   d_head,
+                                   n_head,
+                                   L_q,
+                                   N,
+                                   kqv->nb[1],
+                                   kqv->nb[2],
+                                   kqv->nb[1] * n_head,
+                                   0);
+            }
+        }
+    }
+
+    if (kqv == nullptr) {
+        // if (flash_attn) {
+        //     LOG_VERBOSE("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
+        // }
+        v = ggml_ext_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));  // [N, n_kv_head, d_head, L_k]
+        v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N);   // [N * n_kv_head, d_head, L_k]
+
+        auto kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, L_q, L_k]
+        ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
+        kq = ggml_scale_inplace(ctx, kq, scale);
+        if (mask) {
+            kq = ggml_add_inplace(ctx, kq, mask);
+        }
+        kq = ggml_soft_max_inplace(ctx, kq);
+
+        kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, L_q, d_head]
+
+        kqv = ggml_reshape_4d(ctx, kqv, d_head, L_q, n_head, N);  // [N, n_head, L_q, d_head]
+        kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3);                 // [N, L_q, n_head, d_head]
+    }
+
+    kqv = ggml_ext_cont(ctx, kqv);
+    kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N);  // [N, L_q, C]
+
+    return kqv;
+}
+
+ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 float eps) {
+    x = ggml_norm(ctx, x, eps);
+    if (w != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        if (b != nullptr) {
+            x = ggml_add_inplace(ctx, x, b);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 int num_groups) {
+    if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+    }
+
+    const float eps = 1e-6f;  // default eps parameter
+    x               = ggml_group_norm(ctx, x, num_groups, eps);
+    if (w != nullptr && b != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        // b = ggml_repeat(ctx, b, x);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_timestep_embedding(
+    ggml_context* ctx,
+    ggml_tensor* timesteps,
+    int dim,
+    int max_period,
+    float time_factor) {
+    timesteps = ggml_ext_scale(ctx, timesteps, time_factor);
+    return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
+}
+
+ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                 std::vector<ggml_tensor*>& tensors,
+                                 int dim) {
+    while (tensors.size() > 1) {
+        std::vector<ggml_tensor*> next_level;
+        for (size_t i = 0; i < tensors.size(); i += 2) {
+            if (i + 1 < tensors.size()) {
+                next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim));
+            } else {
+                next_level.push_back(tensors[i]);
+            }
+        }
+        tensors = std::move(next_level);
+    }
+    return tensors[0];
+}

src/core/ggml_extend.h

@@ -0,0 +1,235 @@
+#ifndef __SD_CORE_GGML_EXTEND_H__
+#define __SD_CORE_GGML_EXTEND_H__
+
+#include <cstdint>
+#include <vector>
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#define EPS 1e-05f
+
+static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
+
+// n-mode tensor-matrix product
+// example: 2-mode product
+// A: [ne03, k, ne01, ne00]
+// B: k rows, m columns => [k, m]
+// result is [ne03, m, ne01, ne00]
+ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0);
+
+// Kronecker product
+// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10]
+ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b);
+
+ggml_tensor* ggml_ext_cont(ggml_context* ctx,
+                           ggml_tensor* x);
+
+// torch like permute
+ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int axis0,
+                                    int axis1,
+                                    int axis2,
+                                    int axis3);
+
+ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+                            ggml_tensor* x,
+                            int dim,
+                            int64_t start,
+                            int64_t end,
+                            bool cont = true);
+
+// example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W])
+std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+                                         ggml_tensor* x,
+                                         int num,
+                                         int64_t dim,
+                                         bool cont = true);
+
+ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true);
+
+ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
+                                    ggml_tensor* a);
+
+ggml_tensor* ggml_ext_scale(ggml_context* ctx,
+                            ggml_tensor* x,
+                            float factor,
+                            bool inplace = false);
+
+ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
+                           ggml_tensor* x,
+                           bool inplace = false);
+
+ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 bool inplace = false);
+
+ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* w,
+                             ggml_tensor* b,
+                             bool force_prec_f32 = false,
+                             float scale         = 1.f);
+
+ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* weight_scale,
+                                           ggml_tensor* b,
+                                           int convrot_group_size,
+                                           float scale = 1.f);
+
+ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              int lp0,
+                              int rp0,
+                              int lp1,
+                              int rp1,
+                              int lp2,
+                              int rp2,
+                              int lp3,
+                              int rp3,
+                              bool circular_x = false,
+                              bool circular_y = false);
+
+ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+                          ggml_tensor* x,
+                          int p0,
+                          int p1,
+                          int p2          = 0,
+                          int p3          = 0,
+                          bool circular_x = false,
+                          bool circular_y = false);
+
+// w: [OC，IC, KH, KW]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N, OC, OH, OW]
+ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0          = 1,
+                              int s1          = 1,
+                              int p0          = 0,
+                              int p1          = 0,
+                              int d0          = 1,
+                              int d1          = 1,
+                              bool direct     = false,
+                              bool circular_x = false,
+                              bool circular_y = false,
+                              float scale     = 1.f);
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N*OC, OD, OH, OW]
+ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int64_t IC,
+                              int s0              = 1,
+                              int s1              = 1,
+                              int s2              = 1,
+                              int p0              = 0,
+                              int p1              = 0,
+                              int p2              = 0,
+                              int d0              = 1,
+                              int d1              = 1,
+                              int d2              = 1,
+                              bool force_prec_f32 = false);
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, ID, IH*IW]
+// b: [OC,]
+// result: [N, OC, OD, OH*OW]
+ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* w,
+                                    ggml_tensor* b,
+                                    int s2 = 1,
+                                    int p2 = 1,
+                                    int d2 = 1);
+
+// qkv: [N, L, 3*C]
+// return: ([N, L, C], [N, L, C], [N, L, C])
+std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
+                                    ggml_tensor* qkv);
+
+// qkv: [N, 3*C, H, W]
+// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W])
+std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
+                                          ggml_tensor* qkv);
+
+// Constant and cast helpers require the built-in tensors initialized by GGMLRunner.
+ggml_tensor* ggml_ext_full(ggml_context* ctx,
+                           float value,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3);
+
+ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
+                            int64_t ne0,
+                            int64_t ne1,
+                            int64_t ne2,
+                            int64_t ne3);
+
+ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
+                                 ggml_tensor* x);
+
+ggml_tensor* ggml_ext_ones(ggml_context* ctx,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3);
+
+ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
+                                ggml_tensor* x);
+
+ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a);
+
+// q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head]
+// k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head]
+// v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head]
+// mask: [N, L_q, L_k]
+// return: [N, L_q, C]
+ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask = nullptr,
+                                    bool skip_reshape = false,
+                                    bool flash_attn   = false,
+                                    float kv_scale    = 1.0f);
+
+ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 float eps = EPS);
+
+ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 int num_groups = 32);
+
+ggml_tensor* ggml_ext_timestep_embedding(
+    ggml_context* ctx,
+    ggml_tensor* timesteps,
+    int dim,
+    int max_period    = 10000,
+    float time_factor = 1.0f);
+
+ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                 std::vector<ggml_tensor*>& tensors,
+                                 int dim);
+
+#endif  // __SD_CORE_GGML_EXTEND_H__

src/core/ggml_extend_backend.cpp

@@ -965,3 +965,35 @@ const char* sd_backend_module_name(SDBackendModule module) {
     }
     return "unknown";
 }
+
+void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) {
+    if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) &&
+        !sd_backend_is_cpu(backend)) {
+        ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
+        ggml_backend_synchronize(backend);
+        return;
+    }
+
+    ggml_backend_tensor_get(tensor, data, offset, size);
+}
+
+float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor) {
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_I32 || tensor->type == GGML_TYPE_BF16);
+    float value;
+    if (tensor->type == GGML_TYPE_F32) {
+        ggml_backend_tensor_get(tensor, &value, 0, sizeof(value));
+    } else if (tensor->type == GGML_TYPE_BF16) {
+        ggml_bf16_t bf16_value;
+        ggml_backend_tensor_get(tensor, &bf16_value, 0, sizeof(bf16_value));
+        value = ggml_bf16_to_fp32(bf16_value);
+    } else if (tensor->type == GGML_TYPE_F16) {
+        ggml_fp16_t f16_value;
+        ggml_backend_tensor_get(tensor, &f16_value, 0, sizeof(f16_value));
+        value = ggml_fp16_to_fp32(f16_value);
+    } else {  // GGML_TYPE_I32
+        int int32_value;
+        ggml_backend_tensor_get(tensor, &int32_value, 0, sizeof(int32_value));
+        value = (float)int32_value;
+    }
+    return value;
+}

src/core/ggml_extend_backend.h

@@ -96,4 +96,6 @@ std::string sd_backend_resolve_name(const std::string& name);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
 bool add_rpc_devices(const std::string& servers);
+void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size);
+float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor);
 #endif  // __SD_CORE_GGML_EXTEND_BACKEND_H__

src/core/ggml_runner.cpp

@@ -2,12 +2,669 @@
 #include <map>
 #include <utility>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/layer_split_partition.h"
 #include "core/segment_graph_bindings.h"
 #include "core/segment_weight_pipeline.h"
 
 using namespace sd;
 
+void GGMLRunner::alloc_params_ctx() {
+    ggml_init_params params;
+    params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+
+    params_ctx = ggml_init(params);
+    GGML_ASSERT(params_ctx != nullptr);
+    params_tensor_set_.clear();
+    params_tensor_set_dirty_ = true;
+}
+
+void GGMLRunner::free_params_ctx() {
+    if (params_ctx != nullptr) {
+        ggml_free(params_ctx);
+        params_ctx = nullptr;
+    }
+    params_tensor_set_.clear();
+    params_tensor_set_dirty_ = true;
+}
+
+void GGMLRunner::alloc_compute_ctx() {
+    ggml_init_params params;
+    params.mem_size   = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+
+    compute_ctx = ggml_init(params);
+    GGML_ASSERT(compute_ctx != nullptr);
+}
+
+void GGMLRunner::free_compute_ctx() {
+    debug_tensors.clear();
+    if (compute_ctx != nullptr) {
+        ggml_free(compute_ctx);
+        compute_ctx = nullptr;
+    }
+    backend_tensor_data_map.clear();
+}
+
+void GGMLRunner::rebuild_params_tensor_set() {
+    if (!params_tensor_set_dirty_) {
+        return;
+    }
+    params_tensor_set_.clear();
+    if (params_ctx == nullptr) {
+        return;
+    }
+    for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) {
+        params_tensor_set_.insert(t);
+    }
+    params_tensor_set_dirty_ = false;
+}
+
+ggml_tensor* GGMLRunner::canonical_param_tensor(ggml_tensor* tensor) {
+    for (auto* current = tensor; current != nullptr; current = current->view_src) {
+        if (params_tensor_set_.count(current) != 0)
+            return current;
+    }
+    return nullptr;
+}
+
+std::vector<ggml_tensor*> GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf) {
+    std::vector<ggml_tensor*> used_params;
+    rebuild_params_tensor_set();
+    if (gf == nullptr || params_tensor_set_.empty()) {
+        return used_params;
+    }
+
+    std::unordered_set<const ggml_tensor*> seen_params;
+    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+    seen_params.reserve(static_cast<size_t>(n_leafs));
+    for (int i = 0; i < n_leafs; ++i) {
+        ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
+        ggml_tensor* param_leaf = canonical_param_tensor(leaf);
+        if (param_leaf != nullptr &&
+            seen_params.insert(param_leaf).second) {
+            used_params.push_back(param_leaf);
+        }
+    }
+    return used_params;
+}
+
+void GGMLRunner::evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+    if (tensors.empty()) {
+        return;
+    }
+    auto manager = residency_manager.lock();
+    if (manager != nullptr) {
+        manager->evict_compute_backend_params(tensors);
+    }
+}
+
+void GGMLRunner::prepare_build_in_tensor_before() {
+    one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
+    ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one");
+    set_backend_tensor_data(one_tensor, one_vec.data());
+
+    zero_int_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, 1);
+    ggml_set_name(zero_int_tensor, "ggml_runner_build_in_tensor:zero_int");
+    set_backend_tensor_data(zero_int_tensor, zero_int_vec.data());
+}
+
+void GGMLRunner::prepare_build_in_tensor_after(ggml_cgraph* gf) {
+    ggml_build_forward_expand(gf, one_tensor);
+    ggml_build_forward_expand(gf, zero_int_tensor);
+}
+
+ggml_cgraph* GGMLRunner::new_graph_custom(size_t graph_size) {
+    if (weight_adapter) {
+        graph_size += weight_adapter->get_extra_graph_size();
+    }
+    return ggml_new_graph_custom(compute_ctx, graph_size, false);
+}
+
+ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) {
+    prepare_build_in_tensor_before();
+    ggml_cgraph* gf = get_graph();
+    if (gf == nullptr) {
+        return nullptr;
+    }
+    if (ggml_graph_n_nodes(gf) > 0) {
+        auto result = ggml_graph_node(gf, -1);
+        ggml_set_name(result, final_result_name.c_str());
+    }
+    for (const auto& entry : debug_tensors) {
+        if (entry.first != nullptr) {
+            ggml_build_forward_expand(gf, entry.first);
+        }
+    }
+    for (const auto& entry : cache_.outputs()) {
+        if (entry.second != nullptr) {
+            ggml_build_forward_expand(gf, entry.second);
+        }
+    }
+    prepare_build_in_tensor_after(gf);
+    return gf;
+}
+
+bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph,
+                                       ggml_cgraph** gf_out) {
+    GGML_ASSERT(gf_out != nullptr);
+
+    reset_compute_ctx();
+    ggml_cgraph* gf = get_compute_graph(get_graph);
+    if (gf == nullptr) {
+        free_compute_ctx();
+        return false;
+    }
+
+    *gf_out = gf;
+    return true;
+}
+
+ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const {
+    if (tensor == nullptr || tensor->buffer == nullptr) {
+        return nullptr;
+    }
+    if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+        ggml_backend_buffer_is_host(tensor->buffer)) {
+        return nullptr;
+    }
+    ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+    if (dev == nullptr) {
+        return nullptr;
+    }
+    if (ggml_backend_get_device(runtime_backend) == dev) {
+        return runtime_backend;
+    }
+    for (ggml_backend_t backend : extra_runtime_backends) {
+        if (ggml_backend_get_device(backend) == dev) {
+            return backend;
+        }
+    }
+    return nullptr;
+}
+
+void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
+    if (sched == nullptr || gf == nullptr) {
+        return;
+    }
+    ggml_backend_t current = runtime_backend;
+    const int n_nodes      = ggml_graph_n_nodes(gf);
+    for (int i = 0; i < n_nodes; i++) {
+        ggml_tensor* node    = ggml_graph_node(gf, i);
+        auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
+        if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
+            current = node_assignment->second;
+        }
+        for (int s = 0; s < GGML_MAX_SRC; s++) {
+            ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
+            if (weight_backend != nullptr) {
+                if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
+                    current = weight_backend;
+                }
+            }
+        }
+        if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+            node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+            continue;
+        }
+        if (ggml_backend_supports_op(current, node)) {
+            ggml_backend_sched_set_tensor_backend(sched, node, current);
+        }
+    }
+}
+
+size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const {
+    backend      = backend == nullptr ? runtime_backend : backend;
+    size_t bytes = workspace_.bytes(backend);
+    if (backend == runtime_backend) {
+        const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend));
+        bytes                    = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes;
+        const size_t cut_bytes   = cut_cache_.resident_bytes(ggml_backend_get_device(backend));
+        bytes                    = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes;
+    }
+    return bytes;
+}
+
+void GGMLRunner::sync_runtime_residency() {
+    if (auto manager = residency_manager.lock()) {
+        manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
+                                          runtime_backend, retained_runtime_buffer_bytes());
+        for (auto backend : extra_runtime_backends) {
+            manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
+                                              backend, retained_runtime_buffer_bytes(backend));
+        }
+    }
+}
+
+std::optional<sd::Tensor<float>> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) {
+    if (tensor == nullptr) {
+        LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label);
+        return std::nullopt;
+    }
+    if (tensor->type != GGML_TYPE_F32) {
+        LOG_ERROR("%s %s tensor type mismatch: got %s",
+                  get_desc().c_str(),
+                  label,
+                  ggml_type_name(tensor->type));
+        return std::nullopt;
+    }
+    ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
+    if (buf == nullptr) {
+        LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
+                  get_desc().c_str(),
+                  label,
+                  tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                  ggml_op_name(tensor->op),
+                  tensor->buffer,
+                  tensor->view_src,
+                  tensor->view_src ? tensor->view_src->buffer : nullptr,
+                  tensor->data);
+        return std::nullopt;
+    }
+
+    return sd::make_sd_tensor_from_ggml<float>(tensor);
+}
+
+void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) {
+    GGML_ASSERT(gf != nullptr);
+    std::unordered_set<const ggml_tensor*> graph_tensor_set;
+    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+    const int n_nodes = ggml_graph_n_nodes(gf);
+    graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
+    for (int i = 0; i < n_leafs; ++i) {
+        graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
+    }
+    for (int i = 0; i < n_nodes; ++i) {
+        graph_tensor_set.insert(ggml_graph_node(gf, i));
+    }
+
+    for (auto& kv : backend_tensor_data_map) {
+        auto tensor = kv.first;
+        auto data   = kv.second;
+        if (tensor == nullptr || data == nullptr) {
+            continue;
+        }
+        const char* name = ggml_get_name(tensor);
+        if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
+            continue;
+        }
+        if (tensor->buffer == nullptr) {
+            LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
+                     get_desc().c_str(),
+                     name != nullptr ? name : "",
+                     (long long)tensor->ne[0],
+                     (long long)tensor->ne[1],
+                     (long long)tensor->ne[2],
+                     (long long)tensor->ne[3],
+                     ggml_type_name(tensor->type));
+            continue;
+        }
+
+        ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+        if (buf == nullptr) {
+            LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
+                     get_desc().c_str(),
+                     tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                     tensor ? ggml_op_name(tensor->op) : "<null>",
+                     data,
+                     tensor ? tensor->view_src : nullptr,
+                     (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
+            continue;
+        }
+
+        ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+    }
+
+    if (clear_after_copy) {
+        backend_tensor_data_map.clear();
+    }
+}
+
+bool GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf,
+                                        GraphCutPlan* plan_out) {
+    GGML_ASSERT(plan_out != nullptr);
+    GGML_ASSERT(gf != nullptr);
+    *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                                 gf,
+                                                 &graph_cut_plan_cache_,
+                                                 params_tensor_set_,
+                                                 get_desc().c_str());
+    return true;
+}
+
+bool GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                                    GraphCutPlan* plan_out) {
+    return resolve_graph_cut_plan(gf, plan_out);
+}
+
+bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+    graph_cut_layer_split_node_assignments_.clear();
+    if (!graph_cut_layer_split_enabled) {
+        return true;
+    }
+    if (!is_multi_device()) {
+        LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
+        return false;
+    }
+
+    GraphCutPlan plan;
+    if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
+        return false;
+    }
+    if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+        auto manager = residency_manager.lock();
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+            return false;
+        }
+        std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+        if (!graph_params.empty() &&
+            !manager->assign_compute_backend(graph_params, runtime_backend)) {
+            LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                      get_desc().c_str(),
+                      sd::layer_split_backend_device_display_name(runtime_backend).c_str());
+            return false;
+        }
+        for (ggml_tensor* param : graph_params) {
+            if (param != nullptr) {
+                graph_cut_layer_split_assignments_[param] = runtime_backend;
+            }
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; i++) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr) {
+                graph_cut_layer_split_node_assignments_[node] = runtime_backend;
+            }
+        }
+        if (!graph_cut_layer_split_primary_notice_logged_) {
+            LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                     get_desc().c_str(),
+                     sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                     graph_params.size());
+            graph_cut_layer_split_primary_notice_logged_ = true;
+        } else {
+            LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                        get_desc().c_str(),
+                        sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                        graph_params.size());
+        }
+        return true;
+    }
+
+    std::vector<ggml_backend_t> split_backends;
+    split_backends.reserve(extra_runtime_backends.size() + 1);
+    split_backends.push_back(runtime_backend);
+    for (ggml_backend_t backend : extra_runtime_backends) {
+        if (backend != nullptr) {
+            split_backends.push_back(backend);
+        }
+    }
+
+    auto manager = residency_manager.lock();
+    if (manager == nullptr) {
+        LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+        return false;
+    }
+
+    sd::GraphCutLayerSplitAssignment assignment;
+    auto canonicalize_param = [this](ggml_tensor* tensor) {
+        return canonical_param_tensor(tensor);
+    };
+    if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
+                                             gf,
+                                             plan,
+                                             split_backends,
+                                             graph_cut_layer_split_backend_vram_limits_,
+                                             max_graph_vram_bytes,
+                                             graph_cut_layer_split_assignments_,
+                                             canonicalize_param,
+                                             &assignment)) {
+        return false;
+    }
+
+    for (size_t i = 0; i < split_backends.size(); i++) {
+        if (assignment.tensors_by_backend[i].empty()) {
+            continue;
+        }
+        if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+            LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                      get_desc().c_str(),
+                      sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+            return false;
+        }
+    }
+
+    graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
+    sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
+
+    return true;
+}
+
+bool GGMLRunner::runner_start() {
+    if (runner_started_) {
+        return true;
+    }
+    cache_.clear();
+    workspace_.set_extra_backends(extra_runtime_backends);
+    if (auto manager = residency_manager.lock()) {
+        manager->set_workspace_reclaimer(reinterpret_cast<uintptr_t>(this), [this]() {
+            if (!workspace_.release()) {
+                return false;
+            }
+            sync_runtime_residency();
+            return true;
+        });
+    }
+    runner_started_ = true;
+    return true;
+}
+
+void GGMLRunner::runner_end() {
+    GGML_ASSERT(!graph_active_);
+    if (!runner_started_) {
+        return;
+    }
+    workspace_.release();
+    cache_.clear();
+    logged_compute_bytes_.clear();
+    logged_segment_count_ = 0;
+    if (auto manager = residency_manager.lock()) {
+        manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
+        std::vector<ggml_tensor*> tensors;
+        for (auto tensor : params_tensor_set_) {
+            auto* parameter = manager->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
+            if (parameter != nullptr)
+                tensors.push_back(parameter);
+        }
+        manager->evict_compute_backend_params(tensors);
+        manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
+    }
+    runner_started_ = false;
+}
+
+GGMLRunner::GGMLRunner(ggml_backend_t backend,
+                       std::shared_ptr<DeviceResidencyManager> manager)
+    : runtime_backend(backend),
+      cache_(backend),
+      cut_cache_(backend),
+      workspace_(backend),
+      residency_manager(manager) {
+    GGML_ASSERT(runtime_backend != nullptr);
+    alloc_params_ctx();
+}
+
+GGMLRunner::~GGMLRunner() {
+    runner_end();
+    free_compute_ctx();
+    free_params_ctx();
+}
+
+GGMLRunnerContext GGMLRunner::get_context() {
+    GGMLRunnerContext runner_ctx;
+    runner_ctx.ggml_ctx              = compute_ctx;
+    runner_ctx.backend               = runtime_backend;
+    runner_ctx.flash_attn_enabled    = flash_attn_enabled;
+    runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
+    runner_ctx.circular_x_enabled    = circular_x_enabled;
+    runner_ctx.circular_y_enabled    = circular_y_enabled;
+    runner_ctx.weight_adapter        = weight_adapter;
+    runner_ctx.debug_tensors         = &debug_tensors;
+    runner_ctx.get_cache_tensor      = [this](const std::string& name) {
+        return this->get_cache_tensor_by_name(name);
+    };
+    runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
+        this->cache(name, tensor);
+    };
+    runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
+        this->set_backend_tensor_data(tensor, data);
+    };
+    return runner_ctx;
+}
+
+void GGMLRunner::reset_compute_ctx() {
+    free_compute_ctx();
+    alloc_compute_ctx();
+}
+
+void GGMLRunner::free_cache_ctx_and_buffer() {
+    cache_.clear();
+    sync_runtime_residency();
+}
+
+void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
+    // The scheduler only allocates standalone data tensors when they are
+    // marked as graph inputs. The flag is harmless for single-backend graphs.
+    ggml_set_input(tensor);
+    backend_tensor_data_map[tensor] = data;
+}
+
+ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
+    GGML_ASSERT(compute_ctx != nullptr);
+    if (tensor == nullptr) {
+        return nullptr;
+    }
+    // it's performing a compute, check if backend isn't cpu
+    if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
+        // pass input tensors to gpu memory
+        auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
+
+        set_backend_tensor_data(backend_tensor, tensor->data);
+        return backend_tensor;
+    } else {
+        return tensor;
+    }
+}
+
+void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
+    if (tensor != nullptr && tensor->view_src != nullptr) {
+        tensor = ggml_cont(compute_ctx, tensor);
+    }
+    if (tensor != nullptr) {
+        ggml_set_output(tensor);
+    }
+    cache_.stage(name, tensor);
+}
+
+std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
+                                                     int n_threads,
+                                                     bool auto_runner_end,
+                                                     bool no_return,
+                                                     const std::function<bool()>& read_outputs) {
+    if (graph_active_) {
+        LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str());
+        return std::nullopt;
+    }
+    if (!runner_start()) {
+        runner_end();
+        return std::nullopt;
+    }
+    struct RunnerEndGuard {
+        GGMLRunner& runner;
+        bool enabled;
+        ~RunnerEndGuard() {
+            if (enabled) {
+                runner.runner_end();
+            }
+        }
+    } runner_guard{*this, auto_runner_end};
+    graph_active_ = true;
+    bool success  = false;
+    struct GraphEndGuard {
+        GGMLRunner& runner;
+        const bool& success;
+        ~GraphEndGuard() {
+            runner.workspace_.segment_end();
+            runner.cache_.graph_end(false);
+            runner.cut_cache_.clear();
+            runner.free_compute_ctx();
+            runner.graph_active_ = false;
+            if (!success) {
+                runner.workspace_.release();
+            }
+            runner.sync_runtime_residency();
+        }
+    } graph_guard{*this, success};
+
+    ggml_cgraph* graph = nullptr;
+    if (!prepare_compute_graph(get_graph, &graph)) {
+        return std::nullopt;
+    }
+    params_tensor_set_dirty_ = true;
+    rebuild_params_tensor_set();
+    if (auto manager = residency_manager.lock()) {
+        for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) {
+            auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i));
+            if (parameter != nullptr)
+                params_tensor_set_.insert(parameter);
+        }
+    }
+    auto output = execute_graph(graph, n_threads, no_return, read_outputs);
+    success     = output.has_value();
+    if (success) {
+        cache_.graph_end(true);
+    }
+    return output;
+}
+
+void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) {
+    graph_cut_layer_split_enabled = enabled;
+    if (!enabled) {
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+    }
+}
+
+void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
+    graph_cut_layer_split_backend_vram_limits_ = limits;
+    graph_cut_layer_split_assignments_.clear();
+    graph_cut_layer_split_node_assignments_.clear();
+    graph_cut_layer_split_primary_notice_logged_ = false;
+}
+
+void GGMLRunner::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+    extra_runtime_backends.clear();
+    for (ggml_backend_t backend : backends) {
+        if (backend == nullptr || backend == runtime_backend) {
+            continue;
+        }
+        if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
+            extra_runtime_backends.end()) {
+            extra_runtime_backends.push_back(backend);
+        }
+    }
+    workspace_.set_extra_backends(extra_runtime_backends);
+    graph_cut_layer_split_assignments_.clear();
+    graph_cut_layer_split_node_assignments_.clear();
+    graph_cut_layer_split_primary_notice_logged_ = false;
+}
+
 static size_t add_bytes(size_t a, size_t b) {
     return b > SIZE_MAX - a ? SIZE_MAX : a + b;
 }
@@ -275,7 +932,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             }
             if (!no_return) {
                 auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
-                output      = read_graph_tensor<float>(result, "output");
+                output      = read_graph_tensor(result, "output");
                 if (!output.has_value()) {
                     return fail_segment("output readback");
                 }

src/core/ggml_runner.h

@@ -0,0 +1,350 @@
+#ifndef __SD_CORE_GGML_RUNNER_H__
+#define __SD_CORE_GGML_RUNNER_H__
+
+#include <cstddef>
+#include <functional>
+#include <map>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "core/compute_workspace.h"
+#include "core/ggml_graph_cut.h"
+#include "core/runner_cache.h"
+#include "core/tensor_ggml.hpp"
+#include "core/util.h"
+#include "device_residency_manager.h"
+
+/* SDXL with LoRA requires more space */
+#define MAX_PARAMS_TENSOR_NUM 32768
+#define MAX_GRAPH_SIZE 327680
+
+struct WeightAdapter {
+    struct ForwardParams {
+        enum class op_type_t {
+            OP_LINEAR,
+            OP_CONV2D,
+        } op_type;
+        struct {
+            bool force_prec_f32 = false;
+            float scale         = 1.f;
+        } linear;
+        struct conv2d_params_t {
+            int s0          = 1;
+            int s1          = 1;
+            int p0          = 0;
+            int p1          = 0;
+            int d0          = 1;
+            int d1          = 1;
+            bool direct     = false;
+            bool circular_x = false;
+            bool circular_y = false;
+            float scale     = 1.f;
+        } conv2d;
+    };
+    virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0;
+    virtual ggml_tensor* forward_with_lora(ggml_context* ctx,
+                                           ggml_backend_t backend,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* b,
+                                           const std::string& prefix,
+                                           ForwardParams forward_params)                                                              = 0;
+    virtual ggml_tensor* add_lora_to_output(ggml_context* ctx,
+                                            ggml_backend_t backend,
+                                            ggml_tensor* x,
+                                            ggml_tensor* w,
+                                            ggml_tensor* output,
+                                            const std::string& prefix,
+                                            ForwardParams forward_params)                                                             = 0;
+    virtual size_t get_extra_graph_size()                                                                                             = 0;
+};
+
+struct GGMLRunnerContext {
+    ggml_backend_t backend                                           = nullptr;
+    ggml_context* ggml_ctx                                           = nullptr;
+    bool flash_attn_enabled                                          = false;
+    bool conv2d_direct_enabled                                       = false;
+    bool circular_x_enabled                                          = false;
+    bool circular_y_enabled                                          = false;
+    ggml_tensor* ip_context                                          = nullptr;
+    float ip_scale                                                   = 1.0f;
+    std::shared_ptr<WeightAdapter> weight_adapter                    = nullptr;
+    std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
+    std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
+    std::function<void(const std::string&, ggml_tensor*)> cache_tensor;
+    std::function<void(ggml_tensor*, const void*)> set_backend_tensor_data;
+    std::map<std::pair<ggml_tensor*, int>, ggml_tensor*> int8_convrot_cache;
+
+    void capture_tensor(const std::string& name, ggml_tensor* tensor) {
+        if (debug_tensors == nullptr || tensor == nullptr) {
+            return;
+        }
+        ggml_tensor* snapshot = ggml_cont(ggml_ctx, tensor);
+        ggml_tensor* dst      = ggml_dup_tensor(ggml_ctx, snapshot);
+        snapshot              = ggml_cpy(ggml_ctx, snapshot, dst);
+        ggml_set_output(snapshot);
+        debug_tensors->push_back({snapshot, name});
+    }
+
+    ggml_tensor* load_cache_tensor(const std::string& name) const {
+        if (!get_cache_tensor) {
+            return nullptr;
+        }
+        return get_cache_tensor(name);
+    }
+
+    void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
+        if (!cache_tensor || tensor == nullptr) {
+            return;
+        }
+        cache_tensor(name, tensor);
+    }
+
+    void bind_backend_tensor_data(ggml_tensor* tensor, const void* data) const {
+        if (!set_backend_tensor_data || tensor == nullptr || data == nullptr) {
+            return;
+        }
+        set_backend_tensor_data(tensor, data);
+    }
+};
+
+struct GGMLRunner {
+private:
+    std::map<ggml_backend_t, size_t> logged_compute_bytes_;
+    size_t logged_segment_count_ = 0;
+
+    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
+    std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
+                                                     size_t pending_cache_bytes) const;
+    bool fits(const std::vector<DeviceMemoryRequest>& requests,
+              const std::vector<ggml_tensor*>& params) const;
+    bool execute_segment(ggml_cgraph* graph, int n_threads);
+    std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
+
+protected:
+    typedef std::function<ggml_cgraph*()> get_graph_cb_t;
+    using GraphCutPlan = sd::ggml_graph_cut::Plan;
+
+    ggml_backend_t runtime_backend = nullptr;
+
+    ggml_context* params_ctx = nullptr;
+
+    sd::RunnerCache cache_;
+    sd::GraphCutTensorCache cut_cache_;
+    sd::ComputeWorkspace workspace_;
+    ggml_context* compute_ctx = nullptr;
+    bool runner_started_      = false;
+    bool graph_active_        = false;
+
+    size_t max_graph_vram_bytes        = 0;
+    bool graph_cut_layer_split_enabled = false;
+    std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
+
+    std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
+    bool multi_device_eval_callback_warned = false;
+
+    std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
+    std::weak_ptr<DeviceResidencyManager> residency_manager;
+    bool params_tensor_set_dirty_ = true;
+
+    std::vector<float> one_vec = {1.f};
+    ggml_tensor* one_tensor    = nullptr;
+
+    std::vector<int> zero_int_vec = {0};
+    ggml_tensor* zero_int_tensor  = nullptr;
+
+    std::map<ggml_tensor*, const void*> backend_tensor_data_map;
+    std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
+    const std::string final_result_name = "ggml_runner_final_result_tensor";
+
+    bool flash_attn_enabled    = false;
+    bool conv2d_direct_enabled = false;
+    bool circular_x_enabled    = false;
+    bool circular_y_enabled    = false;
+
+    sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
+    std::unordered_set<const ggml_tensor*> params_tensor_set_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
+    bool graph_cut_layer_split_primary_notice_logged_ = false;
+
+    template <typename T>
+    static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
+        if (!tensor.has_value()) {
+            return {};
+        }
+        return std::move(*tensor);
+    }
+
+    template <typename T>
+    static sd::Tensor<T> restore_trailing_singleton_dims(std::optional<sd::Tensor<T>> tensor,
+                                                         size_t expected_dim) {
+        return restore_trailing_singleton_dims(take_or_empty(std::move(tensor)), expected_dim);
+    }
+
+    template <typename T>
+    static sd::Tensor<T> restore_trailing_singleton_dims(sd::Tensor<T> tensor,
+                                                         size_t expected_dim) {
+        if (tensor.empty()) {
+            return tensor;
+        }
+        while (static_cast<size_t>(tensor.dim()) < expected_dim) {
+            tensor.unsqueeze_(tensor.dim());
+        }
+        return tensor;
+    }
+
+    void alloc_params_ctx();
+
+    void free_params_ctx();
+
+    void alloc_compute_ctx();
+
+    void free_compute_ctx();
+
+    void rebuild_params_tensor_set();
+
+    ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
+
+    std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf);
+
+    void evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors);
+
+    void prepare_build_in_tensor_before();
+
+    void prepare_build_in_tensor_after(ggml_cgraph* gf);
+
+    ggml_cgraph* new_graph_custom(size_t graph_size);
+
+    ggml_cgraph* get_compute_graph(get_graph_cb_t get_graph);
+
+    bool prepare_compute_graph(get_graph_cb_t get_graph,
+                               ggml_cgraph** gf_out);
+
+    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
+
+    // Weightless ops have no scheduler anchor, so pin them to the most recent
+    // weight device. Views must stay unpinned or cross-device copies can be
+    // skipped for their consumers.
+    void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr);
+
+    bool is_multi_device() const {
+        return !extra_runtime_backends.empty();
+    }
+
+    size_t reusable_compute_buffer_bytes() const {
+        return workspace_.bytes(runtime_backend);
+    }
+
+    size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const;
+
+    void sync_runtime_residency();
+
+    std::optional<sd::Tensor<float>> read_graph_tensor(ggml_tensor* tensor, const char* label);
+
+    void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true);
+
+    bool resolve_graph_cut_plan(ggml_cgraph* gf,
+                                GraphCutPlan* plan_out);
+
+    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                            GraphCutPlan* plan_out);
+
+    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf);
+
+public:
+    bool runner_start();
+
+    bool runner_started() const { return runner_started_; }
+
+    void runner_end();
+
+public:
+    virtual std::string get_desc() = 0;
+
+    GGMLRunner(ggml_backend_t backend,
+               std::shared_ptr<DeviceResidencyManager> manager = nullptr);
+
+    virtual ~GGMLRunner();
+
+    virtual GGMLRunnerContext get_context();
+
+    void reset_compute_ctx();
+
+public:
+    void free_cache_ctx_and_buffer();
+
+    // do copy after alloc graph
+    void set_backend_tensor_data(ggml_tensor* tensor, const void* data);
+
+    template <typename T>
+    ggml_tensor* make_input(const sd::Tensor<T>& tensor) {
+        ggml_tensor* input = sd::make_ggml_tensor(compute_ctx, tensor, false);
+        set_backend_tensor_data(input, tensor.data());
+        return input;
+    }
+
+    template <typename T>
+    ggml_tensor* make_optional_input(const sd::Tensor<T>& tensor) {
+        if (tensor.empty()) {
+            return nullptr;
+        }
+        return make_input(tensor);
+    }
+
+    template <typename T>
+    ggml_tensor* make_optional_input(const sd::Tensor<T>* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        return make_input(*tensor);
+    }
+
+    ggml_tensor* to_backend(ggml_tensor* tensor);
+
+    void cache(const std::string name, ggml_tensor* tensor);
+
+    ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
+        return cache_.get(name);
+    }
+
+    std::optional<sd::Tensor<float>> compute(get_graph_cb_t get_graph,
+                                             int n_threads,
+                                             bool auto_runner_end                      = true,
+                                             bool no_return                            = false,
+                                             const std::function<bool()>& read_outputs = {});
+
+    void set_flash_attention_enabled(bool enabled) {
+        flash_attn_enabled = enabled;
+    }
+
+    void set_conv2d_direct_enabled(bool enabled) {
+        conv2d_direct_enabled = enabled;
+    }
+
+    void set_circular_axes(bool circular_x, bool circular_y) {
+        circular_x_enabled = circular_x;
+        circular_y_enabled = circular_y;
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {
+        weight_adapter = adapter;
+    }
+
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) {
+        max_graph_vram_bytes = max_vram_bytes;
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled);
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits);
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends);
+};
+
+#endif  // __SD_CORE_GGML_RUNNER_H__

src/core/ggml_tensor_utils.cpp

@@ -0,0 +1,428 @@
+#include "core/ggml_tensor_utils.h"
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+
+#include "core/ggml_extend_backend.h"
+#include "core/rng.hpp"
+
+void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr<RNG> rng) {
+    uint32_t n                        = (uint32_t)ggml_nelements(tensor);
+    std::vector<float> random_numbers = rng->randn(n);
+    for (uint32_t i = 0; i < n; i++) {
+        ggml_ext_im_set_f32_1d(tensor, i, random_numbers[i]);
+    }
+}
+
+void print_ggml_tensor(ggml_tensor* tensor, bool shape_only, const char* mark) {
+    printf("%s (%s): shape(%zu, %zu, %zu, %zu)\n", mark, ggml_type_name(tensor->type), tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    int range = 3;
+    for (int i3 = 0; i3 < tensor->ne[3]; i3++) {
+        if (i3 >= range && i3 + range < tensor->ne[3]) {
+            continue;
+        }
+        for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+            if (i2 >= range && i2 + range < tensor->ne[2]) {
+                continue;
+            }
+            for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                if (i1 >= range && i1 + range < tensor->ne[1]) {
+                    continue;
+                }
+                for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                    if (i0 >= range && i0 + range < tensor->ne[0]) {
+                        continue;
+                    }
+                    if (tensor->type == GGML_TYPE_F32) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_ext_tensor_get_f32(tensor, i0, i1, i2, i3));
+                    } else if (tensor->type == GGML_TYPE_F16) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_fp16_to_fp32(ggml_ext_tensor_get_f16(tensor, i0, i1, i2, i3)));
+                    } else if (tensor->type == GGML_TYPE_I32) {
+                        printf("  [%d, %d, %d, %d] = %i3\n", i3, i2, i1, i0, ggml_ext_tensor_get_i32(tensor, i0, i1, i2, i3));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i3 = 0; i3 < n3; i3++) {
+        for (int64_t i2 = 0; i2 < n2; i2++) {
+            for (int64_t i1 = 0; i1 < n1; i1++) {
+                for (int64_t i0 = 0; i0 < n0; i0++) {
+                    fn(tensor, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i = 0; i < ggml_nelements(tensor); i++) {
+        fn(tensor, i);
+    }
+}
+
+void ggml_ext_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    ggml_ext_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float a_value = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3);
+        float b_value = ggml_ext_tensor_get_f32(b, i0, i1, i2, i3);
+        if (abs(a_value - b_value) > gap) {
+            LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value);
+        }
+    });
+}
+
+ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        LOG_ERROR("failed to open '%s'", file_path.c_str());
+        return nullptr;
+    }
+    int32_t n_dims;
+    int32_t length;
+    int32_t ttype;
+
+    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));
+
+    LOG_VERBOSE("load_tensor_from_file %d %d %d", n_dims, length, ttype);
+
+    if (file.eof()) {
+        LOG_ERROR("incomplete file '%s'", file_path.c_str());
+        return nullptr;
+    }
+
+    int32_t nelements = 1;
+    int32_t ne[4]     = {1, 1, 1, 1};
+    for (int i = 0; i < n_dims; ++i) {
+        file.read(reinterpret_cast<char*>(&ne[i]), sizeof(ne[i]));
+        nelements *= ne[i];
+    }
+    std::string name(length, 0);
+    file.read(&name[0], length);
+    ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]);
+    const size_t bpe    = ggml_type_size(ggml_type(ttype));
+    file.read(reinterpret_cast<char*>(tensor->data), ggml_nbytes(tensor));
+    return tensor;
+}
+
+// __STATIC_INLINE__ void save_tensor_to_file(const std::string& file_name, ggml_tensor* tensor, const std::string & name) {
+//     std::string file_name_ = file_name + ".tensor";
+//     std::string name_ = name;
+//     std::ofstream file("./" + file_name_, std::ios::binary);
+//     file.write(reinterpret_cast<char*>(&tensor->n_dims), sizeof(tensor->n_dims));
+//     int len = (int)name_.size();
+//     file.write(reinterpret_cast<char*>(&len), sizeof(len));
+//     int ttype = (int)tensor->type;
+//     file.write(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+//     for (int i = 0; i < tensor->n_dims; ++i) {
+//         int ne_ = (int) tensor->ne[i];
+//         file.write(reinterpret_cast<char*>(&ne_), sizeof(ne_));
+//     }
+//     file.write(&name_[0], len);
+//     char* data = nullptr;
+//     file.write((char*)tensor->data, ggml_nbytes(tensor));
+//     file.close();
+// }
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data) {
+    int64_t width    = input->ne[0];
+    int64_t height   = input->ne[1];
+    int64_t channels = input->ne[2];
+    GGML_ASSERT(input->type == GGML_TYPE_F32);
+    if (image_data == nullptr) {
+        image_data = (uint8_t*)malloc(width * height * channels);
+    }
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            for (int k = 0; k < channels; k++) {
+                float value                                               = ggml_ext_tensor_get_f32(input, ix, iy, k);
+                *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video) {
+    int64_t width  = input->ne[0];
+    int64_t height = input->ne[1];
+    int64_t channels;
+    if (video) {
+        channels = input->ne[3];
+    } else {
+        channels = input->ne[2];
+    }
+    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    uint8_t* image_data = (uint8_t*)malloc(width * height * channels);
+    for (int ih = 0; ih < height; ih++) {
+        for (int iw = 0; iw < width; iw++) {
+            for (int ic = 0; ic < channels; ic++) {
+                float value;
+                if (video) {
+                    value = ggml_ext_tensor_get_f32(input, iw, ih, idx, ic);
+                } else {
+                    value = ggml_ext_tensor_get_f32(input, iw, ih, ic, idx);
+                }
+                *(image_data + ih * width * channels + iw * channels + ic) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+void sd_image_to_ggml_tensor(sd_image_t image,
+                             ggml_tensor* tensor,
+                             bool scale) {
+    GGML_ASSERT(image.width == tensor->ne[0]);
+    GGML_ASSERT(image.height == tensor->ne[1]);
+    GGML_ASSERT(image.channel == tensor->ne[2]);
+    GGML_ASSERT(1 == tensor->ne[3]);
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+    ggml_ext_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float value = sd_image_get_f32(image, i0, i1, i2, scale);
+        ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2, i3);
+    });
+}
+
+void ggml_ext_tensor_apply_mask(ggml_tensor* image_data,
+                                ggml_tensor* mask,
+                                ggml_tensor* output,
+                                float masked_value) {
+    int64_t width    = output->ne[0];
+    int64_t height   = output->ne[1];
+    int64_t channels = output->ne[2];
+    float rescale_mx = 1.f * mask->ne[0] / output->ne[0];
+    float rescale_my = 1.f * mask->ne[1] / output->ne[1];
+    GGML_ASSERT(output->type == GGML_TYPE_F32);
+    for (int ix = 0; ix < width; ix++) {
+        for (int iy = 0; iy < height; iy++) {
+            int mx  = (int)(ix * rescale_mx);
+            int my  = (int)(iy * rescale_my);
+            float m = ggml_ext_tensor_get_f32(mask, mx, my);
+            m       = round(m);  // inpaint models need binary masks
+            ggml_ext_tensor_set_f32(mask, m, mx, my);
+            for (int k = 0; k < channels; k++) {
+                float value = ggml_ext_tensor_get_f32(image_data, ix, iy, k);
+                value       = (1 - m) * (value - masked_value) + masked_value;
+                ggml_ext_tensor_set_f32(output, value, ix, iy, k);
+            }
+        }
+    }
+}
+
+float ggml_ext_tensor_mean(ggml_tensor* src) {
+    float mean        = 0.0f;
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        mean += data[i] / nelements * 1.0f;
+    }
+    return mean;
+}
+
+void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    int64_t nelements = ggml_nelements(a);
+    float* vec_a      = (float*)a->data;
+    float* vec_b      = (float*)b->data;
+    for (int i = 0; i < nelements; i++) {
+        vec_a[i] = vec_a[i] + vec_b[i];
+    }
+}
+
+void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        data[i] = data[i] * scale;
+    }
+}
+
+void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val < min ? min : (val > max ? max : val);
+    }
+}
+
+ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx,
+                                    ggml_tensor* a,
+                                    ggml_tensor* b,
+                                    int dim) {
+    int64_t ne[GGML_MAX_DIMS];
+    for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+        if (d == dim) {
+            ne[d] = a->ne[d] + b->ne[d];
+            continue;
+        }
+        GGML_ASSERT(a->ne[d] == b->ne[d]);
+        ne[d] = a->ne[d];
+    }
+    ggml_tensor* result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne);
+    int64_t o[4]        = {0, 0, 0, 0};
+    o[dim]              = a->ne[dim];
+
+    float v;
+    for (int i3 = 0; i3 < result->ne[3]; i3++) {
+        for (int i2 = 0; i2 < result->ne[2]; i2++) {
+            for (int i1 = 0; i1 < result->ne[1]; i1++) {
+                for (int i0 = 0; i0 < result->ne[0]; i0++) {
+                    if (i0 < a->ne[0] && i1 < a->ne[1] && i2 < a->ne[2] && i3 < a->ne[3]) {
+                        v = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3);
+                    } else {
+                        v = ggml_ext_tensor_get_f32(b, i0 - o[0], i1 - o[1], i2 - o[2], i3 - o[3]);
+                    }
+
+                    ggml_ext_tensor_set_f32(result, v, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+    return result;
+}
+
+void scale_to_minus1_1(ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val * 2.0f - 1.0f;
+    }
+}
+
+void scale_to_0_1(ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = (val + 1.0f) * 0.5f;
+    }
+}
+
+ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx,
+                                   const std::vector<float>& vec) {
+    ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx,
+                                       const std::vector<int>& vec) {
+    ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+std::vector<float> arange(float start, float end, float step) {
+    std::vector<float> result;
+
+    for (float value = start; value < end; value += step) {
+        result.push_back(value);
+    }
+
+    return result;
+}
+
+std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                      int dim,
+                                      int max_period,
+                                      bool flip_sin_to_cos,
+                                      float scale) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    size_t N = timesteps.size();
+    std::vector<float> embedding(N * dim, 0.f);
+    int half = dim / 2;
+    std::vector<float> freqs(half);
+    for (int i = 0; i < half; ++i) {
+        freqs[i] = (float)std::exp(-std::log(max_period) * i / half);
+    }
+    for (int i = 0; i < N; ++i) {
+        for (int j = 0; j < half; ++j) {
+            float arg = timesteps[i] * freqs[j] * scale;
+            if (flip_sin_to_cos) {
+                embedding[i * dim + j]        = std::cos(arg);
+                embedding[i * dim + j + half] = std::sin(arg);
+            } else {
+                embedding[i * dim + j]        = std::sin(arg);
+                embedding[i * dim + j + half] = std::cos(arg);
+            }
+        }
+    }
+    return embedding;
+}
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            ggml_tensor* embedding,
+                            int dim,
+                            int max_period) {
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+}
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            sd::Tensor<float>* embedding,
+                            int dim,
+                            int max_period) {
+    GGML_ASSERT(embedding != nullptr);
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    if (embedding->numel() != static_cast<int64_t>(embedding_vec.size())) {
+        embedding->resize({dim, static_cast<int64_t>(timesteps.size())});
+    }
+    std::copy(embedding_vec.begin(), embedding_vec.end(), embedding->values().begin());
+}
+
+ggml_tensor* new_timestep_embedding(ggml_context* ctx,
+                                    std::vector<float> timesteps,
+                                    int dim,
+                                    int max_period) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    ggml_tensor* embedding           = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps.size());
+    if (embedding->data != nullptr) {
+        memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+    } else {
+        ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding));
+    }
+    return embedding;
+}
+
+size_t ggml_tensor_num(ggml_context* ctx) {
+    size_t num = 0;
+    for (ggml_tensor* t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+        num++;
+    }
+    return num;
+}

src/core/ggml_tensor_utils.h

@@ -0,0 +1,210 @@
+#ifndef __SD_CORE_GGML_TENSOR_UTILS_H__
+#define __SD_CORE_GGML_TENSOR_UTILS_H__
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <functional>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "core/tensor.hpp"
+#include "core/util.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+#include "stable-diffusion.h"
+
+class RNG;
+
+__STATIC_INLINE__ int align_up_offset(int n, int multiple) {
+    return (multiple - n % multiple) % multiple;
+}
+
+__STATIC_INLINE__ int align_up(int n, int multiple) {
+    return n + align_up_offset(n, multiple);
+}
+
+void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr<RNG> rng);
+
+__STATIC_INLINE__ void ggml_ext_tensor_set_f32(ggml_tensor* tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]) = value;
+}
+
+__STATIC_INLINE__ float ggml_ext_tensor_get_f32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    if (tensor->buffer != nullptr) {
+        float value;
+        ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(float));
+        return value;
+    }
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    return *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ int ggml_ext_tensor_get_i32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    if (tensor->buffer != nullptr) {
+        int value;
+        ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(int));
+        return value;
+    }
+    GGML_ASSERT(tensor->nb[0] == sizeof(int));
+    return *(int*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ ggml_fp16_t ggml_ext_tensor_get_f16(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t));
+    return *(ggml_fp16_t*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int64_t iw, int64_t ih, int64_t ic, bool scale = true) {
+    float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic);
+    if (scale) {
+        value /= 255.f;
+    }
+    return value;
+}
+
+void print_ggml_tensor(ggml_tensor* tensor, bool shape_only = false, const char* mark = "");
+
+template <typename T>
+__STATIC_INLINE__ void print_sd_tensor(const sd::Tensor<T>& tensor, bool shape_only = false, const char* mark = "") {
+    printf("%s: shape(", mark);
+    for (size_t i = 0; i < static_cast<size_t>(tensor.dim()); ++i) {
+        printf("%s%lld", i == 0 ? "" : ", ", static_cast<long long>(tensor.shape()[i]));
+    }
+    printf(")\n");
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    if (tensor.empty()) {
+        return;
+    }
+    int range                  = 3;
+    std::vector<int64_t> shape = tensor.shape();
+    while (shape.size() < 4) {
+        shape.push_back(1);
+    }
+    for (int64_t i3 = 0; i3 < shape[3]; i3++) {
+        if (i3 >= range && i3 + range < shape[3]) {
+            continue;
+        }
+        for (int64_t i2 = 0; i2 < shape[2]; i2++) {
+            if (i2 >= range && i2 + range < shape[2]) {
+                continue;
+            }
+            for (int64_t i1 = 0; i1 < shape[1]; i1++) {
+                if (i1 >= range && i1 + range < shape[1]) {
+                    continue;
+                }
+                for (int64_t i0 = 0; i0 < shape[0]; i0++) {
+                    if (i0 >= range && i0 + range < shape[0]) {
+                        continue;
+                    }
+                    size_t offset = static_cast<size_t>(i0 + shape[0] * (i1 + shape[1] * (i2 + shape[2] * i3)));
+                    printf("  [%lld, %lld, %lld, %lld] = ", static_cast<long long>(i3), static_cast<long long>(i2), static_cast<long long>(i1), static_cast<long long>(i0));
+                    if constexpr (std::is_same_v<T, float>) {
+                        printf("%f\n", tensor[static_cast<int64_t>(offset)]);
+                    } else if constexpr (std::is_same_v<T, ggml_fp16_t>) {
+                        printf("%f\n", ggml_fp16_to_fp32(tensor[static_cast<int64_t>(offset)]));
+                    } else if constexpr (std::is_same_v<T, int32_t>) {
+                        printf("%d\n", tensor[static_cast<int64_t>(offset)]);
+                    } else if constexpr (std::is_same_v<T, int64_t>) {
+                        printf("%lld\n", static_cast<long long>(tensor[static_cast<int64_t>(offset)]));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn);
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn);
+
+void ggml_ext_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap = 0.1f);
+
+ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path);
+
+__STATIC_INLINE__ float sigmoid(float x) {
+    return 1 / (1.0f + expf(-x));
+}
+
+// SPECIAL OPERATIONS WITH TENSORS
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data = nullptr);
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video = false);
+
+void sd_image_to_ggml_tensor(sd_image_t image,
+                             ggml_tensor* tensor,
+                             bool scale = true);
+
+void ggml_ext_tensor_apply_mask(ggml_tensor* image_data,
+                                ggml_tensor* mask,
+                                ggml_tensor* output,
+                                float masked_value = 0.5f);
+
+float ggml_ext_tensor_mean(ggml_tensor* src);
+
+// a = a+b
+void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b);
+
+void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale);
+
+void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max);
+
+ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx,
+                                    ggml_tensor* a,
+                                    ggml_tensor* b,
+                                    int dim);
+
+// convert values from [0, 1] to [-1, 1]
+void scale_to_minus1_1(ggml_tensor* src);
+
+// convert values from [-1, 1] to [0, 1]
+void scale_to_0_1(ggml_tensor* src);
+
+ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx,
+                                   const std::vector<float>& vec);
+
+ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx,
+                                       const std::vector<int>& vec);
+
+std::vector<float> arange(float start, float end, float step = 1.f);
+
+// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
+std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                      int dim,
+                                      int max_period       = 10000,
+                                      bool flip_sin_to_cos = true,
+                                      float scale          = 1.f);
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            ggml_tensor* embedding,
+                            int dim,
+                            int max_period = 10000);
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            sd::Tensor<float>* embedding,
+                            int dim,
+                            int max_period = 10000);
+
+ggml_tensor* new_timestep_embedding(ggml_context* ctx,
+                                    std::vector<float> timesteps,
+                                    int dim,
+                                    int max_period = 10000);
+
+size_t ggml_tensor_num(ggml_context* ctx);
+
+#endif  // __SD_CORE_GGML_TENSOR_UTILS_H__

src/core/ordered_map.hpp

@@ -38,9 +38,19 @@ public:
             insert(kv);
     }
 
-    OrderedMap(const OrderedMap&)                = default;
-    OrderedMap(OrderedMap&&) noexcept            = default;
-    OrderedMap& operator=(const OrderedMap&)     = default;
+    OrderedMap(const OrderedMap& other) {
+        for (const auto& value : other) {
+            insert(value);
+        }
+    }
+    OrderedMap(OrderedMap&&) noexcept = default;
+    OrderedMap& operator=(const OrderedMap& other) {
+        if (this != &other) {
+            OrderedMap copy(other);
+            swap(copy);
+        }
+        return *this;
+    }
     OrderedMap& operator=(OrderedMap&&) noexcept = default;
 
     // --- element access ---
@@ -174,4 +184,4 @@ public:
     }
 };
 
-#endif  // __SD_CORE_ORDERED_MAP_HPP__
+#endif  // __SD_CORE_ORDERED_MAP_HPP__

src/core/util.cpp

@@ -15,6 +15,7 @@
 #include <thread>
 #include <unordered_set>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 #include "runtime/preprocessing.hpp"
 
 #if defined(__APPLE__) && defined(__MACH__)
@@ -618,6 +619,25 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
     va_end(args);
 }
 
+void sd_ggml_log_callback(ggml_log_level level, const char* text, void*) {
+    switch (level) {
+        case GGML_LOG_LEVEL_DEBUG:
+            LOG_VERBOSE(text);
+            break;
+        case GGML_LOG_LEVEL_INFO:
+            LOG_INFO(text);
+            break;
+        case GGML_LOG_LEVEL_WARN:
+            LOG_WARN(text);
+            break;
+        case GGML_LOG_LEVEL_ERROR:
+            LOG_ERROR(text);
+            break;
+        default:
+            LOG_VERBOSE(text);
+    }
+}
+
 void sd_set_log_callback(sd_log_cb_t cb, void* data) {
     sd_log_cb      = cb;
     sd_log_cb_data = data;

src/core/util.h

@@ -11,6 +11,14 @@
 #include "ggml-backend.h"
 #include "stable-diffusion.h"
 
+#ifndef __STATIC_INLINE__
+#define __STATIC_INLINE__ static inline
+#endif
+
+#ifndef SD_UNUSED
+#define SD_UNUSED(x) (void)(x)
+#endif
+
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
@@ -79,6 +87,7 @@ 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, ...);
+void sd_ggml_log_callback(ggml_log_level level, const char* text, void*);
 
 ggml_type sd_type_to_ggml_type(sd_type_t sdtype);
 

src/detailer.cpp

@@ -676,7 +676,7 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) {
     model_manager = std::make_shared<ModelManager>();
     model_manager->set_n_threads(n_threads);
     model_manager->set_enable_mmap(false);
-    ModelLoader& loader = model_manager->loader();
+    ModelLoader loader;
     if (!loader.init_from_file(detector_path)) {
         LOG_ERROR("failed to load ADetailer detector: '%s'", detector_path.c_str());
         return false;
@@ -696,7 +696,8 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) {
 
     std::map<std::string, ggml_tensor*> tensors;
     detector->get_param_tensors(tensors);
-    if (!model_manager->register_param_tensors("YOLOv8",
+    if (!model_manager->set_loader(loader) ||
+        !model_manager->register_param_tensors(ModelComponent::Detector,
                                                std::move(tensors),
                                                backend_manager.params_backend_is_disk(SDBackendModule::DETECTOR)
                                                    ? ModelManager::ResidencyMode::Disk

src/device_residency_manager.h

@@ -48,9 +48,10 @@ struct DeviceResidencyManager {
                                                const std::vector<ggml_tensor*>& required_params) const = 0;
     virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
                                         ggml_backend_t compute_backend)                                = 0;
-    virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                              = 0;
-    virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors)              = 0;
-    virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors)                = 0;
+    virtual ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const { return nullptr; }
+    virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                 = 0;
+    virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
+    virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors)   = 0;
     virtual WeightResidencyInfo inspect_compute_backend_params(
         const std::vector<ggml_tensor*>& tensors) const          = 0;
     virtual void update_runtime_residency(uintptr_t owner_id,

src/extensions/generation_extension.h

@@ -19,7 +19,7 @@ struct GenerationExtensionInitContext {
     const sd_ctx_params_t* params;
     SDVersion version;
     const String2TensorStorage& tensor_storage_map;
-    ModelLoader& model_loader;
+    bool photomaker_source_available;
     std::shared_ptr<ModelManager> model_manager;
     int n_threads;
     std::function<bool(SDBackendModule)> ensure_backend_pair;
@@ -39,7 +39,8 @@ struct GenerationExtensionConditionContext {
 struct GenerationExtension {
     virtual ~GenerationExtension() = default;
 
-    virtual const char* name() const = 0;
+    virtual ModelComponent component() const = 0;
+    const char* name() const { return model_component_name(component()); }
     virtual bool is_enabled() const {
         return false;
     }

src/extensions/photomaker_extension.cpp

@@ -1,3 +1,4 @@
+#include <cinttypes>
 #include "extensions/generation_extension.h"
 
 #include <algorithm>
@@ -108,8 +109,8 @@ struct PhotoMakerExtension : public GenerationExtension {
     SDCondition id_condition;
     int start_merge_step = -1;
 
-    const char* name() const override {
-        return "photomaker";
+    ModelComponent component() const override {
+        return ModelComponent::PhotoMaker;
     }
 
     bool is_enabled() const override {
@@ -118,7 +119,7 @@ struct PhotoMakerExtension : public GenerationExtension {
 
     bool init(const GenerationExtensionInitContext& ctx) override {
         model_path = SAFE_STR(ctx.params->photo_maker_path);
-        if (model_path.empty()) {
+        if (model_path.empty() || !ctx.photomaker_source_available) {
             return true;
         }
 
@@ -127,13 +128,7 @@ struct PhotoMakerExtension : public GenerationExtension {
         }
 
         PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1;
-        LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", model_path.c_str());
-        if (!ctx.model_loader.init_from_file_and_convert_name(model_path, "pmid.")) {
-            LOG_WARN("loading stacked ID embedding from '%s' failed", model_path.c_str());
-            return true;
-        }
-
-        pmid_model = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
+        pmid_model           = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
                                                            ctx.tensor_storage_map,
                                                            "pmid",
                                                            ctx.version,

src/extensions/pulid_extension.cpp

@@ -79,8 +79,8 @@ struct PuLIDExtension : public GenerationExtension {
     sd::Tensor<float> id_embedding;
     float id_weight = 1.0f;
 
-    const char* name() const override {
-        return "pulid";
+    ModelComponent component() const override {
+        return ModelComponent::PuLID;
     }
 
     bool is_enabled() const override {

src/model/adapter/ip_adapter.hpp

@@ -1,8 +1,10 @@
 #ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
 #define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model_loader.h"
 
 namespace IPAdapter {
@@ -200,7 +202,7 @@ namespace IPAdapter {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(image_embeds);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
         }
     };
 

src/model/adapter/lora.hpp

@@ -2,7 +2,13 @@
 #define __SD_MODEL_ADAPTER_LORA_HPP__
 
 #include <mutex>
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model.h"
+#include "model/adapter/lora_ops.h"
 #include "model_loader.h"
 #include "model_manager.h"
 
@@ -17,25 +23,31 @@ struct LoraModel : public GGMLRunner {
     std::set<std::string> skipped_incompatible_lora_tensors;
     std::set<std::string> warned_incompatible_model_tensors;
     std::string file_path;
-    std::shared_ptr<ModelManager> model_manager;
-    ggml_backend_t params_backend = nullptr;
-    bool load_failed              = false;
-    bool applied                  = false;
-    bool tensor_preprocessed      = false;
+    ggml_backend_t params_backend              = nullptr;
+    bool load_failed                           = false;
+    bool applied                               = false;
+    bool tensor_preprocessed                   = false;
+    ModelLoader::FileId source_file            = 0;
+    SDVersion source_version                   = VERSION_COUNT;
+    ModelManager::ResidencyMode residency_mode = ModelManager::ResidencyMode::ParamBackend;
+    bool params_follow_compute                 = false;
+    std::vector<ggml_tensor*> registered_params;
+    std::map<ggml_tensor*, float> scalar_values;
 
     typedef std::function<bool(const std::string&)> filter_t;
 
-    LoraModel(const std::string& lora_id,
-              ggml_backend_t backend,
-              ggml_backend_t params_backend_,
-              const std::string& file_path          = "",
-              std::string prefix                    = "",
-              SDVersion version                     = VERSION_COUNT,
-              std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>())
-        : GGMLRunner(backend, manager), lora_id(lora_id), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
-        prefix = "lora." + prefix;
-        if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix, version)) {
-            load_failed = true;
+    LoraModel(const std::string& id, ggml_backend_t backend, ggml_backend_t params, std::shared_ptr<ModelManager> manager, ModelLoader::FileId file, SDVersion version, ModelManager::ResidencyMode mode = ModelManager::ResidencyMode::ParamBackend, bool follow_compute = false)
+        : GGMLRunner(backend, manager), lora_id(id), params_backend(params), source_file(file), source_version(version), residency_mode(mode), params_follow_compute(follow_compute) {
+        load_failed = source_file == 0 || manager == nullptr || manager->loader().file_revision(source_file) == 0;
+        if (!load_failed) {
+            file_path = manager->loader().file_path(source_file);
+        }
+    }
+
+    ~LoraModel() override {
+        runner_end();
+        if (auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock())) {
+            GGML_ASSERT(manager->unregister_param_tensors(registered_params));
         }
     }
 
@@ -43,95 +55,65 @@ struct LoraModel : public GGMLRunner {
         return "lora";
     }
 
-    bool load_from_file(int n_threads, filter_t filter = nullptr) {
-        LOG_INFO("loading LoRA from '%s'", file_path.c_str());
-
-        if (load_failed) {
-            LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
+    bool init_params(int n_threads, filter_t filter = nullptr) {
+        auto model_manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
+        if (model_manager == nullptr)
             return false;
-        }
-
-        std::unordered_map<std::string, TensorStorage> tensors_to_create;
-        std::mutex lora_mutex;
-        bool dry_run          = true;
-        auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-            if (dry_run) {
-                const std::string& name = tensor_storage.name;
-
-                if (filter && !filter(name)) {
-                    return true;
-                }
-
-                {
-                    std::lock_guard<std::mutex> lock(lora_mutex);
-                    tensors_to_create[name] = tensor_storage;
-                }
-            } else {
-                const std::string& name = tensor_storage.name;
-                auto iter               = lora_tensors.find(name);
-                if (iter != lora_tensors.end()) {
-                    *dst_tensor = iter->second;
-                }
-            }
-            return true;
-        };
-
-        if (model_manager != nullptr) {
-            model_manager->set_n_threads(n_threads);
-        }
-        ModelLoader& model_loader = model_manager->loader();
-        model_loader.load_tensors(on_new_tensor_cb);
-
-        if (tensors_to_create.empty()) {
-            return true;
-        }
-
-        for (const auto& pair : tensors_to_create) {
-            const auto& name   = pair.first;
-            const auto& ts     = pair.second;
-            ggml_tensor* real  = ggml_new_tensor(params_ctx,
-                                                 ts.type,
-                                                 ts.n_dims,
-                                                 ts.ne);
-            lora_tensors[name] = real;
-        }
-
+        if (load_failed || !registered_params.empty())
+            return false;
+        model_manager->set_n_threads(n_threads);
+        const auto sources = model_manager->loader().file_tensors(source_file, source_version);
         std::map<std::string, ggml_tensor*> tensors;
-        for (const auto& pair : lora_tensors) {
-            tensors[pair.first] = pair.second;
+        std::map<std::string, ggml_tensor*> scalars;
+        std::set<std::string> scalar_names;
+        for (const auto& [name, source] : sources) {
+            if (is_unused_tensor(name) || (filter && !filter(name)))
+                continue;
+            const bool scalar  = source.nelements() == 1 && (ends_with(name, ".alpha") || ends_with(name, ".scale"));
+            auto* tensor       = ggml_new_tensor(params_ctx, scalar ? GGML_TYPE_F32 : source.type, source.n_dims, source.ne);
+            lora_tensors[name] = tensor;
+            if (scalar) {
+                tensor->data  = &scalar_values[tensor];
+                scalars[name] = tensor;
+                scalar_names.insert(name);
+            } else {
+                tensors[name] = tensor;
+            }
         }
-        if (model_manager == nullptr ||
-            !model_manager->register_param_tensors("LoRA",
-                                                   std::move(tensors),
-                                                   ModelManager::ResidencyMode::ParamBackend,
-                                                   runtime_backend,
-                                                   params_backend) ||
-            !model_manager->validate_registered_tensors()) {
-            LOG_ERROR("lora model manager registration failed");
+        // These values are consumed while constructing the graph, before weight preparation.
+        if (!scalars.empty()) {
+            auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
+                auto found = scalars.find(source.name);
+                *dst       = found == scalars.end() ? nullptr : found->second;
+                return true;
+            };
+            if (!model_manager->loader().load_file_tensors(source_file, source_version, callback, scalar_names))
+                return false;
+        }
+        if (!model_manager->register_param_tensors(ModelComponent::LoRA, tensors, residency_mode,
+                                                   runtime_backend, params_backend, nullptr, false, params_follow_compute,
+                                                   nullptr, source_file, source_version))
             return false;
-        }
-        std::vector<ggml_tensor*> lora_params;
-        lora_params.reserve(lora_tensors.size());
-        for (const auto& pair : lora_tensors) {
-            lora_params.push_back(pair.second);
-        }
-        if (!model_manager->prepare_params(lora_params)) {
-            LOG_ERROR("lora model manager prepare params failed");
-            return false;
-        }
+        for (const auto& entry : tensors)
+            registered_params.push_back(entry.second);
+        return model_manager->validate_registered_tensors();
+    }
 
-        LOG_VERBOSE("finished loaded lora");
-        return true;
+    float scalar_value(ggml_tensor* tensor) const {
+        auto found = scalar_values.find(tensor);
+        return found != scalar_values.end() ? found->second : ggml_ext_backend_tensor_get_f32(tensor);
     }
 
     void release_loaded_tensors() {
         runner_end();
-        model_manager.reset();
+        if (auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock())) {
+            GGML_ASSERT(manager->unregister_param_tensors(registered_params));
+        }
+        registered_params.clear();
         free_params_ctx();
         alloc_params_ctx();
-        model_manager     = std::make_shared<ModelManager>();
-        residency_manager = model_manager;
         lora_tensors.clear();
+        scalar_values.clear();
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
         skipped_incompatible_lora_tensors.clear();
@@ -235,12 +217,12 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_value = scalar_value(iter->second);
                 applied_lora_tensors.insert(scale_name);
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha = scalar_value(iter->second);
                     scale_value = alpha / rank;
                     // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
                     applied_lora_tensors.insert(alpha_name);
@@ -389,7 +371,7 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = hada_1_down->ne[ggml_n_dims(hada_1_down) - 1];
             iter         = lora_tensors.find(alpha_name);
             if (iter != lora_tensors.end()) {
-                float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                float alpha = scalar_value(iter->second);
                 scale_value = alpha / rank;
                 applied_lora_tensors.insert(alpha_name);
             }
@@ -502,7 +484,7 @@ struct LoraModel : public GGMLRunner {
             float scale_value = 1.0f;
             iter              = lora_tensors.find(alpha_name);
             if (iter != lora_tensors.end()) {
-                float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                float alpha = scalar_value(iter->second);
                 scale_value = alpha / rank;
                 applied_lora_tensors.insert(alpha_name);
             }
@@ -663,7 +645,7 @@ struct LoraModel : public GGMLRunner {
                 float scale_value = 1.0f;
                 iter              = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha = scalar_value(iter->second);
                     scale_value = alpha / rank;
                 }
 
@@ -790,12 +772,12 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value       = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_value       = scalar_value(iter->second);
                 scale_tensor_name = scale_name;
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha       = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha       = scalar_value(iter->second);
                     scale_value       = alpha / rank;
                     scale_tensor_name = alpha_name;
                     // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
@@ -943,7 +925,7 @@ struct LoraModel : public GGMLRunner {
         return gf;
     }
 
-    void apply(std::map<std::string, ggml_tensor*> model_tensors,
+    bool apply(std::map<std::string, ggml_tensor*> model_tensors,
                const std::set<std::string>& model_tensor_names,
                SDVersion version,
                int n_threads,
@@ -957,17 +939,18 @@ struct LoraModel : public GGMLRunner {
             }
             return true;
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        auto result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs);
         if (!result.has_value()) {
             LOG_ERROR("LoRA graph execution failed");
         }
         stat(!warn_unused);
         original_tensor_to_final_tensor.clear();
         runner_end();
+        return result.has_value();
     }
 
-    void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
-        apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
+    bool apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
+        return apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
     }
 
     void stat(bool at_runntime = false) {

src/model/adapter/lora_ops.cpp

@@ -0,0 +1,207 @@
+#include "model/adapter/lora_ops.h"
+
+#include <cmath>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+
+ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+                                 ggml_tensor* lora_down,
+                                 ggml_tensor* lora_up,
+                                 ggml_tensor* lora_mid) {
+    ggml_tensor* updown;
+    // flat lora tensors to multiply it
+    int64_t lora_up_rows  = lora_up->ne[ggml_n_dims(lora_up) - 1];
+    lora_up               = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
+    auto lora_down_n_dims = ggml_n_dims(lora_down);
+    // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work)
+    lora_down_n_dims       = (lora_down_n_dims + lora_down_n_dims % 2);
+    int64_t lora_down_rows = lora_down->ne[lora_down_n_dims - 1];
+    lora_down              = ggml_reshape_2d(ctx, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
+
+    // ggml_mul_mat requires tensor b transposed
+    lora_down = ggml_cont(ctx, ggml_transpose(ctx, lora_down));
+    if (lora_mid == nullptr) {
+        updown = ggml_mul_mat(ctx, lora_up, lora_down);
+        updown = ggml_cont(ctx, ggml_transpose(ctx, updown));
+    } else {
+        // undoing tucker decomposition for conv layers.
+        // lora_mid  has shape (3,    3,   Rank, Rank)
+        // lora_down has shape (Rank, In,  1,    1)
+        // lora_up   has shape (Rank, Out, 1,    1)
+        // conv layer shape is (3,    3,   Out,  In)
+        updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2);
+        updown = ggml_cont(ctx, updown);
+    }
+    return updown;
+}
+
+ggml_tensor* ggml_ext_lokr_forward(
+    ggml_context* ctx,
+    ggml_backend_t backend,
+    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+    ggml_tensor* w1,   // Outer C (Full rank)
+    ggml_tensor* w1a,  // Outer A (Low rank part 1)
+    ggml_tensor* w1b,  // Outer B (Low rank part 2)
+    ggml_tensor* w2,   // Inner BA (Full rank)
+    ggml_tensor* w2a,  // Inner A (Low rank part 1)
+    ggml_tensor* w2b,  // Inner B (Low rank part 2)
+    bool is_conv,
+    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+    float scale) {
+    GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr)));
+    GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr)));
+
+    int uq = (w1 != nullptr) ? (int)w1->ne[0] : (int)w1a->ne[0];
+    int up = (w1 != nullptr) ? (int)w1->ne[1] : (int)w1b->ne[1];
+
+    int q_actual = is_conv ? (int)h->ne[2] : (int)h->ne[0];
+    int vq       = q_actual / uq;
+
+    int vp = (w2 != nullptr) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1])
+                             : (int)w2a->ne[1];
+    GGML_ASSERT(q_actual == (uq * vq) && "Input dimension mismatch for LoKR split");
+
+    ggml_tensor* hb;
+
+    if (!is_conv) {
+        int batch          = (int)h->ne[1];
+        int merge_batch_uq = batch;
+        int merge_batch_vp = batch;
+
+        if (sd_backend_is(backend, "Vulkan")) {
+            if (batch > 1) {
+                // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend
+                int max_batch    = 65535;
+                int max_batch_uq = max_batch / uq;
+                merge_batch_uq   = 1;
+                for (int i = max_batch_uq; i > 0; i--) {
+                    if (batch % i == 0) {
+                        merge_batch_uq = i;
+                        break;
+                    }
+                }
+
+                int max_batch_vp = max_batch / vp;
+                merge_batch_vp   = 1;
+                for (int i = max_batch_vp; i > 0; i--) {
+                    if (batch % i == 0) {
+                        merge_batch_vp = i;
+                        break;
+                    }
+                }
+            }
+        }
+
+        ggml_tensor* h_split = ggml_reshape_3d(ctx, h, vq, uq * merge_batch_uq, batch / merge_batch_uq);
+        if (w2 != nullptr) {
+            hb = ggml_mul_mat(ctx, w2, h_split);
+        } else {
+            hb = ggml_mul_mat(ctx, w2b, ggml_mul_mat(ctx, w2a, h_split));
+        }
+
+        if (batch > 1) {
+            hb = ggml_reshape_3d(ctx, hb, vp, uq, batch);
+        }
+        ggml_tensor* hb_t = ggml_cont(ctx, ggml_transpose(ctx, hb));
+        hb_t              = ggml_reshape_3d(ctx, hb_t, uq, vp * merge_batch_vp, batch / merge_batch_vp);
+
+        ggml_tensor* hc_t;
+        if (w1 != nullptr) {
+            hc_t = ggml_mul_mat(ctx, w1, hb_t);
+        } else {
+            hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_t));
+        }
+
+        if (batch > 1) {
+            hc_t = ggml_reshape_3d(ctx, hc_t, up, vp, batch);
+        }
+
+        ggml_tensor* hc  = ggml_transpose(ctx, hc_t);
+        ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch);
+        return ggml_ext_scale(ctx, out, scale);
+    } else {
+        int batch = (int)h->ne[3];
+        // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
+        ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch);
+
+        if (w2 != nullptr) {
+            hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr,
+                                  conv_params.s0,
+                                  conv_params.s1,
+                                  conv_params.p0,
+                                  conv_params.p1,
+                                  conv_params.d0,
+                                  conv_params.d1,
+                                  conv_params.direct,
+                                  conv_params.circular_x,
+                                  conv_params.circular_y,
+                                  conv_params.scale);
+        } else {
+            // swap a and b order for conv lora
+            ggml_tensor* a = w2b;
+            ggml_tensor* b = w2a;
+
+            // unpack conv2d weights if needed
+            if (ggml_n_dims(a) < 4) {
+                int k = (int)sqrt(a->ne[0] / h_split->ne[2]);
+                GGML_ASSERT(k * k * h_split->ne[2] == a->ne[0]);
+                a = ggml_reshape_4d(ctx, a, k, k, a->ne[0] / (k * k), a->ne[1]);
+            } else if (a->ne[2] != h_split->ne[2]) {
+                int k = (int)sqrt(a->ne[2] / h_split->ne[2]);
+                GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]);
+                a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]);
+            }
+            ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr,
+                                               conv_params.s0,
+                                               conv_params.s1,
+                                               conv_params.p0,
+                                               conv_params.p1,
+                                               conv_params.d0,
+                                               conv_params.d1,
+                                               conv_params.direct,
+                                               conv_params.circular_x,
+                                               conv_params.circular_y,
+                                               conv_params.scale);
+
+            // not supporting lora_mid here
+            hb = ggml_ext_conv_2d(ctx,
+                                  ha,
+                                  b,
+                                  nullptr,
+                                  1,
+                                  1,
+                                  0,
+                                  0,
+                                  1,
+                                  1,
+                                  conv_params.direct,
+                                  conv_params.circular_x,
+                                  conv_params.circular_y,
+                                  conv_params.scale);
+        }
+
+        // Current hb shape: [W_out, H_out, vp, uq * batch]
+        int w_out = (int)hb->ne[0];
+        int h_out = (int)hb->ne[1];
+
+        // ggml_tensor* hb_cat = ggml_reshape_4d(ctx, hb, w_out , h_out , vp * uq, batch);
+        // [W_out, H_out, vp * uq,  batch]
+        // Now left to compute (W1 kr Id) * hb_cat == (W1 kr W2) cv h
+
+        // merge the uq groups of size vp*w_out*h_out
+        ggml_tensor* hb_merged = ggml_reshape_2d(ctx, hb, w_out * h_out * vp, uq * batch);
+        ggml_tensor* hc_t;
+        ggml_tensor* hb_merged_t = ggml_cont(ctx, ggml_transpose(ctx, hb_merged));
+        if (w1 != nullptr) {
+            // Would be great to be able to transpose w1 instead to avoid transposing both hb and hc
+            hc_t = ggml_mul_mat(ctx, w1, hb_merged_t);
+        } else {
+            hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_merged_t));
+        }
+        ggml_tensor* hc = ggml_transpose(ctx, hc_t);
+        // ungroup
+        ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch);
+        return ggml_ext_scale(ctx, out, scale);
+    }
+}

src/model/adapter/lora_ops.h

@@ -0,0 +1,25 @@
+#ifndef __SD_MODEL_ADAPTER_LORA_OPS_H__
+#define __SD_MODEL_ADAPTER_LORA_OPS_H__
+
+#include "core/ggml_runner.h"
+
+ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+                                 ggml_tensor* lora_down,
+                                 ggml_tensor* lora_up,
+                                 ggml_tensor* lora_mid = nullptr);
+
+ggml_tensor* ggml_ext_lokr_forward(
+    ggml_context* ctx,
+    ggml_backend_t backend,
+    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+    ggml_tensor* w1,   // Outer C (Full rank)
+    ggml_tensor* w1a,  // Outer A (Low rank part 1)
+    ggml_tensor* w1b,  // Outer B (Low rank part 2)
+    ggml_tensor* w2,   // Inner BA (Full rank)
+    ggml_tensor* w2a,  // Inner A (Low rank part 1)
+    ggml_tensor* w2b,  // Inner B (Low rank part 2)
+    bool is_conv,
+    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+    float scale);
+
+#endif  // __SD_MODEL_ADAPTER_LORA_OPS_H__

src/model/adapter/pmid.hpp

@@ -1,12 +1,13 @@
 #ifndef __SD_MODEL_ADAPTER_PMID_HPP__
 #define __SD_MODEL_ADAPTER_PMID_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
-#include "model/adapter/lora.hpp"
 #include "model/common/block.hpp"
 #include "model/te/clip.hpp"
-#include "model_loader.h"
 
 struct FuseBlock : public GGMLBlock {
     // network hparams
@@ -558,97 +559,7 @@ public:
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
-    }
-};
-
-struct PhotoMakerIDEmbed : public GGMLRunner {
-    std::map<std::string, ggml_tensor*> tensors;
-    std::string file_path;
-    std::shared_ptr<ModelManager> model_manager;
-    ggml_backend_t params_backend = nullptr;
-    bool load_failed              = false;
-    bool applied                  = false;
-
-    PhotoMakerIDEmbed(ggml_backend_t backend,
-                      ggml_backend_t params_backend_,
-                      std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>(),
-                      const std::string& file_path          = "",
-                      const std::string& prefix             = "")
-        : GGMLRunner(backend, manager), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
-        if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix)) {
-            load_failed = true;
-        }
-    }
-
-    std::string get_desc() {
-        return "id_embeds";
-    }
-
-    bool load_from_file(bool filter_tensor, int n_threads) {
-        LOG_INFO("loading PhotoMaker ID Embeds from '%s'", file_path.c_str());
-
-        if (load_failed) {
-            LOG_ERROR("init photomaker id embed from file failed: '%s'", file_path.c_str());
-            return false;
-        }
-
-        bool dry_run = true;
-        std::mutex tensor_mutex;
-        auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-            const std::string& name = tensor_storage.name;
-
-            if (filter_tensor && !contains(name, "pmid.id_embeds")) {
-                // LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
-                return true;
-            }
-            if (dry_run) {
-                std::lock_guard<std::mutex> lock(tensor_mutex);
-                ggml_tensor* real = ggml_new_tensor(params_ctx,
-                                                    tensor_storage.type,
-                                                    tensor_storage.n_dims,
-                                                    tensor_storage.ne);
-                tensors[name]     = real;
-            } else {
-                auto real   = tensors[name];
-                *dst_tensor = real;
-            }
-
-            return true;
-        };
-
-        model_manager->set_n_threads(n_threads);
-        ModelLoader& model_loader = model_manager->loader();
-        model_loader.load_tensors(on_new_tensor_cb);
-        if (!model_manager->register_param_tensors("PhotoMaker ID embeds",
-                                                   tensors,
-                                                   ModelManager::ResidencyMode::ParamBackend,
-                                                   runtime_backend,
-                                                   params_backend) ||
-            !model_manager->validate_registered_tensors()) {
-            LOG_ERROR("PhotoMaker ID embeds model manager registration failed");
-            return false;
-        }
-        std::vector<ggml_tensor*> id_embed_params;
-        id_embed_params.reserve(tensors.size());
-        for (const auto& pair : tensors) {
-            id_embed_params.push_back(pair.second);
-        }
-        if (!model_manager->prepare_params(id_embed_params)) {
-            LOG_ERROR("PhotoMaker ID embeds model manager prepare params failed");
-            return false;
-        }
-
-        LOG_VERBOSE("finished loading PhotoMaker ID Embeds ");
-        return true;
-    }
-
-    ggml_tensor* get() {
-        std::map<std::string, ggml_tensor*>::iterator pos;
-        pos = tensors.find("pmid.id_embeds");
-        if (pos != tensors.end())
-            return pos->second;
-        return nullptr;
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
     }
 };
 

src/model/adapter/pulid.hpp

@@ -1,8 +1,10 @@
 #ifndef __PULID_HPP__
 #define __PULID_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 
 class PuLIDPerceiverAttentionCA : public GGMLBlock {
 public:

src/model/common/block.hpp

@@ -1,9 +1,11 @@
 #ifndef __SD_MODEL_COMMON_BLOCK_HPP__
 #define __SD_MODEL_COMMON_BLOCK_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
 #include "ggml-backend.h"
+#include "model/common/ggml_block.hpp"
 
 class DownSampleBlock : public GGMLBlock {
 protected:

src/model/common/ggml_block.hpp

@@ -0,0 +1,880 @@
+#ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__
+#define __SD_MODEL_COMMON_GGML_BLOCK_HPP__
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "model.h"
+
+class GGMLBlock {
+protected:
+    typedef std::unordered_map<std::string, ggml_tensor*> ParameterMap;
+    typedef std::unordered_map<std::string, std::shared_ptr<GGMLBlock>> GGMLBlockMap;
+    GGMLBlockMap blocks;
+    ParameterMap params;
+
+    ggml_type get_type(const std::string& name, const String2TensorStorage& tensor_storage_map, ggml_type default_type) {
+        ggml_type wtype = default_type;
+        auto iter       = tensor_storage_map.find(name);
+        if (iter != tensor_storage_map.end()) {
+            const TensorStorage& tensor_storage = iter->second;
+            if (tensor_storage.expected_type != GGML_TYPE_COUNT) {
+                wtype = tensor_storage.expected_type;
+            } else {
+                wtype = tensor_storage.type;
+            }
+        }
+        return wtype;
+    }
+
+    void init_blocks(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+            block->init(ctx, tensor_storage_map, prefix + pair.first);
+        }
+    }
+
+    virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
+
+    virtual enum ggml_op param_usage_op(const std::string& name) const {
+        (void)name;
+        return GGML_OP_NONE;
+    }
+
+public:
+    void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        init_params(ctx, tensor_storage_map, prefix);
+        init_blocks(ctx, tensor_storage_map, prefix);
+    }
+
+    size_t get_params_num() {
+        size_t num_tensors = params.size();
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            num_tensors += block->get_params_num();
+        }
+        return num_tensors;
+    };
+
+    size_t get_params_mem_size() {
+        size_t mem_size = 0;
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            mem_size += block->get_params_mem_size();
+        }
+
+        for (auto& pair : params) {
+            mem_size += ggml_nbytes(pair.second);
+        }
+
+        return mem_size;
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+            block->get_param_tensors(tensors, prefix + pair.first);
+        }
+
+        for (auto& pair : params) {
+            ggml_tensor* param           = pair.second;
+            tensors[prefix + pair.first] = pair.second;
+            ggml_set_name(param, (prefix + pair.first).c_str());
+        }
+    }
+
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+        for (auto& pair : blocks) {
+            pair.second->get_param_tensor_ops(tensor_ops);
+        }
+        for (auto& pair : params) {
+            enum ggml_op op = param_usage_op(pair.first);
+            if (op != GGML_OP_NONE) {
+                tensor_ops[pair.second] = op;
+            }
+        }
+    }
+
+    virtual std::string get_desc() {
+        return "GGMLBlock";
+    }
+
+    void get_all_blocks(std::vector<GGMLBlock*>& result) {
+        result.push_back(this);
+        for (auto& block_iter : blocks) {
+            if (block_iter.second) {
+                block_iter.second->get_all_blocks(result);
+            }
+        }
+    }
+};
+
+class UnaryBlock : public GGMLBlock {
+public:
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) = 0;
+};
+
+class Identity : public UnaryBlock {
+public:
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        return x;
+    }
+};
+
+class Linear : public UnaryBlock {
+protected:
+    int64_t in_features;
+    int64_t out_features;
+    bool bias;
+    bool force_f32;
+    bool force_prec_f32;
+    bool has_weight_scale       = false;
+    bool int8_convrot           = false;
+    int int8_convrot_group_size = 0;
+    float scale;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix            = prefix;
+        has_weight_scale        = false;
+        int8_convrot            = false;
+        int8_convrot_group_size = 0;
+        enum ggml_type wtype    = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
+        if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_features);
+        }
+        auto weight_storage           = tensor_storage_map.find(prefix + "weight");
+        const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
+        auto weight_scale_storage     = tensor_storage_map.find(prefix + "weight_scale");
+        if (weight_scale_storage != tensor_storage_map.end()) {
+            const int64_t scale_nelements = weight_scale_storage->second.nelements();
+            GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features);
+            params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
+            has_weight_scale       = true;
+        }
+        if (is_int8_tensorwise) {
+            GGML_ASSERT(wtype == GGML_TYPE_I8);
+            GGML_ASSERT(has_weight_scale);
+            int8_convrot            = weight_storage->second.int8_convrot;
+            int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
+        }
+    }
+
+public:
+    Linear(int64_t in_features,
+           int64_t out_features,
+           bool bias           = true,
+           bool force_f32      = false,
+           bool force_prec_f32 = false,
+           float scale         = 1.f)
+        : in_features(in_features),
+          out_features(out_features),
+          bias(bias),
+          force_f32(force_f32),
+          force_prec_f32(force_prec_f32),
+          scale(scale) {}
+
+    void set_scale(float scale_) {
+        scale = scale_;
+    }
+
+    void set_force_prec_f32(bool force_prec_f32_) {
+        force_prec_f32 = force_prec_f32_;
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w            = params["weight"];
+        ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
+        if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
+            bool supports_fp8_matmul = false;
+            if (ctx->backend != nullptr) {
+                ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x);
+                if (force_prec_f32) {
+                    ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32);
+                }
+                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul);
+            }
+            if (!supports_fp8_matmul) {
+                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
+            }
+        }
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+        ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
+        ggml_tensor* out         = nullptr;
+        if (w->type == GGML_TYPE_I8) {
+            if (x->type != GGML_TYPE_F32) {
+                x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x);
+            }
+            if (!ggml_is_contiguous(x)) {
+                x = ggml_cont(ctx->ggml_ctx, x);
+            }
+            ggml_tensor* lora_input = x;
+            if (ctx->weight_adapter && b != nullptr) {
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+            if (int8_convrot && scale == 1.f) {
+                const auto cache_key = std::make_pair(x, int8_convrot_group_size);
+                auto cached          = ctx->int8_convrot_cache.find(cache_key);
+                if (cached == ctx->int8_convrot_cache.end()) {
+                    x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size);
+                    ctx->int8_convrot_cache.emplace(cache_key, x);
+                } else {
+                    x = cached->second;
+                }
+            }
+            out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+                                                x,
+                                                w,
+                                                weight_scale,
+                                                b,
+                                                int8_convrot ? int8_convrot_group_size : 0,
+                                                scale);
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+                forward_params.linear.force_prec_f32 = force_prec_f32;
+                forward_params.linear.scale          = scale;
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
+                                                                                               ctx->backend,
+                                                                                               lora_input,
+                                                                                               w,
+                                                                                               out,
+                                                                                               prefix,
+                                                                                               forward_params);
+            }
+            return out;
+        }
+        if (has_weight_scale) {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale);
+            out = ggml_mul(ctx->ggml_ctx, out, weight_scale);
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+                forward_params.linear.force_prec_f32 = force_prec_f32;
+                forward_params.linear.scale          = scale;
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
+                                                                                               ctx->backend,
+                                                                                               x,
+                                                                                               w,
+                                                                                               out,
+                                                                                               prefix,
+                                                                                               forward_params);
+                if (b != nullptr) {
+                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                }
+            }
+            if (b != nullptr) {
+                out = ggml_add_inplace(ctx->ggml_ctx, out, b);
+            }
+            return out;
+        }
+        if (ctx->weight_adapter) {
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+            forward_params.linear.force_prec_f32 = force_prec_f32;
+            forward_params.linear.scale          = scale;
+            out                                  = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params);
+        } else {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
+        }
+        return out;
+    }
+};
+
+__STATIC_INLINE__ bool support_get_rows(ggml_type wtype) {
+    switch (wtype) {
+        case GGML_TYPE_F16:
+        case GGML_TYPE_BF16:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+            return true;
+        default:
+            return false;
+    }
+}
+
+class Embedding : public UnaryBlock {
+protected:
+    int64_t embedding_dim;
+    int64_t num_embeddings;
+    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 (!support_get_rows(wtype)) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
+    }
+
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
+public:
+    Embedding(int64_t num_embeddings, int64_t embedding_dim)
+        : embedding_dim(embedding_dim),
+          num_embeddings(num_embeddings) {
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* input_ids) override {
+        // input_ids: [N, n_token]
+        auto weight = params["weight"];
+
+        // There are issues with ggml batch inference, so we are expanding it here first.
+        // TODO: fix ggml batch inference
+        int64_t n = input_ids->ne[1];
+        input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]);
+
+        input_ids      = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
+        auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
+        embedding      = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n);
+
+        // [N, n_token, embedding_dim]
+        return embedding;
+    }
+};
+
+class Conv2d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d(int64_t in_channels,
+           int64_t out_channels,
+           std::pair<int, int> kernel_size,
+           std::pair<int, int> stride   = {1, 1},
+           std::pair<int, int> padding  = {0, 0},
+           std::pair<int, int> dilation = {1, 1},
+           bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() override {
+        return "Conv2d";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+        if (ctx->weight_adapter) {
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+            forward_params.conv2d.s0         = stride.second;
+            forward_params.conv2d.s1         = stride.first;
+            forward_params.conv2d.p0         = padding.second;
+            forward_params.conv2d.p1         = padding.first;
+            forward_params.conv2d.d0         = dilation.second;
+            forward_params.conv2d.d1         = dilation.first;
+            forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+            forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+            forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+            forward_params.conv2d.scale      = scale;
+            return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+        }
+        return ggml_ext_conv_2d(ctx->ggml_ctx,
+                                x,
+                                w,
+                                b,
+                                stride.second,
+                                stride.first,
+                                padding.second,
+                                padding.first,
+                                dilation.second,
+                                dilation.first,
+                                ctx->conv2d_direct_enabled,
+                                ctx->circular_x_enabled,
+                                ctx->circular_y_enabled,
+                                scale);
+    }
+};
+
+class Conv2d_grouped : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    int groups;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d_grouped(int64_t in_channels,
+                   int64_t out_channels,
+                   int groups,
+                   std::pair<int, int> kernel_size,
+                   std::pair<int, int> stride   = {1, 1},
+                   std::pair<int, int> padding  = {0, 0},
+                   std::pair<int, int> dilation = {1, 1},
+                   bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          groups(groups),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() override {
+        return "Conv2d_grouped";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+
+        if (groups == 1) {
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+            }
+            return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
+                                    stride.second, stride.first,
+                                    padding.second, padding.first,
+                                    dilation.second, dilation.first,
+                                    ctx->conv2d_direct_enabled,
+                                    ctx->circular_x_enabled,
+                                    ctx->circular_y_enabled,
+                                    scale);
+        }
+
+        if (groups == in_channels && groups == out_channels) {
+            ggml_tensor* res;
+            if (ctx->conv2d_direct_enabled) {
+                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
+                                             stride.second, stride.first,
+                                             padding.second, padding.first,
+                                             dilation.second, dilation.first);
+            } else {
+                res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
+                                      stride.second, stride.first,
+                                      padding.second, padding.first,
+                                      dilation.second, dilation.first);
+            }
+            if (b) {
+                b   = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
+                res = ggml_add_inplace(ctx->ggml_ctx, res, b);
+            }
+            return res;
+        }
+
+        int64_t ic_g = in_channels / groups;
+        int64_t oc_g = out_channels / groups;
+
+        std::vector<ggml_tensor*> out_slices(groups);
+
+        for (int i = 0; i < groups; ++i) {
+            size_t x_offset  = i * ic_g * x->nb[2];
+            ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x,
+                                            x->ne[0], x->ne[1], ic_g, x->ne[3],
+                                            x->nb[1], x->nb[2], x->nb[3],
+                                            x_offset);
+
+            size_t w_offset  = i * oc_g * w->nb[3];
+            ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w,
+                                            w->ne[0], w->ne[1], w->ne[2], oc_g,
+                                            w->nb[1], w->nb[2], w->nb[3],
+                                            w_offset);
+
+            ggml_tensor* b_i = nullptr;
+            if (b) {
+                size_t b_offset = i * oc_g * b->nb[0];
+                b_i             = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset);
+            }
+
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                out_slices[i]                    = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params);
+            } else {
+                out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i,
+                                                 stride.second, stride.first,
+                                                 padding.second, padding.first,
+                                                 dilation.second, dilation.first,
+                                                 ctx->conv2d_direct_enabled,
+                                                 ctx->circular_x_enabled,
+                                                 ctx->circular_y_enabled,
+                                                 scale);
+            }
+        }
+
+        ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
+
+        return out;
+    }
+};
+
+class Conv3d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    std::tuple<int, int, int> kernel_size;
+    std::tuple<int, int, int> stride;
+    std::tuple<int, int, int> padding;
+    std::tuple<int, int, int> dilation;
+    bool bias;
+    bool force_prec_f32;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx,
+                                                  wtype,
+                                                  std::get<2>(kernel_size),
+                                                  std::get<1>(kernel_size),
+                                                  std::get<0>(kernel_size),
+                                                  in_channels * out_channels);
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+public:
+    Conv3d(int64_t in_channels,
+           int64_t out_channels,
+           std::tuple<int, int, int> kernel_size,
+           std::tuple<int, int, int> stride   = {1, 1, 1},
+           std::tuple<int, int, int> padding  = {0, 0, 0},
+           std::tuple<int, int, int> dilation = {1, 1, 1},
+           bool bias                          = true,
+           bool force_prec_f32                = false)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias),
+          force_prec_f32(force_prec_f32) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (ctx->weight_adapter) {
+            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            if (w->type != GGML_TYPE_F16) {
+                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16);
+            }
+        }
+        if (bias) {
+            b = params["bias"];
+            if (ctx->weight_adapter) {
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+        }
+        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
+                                std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
+                                std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
+                                std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
+                                force_prec_f32);
+    }
+};
+
+class LayerNorm : public UnaryBlock {
+protected:
+    int64_t normalized_shape;
+    float eps;
+    bool elementwise_affine;
+    bool bias;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix = prefix;
+        if (elementwise_affine) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["weight"]     = ggml_new_tensor_1d(ctx, wtype, normalized_shape);
+            if (bias) {
+                enum ggml_type wtype = GGML_TYPE_F32;
+                params["bias"]       = ggml_new_tensor_1d(ctx, wtype, normalized_shape);
+            }
+        }
+    }
+
+public:
+    LayerNorm(int64_t normalized_shape,
+              float eps               = 1e-05f,
+              bool elementwise_affine = true,
+              bool bias               = true)
+        : normalized_shape(normalized_shape),
+          eps(eps),
+          elementwise_affine(elementwise_affine),
+          bias(bias) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = nullptr;
+        ggml_tensor* b = nullptr;
+
+        if (elementwise_affine) {
+            w = params["weight"];
+            if (ctx->weight_adapter) {
+                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            }
+            if (bias) {
+                b = params["bias"];
+                if (ctx->weight_adapter) {
+                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                }
+            }
+        }
+        return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps);
+    }
+};
+
+class GroupNorm : public GGMLBlock {
+protected:
+    int num_groups;
+    int64_t num_channels;
+    float eps;
+    bool affine;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix = prefix;
+        if (affine) {
+            enum ggml_type wtype      = GGML_TYPE_F32;
+            enum ggml_type bias_wtype = GGML_TYPE_F32;
+            params["weight"]          = ggml_new_tensor_1d(ctx, wtype, num_channels);
+            params["bias"]            = ggml_new_tensor_1d(ctx, bias_wtype, num_channels);
+        }
+    }
+
+public:
+    GroupNorm(int num_groups,
+              int64_t num_channels,
+              float eps   = 1e-05f,
+              bool affine = true)
+        : num_groups(num_groups),
+          num_channels(num_channels),
+          eps(eps),
+          affine(affine) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        ggml_tensor* w = nullptr;
+        ggml_tensor* b = nullptr;
+        if (affine) {
+            w = params["weight"];
+            b = params["bias"];
+            if (ctx->weight_adapter) {
+                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+        }
+        return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups);
+    }
+};
+
+class GroupNorm32 : public GroupNorm {
+public:
+    GroupNorm32(int64_t num_channels)
+        : GroupNorm(32, num_channels, 1e-06f) {}
+};
+
+class RMSNorm : public UnaryBlock {
+protected:
+    int64_t hidden_size;
+    float eps;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F32;
+        params["weight"]     = ggml_new_tensor_1d(ctx, wtype, hidden_size);
+    }
+
+public:
+    RMSNorm(int64_t hidden_size,
+            float eps = 1e-06f)
+        : hidden_size(hidden_size),
+          eps(eps) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        if (ctx->weight_adapter) {
+            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+        }
+        x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+        x = ggml_mul_inplace(ctx->ggml_ctx, x, w);
+        return x;
+    }
+};
+
+class MultiheadAttention : public GGMLBlock {
+protected:
+    int64_t embed_dim;
+    int64_t n_head;
+    bool proj_in;
+    std::string q_proj_name;
+    std::string k_proj_name;
+    std::string v_proj_name;
+    std::string in_proj_name;
+    std::string out_proj_name;
+
+public:
+    MultiheadAttention(int64_t embed_dim,
+                       int64_t n_head,
+                       bool qkv_proj_bias        = true,
+                       bool out_proj_bias        = true,
+                       bool proj_in              = false,
+                       std::string q_proj_name   = "q_proj",
+                       std::string k_proj_name   = "k_proj",
+                       std::string v_proj_name   = "v_proj",
+                       std::string in_proj_name  = "in_proj",
+                       std::string out_proj_name = "out_proj")
+        : embed_dim(embed_dim),
+          n_head(n_head),
+          proj_in(proj_in),
+          q_proj_name(q_proj_name),
+          k_proj_name(k_proj_name),
+          v_proj_name(v_proj_name),
+          in_proj_name(in_proj_name),
+          out_proj_name(out_proj_name) {
+        if (proj_in) {
+            blocks[in_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim * 3, qkv_proj_bias));
+        } else {
+            blocks[q_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+            blocks[k_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+            blocks[v_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+        }
+        blocks[out_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, out_proj_bias));
+    }
+
+    // x: [N, n_token, embed_dim]
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* mask = nullptr) {
+        auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[out_proj_name]);
+
+        ggml_tensor* q;
+        ggml_tensor* k;
+        ggml_tensor* v;
+        if (proj_in) {
+            auto in_proj = std::dynamic_pointer_cast<Linear>(blocks[in_proj_name]);
+            auto qkv     = in_proj->forward(ctx, x);
+            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            q            = qkv_vec[0];
+            k            = qkv_vec[1];
+            v            = qkv_vec[2];
+        } else {
+            auto q_proj = std::dynamic_pointer_cast<Linear>(blocks[q_proj_name]);
+            auto k_proj = std::dynamic_pointer_cast<Linear>(blocks[k_proj_name]);
+            auto v_proj = std::dynamic_pointer_cast<Linear>(blocks[v_proj_name]);
+
+            q = q_proj->forward(ctx, x);
+            k = k_proj->forward(ctx, x);
+            v = v_proj->forward(ctx, x);
+        }
+
+        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false);  // [N, n_token, embed_dim]
+
+        x = out_proj->forward(ctx, x);  // [N, n_token, embed_dim]
+        return x;
+    }
+};
+
+#endif  // __SD_MODEL_COMMON_GGML_BLOCK_HPP__

src/model/common/rope.hpp

@@ -2,9 +2,13 @@
 #define __SD_MODEL_COMMON_ROPE_HPP__
 
 #include <algorithm>
+#include <cassert>
 #include <cmath>
+#include <set>
 #include <vector>
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 
 namespace Rope {
     enum class EmbedNDLayout {

src/model/detector/yolov8.h

@@ -8,8 +8,9 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
 struct YOLOv8Config {
     std::array<int, 23> out_channels{};
@@ -355,7 +356,7 @@ struct YOLOv8Runner : public GGMLRunner {
 
     sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
         auto get_graph = [&]() { return build_graph(input); };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false));
     }
 };
 

src/model/diffusion/anima.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_ANIMA_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <memory>
 #include <utility>
@@ -717,7 +718,7 @@ namespace Anima {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/animatediff.hpp

@@ -1,8 +1,10 @@
 #ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
 #define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 
 // AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules.
 namespace AnimateDiff {

src/model/diffusion/boogu.hpp

@@ -2,11 +2,15 @@
 #define __SD_MODEL_DIFFUSION_BOOGU_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <tuple>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
@@ -815,7 +819,7 @@ namespace Boogu {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/control.hpp

@@ -2,8 +2,6 @@
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"
-#include "model_loader.h"
-#include "model_manager.h"
 
 // Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
 #define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
@@ -317,20 +315,17 @@ struct ControlNet : public GGMLRunner {
     ggml_tensor* guided_hint_output_ggml = nullptr;
     std::vector<sd::Tensor<float>> controls;
     bool guided_hint_cached = false;
-    std::shared_ptr<ModelManager> owned_model_manager;
-    ggml_backend_t params_backend = nullptr;
 
     static const char* guided_hint_cache_name() {
         return "controlnet.guided_hint";
     }
 
     ControlNet(ggml_backend_t backend,
-               ggml_backend_t params_backend_,
                const String2TensorStorage& tensor_storage_map      = {},
                SDVersion version                                   = VERSION_SD1,
                const std::string& prefix                           = "",
                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix), params_backend(params_backend_) {
+        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix) {
         control_net.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -435,7 +430,7 @@ struct ControlNet : public GGMLRunner {
             }
             return true;
         };
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs);
         control_outputs_ggml.clear();
         guided_hint_output_ggml = nullptr;
         if (!compute_result.has_value()) {
@@ -445,39 +440,6 @@ struct ControlNet : public GGMLRunner {
         guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
         return controls;
     }
-
-    bool load_from_file(const std::string& file_path, int n_threads) {
-        LOG_INFO("loading control net from '%s'", file_path.c_str());
-        std::map<std::string, ggml_tensor*> tensors;
-        control_net.get_param_tensors(tensors);
-
-        auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
-        if (manager == nullptr) {
-            owned_model_manager = std::make_shared<ModelManager>();
-            residency_manager   = owned_model_manager;
-            manager             = owned_model_manager;
-        }
-
-        ModelLoader& model_loader = manager->loader();
-        if (!model_loader.init_from_file_and_convert_name(file_path)) {
-            LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str());
-            return false;
-        }
-
-        manager->set_n_threads(n_threads);
-        if (!manager->register_param_tensors("ControlNet",
-                                             std::move(tensors),
-                                             ModelManager::ResidencyMode::ParamBackend,
-                                             runtime_backend,
-                                             params_backend) ||
-            !manager->validate_registered_tensors()) {
-            LOG_ERROR("register control net tensors with model manager failed");
-            return false;
-        }
-
-        LOG_INFO("control net model loaded");
-        return true;
-    }
 };
 
 #endif  // __SD_MODEL_DIFFUSION_CONTROL_HPP__

src/model/diffusion/dit.hpp

@@ -1,7 +1,8 @@
 #ifndef __SD_MODEL_DIFFUSION_DIT_HPP__
 #define __SD_MODEL_DIFFUSION_DIT_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 
 namespace DiT {
     inline ggml_tensor* patchify(ggml_context* ctx,

src/model/diffusion/ernie_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
 
@@ -440,7 +441,7 @@ namespace ErnieImage {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/flux.hpp

@@ -1,8 +1,11 @@
 #ifndef __SD_MODEL_DIFFUSION_FLUX_HPP__
 #define __SD_MODEL_DIFFUSION_FLUX_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "core/util.h"
 #include "model/adapter/pulid.hpp"
@@ -1626,7 +1629,7 @@ namespace Flux {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
             };
 
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return result;
         }
 
@@ -1711,8 +1714,8 @@ namespace Flux {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -1733,7 +1736,8 @@ namespace Flux {
                                                                             VERSION_FLUX2,
                                                                             model_manager);
 
-            if (!model_manager->register_runner_params("Flux test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *flux,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/hidream_o1.hpp

@@ -329,7 +329,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
+            auto output = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
@@ -457,7 +457,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/hunyuan.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
 #define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
 
+#include <cinttypes>
 #include <memory>
 
 #include "model/common/block.hpp"
@@ -654,7 +655,7 @@ namespace Hunyuan {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ideogram4.hpp

@@ -2,14 +2,18 @@
 #define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/model.hpp"
 
@@ -537,7 +541,7 @@ namespace Ideogram4 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, use_uncond_model);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/krea2.hpp

@@ -12,8 +12,11 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
@@ -775,7 +778,7 @@ namespace Krea2 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, ref_image_params);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lens.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_LENS_HPP__
 #define __SD_MODEL_DIFFUSION_LENS_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
 
@@ -408,7 +409,7 @@ namespace Lens {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lingbot_video.hpp

@@ -674,7 +674,7 @@ namespace LingBotVideo {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ltxv.hpp

@@ -2,12 +2,15 @@
 #define __SD_MODEL_DIFFUSION_LTXV_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <memory>
 #include <string>
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -1998,7 +2001,7 @@ namespace LTXV {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
             };
-            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return out;
         }
 
@@ -2084,8 +2087,8 @@ namespace LTXV {
             ggml_backend_t backend = sd_backend_cpu_init();
             LOG_INFO("loading ltxav from '%s'", model_path.c_str());
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(model_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -2104,7 +2107,8 @@ namespace LTXV {
                                                                                "model.diffusion_model",
                                                                                model_manager);
 
-            if (!model_manager->register_runner_params("LTXAV test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *ltxav,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/mage_flow.hpp

@@ -142,7 +142,7 @@ namespace MageFlow {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/minimax_h3.hpp

@@ -2,12 +2,14 @@
 #define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <set>
 #include <string>
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 
 #include "core/ggml_graph_cut.h"
 #include "model/diffusion/dit.hpp"
@@ -1166,9 +1168,9 @@ namespace MiniMaxH3 {
                                    extra->video_sigma_shift,
                                    extra->audio_sigma_shift);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
-                                                                              n_threads,
-                                                                              false),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
+                                                                       n_threads,
+                                                                       false),
                                                    params.x->dim());
         }
     };

src/model/diffusion/minit2i.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
@@ -9,7 +10,10 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
@@ -589,7 +593,7 @@ namespace MiniT2I {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, mask);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/mmdit.hpp

@@ -2,12 +2,18 @@
 #define __SD_MODEL_DIFFUSION_MMDIT_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/model.hpp"
 #include "model_loader.h"
 
@@ -987,7 +993,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,
@@ -1058,13 +1064,14 @@ struct MMDiTRunner : public DiffusionModelRunner {
         {
             LOG_INFO("loading from '%s'", file_path.c_str());
 
-            ModelLoader& model_loader = model_manager->loader();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
             }
 
-            if (!model_manager->register_runner_params("MMDiT test",
+            if (!model_manager->set_loader(std::move(model_loader)) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *mmdit,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/model.hpp

@@ -5,7 +5,7 @@
 #include <utility>
 #include <variant>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
 #include "core/tensor_ggml.hpp"
 #include "model/common/rope.hpp"
 #include "model_manager.h"

src/model/diffusion/pid.hpp

@@ -1,13 +1,18 @@
 #ifndef __SD_MODEL_DIFFUSION_PID_HPP__
 #define __SD_MODEL_DIFFUSION_PID_HPP__
 
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/mmdit.hpp"
@@ -938,7 +943,7 @@ namespace Pid {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/qwen_image.hpp

@@ -2,6 +2,8 @@
 #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
 
 #include <memory>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "core/util.h"
 #include "model/common/block.hpp"
@@ -707,7 +709,7 @@ namespace Qwen {
                 return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -771,8 +773,8 @@ namespace Qwen {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -791,7 +793,8 @@ namespace Qwen {
                                                                                             VERSION_QWEN_IMAGE,
                                                                                             model_manager);
 
-            if (!model_manager->register_runner_params("Qwen image test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *qwen_image,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/sefi_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
 
+#include <cinttypes>
 #include <memory>
 
 #include "model/common/block.hpp"

src/model/diffusion/unet.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 
 #include "model.h"
 #include "model/common/block.hpp"
@@ -835,7 +836,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,

src/model/diffusion/wan.hpp

@@ -1,9 +1,12 @@
 #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
 #define __SD_MODEL_DIFFUSION_WAN_HPP__
 
+#include <cinttypes>
 #include <map>
 #include <memory>
 #include <utility>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -950,7 +953,7 @@ namespace WAN {
                 return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -1017,8 +1020,8 @@ namespace WAN {
             ggml_type model_data_type = GGML_TYPE_F16;
             LOG_INFO("loading from '%s'", file_path.c_str());
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -1037,7 +1040,8 @@ namespace WAN {
                                                                          VERSION_WAN2_2_TI2V,
                                                                          model_manager);
 
-            if (!model_manager->register_runner_params("Wan test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *wan,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/z_image.hpp

@@ -2,8 +2,14 @@
 #define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
@@ -636,7 +642,7 @@ namespace ZImage {
                 return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -700,8 +706,8 @@ namespace ZImage {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -722,7 +728,8 @@ namespace ZImage {
                                                                                    VERSION_QWEN_IMAGE,
                                                                                    model_manager);
 
-            if (!model_manager->register_runner_params("ZImage test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *z_image,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/te/clip.hpp

@@ -1,8 +1,11 @@
 #ifndef __SD_MODEL_TE_CLIP_HPP__
 #define __SD_MODEL_TE_CLIP_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 #include "model.h"
+#include "model/common/ggml_block.hpp"
 #include "tokenizers/clip_tokenizer.h"
 
 /*================================================ FrozenCLIPEmbedder ================================================*/
@@ -142,6 +145,10 @@ protected:
         params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions);
     }
 
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "token_embedding.weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
 public:
     CLIPEmbeddings(int64_t embed_dim,
                    int64_t vocab_size    = 49408,
@@ -572,7 +579,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
+        auto result = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
         if (return_pooled) {
             return take_or_empty(std::move(result));
         }

src/model/te/llm.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <array>
+#include <cinttypes>
 #include <cmath>
 #include <fstream>
 #include <functional>
@@ -18,8 +19,13 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
 #include "json.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
@@ -2091,7 +2097,7 @@ namespace LLM {
                                    out_layers,
                                    return_all_hidden_states);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end),
                                                    input_ids.dim() + 1);
         }
 
@@ -2175,7 +2181,7 @@ namespace LLM {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
         }
 
         ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
@@ -2287,7 +2293,7 @@ namespace LLM {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_outputs_graph(image);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }
@@ -2313,7 +2319,7 @@ namespace LLM {
             auto get_graph    = [&]() -> ggml_cgraph* {
                 return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }
@@ -2570,8 +2576,8 @@ namespace LLM {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "text_encoders.llm.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -2595,7 +2601,8 @@ namespace LLM {
                                                                              true,
                                                                              model_manager);
 
-            if (!model_manager->register_runner_params("LLM test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Conditioner,
                                                        *llm,
                                                        "text_encoders.llm",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/te/t5.hpp

@@ -10,7 +10,12 @@
 #include <string>
 #include <unordered_map>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
 #include "tokenizers/t5_unigram_tokenizer.h"
@@ -455,7 +460,7 @@ struct T5Runner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end), 3);
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), 3);
     }
 
     static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
@@ -626,8 +631,8 @@ struct T5Embedder {
         ggml_backend_t backend    = sd_backend_cpu_init();
         ggml_type model_data_type = GGML_TYPE_F16;
 
-        auto model_manager        = std::make_shared<ModelManager>();
-        ModelLoader& model_loader = model_manager->loader();
+        auto model_manager = std::make_shared<ModelManager>();
+        ModelLoader model_loader;
         if (!model_loader.init_from_file_and_convert_name(file_path)) {
             LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
             return;
@@ -642,7 +647,8 @@ struct T5Embedder {
 
         std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, tensor_storage_map, "", true, model_manager);
 
-        if (!model_manager->register_runner_params("T5 test",
+        if (!model_manager->set_loader(model_loader) ||
+            !model_manager->register_runner_params(ModelComponent::Conditioner,
                                                    *t5,
                                                    "",
                                                    ModelManager::ResidencyMode::ParamBackend,

src/model/upscaler/esrgan.hpp

@@ -7,8 +7,10 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
 /*
     ===================================    ESRGAN  ===================================
@@ -265,7 +267,7 @@ struct ESRGAN : public GGMLRunner {
     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());
+        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         return result;
     }
 };

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -11,9 +11,11 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model_loader.h"
 
@@ -431,12 +433,15 @@ namespace LTXVUpsampler {
     struct LatentUpsamplerRunner : public GGMLRunner {
         LatentUpsamplerConfig config;
         std::unique_ptr<LatentUpsampler> model;
+        std::string weight_prefix;
 
         LatentUpsamplerRunner(ggml_backend_t backend,
                               const String2TensorStorage& tensor_storage_map,
+                              const std::string& prefix                           = "",
                               std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : GGMLRunner(backend, weight_manager),
-              config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map)) {
+              config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map, prefix)),
+              weight_prefix(prefix) {
             if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
                 config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
                 LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
@@ -450,7 +455,7 @@ namespace LTXVUpsampler {
             }
 
             model = std::make_unique<LatentUpsampler>(config);
-            model->init(params_ctx, tensor_storage_map, "");
+            model->init(params_ctx, tensor_storage_map, prefix);
         }
 
         std::string get_desc() override {
@@ -459,7 +464,7 @@ namespace LTXVUpsampler {
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
             if (model) {
-                model->get_param_tensors(tensors);
+                model->get_param_tensors(tensors, weight_prefix);
             }
         }
 
@@ -499,7 +504,7 @@ namespace LTXVUpsampler {
             }
             size_t expected_dim = static_cast<size_t>(x.dim());
             auto get_graph      = [&]() -> ggml_cgraph* { return build_graph(x); };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim);
         }
     };
 

src/model/vae/audio_vae.hpp

@@ -1,7 +1,8 @@
 #ifndef __SD_MODEL_VAE_AUDIO_VAE_HPP__
 #define __SD_MODEL_VAE_AUDIO_VAE_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 
 struct AudioVAERunner : public GGMLRunner {
     AudioVAERunner(ggml_backend_t backend,

src/model/vae/auto_encoder_kl.hpp

@@ -1,6 +1,8 @@
 #ifndef __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
 #define __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
 
+#include <cinttypes>
+#include "core/ggml_tensor_utils.h"
 #include "model/vae/vae.hpp"
 
 /*================================================== AutoEncoderKL ===================================================*/
@@ -744,7 +746,7 @@ struct AutoEncoderKL : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z.dim());
     }
 
     sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {

src/model/vae/hunyuan_vae.hpp

@@ -825,9 +825,9 @@ namespace Hunyuan {
             auto get_graph          = [&]() -> ggml_cgraph* {
                 return build_graph(graph_input, decode_graph);
             };
-            auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
-                                                                                     n_threads,
-                                                                                     false),
+            auto output = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
+                                                                              n_threads,
+                                                                              false),
                                                           graph_input.dim());
             if (!output.empty() && input.dim() == 4) {
                 output.squeeze_(2);

src/model/vae/ltx_audio_vae.hpp

@@ -7,7 +7,12 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/vae/audio_vae.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
@@ -1042,7 +1047,7 @@ namespace LTXV {
                 ggml_build_forward_expand(gf, waveform);
                 return gf;
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4);
             int64_t t1  = ggml_time_ms();
             LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
             return result;
@@ -1073,8 +1078,8 @@ namespace LTXV {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file(model_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -1086,7 +1091,8 @@ namespace LTXV {
                                                                      prefix,
                                                                      model_manager);
 
-            if (!model_manager->register_runner_params("LTX audio VAE test",
+            if (!model_manager->set_loader(std::move(model_loader)) ||
+                !model_manager->register_runner_params(ModelComponent::AudioVAE,
                                                        *ltx_audio_vae,
                                                        ModelManager::ResidencyMode::ParamBackend,
                                                        backend,

src/model/vae/ltx_vae.hpp

@@ -8,6 +8,8 @@
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/diffusion/ltxv.hpp"
 #include "model/vae/vae.hpp"
@@ -1348,7 +1350,7 @@ struct LTXVideoVAE : public VAE {
                                                  static_cast<int>(tile.start),
                                                  tile.overlap);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
                                                    expected_dim);
         });
 
@@ -1405,7 +1407,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input, decode_graph);
         };
-        auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
+        auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim);
         if (result.empty()) {
             return {};
         }
@@ -1418,7 +1420,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_latent_statistics_graph(z, normalize);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
                                                static_cast<size_t>(z.dim()));
     }
 
@@ -1474,8 +1476,8 @@ struct LTXVideoVAE : public VAE {
         ggml_backend_t backend = sd_backend_cpu_init();
         LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
 
-        auto model_manager        = std::make_shared<ModelManager>();
-        ModelLoader& model_loader = model_manager->loader();
+        auto model_manager = std::make_shared<ModelManager>();
+        ModelLoader model_loader;
         if (!model_loader.init_from_file_and_convert_name(model_path, "vae.")) {
             LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
             return;
@@ -1489,7 +1491,8 @@ struct LTXVideoVAE : public VAE {
                                                                          VERSION_LTXAV,
                                                                          model_manager);
 
-        if (!model_manager->register_runner_params("LTX VAE test",
+        if (!model_manager->set_loader(model_loader) ||
+            !model_manager->register_runner_params(ModelComponent::VAE,
                                                    *vae,
                                                    ModelManager::ResidencyMode::ParamBackend,
                                                    backend,

src/model/vae/mage_vae.hpp

@@ -490,7 +490,7 @@ namespace MageVAE {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input, decode_graph);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), input.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), input.dim());
         }
 
         int get_encoder_output_channels(int input_channels) override {

src/model/vae/minimax_h3_audio_vae.hpp

@@ -480,7 +480,7 @@ namespace MiniMaxH3 {
                 return graph;
             };
             auto result = restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph, n_threads, false),
+                GGMLRunner::compute(get_graph, n_threads, false),
                 4);
             int64_t t1 = ggml_time_ms();
             LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
@@ -500,7 +500,7 @@ namespace MiniMaxH3 {
                 return graph;
             };
             auto result = restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph, n_threads, false),
+                GGMLRunner::compute(get_graph, n_threads, false),
                 4);
             int64_t t1 = ggml_time_ms();
             LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",

src/model/vae/minimax_h3_vae.hpp

@@ -791,9 +791,9 @@ namespace MiniMaxH3VAE {
                 return graph;
             };
             return restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph,
-                                           n_threads,
-                                           false),
+                GGMLRunner::compute(get_graph,
+                                    n_threads,
+                                    false),
                 5);
         }
     };

src/model/vae/tae.hpp

@@ -1,8 +1,12 @@
 #ifndef __SD_MODEL_VAE_TAE_HPP__
 #define __SD_MODEL_VAE_TAE_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/rng.hpp"
+#include "core/util.h"
 #include "model.h"
+#include "model/common/ggml_block.hpp"
 
 /*
     ===================================    TinyAutoEncoder  ===================================
@@ -787,7 +791,7 @@ struct TinyImageAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 
@@ -872,7 +876,7 @@ struct TinyVideoAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 

src/model/vae/vae.hpp

@@ -5,6 +5,7 @@
 #include "model/common/block.hpp"
 #include "model/vae/vae_tiling.hpp"
 #include "model_manager.h"
+#include "runtime/tiling.h"
 
 struct VAE : public GGMLRunner {
 protected:

src/model/vae/wan_vae.hpp

@@ -4,6 +4,8 @@
 #include <map>
 #include <memory>
 #include <utility>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/vae/vae.hpp"
@@ -1427,7 +1429,7 @@ namespace WAN {
                     return build_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
                 };
                 return restore_trailing_singleton_dims(
-                    GGMLRunner::compute<float>(get_graph, n_threads, false),
+                    GGMLRunner::compute(get_graph, n_threads, false),
                     static_cast<size_t>(input.dim()));
             });
 
@@ -1446,7 +1448,7 @@ namespace WAN {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input.empty() ? z : input, decode_graph);
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
                                                           input.empty() ? z.dim() : input.dim());
             if (!result.empty() && z.dim() == 4) {
                 result.squeeze_(2);
@@ -1492,13 +1494,14 @@ namespace WAN {
             {
                 LOG_INFO("loading from '%s'", file_path.c_str());
 
-                ModelLoader& model_loader = model_manager->loader();
+                ModelLoader model_loader;
                 if (!model_loader.init_from_file_and_convert_name(file_path, "vae.")) {
                     LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                     return;
                 }
 
-                if (!model_manager->register_runner_params("Wan VAE test",
+                if (!model_manager->set_loader(model_loader) ||
+                    !model_manager->register_runner_params(ModelComponent::VAE,
                                                            *vae,
                                                            ModelManager::ResidencyMode::ParamBackend,
                                                            backend,

src/model_component.h

@@ -0,0 +1,61 @@
+#ifndef __SD_MODEL_COMPONENT_H__
+#define __SD_MODEL_COMPONENT_H__
+
+enum class ModelComponent {
+    Conditioner,
+    Diffusion,
+    HighNoiseDiffusion,
+    CLIPVision,
+    IPAdapter,
+    VAE,
+    PreviewVAE,
+    AudioVAE,
+    ControlNet,
+    PhotoMaker,
+    PuLID,
+    LoRA,
+    Upscaler,
+    Detector,
+    LatentUpsampler,
+    Count,
+};
+
+inline const char* model_component_name(ModelComponent component) {
+    switch (component) {
+        case ModelComponent::Conditioner:
+            return "Conditioner model";
+        case ModelComponent::Diffusion:
+            return "Diffusion model";
+        case ModelComponent::HighNoiseDiffusion:
+            return "High noise diffusion model";
+        case ModelComponent::CLIPVision:
+            return "CLIP vision";
+        case ModelComponent::IPAdapter:
+            return "IP-Adapter";
+        case ModelComponent::VAE:
+            return "VAE";
+        case ModelComponent::PreviewVAE:
+            return "preview VAE";
+        case ModelComponent::AudioVAE:
+            return "audio VAE";
+        case ModelComponent::ControlNet:
+            return "ControlNet";
+        case ModelComponent::PhotoMaker:
+            return "photomaker";
+        case ModelComponent::PuLID:
+            return "pulid";
+        case ModelComponent::LoRA:
+            return "LoRA";
+        case ModelComponent::Upscaler:
+            return "ESRGAN";
+        case ModelComponent::Detector:
+            return "YOLOv8";
+        case ModelComponent::LatentUpsampler:
+            return "LTX latent upsampler";
+        case ModelComponent::Count:
+            break;
+    }
+    return "unknown";
+}
+
+#endif  // __SD_MODEL_COMPONENT_H__

src/model_io/safetensors_io.cpp

@@ -144,7 +144,8 @@ static bool read_comfy_quant_config(std::ifstream& file,
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error,
-                           std::map<std::string, std::string>* metadata) {
+                           std::map<std::string, std::string>* metadata,
+                           std::set<std::string>* tensor_names) {
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
         set_error(error, "failed to open '" + file_path + "'");
@@ -246,10 +247,6 @@ bool read_safetensors_file(const std::string& file_path,
         std::string dtype    = tensor_info["dtype"];
         nlohmann::json shape = tensor_info["shape"];
 
-        if (dtype == "U8") {
-            continue;
-        }
-
         size_t begin = tensor_info["data_offsets"][0].get<size_t>();
         size_t end   = tensor_info["data_offsets"][1].get<size_t>();
         if (begin > end || end > file_size_ - data_start) {
@@ -257,6 +254,26 @@ bool read_safetensors_file(const std::string& file_path,
             return false;
         }
 
+        if (tensor_names != nullptr) {
+            tensor_names->insert(name);
+        }
+        if (dtype == "U8") {
+            uint64_t bytes = 1;
+            for (const auto& dimension : shape) {
+                const int64_t size = dimension.get<int64_t>();
+                if (size < 0 || (bytes != 0 && static_cast<uint64_t>(size) > UINT64_MAX / bytes)) {
+                    set_error(error, "invalid dimensions for tensor '" + name + "'");
+                    return false;
+                }
+                bytes *= size;
+            }
+            if (bytes != end - begin) {
+                set_error(error, "size mismatch for tensor '" + name + "'");
+                return false;
+            }
+            continue;
+        }
+
         ggml_type type = safetensors_dtype_to_ggml_type(dtype);
         if (type == GGML_TYPE_COUNT) {
             set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
@@ -270,8 +287,20 @@ bool read_safetensors_file(const std::string& file_path,
 
         int n_dims              = (int)shape.size();
         int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        uint64_t elements       = 1;
         for (int i = 0; i < n_dims; i++) {
             ne[i] = shape[i].get<int64_t>();
+            if (ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(ne[i]) > INT64_MAX / elements)) {
+                set_error(error, "invalid dimensions for tensor '" + name + "'");
+                return false;
+            }
+            elements *= ne[i];
+        }
+        const uint64_t storage_size = ggml_type_size(type) * ((dtype == "F64" || dtype == "I64") ? 2 : 1);
+        if (elements % ggml_blck_size(type) != 0 ||
+            elements / ggml_blck_size(type) > INT64_MAX / storage_size) {
+            set_error(error, "invalid storage size for tensor '" + name + "'");
+            return false;
         }
 
         if (n_dims == 5) {

src/model_io/safetensors_io.h

@@ -2,6 +2,7 @@
 #define __SD_MODEL_IO_SAFETENSORS_IO_H__
 
 #include <map>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -12,7 +13,8 @@ bool is_safetensors_file(const std::string& file_path);
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error                           = nullptr,
-                           std::map<std::string, std::string>* metadata = nullptr);
+                           std::map<std::string, std::string>* metadata = nullptr,
+                           std::set<std::string>* tensor_names          = nullptr);
 bool read_safetensors_index_file(const std::string& file_path,
                                  std::vector<std::string>& shard_paths,
                                  std::string* error = nullptr);

src/model_io/tensor_storage.h

@@ -28,9 +28,11 @@ struct TensorStorage {
     int n_dims                  = 0;
 
     std::string storage_key;
-    size_t file_index = 0;
-    int index_in_zip  = -1;  // >= means stored in a zip file
-    uint64_t offset   = 0;   // offset in file
+    size_t file_index      = 0;
+    uint64_t file_id       = 0;
+    uint64_t file_revision = 0;
+    int index_in_zip       = -1;  // >= means stored in a zip file
+    uint64_t offset        = 0;   // offset in file
 
     TensorStorage() = default;
 

src/model_loader.cpp

@@ -9,6 +9,7 @@
 #include <mutex>
 #include <regex>
 #include <set>
+#include <stdexcept>
 #include <string>
 #include <thread>
 #include <unordered_map>
@@ -27,6 +28,7 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "json.hpp"
 #include "zip.h"
 
 #include "name_conversion.h"
@@ -151,15 +153,19 @@ ModelLoader::ModelLoader()
 }
 
 size_t ModelLoader::add_file_path(const std::string& file_path) {
-    if (model_files_processed) {
-        file_data.clear();
-        model_files_processed = false;
+    auto it = std::find(file_paths_.begin(), file_paths_.end(), file_path);
+    if (it != file_paths_.end()) {
+        return static_cast<size_t>(it - file_paths_.begin());
     }
+    invalidate_file_data();
     file_paths_.push_back(file_path);
     return file_paths_.size() - 1;
 }
 
 void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
+    if (tensor_storage_map.count(tensor_storage.name) != 0) {
+        throw std::runtime_error("duplicate tensor in model source: " + tensor_storage.name);
+    }
     tensor_storage_map[tensor_storage.name] = tensor_storage;
 }
 
@@ -169,6 +175,15 @@ void ModelLoader::set_n_threads(int n_threads) {
 }
 
 bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
+    return add_file(file_path, prefix);
+}
+
+bool ModelLoader::parse_file(const std::string& file_path, const std::string& prefix) {
+    FileStamp stamp;
+    if (!read_file_stamp(file_path, stamp)) {
+        return false;
+    }
+    parsed_dependencies_.push_back(stamp);
     if (is_directory(file_path)) {
         LOG_INFO("load %s using diffusers format", file_path.c_str());
         return init_from_diffusers_file(file_path, prefix);
@@ -198,17 +213,11 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
 }
 
 void ModelLoader::convert_tensors_name() {
-    SDVersion version = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
-    String2TensorStorage new_map;
-
-    for (auto& [_, tensor_storage] : tensor_storage_map) {
-        auto new_name = convert_tensor_name(tensor_storage.name, version);
-        // LOG_VERBOSE("%s -> %s", tensor_storage.name.c_str(), new_name.c_str());
-        tensor_storage.name = new_name;
-        new_map[new_name]   = std::move(tensor_storage);
+    if (names_converted_) {
+        return;
     }
-
-    tensor_storage_map.swap(new_map);
+    names_converted_ = true;
+    rebuild_catalog();
 }
 
 bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path, const std::string& prefix, SDVersion version) {
@@ -257,7 +266,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
 
     std::vector<TensorStorage> tensor_storages;
     std::string error;
-    if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) {
+    if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_, &parsed_tensor_names_[file_path])) {
         LOG_ERROR("%s", error.c_str());
         return false;
     }
@@ -293,7 +302,26 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
     }
 
     for (const std::string& shard_path : shard_paths) {
-        if (!init_from_file(shard_path, prefix)) {
+        if (!parse_file(shard_path, prefix)) {
+            return false;
+        }
+    }
+
+    std::ifstream index_file(file_path);
+    const auto index = nlohmann::json::parse(index_file);
+    for (const auto& entry : index.at("weight_map").items()) {
+        const auto expected = (std::filesystem::u8path(file_path).parent_path() /
+                               std::filesystem::u8path(entry.value().get<std::string>()))
+                                  .lexically_normal();
+        bool found = false;
+        for (const auto& shard : parsed_tensor_names_) {
+            if (std::filesystem::u8path(shard.first).lexically_normal() == expected) {
+                found = shard.second.count(entry.key()) != 0;
+                break;
+            }
+        }
+        if (!found) {
+            LOG_ERROR("safetensors index tensor '%s' is missing from its declared shard", entry.key().c_str());
             return false;
         }
     }
@@ -369,25 +397,23 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
     std::string clip_path   = path_join(file_path, "text_encoder/model.safetensors");
     std::string clip_g_path = path_join(file_path, "text_encoder_2/model.safetensors");
 
-    if (!init_from_safetensors_file(unet_path, "unet.")) {
+    if (!parse_file(unet_path, prefix + "unet.")) {
         return false;
     }
 
-    if (!init_from_safetensors_file(vae_path, "vae.")) {
-        LOG_WARN("Couldn't find working VAE in %s", file_path.c_str());
-        // return false;
+    if (file_exists(vae_path) && !parse_file(vae_path, prefix + "vae.")) {
+        return false;
     }
-    if (!init_from_safetensors_file(clip_path, "te.")) {
-        LOG_WARN("Couldn't find working text encoder in %s", file_path.c_str());
-        // return false;
+    if (file_exists(clip_path) && !parse_file(clip_path, prefix + "te.")) {
+        return false;
     }
-    if (!init_from_safetensors_file(clip_g_path, "te.1.")) {
-        LOG_VERBOSE("Couldn't find working second text encoder in %s", file_path.c_str());
+    if (file_exists(clip_g_path) && !parse_file(clip_g_path, prefix + "te.1.")) {
+        return false;
     }
     return true;
 }
 
-SDVersion ModelLoader::get_sd_version() {
+SDVersion ModelLoader::get_sd_version() const {
     TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
 
     bool has_multiple_encoders = false;
@@ -623,7 +649,7 @@ SDVersion ModelLoader::get_sd_version() {
     return VERSION_COUNT;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
+std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -640,7 +666,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
     return wtype_stat;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
+std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -664,7 +690,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
     return wtype_stat;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
+std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -685,7 +711,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
     return wtype_stat;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() {
+std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -743,9 +769,12 @@ TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules) {
 }
 
 void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_rules) {
-    auto map_rules = parse_tensor_type_rules(tensor_type_rules);
+    wtype_override_    = wtype;
+    tensor_type_rules_ = tensor_type_rules;
+    auto map_rules     = parse_tensor_type_rules(tensor_type_rules);
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        ggml_type dst_type = wtype;
+        tensor_storage.expected_type = GGML_TYPE_COUNT;
+        ggml_type dst_type           = wtype;
         for (const auto& tensor_type_rule : map_rules) {
             std::regex pattern(tensor_type_rule.first);
             if (std::regex_search(name, pattern)) {
@@ -761,6 +790,8 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
         }
         tensor_storage.expected_type = dst_type;
     }
+    invalidate_file_data();
+    ++revision_;
 }
 
 void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
@@ -829,6 +860,13 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
 std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                                        std::set<std::string> ignore_tensors,
                                                        bool writable_mmap) {
+    std::set<std::string> names;
+    for (const auto& entry : tensors) {
+        names.insert(entry.first);
+    }
+    if (!validate_sources(&names)) {
+        return {};
+    }
     process_model_files(true, writable_mmap);
 
     std::vector<MmapTensorStore> result;
@@ -919,6 +957,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                                bool enable_mmap,
                                const std::set<std::string>* target_tensor_names,
                                bool log_progress) {
+    if (!validate_sources(target_tensor_names)) {
+        return false;
+    }
     process_model_files(enable_mmap, false);
 
     std::atomic<int64_t> read_time_ms(0);
@@ -1242,7 +1283,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                  (convert_time_ms.load() / (float)last_n_threads) / 1000.f,
                  (copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
     }
-    return success;
+    return success && validate_sources(target_tensor_names);
 }
 
 bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor) {
@@ -1259,7 +1300,9 @@ bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor*
             return true;
         }
 
-        if (current_tensor_storage.file_index != tensor_storage.file_index ||
+        if (current_tensor_storage.file_id != tensor_storage.file_id ||
+            current_tensor_storage.file_revision != tensor_storage.file_revision ||
+            current_tensor_storage.file_index != tensor_storage.file_index ||
             current_tensor_storage.offset != tensor_storage.offset ||
             current_tensor_storage.index_in_zip != tensor_storage.index_in_zip) {
             LOG_ERROR("load tensor failed: storage mismatch for '%s'", tensor_storage.name.c_str());
@@ -1440,7 +1483,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
     return true;
 }
 
-bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) {
+bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const {
     const std::string& name = tensor_storage.name;
     if (tensor_storage.is_int8_tensorwise) {
         return false;
@@ -1478,7 +1521,7 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
     return false;
 }
 
-int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) {
+int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) const {
     size_t alignment = 128;
     if (backend != nullptr) {
         alignment = ggml_backend_get_alignment(backend);

src/model_loader.h

@@ -2,6 +2,7 @@
 #define __MODEL_LOADER_H__
 
 #include <cstdint>
+#include <filesystem>
 #include <map>
 #include <memory>
 #include <set>
@@ -30,6 +31,46 @@ struct MmapTensorStore {
 bool is_unused_tensor(const std::string& name);
 
 class ModelLoader {
+public:
+    using FileId       = uint64_t;
+    using FileVersions = std::map<FileId, uint64_t>;
+    enum class FileScope { Catalog,
+                           Isolated };
+
+private:
+    struct FileStamp {
+        std::string path;
+        uintmax_t size = 0;
+        std::filesystem::file_time_type modified;
+    };
+
+    struct FileRecord {
+        FileId id         = 0;
+        uint64_t revision = 0;
+        std::string path;
+        std::string prefix;
+        FileScope scope = FileScope::Catalog;
+        std::vector<FileStamp> dependencies;
+        String2TensorStorage tensors;
+        std::map<std::string, std::string> metadata;
+    };
+
+    std::vector<FileRecord> files_;
+    uint64_t revision_        = 0;
+    bool names_converted_     = false;
+    ggml_type wtype_override_ = GGML_TYPE_COUNT;
+    std::string tensor_type_rules_;
+    std::vector<FileStamp> parsed_dependencies_;
+    std::map<std::string, std::set<std::string>> parsed_tensor_names_;
+
+    static bool read_file_stamp(const std::string& path, FileStamp& stamp);
+    static bool file_unchanged(const FileStamp& stamp);
+    bool parse_file(const std::string& path, const std::string& prefix);
+    bool add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope);
+    ModelLoader file_reader(FileId id, SDVersion version) const;
+    void rebuild_catalog();
+    void invalidate_file_data();
+
 protected:
     SDVersion version_ = VERSION_COUNT;
     std::vector<std::string> file_paths_;
@@ -52,16 +93,27 @@ protected:
 public:
     ModelLoader();
 
+    bool add_file(const std::string& path, const std::string& prefix = "", FileId* id = nullptr, bool force = false, FileScope scope = FileScope::Catalog);
+    bool del_file(FileId id);
+    uint64_t file_revision(FileId id) const;
+    std::string file_path(FileId id) const;
+    String2TensorStorage file_tensors(FileId id, SDVersion version) const;
+    bool load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap = false) const;
+    bool refresh_files(bool include_isolated = true);
+    bool files_changed(bool& changed, bool include_isolated = true) const;
+    bool validate_sources(const std::set<std::string>* tensor_names = nullptr) const;
+    uint64_t revision() const { return revision_; }
+    FileVersions file_versions(const std::vector<std::string>& prefixes = {}) const;
     bool init_from_file(const std::string& file_path, const std::string& prefix = "");
     void convert_tensors_name();
     bool init_from_file_and_convert_name(const std::string& file_path,
                                          const std::string& prefix = "",
                                          SDVersion version         = VERSION_COUNT);
-    SDVersion get_sd_version();
-    std::map<ggml_type, uint32_t> get_wtype_stat();
-    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat();
-    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat();
-    std::map<ggml_type, uint32_t> get_vae_wtype_stat();
+    SDVersion get_sd_version() const;
+    std::map<ggml_type, uint32_t> get_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_vae_wtype_stat() const;
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
     const std::map<std::string, std::string>& get_metadata() const { return metadata_; }
@@ -92,8 +144,8 @@ public:
         return names;
     }
 
-    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
-    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
+    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const;
+    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT) const;
     ~ModelLoader() = default;
 };
 

src/model_loader_files.cpp

@@ -0,0 +1,336 @@
+#include "model_loader.h"
+
+#include <algorithm>
+#include <atomic>
+#include <limits>
+
+#include "core/util.h"
+#include "name_conversion.h"
+
+static uint64_t next_source_revision() {
+    static std::atomic<uint64_t> revision{0};
+    return revision.fetch_add(1, std::memory_order_relaxed) + 1;
+}
+
+bool ModelLoader::read_file_stamp(const std::string& path, FileStamp& stamp) {
+    std::error_code error;
+    const auto file_path = std::filesystem::u8path(path);
+    stamp.path           = path;
+    stamp.size           = 0;
+    stamp.modified       = std::filesystem::last_write_time(file_path, error);
+    if (!error && std::filesystem::is_regular_file(file_path, error)) {
+        stamp.size = std::filesystem::file_size(file_path, error);
+    }
+    if (error) {
+        LOG_ERROR("cannot inspect model source '%s': %s", path.c_str(), error.message().c_str());
+        return false;
+    }
+    return true;
+}
+
+bool ModelLoader::file_unchanged(const FileStamp& stamp) {
+    std::error_code error;
+    if (!std::filesystem::exists(std::filesystem::u8path(stamp.path), error)) {
+        return false;
+    }
+    FileStamp current;
+    return read_file_stamp(stamp.path, current) &&
+           current.size == stamp.size && current.modified == stamp.modified;
+}
+
+void ModelLoader::invalidate_file_data() {
+    file_data.clear();
+    model_files_processed = false;
+}
+
+void ModelLoader::rebuild_catalog() {
+    tensor_storage_map.clear();
+    metadata_.clear();
+    for (const auto& file : files_) {
+        if (file.scope == FileScope::Isolated)
+            continue;
+        for (const auto& entry : file.tensors) {
+            tensor_storage_map[entry.first] = entry.second;
+        }
+        for (const auto& entry : file.metadata) {
+            metadata_[entry.first] = entry.second;
+        }
+    }
+    if (names_converted_) {
+        const SDVersion version = version_ == VERSION_COUNT ? get_sd_version() : version_;
+        tensor_storage_map.clear();
+        for (const auto& file : files_) {
+            if (file.scope == FileScope::Isolated)
+                continue;
+            for (const auto& entry : file.tensors) {
+                TensorStorage tensor            = entry.second;
+                tensor.name                     = convert_tensor_name(tensor.name, version);
+                tensor_storage_map[tensor.name] = std::move(tensor);
+            }
+        }
+    }
+    std::set<size_t> used_files;
+    for (const auto& file : files_) {
+        for (const auto& entry : file.tensors) {
+            used_files.insert(entry.second.file_index);
+        }
+    }
+    for (size_t i = 0; i < file_paths_.size(); ++i) {
+        if (used_files.count(i) == 0) {
+            file_paths_[i].clear();
+        }
+    }
+    set_wtype_override(wtype_override_, tensor_type_rules_);
+}
+
+bool ModelLoader::add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope) {
+    FileStamp root;
+    if (!read_file_stamp(path, root)) {
+        return false;
+    }
+    auto existing = std::find_if(files_.begin(), files_.end(), [&](const FileRecord& file) {
+        return file.path == root.path && file.prefix == prefix && file.scope == scope;
+    });
+    if (existing != files_.end() && !force &&
+        std::all_of(existing->dependencies.begin(), existing->dependencies.end(), file_unchanged)) {
+        if (id != nullptr) {
+            *id = existing->id;
+        }
+        return true;
+    }
+
+    ModelLoader parsed;
+    try {
+        if (!parsed.parse_file(root.path, prefix)) {
+            return false;
+        }
+    } catch (const std::exception& error) {
+        LOG_ERROR("invalid model source '%s': %s", path.c_str(), error.what());
+        return false;
+    }
+
+    std::vector<size_t> file_indices;
+    std::vector<FileStamp> physical_files;
+    for (const auto& physical_path : parsed.file_paths_) {
+        FileStamp stamp;
+        if (!read_file_stamp(physical_path, stamp)) {
+            return false;
+        }
+        parsed.parsed_dependencies_.push_back(stamp);
+        file_indices.push_back(add_file_path(stamp.path));
+        physical_files.push_back(std::move(stamp));
+    }
+    for (auto& entry : parsed.tensor_storage_map) {
+        auto& tensor = entry.second;
+        // Pickle preserves rank-zero scalars; GGML uses a one-element dimension.
+        if (tensor.n_dims == 0) {
+            tensor.n_dims = 1;
+        }
+        if (tensor.n_dims < 1 || tensor.n_dims > SD_MAX_DIMS || tensor.type < 0 ||
+            tensor.type >= GGML_TYPE_COUNT || tensor.file_index >= parsed.file_paths_.size()) {
+            LOG_ERROR("invalid tensor metadata for '%s'", tensor.name.c_str());
+            return false;
+        }
+        uint64_t elements = 1;
+        for (int i = 0; i < tensor.n_dims; ++i) {
+            if (tensor.ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(tensor.ne[i]) > INT64_MAX / elements)) {
+                LOG_ERROR("invalid tensor dimensions for '%s'", tensor.name.c_str());
+                return false;
+            }
+            elements *= tensor.ne[i];
+        }
+        const uint64_t block_size = ggml_blck_size(tensor.type);
+        const uint64_t type_size  = ggml_type_size(tensor.type) * ((tensor.is_f64 || tensor.is_i64) ? 2 : 1);
+        if (block_size == 0 || type_size == 0 || elements % block_size != 0 || elements / block_size > INT64_MAX / type_size) {
+            LOG_ERROR("invalid tensor storage size for '%s'", tensor.name.c_str());
+            return false;
+        }
+        if (tensor.index_in_zip < 0) {
+            const auto& stamp = physical_files[tensor.file_index];
+            if (tensor.offset > stamp.size || elements / block_size * type_size > stamp.size - tensor.offset) {
+                LOG_ERROR("tensor '%s' extends beyond its model file", tensor.name.c_str());
+                return false;
+            }
+        }
+    }
+    if (!std::all_of(parsed.parsed_dependencies_.begin(), parsed.parsed_dependencies_.end(), file_unchanged)) {
+        LOG_ERROR("model source changed while reading metadata: '%s'", path.c_str());
+        return false;
+    }
+
+    FileRecord record;
+    // Snapshots and independently created loaders must never alias different versions.
+    record.revision = next_source_revision();
+    record.id       = existing == files_.end() ? record.revision : existing->id;
+    ++revision_;
+    record.path   = root.path;
+    record.prefix = prefix;
+    record.scope  = scope;
+    std::set<std::string> seen_dependencies;
+    for (auto& stamp : parsed.parsed_dependencies_) {
+        if (seen_dependencies.insert(stamp.path).second) {
+            record.dependencies.push_back(std::move(stamp));
+        }
+    }
+    record.metadata = std::move(parsed.metadata_);
+    record.tensors  = std::move(parsed.tensor_storage_map);
+    for (auto& entry : record.tensors) {
+        entry.second.file_index    = file_indices[entry.second.file_index];
+        entry.second.file_id       = record.id;
+        entry.second.file_revision = record.revision;
+    }
+    if (id != nullptr) {
+        *id = record.id;
+    }
+    if (existing == files_.end()) {
+        files_.push_back(std::move(record));
+    } else {
+        *existing = std::move(record);
+    }
+    rebuild_catalog();
+    return true;
+}
+
+bool ModelLoader::add_file(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope) {
+    ModelLoader candidate = *this;
+    FileId added_id       = 0;
+    if (!candidate.add_file_impl(path, prefix, &added_id, force, scope)) {
+        return false;
+    }
+    *this = std::move(candidate);
+    if (id != nullptr) {
+        *id = added_id;
+    }
+    return true;
+}
+
+bool ModelLoader::del_file(FileId id) {
+    auto it = std::find_if(files_.begin(), files_.end(), [id](const FileRecord& file) { return file.id == id; });
+    if (it == files_.end()) {
+        return false;
+    }
+    files_.erase(it);
+    ++revision_;
+    rebuild_catalog();
+    return true;
+}
+
+bool ModelLoader::files_changed(bool& changed, bool include_isolated) const {
+    changed = false;
+    for (const auto& file : files_) {
+        if (!include_isolated && file.scope == FileScope::Isolated)
+            continue;
+        for (const auto& stamp : file.dependencies) {
+            std::error_code error;
+            if (!std::filesystem::exists(std::filesystem::u8path(stamp.path), error) && !error) {
+                // An updated index may no longer reference this dependency.
+                changed = true;
+                continue;
+            }
+            FileStamp current;
+            if (!read_file_stamp(stamp.path, current)) {
+                return false;
+            }
+            changed |= current.size != stamp.size || current.modified != stamp.modified;
+        }
+    }
+    return true;
+}
+
+bool ModelLoader::refresh_files(bool include_isolated) {
+    bool changed;
+    if (!files_changed(changed, include_isolated)) {
+        return false;
+    }
+    if (!changed) {
+        return true;
+    }
+    ModelLoader candidate = *this;
+    for (const auto& file : files_) {
+        if (!include_isolated && file.scope == FileScope::Isolated)
+            continue;
+        if (!candidate.add_file_impl(file.path, file.prefix, nullptr, false, file.scope)) {
+            return false;
+        }
+    }
+    *this = std::move(candidate);
+    return true;
+}
+
+bool ModelLoader::validate_sources(const std::set<std::string>* tensor_names) const {
+    std::set<FileId> required;
+    if (tensor_names != nullptr) {
+        for (const auto& name : *tensor_names) {
+            auto it = tensor_storage_map.find(name);
+            if (it != tensor_storage_map.end()) {
+                required.insert(it->second.file_id);
+            }
+        }
+    }
+    for (const auto& file : files_) {
+        if (tensor_names != nullptr && required.count(file.id) == 0) {
+            continue;
+        }
+        if (!std::all_of(file.dependencies.begin(), file.dependencies.end(), file_unchanged)) {
+            LOG_ERROR("model source changed; refresh it before execution: '%s'", file.path.c_str());
+            return false;
+        }
+    }
+    return true;
+}
+
+ModelLoader::FileVersions ModelLoader::file_versions(const std::vector<std::string>& prefixes) const {
+    FileVersions versions;
+    for (const auto& entry : tensor_storage_map) {
+        if (prefixes.empty() || std::any_of(prefixes.begin(), prefixes.end(), [&](const std::string& prefix) {
+                return starts_with(entry.first, prefix);
+            })) {
+            versions[entry.second.file_id] = entry.second.file_revision;
+        }
+    }
+    return versions;
+}
+
+uint64_t ModelLoader::file_revision(FileId id) const {
+    for (const auto& file : files_) {
+        if (file.id == id)
+            return file.revision;
+    }
+    return 0;
+}
+
+std::string ModelLoader::file_path(FileId id) const {
+    for (const auto& file : files_) {
+        if (file.id == id)
+            return file.path;
+    }
+    return {};
+}
+
+ModelLoader ModelLoader::file_reader(FileId id, SDVersion version) const {
+    ModelLoader reader;
+    reader.file_paths_      = file_paths_;
+    reader.n_threads_       = n_threads_;
+    reader.version_         = version;
+    reader.names_converted_ = true;
+    for (const auto& file : files_) {
+        if (file.id == id) {
+            reader.files_.push_back(file);
+            reader.files_.back().scope = FileScope::Catalog;
+            break;
+        }
+    }
+    reader.rebuild_catalog();
+    return reader;
+}
+
+String2TensorStorage ModelLoader::file_tensors(FileId id, SDVersion version) const {
+    return file_reader(id, version).tensor_storage_map;
+}
+
+bool ModelLoader::load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap) const {
+    if (file_revision(id) == 0)
+        return false;
+    auto reader = file_reader(id, version);
+    return reader.load_tensors(callback, use_mmap, &names, false);
+}

src/model_manager.cpp

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <iterator>
 #include <mutex>
+#include <tuple>
 #include <unordered_set>
 
 #include "core/ggml_extend_backend.h"
@@ -26,7 +27,8 @@ static bool lora_specs_equal(const std::vector<ModelManager::LoraSpec>& lhs,
             lhs[i].multiplier != rhs[i].multiplier ||
             lhs[i].is_high_noise != rhs[i].is_high_noise ||
             lhs[i].tensor_name_prefix_filter != rhs[i].tensor_name_prefix_filter ||
-            lhs[i].required != rhs[i].required) {
+            lhs[i].required != rhs[i].required ||
+            lhs[i].file_id != rhs[i].file_id || lhs[i].file_revision != rhs[i].file_revision) {
             return false;
         }
     }
@@ -104,25 +106,61 @@ void ModelManager::set_common_ignore_tensors(std::set<std::string> ignore_tensor
     common_ignore_tensors_ = std::move(ignore_tensors);
 }
 
-void ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
-    if (loras.empty() && loras_.empty()) {
-        lora_version_ = version;
-        return;
+bool ModelManager::prepare_lora_sources(std::vector<LoraSpec>& loras) {
+    ModelLoader candidate = model_loader_;
+    std::vector<LoraSpec> resolved;
+    std::set<ModelLoader::FileId> sources;
+    for (auto spec : loras) {
+        const std::string prefix = spec.is_high_noise ? "lora.model.high_noise_" : "lora.";
+        if (!candidate.add_file(spec.path, prefix, &spec.file_id, false, ModelLoader::FileScope::Isolated)) {
+            if (spec.required)
+                return false;
+            LOG_WARN("cannot register LoRA source '%s'", spec.path.c_str());
+            continue;
+        }
+        spec.file_revision = candidate.file_revision(spec.file_id);
+        sources.insert(spec.file_id);
+        resolved.push_back(std::move(spec));
     }
-    if (lora_version_ == version && lora_specs_equal(loras_, loras)) {
-        return;
+    for (auto id : lora_sources_) {
+        if (sources.count(id) == 0)
+            candidate.del_file(id);
     }
+    if (!set_loader(std::move(candidate)))
+        return false;
+    lora_sources_ = std::move(sources);
+    loras         = std::move(resolved);
+    return true;
+}
 
+bool ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
+    if (std::any_of(loras.begin(), loras.end(), [](const LoraSpec& spec) { return spec.file_id == 0; }) &&
+        !prepare_lora_sources(loras))
+        return false;
+    for (auto& spec : loras) {
+        spec.file_revision = model_loader_.file_revision(spec.file_id);
+        if (spec.file_revision == 0)
+            return false;
+    }
+    if (lora_version_ == version && lora_specs_equal(loras_, loras))
+        return true;
+    if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
+            return state->pin_count != 0;
+        })) {
+        LOG_ERROR("cannot change LoRA configuration during execution");
+        return false;
+    }
     loras_        = std::move(loras);
     lora_version_ = version;
     current_lora_epoch_++;
     reset_lora_applied_params();
+    return true;
 }
 
 std::set<std::string> ModelManager::tensor_names() const {
     std::set<std::string> names;
     for (const auto& state : tensor_states_) {
-        if (state != nullptr) {
+        if (state != nullptr && state->component != ModelComponent::LoRA) {
             names.insert(state->name);
         }
     }
@@ -171,7 +209,7 @@ ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState
     return state.split_buffer_type;
 }
 
-bool ModelManager::register_param_tensors(const std::string& desc,
+bool ModelManager::register_param_tensors(ModelComponent component,
                                           std::map<std::string, ggml_tensor*> tensors,
                                           ResidencyMode residency_mode,
                                           ggml_backend_t compute_backend,
@@ -179,15 +217,20 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           size_t* registered_tensor_size,
                                           bool allow_split_buffer,
                                           bool params_follow_compute_backend,
-                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
-    if (desc.empty()) {
-        LOG_ERROR("model manager tensor desc is empty");
+                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops,
+                                          ModelLoader::FileId source_file,
+                                          SDVersion source_version) {
+    if (component == ModelComponent::Count) {
+        LOG_ERROR("model manager tensor component is invalid");
         return false;
     }
     if (registered_tensor_size != nullptr) {
         *registered_tensor_size += estimate_tensors_size(tensors);
     }
 
+    const auto scoped_sources = source_file != 0 ? model_loader_.file_tensors(source_file, source_version) : String2TensorStorage{};
+    const auto& sources       = source_file != 0 ? scoped_sources : model_loader_.get_tensor_storage_map();
+    std::unordered_set<ggml_tensor*> new_tensors;
     std::vector<std::unique_ptr<TensorState>> new_states;
     new_states.reserve(tensors.size());
 
@@ -197,16 +240,23 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         if (tensor == nullptr) {
             continue;
         }
-        if (tensor_states_by_name_.find(name) != tensor_states_by_name_.end()) {
+        if (tensor_states_by_tensor_.count(tensor) != 0 || !new_tensors.insert(tensor).second) {
             LOG_ERROR("model manager tensor name '%s' is already registered", name.c_str());
             return false;
         }
         ggml_set_name(tensor, name.c_str());
 
-        auto state             = std::make_unique<TensorState>();
-        state->name            = name;
-        state->tensor          = tensor;
-        state->desc            = desc;
+        auto state            = std::make_unique<TensorState>();
+        state->name           = name;
+        state->tensor         = tensor;
+        state->component      = component;
+        state->source_file    = source_file;
+        state->source_version = source_version;
+        auto source           = sources.find(name);
+        if (source != sources.end()) {
+            state->source     = source->second;
+            state->has_source = true;
+        }
         state->residency_mode  = residency_mode;
         state->compute_backend = compute_backend;
         state->params_backend  = params_backend;
@@ -225,31 +275,45 @@ bool ModelManager::register_param_tensors(const std::string& desc,
     }
 
     for (auto& state : new_states) {
-        TensorState* registered_state                  = state.get();
-        tensor_states_by_name_[registered_state->name] = registered_state;
+        TensorState* registered_state                      = state.get();
+        tensor_states_by_tensor_[registered_state->tensor] = registered_state;
         tensor_states_.push_back(std::move(state));
     }
     return true;
 }
 
-bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) {
-    if (desc.empty()) {
-        return true;
+bool ModelManager::unregister_param_tensors(ModelComponent component, size_t* registered_tensor_size) {
+    std::unordered_set<TensorState*> states;
+    for (auto& state : tensor_states_) {
+        if (state->component == component)
+            states.insert(state.get());
     }
+    return unregister_tensor_states(states, registered_tensor_size);
+}
 
-    std::unordered_set<TensorState*> target_states;
+bool ModelManager::unregister_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+    std::unordered_set<TensorState*> states;
+    for (auto tensor : tensors) {
+        auto found = tensor_states_by_tensor_.find(tensor);
+        if (found != tensor_states_by_tensor_.end())
+            states.insert(found->second);
+    }
+    return unregister_tensor_states(states, nullptr);
+}
+
+bool ModelManager::unregister_tensor_states(const std::unordered_set<TensorState*>& target_states,
+                                            size_t* registered_tensor_size) {
     size_t released_size = 0;
     for (auto& state : tensor_states_) {
-        if (state == nullptr || state->desc != desc) {
+        if (state == nullptr || target_states.count(state.get()) == 0) {
             continue;
         }
         if (state->pin_count > 0) {
             LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
-                      desc.c_str(),
+                      model_component_name(state->component),
                       state->name.c_str());
             return false;
         }
-        target_states.insert(state.get());
         if (state->tensor != nullptr) {
             released_size += ggml_nbytes(state->tensor);
         }
@@ -260,7 +324,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
     }
 
     clear_all_prefetched_params();
-    release_compute_staging_blocks(false);
+    release_compute_staging_blocks(false, &target_states);
 
     std::vector<ParamsStorageBlock*> storage_blocks_to_release;
     std::unordered_set<TensorState*> affected_storage_states;
@@ -292,7 +356,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
         }
         if (state->pin_count > 0 || state->staged_to_compute_backend) {
             LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
-                      desc.c_str(),
+                      model_component_name(state->component),
                       state->name.c_str());
             return false;
         }
@@ -305,9 +369,9 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
         }
     }
 
-    for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) {
+    for (auto it = tensor_states_by_tensor_.begin(); it != tensor_states_by_tensor_.end();) {
         if (target_states.count(it->second) > 0) {
-            it = tensor_states_by_name_.erase(it);
+            it = tensor_states_by_tensor_.erase(it);
         } else {
             ++it;
         }
@@ -559,19 +623,24 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
 }
 
 bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states) {
-    if (loras_.empty()) {
+    if (loras_.empty() || applying_loras_)
         return true;
-    }
+    applying_loras_ = true;
+    struct ApplyGuard {
+        bool& active;
+        ~ApplyGuard() { active = false; }
+    } guard{applying_loras_};
 
     struct LoraApplyGroup {
         std::map<std::string, ggml_tensor*> model_tensors;
         std::vector<TensorState*> states;
     };
 
-    std::map<ggml_backend_t, LoraApplyGroup> groups;
+    using ApplyTarget = std::tuple<ggml_backend_t, ggml_backend_t, ResidencyMode>;
+    std::map<ApplyTarget, LoraApplyGroup> groups;
     for (TensorState* state : states) {
-        if (state == nullptr || state->tensor == nullptr ||
-            should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+        if (state == nullptr || state->tensor == nullptr || state->component == ModelComponent::LoRA ||
+            state->component == ModelComponent::LatentUpsampler || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
             continue;
         }
         if (state->applied_lora_epoch == current_lora_epoch_) {
@@ -596,7 +665,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
             LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
             return false;
         }
-        LoraApplyGroup& group            = groups[state->compute_backend];
+        LoraApplyGroup& group            = groups[{state->compute_backend, state->params_backend, state->residency_mode}];
         group.model_tensors[state->name] = state->tensor;
         group.states.push_back(state);
     }
@@ -607,20 +676,20 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
 
     std::set<std::string> all_tensor_names = tensor_names();
     for (auto& group_pair : groups) {
-        ggml_backend_t compute_backend = group_pair.first;
+        ggml_backend_t compute_backend = std::get<0>(group_pair.first);
         LoraApplyGroup& group          = group_pair.second;
         for (const LoraSpec& lora_spec : loras_) {
             if (group.model_tensors.empty()) {
                 continue;
             }
 
-            std::string id = lora_id(lora_spec);
-            auto lora      = std::make_shared<LoraModel>(id,
-                                                    compute_backend,
-                                                    compute_backend,
-                                                    lora_spec.path,
-                                                    lora_spec.is_high_noise ? "model.high_noise_" : "",
-                                                    lora_version_);
+            std::string id     = lora_id(lora_spec);
+            const auto* target = group.states.front();
+            // The temporary runner is destroyed before this manager call returns.
+            auto borrowed_manager = std::shared_ptr<ModelManager>(this, [](ModelManager*) {});
+            auto lora             = std::make_shared<LoraModel>(id, compute_backend, target->params_backend,
+                                                    borrowed_manager, lora_spec.file_id, lora_version_,
+                                                    target->residency_mode);
 
             LoraModel::filter_t lora_tensor_filter = nullptr;
             if (!lora_spec.tensor_name_prefix_filter.empty()) {
@@ -628,7 +697,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
                     return starts_with(tensor_name, lora_spec.tensor_name_prefix_filter);
                 };
             }
-            if (!lora->load_from_file(n_threads_, lora_tensor_filter)) {
+            if (!lora->init_params(n_threads_, lora_tensor_filter)) {
                 LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str());
                 if (lora_spec.required) {
                     return false;
@@ -643,7 +712,8 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
                 continue;
             }
             lora->multiplier = lora_spec.multiplier;
-            lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false);
+            if (!lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false))
+                return false;
             lora->release_loaded_tensors();
         }
 
@@ -657,12 +727,13 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
 }
 
 void ModelManager::reset_lora_applied_params() {
-    clear_all_prefetched_params();
-    release_compute_staging_blocks(true);
-    release_params_storage_blocks(true);
+    std::unordered_set<TensorState*> affected;
     for (auto& state : tensor_states_) {
-        state->applied_lora_epoch = UINT64_MAX;
+        if (state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX) {
+            affected.insert(state.get());
+        }
     }
+    invalidate_sources(affected);
 }
 
 bool ModelManager::should_ignore(const TensorState& state) const {
@@ -684,21 +755,19 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
         return true;
     }
 
-    const auto& tensor_storage_map = model_loader_.get_tensor_storage_map();
-    auto ts_it                     = tensor_storage_map.find(state.name);
-    if (ts_it == tensor_storage_map.end()) {
-        LOG_ERROR("%s tensor '%s' not in model metadata", state.desc.c_str(), state.name.c_str());
+    if (!state.has_source) {
+        LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str());
         return false;
     }
 
-    const TensorStorage& tensor_storage = ts_it->second;
+    const TensorStorage& tensor_storage = state.source;
     if (state.tensor->ne[0] != tensor_storage.ne[0] ||
         state.tensor->ne[1] != tensor_storage.ne[1] ||
         state.tensor->ne[2] != tensor_storage.ne[2] ||
         state.tensor->ne[3] != tensor_storage.ne[3]) {
         LOG_ERROR(
             "%s tensor '%s' has wrong shape in model metadata: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
-            state.desc.c_str(),
+            model_component_name(state.component),
             state.name.c_str(),
             (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
             (int)state.tensor->ne[0], (int)state.tensor->ne[1], (int)state.tensor->ne[2], (int)state.tensor->ne[3]);
@@ -746,7 +815,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
 }
 
 bool ModelManager::can_mmap_storage(const TensorState& state) const {
-    if (!enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
+    if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
         return false;
     }
     if (state.compute_backend == nullptr || state.params_backend == nullptr) {
@@ -857,75 +926,55 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
 }
 
 bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
-    std::map<std::string, TensorState*> states_by_name;
-    std::set<std::string> target_tensor_names;
-    for (TensorState* state : states) {
-        if (state == nullptr) {
+    using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
+    using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
+    std::map<ReadGroup, std::vector<ReadBatch>> groups;
+    for (auto* state : states) {
+        if (state == nullptr)
             continue;
+        auto& batches = groups[{state->source_file, state->source_version}];
+        // The loader supplies one destination per name; only conflicting types need another batch.
+        auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
+            auto found = candidate.find(state->name);
+            return found == candidate.end() || found->second.front()->tensor->type == state->tensor->type;
+        });
+        if (batch == batches.end()) {
+            batches.emplace_back();
+            batch = std::prev(batches.end());
         }
-        states_by_name[state->name] = state;
-        target_tensor_names.insert(state->name);
+        (*batch)[state->name].push_back(state);
     }
-    if (states_by_name.empty()) {
-        return true;
-    }
-
-    std::set<std::string> loaded_names;
-    std::mutex loaded_names_mutex;
-    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-        const std::string& name = tensor_storage.name;
-        *dst_tensor             = nullptr;
-
-        auto state_it = states_by_name.find(name);
-        if (state_it == states_by_name.end()) {
-            return true;
+    for (auto& group : groups) {
+        for (auto& batch : group.second) {
+            std::set<std::string> names;
+            std::set<std::string> loaded;
+            std::mutex mutex;
+            for (const auto& entry : batch)
+                names.insert(entry.first);
+            auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
+                *dst       = nullptr;
+                auto found = batch.find(source.name);
+                if (found == batch.end())
+                    return true;
+                *dst = found->second.front()->tensor;
+                std::lock_guard<std::mutex> lock(mutex);
+                loaded.insert(source.name);
+                return true;
+            };
+            const auto file = group.first.first;
+            bool success    = file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names)
+                                        : model_loader_.load_file_tensors(file, group.first.second, callback, names, enable_mmap_);
+            if (!success || loaded != names)
+                return false;
+            for (auto& entry : batch) {
+                auto* first = entry.second.front()->tensor;
+                for (auto* state : entry.second) {
+                    if (state->tensor != first)
+                        ggml_backend_tensor_copy(first, state->tensor);
+                    state->loaded_to_params_backend = true;
+                }
+            }
         }
-
-        TensorState* state = state_it->second;
-        if (state == nullptr || state->tensor == nullptr) {
-            LOG_ERROR("model manager tensor '%s' is null", name.c_str());
-            return false;
-        }
-
-        if (state->tensor->ne[0] != tensor_storage.ne[0] ||
-            state->tensor->ne[1] != tensor_storage.ne[1] ||
-            state->tensor->ne[2] != tensor_storage.ne[2] ||
-            state->tensor->ne[3] != tensor_storage.ne[3]) {
-            LOG_ERROR(
-                "model manager tensor '%s' has wrong shape in model file: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
-                name.c_str(),
-                (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
-                (int)state->tensor->ne[0], (int)state->tensor->ne[1], (int)state->tensor->ne[2], (int)state->tensor->ne[3]);
-            return false;
-        }
-
-        {
-            std::lock_guard<std::mutex> lock(loaded_names_mutex);
-            loaded_names.insert(name);
-        }
-        *dst_tensor = state->tensor;
-        return true;
-    };
-
-    if (!model_loader_.load_tensors(on_new_tensor_cb, enable_mmap_, &target_tensor_names)) {
-        LOG_ERROR("model manager load tensors failed");
-        return false;
-    }
-
-    bool missing = false;
-    for (const auto& pair : states_by_name) {
-        const std::string& name = pair.first;
-        if (loaded_names.find(name) == loaded_names.end()) {
-            LOG_ERROR("model manager tensor '%s' was not loaded", name.c_str());
-            missing = true;
-        }
-    }
-    if (missing) {
-        return false;
-    }
-
-    for (const auto& pair : states_by_name) {
-        pair.second->loaded_to_params_backend = true;
     }
     return true;
 }
@@ -1138,6 +1187,14 @@ void ModelManager::release_all() {
     release_params_storage_blocks(true);
 }
 
+ggml_tensor* ModelManager::resolve_param_tensor(ggml_tensor* tensor) const {
+    for (auto* current = tensor; current != nullptr; current = current->view_src) {
+        if (tensor_states_by_tensor_.count(current) != 0)
+            return current;
+    }
+    return nullptr;
+}
+
 bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
                                                   std::vector<TensorState*>& required_states,
                                                   ggml_backend_t compute_backend) const {
@@ -1147,21 +1204,13 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
         if (tensor == nullptr) {
             continue;
         }
-        const char* raw_name = ggml_get_name(tensor);
-        if (raw_name == nullptr || raw_name[0] == '\0') {
-            LOG_ERROR("model manager unnamed tensor is not registered");
-            return false;
-        }
-        auto state_it = tensor_states_by_name_.find(raw_name);
-        if (state_it == tensor_states_by_name_.end()) {
-            LOG_ERROR("model manager tensor '%s' is not registered", raw_name);
-            return false;
-        }
-        TensorState* state = state_it->second;
-        if (state == nullptr) {
-            LOG_ERROR("model manager tensor '%s' has no tensor state", raw_name);
+        auto param = resolve_param_tensor(tensor);
+        auto found = tensor_states_by_tensor_.find(param);
+        if (found == tensor_states_by_tensor_.end()) {
+            LOG_ERROR("model manager tensor '%s' is not registered", ggml_get_name(tensor));
             return false;
         }
+        TensorState* state = found->second;
         if ((compute_backend == nullptr || state->compute_backend == nullptr ||
              state->compute_backend == compute_backend) &&
             seen.insert(state).second) {
@@ -1375,8 +1424,8 @@ bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
     }
     if (!apply_loras_to_params(required_states)) {
         finish_compute_backend_usage(required_states);
-        release_compute_staging_blocks(false);
-        release_params_storage_blocks(false);
+        std::unordered_set<TensorState*> failed(required_states.begin(), required_states.end());
+        invalidate_sources(failed);
         return false;
     }
     return true;

src/model_manager.h

@@ -10,6 +10,7 @@
 #include <vector>
 
 #include "device_residency_manager.h"
+#include "model_component.h"
 #include "model_loader.h"
 
 class ModelManager : public DeviceResidencyManager {
@@ -24,7 +25,9 @@ public:
         float multiplier   = 1.0f;
         bool is_high_noise = false;
         std::string tensor_name_prefix_filter;
-        bool required = false;
+        bool required               = false;
+        ModelLoader::FileId file_id = 0;
+        uint64_t file_revision      = 0;
     };
 
 private:
@@ -32,8 +35,12 @@ private:
 
     struct TensorState {
         std::string name;
-        ggml_tensor* tensor = nullptr;
-        std::string desc;
+        ggml_tensor* tensor      = nullptr;
+        ModelComponent component = ModelComponent::Count;
+        TensorStorage source;
+        bool has_source                 = false;
+        ModelLoader::FileId source_file = 0;
+        SDVersion source_version        = VERSION_COUNT;
 
         ResidencyMode residency_mode                 = ResidencyMode::ParamBackend;
         ggml_backend_t compute_backend               = nullptr;
@@ -79,7 +86,7 @@ private:
 
     ModelLoader model_loader_;
     std::vector<std::unique_ptr<TensorState>> tensor_states_;
-    std::map<std::string, TensorState*> tensor_states_by_name_;
+    std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
     std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
     std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
     std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
@@ -91,6 +98,8 @@ private:
     bool warned_split_lora_skip_ = false;
     std::set<std::string> common_ignore_tensors_;
     std::vector<LoraSpec> loras_;
+    std::set<ModelLoader::FileId> lora_sources_;
+    bool applying_loras_             = false;
     SDVersion lora_version_          = VERSION_COUNT;
     uint64_t current_lora_epoch_     = 0;
     uint64_t residency_epoch_        = 0;
@@ -102,6 +111,7 @@ private:
 
     void finish_compute_backend_usage(const std::vector<TensorState*>& states);
     void release_all();
+    void invalidate_sources(const std::unordered_set<TensorState*>& states);
 
     ggml_backend_t prefetch_backend_for(ggml_backend_t compute_backend);
     bool populate_prefetch_block(PrefetchBlock& block);
@@ -152,15 +162,27 @@ private:
     void free_params_storage_block(ParamsStorageBlock& block);
     void erase_params_storage_block(ParamsStorageBlock* block);
     void reset_lora_applied_params();
+    bool unregister_tensor_states(const std::unordered_set<TensorState*>& states, size_t* size);
     size_t other_runtime_resident_bytes(uintptr_t owner_id,
                                         ggml_backend_t compute_backend) const;
 
 public:
     ~ModelManager() override;
 
-    ModelLoader& loader() { return model_loader_; }
     const ModelLoader& loader() const { return model_loader_; }
 
+    bool set_loader(ModelLoader loader);
+    bool add_file(const std::string& path, const std::string& prefix = "", ModelLoader::FileId* id = nullptr, bool force = false);
+    bool del_file(ModelLoader::FileId id);
+    bool refresh_files();
+    ModelLoader::FileVersions source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const;
+    size_t registered_params_size(const std::set<ModelComponent>& components) const;
+
+    void prepare_file_io() { model_loader_.process_model_files(enable_mmap_, writable_mmap_); }
+    bool load_float_tensor(const std::string& name, std::vector<float>& data) {
+        return model_loader_.load_float_tensor(name, data, n_threads_, enable_mmap_);
+    }
+
     void set_n_threads(int n_threads) {
         n_threads_ = n_threads;
         model_loader_.set_n_threads(n_threads);
@@ -172,14 +194,15 @@ public:
     void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; }
     void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
     void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
-    void set_loras(std::vector<LoraSpec> loras, SDVersion version);
+    bool prepare_lora_sources(std::vector<LoraSpec>& loras);
+    bool set_loras(std::vector<LoraSpec> loras, SDVersion version);
     void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits);
 
     static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
 
     std::set<std::string> tensor_names() const;
 
-    bool register_param_tensors(const std::string& desc,
+    bool register_param_tensors(ModelComponent component,
                                 std::map<std::string, ggml_tensor*> tensors,
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
@@ -187,13 +210,18 @@ public:
                                 size_t* registered_tensor_size                         = nullptr,
                                 bool allow_split_buffer                                = false,
                                 bool params_follow_compute_backend                     = false,
-                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
+                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr,
+                                ModelLoader::FileId source_file                        = 0,
+                                SDVersion source_version                               = VERSION_COUNT);
 
-    bool unregister_param_tensors(const std::string& desc,
+    ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const override;
+    bool unregister_param_tensors(const std::vector<ggml_tensor*>& tensors);
+
+    bool unregister_param_tensors(ModelComponent component,
                                   size_t* registered_tensor_size = nullptr);
 
     template <typename Runner>
-    bool register_runner_params(const std::string& desc,
+    bool register_runner_params(ModelComponent component,
                                 Runner& runner,
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
@@ -201,7 +229,7 @@ public:
                                 size_t* registered_tensor_size = nullptr) {
         std::map<std::string, ggml_tensor*> tensors;
         runner.get_param_tensors(tensors);
-        return register_param_tensors(desc,
+        return register_param_tensors(component,
                                       std::move(tensors),
                                       residency_mode,
                                       compute_backend,
@@ -210,7 +238,7 @@ public:
     }
 
     template <typename Runner>
-    bool register_runner_params(const std::string& desc,
+    bool register_runner_params(ModelComponent component,
                                 Runner& runner,
                                 const std::string& prefix,
                                 ResidencyMode residency_mode,
@@ -219,7 +247,7 @@ public:
                                 size_t* registered_tensor_size = nullptr) {
         std::map<std::string, ggml_tensor*> tensors;
         runner.get_param_tensors(tensors, prefix);
-        return register_param_tensors(desc,
+        return register_param_tensors(component,
                                       std::move(tensors),
                                       residency_mode,
                                       compute_backend,

src/model_manager_files.cpp

@@ -0,0 +1,154 @@
+#include "model_manager.h"
+
+#include <algorithm>
+
+#include "core/util.h"
+
+static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
+    return a.file_id == b.file_id && a.file_revision == b.file_revision &&
+           a.file_index == b.file_index && a.offset == b.offset && a.index_in_zip == b.index_in_zip &&
+           a.storage_key == b.storage_key && a.type == b.type && a.expected_type == b.expected_type &&
+           a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) &&
+           a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 &&
+           a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 &&
+           a.is_int8_tensorwise == b.is_int8_tensorwise && a.int8_convrot == b.int8_convrot &&
+           a.int8_convrot_group_size == b.int8_convrot_group_size;
+}
+
+void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
+    auto affected = states;
+    for (const auto& block : params_storage_blocks_) {
+        if (std::any_of(block->states.begin(), block->states.end(), [&](TensorState* state) { return states.count(state) != 0; })) {
+            affected.insert(block->states.begin(), block->states.end());
+        }
+    }
+    for (auto it = prefetch_blocks_.begin(); it != prefetch_blocks_.end();) {
+        if (std::any_of(it->second->states.begin(), it->second->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
+            free_prefetch_block(*it->second);
+            it = prefetch_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) {
+        if (std::any_of((*it)->staged_tensors.begin(), (*it)->staged_tensors.end(), [&](const auto& entry) { return affected.count(entry.first) != 0; })) {
+            ggml_backend_synchronize((*it)->compute_backend);
+            free_compute_staging_block(**it);
+            it = compute_staging_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) {
+        if (std::any_of((*it)->states.begin(), (*it)->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
+            free_params_storage_block(**it);
+            it = params_storage_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    for (auto* state : affected) {
+        state->metadata_validated = false;
+        state->applied_lora_epoch = UINT64_MAX;
+    }
+}
+
+bool ModelManager::set_loader(ModelLoader loader) {
+    if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
+            return state->pin_count != 0;
+        })) {
+        LOG_ERROR("cannot update model sources during execution");
+        return false;
+    }
+    std::map<std::pair<ModelLoader::FileId, SDVersion>, String2TensorStorage> scoped;
+    auto sources_for = [&](const TensorState& state) -> const String2TensorStorage& {
+        if (state.source_file == 0)
+            return loader.get_tensor_storage_map();
+        auto key   = std::make_pair(state.source_file, state.source_version);
+        auto found = scoped.find(key);
+        if (found == scoped.end())
+            found = scoped.emplace(key, loader.file_tensors(key.first, key.second)).first;
+        return found->second;
+    };
+    bool lora_changed = false;
+    for (const auto& spec : loras_) {
+        lora_changed |= loader.file_revision(spec.file_id) != spec.file_revision;
+    }
+    std::unordered_set<TensorState*> changed;
+    for (const auto& state : tensor_states_) {
+        const auto& sources = sources_for(*state);
+        auto source         = sources.find(state->name);
+        const bool found    = source != sources.end();
+        if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
+            (lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
+            changed.insert(state.get());
+        }
+    }
+    invalidate_sources(changed);
+    for (auto* state : changed) {
+        const auto& sources = sources_for(*state);
+        auto source         = sources.find(state->name);
+        state->has_source   = source != sources.end();
+        state->source       = state->has_source ? source->second : TensorStorage{};
+    }
+    if (lora_changed) {
+        ++current_lora_epoch_;
+        for (auto& spec : loras_)
+            spec.file_revision = loader.file_revision(spec.file_id);
+    }
+    model_loader_ = std::move(loader);
+    model_loader_.set_n_threads(n_threads_);
+    return true;
+}
+
+bool ModelManager::add_file(const std::string& path, const std::string& prefix, ModelLoader::FileId* id, bool force) {
+    ModelLoader candidate = model_loader_;
+    ModelLoader::FileId added_id;
+    if (!candidate.add_file(path, prefix, &added_id, force) || !set_loader(std::move(candidate))) {
+        return false;
+    }
+    if (id != nullptr) {
+        *id = added_id;
+    }
+    return true;
+}
+
+bool ModelManager::del_file(ModelLoader::FileId id) {
+    ModelLoader candidate = model_loader_;
+    return candidate.del_file(id) && set_loader(std::move(candidate));
+}
+
+bool ModelManager::refresh_files() {
+    ModelLoader candidate = model_loader_;
+    return candidate.refresh_files() && set_loader(std::move(candidate));
+}
+
+ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const {
+    ModelLoader::FileVersions versions;
+    const auto& sources = loader.get_tensor_storage_map();
+    for (const auto& state : tensor_states_) {
+        if (components.count(state->component) == 0) {
+            continue;
+        }
+        if (state->source_file != 0) {
+            versions[state->source_file] = loader.file_revision(state->source_file);
+            continue;
+        }
+        auto source = sources.find(state->name);
+        if (source != sources.end()) {
+            versions[source->second.file_id] = source->second.file_revision;
+        }
+    }
+    return versions;
+}
+
+size_t ModelManager::registered_params_size(const std::set<ModelComponent>& components) const {
+    size_t bytes = 0;
+    std::unordered_set<const ggml_tensor*> seen;
+    for (const auto& state : tensor_states_) {
+        if (components.count(state->component) != 0 && state->tensor != nullptr && seen.insert(state->tensor).second) {
+            bytes += ggml_nbytes(state->tensor);
+        }
+    }
+    return bytes;
+}

src/pipeline/diffusion_engine.h

@@ -0,0 +1,482 @@
+#ifndef __SD_PIPELINE_DIFFUSION_ENGINE_H__
+#define __SD_PIPELINE_DIFFUSION_ENGINE_H__
+
+#include <atomic>
+#include <cmath>
+#include <functional>
+#include <list>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_graph_cut.h"
+#include "core/tensor.hpp"
+#include "core/util.h"
+#include "model/adapter/lora.hpp"
+#include "model_builders.h"
+#include "model_manager.h"
+#include "stable-diffusion.h"
+
+class RNG;
+struct Denoiser;
+struct LoraModel;
+struct ConditionerParams;
+struct SDCondition;
+struct RefImageParams;
+
+extern const char* model_version_to_str[];
+
+static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
+    return version == VERSION_FLUX ||
+           sd_version_is_flux2(version) ||
+           sd_version_is_qwen_image(version) ||
+           sd_version_is_mage_flow(version) ||
+           sd_version_is_longcat(version) ||
+           sd_version_is_z_image(version) ||
+           sd_version_is_boogu_image(version);
+}
+
+class StableDiffusionGGML {
+public:
+    SDBackendManager backend_manager;
+
+    SDVersion version;
+    bool external_vae_is_invalid = false;
+
+    bool circular_x = false;
+    bool circular_y = false;
+
+    std::shared_ptr<RNG> rng;
+    std::shared_ptr<RNG> sampler_rng = nullptr;
+    int n_threads                    = -1;
+    float default_flow_shift         = INFINITY;
+    float active_flow_shift          = INFINITY;
+
+    std::shared_ptr<Conditioner> cond_stage_model;
+    std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;  // for svd or wan2.1 i2v
+    std::shared_ptr<DiffusionModelRunner> diffusion_model;
+    std::shared_ptr<DiffusionModelRunner> high_noise_diffusion_model;
+    std::shared_ptr<VAE> first_stage_model;
+    std::shared_ptr<VAE> preview_vae;
+    std::shared_ptr<AudioVAERunner> audio_vae_model;
+    std::shared_ptr<ControlNet> control_net;
+    std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
+    sd::Tensor<float> ip_adapter_tokens;
+    sd::Tensor<float> ip_adapter_uncond_tokens;
+    float ip_adapter_strength = 1.0f;
+    std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
+    struct RuntimeLora {
+        ModelManager::LoraSpec spec;
+        SDBackendModule module;
+        std::shared_ptr<LoraModel> model;
+
+        bool matches(const ModelManager::LoraSpec& other) const {
+            return spec.file_id == other.file_id && spec.file_revision == other.file_revision &&
+                   spec.tensor_name_prefix_filter == other.tensor_name_prefix_filter;
+        }
+    };
+    std::vector<RuntimeLora> runtime_lora_models;
+    bool apply_lora_immediately = false;
+    int animatediff_num_frames  = 0;
+
+    std::string taesd_path;
+    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
+    bool enable_mmap                     = false;
+    sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
+    bool disable_prefetch          = false;
+    bool disable_segmented_compute = false;
+    bool eager_load                = false;
+    std::string backend_spec;
+    std::string params_backend_spec;
+    std::string split_mode_spec;
+    bool auto_fit_enabled = false;
+
+    bool diffusion_conv_direct = false;
+
+    bool is_using_v_parameterization     = false;
+    bool is_using_edm_v_parameterization = false;
+
+    std::shared_ptr<ModelManager> model_manager;
+
+    enum class RunnerGroup { Core,
+                             VAE,
+                             ControlNet,
+                             Extensions };
+    using RunnerGroups = std::set<RunnerGroup>;
+
+    struct ModelConfig {
+        sd_ctx_params_t params{};
+        std::list<std::string> strings;
+        std::vector<sd_embedding_t> embeddings;
+        ModelLoader::FileId control_net_file = 0;
+        bool use_tae                         = false;
+        bool use_audio_vae                   = false;
+        bool photomaker_source_available     = false;
+        bool animatediff_loaded              = false;
+
+        explicit ModelConfig(const sd_ctx_params_t& initial)
+            : params(initial) {
+            for (auto member : {&sd_ctx_params_t::model_path, &sd_ctx_params_t::clip_l_path,
+                                &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path,
+                                &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path,
+                                &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path,
+                                &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path,
+                                &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path,
+                                &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path,
+                                &sd_ctx_params_t::control_net_path, &sd_ctx_params_t::ip_adapter_path,
+                                &sd_ctx_params_t::motion_module_path, &sd_ctx_params_t::photo_maker_path,
+                                &sd_ctx_params_t::pulid_weights_path, &sd_ctx_params_t::tensor_type_rules,
+                                &sd_ctx_params_t::max_vram, &sd_ctx_params_t::backend,
+                                &sd_ctx_params_t::params_backend, &sd_ctx_params_t::split_mode,
+                                &sd_ctx_params_t::rpc_servers, &sd_ctx_params_t::model_args}) {
+                strings.emplace_back(SAFE_STR(initial.*member));
+                params.*member = strings.back().c_str();
+            }
+            for (uint32_t i = 0; i < initial.embedding_count; ++i) {
+                strings.emplace_back(SAFE_STR(initial.embeddings[i].name));
+                const char* name = strings.back().c_str();
+                strings.emplace_back(SAFE_STR(initial.embeddings[i].path));
+                embeddings.push_back({name, strings.back().c_str()});
+            }
+            params.embeddings = embeddings.data();
+        }
+
+        ModelConfig(const ModelConfig& other)
+            : ModelConfig(other.params) {
+            control_net_file            = other.control_net_file;
+            use_tae                     = other.use_tae;
+            use_audio_vae               = other.use_audio_vae;
+            photomaker_source_available = other.photomaker_source_available;
+            animatediff_loaded          = other.animatediff_loaded;
+        }
+        ModelConfig& operator=(const ModelConfig&) = delete;
+
+        void set_control_net(ModelLoader::FileId id, const std::string& path) {
+            control_net_file = id;
+            strings.push_back(path);
+            params.control_net_path = strings.back().c_str();
+        }
+    };
+
+    struct RunnerState {
+        bool ready                = false;
+        uint64_t catalog_revision = 0;
+        std::map<RunnerGroup, ModelLoader::FileVersions> sources;
+    };
+
+    std::recursive_mutex execution_mutex;
+    std::unique_ptr<ModelConfig> config_;
+    RunnerState runner_state_;
+    bool executing_ = false;
+
+    std::shared_ptr<Denoiser> denoiser;
+    std::vector<float> file_alphas_cumprod;
+
+    StableDiffusionGGML();
+    ~StableDiffusionGGML();
+
+    static const std::map<RunnerGroup, std::set<ModelComponent>>& runner_components();
+
+    static RunnerGroups all_runner_groups();
+
+    ModelLoader::FileVersions runner_source_versions(RunnerGroup group, const ModelLoader& loader) const;
+
+    void capture_runner_sources();
+
+    void end_runners();
+
+    bool reset_runners(const RunnerGroups& groups);
+
+    bool refresh_model_sources();
+
+    bool apply_model_update(ModelLoader candidate,
+                            std::unique_ptr<ModelConfig> next_config = nullptr,
+                            RunnerGroups groups                      = {});
+
+    struct ContextOperation {
+        StableDiffusionGGML& sd;
+        std::unique_lock<std::recursive_mutex> lock;
+        bool acquired = false;
+
+        explicit ContextOperation(StableDiffusionGGML& sd)
+            : sd(sd), lock(sd.execution_mutex, std::try_to_lock) {
+            if (!lock.owns_lock() || sd.executing_) {
+                // The caller may be a log callback, so rejecting it must not log.
+                return;
+            }
+            sd.executing_ = true;
+            acquired      = true;
+        }
+
+        ~ContextOperation() {
+            if (acquired) {
+                sd.executing_ = false;
+            }
+        }
+    };
+
+    struct ExecutionScope {
+        ContextOperation operation;
+        bool ready = false;
+
+        explicit ExecutionScope(StableDiffusionGGML& sd)
+            : operation(sd) {
+            ready = operation.acquired && sd.refresh_model_sources();
+        }
+
+        ~ExecutionScope() {
+            if (ready) {
+                operation.sd.end_runners();
+            }
+        }
+    };
+
+    ggml_backend_t backend_for(SDBackendModule module);
+
+    ggml_backend_t params_backend_for(SDBackendModule module);
+
+    std::atomic<sd_cancel_mode_t> cancellation_flag = SD_CANCEL_RESET;
+
+    void set_cancel_flag(enum sd_cancel_mode_t flag);
+
+    void reset_cancel_flag();
+
+    enum sd_cancel_mode_t get_cancel_flag();
+
+    size_t max_graph_vram_bytes_for_module(SDBackendModule module);
+
+    std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends);
+
+    bool ensure_backend_pair(SDBackendModule module);
+
+    template <typename T>
+    bool register_runner_params(ModelComponent component,
+                                const std::shared_ptr<T>& model,
+                                SDBackendModule module,
+                                size_t* params_mem_size = nullptr);
+
+    template <typename T>
+    bool register_row_split_runner_params(ModelComponent component,
+                                          const std::shared_ptr<T>& model,
+                                          SDBackendModule module,
+                                          const std::vector<ggml_backend_t>& module_backends,
+                                          std::map<std::string, ggml_tensor*> group_tensors,
+                                          const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
+                                          ModelManager::ResidencyMode residency_mode,
+                                          size_t* params_mem_size);
+
+    // Register graph-cut layer-split tensors on the primary backend first.
+    // The first real graph assigns each param tensor to a runtime backend
+    // before weights are loaded or staged.
+    template <typename T>
+    bool register_layer_split_runner_params(ModelComponent component,
+                                            const std::shared_ptr<T>& model,
+                                            SDBackendModule module,
+                                            const std::vector<ggml_backend_t>& module_backends,
+                                            std::map<std::string, ggml_tensor*> group_tensors,
+                                            const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
+                                            ModelManager::ResidencyMode residency_mode,
+                                            size_t* params_mem_size);
+
+    bool unload_control_net();
+
+    bool load_control_net_from_file(const std::string& path);
+
+    void apply_circular_axes(bool circular_x, bool circular_y);
+
+    bool init_backend();
+
+    bool row_split_active();
+
+    bool graph_cut_layer_split_active();
+
+    std::shared_ptr<RNG> get_rng(rng_type_t rng_type);
+
+    void refresh_compvis_denoiser_sigmas();
+
+    void load_alphas_cumprod();
+
+    bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration);
+
+    bool init(const sd_ctx_params_t* sd_ctx_params);
+
+    bool uses_tae() const;
+
+    bool tae_preview_only() const;
+
+    void configure_weight_loading();
+
+    sd::model_builders::Context model_build_context();
+
+    bool build_core_runners();
+
+    bool build_vae_runners();
+
+    bool build_control_net_runner();
+
+    bool build_extension_runners();
+
+    bool validate_and_load_runners();
+
+    bool build_denoiser();
+
+    bool build_runners(const RunnerGroups& groups);
+
+    bool is_using_v_parameterization_for_sd2(bool is_inpaint = false);
+
+    static std::string lora_log_id(const ModelManager::LoraSpec& lora);
+
+    std::shared_ptr<LoraModel> load_lora_model(const ModelManager::LoraSpec& lora_spec,
+                                               SDBackendModule module,
+                                               LoraModel::filter_t module_filter = nullptr);
+
+    void clear_lora_adapters();
+
+    std::vector<std::shared_ptr<LoraModel>> load_runtime_loras_for_module(const std::vector<ModelManager::LoraSpec>& loras,
+                                                                          const std::set<std::string>& model_tensor_names,
+                                                                          SDBackendModule module,
+                                                                          LoraModel::filter_t module_filter,
+                                                                          bool& success,
+                                                                          std::vector<RuntimeLora>& next_models);
+
+    bool apply_loras_immediately(const std::vector<ModelManager::LoraSpec>& loras);
+
+    bool apply_loras_at_runtime(const std::vector<ModelManager::LoraSpec>& loras);
+
+    void lora_stat();
+
+    bool apply_loras(const sd_lora_t* loras, uint32_t lora_count);
+
+    void reset_generation_extensions();
+
+    void prepare_generation_extensions(const sd_pm_params_t& pm_params,
+                                       const sd_pulid_params_t& pulid_params,
+                                       ConditionerParams& condition_params,
+                                       int total_steps);
+
+    sd::Tensor<float> get_clip_vision_output(const sd::Tensor<float>& image,
+                                             bool return_pooled   = true,
+                                             int clip_skip        = -1,
+                                             bool zero_out_masked = false);
+
+    void compute_ip_adapter_tokens(const sd_image_t& image, float strength);
+
+    std::vector<float> process_timesteps(const std::vector<float>& timesteps,
+                                         const sd::Tensor<float>& init_latent,
+                                         const sd::Tensor<float>& denoise_mask,
+                                         int step);
+
+    std::vector<float> process_ltxav_video_timesteps(const std::vector<float>& timesteps,
+                                                     const sd::Tensor<float>& init_latent,
+                                                     const sd::Tensor<float>& denoise_mask);
+
+    void preview_image(int step,
+                       const sd::Tensor<float>& latents,
+                       enum SDVersion version,
+                       preview_t preview_mode,
+                       std::function<void(int, int, sd_image_t*, bool, void*)> step_callback,
+                       void* step_callback_data,
+                       bool is_noisy);
+
+    std::vector<float> prepare_sample_timesteps(float sigma,
+                                                int shifted_timestep);
+
+    void adjust_sample_step_scalings(int shifted_timestep,
+                                     const std::vector<float>& timesteps_vec,
+                                     float c_in,
+                                     float* c_skip,
+                                     float* c_out);
+
+    struct SamplePreviewContext {
+        sd_preview_cb_t callback = nullptr;
+        void* data               = nullptr;
+        preview_t mode           = PREVIEW_NONE;
+    };
+
+    SamplePreviewContext prepare_sample_preview_context();
+
+    void report_sample_progress(int step,
+                                size_t total_steps,
+                                bool terminal_sigma_is_zero,
+                                int64_t* last_progress_us);
+
+    void compute_sample_controls(const sd::Tensor<float>& control_image,
+                                 const sd::Tensor<float>& noised_input,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const SDCondition& condition,
+                                 std::vector<sd::Tensor<float>>* controls);
+
+    sd::Tensor<float> sample(const std::shared_ptr<DiffusionModelRunner>& work_diffusion_model,
+                             bool inverse_noise_scaling,
+                             const sd::Tensor<float>& init_latent,
+                             sd::Tensor<float> noise,
+                             const SDCondition& cond,
+                             const SDCondition& uncond,
+                             const SDCondition& img_uncond,
+                             const sd::Tensor<float>& control_image,
+                             float control_strength,
+                             const sd_guidance_params_t& guidance,
+                             float eta,
+                             int shifted_timestep,
+                             sample_method_t method,
+                             bool is_flow_denoiser,
+                             const char* extra_sample_args,
+                             const std::vector<float>& sigmas,
+                             const std::vector<sd::Tensor<float>>& ref_latents,
+                             const RefImageParams& ref_image_params,
+                             const sd::Tensor<float>& denoise_mask,
+                             const sd::Tensor<float>& vace_context,
+                             float vace_strength,
+                             int audio_length,
+                             float frame_rate,
+                             const sd_cache_params_t* cache_params,
+                             bool preview_final_step,
+                             const sd::Tensor<float>& video_positions = {});
+
+    int get_vae_scale_factor();
+
+    int get_diffusion_model_down_factor();
+
+    int get_latent_channel();
+
+    int get_image_channels() const;
+
+    int get_image_seq_len(int h, int w);
+
+    sd::Tensor<float> generate_init_latent(int width,
+                                           int height,
+                                           int frames = 1,
+                                           bool video = false);
+
+    int video_frames_to_latent_frames(int frames);
+
+    int latent_frames_to_video_frames(int latent_frames);
+
+    int align_video_frames(int frames);
+
+    sd::Tensor<float> encode_to_vae_latents(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> encode_first_stage(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false);
+
+    sd::Tensor<float> normalize_ltx_video_latents(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> un_normalize_ltx_video_latents(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> decode_ltx_audio_latent(const sd::Tensor<float>& audio_latent);
+
+    void set_flow_shift(float flow_shift = INFINITY);
+
+    bool is_flow_denoiser();
+
+    std::string get_default_ref_image_preset(SDVersion version) const;
+
+    RefImageParams resolve_ref_image_params(const char* ref_image_args) const;
+};
+
+#endif  // __SD_PIPELINE_DIFFUSION_ENGINE_H__

src/pipeline/generation.h

@@ -0,0 +1,73 @@
+#ifndef __SD_PIPELINE_GENERATION_H__
+#define __SD_PIPELINE_GENERATION_H__
+
+#include "conditioning/conditioner.hpp"
+#include "stable-diffusion.h"
+
+class StableDiffusionGGML;
+
+static inline bool sd_version_supports_animatediff(SDVersion version) {
+    return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX;
+}
+
+namespace sd::pipeline {
+
+    struct ImageGenerationLatents {
+        sd::Tensor<float> init_latent;
+        sd::Tensor<float> concat_latent;
+        sd::Tensor<float> img_uncond_concat_latent;
+        sd::Tensor<float> audio_latent;
+        sd::Tensor<float> video_positions;
+        sd::Tensor<float> control_image;
+        std::vector<sd::Tensor<float>> ref_images;
+        std::vector<sd::Tensor<float>> ref_latents;
+        std::vector<sd::Tensor<float>> reference_audio_latents;
+        std::vector<MiniMaxH3ReferenceBlock> minimax_reference_blocks;
+        std::vector<MiniMaxH3PresentationItem> minimax_presentation_refs;
+        std::vector<int32_t> keyframe_indices;
+        sd::Tensor<float> denoise_mask;
+        sd::Tensor<float> clip_vision_output;
+        sd::Tensor<float> vace_context;
+        int64_t ref_image_num                  = 0;
+        int64_t video_conditioning_frame_count = 0;
+        int64_t video_target_frame_count       = 0;
+        int audio_length                       = 0;
+    };
+
+    struct ImageGenerationEmbeds {
+        SDCondition cond;
+        SDCondition uncond;
+        SDCondition img_uncond;
+    };
+
+    struct ConditionerRunnerEndOnExit {
+        Conditioner* conditioner = nullptr;
+        ~ConditionerRunnerEndOnExit() {
+            if (conditioner != nullptr) {
+                conditioner->runner_end();
+            }
+        }
+    };
+
+    // Callers hold ExecutionScope; AnimateDiff reuses the image path within the same scope.
+    bool generate_image(StableDiffusionGGML* sd,
+                        const sd_img_gen_params_t* sd_img_gen_params,
+                        sd_image_t** images_out,
+                        int* num_images_out);
+
+    bool generate_video(StableDiffusionGGML* sd,
+                        const sd_vid_gen_params_t* sd_vid_gen_params,
+                        sd_image_t** frames_out,
+                        int* num_frames_out,
+                        sd_audio_t** audio_out);
+
+    sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
+                                                       const char* model_path,
+                                                       const sd::Tensor<float>& packed_latent,
+                                                       int audio_length);
+
+    sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels);
+
+}  // namespace sd::pipeline
+
+#endif  // __SD_PIPELINE_GENERATION_H__

src/pipeline/model_builders.cpp

@@ -0,0 +1,574 @@
+#include "model_builders.h"
+
+#include <cstring>
+#include <utility>
+
+#include "conditioning/conditioner.hpp"
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+#include "extensions/generation_extension.h"
+#include "model/adapter/ip_adapter.hpp"
+#include "model/diffusion/anima.hpp"
+#include "model/diffusion/boogu.hpp"
+#include "model/diffusion/control.hpp"
+#include "model/diffusion/ernie_image.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/hidream_o1.hpp"
+#include "model/diffusion/hunyuan.hpp"
+#include "model/diffusion/ideogram4.hpp"
+#include "model/diffusion/krea2.hpp"
+#include "model/diffusion/lens.hpp"
+#include "model/diffusion/lingbot_video.hpp"
+#include "model/diffusion/ltxv.hpp"
+#include "model/diffusion/mage_flow.hpp"
+#include "model/diffusion/minimax_h3.hpp"
+#include "model/diffusion/minit2i.hpp"
+#include "model/diffusion/mmdit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/pid.hpp"
+#include "model/diffusion/qwen_image.hpp"
+#include "model/diffusion/unet.hpp"
+#include "model/diffusion/wan.hpp"
+#include "model/diffusion/z_image.hpp"
+#include "model/vae/auto_encoder_kl.hpp"
+#include "model/vae/hunyuan_vae.hpp"
+#include "model/vae/ltx_audio_vae.hpp"
+#include "model/vae/ltx_vae.hpp"
+#include "model/vae/mage_vae.hpp"
+#include "model/vae/minimax_h3_audio_vae.hpp"
+#include "model/vae/minimax_h3_vae.hpp"
+#include "model/vae/tae.hpp"
+#include "model/vae/vae.hpp"
+#include "model/vae/wan_vae.hpp"
+
+namespace sd::model_builders {
+
+    static bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) {
+        if (backends.runtime_backend(module) == nullptr) {
+            LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
+            return false;
+        }
+        if (backends.params_backend(module) == nullptr) {
+            LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
+            return false;
+        }
+        return true;
+    }
+
+    static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback) {
+        switch (format) {
+            case SD_VAE_FORMAT_FLUX:
+                return VERSION_FLUX;
+            case SD_VAE_FORMAT_SD3:
+                return VERSION_SD3;
+            case SD_VAE_FORMAT_FLUX2:
+                return VERSION_FLUX2;
+            case SD_VAE_FORMAT_WAN:
+                return VERSION_WAN2;
+            case SD_VAE_FORMAT_AUTO:
+            default:
+                return fallback;
+        }
+    }
+
+    bool build_core_runners(const Context& ctx, CoreRunners& runners) {
+        const auto* sd_ctx_params      = &ctx.params;
+        const auto& tensor_storage_map = ctx.tensor_storage_map;
+        const auto version             = ctx.version;
+        const auto& weight_manager     = ctx.weight_manager;
+        CoreRunners result;
+        if (!ensure_backend_pair(ctx.backends, SDBackendModule::TE) ||
+            !ensure_backend_pair(ctx.backends, SDBackendModule::DIFFUSION)) {
+            return false;
+        }
+
+        if (sd_version_is_sd3(version)) {
+            result.conditioner = std::make_shared<SD3CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                   tensor_storage_map,
+                                                                   weight_manager);
+            result.diffusion   = std::make_shared<MMDiTRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                             tensor_storage_map,
+                                                             "model.diffusion_model",
+                                                             weight_manager);
+        } else if (sd_version_is_pid(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Pid::PiDRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                tensor_storage_map,
+                                                                "model.diffusion_model.net",
+                                                                weight_manager);
+        } else if (sd_version_is_ideogram4(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Ideogram4::Ideogram4Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model.diffusion_model",
+                                                                            weight_manager);
+        } else if (sd_version_is_krea2(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Krea2::Krea2Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                    tensor_storage_map,
+                                                                    "model.diffusion_model",
+                                                                    weight_manager);
+        } else if (sd_version_is_flux(version)) {
+            bool is_chroma = false;
+            for (auto pair : tensor_storage_map) {
+                if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
+                    is_chroma = true;
+                    break;
+                }
+            }
+            if (is_chroma) {
+                result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                      tensor_storage_map,
+                                                                      false,
+                                                                      1,
+                                                                      false,
+                                                                      weight_manager,
+                                                                      sd_ctx_params->model_args);
+            } else if (version == VERSION_OVIS_IMAGE) {
+                result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                   tensor_storage_map,
+                                                                   version,
+                                                                   "",
+                                                                   false,
+                                                                   weight_manager);
+            } else {
+                result.conditioner = std::make_shared<FluxCLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                        tensor_storage_map,
+                                                                        weight_manager);
+            }
+            result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map,
+                                                                  "model.diffusion_model",
+                                                                  version,
+                                                                  weight_manager,
+                                                                  sd_ctx_params->model_args);
+        } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
+            bool is_chroma     = false;
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map,
+                                                                  "model.diffusion_model",
+                                                                  version,
+                                                                  weight_manager,
+                                                                  sd_ctx_params->model_args);
+        } else if (sd_version_is_ltxav(version)) {
+            result.conditioner = std::make_shared<LTXAVEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 "text_encoders.llm",
+                                                                 "text_embedding_projection",
+                                                                 weight_manager);
+            result.diffusion   = std::make_shared<LTXV::LTXAVRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                   tensor_storage_map,
+                                                                   "model.diffusion_model",
+                                                                   weight_manager);
+        } else if (sd_version_is_minimax_h3(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model.diffusion_model",
+                                                                            weight_manager);
+        } else if (sd_version_is_hunyuan_video(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Hunyuan::HunyuanVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "model.diffusion_model",
+                                                                             version,
+                                                                             weight_manager);
+        } else if (sd_version_is_wan(version)) {
+            result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                  tensor_storage_map,
+                                                                  true,
+                                                                  0,
+                                                                  true,
+                                                                  weight_manager);
+            result.diffusion   = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                tensor_storage_map,
+                                                                "model.diffusion_model",
+                                                                version,
+                                                                weight_manager);
+            if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {
+                result.high_noise_diffusion = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                               tensor_storage_map,
+                                                                               "model.high_noise_diffusion_model",
+                                                                               version,
+                                                                               weight_manager);
+            }
+            if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" ||
+                result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" ||
+                result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") {
+                if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
+                    return false;
+                }
+                result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
+                                                                                tensor_storage_map,
+                                                                                weight_manager);
+            }
+        } else if (sd_version_is_lingbot_video(version)) {
+            bool enable_vision = false;
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (starts_with(name, "text_encoders.llm.visual.")) {
+                    enable_vision = true;
+                    break;
+                }
+            }
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               enable_vision,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<LingBotVideo::LingBotVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                                  tensor_storage_map,
+                                                                                  "model.diffusion_model",
+                                                                                  weight_manager,
+                                                                                  sd_ctx_params->model_args);
+        } else if (sd_version_is_qwen_image(version)) {
+            bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               enable_vision,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                       tensor_storage_map,
+                                                                       "model.diffusion_model",
+                                                                       version,
+                                                                       weight_manager,
+                                                                       sd_ctx_params->model_args);
+        } else if (sd_version_is_mage_flow(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<MageFlow::MageFlowRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                          tensor_storage_map,
+                                                                          "model.diffusion_model",
+                                                                          weight_manager);
+        } else if (sd_version_is_longcat(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map,
+                                                                  "model.diffusion_model",
+                                                                  version,
+                                                                  weight_manager,
+                                                                  sd_ctx_params->model_args);
+        } else if (version == VERSION_HIDREAM_O1) {
+            result.conditioner = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                                   tensor_storage_map,
+                                                                                   weight_manager);
+            result.diffusion   = std::make_shared<HiDreamO1::HiDreamO1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model",
+                                                                            weight_manager);
+        } else if (sd_version_is_minit2i(version)) {
+            result.conditioner = std::make_shared<MiniT2IConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                      tensor_storage_map,
+                                                                      weight_manager);
+            result.diffusion   = std::make_shared<MiniT2I::MiniT2IRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                        tensor_storage_map,
+                                                                        "model.diffusion_model.model.net",
+                                                                        weight_manager);
+        } else if (sd_version_is_anima(version)) {
+            result.conditioner = std::make_shared<AnimaConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                    tensor_storage_map,
+                                                                    weight_manager);
+            result.diffusion   = std::make_shared<Anima::AnimaRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                    tensor_storage_map,
+                                                                    "model.diffusion_model",
+                                                                    weight_manager);
+        } else if (sd_version_is_z_image(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<ZImage::ZImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                      tensor_storage_map,
+                                                                      "model.diffusion_model",
+                                                                      version,
+                                                                      weight_manager);
+        } else if (sd_version_is_boogu_image(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Boogu::BooguImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                         tensor_storage_map,
+                                                                         "model.diffusion_model",
+                                                                         version,
+                                                                         weight_manager);
+        } else if (sd_version_is_ernie_image(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<ErnieImage::ErnieImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                              tensor_storage_map,
+                                                                              "model.diffusion_model",
+                                                                              weight_manager);
+        } else if (sd_version_is_lens(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager);
+            result.diffusion   = std::make_shared<Lens::LensRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map,
+                                                                  "model.diffusion_model",
+                                                                  weight_manager);
+        } else {  // SD1.x SD2.x SDXL
+            std::map<std::string, std::string> embbeding_map;
+            for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
+                embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path));
+            }
+            result.conditioner = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                                     tensor_storage_map,
+                                                                                     embbeding_map,
+                                                                                     version,
+                                                                                     weight_manager);
+            result.diffusion   = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                 tensor_storage_map,
+                                                                 "model.diffusion_model",
+                                                                 version,
+                                                                 weight_manager);
+            if (sd_ctx_params->diffusion_conv_direct) {
+                LOG_INFO("Using Conv2d direct in the diffusion model");
+                result.diffusion->set_conv2d_direct_enabled(true);
+            }
+        }
+
+        if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) {
+            if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
+                return false;
+            }
+            result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
+                                                                            tensor_storage_map,
+                                                                            weight_manager);
+        }
+
+        if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
+            result.ip_adapter = std::make_shared<IPAdapter::IPAdapterRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "ip_adapter",
+                                                                             weight_manager);
+        }
+        runners = std::move(result);
+        return true;
+    }
+
+    bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) {
+        const auto* sd_ctx_params      = &ctx.params;
+        const auto& tensor_storage_map = ctx.tensor_storage_map;
+        const auto version             = ctx.version;
+        const auto& weight_manager     = ctx.weight_manager;
+        VAERunners result;
+        if (!ensure_backend_pair(ctx.backends, SDBackendModule::VAE)) {
+            return false;
+        }
+
+        auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
+            if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
+                return std::make_shared<TinyVideoAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                              tensor_storage_map,
+                                                              "decoder",
+                                                              decode_only,
+                                                              version,
+                                                              weight_manager);
+
+            } else {
+                auto model = std::make_shared<TinyImageAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                    tensor_storage_map,
+                                                                    "decoder.layers",
+                                                                    decode_only,
+                                                                    version,
+                                                                    weight_manager);
+                return model;
+            }
+        };
+
+        sd_vae_format_t vae_format = sd_ctx_params->vae_format;
+        if (vae_format < SD_VAE_FORMAT_AUTO || vae_format >= SD_VAE_FORMAT_COUNT) {
+            LOG_WARN("invalid VAE format override, using auto");
+            vae_format = SD_VAE_FORMAT_AUTO;
+        }
+        SDVersion vae_version = version;
+        if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) {
+            vae_version = sd_vae_format_to_version(vae_format, vae_version);
+        }
+
+        auto create_vae = [&]() -> std::shared_ptr<VAE> {
+            if (sd_version_is_ltxav(version)) {
+                return std::make_shared<LTXVideoVAE>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                     tensor_storage_map,
+                                                     "first_stage_model",
+                                                     false,
+                                                     version,
+                                                     weight_manager);
+            } else if (sd_version_is_minimax_h3(version)) {
+                return std::make_shared<MiniMaxH3VAE::MiniMaxH3VideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                               tensor_storage_map,
+                                                                               "first_stage_model",
+                                                                               weight_manager);
+            } else if (sd_version_is_mage_flow(vae_version)) {
+                return std::make_shared<MageVAE::MageVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                tensor_storage_map,
+                                                                "first_stage_model",
+                                                                weight_manager);
+            } else if (sd_version_uses_hunyuan_video_vae(vae_version)) {
+                return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                        tensor_storage_map,
+                                                                        "first_stage_model",
+                                                                        false,
+                                                                        vae_version,
+                                                                        weight_manager);
+            } else if (sd_version_uses_wan_vae(vae_version)) {
+                return std::make_shared<WAN::WanVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                           tensor_storage_map,
+                                                           "first_stage_model",
+                                                           false,
+                                                           vae_version,
+                                                           weight_manager);
+            } else {
+                auto model = std::make_shared<AutoEncoderKL>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                             tensor_storage_map,
+                                                             "first_stage_model",
+                                                             false,
+                                                             false,
+                                                             vae_version,
+                                                             weight_manager);
+                if (sd_version_is_sdxl(version) &&
+                    (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) {
+                    float vae_conv_2d_scale = 1.f / 32.f;
+                    LOG_WARN(
+                        "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, "
+                        "using Conv2D scale %.3f",
+                        vae_conv_2d_scale);
+                    model->set_conv2d_scale(vae_conv_2d_scale);
+                }
+                return model;
+            }
+        };
+
+        if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
+            LOG_INFO("using FakeVAE");
+            result.vae = std::make_shared<FakeVAE>(version,
+                                                   ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                   weight_manager);
+        } else if (options.use_tae && !options.tae_preview_only) {
+            LOG_INFO("using TAE for encoding / decoding");
+            result.vae = create_tae(false);
+        } else {
+            LOG_INFO("using VAE for encoding / decoding");
+            result.vae = create_vae();
+            if (options.use_tae && options.tae_preview_only) {
+                LOG_INFO("using TAE for preview");
+                result.preview = create_tae(true);
+            }
+        }
+
+        if (options.use_audio_vae) {
+            if (sd_version_is_minimax_h3(version)) {
+                result.audio = std::make_shared<MiniMaxH3::AudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                           tensor_storage_map,
+                                                                           "",
+                                                                           weight_manager);
+            } else {
+                result.audio = std::make_shared<LTXV::LTXAudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                         tensor_storage_map,
+                                                                         "",
+                                                                         weight_manager);
+            }
+        }
+
+        if (sd_ctx_params->vae_conv_direct) {
+            LOG_INFO("Using Conv2d direct in the vae model");
+            result.vae->set_conv2d_direct_enabled(true);
+            if (result.preview) {
+                result.preview->set_conv2d_direct_enabled(true);
+            }
+        }
+        runners = std::move(result);
+        return true;
+    }
+
+    bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner) {
+        const auto* sd_ctx_params      = &ctx.params;
+        const auto& tensor_storage_map = ctx.tensor_storage_map;
+        const auto version             = ctx.version;
+        const auto& weight_manager     = ctx.weight_manager;
+        if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) {
+            return false;
+        }
+        auto control_net = std::make_shared<ControlNet>(ctx.backends.runtime_backend(SDBackendModule::CONTROL_NET),
+                                                        tensor_storage_map,
+                                                        version,
+                                                        "",
+                                                        weight_manager);
+        if (sd_ctx_params->diffusion_conv_direct) {
+            LOG_INFO("Using Conv2d direct in the control net");
+            control_net->set_conv2d_direct_enabled(true);
+        }
+        runner = std::move(control_net);
+        return true;
+    }
+
+    bool build_extension_runners(const GenerationExtensionInitContext& ctx,
+                                 std::vector<std::shared_ptr<GenerationExtension>>& extensions) {
+        std::vector<std::shared_ptr<GenerationExtension>> result;
+        for (auto extension : {create_photomaker_extension(), create_pulid_extension()}) {
+            if (!extension->init(ctx)) {
+                return false;
+            }
+            if (extension->is_enabled()) {
+                result.push_back(std::move(extension));
+            }
+        }
+        extensions = std::move(result);
+        return true;
+    }
+
+}  // namespace sd::model_builders

src/pipeline/model_builders.h

@@ -0,0 +1,63 @@
+#ifndef __SD_PIPELINE_MODEL_BUILDERS_H__
+#define __SD_PIPELINE_MODEL_BUILDERS_H__
+
+#include <memory>
+#include <vector>
+
+#include "model.h"
+#include "stable-diffusion.h"
+
+class SDBackendManager;
+struct DeviceResidencyManager;
+struct Conditioner;
+struct FrozenCLIPVisionEmbedder;
+struct DiffusionModelRunner;
+struct VAE;
+struct AudioVAERunner;
+struct ControlNet;
+struct GenerationExtension;
+struct GenerationExtensionInitContext;
+namespace IPAdapter {
+    struct IPAdapterRunner;
+}
+
+namespace sd::model_builders {
+
+    struct Context {
+        const sd_ctx_params_t& params;
+        SDVersion version;
+        const String2TensorStorage& tensor_storage_map;
+        SDBackendManager& backends;
+        std::shared_ptr<DeviceResidencyManager> weight_manager;
+    };
+
+    struct CoreRunners {
+        std::shared_ptr<Conditioner> conditioner;
+        std::shared_ptr<DiffusionModelRunner> diffusion;
+        std::shared_ptr<DiffusionModelRunner> high_noise_diffusion;
+        std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;
+        std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
+    };
+
+    struct VAEOptions {
+        bool use_tae                 = false;
+        bool tae_preview_only        = false;
+        bool use_audio_vae           = false;
+        bool external_vae_is_invalid = false;
+    };
+
+    struct VAERunners {
+        std::shared_ptr<VAE> vae;
+        std::shared_ptr<VAE> preview;
+        std::shared_ptr<AudioVAERunner> audio;
+    };
+
+    bool build_core_runners(const Context& ctx, CoreRunners& runners);
+    bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners);
+    bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner);
+    bool build_extension_runners(const GenerationExtensionInitContext& ctx,
+                                 std::vector<std::shared_ptr<GenerationExtension>>& extensions);
+
+}  // namespace sd::model_builders
+
+#endif  // __SD_PIPELINE_MODEL_BUILDERS_H__

src/pipeline/request.cpp

@@ -0,0 +1,471 @@
+#include "request.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <ctime>
+
+#include "diffusion_engine.h"
+#include "runtime/denoiser.hpp"
+
+namespace sd::pipeline {
+
+    const char* sampling_methods_str[] = {
+        "Euler",
+        "Euler A",
+        "Heun",
+        "DPM2",
+        "DPM++ (2s)",
+        "DPM++ (2M)",
+        "modified DPM++ (2M)",
+        "iPNDM",
+        "iPNDM_v",
+        "LCM",
+        "DDIM \"trailing\"",
+        "TCD",
+        "Res Multistep",
+        "Res 2s",
+        "ER-SDE",
+        "Euler CFG++",
+        "Euler A CFG++",
+        "Euler GE",
+        "DPM++ (2M) SDE",
+        "DPM++ (2M) SDE BT",
+        "LMS",
+    };
+
+    static_assert(SAMPLE_METHOD_COUNT == sizeof(sampling_methods_str) / sizeof(sampling_methods_str[0]),
+                  "\nnumber of elements in sampling_methods_str[] != SAMPLE_METHOD_COUNT");
+
+    static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) {
+        return sd_version_is_inpaint_or_unet_edit(version) ||
+               (has_ref_images && sd_version_supports_ref_latent_img_cfg(version));
+    }
+
+    enum sample_method_t default_sample_method(const StableDiffusionGGML* sd) {
+        if (sd != nullptr) {
+            if (sd_version_is_pid(sd->version)) {
+                return LCM_SAMPLE_METHOD;
+            }
+            if (sd_version_is_dit(sd->version)) {
+                return EULER_SAMPLE_METHOD;
+            }
+        }
+        return EULER_A_SAMPLE_METHOD;
+    }
+
+    enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method) {
+        if (sd != nullptr) {
+            auto edm_v_denoiser = std::dynamic_pointer_cast<EDMVDenoiser>(sd->denoiser);
+            if (edm_v_denoiser) {
+                return EXPONENTIAL_SCHEDULER;
+            }
+        }
+        if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) {
+            return LCM_SCHEDULER;
+        } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
+            return SIMPLE_SCHEDULER;
+        } else if (sd != nullptr && sd_version_is_flux(sd->version)) {
+            return FLUX_SCHEDULER;
+        } else if (sd != nullptr && sd_version_is_flux2(sd->version)) {
+            return FLUX2_SCHEDULER;
+        } else if (sd != nullptr && sd_version_is_ltxav(sd->version)) {
+            return LTX2_SCHEDULER;
+        } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) {
+            return LOGIT_NORMAL_SCHEDULER;
+        }
+        return DISCRETE_SCHEDULER;
+    }
+
+    static int64_t resolve_seed(int64_t seed) {
+        if (seed >= 0) {
+            return seed;
+        }
+        srand((int)time(nullptr));
+        return rand();
+    }
+
+    static enum sample_method_t resolve_sample_method(StableDiffusionGGML* sd, enum sample_method_t sample_method) {
+        if (sample_method == SAMPLE_METHOD_COUNT) {
+            return default_sample_method(sd);
+        }
+        return sample_method;
+    }
+
+    static scheduler_t resolve_scheduler(StableDiffusionGGML* sd,
+                                         scheduler_t scheduler,
+                                         enum sample_method_t sample_method) {
+        if (scheduler == SCHEDULER_COUNT) {
+            return default_scheduler(sd, sample_method);
+        }
+        return scheduler;
+    }
+
+    float resolve_eta(StableDiffusionGGML* sd,
+                      float eta,
+                      enum sample_method_t sample_method) {
+        if (eta == INFINITY) {
+            if (sd->version == VERSION_HIDREAM_O1) {
+                return 8.f;
+            }
+            switch (sample_method) {
+                case DDIM_TRAILING_SAMPLE_METHOD:
+                case TCD_SAMPLE_METHOD:
+                case RES_MULTISTEP_SAMPLE_METHOD:
+                case RES_2S_SAMPLE_METHOD:
+                    return 0.0f;
+                case EULER_A_SAMPLE_METHOD:
+                case DPMPP2S_A_SAMPLE_METHOD:
+                case ER_SDE_SAMPLE_METHOD:
+                case EULER_A_CFG_PP_SAMPLE_METHOD:
+                case DPMPP2M_SDE_SAMPLE_METHOD:
+                case DPMPP2M_SDE_BT_SAMPLE_METHOD:
+                    return 1.0f;
+                default:;
+            }
+            return 0.0f;
+        }
+        return eta;
+    }
+
+    GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params) {
+        prompt                      = SAFE_STR(sd_img_gen_params->prompt);
+        negative_prompt             = SAFE_STR(sd_img_gen_params->negative_prompt);
+        width                       = sd_img_gen_params->width;
+        height                      = sd_img_gen_params->height;
+        vae_scale_factor            = sd->get_vae_scale_factor();
+        diffusion_model_down_factor = sd->get_diffusion_model_down_factor();
+        seed                        = sd_img_gen_params->seed;
+        batch_count                 = sd_img_gen_params->batch_count;
+        qwen_image_layers           = std::max(0, sd_img_gen_params->qwen_image_layers);
+        clip_skip                   = sd_img_gen_params->clip_skip;
+        shifted_timestep            = sd_img_gen_params->sample_params.shifted_timestep;
+        strength                    = sd_img_gen_params->strength;
+        control_strength            = sd_img_gen_params->control_strength;
+        eta                         = sd_img_gen_params->sample_params.eta;
+        has_ref_images              = sd_img_gen_params->ref_images_count > 0;
+        guidance                    = sd_img_gen_params->sample_params.guidance;
+        pm_params                   = sd_img_gen_params->pm_params;
+        pulid_params                = sd_img_gen_params->pulid_params;
+        hires                       = sd_img_gen_params->hires;
+        cache_params                = &sd_img_gen_params->cache;
+        resolve(sd);
+    }
+
+    GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params) {
+        prompt           = SAFE_STR(sd_vid_gen_params->prompt);
+        negative_prompt  = SAFE_STR(sd_vid_gen_params->negative_prompt);
+        width            = sd_vid_gen_params->width;
+        height           = sd_vid_gen_params->height;
+        requested_frames = std::max(1, sd_vid_gen_params->video_frames);
+        frames           = sd->align_video_frames(requested_frames);
+        clip_skip        = sd_vid_gen_params->clip_skip;
+        fps              = std::max(1, sd_vid_gen_params->fps);
+        if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
+            LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps);
+            fps = 24;
+        }
+        vae_scale_factor            = sd->get_vae_scale_factor();
+        diffusion_model_down_factor = sd->get_diffusion_model_down_factor();
+        seed                        = sd_vid_gen_params->seed;
+        strength                    = sd_vid_gen_params->strength;
+        cache_params                = &sd_vid_gen_params->cache;
+        vace_strength               = sd_vid_gen_params->vace_strength;
+        guidance                    = sd_vid_gen_params->sample_params.guidance;
+        high_noise_guidance         = sd_vid_gen_params->high_noise_sample_params.guidance;
+        hires                       = sd_vid_gen_params->hires;
+        resolve(sd);
+        if (frames != requested_frames) {
+            LOG_WARN("align video frames from %d to %d for %s",
+                     requested_frames,
+                     frames,
+                     model_version_to_str[sd->version]);
+        }
+    }
+
+    void GenerationRequest::align_generation_request_size() {
+        align_image_size(&width, &height, "generation request");
+    }
+
+    void GenerationRequest::align_image_size(int* target_width, int* target_height, const char* label) {
+        int spatial_multiple = vae_scale_factor * diffusion_model_down_factor;
+        int width_offset     = align_up_offset(*target_width, spatial_multiple);
+        int height_offset    = align_up_offset(*target_height, spatial_multiple);
+        if (width_offset <= 0 && height_offset <= 0) {
+            return;
+        }
+
+        int original_width  = *target_width;
+        int original_height = *target_height;
+
+        *target_width += width_offset;
+        *target_height += height_offset;
+        LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)",
+                 label,
+                 original_width,
+                 original_height,
+                 *target_width,
+                 *target_height,
+                 spatial_multiple);
+    }
+
+    void GenerationRequest::resolve_hires() {
+        if (!hires.enabled) {
+            return;
+        }
+        if (hires.upscaler == SD_HIRES_UPSCALER_NONE) {
+            hires.enabled = false;
+            return;
+        }
+        if (hires.upscaler < SD_HIRES_UPSCALER_NONE || hires.upscaler >= SD_HIRES_UPSCALER_COUNT) {
+            LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler);
+            hires.enabled = false;
+            return;
+        }
+        if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) {
+            LOG_WARN("hires model upscaler requires a model path, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) {
+            LOG_WARN("hires scale must be positive when no target size is set, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        if (hires.custom_sigmas_count < 0) {
+            LOG_WARN("hires custom sigmas count is negative, ignoring custom sigmas");
+            hires.custom_sigmas       = nullptr;
+            hires.custom_sigmas_count = 0;
+        }
+        if (hires.custom_sigmas_count > 0 && hires.custom_sigmas == nullptr) {
+            LOG_WARN("hires custom sigmas count is positive but custom sigmas are null, ignoring custom sigmas");
+            hires.custom_sigmas_count = 0;
+        }
+        if (hires.custom_sigmas_count == 1) {
+            LOG_WARN("hires custom sigmas requires at least two values, ignoring custom sigmas");
+            hires.custom_sigmas       = nullptr;
+            hires.custom_sigmas_count = 0;
+        }
+        hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f);
+        hires.steps              = std::max(0, hires.steps);
+
+        if (hires.target_width > 0 && hires.target_height > 0) {
+            // pass
+        } else if (hires.target_width > 0) {
+            hires.target_height = hires.target_width;
+        } else if (hires.target_height > 0) {
+            hires.target_width = hires.target_height;
+        } else {
+            hires.target_width  = static_cast<int>(std::round(width * hires.scale));
+            hires.target_height = static_cast<int>(std::round(height * hires.scale));
+        }
+
+        if (hires.target_width <= 0 || hires.target_height <= 0) {
+            LOG_WARN("hires target size is not positive, disabling hires");
+            hires.enabled = false;
+            return;
+        }
+        align_image_size(&hires.target_width, &hires.target_height, "hires target");
+    }
+
+    void GenerationRequest::resolve_guidance(StableDiffusionGGML* sd,
+                                             sd_guidance_params_t* guidance,
+                                             bool* use_uncond,
+                                             bool* use_img_uncond,
+                                             bool has_ref_images,
+                                             const char* stage_name) {
+        GGML_ASSERT(guidance != nullptr);
+        GGML_ASSERT(use_uncond != nullptr);
+        GGML_ASSERT(use_img_uncond != nullptr);
+        // out_img_uncond + text_cfg_scale * (out_cond - out_uncond) + image_cfg_scale * (out_uncond - out_img_uncond)
+        // -> text_cfg_scale * out_cond + (image_cfg_scale - text_cfg_scale) * out_uncond + (1 - image_cfg_scale) * out_img_uncond
+        // out_cond       : prompt, image latent
+        // out_uncond     : negative prompt, image latent
+        // out_img_uncond : negative prompt, zero image latent
+        // image_cfg_scale == 1 reduces 3-cond CFG to 2-cond CFG.
+        bool img_cfg_was_set = std::isfinite(guidance->img_cfg);
+        if (!img_cfg_was_set) {
+            guidance->img_cfg = 1.f;
+        }
+
+        if (!sd_version_supports_img_cfg(sd->version, has_ref_images)) {
+            if (img_cfg_was_set && guidance->img_cfg != 1.f) {
+                LOG_WARN("3-conditioning CFG is not supported with this model, disabling it for better performance");
+            }
+            guidance->img_cfg = 1.f;
+        }
+
+        if (guidance->img_cfg != guidance->txt_cfg) {
+            *use_uncond = true;
+        }
+
+        if (guidance->img_cfg != 1.f) {
+            *use_img_uncond = true;
+        }
+
+        if (guidance->txt_cfg < 1.f) {
+            const char* prefix = stage_name == nullptr ? "" : stage_name;
+            if (guidance->txt_cfg == 0.f) {
+                LOG_WARN("%sunconditioned mode, images won't follow the prompt (use cfg-scale=1 for distilled models)",
+                         prefix);
+            } else {
+                LOG_WARN("%scfg value out of expected range may produce unexpected results", prefix);
+            }
+        }
+    }
+
+    void GenerationRequest::resolve(StableDiffusionGGML* sd) {
+        align_generation_request_size();
+        resolve_hires();
+        seed = resolve_seed(seed);
+
+        resolve_guidance(sd, &guidance, &use_uncond, &use_img_uncond, has_ref_images);
+        if (sd->high_noise_diffusion_model) {
+            resolve_guidance(sd,
+                             &high_noise_guidance,
+                             &use_high_noise_uncond,
+                             &use_high_noise_img_uncond,
+                             has_ref_images,
+                             "high noise: ");
+        }
+
+        if (shifted_timestep > 0 && !sd_version_is_sdxl(sd->version)) {
+            LOG_WARN("timestep shifting is only supported for SDXL models!");
+            shifted_timestep = 0;
+        }
+    }
+
+    SamplePlan::SamplePlan(StableDiffusionGGML* sd,
+                           const sd_img_gen_params_t* sd_img_gen_params,
+                           const GenerationRequest& request) {
+        sample_method     = sd_img_gen_params->sample_params.sample_method;
+        extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args;
+        eta               = sd_img_gen_params->sample_params.eta;
+        sample_steps      = sd_img_gen_params->sample_params.sample_steps;
+        resolve(sd, &request, &sd_img_gen_params->sample_params);
+    }
+
+    SamplePlan::SamplePlan(StableDiffusionGGML* sd,
+                           const sd_vid_gen_params_t* sd_vid_gen_params,
+                           const GenerationRequest& request) {
+        sample_method     = sd_vid_gen_params->sample_params.sample_method;
+        extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args;
+        eta               = sd_vid_gen_params->sample_params.eta;
+        sample_steps      = sd_vid_gen_params->sample_params.sample_steps;
+        if (sd->high_noise_diffusion_model) {
+            high_noise_sample_steps      = sd_vid_gen_params->high_noise_sample_params.sample_steps;
+            high_noise_sample_method     = sd_vid_gen_params->high_noise_sample_params.sample_method;
+            high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args;
+            high_noise_eta               = sd_vid_gen_params->high_noise_sample_params.eta;
+        }
+        moe_boundary = sd_vid_gen_params->moe_boundary;
+        resolve(sd, &request, &sd_vid_gen_params->sample_params);
+    }
+
+    void SamplePlan::resolve(StableDiffusionGGML* sd,
+                             const GenerationRequest* request,
+                             const sd_sample_params_t* sample_params) {
+        sample_method = resolve_sample_method(sd, sample_method);
+
+        total_steps = sample_steps + std::max(0, high_noise_sample_steps);
+
+        if (sample_params->custom_sigmas_count > 0) {
+            sigmas      = std::vector<float>(sample_params->custom_sigmas,
+                                        sample_params->custom_sigmas + sample_params->custom_sigmas_count);
+            total_steps = static_cast<int>(sigmas.size()) - 1;
+            LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps);
+            if (sample_steps >= total_steps) {
+                sample_steps = total_steps;
+                LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps);
+            }
+            if (high_noise_sample_steps > 0) {
+                high_noise_sample_steps = total_steps - sample_steps;
+                LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps);
+            }
+        } else {
+            scheduler_t scheduler = resolve_scheduler(sd,
+                                                      sample_params->scheduler,
+                                                      sample_method);
+            int sample_seq_len    = sd->get_image_seq_len(request->height, request->width);
+            if (sd_version_is_ltxav(sd->version) && request->frames > 0) {
+                int latent_frames = ((request->frames - 1) / 8) + 1;
+                sample_seq_len *= latent_frames;
+            } else if (sd_version_is_minimax_h3(sd->version) && request->frames > 0) {
+                sample_seq_len *= sd->video_frames_to_latent_frames(request->frames);
+            }
+            sigmas = sd->denoiser->get_sigmas(total_steps,
+                                              sample_seq_len,
+                                              scheduler,
+                                              sd->version,
+                                              sample_params->extra_sample_args);
+        }
+
+        eta = resolve_eta(sd, eta, sample_method);
+
+        if (high_noise_sample_steps < 0) {
+            for (size_t i = 0; i < sigmas.size(); ++i) {
+                if (sigmas[i] < moe_boundary) {
+                    high_noise_sample_steps = static_cast<int>(i);
+                    break;
+                }
+            }
+            LOG_VERBOSE("switching from high noise model at step %d", high_noise_sample_steps);
+        }
+
+        LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]);
+        if (high_noise_sample_steps > 0) {
+            high_noise_sample_method = resolve_sample_method(sd,
+                                                             high_noise_sample_method);
+            high_noise_eta           = resolve_eta(sd, high_noise_eta, high_noise_sample_method);
+            LOG_INFO("sampling(high noise) using %s method", sampling_methods_str[high_noise_sample_method]);
+        }
+    }
+
+    std::vector<float> make_hires_sigma_schedule(StableDiffusionGGML* sd,
+                                                 const sd_hires_params_t& hires,
+                                                 const sd_sample_params_t& sample_params,
+                                                 sample_method_t sample_method,
+                                                 int default_steps,
+                                                 int sample_seq_len,
+                                                 int* scheduler_steps_out) {
+        if (scheduler_steps_out != nullptr) {
+            *scheduler_steps_out = 0;
+        }
+
+        if (hires.custom_sigmas_count > 0 && hires.custom_sigmas != nullptr) {
+            std::vector<float> custom_sigmas(hires.custom_sigmas,
+                                             hires.custom_sigmas + hires.custom_sigmas_count);
+            if (scheduler_steps_out != nullptr) {
+                *scheduler_steps_out = static_cast<int>(custom_sigmas.size()) - 1;
+            }
+            return custom_sigmas;
+        }
+
+        int effective_steps = hires.steps > 0 ? hires.steps : default_steps;
+        effective_steps     = std::max(1, effective_steps);
+
+        // sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps,
+        // unlike img2img which trims from a fixed step count.
+        int scheduler_steps = static_cast<int>(effective_steps / hires.denoising_strength);
+        scheduler_steps     = std::max(1, scheduler_steps);
+
+        scheduler_t scheduler     = resolve_scheduler(sd,
+                                                      sample_params.scheduler,
+                                                      sample_method);
+        std::vector<float> sigmas = sd->denoiser->get_sigmas(scheduler_steps,
+                                                             sample_seq_len,
+                                                             scheduler,
+                                                             sd->version,
+                                                             sample_params.extra_sample_args);
+        size_t t_enc              = static_cast<size_t>(scheduler_steps * hires.denoising_strength);
+        if (t_enc >= static_cast<size_t>(scheduler_steps)) {
+            t_enc = static_cast<size_t>(scheduler_steps) - 1;
+        }
+        if (scheduler_steps_out != nullptr) {
+            *scheduler_steps_out = scheduler_steps;
+        }
+        return std::vector<float>(sigmas.begin() + scheduler_steps - static_cast<int>(t_enc) - 1,
+                                  sigmas.end());
+    }
+
+}  // namespace sd::pipeline

src/pipeline/request.h

@@ -0,0 +1,110 @@
+#ifndef __SD_PIPELINE_REQUEST_H__
+#define __SD_PIPELINE_REQUEST_H__
+
+#include <string>
+#include <vector>
+
+#include "stable-diffusion.h"
+
+class StableDiffusionGGML;
+
+namespace sd::pipeline {
+
+    extern const char* sampling_methods_str[];
+
+    enum sample_method_t default_sample_method(const StableDiffusionGGML* sd);
+
+    enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method);
+
+    float resolve_eta(StableDiffusionGGML* sd,
+                      float eta,
+                      enum sample_method_t sample_method);
+
+    struct GenerationRequest {
+        std::string prompt;
+        std::string negative_prompt;
+        int width                                = -1;
+        int height                               = -1;
+        int clip_skip                            = -1;
+        int vae_scale_factor                     = -1;
+        int diffusion_model_down_factor          = -1;
+        int64_t seed                             = -1;
+        bool use_uncond                          = false;
+        bool use_img_uncond                      = false;
+        bool use_high_noise_uncond               = false;
+        bool use_high_noise_img_uncond           = false;
+        bool has_ref_images                      = false;
+        const sd_cache_params_t* cache_params    = nullptr;
+        int batch_count                          = 1;
+        int qwen_image_layers                    = 3;
+        int shifted_timestep                     = 0;
+        float strength                           = 1.f;
+        float control_strength                   = 0.f;
+        float eta                                = 0.f;
+        sd_guidance_params_t guidance            = {};
+        sd_guidance_params_t high_noise_guidance = {};
+        sd_pm_params_t pm_params                 = {};
+        sd_pulid_params_t pulid_params           = {};
+        sd_hires_params_t hires                  = {};
+        int frames                               = -1;
+        int requested_frames                     = -1;
+        int fps                                  = 16;
+        float vace_strength                      = 1.f;
+
+        GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params);
+
+        GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params);
+
+        void align_generation_request_size();
+
+        void align_image_size(int* target_width, int* target_height, const char* label);
+
+        void resolve_hires();
+
+        static void resolve_guidance(StableDiffusionGGML* sd,
+                                     sd_guidance_params_t* guidance,
+                                     bool* use_uncond,
+                                     bool* use_img_uncond,
+                                     bool has_ref_images,
+                                     const char* stage_name = nullptr);
+
+        void resolve(StableDiffusionGGML* sd);
+    };
+
+    struct SamplePlan {
+        enum sample_method_t sample_method            = SAMPLE_METHOD_COUNT;
+        enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT;
+        const char* extra_sample_args                 = nullptr;
+        const char* high_noise_extra_sample_args      = nullptr;
+        float eta                                     = 0.f;
+        float high_noise_eta                          = 0.f;
+        int sample_steps                              = 0;
+        int high_noise_sample_steps                   = 0;
+        int total_steps                               = 0;
+        float moe_boundary                            = 0.f;
+        std::vector<float> sigmas;
+
+        SamplePlan(StableDiffusionGGML* sd,
+                   const sd_img_gen_params_t* sd_img_gen_params,
+                   const GenerationRequest& request);
+
+        SamplePlan(StableDiffusionGGML* sd,
+                   const sd_vid_gen_params_t* sd_vid_gen_params,
+                   const GenerationRequest& request);
+
+        void resolve(StableDiffusionGGML* sd,
+                     const GenerationRequest* request,
+                     const sd_sample_params_t* sample_params);
+    };
+
+    std::vector<float> make_hires_sigma_schedule(StableDiffusionGGML* sd,
+                                                 const sd_hires_params_t& hires,
+                                                 const sd_sample_params_t& sample_params,
+                                                 sample_method_t sample_method,
+                                                 int default_steps,
+                                                 int sample_seq_len,
+                                                 int* scheduler_steps_out);
+
+}  // namespace sd::pipeline
+
+#endif  // __SD_PIPELINE_REQUEST_H__

src/runtime/cache_dit.hpp

@@ -8,8 +8,8 @@
 #include <unordered_map>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 #include "runtime/condition_cache_utils.hpp"
 
 struct DBCacheConfig {

src/runtime/denoiser.hpp

@@ -11,8 +11,10 @@
 #include <string>
 #include <utility>
 
-#include "core/ggml_extend.hpp"
+#include "core/rng.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
+#include "model.h"
 #include "runtime/gits_noise.h"
 #include "runtime/guidance.h"
 

src/runtime/easycache.hpp

@@ -6,8 +6,8 @@
 #include <unordered_map>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 #include "runtime/condition_cache_utils.hpp"
 #include "runtime/denoiser.hpp"
 

src/runtime/preprocessing.hpp

@@ -4,7 +4,7 @@
 #include <cmath>
 #include <limits>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_tensor_utils.h"
 
 #define M_PI_ 3.14159265358979323846f
 

src/runtime/spectrum.hpp

@@ -5,8 +5,8 @@
 #include <cstring>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 
 struct SpectrumConfig {
     float w            = 0.40f;

src/runtime/tiling.cpp

@@ -0,0 +1,268 @@
+#include "runtime/tiling.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "core/util.h"
+#include "ggml.h"
+
+static void sd_tiling_calc_tiles(int& num_tiles_dim,
+                                 float& tile_overlap_factor_dim,
+                                 int small_dim,
+                                 int tile_size,
+                                 const float tile_overlap_factor,
+                                 bool circular) {
+    int tile_overlap     = static_cast<int>(tile_size * tile_overlap_factor);
+    int non_tile_overlap = tile_size - tile_overlap;
+
+    if (circular) {
+        // circular means the last and first tile are overlapping (wraping around)
+        num_tiles_dim = small_dim / non_tile_overlap;
+
+        if (num_tiles_dim < 1) {
+            num_tiles_dim = 1;
+        }
+
+        tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size;
+
+        // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles
+        if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) {
+            num_tiles_dim++;
+            tile_overlap_factor_dim = 0.5;
+        }
+
+        return;
+    }
+    // else, non-circular means the last and first tile are not overlapping
+
+    num_tiles_dim     = (small_dim - tile_overlap) / non_tile_overlap;
+    int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim;
+
+    if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) {
+        // if tiles don't fit perfectly using the desired overlap
+        // and there is enough room to squeeze an extra tile without overlap becoming >0.5
+        num_tiles_dim++;
+    }
+
+    tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1));
+    if (num_tiles_dim <= 2) {
+        if (small_dim <= tile_size) {
+            num_tiles_dim           = 1;
+            tile_overlap_factor_dim = 0;
+        } else {
+            num_tiles_dim           = 2;
+            tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size;
+        }
+    }
+}
+
+static int64_t sd_tensor_plane_size(const sd::Tensor<float>& tensor) {
+    GGML_ASSERT(tensor.dim() >= 2);
+    return tensor.shape()[0] * tensor.shape()[1];
+}
+
+static sd::Tensor<float> sd_tensor_split_2d(const sd::Tensor<float>& input, int width, int height, int x, int y) {
+    GGML_ASSERT(input.dim() >= 4);
+    std::vector<int64_t> output_shape = input.shape();
+    output_shape[0]                   = width;
+    output_shape[1]                   = height;
+    sd::Tensor<float> output(std::move(output_shape));
+    int64_t input_width  = input.shape()[0];
+    int64_t input_height = input.shape()[1];
+    int64_t input_plane  = sd_tensor_plane_size(input);
+    int64_t output_plane = sd_tensor_plane_size(output);
+    int64_t plane_count  = input.numel() / input_plane;
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            int64_t src_xy = (ix + x) % input_width + input_width * ((iy + y) % input_height);
+            int64_t dst_xy = ix + width * iy;
+            for (int64_t plane = 0; plane < plane_count; ++plane) {
+                output[plane * output_plane + dst_xy] = input[plane * input_plane + src_xy];
+            }
+        }
+    }
+    return output;
+}
+
+static void sd_tensor_merge_2d(const sd::Tensor<float>& input,
+                               sd::Tensor<float>* output,
+                               int x,
+                               int y,
+                               int overlap_x,
+                               int overlap_y,
+                               bool circular_x,
+                               bool circular_y,
+                               int x_skip,
+                               int y_skip) {
+    GGML_ASSERT(output != nullptr);
+    int64_t width        = input.shape()[0];
+    int64_t height       = input.shape()[1];
+    int64_t img_width    = output->shape()[0];
+    int64_t img_height   = output->shape()[1];
+    int64_t input_plane  = sd_tensor_plane_size(input);
+    int64_t output_plane = sd_tensor_plane_size(*output);
+    int64_t plane_count  = input.numel() / input_plane;
+    GGML_ASSERT(output->numel() / output_plane == plane_count);
+
+    // unclamped -> expects x in the range [0-1]
+    auto smootherstep_f32 = [](const float x) -> float {
+        GGML_ASSERT(x >= 0.f && x <= 1.f);
+        return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f);
+    };
+
+    for (int iy = y_skip; iy < height; iy++) {
+        for (int ix = x_skip; ix < width; ix++) {
+            int64_t src_xy = ix + width * iy;
+            int64_t ox     = (x + ix) % img_width;
+            int64_t oy     = (y + iy) % img_height;
+            int64_t dst_xy = ox + img_width * oy;
+            for (int64_t plane = 0; plane < plane_count; ++plane) {
+                float new_value = input[plane * input_plane + src_xy];
+                if (overlap_x > 0 || overlap_y > 0) {
+                    float old_value   = (*output)[plane * output_plane + dst_xy];
+                    const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f;
+                    const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f;
+                    const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f;
+                    const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f;
+                    const float x_f   = std::min(std::min(x_f_0, x_f_1), 1.f);
+                    const float y_f   = std::min(std::min(y_f_0, y_f_1), 1.f);
+                    (*output)[plane * output_plane + dst_xy] =
+                        old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f);
+                } else {
+                    (*output)[plane * output_plane + dst_xy] = new_value;
+                }
+            }
+        }
+    }
+}
+
+sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
+                                   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,
+                                   const TileProcessCallback& on_processing,
+                                   bool silent) {
+    sd::Tensor<float> output;
+    int input_width  = static_cast<int>(input.shape()[0]);
+    int input_height = static_cast<int>(input.shape()[1]);
+
+    GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) &&
+                ((output_width / input_width) == (output_height / input_height)));
+    GGML_ASSERT(((input_width / output_width) == scale) ||
+                ((output_width / input_width) == scale));
+
+    int small_width  = output_width;
+    int small_height = output_height;
+    bool decode      = output_width > input_width;
+    if (decode) {
+        small_width  = input_width;
+        small_height = input_height;
+    }
+
+    int num_tiles_x;
+    float tile_overlap_factor_x;
+    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x);
+
+    int num_tiles_y;
+    float tile_overlap_factor_y;
+    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y);
+
+    int tile_overlap_x     = static_cast<int32_t>(p_tile_size_x * tile_overlap_factor_x);
+    int non_tile_overlap_x = p_tile_size_x - tile_overlap_x;
+    int tile_overlap_y     = static_cast<int32_t>(p_tile_size_y * tile_overlap_factor_y);
+    int non_tile_overlap_y = p_tile_size_y - tile_overlap_y;
+    int tile_size_x        = p_tile_size_x < small_width ? p_tile_size_x : small_width;
+    int tile_size_y        = p_tile_size_y < small_height ? p_tile_size_y : small_height;
+    int input_tile_size_x  = tile_size_x;
+    int input_tile_size_y  = tile_size_y;
+    int output_tile_size_x = tile_size_x;
+    int output_tile_size_y = tile_size_y;
+    if (decode) {
+        output_tile_size_x *= scale;
+        output_tile_size_y *= scale;
+    } else {
+        input_tile_size_x *= scale;
+        input_tile_size_y *= scale;
+    }
+
+    int num_tiles   = num_tiles_x * num_tiles_y;
+    int tile_count  = 1;
+    bool last_y     = false;
+    bool last_x     = false;
+    float last_time = 0.0f;
+    if (!silent) {
+        LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y);
+        LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor);
+        LOG_VERBOSE("processing %i tiles", num_tiles);
+        pretty_progress(0, num_tiles, 0.0f);
+    }
+    for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) {
+        int dy = 0;
+        if (!circular_y && y + tile_size_y >= small_height) {
+            int original_y = y;
+            y              = small_height - tile_size_y;
+            dy             = original_y - y;
+            if (decode) {
+                dy *= scale;
+            }
+            last_y = true;
+        }
+        for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) {
+            int dx = 0;
+            if (!circular_x && x + tile_size_x >= small_width) {
+                int original_x = x;
+                x              = small_width - tile_size_x;
+                dx             = original_x - x;
+                if (decode) {
+                    dx *= scale;
+                }
+                last_x = true;
+            }
+
+            int x_in  = decode ? x : scale * x;
+            int y_in  = decode ? y : scale * y;
+            int x_out = decode ? x * scale : x;
+            int y_out = decode ? y * scale : y;
+
+            int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x;
+            int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y;
+
+            int64_t t1       = ggml_time_ms();
+            auto input_tile  = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in);
+            auto output_tile = on_processing(input_tile);
+            if (output_tile.empty()) {
+                return {};
+            }
+            GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y);
+            if (output.empty()) {
+                std::vector<int64_t> output_shape = output_tile.shape();
+                output_shape[0]                   = output_width;
+                output_shape[1]                   = output_height;
+                output                            = sd::Tensor<float>::zeros(std::move(output_shape));
+            }
+            sd_tensor_merge_2d(output_tile, &output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy);
+
+            if (!silent) {
+                int64_t t2 = ggml_time_ms();
+                last_time  = (t2 - t1) / 1000.0f;
+                pretty_progress(tile_count, num_tiles, last_time);
+            }
+            tile_count++;
+        }
+        last_x = false;
+    }
+    if (!silent && tile_count < num_tiles) {
+        pretty_progress(num_tiles, num_tiles, last_time);
+    }
+    if (output.empty()) {
+        return {};
+    }
+    return output;
+}