trending/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-10-29 07:33:29 +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
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;