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

refactor: centralize circular RoPE and extend image model support (#2039)

Browse Source
This commit is contained in:
leejet 2026-09-24 01:45:14 +08:00 committed by GitHub
parent 500ef5fa7c
commit 88411ef1e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 611 additions and 564 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

482
src/model/common/rope.hpp
View File

@ -5,6 +5,7 @@
#include <cassert>
#include <cmath>
#include <set>
#include <utility>
#include <vector>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@ -16,6 +17,45 @@ namespace Rope {
ErnieImage,
};
struct SpatialRegion {
size_t begin;
size_t count;
float height_period;
float width_period;
int height_axis = 1;
int width_axis = 2;
};
struct PositionLayout {
// Token ranges are relative to one batch item.
std::vector<SpatialRegion> images;
size_t token_count = 0;
void append_tokens(size_t count) {
token_count += count;
}
void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) {
size_t count = static_cast<size_t>(height) * width * frames;
images.push_back({token_count, count, height * height_step, width * width_step});
append_tokens(count);
}
};
struct Frequency {
size_t axis;
float omega;
};
struct Embedding {
std::vector<float> values;
std::vector<std::vector<float>> ids;
PositionLayout positions;
std::vector<Frequency> frequencies;
EmbedNDLayout layout = EmbedNDLayout::Matrix;
int batch_size = 1;
};
enum class RefIndexMode {
FIXED,
INCREASE,
@ -56,40 +96,25 @@ namespace Rope {
return flat_vec;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
int dim,
float theta,
const std::vector<int>& axis_wrap_dims = {}) {
__STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
assert(dim % 2 == 0);
int half_dim = dim / 2;
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
std::vector<float> omega(half_dim);
for (int i = 0; i < half_dim; ++i) {
omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
}
return omega;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
const std::vector<float>& omega) {
int half_dim = static_cast<int>(omega.size());
size_t pos_size = pos.size();
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
for (size_t i = 0; i < pos_size; ++i) {
for (size_t j = 0; j < half_dim; ++j) {
float angle = pos[i] * omega[j];
if (!axis_wrap_dims.empty()) {
size_t wrap_size = axis_wrap_dims.size();
// mod batch size since we only store this for one item in the batch
size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
int wrap_dim = axis_wrap_dims[wrap_idx];
if (wrap_dim > 0) {
constexpr float TWO_PI = 6.28318530717958647692f;
float cycles = omega[j] * wrap_dim / TWO_PI;
// closest periodic harmonic, necessary to ensure things neatly tile
// without this round, things don't tile at the boundaries and you end up
// with the model knowing what is "center"
float rounded = std::round(cycles);
angle = pos[i] * TWO_PI * rounded / wrap_dim;
}
}
out[i][j] = angle;
}
@ -108,6 +133,12 @@ namespace Rope {
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
int dim,
float theta) {
return rope(pos, rope_frequencies(dim, theta));
}
// Generate IDs for image patches and text
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
@ -139,9 +170,13 @@ namespace Rope {
int index = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false) {
bool scale_rope = false,
PositionLayout* layout = nullptr) {
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
if (layout) {
layout->append_image(h_len, w_len);
}
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
int h_start = h_offset;
@ -192,8 +227,8 @@ namespace Rope {
int bs,
const std::vector<float>& axis_thetas,
const std::vector<int>& axes_dim,
const std::vector<std::vector<int>>& wrap_dims = {},
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
EmbedNDLayout layout = EmbedNDLayout::Matrix,
std::vector<Frequency>* frequencies = nullptr) {
std::vector<std::vector<float>> trans_ids = transpose(ids);
size_t pos_len = ids.size() / bs;
size_t num_axes = axes_dim.size();
@ -205,19 +240,25 @@ namespace Rope {
for (int d : axes_dim)
emb_dim += d / 2;
if (frequencies) {
frequencies->clear();
frequencies->reserve(emb_dim);
}
std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
size_t offset = 0;
for (size_t i = 0; i < num_axes; ++i) {
std::vector<int> axis_wrap_dims;
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
axis_wrap_dims = wrap_dims[i];
}
float axis_theta = 10000.0f;
if (!axis_thetas.empty()) {
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
}
auto omega = rope_frequencies(axes_dim[i], axis_theta);
if (frequencies) {
for (float frequency : omega) {
frequencies->push_back({i, frequency});
}
}
std::vector<std::vector<float>> rope_emb =
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
rope(trans_ids[i], omega); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
for (int b = 0; b < bs; ++b) {
for (int j = 0; j < pos_len; ++j) {
for (int k = 0; k < rope_emb[0].size(); ++k) {
@ -253,10 +294,10 @@ namespace Rope {
int bs,
float theta,
const std::vector<int>& axes_dim,
const std::vector<std::vector<int>>& wrap_dims = {},
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
EmbedNDLayout layout = EmbedNDLayout::Matrix,
std::vector<Frequency>* frequencies = nullptr) {
std::vector<float> axis_thetas(axes_dim.size(), theta);
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies);
}
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
@ -264,7 +305,7 @@ namespace Rope {
float theta,
int head_dim,
const std::vector<int>& mrope_section,
const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
std::vector<Frequency>* frequencies = nullptr) {
GGML_ASSERT(bs > 0);
GGML_ASSERT(head_dim % 2 == 0);
GGML_ASSERT(mrope_section.size() >= 3);
@ -273,20 +314,26 @@ namespace Rope {
size_t pos_len = ids.size() / bs;
int half_dim = head_dim / 2;
auto omega = rope_frequencies(head_dim, theta);
if (frequencies) {
frequencies->clear();
for (float frequency : omega) {
frequencies->push_back({0, frequency});
}
}
std::vector<std::vector<std::vector<float>>> axis_embs;
axis_embs.reserve(3);
for (int axis = 0; axis < 3; ++axis) {
std::vector<int> axis_wrap;
if (axis < static_cast<int>(axis_wrap_dims.size())) {
axis_wrap = axis_wrap_dims[axis];
}
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
axis_embs.push_back(rope(trans_ids[axis], omega));
}
std::vector<std::vector<float>> emb = axis_embs[0];
for (int axis = 1; axis < 3; ++axis) {
int length = std::min<int>(mrope_section[axis] * 3, half_dim);
for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
if (frequencies) {
(*frequencies)[freq_idx].axis = axis;
}
for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
for (int k = 0; k < 4; ++k) {
emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
@ -298,7 +345,7 @@ namespace Rope {
return flatten(emb);
}
__STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height,
__STATIC_INLINE__ Embedding embed_2d_interleaved(int height,
int width,
int dim,
float theta = 10000.f,
@ -318,6 +365,10 @@ namespace Rope {
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
}
Embedding result;
result.positions.append_image(height, width, 1,
height > 1 ? scale / (height - 1) : 1.f,
width > 1 ? scale / (width - 1) : 1.f);
std::vector<float> x_pos;
std::vector<float> y_pos;
x_pos.reserve(static_cast<size_t>(height) * width);
@ -326,13 +377,20 @@ namespace Rope {
float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
for (int ix = 0; ix < width; ++ix) {
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
result.ids.push_back({0.f, y, x});
x_pos.push_back(x);
y_pos.push_back(y);
}
}
auto x_emb = rope(x_pos, dim_axis, theta * w_ntk);
auto y_emb = rope(y_pos, dim_axis, theta * h_ntk);
auto x_freq = rope_frequencies(dim_axis, theta * w_ntk);
auto y_freq = rope_frequencies(dim_axis, theta * h_ntk);
auto x_emb = rope(x_pos, x_freq);
auto y_emb = rope(y_pos, y_freq);
for (int i = 0; i < axis_half_dim; ++i) {
result.frequencies.push_back({2, x_freq[i]});
result.frequencies.push_back({1, y_freq[i]});
}
std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
for (int pos = 0; pos < height * width; ++pos) {
@ -348,7 +406,8 @@ namespace Rope {
}
}
}
return out;
result.values = std::move(out);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@ -359,7 +418,8 @@ namespace Rope {
RefIndexMode ref_index_mode,
float ref_index_scale,
bool scale_rope,
int base_offset = 0) {
int base_offset = 0,
PositionLayout* layout = nullptr) {
std::vector<std::vector<float>> ids;
int curr_h_offset = 0;
int curr_w_offset = 0;
@ -386,7 +446,8 @@ namespace Rope {
static_cast<int>(index * ref_index_scale),
h_offset + base_offset,
w_offset + base_offset,
scale_rope);
scale_rope,
layout);
ids = concat_ids(ids, ref_ids, bs);
if (ref_index_mode == RefIndexMode::INCREASE) {
@ -409,23 +470,27 @@ namespace Rope {
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
float ref_index_scale,
bool is_longcat) {
bool is_longcat,
PositionLayout* layout = nullptr) {
if (layout) {
layout->append_tokens(context_len);
}
int x_index = is_longcat ? 1 : 0;
auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
int offset = is_longcat ? context_len : 0;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset, false, layout);
auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset, layout);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
}
// Generate flux positional embeddings
__STATIC_INLINE__ std::vector<float> gen_flux_pe(int h,
__STATIC_INLINE__ Embedding gen_flux_pe(int h,
int w,
int patch_size,
int bs,
@ -435,11 +500,11 @@ namespace Rope {
RefIndexMode ref_index_mode,
float ref_index_scale,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim,
bool is_longcat) {
std::vector<std::vector<float>> ids = gen_flux_ids(h,
Embedding result;
result.batch_size = bs;
result.ids = gen_flux_ids(h,
w,
patch_size,
bs,
@ -449,48 +514,9 @@ namespace Rope {
ref_latents,
ref_index_mode,
ref_index_scale,
is_longcat);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len; // text first
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
cursor += img_tokens;
// reference latents
for (ggml_tensor* ref : ref_latents) {
if (ref == nullptr) {
continue;
}
int ref_h = static_cast<int>(ref->ne[1]);
int ref_w = static_cast<int>(ref->ne[0]);
int ref_h_l = (ref_h + (patch_size / 2)) / patch_size;
int ref_w_l = (ref_w + (patch_size / 2)) / patch_size;
size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = ref_h_l;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = ref_w_l;
}
}
cursor += ref_tokens;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
is_longcat, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@ -503,11 +529,15 @@ namespace Rope {
int t_offset = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false) {
bool scale_rope = false,
PositionLayout* layout = nullptr) {
int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw;
if (layout) {
layout->append_image(h_len, w_len, t_len);
}
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
if (scale_rope) {
@ -573,7 +603,11 @@ namespace Rope {
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) {
RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
if (layout) {
layout->append_tokens(context_len);
}
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
int txt_id_start = std::max(h_len, w_len) / 2;
@ -585,18 +619,18 @@ namespace Rope {
}
}
int axes_dim_num = 3;
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true, layout);
auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) {
int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true, 0, layout);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
}
// Generate qwen_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
__STATIC_INLINE__ Embedding gen_qwen_image_pe(int t,
int h,
int w,
int patch_size,
@ -605,70 +639,29 @@ namespace Rope {
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
const size_t total_tokens = ids.size();
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
size_t cursor = context_len; // ignore text tokens
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
cursor += img_tokens;
// For each reference image, store wrap sizes as well
for (ggml_tensor* ref : ref_latents) {
if (ref == nullptr) {
continue;
}
int ref_h = static_cast<int>(ref->ne[1]);
int ref_w = static_cast<int>(ref->ne[0]);
int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size;
int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size;
int ref_h_len = (ref_h + ref_pad_h) / patch_size;
int ref_w_len = (ref_w + ref_pad_w) / patch_size;
size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = ref_h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = ref_w_len;
}
}
cursor += ref_n_tokens;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
Embedding result;
result.batch_size = bs;
result.ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
__STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
int w,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
const int axes_dim_num = static_cast<int>(axes_dim.size());
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) {
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
std::vector<float>(axes_dim_num, 0.f));
result.positions.append_image(image_h, image_w);
int h_start = -(image_h - image_h / 2);
int w_start = -(image_w - image_w / 2);
for (int b = 0; b < bs; ++b) {
@ -692,15 +685,18 @@ namespace Rope {
static_cast<int>(i + 1));
ids = concat_ids(ids, ref_ids, bs);
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
result.ids = std::move(ids);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
int w,
int bs,
int context_len,
bool scale_rope = true) {
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
bool scale_rope = true,
PositionLayout* layout = nullptr) {
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope, layout);
int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
@ -711,44 +707,37 @@ namespace Rope {
}
}
if (layout) {
layout->append_tokens(context_len);
}
return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
}
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
__STATIC_INLINE__ Embedding gen_lens_pe(int h,
int w,
int bs,
int context_len,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][token_i] = h;
}
if (circular_w) {
wrap_dims[2][token_i] = w;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
Embedding result;
result.batch_size = bs;
result.ids = gen_lens_ids(h, w, bs, context_len, true, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
int w,
int patch_size,
int bs,
int context_len) {
int context_len,
PositionLayout* layout = nullptr) {
int h_len = h / patch_size;
int w_len = w / patch_size;
if (layout) {
layout->append_image(h_len, w_len);
}
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
@ -774,39 +763,25 @@ namespace Rope {
}
}
if (layout) {
layout->append_tokens(context_len);
}
return concat_ids(img_ids_repeated, txt_ids, bs);
}
__STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
__STATIC_INLINE__ Embedding gen_ernie_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int h_len = h / patch_size;
int w_len = w / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
Embedding result;
result.batch_size = bs;
result.layout = EmbedNDLayout::ErnieImage;
result.ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
// Generate wan positional embeddings
@ -905,7 +880,8 @@ namespace Rope {
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) {
RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
SD_UNUSED(ref_index_mode);
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@ -913,11 +889,17 @@ namespace Rope {
txt_ids[i][0] = (i % padded_context_len) + 1.f;
}
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3;
int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (layout) {
layout->append_tokens(img_pad_len);
}
if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -936,7 +918,8 @@ namespace Rope {
int patch_size,
int bs,
int context_len,
int seq_multi_of) {
int seq_multi_of,
PositionLayout* layout = nullptr) {
int context_pad_len = bound_mod(context_len, seq_multi_of);
int padded_context_len = context_len + context_pad_len;
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@ -947,11 +930,17 @@ namespace Rope {
}
}
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3;
int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (layout) {
layout->append_tokens(img_pad_len);
}
if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -968,7 +957,8 @@ namespace Rope {
int patch_size,
int context_len,
int sigvq_len,
int seq_multi_of) {
int seq_multi_of,
PositionLayout* layout = nullptr) {
const int context_pad = bound_mod(context_len, seq_multi_of);
const int padded_context = context_len + context_pad;
const int h_len = (h + (patch_size / 2)) / patch_size;
@ -994,11 +984,17 @@ namespace Rope {
cursor += 2;
}
if (layout) {
layout->append_tokens(cap_ids.size());
}
std::vector<std::vector<float>> img_ids;
for (int copy = 0; copy < 2; ++copy) {
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]);
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy], 0, 0, false, layout);
img_ids.insert(img_ids.end(), ids.begin(), ids.end());
img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
if (layout) {
layout->append_tokens(image_pad);
}
}
const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
@ -1016,11 +1012,14 @@ namespace Rope {
ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
ids.insert(ids.end(), img_ids.begin(), img_ids.end());
ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
if (layout) {
layout->append_tokens(sigvq_ids.size());
}
SD_UNUSED(padded_image);
return ids;
}
__STATIC_INLINE__ std::vector<float> gen_llada_image_edit_pe(int h,
__STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h,
int w,
int patch_size,
int context_len,
@ -1028,48 +1027,30 @@ namespace Rope {
int seq_multi_of,
int theta,
const std::vector<int>& axes_dim) {
auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of);
return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {});
Embedding result;
result.batch_size = 1;
result.ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of, &result.positions);
result.values = embed_nd(result.ids, 1, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h,
__STATIC_INLINE__ Embedding gen_llada_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len + bound_mod(context_len, seq_multi_of);
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
Embedding result;
result.batch_size = bs;
result.ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
// Generate z_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
__STATIC_INLINE__ Embedding gen_z_image_pe(int h,
int w,
int patch_size,
int bs,
@ -1078,33 +1059,12 @@ namespace Rope {
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding)
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
Embedding result;
result.batch_size = bs;
result.ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,

65
src/model/common/rope_circular.hpp Normal file
View File

@ -0,0 +1,65 @@
#ifndef __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
#define __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
#include "model/common/rope.hpp"
namespace Rope {
__STATIC_INLINE__ void apply_circular(Embedding& embedding, bool circular_x, bool circular_y) {
if (!circular_x && !circular_y) {
return;
}
GGML_ASSERT(embedding.batch_size > 0);
GGML_ASSERT(embedding.ids.size() % embedding.batch_size == 0);
size_t pos_len = embedding.ids.size() / embedding.batch_size;
size_t half_dim = embedding.frequencies.size();
GGML_ASSERT(embedding.positions.token_count == pos_len);
GGML_ASSERT(embedding.values.size() == embedding.ids.size() * half_dim * 4);
constexpr float TWO_PI = 6.28318530717958647692f;
for (const auto& region : embedding.positions.images) {
GGML_ASSERT(region.begin <= pos_len && region.count <= pos_len - region.begin);
for (size_t j = 0; j < half_dim; ++j) {
const auto& frequency = embedding.frequencies[j];
float period = 0.f;
if (circular_y && frequency.axis == static_cast<size_t>(region.height_axis)) {
period = region.height_period;
} else if (circular_x && frequency.axis == static_cast<size_t>(region.width_axis)) {
period = region.width_period;
}
if (period <= 0) {
continue;
}
// Quantize to periodic harmonics while preserving the original coordinate offsets.
float rounded = std::round(frequency.omega * period / TWO_PI);
for (int b = 0; b < embedding.batch_size; ++b) {
size_t begin = b * pos_len + region.begin;
for (size_t i = begin; i < begin + region.count; ++i) {
GGML_ASSERT(frequency.axis < embedding.ids[i].size());
float angle = embedding.ids[i][frequency.axis] * TWO_PI * rounded / period;
float cos_val = std::cos(angle);
float sin_val = std::sin(angle);
if (embedding.layout == EmbedNDLayout::ErnieImage) {
size_t cos_offset = (i * half_dim + j) * 2;
size_t sin_offset = embedding.ids.size() * half_dim * 2 + cos_offset;
embedding.values[cos_offset] = cos_val;
embedding.values[cos_offset + 1] = cos_val;
embedding.values[sin_offset] = sin_val;
embedding.values[sin_offset + 1] = sin_val;
} else {
size_t offset = (i * half_dim + j) * 4;
embedding.values[offset] = cos_val;
embedding.values[offset + 1] = -sin_val;
embedding.values[offset + 2] = sin_val;
embedding.values[offset + 3] = cos_val;
}
}
}
}
}
}
} // namespace Rope
#endif // __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__

15
src/model/diffusion/anima.hpp
View File

@ -603,7 +603,7 @@ namespace Anima {
return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
}
static std::vector<float> gen_anima_image_pe_vec(int bs,
static Rope::Embedding gen_anima_image_pe_vec(int bs,
int h,
int w,
int patch_size,
@ -613,7 +613,9 @@ namespace Anima {
float w_extrapolation_ratio,
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h,
Rope::Embedding result;
result.batch_size = bs;
result.ids = Rope::gen_flux_ids(h,
w,
patch_size,
bs,
@ -623,14 +625,15 @@ namespace Anima {
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
false);
false, &result.positions);
std::vector<float> axis_thetas = {
static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
};
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
result.values = Rope::embed_nd(result.ids, bs, axis_thetas, axes_dim, result.layout, &result.frequencies);
return result;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@ -657,7 +660,7 @@ namespace Anima {
int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w;
image_pe_vec = gen_anima_image_pe_vec(1,
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),
@ -666,7 +669,7 @@ namespace Anima {
4.0f,
4.0f,
1.0f,
ref_latents);
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());

16
src/model/diffusion/boogu.hpp
View File

@ -720,7 +720,7 @@ namespace Boogu {
}
}
__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
__STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
int w,
int patch_size,
int bs,
@ -728,7 +728,10 @@ namespace Boogu {
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids;
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
auto& ids = result.ids;
ids.reserve(static_cast<size_t>(bs) * context_len);
for (int b = 0; b < bs; b++) {
for (int i = 0; i < context_len; i++) {
@ -741,15 +744,18 @@ namespace Boogu {
for (ggml_tensor* ref : ref_latents) {
int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
result.positions.append_image(ref_h_tokens, ref_w_tokens);
append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
pe_shift += std::max(ref_h_tokens, ref_w_tokens);
}
int h_tokens = patched_token_count(h, patch_size);
int w_tokens = patched_token_count(w, patch_size);
result.positions.append_image(h_tokens, w_tokens);
append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
result.values = Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
struct BooguImageRunner : public DiffusionModelRunner {
@ -793,14 +799,14 @@ namespace Boogu {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = gen_boogu_pe(static_cast<int>(x->ne[1]),
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);
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());

6
src/model/diffusion/ernie_image.hpp
View File

@ -415,15 +415,13 @@ namespace ErnieImage {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
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,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
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());

6
src/model/diffusion/flux.hpp
View File

@ -1548,7 +1548,7 @@ namespace Flux {
} else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2};
}
pe_vec = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
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]),
@ -1558,10 +1558,8 @@ namespace Flux {
ref_index_mode,
config.ref_index_scale,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim,
sd_version_is_longcat(version));
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);

41
src/model/diffusion/ideogram4.hpp
View File

@ -149,18 +149,21 @@ namespace Ideogram4 {
return std::make_shared<Linear>(in_features, out_features, bias);
}
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
__STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h,
int grid_w,
int bs,
int context_len,
int head_dim,
int rope_theta,
const std::vector<int>& mrope_section,
bool circular_x = false,
bool circular_y = false) {
const std::vector<int>& mrope_section) {
GGML_ASSERT(bs == 1);
std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
result.positions.append_image(grid_h, grid_w);
result.ids.assign(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
std::vector<float>(3, 0.f));
auto& ids = result.ids;
for (int i = 0; i < context_len; ++i) {
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
@ -175,29 +178,13 @@ namespace Ideogram4 {
}
}
std::vector<std::vector<int>> axis_wrap_dims(3);
if (circular_y || circular_x) {
size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
axis_wrap_dims[1].assign(total_len, 0);
axis_wrap_dims[2].assign(total_len, 0);
if (circular_y) {
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
axis_wrap_dims[1][idx] = grid_h;
}
}
if (circular_x) {
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
axis_wrap_dims[2][idx] = grid_w;
}
}
}
return Rope::embed_interleaved_mrope(ids,
result.values = Rope::embed_interleaved_mrope(ids,
bs,
static_cast<float>(rope_theta),
head_dim,
mrope_section,
axis_wrap_dims);
&result.frequencies);
return result;
}
class Ideogram4Attention : public GGMLBlock {
@ -509,15 +496,13 @@ namespace Ideogram4 {
int64_t head_dim = config.emb_dim / config.num_heads;
auto runner_ctx = get_context();
pe_vec = gen_ideogram4_pe(static_cast<int>(grid_h),
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,
runner_ctx.circular_x_enabled,
runner_ctx.circular_y_enabled);
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());

17
src/model/diffusion/krea2.hpp
View File

@ -689,7 +689,7 @@ namespace Krea2 {
}
};
__STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h,
__STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
int w,
int patch_size,
int bs,
@ -698,14 +698,19 @@ namespace Krea2 {
const std::vector<int>& axes_dim,
const std::vector<ggml_tensor*>& ref_latents,
Rope::RefIndexMode ref_index_mode) {
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false, &result.positions);
auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0, &result.positions);
ids = Rope::concat_ids(ids, refs_ids, bs);
}
return Rope::embed_nd(ids, bs, theta, axes_dim);
result.ids = std::move(ids);
result.values = Rope::embed_nd(result.ids, bs, theta, axes_dim, result.layout, &result.frequencies);
return result;
}
struct Krea2Runner : public DiffusionModelRunner {
@ -749,7 +754,7 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
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]),
@ -757,7 +762,7 @@ namespace Krea2 {
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode);
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());

6
src/model/diffusion/lens.hpp
View File

@ -384,14 +384,12 @@ namespace Lens {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
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,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
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());

10
src/model/diffusion/llada_image.hpp
View File

@ -412,16 +412,14 @@ namespace LLaDAImage {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
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,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
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());
@ -461,14 +459,14 @@ namespace LLaDAImage {
ggml_tensor* source = make_input(source_tensor);
GGML_ASSERT(x->ne[3] == 1);
pe_vec = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[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);
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());

4
src/model/diffusion/mage_flow.hpp
View File

@ -110,13 +110,13 @@ namespace MageFlow {
}
int batch_size = static_cast<int>(x->ne[3]);
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
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);
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());

21
src/model/diffusion/minit2i.hpp
View File

@ -154,18 +154,26 @@ namespace MiniT2I {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
}
inline std::vector<float> make_vision_rope(int side, int head_dim) {
inline Rope::Embedding make_vision_rope(int side, int head_dim) {
GGML_ASSERT(head_dim % 4 == 0);
int dim = head_dim / 2;
int quarter = dim / 2;
int length = side * side;
Rope::Embedding result;
result.positions.append_image(side, side);
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
std::vector<float> freqs(quarter);
for (int i = 0; i < quarter; ++i) {
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
}
for (int axis : {1, 2}) {
for (float frequency : freqs) {
result.frequencies.push_back({static_cast<size_t>(axis), frequency});
}
}
for (int y = 0; y < side; ++y) {
for (int x = 0; x < side; ++x) {
result.ids.push_back({0.f, static_cast<float>(y), static_cast<float>(x)});
int pos = y * side + x;
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
for (int i = 0; i < quarter; ++i) {
@ -182,7 +190,8 @@ namespace MiniT2I {
}
}
}
return out;
result.values = std::move(out);
return result;
}
struct SwiGLUMlp : public GGMLBlock {
@ -475,6 +484,8 @@ namespace MiniT2I {
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 = {},
@ -521,6 +532,8 @@ namespace MiniT2I {
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) {
@ -531,7 +544,7 @@ namespace MiniT2I {
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 = make_vision_rope(static_cast<int>(img_side), 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());
@ -561,6 +574,8 @@ namespace MiniT2I {
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,

7
src/model/diffusion/model.hpp
View File

@ -7,7 +7,7 @@
#include "core/ggml_runner.h"
#include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp"
#include "model/common/rope_circular.hpp"
#include "model_manager.h"
enum class RefImageResizeMode {
@ -184,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
protected:
std::string prefix;
std::vector<float> finish_rope_pe(Rope::Embedding embedding) {
Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled);
return std::move(embedding.values);
}
public:
DiffusionModelRunner(ggml_backend_t backend,
const std::string& prefix,

10
src/model/diffusion/pid.hpp
View File

@ -135,7 +135,7 @@ namespace Pid {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
}
inline std::vector<float> make_rope_2d(int height,
inline Rope::Embedding make_rope_2d(int height,
int width,
int dim,
float theta = 10000.f,
@ -867,13 +867,13 @@ namespace Pid {
int64_t Hs = Hp / config.patch_size;
int64_t Ws = Wp / config.patch_size;
pos_img_vec = make_rope_2d(static_cast<int>(Hs),
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));
static_cast<int>(config.rope_ref_grid_w)));
auto pos_img = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
@ -904,13 +904,13 @@ namespace Pid {
1);
set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
pixel_pos_comp_vec = make_rope_2d(static_cast<int>(Hs),
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));
static_cast<int>(config.rope_ref_grid_w)));
auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,

6
src/model/diffusion/qwen_image.hpp
View File

@ -635,7 +635,7 @@ namespace Qwen {
ref_index_mode = Rope::RefIndexMode::DECREASE;
}
pe_vec = Rope::gen_qwen_image_pe(time_len,
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,
@ -644,9 +644,7 @@ namespace Qwen {
ref_latents,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
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);

19
src/model/diffusion/qwen_image_2_1.hpp
View File

@ -68,6 +68,7 @@ namespace Qwen {
std::vector<QwenImage21Segment> segments;
std::vector<std::vector<float>> positions;
int64_t prefix_length = 0;
Rope::PositionLayout rope_layout;
static QwenImage21Layout build(int64_t text_length,
const sd::Tensor<int32_t>& image_slots,
@ -82,6 +83,7 @@ namespace Qwen {
auto [height, width] = image_shapes[index];
int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + height * width, context_start, index});
layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
for (int64_t h = 0; h < height; ++h) {
for (int64_t w = 0; w < width; ++w) {
layout.positions.push_back({static_cast<float>(position),
@ -106,6 +108,7 @@ namespace Qwen {
} else {
int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + i - begin, begin, -1});
layout.rope_layout.append_tokens(i - begin);
for (int64_t j = begin; j < i; ++j, ++position) {
float p = static_cast<float>(position);
layout.positions.push_back({p, p, p});
@ -418,14 +421,26 @@ namespace Qwen {
}
QwenImage21PrefixCache cache;
if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id);
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
cache.prefix_length = layout.prefix_length;
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
}
auto run = [&](const QwenImage21PrefixCache& active_cache) {
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
pe_data = Rope::embed_nd(std::vector<std::vector<float>>(first_position, layout.positions.end()), 1, 10000.f, config.axes_dim);
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));
mask_data.clear();
if (!cached) {
for (const auto& segment : layout.segments) {

6
src/model/diffusion/z_image.hpp
View File

@ -642,7 +642,7 @@ namespace ZImage {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
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]),
@ -651,9 +651,7 @@ namespace ZImage {
ref_latents,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
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);