Commit ea209db for stable-diffusion.cpp
commit ea209dbfcf1ab96ca25e3612820dc94bbc66f00d
Author: leejet <leejet714@gmail.com>
Date: Sun Oct 25 05:18:30 2026 +0800
perf: cache CPU position embeddings across DiT runners (#2211)
diff --git a/docs/pe_cache.md b/docs/pe_cache.md
new file mode 100644
index 0000000..e181959
--- /dev/null
+++ b/docs/pe_cache.md
@@ -0,0 +1,60 @@
+# Position embedding cache
+
+`sd::PECache` stores CPU-side `sd::Tensor<float>` data. Each model runner owns
+its cache and calls the existing `make_input()` on every graph build. Cache
+hits skip PE generation; allocation and upload remain ordinary graph-input
+operations. `GGMLRunner` has no PE-specific code or resource hooks.
+
+The cache keeps the most recent key and tensor for each named PE input. A
+matching key returns the existing tensor by reference. Changing its parameters
+replaces that input's cached result. Different named inputs have independent
+storage. Alternating two layouts under the same name recomputes their PE;
+there is no LRU, device-buffer cache, memory budget, or graph lifecycle callback.
+CPU data lives with the model runner, including across `runner_end()` calls,
+and is released when its named entry is replaced or the cache is destroyed.
+
+## Adding a model
+
+Include `core/pe_cache.h` and declare a member:
+
+```cpp
+sd::PECache pe_cache_;
+```
+
+Construct the key before generating PE and place the generation code in the
+lazy callback:
+
+```cpp
+sd::PECacheKey key("model.image_pe");
+key.add(height).add(width).add(head_dim).add(theta)
+ .add(circular_x_enabled).add(circular_y_enabled);
+auto pe = make_input(pe_cache_.get_or_create(key, [&]() {
+ auto values = generate_position_embedding(/* existing arguments */);
+ return sd::Tensor<float>({2, 2, head_dim / 2, token_count}, std::move(values));
+}));
+```
+
+Keys must cover every variable affecting PE values or shape: geometry, text
+length, frequency settings, reference-image ordering and offsets, circular
+settings, and video/audio positions or frame rates as applicable.
+`add_shape()` records all four ggml dimensions, including trailing singleton
+dimensions. `add_shapes()` records an ordered list of reference shapes.
+`add(sd::Tensor<T>)` records both shape and values for explicit position inputs.
+Keys compare their complete encoded fields rather than a data pointer or hash.
+
+Use distinct names for PE inputs that coexist in a graph, such as Anima's
+adapter query and key PE. `make_input()` defers reading the CPU data until graph
+execution: do not replace a named entry while an unfinished graph references
+it. Adding or updating other named entries does not invalidate its storage.
+Return tensor shapes with the contiguous dimension first, and preserve each
+model's existing F32 values, layout, and final transformations.
+
+## Current integrations
+
+The cache is used by MiniT2I, the Flux runner and its variants, Flux3, Qwen Image,
+Qwen Image 2.1, Z-Image and L2P, Anima, Boogu, ERNIE Image, Fibo2, Ideogram4,
+Iris, Krea2, Lens, LLaDA Image, MAGE Flow, Ming Image, MOSS Image, PiD, PixArt,
+Wan, Hunyuan Video, LingBot Video, and LTX audio/video.
+
+Position embeddings stored as model weights and graph-side RoPE operations
+remain on their existing paths, including MiniMax H3 and SenseNova U1.
diff --git a/src/core/pe_cache.cpp b/src/core/pe_cache.cpp
new file mode 100644
index 0000000..df4aa62
--- /dev/null
+++ b/src/core/pe_cache.cpp
@@ -0,0 +1,30 @@
+#include "core/pe_cache.h"
+
+namespace sd {
+ PECacheKey& PECacheKey::add_shape(const ggml_tensor* tensor) {
+ add(tensor != nullptr);
+ if (tensor != nullptr) {
+ for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+ add(tensor->ne[i]);
+ }
+ }
+ return *this;
+ }
+
+ PECacheKey& PECacheKey::add_shapes(const std::vector<ggml_tensor*>& tensors) {
+ add(tensors.size());
+ for (const auto* tensor : tensors) {
+ add_shape(tensor);
+ }
+ return *this;
+ }
+
+ const Tensor<float>& PECache::get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build) {
+ auto entry = entries_.find(key.name());
+ if (entry != entries_.end() && entry->second.key == key) {
+ return entry->second.value;
+ }
+ auto value = build();
+ return entries_.insert_or_assign(key.name(), Entry{key, std::move(value)}).first->second.value;
+ }
+}
diff --git a/src/core/pe_cache.h b/src/core/pe_cache.h
new file mode 100644
index 0000000..3539707
--- /dev/null
+++ b/src/core/pe_cache.h
@@ -0,0 +1,82 @@
+#ifndef __SD_CORE_PE_CACHE_H__
+#define __SD_CORE_PE_CACHE_H__
+
+#include <functional>
+#include <map>
+#include <set>
+#include <type_traits>
+
+#include "core/tensor.hpp"
+#include "ggml.h"
+
+namespace sd {
+ class PECacheKey {
+ std::string name_;
+ std::string fields_;
+
+ public:
+ explicit PECacheKey(std::string name)
+ : name_(std::move(name)) {}
+ const std::string& name() const { return name_; }
+
+ template <typename T, std::enable_if_t<std::is_arithmetic_v<T> || std::is_enum_v<T>, int> = 0>
+ PECacheKey& add(T value) {
+ if constexpr (std::is_floating_point_v<T>) {
+ const double field = static_cast<double>(value);
+ fields_.push_back('f');
+ fields_.append(reinterpret_cast<const char*>(&field), sizeof(field));
+ } else {
+ const uint64_t field = static_cast<uint64_t>(value);
+ fields_.push_back('i');
+ fields_.append(reinterpret_cast<const char*>(&field), sizeof(field));
+ }
+ return *this;
+ }
+
+ template <typename T>
+ PECacheKey& add(const std::vector<T>& values) {
+ fields_.push_back('v');
+ add(values.size());
+ for (const auto& value : values) {
+ add(value);
+ }
+ return *this;
+ }
+
+ template <typename T>
+ PECacheKey& add(const std::set<T>& values) {
+ return add(std::vector<T>(values.begin(), values.end()));
+ }
+
+ template <typename T>
+ PECacheKey& add(const Tensor<T>& values) {
+ add(values.shape());
+ for (int64_t i = 0; i < values.numel(); ++i) {
+ add(values.data()[i]);
+ }
+ return *this;
+ }
+
+ PECacheKey& add_shape(const ggml_tensor* tensor);
+ PECacheKey& add_shapes(const std::vector<ggml_tensor*>& tensors);
+
+ bool operator==(const PECacheKey& other) const {
+ return name_ == other.name_ && fields_ == other.fields_;
+ }
+ };
+
+ class PECache {
+ struct Entry {
+ PECacheKey key;
+ Tensor<float> value;
+ };
+
+ std::map<std::string, Entry> entries_;
+
+ public:
+ const Tensor<float>& get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build);
+ };
+
+}
+
+#endif // __SD_CORE_PE_CACHE_H__
diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp
index a414d96..66e1a0b 100644
--- a/src/model/diffusion/anima.hpp
+++ b/src/model/diffusion/anima.hpp
@@ -7,6 +7,7 @@
#include <memory>
#include <utility>
#include <vector>
+#include "core/pe_cache.h"
#include "model/common/block.hpp"
#include "model/common/rope.hpp"
@@ -562,10 +563,9 @@ namespace Anima {
};
struct AnimaRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
- std::vector<float> image_pe_vec;
- std::vector<float> adapter_q_pe_vec;
- std::vector<float> adapter_k_pe_vec;
AnimaConfig config;
AnimaNet net;
@@ -660,19 +660,23 @@ namespace Anima {
int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w;
- image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1,
- static_cast<int>(h_pad),
- static_cast<int>(w_pad),
- static_cast<int>(config.patch_size),
- config.theta,
- config.axes_dim,
- 4.0f,
- 4.0f,
- 1.0f,
- ref_latents));
- int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
- auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
- set_backend_tensor_data(image_pe, image_pe_vec.data());
+ sd::PECacheKey image_pe_key("anima.image_pe");
+ image_pe_key.add(h_pad).add(w_pad).add(config.patch_size).add(config.theta).add(config.axes_dim).add(config.head_dim);
+ image_pe_key.add_shapes(ref_latents).add(circular_x_enabled).add(circular_y_enabled);
+ auto image_pe = make_input(pe_cache_.get_or_create(image_pe_key, [&]() {
+ auto image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1,
+ static_cast<int>(h_pad),
+ static_cast<int>(w_pad),
+ static_cast<int>(config.patch_size),
+ config.theta,
+ config.axes_dim,
+ 4.0f,
+ 4.0f,
+ 1.0f,
+ ref_latents));
+ int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
+ return sd::Tensor<float>({2, 2, config.head_dim / 2, image_pos_len}, std::move(image_pe_vec));
+ }));
ggml_tensor* adapter_q_pe = nullptr;
ggml_tensor* adapter_k_pe = nullptr;
@@ -680,16 +684,16 @@ namespace Anima {
int64_t target_len = t5_ids->ne[0];
int64_t source_len = context->ne[1];
- adapter_q_pe_vec = gen_1d_rope_pe_vec(target_len, 64, 10000.f);
- adapter_k_pe_vec = gen_1d_rope_pe_vec(source_len, 64, 10000.f);
-
- int64_t target_pos_len = static_cast<int64_t>(adapter_q_pe_vec.size()) / (2 * 2 * 32);
- int64_t source_pos_len = static_cast<int64_t>(adapter_k_pe_vec.size()) / (2 * 2 * 32);
-
- adapter_q_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, target_pos_len);
- adapter_k_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, source_pos_len);
- set_backend_tensor_data(adapter_q_pe, adapter_q_pe_vec.data());
- set_backend_tensor_data(adapter_k_pe, adapter_k_pe_vec.data());
+ sd::PECacheKey q_key("anima.adapter_q_pe");
+ q_key.add(target_len);
+ adapter_q_pe = make_input(pe_cache_.get_or_create(q_key, [&]() {
+ return sd::Tensor<float>({2, 2, 32, target_len}, gen_1d_rope_pe_vec(target_len, 64, 10000.f));
+ }));
+ sd::PECacheKey k_key("anima.adapter_k_pe");
+ k_key.add(source_len);
+ adapter_k_pe = make_input(pe_cache_.get_or_create(k_key, [&]() {
+ return sd::Tensor<float>({2, 2, 32, source_len}, gen_1d_rope_pe_vec(source_len, 64, 10000.f));
+ }));
}
auto runner_ctx = get_context();
diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp
index 3ee6714..ac67ff9 100644
--- a/src/model/diffusion/boogu.hpp
+++ b/src/model/diffusion/boogu.hpp
@@ -6,6 +6,7 @@
#include <cmath>
#include <tuple>
#include <vector>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@@ -759,9 +760,9 @@ namespace Boogu {
}
struct BooguImageRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
BooguConfig config;
BooguImageModel boogu;
- std::vector<float> pe_vec;
BooguImageRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -799,17 +800,21 @@ namespace Boogu {
ref_latents.push_back(make_input(ref_latent_tensor));
}
- pe_vec = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- ref_latents,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("boogu.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ ref_latents,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp
index b284cb8..d4f774f 100644
--- a/src/model/diffusion/ernie_image.hpp
+++ b/src/model/diffusion/ernie_image.hpp
@@ -4,6 +4,7 @@
#include <cinttypes>
#include <memory>
#include <vector>
+#include "core/pe_cache.h"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
@@ -383,9 +384,9 @@ namespace ErnieImage {
};
struct ErnieImageRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
ErnieImageConfig config;
ErnieImageModel ernie_image;
- std::vector<float> pe_vec;
ErnieImageRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -415,16 +416,20 @@ namespace ErnieImage {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
- pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("ernie_image.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({config.axes_dim_sum, 1, pos_len, 2}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe);
diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp
index f53b6ec..a642cb4 100644
--- a/src/model/diffusion/fibo2.hpp
+++ b/src/model/diffusion/fibo2.hpp
@@ -1,6 +1,7 @@
#ifndef __SD_MODEL_DIFFUSION_FIBO2_HPP__
#define __SD_MODEL_DIFFUSION_FIBO2_HPP__
+#include "core/pe_cache.h"
#include "model/diffusion/z_image.hpp"
// https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_bria_fibo2.py
@@ -333,9 +334,9 @@ namespace Fibo2 {
};
struct Fibo2Runner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
Fibo2Config config;
Fibo2Model model;
- std::vector<float> pe_values;
Fibo2Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
@@ -360,28 +361,35 @@ namespace Fibo2 {
auto x = make_input(*p.x);
GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels);
int count = config.gist_count(static_cast<int>(p.context->shape()[1]));
- Rope::Embedding embedding;
- for (int i = 0; i < count; ++i) {
- embedding.ids.push_back({static_cast<float>(i + 1), 0.f, 0.f});
- }
- embedding.positions.append_tokens(count);
- auto append_image = [&](ggml_tensor* image, int plane) {
- auto ids = Rope::gen_flux_img_ids(static_cast<int>(image->ne[1]), static_cast<int>(image->ne[0]),
- 1, 1, 3, plane, 0, 0, false, &embedding.positions);
- embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end());
- };
- append_image(x, static_cast<int>(config.max_gist + 1));
std::vector<ggml_tensor*> refs;
if (p.ref_latents && p.ref_image_params.pass_to_dit) {
for (const auto& ref : *p.ref_latents) {
refs.push_back(make_input(ref));
- append_image(refs.back(), static_cast<int>(config.max_gist + 1 + refs.size()));
}
}
- embedding.values = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies);
- pe_values = finish_rope_pe(std::move(embedding));
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, static_cast<int64_t>(pe_values.size() / 256));
- set_backend_tensor_data(pe, pe_values.data());
+ sd::PECacheKey pe_key("fibo2.pe");
+ pe_key.add_shape(x).add(count).add(config.max_gist).add_shapes(refs).add(circular_x_enabled);
+ pe_key.add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ Rope::Embedding embedding;
+ for (int i = 0; i < count; ++i) {
+ embedding.ids.push_back({static_cast<float>(i + 1), 0.f, 0.f});
+ }
+ embedding.positions.append_tokens(count);
+ auto append_image = [&](ggml_tensor* image, int plane) {
+ auto ids = Rope::gen_flux_img_ids(static_cast<int>(image->ne[1]), static_cast<int>(image->ne[0]),
+ 1, 1, 3, plane, 0, 0, false, &embedding.positions);
+ embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end());
+ };
+ append_image(x, static_cast<int>(config.max_gist + 1));
+ for (size_t i = 0; i < refs.size(); ++i) {
+ append_image(refs[i], static_cast<int>(config.max_gist + 2 + i));
+ }
+ embedding.values = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies);
+ auto values = finish_rope_pe(std::move(embedding));
+ const int64_t tokens = static_cast<int64_t>(values.size() / 256);
+ return sd::Tensor<float>({2, 2, 64, tokens}, std::move(values));
+ }));
auto ctx = get_context();
auto out = model.forward(&ctx, x, make_input(*p.timesteps), make_input(*p.context), pe, count, refs);
ggml_build_forward_expand(graph, out);
diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp
index 825dabe..2f4272e 100644
--- a/src/model/diffusion/flux.hpp
+++ b/src/model/diffusion/flux.hpp
@@ -6,6 +6,7 @@
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
#include "core/util.h"
#include "model/adapter/pulid.hpp"
@@ -1398,10 +1399,11 @@ namespace Flux {
};
struct FluxRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
FluxConfig config;
Flux flux;
- std::vector<float> pe_vec;
std::vector<float> mod_index_arange_vec;
std::vector<float> dct_vec;
sd::Tensor<float> guidance_tensor;
@@ -1548,25 +1550,26 @@ namespace Flux {
} else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2};
}
- pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- txt_arange_dims,
- ref_latents,
- ref_index_mode,
- config.ref_index_scale,
- config.theta,
- config.axes_dim,
- sd_version_is_longcat(version)));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- // LOG_VERBOSE("pos_len %d", pos_len);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- // pe->data = pe_vec.data();
- // print_ggml_tensor(pe);
- // pe->data = nullptr;
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("flux.pe");
+ pe_key.add(version).add_shape(x).add_shape(context).add(config.patch_size).add(txt_arange_dims);
+ pe_key.add_shapes(ref_latents).add(ref_index_mode).add(config.ref_index_scale).add(config.theta).add(config.axes_dim);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ txt_arange_dims,
+ ref_latents,
+ ref_index_mode,
+ config.ref_index_scale,
+ config.theta,
+ config.axes_dim,
+ sd_version_is_longcat(version)));
+ int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
if (version == VERSION_CHROMA_RADIANCE) {
int patch_size = config.patch_size;
diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp
index 352ed94..223fc62 100644
--- a/src/model/diffusion/flux3.hpp
+++ b/src/model/diffusion/flux3.hpp
@@ -1,6 +1,7 @@
#ifndef __SD_MODEL_DIFFUSION_FLUX3_HPP__
#define __SD_MODEL_DIFFUSION_FLUX3_HPP__
+#include "core/pe_cache.h"
#include "model/diffusion/flux.hpp"
namespace Flux3 {
@@ -183,9 +184,9 @@ namespace Flux3 {
};
struct Flux3Runner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
Flux3Config config;
Model model;
- std::vector<float> pe_values;
sd::Tensor<float> zero = sd::Tensor<float>::zeros({1});
Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
@@ -213,34 +214,41 @@ namespace Flux3 {
auto x = make_input(*params.x);
auto text = make_input(*params.context);
GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1);
- Rope::Embedding embedding;
- auto append_image = [&](int64_t w, int64_t h, float index) {
- embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w), 1, 0.5f, 0.5f);
- for (int64_t y = 0; y < h; ++y) {
- for (int64_t x = 0; x < w; ++x) {
- embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index});
- }
- }
- };
- append_image(x->ne[0], x->ne[1], 0.f);
+ std::vector<ggml_tensor*> refs;
ggml_tensor* references = nullptr;
if (params.ref_latents && params.ref_image_params.pass_to_dit) {
- int index = 0;
for (const auto& ref : *params.ref_latents) {
auto r = make_input(ref);
- append_image(r->ne[0], r->ne[1], -float(++index));
+ refs.push_back(r);
r = DiT::patchify(compute_ctx, r, 1, 1);
references = references ? ggml_concat(compute_ctx, references, r, 1) : r;
}
}
- for (int64_t i = 0; i < text->ne[1]; ++i) {
- embedding.ids.push_back({0.f, 0.f, 0.f, static_cast<float>(i)});
- }
- embedding.positions.append_tokens(text->ne[1]);
- embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies);
- pe_values = finish_rope_pe(std::move(embedding));
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, pe_values.size() / 256);
- set_backend_tensor_data(pe, pe_values.data());
+ sd::PECacheKey pe_key("flux3.pe");
+ pe_key.add_shape(x).add_shape(text).add_shapes(refs).add(circular_x_enabled).add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ Rope::Embedding embedding;
+ auto append_image = [&](int64_t w, int64_t h, float index) {
+ embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w), 1, 0.5f, 0.5f);
+ for (int64_t y = 0; y < h; ++y) {
+ for (int64_t x = 0; x < w; ++x) {
+ embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index});
+ }
+ }
+ };
+ append_image(x->ne[0], x->ne[1], 0.f);
+ for (size_t i = 0; i < refs.size(); ++i) {
+ append_image(refs[i]->ne[0], refs[i]->ne[1], -float(i + 1));
+ }
+ for (int64_t i = 0; i < text->ne[1]; ++i) {
+ embedding.ids.push_back({0.f, 0.f, 0.f, static_cast<float>(i)});
+ }
+ embedding.positions.append_tokens(text->ne[1]);
+ embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies);
+ auto values = finish_rope_pe(std::move(embedding));
+ const int64_t tokens = static_cast<int64_t>(values.size() / 256);
+ return sd::Tensor<float>({2, 2, 64, tokens}, std::move(values));
+ }));
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);
diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp
index 54c7eff..5fc29b6 100644
--- a/src/model/diffusion/hunyuan.hpp
+++ b/src/model/diffusion/hunyuan.hpp
@@ -3,6 +3,7 @@
#include <cinttypes>
#include <memory>
+#include "core/pe_cache.h"
#include "model/common/block.hpp"
#include "model/diffusion/flux.hpp"
@@ -540,10 +541,11 @@ namespace Hunyuan {
};
struct HunyuanVideoRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
HunyuanVideoConfig config;
HunyuanVideoModel hunyuan_video;
- std::vector<float> pe_vec;
SDVersion version;
HunyuanVideoRunner(ggml_backend_t backend,
@@ -605,23 +607,24 @@ namespace Hunyuan {
if (vision != nullptr) {
text_len += static_cast<int>(vision->ne[1]);
}
- pe_vec = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
- static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- std::get<0>(config.patch_size),
- std::get<1>(config.patch_size),
- std::get<2>(config.patch_size),
- 1,
- text_len,
- config.theta,
- config.axes_dim);
- int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
- // LOG_VERBOSE("pos_len %d", pos_len);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- // pe->data = pe_vec.data();
- // print_ggml_tensor(pe, true, "pe");
- // pe->data = nullptr;
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("hunyuan.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(std::get<0>(config.patch_size));
+ pe_key.add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size)).add(text_len);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
+ static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ std::get<0>(config.patch_size),
+ std::get<1>(config.patch_size),
+ std::get<2>(config.patch_size),
+ 1,
+ text_len,
+ config.theta,
+ config.axes_dim);
+ int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp
index 3254571..664c827 100644
--- a/src/model/diffusion/ideogram4.hpp
+++ b/src/model/diffusion/ideogram4.hpp
@@ -8,6 +8,7 @@
#include <memory>
#include <string>
#include <vector>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_graph_cut.h"
@@ -422,6 +423,8 @@ namespace Ideogram4 {
};
class Ideogram4Runner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
protected:
bool should_use_uncond_model(const DiffusionParams& diffusion_params) const {
return has_uncond_model &&
@@ -436,7 +439,6 @@ namespace Ideogram4 {
Ideogram4Transformer uncond_model;
bool has_uncond_model = false;
std::string uncond_prefix;
- std::vector<float> pe_vec;
std::vector<int32_t> image_indicator_vec;
Ideogram4Runner(ggml_backend_t backend,
@@ -496,15 +498,19 @@ namespace Ideogram4 {
int64_t head_dim = config.emb_dim / config.num_heads;
auto runner_ctx = get_context();
- pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
+ sd::PECacheKey pe_key("ideogram4.pe");
+ pe_key.add(grid_h).add(grid_w).add(x->ne[3]).add(context_len).add(head_dim).add(pos_len).add(config.rope_theta);
+ pe_key.add(config.mrope_section).add(circular_x_enabled).add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
static_cast<int>(grid_w),
static_cast<int>(x->ne[3]),
static_cast<int>(context_len),
static_cast<int>(head_dim),
static_cast<int>(config.rope_theta),
config.mrope_section));
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ return sd::Tensor<float>({2, 2, head_dim / 2, pos_len}, std::move(pe_vec));
+ }));
image_indicator_vec.assign(static_cast<size_t>(pos_len), 1);
for (int64_t i = 0; i < context_len; ++i) {
diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp
index b90fbe9..3681c65 100644
--- a/src/model/diffusion/iris.hpp
+++ b/src/model/diffusion/iris.hpp
@@ -3,6 +3,7 @@
#include <array>
#include <stdexcept>
+#include "core/pe_cache.h"
#include "model/diffusion/pid.hpp"
namespace Iris {
@@ -382,11 +383,9 @@ namespace Iris {
}
struct IrisRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
IrisConfig config;
IrisModel model;
- std::vector<float> joint_pe_data;
- std::vector<float> pixel_pe_data;
- std::vector<float> pixel_pos_data;
std::vector<float> text_mask_data;
IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
@@ -418,24 +417,35 @@ namespace Iris {
GGML_ASSERT(h % config.patch_size == 0 && w % config.patch_size == 0);
GGML_ASSERT(text->ne[0] == config.text_dim * config.text_layers && text->ne[1] == nt);
GGML_ASSERT(inputs.y->numel() == nt && x->ne[3] == 1);
- joint_pe_data = Pid::make_rope_1d(nt, static_cast<int>(config.head_dim), 10000.f);
- auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim);
- joint_pe_data.insert(joint_pe_data.end(), img.begin(), img.end());
- pixel_pe_data = image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads);
- pixel_pos_data = Pid::make_pixel_abs_pos(h, w, static_cast<int>(config.pixel_dim));
+ sd::PECacheKey joint_key("iris.joint_pe");
+ joint_key.add(h).add(w).add(nt).add(config.patch_size).add(config.head_dim);
+ auto pe = make_input(pe_cache_.get_or_create(joint_key, [&]() {
+ auto values = Pid::make_rope_1d(nt, static_cast<int>(config.head_dim), 10000.f);
+ auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim);
+ values.insert(values.end(), img.begin(), img.end());
+ const int64_t tokens = nt + static_cast<int64_t>(img.size() / (2 * config.head_dim));
+ return sd::Tensor<float>({2, 2, config.head_dim / 2, tokens}, std::move(values));
+ }));
+ sd::PECacheKey pixel_key("iris.pixel_pe");
+ pixel_key.add(h).add(w).add(config.patch_size).add(config.pixel_attn_dim).add(config.pixel_heads);
+ auto pp = make_input(pe_cache_.get_or_create(pixel_key, [&]() {
+ return sd::Tensor<float>({2, 2, config.pixel_attn_dim / config.pixel_heads / 2,
+ (h / config.patch_size) * (w / config.patch_size)},
+ image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads));
+ }));
+ sd::PECacheKey pos_key("iris.pixel_pos");
+ pos_key.add(h).add(w).add(config.pixel_dim);
+ auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() {
+ return sd::Tensor<float>({config.pixel_dim, h * w, 1},
+ Pid::make_pixel_abs_pos(h, w, static_cast<int>(config.pixel_dim)));
+ }));
text_mask_data.resize(nt * nt);
for (int q = 0; q < nt; ++q) {
for (int k = 0; k < nt; ++k) {
text_mask_data[q * nt + k] = inputs.y->values()[k] != 0.f || q == k ? 0.f : -INFINITY;
}
}
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, nt + img.size() / (2 * config.head_dim));
- auto pp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.pixel_attn_dim / config.pixel_heads / 2, (h / config.patch_size) * (w / config.patch_size));
- auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.pixel_dim, h * w, 1);
auto mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nt, nt);
- set_backend_tensor_data(pe, joint_pe_data.data());
- set_backend_tensor_data(pp, pixel_pe_data.data());
- set_backend_tensor_data(pos, pixel_pos_data.data());
set_backend_tensor_data(mask, text_mask_data.data());
auto ctx = get_context();
auto out = model.forward(&ctx, x, t, text, mask, pe, pp, pos);
diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp
index 47a7a34..1558e4e 100644
--- a/src/model/diffusion/krea2.hpp
+++ b/src/model/diffusion/krea2.hpp
@@ -11,6 +11,7 @@
#include <string>
#include <utility>
#include <vector>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_graph_cut.h"
@@ -714,9 +715,9 @@ namespace Krea2 {
}
struct Krea2Runner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
Krea2Config config;
Krea2Model model;
- std::vector<float> pe_vec;
Krea2Runner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -754,18 +755,23 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor));
}
- pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- config.theta,
- config.axes_dim,
- ref_latents,
- ref_image_params.ref_index_mode));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("krea2.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
+ pe_key.add(ref_image_params.ref_index_mode);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ config.theta,
+ config.axes_dim,
+ ref_latents,
+ ref_image_params.ref_index_mode));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);
diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp
index 8fb1512..76f7bee 100644
--- a/src/model/diffusion/lens.hpp
+++ b/src/model/diffusion/lens.hpp
@@ -4,6 +4,7 @@
#include <cinttypes>
#include <memory>
#include <vector>
+#include "core/pe_cache.h"
#include "model/common/block.hpp"
#include "model/common/rope.hpp"
@@ -352,9 +353,9 @@ namespace Lens {
};
struct LensRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
LensConfig config;
LensModel lens;
- std::vector<float> pe_vec;
LensRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -384,15 +385,19 @@ namespace Lens {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
- pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("lens.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = lens.forward(&runner_ctx, x, timesteps, context, pe);
diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp
index fd43dfc..194bf79 100644
--- a/src/model/diffusion/lingbot_video.hpp
+++ b/src/model/diffusion/lingbot_video.hpp
@@ -9,6 +9,7 @@
#include <set>
#include <string>
#include <vector>
+#include "core/pe_cache.h"
#include "core/util.h"
#include "model/common/rope.hpp"
@@ -606,9 +607,9 @@ namespace LingBotVideo {
};
struct LingBotVideoRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
LingBotVideoConfig config;
LingBotVideoModel lingbot_video;
- std::vector<float> pe_vec;
LingBotVideoRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -647,19 +648,23 @@ namespace LingBotVideo {
int64_t T = x->ne[2];
int64_t N = x_tensor.shape()[4];
GGML_ASSERT(N == 1);
- pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
- static_cast<int>(H),
- static_cast<int>(W),
- config.patch_t,
- config.patch_h,
- config.patch_w,
- static_cast<int>(N),
- static_cast<int>(context->ne[1]),
- config.rope_theta,
- config.axes_dim);
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("lingbot_video.pe");
+ pe_key.add(T).add(H).add(W).add(N).add_shape(context).add(config.patch_t).add(config.patch_h).add(config.patch_w);
+ pe_key.add(config.rope_theta).add(config.axes_dim).add(config.axes_dim_sum);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
+ static_cast<int>(H),
+ static_cast<int>(W),
+ config.patch_t,
+ config.patch_h,
+ config.patch_w,
+ static_cast<int>(N),
+ static_cast<int>(context->ne[1]),
+ config.rope_theta,
+ config.axes_dim);
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);
diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp
index 85ec08e..4a9ea1e 100644
--- a/src/model/diffusion/llada_image.hpp
+++ b/src/model/diffusion/llada_image.hpp
@@ -3,6 +3,7 @@
#include <algorithm>
#include <cinttypes>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@@ -379,10 +380,11 @@ namespace LLaDAImage {
};
struct LLaDAImageRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
LLaDAImageConfig config;
LLaDAImageModel llada_image;
- std::vector<float> pe_vec;
LLaDAImageRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -412,17 +414,21 @@ namespace LLaDAImage {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
- pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- ZImage::SEQ_MULTI_OF,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("llada_image.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ ZImage::SEQ_MULTI_OF,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe);
@@ -459,17 +465,21 @@ namespace LLaDAImage {
ggml_tensor* source = make_input(source_tensor);
GGML_ASSERT(x->ne[3] == 1);
- pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(context->ne[1]),
- semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
- ZImage::SEQ_MULTI_OF,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("llada_image.pe1");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shape(semantic);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(context->ne[1]),
+ semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
+ ZImage::SEQ_MULTI_OF,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
int64_t W = x->ne[0];
diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp
index 17b882c..23c5f5f 100644
--- a/src/model/diffusion/ltxv.hpp
+++ b/src/model/diffusion/ltxv.hpp
@@ -11,6 +11,7 @@
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
#include "model/common/block.hpp"
#include "model/common/rope.hpp"
@@ -1736,14 +1737,9 @@ namespace LTXV {
};
struct LTXAVRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
LTXAVConfig config;
LTXAVModelBlock model;
- std::vector<float> video_pe_vec;
- std::vector<float> audio_pe_vec;
- std::vector<float> video_cross_pe_vec;
- std::vector<float> audio_cross_pe_vec;
- std::vector<float> connector_pe_vec;
- std::vector<float> audio_connector_pe_vec;
sd::Tensor<float> vx_input_cache;
sd::Tensor<float> ax_input_cache;
@@ -1853,79 +1849,101 @@ namespace LTXV {
float video_frame_rate = frame_rate > 0.f ? frame_rate : 24.f;
int64_t video_token_count = vx->ne[0] * vx->ne[1] * vx->ne[2];
bool has_video_positions = !video_positions_tensor.empty();
- if (has_video_positions) {
- GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count);
- video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor,
- static_cast<int>(config.hidden_size),
- static_cast<int>(config.num_attention_heads),
- config.positional_embedding_theta,
- config.positional_embedding_max_pos,
- config.use_middle_indices_grid);
- } else {
- video_pe_vec = build_video_rope_matrix(vx->ne[0],
- vx->ne[1],
- vx->ne[2],
- static_cast<int>(config.hidden_size),
- static_cast<int>(config.num_attention_heads),
- video_frame_rate,
- config.positional_embedding_theta,
- config.positional_embedding_max_pos,
- config.vae_scale_factors,
- config.causal_temporal_positioning,
- config.use_middle_indices_grid);
- }
- auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads);
- ggml_set_name(video_pe, "ltxav_video_pe");
- set_backend_tensor_data(video_pe, video_pe_vec.data());
+ sd::PECacheKey video_pe_key("ltxav_video_pe");
+ video_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor).add(config.positional_embedding_theta);
+ video_pe_key.add(config.causal_temporal_positioning).add(config.use_middle_indices_grid).add(config.hidden_size);
+ video_pe_key.add(config.num_attention_heads).add(config.attention_head_dim).add(config.positional_embedding_max_pos);
+ video_pe_key.add(std::get<0>(config.vae_scale_factors)).add(std::get<1>(config.vae_scale_factors));
+ video_pe_key.add(std::get<2>(config.vae_scale_factors));
+ auto video_pe = make_input(pe_cache_.get_or_create(video_pe_key, [&]() {
+ std::vector<float> video_pe_vec;
+ if (has_video_positions) {
+ GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count);
+ video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor,
+ static_cast<int>(config.hidden_size),
+ static_cast<int>(config.num_attention_heads),
+ config.positional_embedding_theta,
+ config.positional_embedding_max_pos,
+ config.use_middle_indices_grid);
+ } else {
+ video_pe_vec = build_video_rope_matrix(vx->ne[0],
+ vx->ne[1],
+ vx->ne[2],
+ static_cast<int>(config.hidden_size),
+ static_cast<int>(config.num_attention_heads),
+ video_frame_rate,
+ config.positional_embedding_theta,
+ config.positional_embedding_max_pos,
+ config.vae_scale_factors,
+ config.causal_temporal_positioning,
+ config.use_middle_indices_grid);
+ }
+ return sd::Tensor<float>({2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads}, std::move(video_pe_vec));
+ }));
ggml_tensor* audio_pe = nullptr;
ggml_tensor* video_cross_pe = nullptr;
ggml_tensor* audio_cross_pe = nullptr;
if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) {
- audio_pe_vec = build_audio_rope_matrix(ax->ne[1],
- static_cast<int>(config.audio_hidden_size),
- static_cast<int>(config.audio_num_attention_heads),
- config.positional_embedding_theta,
- config.audio_positional_embedding_max_pos[0],
- config.use_middle_indices_grid);
- audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
- ggml_set_name(audio_pe, "ltxav_audio_pe");
- set_backend_tensor_data(audio_pe, audio_pe_vec.data());
+ sd::PECacheKey audio_pe_key("ltxav_audio_pe");
+ audio_pe_key.add(ax->ne[1]).add(config.audio_hidden_size).add(config.audio_num_attention_heads);
+ audio_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta);
+ audio_pe_key.add(config.audio_positional_embedding_max_pos[0]).add(config.use_middle_indices_grid);
+ audio_pe = make_input(pe_cache_.get_or_create(audio_pe_key, [&]() {
+ auto audio_pe_vec = build_audio_rope_matrix(ax->ne[1],
+ static_cast<int>(config.audio_hidden_size),
+ static_cast<int>(config.audio_num_attention_heads),
+ config.positional_embedding_theta,
+ config.audio_positional_embedding_max_pos[0],
+ config.use_middle_indices_grid);
+ return sd::Tensor<float>({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_pe_vec));
+ }));
int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]);
- if (has_video_positions) {
- video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor,
- static_cast<int>(config.audio_cross_attention_dim),
- static_cast<int>(config.audio_num_attention_heads),
- config.positional_embedding_theta,
- temporal_max_pos,
- true);
- } else {
- video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0],
- vx->ne[1],
- vx->ne[2],
- static_cast<int>(config.audio_cross_attention_dim),
- static_cast<int>(config.audio_num_attention_heads),
- video_frame_rate,
- config.positional_embedding_theta,
- temporal_max_pos,
- std::get<0>(config.vae_scale_factors),
- config.causal_temporal_positioning,
- true);
- }
- video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads);
- ggml_set_name(video_cross_pe, "ltxav_video_cross_pe");
- set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data());
-
- audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1],
- static_cast<int>(config.audio_cross_attention_dim),
- static_cast<int>(config.audio_num_attention_heads),
- config.positional_embedding_theta,
- temporal_max_pos,
- true);
- audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
- ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe");
- set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data());
+ sd::PECacheKey video_cross_pe_key("ltxav_video_cross_pe");
+ video_cross_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor);
+ video_cross_pe_key.add(config.positional_embedding_theta).add(config.causal_temporal_positioning);
+ video_cross_pe_key.add(config.use_middle_indices_grid).add(config.audio_cross_attention_dim);
+ video_cross_pe_key.add(config.audio_num_attention_heads).add(config.audio_attention_head_dim);
+ video_cross_pe_key.add(temporal_max_pos).add(std::get<0>(config.vae_scale_factors));
+ video_cross_pe = make_input(pe_cache_.get_or_create(video_cross_pe_key, [&]() {
+ std::vector<float> video_cross_pe_vec;
+ if (has_video_positions) {
+ video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor,
+ static_cast<int>(config.audio_cross_attention_dim),
+ static_cast<int>(config.audio_num_attention_heads),
+ config.positional_embedding_theta,
+ temporal_max_pos,
+ true);
+ } else {
+ video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0],
+ vx->ne[1],
+ vx->ne[2],
+ static_cast<int>(config.audio_cross_attention_dim),
+ static_cast<int>(config.audio_num_attention_heads),
+ video_frame_rate,
+ config.positional_embedding_theta,
+ temporal_max_pos,
+ std::get<0>(config.vae_scale_factors),
+ config.causal_temporal_positioning,
+ true);
+ }
+ return sd::Tensor<float>({2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads}, std::move(video_cross_pe_vec));
+ }));
+
+ sd::PECacheKey audio_cross_pe_key("ltxav_audio_cross_pe");
+ audio_cross_pe_key.add(ax->ne[1]).add(config.audio_cross_attention_dim).add(config.audio_num_attention_heads);
+ audio_cross_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta);
+ audio_cross_pe_key.add(temporal_max_pos);
+ audio_cross_pe = make_input(pe_cache_.get_or_create(audio_cross_pe_key, [&]() {
+ auto audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1],
+ static_cast<int>(config.audio_cross_attention_dim),
+ static_cast<int>(config.audio_num_attention_heads),
+ config.positional_embedding_theta,
+ temporal_max_pos,
+ true);
+ return sd::Tensor<float>({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_cross_pe_vec));
+ }));
}
bool needs_video_connector_pe =
@@ -1941,10 +1959,13 @@ namespace LTXV {
int64_t target_len = std::max<int64_t>(1024, seq_len);
int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers;
int64_t full_len = seq_len + duplications * config.connector_num_registers - seq_len;
- connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.connector_hidden_size), static_cast<int>(config.connector_num_heads), 10000.f, 4096.f, true);
- video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads);
- ggml_set_name(video_connector_pe, "ltxav_video_connector_pe");
- set_backend_tensor_data(video_connector_pe, connector_pe_vec.data());
+ sd::PECacheKey video_connector_pe_key("ltxav_video_connector_pe");
+ video_connector_pe_key.add(full_len).add(config.connector_hidden_size).add(config.connector_num_heads);
+ video_connector_pe_key.add(config.connector_head_dim);
+ video_connector_pe = make_input(pe_cache_.get_or_create(video_connector_pe_key, [&]() {
+ auto connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.connector_hidden_size), static_cast<int>(config.connector_num_heads), 10000.f, 4096.f, true);
+ return sd::Tensor<float>({2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads}, std::move(connector_pe_vec));
+ }));
}
bool run_audio_context =
@@ -1961,14 +1982,17 @@ namespace LTXV {
context->ne[1] < 1024));
ggml_tensor* audio_connector_pe = nullptr;
if (needs_audio_connector_pe) {
- int64_t seq_len = context->ne[1];
- int64_t target_len = std::max<int64_t>(1024, seq_len);
- int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers;
- int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len;
- audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.audio_connector_hidden_size), static_cast<int>(config.audio_connector_num_heads), 10000.f, 4096.f, true);
- audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads);
- ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe");
- set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data());
+ int64_t seq_len = context->ne[1];
+ int64_t target_len = std::max<int64_t>(1024, seq_len);
+ int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers;
+ int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len;
+ sd::PECacheKey audio_connector_pe_key("ltxav_audio_connector_pe");
+ audio_connector_pe_key.add(full_len).add(config.audio_connector_hidden_size).add(config.audio_connector_num_heads);
+ audio_connector_pe_key.add(config.audio_connector_head_dim);
+ audio_connector_pe = make_input(pe_cache_.get_or_create(audio_connector_pe_key, [&]() {
+ auto audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.audio_connector_hidden_size), static_cast<int>(config.audio_connector_num_heads), 10000.f, 4096.f, true);
+ return sd::Tensor<float>({2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads}, std::move(audio_connector_pe_vec));
+ }));
}
auto runner_ctx = get_context();
diff --git a/src/model/diffusion/mage_flow.hpp b/src/model/diffusion/mage_flow.hpp
index 1b9b9f9..415c0f7 100644
--- a/src/model/diffusion/mage_flow.hpp
+++ b/src/model/diffusion/mage_flow.hpp
@@ -3,6 +3,7 @@
#include <cmath>
#include <memory>
+#include "core/pe_cache.h"
#include "model/diffusion/qwen_image.hpp"
@@ -57,11 +58,12 @@ namespace MageFlow {
};
struct MageFlowRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
Qwen::QwenImageConfig config;
Qwen::QwenImageModel mage_flow;
std::shared_ptr<MageFlowTimestepProjEmbeddings> time_text_embed;
- std::vector<float> pe_vec;
MageFlowRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -110,16 +112,20 @@ namespace MageFlow {
}
int batch_size = static_cast<int>(x->ne[3]);
- pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- batch_size,
- static_cast<int>(context->ne[1]),
- ref_latents,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("mage_flow.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(batch_size).add_shapes(ref_latents);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ batch_size,
+ static_cast<int>(context->ne[1]),
+ ref_latents,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
time_text_embed->prepare(timesteps_tensor);
auto runner_ctx = get_context();
diff --git a/src/model/diffusion/ming_image.hpp b/src/model/diffusion/ming_image.hpp
index 2754aa9..1de9ac1 100644
--- a/src/model/diffusion/ming_image.hpp
+++ b/src/model/diffusion/ming_image.hpp
@@ -1,6 +1,7 @@
#ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__
+#include "core/pe_cache.h"
#include "z_image.hpp"
namespace MingImage {
@@ -72,9 +73,9 @@ namespace MingImage {
};
struct MingImageRunner : DiffusionModelRunner {
+ sd::PECache pe_cache_;
MingImageConfig config;
MingImageModel model;
- std::vector<float> pe_values;
MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
@@ -109,19 +110,23 @@ namespace MingImage {
const int64_t n_txt = context->ne[1] + direct->ne[1];
const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) *
((x->ne[1] + config.patch_size - 1) / config.patch_size);
- auto padded = finish_rope_pe(Rope::gen_z_image_pe(
- static_cast<int>(x->ne[1]), static_cast<int>(x->ne[0]), config.patch_size, 1,
- static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED,
- config.theta, config.axes_dim));
- // Zero-masked alignment tokens cannot affect valid queries. Omit them while
- // retaining the padded caption length used to position image tokens.
- const size_t stride = config.axes_dim_sum * 2;
- const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF);
- pe_values.assign(padded.begin(), padded.begin() + n_txt * stride);
- pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride,
- padded.begin() + (padded_txt + n_img) * stride);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img);
- set_backend_tensor_data(pe, pe_values.data());
+ sd::PECacheKey pe_key("ming_image.pe");
+ pe_key.add_shape(x).add(n_txt).add(config.patch_size).add(config.theta).add(config.axes_dim);
+ pe_key.add(config.axes_dim_sum).add(circular_x_enabled).add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto padded = finish_rope_pe(Rope::gen_z_image_pe(
+ static_cast<int>(x->ne[1]), static_cast<int>(x->ne[0]), config.patch_size, 1,
+ static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED,
+ config.theta, config.axes_dim));
+ // Zero-masked alignment tokens cannot affect valid queries. Omit them while
+ // retaining the padded caption length used to position image tokens.
+ const size_t stride = config.axes_dim_sum * 2;
+ const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF);
+ std::vector<float> pe_values(padded.begin(), padded.begin() + n_txt * stride);
+ pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride,
+ padded.begin() + (padded_txt + n_img) * stride);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, n_txt + n_img}, std::move(pe_values));
+ }));
auto ctx = get_context();
auto out = model.forward(&ctx, x, t, context, direct, pe);
ggml_build_forward_expand(gf, out);
diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp
index 0ddacf0..6dc9c17 100644
--- a/src/model/diffusion/minit2i.hpp
+++ b/src/model/diffusion/minit2i.hpp
@@ -9,6 +9,7 @@
#include <memory>
#include <string>
#include <vector>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@@ -473,20 +474,9 @@ namespace MiniT2I {
};
struct MiniT2IRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
MiniT2IConfig config;
MMJiT model;
- ggml_context* position_cache_ctx = nullptr;
- ggml_backend_buffer_t position_cache_buffer = nullptr;
- ggml_tensor* cached_pos_embed = nullptr;
- ggml_tensor* cached_txt_pe = nullptr;
- ggml_tensor* cached_joint_pe = nullptr;
- int64_t cached_img_side = -1;
- int64_t cached_txt_len = -1;
- int64_t cached_hidden_size = -1;
- int64_t cached_head_dim = -1;
- bool cached_circular_x = false;
- bool cached_circular_y = false;
-
MiniT2IRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
@@ -497,10 +487,6 @@ namespace MiniT2I {
model.init(params_ctx, tensor_storage_map, this->prefix);
}
- ~MiniT2IRunner() override {
- free_position_cache();
- }
-
std::string get_desc() override {
return "MiniT2I";
}
@@ -509,75 +495,6 @@ namespace MiniT2I {
model.get_param_tensors(tensors, prefix);
}
- void free_position_cache() {
- if (position_cache_buffer != nullptr) {
- ggml_backend_buffer_free(position_cache_buffer);
- position_cache_buffer = nullptr;
- }
- if (position_cache_ctx != nullptr) {
- ggml_free(position_cache_ctx);
- position_cache_ctx = nullptr;
- }
- cached_pos_embed = nullptr;
- cached_txt_pe = nullptr;
- cached_joint_pe = nullptr;
- cached_img_side = -1;
- cached_txt_len = -1;
- cached_hidden_size = -1;
- cached_head_dim = -1;
- }
-
- void ensure_position_cache(int64_t img_side, int64_t txt_len) {
- if (cached_img_side == img_side &&
- cached_txt_len == txt_len &&
- cached_hidden_size == config.hidden_size &&
- cached_head_dim == config.head_dim &&
- cached_circular_x == circular_x_enabled &&
- cached_circular_y == circular_y_enabled &&
- cached_pos_embed != nullptr &&
- cached_txt_pe != nullptr &&
- cached_joint_pe != nullptr) {
- return;
- }
-
- free_position_cache();
-
- auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
- auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
- auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
- auto joint_pe_vec = txt_pe_vec;
- joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
-
- ggml_init_params params;
- params.mem_size = static_cast<size_t>(3 * ggml_tensor_overhead());
- params.mem_buffer = nullptr;
- params.no_alloc = true;
- position_cache_ctx = ggml_init(params);
- GGML_ASSERT(position_cache_ctx != nullptr);
-
- cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1);
- ggml_set_name(cached_pos_embed, "minit2i.pos_embed");
- cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len);
- ggml_set_name(cached_txt_pe, "minit2i.txt_pe");
- cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side);
- ggml_set_name(cached_joint_pe, "minit2i.joint_pe");
-
- position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend);
- GGML_ASSERT(position_cache_buffer != nullptr);
- ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
- ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed));
- ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe));
- ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe));
- ggml_backend_synchronize(runtime_backend);
-
- cached_img_side = img_side;
- cached_txt_len = txt_len;
- cached_hidden_size = config.hidden_size;
- cached_head_dim = config.head_dim;
- cached_circular_x = circular_x_enabled;
- cached_circular_y = circular_y_enabled;
- }
-
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
@@ -592,10 +509,29 @@ namespace MiniT2I {
int64_t H = x->ne[1];
int64_t img_side = H / config.patch_size;
int64_t txt_len = context->ne[1];
- ensure_position_cache(img_side, txt_len);
+ sd::PECacheKey pos_key("minit2i.pos_embed");
+ pos_key.add(img_side).add(config.hidden_size);
+ auto pos_embed = make_input(pe_cache_.get_or_create(pos_key, [&]() {
+ return sd::Tensor<float>({config.hidden_size, img_side * img_side, 1},
+ make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size)));
+ }));
+ sd::PECacheKey txt_key("minit2i.txt_pe");
+ txt_key.add(txt_len).add(config.head_dim);
+ auto txt_pe = make_input(pe_cache_.get_or_create(txt_key, [&]() {
+ return sd::Tensor<float>({2, 2, config.head_dim / 2, txt_len},
+ make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim)));
+ }));
+ sd::PECacheKey joint_key("minit2i.joint_pe");
+ joint_key.add(img_side).add(txt_len).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled);
+ auto joint_pe = make_input(pe_cache_.get_or_create(joint_key, [&]() {
+ auto values = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
+ auto image = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
+ values.insert(values.end(), image.begin(), image.end());
+ return sd::Tensor<float>({2, 2, config.head_dim / 2, txt_len + img_side * img_side}, std::move(values));
+ }));
auto runner_ctx = get_context();
- auto out = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe);
+ auto out = model.forward(&runner_ctx, x, context, mask, pos_embed, txt_pe, joint_pe);
ggml_build_forward_expand(gf, out);
return gf;
}
diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp
index 036c7a4..4ff91c4 100644
--- a/src/model/diffusion/model.hpp
+++ b/src/model/diffusion/model.hpp
@@ -173,7 +173,7 @@ struct MingImageDiffusionExtra {
struct ZImageDiffusionExtra {
const sd::Tensor<float>* control_latent = nullptr;
- float strength = 1.f;
+ float strength = 1.f;
};
using DiffusionExtraParams = std::variant<std::monostate,
diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp
index f1cbbc7..0eace33 100644
--- a/src/model/diffusion/moss_image.hpp
+++ b/src/model/diffusion/moss_image.hpp
@@ -3,6 +3,7 @@
#include <algorithm>
#include <cinttypes>
+#include "core/pe_cache.h"
#include "model/common/ggml_block.hpp"
#include "model/common/rope.hpp"
@@ -193,9 +194,10 @@ namespace MOSSImage {
};
struct MOSSImageRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
MOSSImageConfig config;
Model model;
- std::vector<float> pe_data, time_data;
+ std::vector<float> time_data;
MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) {
@@ -224,31 +226,34 @@ namespace MOSSImage {
auto text = make_input(*inputs.context);
GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim);
int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1];
- Rope::Embedding embedding;
- embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w));
- embedding.positions.images.back().height_axis = 2;
- embedding.positions.images.back().width_axis = 3;
- embedding.positions.append_tokens(nt);
- for (int64_t token = 0; token < h * w + nt; ++token) {
- float pos[4] = {0.f, float(token - h * w), 0.f, 0.f};
- if (token < h * w) {
- pos[1] = float(nt);
- pos[2] = float(token / w) - float(h - 1) / 2.f;
- pos[3] = float(token % w) - float(w - 1) / 2.f;
+ sd::PECacheKey pe_key("moss_image.pe");
+ pe_key.add(h).add(w).add(nt).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ Rope::Embedding embedding;
+ embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w));
+ embedding.positions.images.back().height_axis = 2;
+ embedding.positions.images.back().width_axis = 3;
+ embedding.positions.append_tokens(nt);
+ for (int64_t token = 0; token < h * w + nt; ++token) {
+ float pos[4] = {0.f, float(token - h * w), 0.f, 0.f};
+ if (token < h * w) {
+ pos[1] = float(nt);
+ pos[2] = float(token / w) - float(h - 1) / 2.f;
+ pos[3] = float(token % w) - float(w - 1) / 2.f;
+ }
+ embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]});
}
- embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]});
- }
- embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector<int>(4, static_cast<int>(config.head_dim / 4)), embedding.layout, &embedding.frequencies);
- pe_data = finish_rope_pe(std::move(embedding));
+ embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector<int>(4, static_cast<int>(config.head_dim / 4)), embedding.layout, &embedding.frequencies);
+ auto values = finish_rope_pe(std::move(embedding));
+ return sd::Tensor<float>({2, 2, config.head_dim / 2, h * w + nt}, std::move(values));
+ }));
time_data.resize(config.freq_dim);
for (int64_t i = 0; i < config.freq_dim / 2; ++i) {
double angle = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2));
time_data[i] = float(std::cos(angle));
time_data[i + config.freq_dim / 2] = float(std::sin(angle));
}
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, h * w + nt);
auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1);
- set_backend_tensor_data(pe, pe_data.data());
set_backend_tensor_data(time, time_data.data());
auto ctx = get_context();
ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip));
diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp
index 8aa9dc1..6c5a59b 100644
--- a/src/model/diffusion/pid.hpp
+++ b/src/model/diffusion/pid.hpp
@@ -7,6 +7,7 @@
#include <memory>
#include <string>
#include <vector>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@@ -822,12 +823,9 @@ namespace Pid {
};
struct PiDRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
PixelDiTConfig config;
PixelDiT model;
- std::vector<float> pos_img_vec;
- std::vector<float> pos_txt_vec;
- std::vector<float> pixel_pos_vec;
- std::vector<float> pixel_pos_comp_vec;
PiDRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
@@ -867,57 +865,53 @@ namespace Pid {
int64_t Hs = Hp / config.patch_size;
int64_t Ws = Wp / config.patch_size;
- pos_img_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
- static_cast<int>(Ws),
- static_cast<int>(config.hidden_size / config.num_groups),
- 10000.f,
- 16.f,
- static_cast<int>(config.rope_ref_grid_h),
- static_cast<int>(config.rope_ref_grid_w)));
- auto pos_img = ggml_new_tensor_4d(compute_ctx,
- GGML_TYPE_F32,
- 2,
- 2,
- config.hidden_size / config.num_groups / 2,
- Hs * Ws);
- set_backend_tensor_data(pos_img, pos_img_vec.data());
+ sd::PECacheKey pos_img_key("pid.pos_img");
+ pos_img_key.add(Hs).add(Ws).add(config.hidden_size).add(config.num_groups).add(config.rope_ref_grid_h);
+ pos_img_key.add(config.rope_ref_grid_w).add(circular_x_enabled).add(circular_y_enabled);
+ auto pos_img = make_input(pe_cache_.get_or_create(pos_img_key, [&]() {
+ auto pos_img_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
+ static_cast<int>(Ws),
+ static_cast<int>(config.hidden_size / config.num_groups),
+ 10000.f,
+ 16.f,
+ static_cast<int>(config.rope_ref_grid_h),
+ static_cast<int>(config.rope_ref_grid_w)));
+ return sd::Tensor<float>({2, 2, config.hidden_size / config.num_groups / 2, Hs * Ws}, std::move(pos_img_vec));
+ }));
int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
- pos_txt_vec = make_rope_1d(static_cast<int>(Ltxt),
- static_cast<int>(config.hidden_size / config.num_groups),
- config.text_rope_theta);
- auto pos_txt = ggml_new_tensor_4d(compute_ctx,
- GGML_TYPE_F32,
- 2,
- 2,
- config.hidden_size / config.num_groups / 2,
- Ltxt);
- set_backend_tensor_data(pos_txt, pos_txt_vec.data());
-
- pixel_pos_vec = make_pixel_abs_pos(static_cast<int>(Hp),
- static_cast<int>(Wp),
- static_cast<int>(config.pixel_hidden_size));
- auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
- GGML_TYPE_F32,
- config.pixel_hidden_size,
- Wp * Hp,
- 1);
- set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
-
- pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
- static_cast<int>(Ws),
- static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
- 10000.f,
- 16.f,
- static_cast<int>(config.rope_ref_grid_h),
- static_cast<int>(config.rope_ref_grid_w)));
- auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
- GGML_TYPE_F32,
- 2,
- 2,
- config.pixel_attn_hidden_size / config.pixel_num_groups / 2,
- Hs * Ws);
- set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data());
+ sd::PECacheKey pos_txt_key("pid.pos_txt");
+ pos_txt_key.add(Ltxt).add(config.hidden_size).add(config.num_groups).add(config.text_rope_theta);
+ auto pos_txt = make_input(pe_cache_.get_or_create(pos_txt_key, [&]() {
+ auto pos_txt_vec = make_rope_1d(static_cast<int>(Ltxt),
+ static_cast<int>(config.hidden_size / config.num_groups),
+ config.text_rope_theta);
+ return sd::Tensor<float>({2, 2, config.hidden_size / config.num_groups / 2, Ltxt}, std::move(pos_txt_vec));
+ }));
+
+ sd::PECacheKey pixel_pos_key("pid.pixel_pos");
+ pixel_pos_key.add(Hp).add(Wp).add(config.pixel_hidden_size);
+ auto pixel_pos = make_input(pe_cache_.get_or_create(pixel_pos_key, [&]() {
+ auto pixel_pos_vec = make_pixel_abs_pos(static_cast<int>(Hp),
+ static_cast<int>(Wp),
+ static_cast<int>(config.pixel_hidden_size));
+ return sd::Tensor<float>({config.pixel_hidden_size, Wp * Hp, 1}, std::move(pixel_pos_vec));
+ }));
+
+ sd::PECacheKey pixel_pos_comp_key("pid.pixel_pos_comp");
+ pixel_pos_comp_key.add(Hs).add(Ws).add(config.pixel_attn_hidden_size).add(config.pixel_num_groups);
+ pixel_pos_comp_key.add(config.rope_ref_grid_h).add(config.rope_ref_grid_w).add(circular_x_enabled);
+ pixel_pos_comp_key.add(circular_y_enabled);
+ auto pixel_pos_comp = make_input(pe_cache_.get_or_create(pixel_pos_comp_key, [&]() {
+ auto pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
+ static_cast<int>(Ws),
+ static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
+ 10000.f,
+ 16.f,
+ static_cast<int>(config.rope_ref_grid_h),
+ static_cast<int>(config.rope_ref_grid_w)));
+ return sd::Tensor<float>({2, 2, config.pixel_attn_hidden_size / config.pixel_num_groups / 2, Hs * Ws}, std::move(pixel_pos_comp_vec));
+ }));
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx,
diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp
index c7e29b7..45f9e43 100644
--- a/src/model/diffusion/pixart.hpp
+++ b/src/model/diffusion/pixart.hpp
@@ -5,6 +5,7 @@
#include <cinttypes>
#include <cmath>
#include <vector>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@@ -295,9 +296,9 @@ namespace PixArt {
};
struct PixArtRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
PixArtConfig config;
PixArtModel model;
- std::vector<float> pos_vec;
PixArtRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -357,10 +358,13 @@ namespace PixArt {
int64_t wp = W / config.patch_size;
int64_t hp = H / config.patch_size;
- pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp,
- config.pos_embed_base_size, config.interpolation_scale);
- auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1);
- set_backend_tensor_data(pos, pos_vec.data());
+ sd::PECacheKey pos_key("pixart.pos");
+ pos_key.add(hp).add(wp).add(config.hidden_size).add(config.pos_embed_base_size).add(config.interpolation_scale);
+ auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() {
+ auto pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp,
+ config.pos_embed_base_size, config.interpolation_scale);
+ return sd::Tensor<float>({config.hidden_size, wp * hp, 1}, std::move(pos_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask);
diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp
index 82d1730..4c473da 100644
--- a/src/model/diffusion/qwen_image.hpp
+++ b/src/model/diffusion/qwen_image.hpp
@@ -4,6 +4,7 @@
#include <memory>
#include "core/ggml_extend_backend.h"
#include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
#include "core/util.h"
#include "model/common/block.hpp"
@@ -561,10 +562,11 @@ namespace Qwen {
};
struct QwenImageRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
QwenImageConfig config;
QwenImageModel qwen_image;
- std::vector<float> pe_vec;
std::vector<float> modulate_index_vec;
std::vector<int32_t> additional_t_cond_vec;
SDVersion version;
@@ -635,23 +637,24 @@ namespace Qwen {
ref_index_mode = Rope::RefIndexMode::DECREASE;
}
- pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
- static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- batch_size,
- static_cast<int>(context->ne[1]),
- ref_latents,
- ref_index_mode,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- // LOG_VERBOSE("pos_len %d", pos_len);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- // pe->data = pe_vec.data();
- // print_ggml_tensor(pe, true, "pe");
- // pe->data = nullptr;
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("qwen_image.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(time_len).add(batch_size).add(config.patch_size);
+ pe_key.add_shapes(ref_latents).add(ref_index_mode);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
+ static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ batch_size,
+ static_cast<int>(context->ne[1]),
+ ref_latents,
+ ref_index_mode,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
ggml_tensor* modulate_index = nullptr;
if (config.zero_cond_t) {
diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp
index ec43477..042544a 100644
--- a/src/model/diffusion/qwen_image_2_1.hpp
+++ b/src/model/diffusion/qwen_image_2_1.hpp
@@ -1,6 +1,7 @@
#ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
#define __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
+#include "core/pe_cache.h"
#include "model/diffusion/qwen_image.hpp"
namespace Qwen {
@@ -370,9 +371,9 @@ namespace Qwen {
};
struct QwenImage21Runner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
QwenImage21Config config;
QwenImage21Model model;
- std::vector<float> pe_data;
std::vector<sd::Tensor<float>> mask_data;
ggml_type prefix_cache_type = GGML_TYPE_COUNT;
bool prefix_cache_enabled = true;
@@ -504,20 +505,7 @@ namespace Qwen {
if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
checked_cache.flash_attn_used = &flash_attn_used;
}
- const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
- const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
- Rope::Embedding embedding;
- embedding.ids.assign(first_position, layout.positions.end());
- const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
- embedding.positions.token_count = embedding.ids.size();
- for (auto region : layout.rope_layout.images) {
- if (region.begin >= offset) {
- region.begin -= offset;
- embedding.positions.images.push_back(region);
- }
- }
- embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
- pe_data = finish_rope_pe(std::move(embedding));
+ const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
mask_data.clear();
if (!cached) {
for (const auto& segment : layout.segments) {
@@ -535,9 +523,30 @@ namespace Qwen {
}
auto build = [&]() {
auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
- layout.positions.size() - (cached ? layout.prefix_length : 0));
- set_backend_tensor_data(pe, pe_data.data());
+ sd::PECacheKey pe_key("qwen_image_2_1.pe");
+ pe_key.add(layout.positions).add(cached).add(layout.prefix_length).add(config.axes_dim).add(config.head_dim);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled);
+ for (const auto& region : layout.rope_layout.images) {
+ pe_key.add(region.begin).add(region.count).add(region.height_period).add(region.width_period);
+ pe_key.add(region.height_axis).add(region.width_axis);
+ }
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
+ Rope::Embedding embedding;
+ embedding.ids.assign(first_position, layout.positions.end());
+ const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
+ embedding.positions.token_count = embedding.ids.size();
+ for (auto region : layout.rope_layout.images) {
+ if (region.begin >= offset) {
+ region.begin -= offset;
+ embedding.positions.images.push_back(region);
+ }
+ }
+ embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
+ auto values = finish_rope_pe(std::move(embedding));
+ const int64_t tokens = static_cast<int64_t>(layout.positions.size()) - (cached ? layout.prefix_length : 0);
+ return sd::Tensor<float>({2, 2, config.head_dim / 2, tokens}, std::move(values));
+ }));
std::vector<ggml_tensor*> masks, ref_inputs;
for (const auto& mask : mask_data) {
masks.push_back(mask.empty() ? nullptr : make_input(mask));
diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp
index 4b75203..3c76c4e 100644
--- a/src/model/diffusion/wan.hpp
+++ b/src/model/diffusion/wan.hpp
@@ -8,6 +8,7 @@
#include <utility>
#include "core/ggml_extend_backend.h"
#include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
#include "model/common/block.hpp"
#include "model/common/rope.hpp"
@@ -885,11 +886,12 @@ namespace WAN {
};
struct WanRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
std::string desc = "wan";
WanConfig config;
Wan wan;
- std::vector<float> pe_vec;
SDVersion version;
WanRunner(ggml_backend_t backend,
@@ -1000,37 +1002,36 @@ namespace WAN {
ggml_tensor* audio_embed = make_optional_input(audio_embed_tensor);
ggml_tensor* ref_latent = make_optional_input(ref_latent_tensor);
- pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
- static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- std::get<0>(config.patch_size),
- std::get<1>(config.patch_size),
- std::get<2>(config.patch_size),
- 1,
- config.theta,
- config.axes_dim);
- if (ref_latent != nullptr) {
- // Match S2V's reference-frame temporal offset.
- int t_start = std::max(30, static_cast<int>(x->ne[2]) + 9);
- auto ref_pe = Rope::gen_wan_pe(static_cast<int>(ref_latent->ne[2]),
- static_cast<int>(ref_latent->ne[1]),
- static_cast<int>(ref_latent->ne[0]),
+ sd::PECacheKey pe_key("wan.pe");
+ pe_key.add_shape(x).add_shape(ref_latent).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(std::get<0>(config.patch_size)).add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size));
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
+ static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size),
1,
config.theta,
- config.axes_dim,
- t_start);
- pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end());
- }
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- // LOG_VERBOSE("pos_len %d", pos_len);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- // pe->data = pe_vec.data();
- // print_ggml_tensor(pe);
- // pe->data = nullptr;
- set_backend_tensor_data(pe, pe_vec.data());
+ config.axes_dim);
+ if (ref_latent != nullptr) {
+ int t_start = std::max(30, static_cast<int>(x->ne[2]) + 9);
+ auto ref_pe = Rope::gen_wan_pe(static_cast<int>(ref_latent->ne[2]),
+ static_cast<int>(ref_latent->ne[1]),
+ static_cast<int>(ref_latent->ne[0]),
+ std::get<0>(config.patch_size),
+ std::get<1>(config.patch_size),
+ std::get<2>(config.patch_size),
+ 1,
+ config.theta,
+ config.axes_dim,
+ t_start);
+ pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end());
+ }
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
if (c_concat != nullptr) {
x = ggml_concat(compute_ctx, x, c_concat, 3);
diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp
index b479507..229c6aa 100644
--- a/src/model/diffusion/z_image.hpp
+++ b/src/model/diffusion/z_image.hpp
@@ -3,6 +3,7 @@
#include <algorithm>
#include <cinttypes>
+#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_extend_backend.h"
@@ -23,10 +24,10 @@
#endif
namespace ZImage {
- constexpr int Z_IMAGE_GRAPH_SIZE = 24576;
- constexpr int ADALN_EMBED_DIM = 256;
- constexpr int SEQ_MULTI_OF = 32;
- constexpr int CONTROL_IN_CHANNELS = 33;
+ constexpr int Z_IMAGE_GRAPH_SIZE = 24576;
+ constexpr int ADALN_EMBED_DIM = 256;
+ constexpr int SEQ_MULTI_OF = 32;
+ constexpr int CONTROL_IN_CHANNELS = 33;
struct ZImageConfig {
int patch_size = 2;
@@ -649,15 +650,15 @@ namespace ZImage {
for (int i = 0; i < control_num_refiner_layers; i++) {
std::string name = "control_noise_refiner." + std::to_string(i);
auto block = std::make_shared<ZImageControlTransformerBlock>(i,
- config.hidden_size,
- config.head_dim,
- config.num_heads,
- config.num_kv_heads,
- config.multiple_of,
- config.ffn_dim_multiplier,
- config.norm_eps,
- config.qk_norm,
- tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end());
+ config.hidden_size,
+ config.head_dim,
+ config.num_heads,
+ config.num_kv_heads,
+ config.multiple_of,
+ config.ffn_dim_multiplier,
+ config.norm_eps,
+ config.qk_norm,
+ tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end());
control_blocks[name] = block;
}
@@ -665,15 +666,15 @@ namespace ZImage {
for (int i = 0; i < control_num_layers; i++) {
std::string name = "control_layers." + std::to_string(i);
auto block = std::make_shared<ZImageControlTransformerBlock>(i,
- config.hidden_size,
- config.head_dim,
- config.num_heads,
- config.num_kv_heads,
- config.multiple_of,
- config.ffn_dim_multiplier,
- config.norm_eps,
- config.qk_norm,
- tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end());
+ config.hidden_size,
+ config.head_dim,
+ config.num_heads,
+ config.num_kv_heads,
+ config.multiple_of,
+ config.ffn_dim_multiplier,
+ config.norm_eps,
+ config.qk_norm,
+ tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end());
control_blocks[name] = block;
}
@@ -705,7 +706,7 @@ namespace ZImage {
if (n_pad_token > 0) {
auto img_pad_token = params["x_pad_token"];
auto ctrl_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_pad_token, N, 1);
- c = ggml_concat(ctx->ggml_ctx, c, ctrl_pad_tokens, 1);
+ c = ggml_concat(ctx->ggml_ctx, c, ctrl_pad_tokens, 1);
}
for (int i = 0; i < control_num_refiner_layers; i++) {
@@ -755,7 +756,7 @@ namespace ZImage {
for (int i = 0; i < control_num_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImageControlTransformerBlock>(control_blocks["control_layers." + std::to_string(i)]);
- auto out = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb);
+ auto out = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb);
ctrl_unified = out.second;
main_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control");
@@ -769,8 +770,8 @@ namespace ZImage {
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
- ggml_tensor* control = nullptr,
- float strength = 0.f) {
+ ggml_tensor* control = nullptr,
+ float strength = 0.f) {
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
@@ -869,7 +870,7 @@ namespace ZImage {
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {},
- ggml_tensor* control_latent = nullptr,
+ ggml_tensor* control_latent = nullptr,
float control_strength = 0.f) {
// Forward pass of DiT.
// x: [N, C, H, W]
@@ -916,10 +917,11 @@ namespace ZImage {
};
struct ZImageRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
+
public:
ZImageConfig config;
ZImageModel z_image;
- std::vector<float> pe_vec;
std::vector<float> timestep_vec;
SDVersion version;
@@ -953,10 +955,10 @@ namespace ZImage {
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
- const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
- Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
- const sd::Tensor<float>& control_latent_tensor = {},
- float control_strength = 0.f) {
+ const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
+ Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
+ const sd::Tensor<float>& control_latent_tensor = {},
+ float control_strength = 0.f) {
ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -971,23 +973,24 @@ namespace ZImage {
ggml_tensor* control_latent = make_optional_input(control_latent_tensor);
- pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- SEQ_MULTI_OF,
- ref_latents,
- ref_index_mode,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- // LOG_VERBOSE("pos_len %d", pos_len);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- // pe->data = pe_vec.data();
- // print_ggml_tensor(pe, true, "pe");
- // pe->data = nullptr;
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("z_image.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
+ pe_key.add(ref_index_mode);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ SEQ_MULTI_OF,
+ ref_latents,
+ ref_index_mode,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto runner_ctx = get_context();
ggml_tensor* out = z_image.forward(&runner_ctx,
@@ -1011,7 +1014,7 @@ namespace ZImage {
const std::vector<sd::Tensor<float>>& ref_latents = {},
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
const sd::Tensor<float>& control_latent = {},
- float control_strength = 0.f) {
+ float control_strength = 0.f) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
@@ -1029,10 +1032,10 @@ namespace ZImage {
static const std::vector<sd::Tensor<float>> empty_ref_latents;
static const sd::Tensor<float> empty_control_latent;
- const auto* z_image_extra = std::get_if<ZImageDiffusionExtra>(&diffusion_params.extra);
+ const auto* z_image_extra = std::get_if<ZImageDiffusionExtra>(&diffusion_params.extra);
const sd::Tensor<float>& control_latent = (z_image_extra != nullptr && z_image_extra->control_latent != nullptr) ? *z_image_extra->control_latent
- : empty_control_latent;
- float control_strength = z_image_extra != nullptr ? z_image_extra->strength : 0.f;
+ : empty_control_latent;
+ float control_strength = z_image_extra != nullptr ? z_image_extra->strength : 0.f;
return compute(n_threads,
*diffusion_params.x,
diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp
index 0eb0b14..f44f2ff 100644
--- a/src/model/diffusion/z_image_l2p.hpp
+++ b/src/model/diffusion/z_image_l2p.hpp
@@ -2,6 +2,7 @@
#define __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__
#include <cmath>
+#include "core/pe_cache.h"
#include "z_image.hpp"
@@ -183,9 +184,9 @@ namespace ZImageL2P {
};
struct ZImageL2PRunner : public DiffusionModelRunner {
+ sd::PECache pe_cache_;
ZImageL2PConfig config;
ZImageL2PModel model;
- std::vector<float> pe_vec;
ZImageL2PRunner(ggml_backend_t backend,
const String2TensorStorage& tensors,
@@ -221,19 +222,23 @@ namespace ZImageL2P {
auto context = make_input(*inputs.context);
GGML_ASSERT(x->ne[3] == 1);
GGML_ASSERT(x->ne[0] % config.patch_size == 0 && x->ne[1] % config.patch_size == 0);
- pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
- static_cast<int>(x->ne[0]),
- config.patch_size,
- static_cast<int>(x->ne[3]),
- static_cast<int>(context->ne[1]),
- ZImage::SEQ_MULTI_OF,
- {},
- Rope::RefIndexMode::FIXED,
- config.theta,
- config.axes_dim));
- int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
- auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
- set_backend_tensor_data(pe, pe_vec.data());
+ sd::PECacheKey pe_key("z_image_l2p.pe");
+ pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+ pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size);
+ auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+ auto pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
+ static_cast<int>(x->ne[0]),
+ config.patch_size,
+ static_cast<int>(x->ne[3]),
+ static_cast<int>(context->ne[1]),
+ ZImage::SEQ_MULTI_OF,
+ {},
+ Rope::RefIndexMode::FIXED,
+ config.theta,
+ config.axes_dim));
+ int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+ return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+ }));
auto ctx = get_context();
auto out = model.forward(&ctx, x, t, context, pe);
ggml_build_forward_expand(gf, out);
diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp
index f44580c..d7a0f00 100644
--- a/src/pipeline/diffusion_engine.cpp
+++ b/src/pipeline/diffusion_engine.cpp
@@ -1150,8 +1150,8 @@ bool StableDiffusionGGML::build_control_net_runner() {
return false;
}
ModelManager::ResidencyMode residency_mode = backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION)
- ? ModelManager::ResidencyMode::Disk
- : ModelManager::ResidencyMode::ParamBackend;
+ ? ModelManager::ResidencyMode::Disk
+ : ModelManager::ResidencyMode::ParamBackend;
return model_manager->register_param_tensors(ModelComponent::ControlNet,
std::move(control_tensors),
residency_mode,
diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp
index c549280..7daec86 100644
--- a/src/pipeline/image.cpp
+++ b/src/pipeline/image.cpp
@@ -130,15 +130,15 @@ namespace sd::pipeline {
}
// TAE encode already returns the latent; KL/Flux VAE returns a mean/logvar pair.
sd::Tensor<float> ctrl_latent = sd->first_stage_model->get_desc() == "taesd"
- ? vae_output
- : sd::ops::chunk(vae_output, 2, 2)[0];
- ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent);
+ ? vae_output
+ : sd::ops::chunk(vae_output, 2, 2)[0];
+ ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent);
if (ctrl_latent.empty() || ctrl_latent.dim() != 4 || ctrl_latent.shape()[2] > 33) {
LOG_ERROR("failed to encode control image");
return {};
}
std::vector<int64_t> padded_shape = ctrl_latent.shape();
- padded_shape[2] = 33;
+ padded_shape[2] = 33;
sd::Tensor<float> padded(padded_shape);
sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent);
return padded;
@@ -1028,7 +1028,7 @@ namespace sd::pipeline {
std::vector<int64_t> ctrl_shape = latents.control_image.shape();
ctrl_shape[0] = hires_latent_w;
ctrl_shape[1] = hires_latent_h;
- hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic);
+ hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic);
}
} else {
hires_control_image = latents.control_image;