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