DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 11:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: cache CPU position embeddings across DiT runners (#2211)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-10-25 05:18:30 +08:00
1 parent a85db96aa5
commit ea209dbfcf
31 files changed
+855 -610
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/diffusion/krea2.hpp
+19 -13
View File
@@ -11,6 +11,7 @@
#include <string>
#include <utility>
#include <vector>
#include "core/pe_cache.h"
#include "core/ggml_extend.h"
#include "core/ggml_graph_cut.h"
@@ -714,9 +715,9 @@ namespace Krea2 {
}
struct Krea2Runner : public DiffusionModelRunner {
sd::PECache pe_cache_;
Krea2Config config;
Krea2Model model;
std::vector<float> pe_vec;
Krea2Runner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -754,18 +755,23 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
sd::PECacheKey pe_key("krea2.pe");
pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
pe_key.add(ref_image_params.ref_index_mode);
auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
auto pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
}));
auto runner_ctx = get_context();
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);