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/moss_image.hpp
+23 -18
View File
@@ -3,6 +3,7 @@
#include <algorithm>
#include <cinttypes>
#include "core/pe_cache.h"
#include "model/common/ggml_block.hpp"
#include "model/common/rope.hpp"
@@ -193,9 +194,10 @@ namespace MOSSImage {
};
struct MOSSImageRunner : public DiffusionModelRunner {
sd::PECache pe_cache_;
MOSSImageConfig config;
Model model;
std::vector<float> pe_data, time_data;
std::vector<float> time_data;
MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) {
@@ -224,31 +226,34 @@ namespace MOSSImage {
auto text = make_input(*inputs.context);
GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim);
int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1];
Rope::Embedding embedding;
embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w));
embedding.positions.images.back().height_axis = 2;
embedding.positions.images.back().width_axis = 3;
embedding.positions.append_tokens(nt);
for (int64_t token = 0; token < h * w + nt; ++token) {
float pos[4] = {0.f, float(token - h * w), 0.f, 0.f};
if (token < h * w) {
pos[1] = float(nt);
pos[2] = float(token / w) - float(h - 1) / 2.f;
pos[3] = float(token % w) - float(w - 1) / 2.f;
sd::PECacheKey pe_key("moss_image.pe");
pe_key.add(h).add(w).add(nt).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled);
auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
Rope::Embedding embedding;
embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w));
embedding.positions.images.back().height_axis = 2;
embedding.positions.images.back().width_axis = 3;
embedding.positions.append_tokens(nt);
for (int64_t token = 0; token < h * w + nt; ++token) {
float pos[4] = {0.f, float(token - h * w), 0.f, 0.f};
if (token < h * w) {
pos[1] = float(nt);
pos[2] = float(token / w) - float(h - 1) / 2.f;
pos[3] = float(token % w) - float(w - 1) / 2.f;
}
embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]});
}
embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]});
}
embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector<int>(4, static_cast<int>(config.head_dim / 4)), embedding.layout, &embedding.frequencies);
pe_data = finish_rope_pe(std::move(embedding));
embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector<int>(4, static_cast<int>(config.head_dim / 4)), embedding.layout, &embedding.frequencies);
auto values = finish_rope_pe(std::move(embedding));
return sd::Tensor<float>({2, 2, config.head_dim / 2, h * w + nt}, std::move(values));
}));
time_data.resize(config.freq_dim);
for (int64_t i = 0; i < config.freq_dim / 2; ++i) {
double angle = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2));
time_data[i] = float(std::cos(angle));
time_data[i + config.freq_dim / 2] = float(std::sin(angle));
}
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, h * w + nt);
auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1);
set_backend_tensor_data(pe, pe_data.data());
set_backend_tensor_data(time, time_data.data());
auto ctx = get_context();
ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip));