From ea209dbfcf1ab96ca25e3612820dc94bbc66f00d Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 25 Oct 2026 05:18:30 +0800 Subject: [PATCH] perf: cache CPU position embeddings across DiT runners (#2211) --- docs/pe_cache.md | 60 ++++++++ src/core/pe_cache.cpp | 30 ++++ src/core/pe_cache.h | 82 +++++++++++ src/model/diffusion/anima.hpp | 56 ++++---- src/model/diffusion/boogu.hpp | 29 ++-- src/model/diffusion/ernie_image.hpp | 27 ++-- src/model/diffusion/fibo2.hpp | 42 +++--- src/model/diffusion/flux.hpp | 43 +++--- src/model/diffusion/flux3.hpp | 50 ++++--- src/model/diffusion/hunyuan.hpp | 39 ++--- src/model/diffusion/ideogram4.hpp | 14 +- src/model/diffusion/iris.hpp | 38 +++-- src/model/diffusion/krea2.hpp | 32 +++-- src/model/diffusion/lens.hpp | 25 ++-- src/model/diffusion/lingbot_video.hpp | 33 +++-- src/model/diffusion/llada_image.hpp | 56 +++++--- src/model/diffusion/ltxv.hpp | 190 ++++++++++++++----------- src/model/diffusion/mage_flow.hpp | 28 ++-- src/model/diffusion/ming_image.hpp | 33 +++-- src/model/diffusion/minit2i.hpp | 110 +++----------- src/model/diffusion/model.hpp | 2 +- src/model/diffusion/moss_image.hpp | 41 +++--- src/model/diffusion/pid.hpp | 96 ++++++------- src/model/diffusion/pixart.hpp | 14 +- src/model/diffusion/qwen_image.hpp | 39 ++--- src/model/diffusion/qwen_image_2_1.hpp | 45 +++--- src/model/diffusion/wan.hpp | 55 +++---- src/model/diffusion/z_image.hpp | 109 +++++++------- src/model/diffusion/z_image_l2p.hpp | 33 +++-- src/pipeline/diffusion_engine.cpp | 4 +- src/pipeline/image.cpp | 10 +- 31 files changed, 855 insertions(+), 610 deletions(-) create mode 100644 docs/pe_cache.md create mode 100644 src/core/pe_cache.cpp create mode 100644 src/core/pe_cache.h diff --git a/docs/pe_cache.md b/docs/pe_cache.md new file mode 100644 index 00000000..e1819599 --- /dev/null +++ b/docs/pe_cache.md @@ -0,0 +1,60 @@ +# Position embedding cache + +`sd::PECache` stores CPU-side `sd::Tensor` 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({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)` 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 00000000..df4aa62c --- /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& tensors) { + add(tensors.size()); + for (const auto* tensor : tensors) { + add_shape(tensor); + } + return *this; + } + + const Tensor& PECache::get_or_create(const PECacheKey& key, const std::function()>& 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 00000000..3539707b --- /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 +#include +#include +#include + +#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 || std::is_enum_v, int> = 0> + PECacheKey& add(T value) { + if constexpr (std::is_floating_point_v) { + const double field = static_cast(value); + fields_.push_back('f'); + fields_.append(reinterpret_cast(&field), sizeof(field)); + } else { + const uint64_t field = static_cast(value); + fields_.push_back('i'); + fields_.append(reinterpret_cast(&field), sizeof(field)); + } + return *this; + } + + template + PECacheKey& add(const std::vector& values) { + fields_.push_back('v'); + add(values.size()); + for (const auto& value : values) { + add(value); + } + return *this; + } + + template + PECacheKey& add(const std::set& values) { + return add(std::vector(values.begin(), values.end())); + } + + template + PECacheKey& add(const Tensor& 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& tensors); + + bool operator==(const PECacheKey& other) const { + return name_ == other.name_ && fields_ == other.fields_; + } + }; + + class PECache { + struct Entry { + PECacheKey key; + Tensor value; + }; + + std::map entries_; + + public: + const Tensor& get_or_create(const PECacheKey& key, const std::function()>& build); + }; + +} + +#endif // __SD_CORE_PE_CACHE_H__ diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index a414d96e..66e1a0ba 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -7,6 +7,7 @@ #include #include #include +#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 image_pe_vec; - std::vector adapter_q_pe_vec; - std::vector 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(h_pad), - static_cast(w_pad), - static_cast(config.patch_size), - config.theta, - config.axes_dim, - 4.0f, - 4.0f, - 1.0f, - ref_latents)); - int64_t image_pos_len = static_cast(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(h_pad), + static_cast(w_pad), + static_cast(config.patch_size), + config.theta, + config.axes_dim, + 4.0f, + 4.0f, + 1.0f, + ref_latents)); + int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); + return sd::Tensor({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(adapter_q_pe_vec.size()) / (2 * 2 * 32); - int64_t source_pos_len = static_cast(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({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({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 3ee6714f..ac67ff9c 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -6,6 +6,7 @@ #include #include #include +#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 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 b284cb83..d4f774fc 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -4,6 +4,7 @@ #include #include #include +#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 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 f53b6ecd..a642cb4d 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 pe_values; Fibo2Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr 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(p.context->shape()[1])); - Rope::Embedding embedding; - for (int i = 0; i < count; ++i) { - embedding.ids.push_back({static_cast(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(image->ne[1]), static_cast(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(config.max_gist + 1)); std::vector 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(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(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(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(image->ne[1]), static_cast(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(config.max_gist + 1)); + for (size_t i = 0; i < refs.size(); ++i) { + append_image(refs[i], static_cast(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(values.size() / 256); + return sd::Tensor({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 825dabe1..2f4272eb 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -6,6 +6,7 @@ #include #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 pe_vec; std::vector mod_index_arange_vec; std::vector dct_vec; sd::Tensor 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(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(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(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(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 352ed949..223fc62d 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 pe_values; sd::Tensor zero = sd::Tensor::zeros({1}); Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr 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(h), static_cast(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 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(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(h), static_cast(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(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(values.size() / 256); + return sd::Tensor({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 54c7eff1..5fc29b6b 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -3,6 +3,7 @@ #include #include +#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 pe_vec; SDVersion version; HunyuanVideoRunner(ggml_backend_t backend, @@ -605,23 +607,24 @@ namespace Hunyuan { if (vision != nullptr) { text_len += static_cast(vision->ne[1]); } - pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(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(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(x->ne[2]), + static_cast(x->ne[1]), + static_cast(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(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 3254571d..664c827d 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -8,6 +8,7 @@ #include #include #include +#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 pe_vec; std::vector 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(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(grid_h), static_cast(grid_w), static_cast(x->ne[3]), static_cast(context_len), static_cast(head_dim), static_cast(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({2, 2, head_dim / 2, pos_len}, std::move(pe_vec)); + })); image_indicator_vec.assign(static_cast(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 b90fbe9d..3681c65a 100644 --- a/src/model/diffusion/iris.hpp +++ b/src/model/diffusion/iris.hpp @@ -3,6 +3,7 @@ #include #include +#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 joint_pe_data; - std::vector pixel_pe_data; - std::vector pixel_pos_data; std::vector text_mask_data; IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr 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(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(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(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(img.size() / (2 * config.head_dim)); + return sd::Tensor({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({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({config.pixel_dim, h * w, 1}, + Pid::make_pixel_abs_pos(h, w, static_cast(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 47a7a340..1558e4e9 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -11,6 +11,7 @@ #include #include #include +#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 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim, - ref_latents, - ref_image_params.ref_index_mode)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim, + ref_latents, + ref_image_params.ref_index_mode)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 8fb1512b..76f7bee1 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -4,6 +4,7 @@ #include #include #include +#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 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(x->ne[1]), - static_cast(x->ne[0]), - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 fd43dfc4..194bf791 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -9,6 +9,7 @@ #include #include #include +#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 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(T), - static_cast(H), - static_cast(W), - config.patch_t, - config.patch_h, - config.patch_w, - static_cast(N), - static_cast(context->ne[1]), - config.rope_theta, - config.axes_dim); - int pos_len = static_cast(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(T), + static_cast(H), + static_cast(W), + config.patch_t, + config.patch_h, + config.patch_w, + static_cast(N), + static_cast(context->ne[1]), + config.rope_theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 85ec08ea..4a9ea1e5 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -3,6 +3,7 @@ #include #include +#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 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(context->ne[1]), - semantic != nullptr ? static_cast(semantic->ne[1]) : 0, - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + semantic != nullptr ? static_cast(semantic->ne[1]) : 0, + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 17b882cf..23c5f5f2 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -11,6 +11,7 @@ #include #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 video_pe_vec; - std::vector audio_pe_vec; - std::vector video_cross_pe_vec; - std::vector audio_cross_pe_vec; - std::vector connector_pe_vec; - std::vector audio_connector_pe_vec; sd::Tensor vx_input_cache; sd::Tensor 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(config.hidden_size), - static_cast(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(config.hidden_size), - static_cast(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 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(config.hidden_size), + static_cast(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(config.hidden_size), + static_cast(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({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(config.audio_hidden_size), - static_cast(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(config.audio_hidden_size), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + config.audio_positional_embedding_max_pos[0], + config.use_middle_indices_grid); + return sd::Tensor({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(config.audio_cross_attention_dim), - static_cast(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(config.audio_cross_attention_dim), - static_cast(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()); + 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 video_cross_pe_vec; + if (has_video_positions) { + video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor, + static_cast(config.audio_cross_attention_dim), + static_cast(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(config.audio_cross_attention_dim), + static_cast(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({2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads}, std::move(video_cross_pe_vec)); + })); - audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(config.audio_cross_attention_dim), - static_cast(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 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(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + temporal_max_pos, + true); + return sd::Tensor({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(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(config.connector_hidden_size), static_cast(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(config.connector_hidden_size), static_cast(config.connector_num_heads), 10000.f, 4096.f, true); + return sd::Tensor({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(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(config.audio_connector_hidden_size), static_cast(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(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(config.audio_connector_hidden_size), static_cast(config.audio_connector_num_heads), 10000.f, 4096.f, true); + return sd::Tensor({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 1b9b9f9b..415c0f72 100644 --- a/src/model/diffusion/mage_flow.hpp +++ b/src/model/diffusion/mage_flow.hpp @@ -3,6 +3,7 @@ #include #include +#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 time_text_embed; - std::vector pe_vec; MageFlowRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -110,16 +112,20 @@ namespace MageFlow { } int batch_size = static_cast(x->ne[3]); - pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - batch_size, - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + batch_size, + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 2754aa94..1de9ac15 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 pe_values; MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr 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(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, - static_cast(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(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(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, + static_cast(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(n_txt), ZImage::SEQ_MULTI_OF); + std::vector 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({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 0ddacf02..6dc9c179 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -9,6 +9,7 @@ #include #include #include +#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(img_side), static_cast(config.hidden_size)); - auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); - auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(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(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& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& 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({config.hidden_size, img_side * img_side, 1}, + make_2d_sincos_pos_embed(static_cast(img_side), static_cast(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({2, 2, config.head_dim / 2, txt_len}, + make_text_rope(static_cast(txt_len), static_cast(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(txt_len), static_cast(config.head_dim)); + auto image = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); + values.insert(values.end(), image.begin(), image.end()); + return sd::Tensor({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 036c7a43..4ff91c41 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -173,7 +173,7 @@ struct MingImageDiffusionExtra { struct ZImageDiffusionExtra { const sd::Tensor* control_latent = nullptr; - float strength = 1.f; + float strength = 1.f; }; using DiffusionExtraParams = std::variant #include +#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 pe_data, time_data; + std::vector time_data; MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr 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(h), static_cast(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(h), static_cast(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(4, static_cast(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(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + return sd::Tensor({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 8aa9dc16..6c5a59be 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -7,6 +7,7 @@ #include #include #include +#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 pos_img_vec; - std::vector pos_txt_vec; - std::vector pixel_pos_vec; - std::vector 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(Hs), - static_cast(Ws), - static_cast(config.hidden_size / config.num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(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(Hs), + static_cast(Ws), + static_cast(config.hidden_size / config.num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); + return sd::Tensor({2, 2, config.hidden_size / config.num_groups / 2, Hs * Ws}, std::move(pos_img_vec)); + })); int64_t Ltxt = std::min(context->ne[1], config.txt_max_length); - pos_txt_vec = make_rope_1d(static_cast(Ltxt), - static_cast(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()); + 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(Ltxt), + static_cast(config.hidden_size / config.num_groups), + config.text_rope_theta); + return sd::Tensor({2, 2, config.hidden_size / config.num_groups / 2, Ltxt}, std::move(pos_txt_vec)); + })); - pixel_pos_vec = make_pixel_abs_pos(static_cast(Hp), - static_cast(Wp), - static_cast(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()); + 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(Hp), + static_cast(Wp), + static_cast(config.pixel_hidden_size)); + return sd::Tensor({config.pixel_hidden_size, Wp * Hp, 1}, std::move(pixel_pos_vec)); + })); - pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(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 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(Hs), + static_cast(Ws), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); + return sd::Tensor({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 c7e29b78..45f9e431 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -5,6 +5,7 @@ #include #include #include +#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 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({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 82d17300..4c473da3 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -4,6 +4,7 @@ #include #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 pe_vec; std::vector modulate_index_vec; std::vector 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - batch_size, - static_cast(context->ne[1]), - ref_latents, - ref_index_mode, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + batch_size, + static_cast(context->ne[1]), + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 ec434773..042544ac 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 pe_data; std::vector> 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(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(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(layout.positions.size()) - (cached ? layout.prefix_length : 0); + return sd::Tensor({2, 2, config.head_dim / 2, tokens}, std::move(values)); + })); std::vector 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 4b75203d..3c76c4ea 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -8,6 +8,7 @@ #include #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 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(x->ne[2]), - static_cast(x->ne[1]), - static_cast(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(x->ne[2]) + 9); - auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), - static_cast(ref_latent->ne[1]), - static_cast(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(x->ne[2]), + static_cast(x->ne[1]), + static_cast(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(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(x->ne[2]) + 9); + auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), + static_cast(ref_latent->ne[1]), + static_cast(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(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 b4795073..229c6aae 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -3,6 +3,7 @@ #include #include +#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(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(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(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(blocks["x_embedder"]); auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); @@ -869,7 +870,7 @@ namespace ZImage { ggml_tensor* context, ggml_tensor* pe, std::vector 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 pe_vec; std::vector timestep_vec; SDVersion version; @@ -953,10 +955,10 @@ namespace ZImage { ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, - const std::vector>& ref_latents_tensor = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, - const sd::Tensor& control_latent_tensor = {}, - float control_strength = 0.f) { + const std::vector>& ref_latents_tensor = {}, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - SEQ_MULTI_OF, - ref_latents, - ref_index_mode, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + SEQ_MULTI_OF, + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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>& ref_latents = {}, Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, const sd::Tensor& 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> empty_ref_latents; static const sd::Tensor empty_control_latent; - const auto* z_image_extra = std::get_if(&diffusion_params.extra); + const auto* z_image_extra = std::get_if(&diffusion_params.extra); const sd::Tensor& 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 0eb0b14d..f44f2ff9 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 +#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 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - {}, - Rope::RefIndexMode::FIXED, - config.theta, - config.axes_dim)); - int pos_len = static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + {}, + Rope::RefIndexMode::FIXED, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({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 f44580cf..d7a0f005 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 c549280a..7daec861 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 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 padded_shape = ctrl_latent.shape(); - padded_shape[2] = 33; + padded_shape[2] = 33; sd::Tensor 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 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;