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

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

@ -5,6 +5,7 @@
#include <cassert> #include <cassert>
#include <cmath> #include <cmath>
#include <set> #include <set>
#include <utility>
#include <vector> #include <vector>
#include "core/ggml_extend.h" #include "core/ggml_extend.h"
#include "core/ggml_runner.h" #include "core/ggml_runner.h"
@ -16,6 +17,45 @@ namespace Rope {
ErnieImage, 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 { enum class RefIndexMode {
FIXED, FIXED,
INCREASE, INCREASE,
@ -56,40 +96,25 @@ namespace Rope {
return flat_vec; return flat_vec;
} }
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos, __STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
int dim,
float theta,
const std::vector<int>& axis_wrap_dims = {}) {
assert(dim % 2 == 0); assert(dim % 2 == 0);
int half_dim = dim / 2; int half_dim = dim / 2;
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim); std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
std::vector<float> omega(half_dim); std::vector<float> omega(half_dim);
for (int i = 0; i < half_dim; ++i) { for (int i = 0; i < half_dim; ++i) {
omega[i] = 1.0f / ::powf(1.f * theta, scale[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(); size_t pos_size = pos.size();
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim)); 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 i = 0; i < pos_size; ++i) {
for (size_t j = 0; j < half_dim; ++j) { for (size_t j = 0; j < half_dim; ++j) {
float angle = pos[i] * omega[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; out[i][j] = angle;
} }
@ -108,6 +133,12 @@ namespace Rope {
return result; 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 // 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) { __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)); auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
@ -136,12 +167,16 @@ namespace Rope {
int patch_size, int patch_size,
int bs, int bs,
int axes_dim_num, int axes_dim_num,
int index = 0, int index = 0,
int h_offset = 0, int h_offset = 0,
int w_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 h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (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)); std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
int h_start = h_offset; int h_start = h_offset;
@ -192,8 +227,8 @@ namespace Rope {
int bs, int bs,
const std::vector<float>& axis_thetas, const std::vector<float>& axis_thetas,
const std::vector<int>& axes_dim, 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); std::vector<std::vector<float>> trans_ids = transpose(ids);
size_t pos_len = ids.size() / bs; size_t pos_len = ids.size() / bs;
size_t num_axes = axes_dim.size(); size_t num_axes = axes_dim.size();
@ -205,19 +240,25 @@ namespace Rope {
for (int d : axes_dim) for (int d : axes_dim)
emb_dim += d / 2; 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)); std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
size_t offset = 0; size_t offset = 0;
for (size_t i = 0; i < num_axes; ++i) { 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; float axis_theta = 10000.0f;
if (!axis_thetas.empty()) { if (!axis_thetas.empty()) {
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)]; 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 = 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 b = 0; b < bs; ++b) {
for (int j = 0; j < pos_len; ++j) { for (int j = 0; j < pos_len; ++j) {
for (int k = 0; k < rope_emb[0].size(); ++k) { for (int k = 0; k < rope_emb[0].size(); ++k) {
@ -253,10 +294,10 @@ namespace Rope {
int bs, int bs,
float theta, float theta,
const std::vector<int>& axes_dim, 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); 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, __STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
@ -264,7 +305,7 @@ namespace Rope {
float theta, float theta,
int head_dim, int head_dim,
const std::vector<int>& mrope_section, 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(bs > 0);
GGML_ASSERT(head_dim % 2 == 0); GGML_ASSERT(head_dim % 2 == 0);
GGML_ASSERT(mrope_section.size() >= 3); GGML_ASSERT(mrope_section.size() >= 3);
@ -273,20 +314,26 @@ namespace Rope {
size_t pos_len = ids.size() / bs; size_t pos_len = ids.size() / bs;
int half_dim = head_dim / 2; 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; std::vector<std::vector<std::vector<float>>> axis_embs;
axis_embs.reserve(3); axis_embs.reserve(3);
for (int axis = 0; axis < 3; ++axis) { for (int axis = 0; axis < 3; ++axis) {
std::vector<int> axis_wrap; axis_embs.push_back(rope(trans_ids[axis], omega));
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));
} }
std::vector<std::vector<float>> emb = axis_embs[0]; std::vector<std::vector<float>> emb = axis_embs[0];
for (int axis = 1; axis < 3; ++axis) { for (int axis = 1; axis < 3; ++axis) {
int length = std::min<int>(mrope_section[axis] * 3, half_dim); int length = std::min<int>(mrope_section[axis] * 3, half_dim);
for (int freq_idx = axis; freq_idx < length; freq_idx += 3) { 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 (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
for (int k = 0; k < 4; ++k) { for (int k = 0; k < 4; ++k) {
emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k]; emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
@ -298,13 +345,13 @@ namespace Rope {
return flatten(emb); return flatten(emb);
} }
__STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height, __STATIC_INLINE__ Embedding embed_2d_interleaved(int height,
int width, int width,
int dim, int dim,
float theta = 10000.f, float theta = 10000.f,
float scale = 16.f, float scale = 16.f,
int ref_grid_h = 0, int ref_grid_h = 0,
int ref_grid_w = 0) { int ref_grid_w = 0) {
assert(dim % 4 == 0); assert(dim % 4 == 0);
int half_dim = dim / 2; int half_dim = dim / 2;
int dim_axis = dim / 2; int dim_axis = dim / 2;
@ -318,6 +365,10 @@ namespace Rope {
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power); 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> x_pos;
std::vector<float> y_pos; std::vector<float> y_pos;
x_pos.reserve(static_cast<size_t>(height) * width); 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); float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
for (int ix = 0; ix < width; ++ix) { for (int ix = 0; ix < width; ++ix) {
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1); 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); x_pos.push_back(x);
y_pos.push_back(y); y_pos.push_back(y);
} }
} }
auto x_emb = rope(x_pos, dim_axis, theta * w_ntk); auto x_freq = rope_frequencies(dim_axis, theta * w_ntk);
auto y_emb = rope(y_pos, dim_axis, theta * h_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); std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
for (int pos = 0; pos < height * width; ++pos) { 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, __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@ -359,7 +418,8 @@ namespace Rope {
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
bool scale_rope, bool scale_rope,
int base_offset = 0) { int base_offset = 0,
PositionLayout* layout = nullptr) {
std::vector<std::vector<float>> ids; std::vector<std::vector<float>> ids;
int curr_h_offset = 0; int curr_h_offset = 0;
int curr_w_offset = 0; int curr_w_offset = 0;
@ -386,7 +446,8 @@ namespace Rope {
static_cast<int>(index * ref_index_scale), static_cast<int>(index * ref_index_scale),
h_offset + base_offset, h_offset + base_offset,
w_offset + base_offset, w_offset + base_offset,
scale_rope); scale_rope,
layout);
ids = concat_ids(ids, ref_ids, bs); ids = concat_ids(ids, ref_ids, bs);
if (ref_index_mode == RefIndexMode::INCREASE) { if (ref_index_mode == RefIndexMode::INCREASE) {
@ -409,88 +470,53 @@ namespace Rope {
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
float ref_index_scale, 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; 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); 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; 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); auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) { 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); ids = concat_ids(ids, refs_ids, bs);
} }
return ids; return ids;
} }
// Generate flux positional embeddings // 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 w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
std::set<int> txt_arange_dims, std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim,
bool circular_w, bool is_longcat) {
const std::vector<int>& axes_dim, Embedding result;
bool is_longcat) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_flux_ids(h, result.ids = gen_flux_ids(h,
w, w,
patch_size, patch_size,
bs, bs,
static_cast<int>(axes_dim.size()), static_cast<int>(axes_dim.size()),
context_len, context_len,
txt_arange_dims, txt_arange_dims,
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
ref_index_scale, ref_index_scale,
is_longcat); is_longcat, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
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);
} }
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t, __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@ -500,14 +526,18 @@ namespace Rope {
int ph, int ph,
int pw, int pw,
int bs, int bs,
int t_offset = 0, int t_offset = 0,
int h_offset = 0, int h_offset = 0,
int w_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 t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph; int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw; 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)); std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
if (scale_rope) { if (scale_rope) {
@ -573,7 +603,11 @@ namespace Rope {
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, 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 h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (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; int txt_id_start = std::max(h_len, w_len) / 2;
@ -585,90 +619,49 @@ namespace Rope {
} }
} }
int axes_dim_num = 3; 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); auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) { if (ref_latents.size() > 0) {
int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1; 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); ids = concat_ids(ids, refs_ids, bs);
} }
return ids; return ids;
} }
// Generate qwen_image positional embeddings // 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 h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode); result.ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles return result;
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);
} }
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h, __STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
int w, int w,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
int theta, int theta,
const std::vector<int>& axes_dim) { 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()); 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<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
std::vector<float>(axes_dim_num, 0.f)); 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 h_start = -(image_h - image_h / 2);
int w_start = -(image_w - image_w / 2); int w_start = -(image_w - image_w / 2);
for (int b = 0; b < bs; ++b) { for (int b = 0; b < bs; ++b) {
@ -692,15 +685,18 @@ namespace Rope {
static_cast<int>(i + 1)); static_cast<int>(i + 1));
ids = concat_ids(ids, ref_ids, bs); 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, __STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
int w, int w,
int bs, int bs,
int context_len, int context_len,
bool scale_rope = true) { bool scale_rope = true,
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope); 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; 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); 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); 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 w,
int bs, int bs,
int context_len, int context_len,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true); result.ids = gen_lens_ids(h, w, bs, context_len, true, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
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);
} }
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h, __STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len) { int context_len,
PositionLayout* layout = nullptr) {
int h_len = h / patch_size; int h_len = h / patch_size;
int w_len = w / 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<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> 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); 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); 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 w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len); result.layout = EmbedNDLayout::ErnieImage;
std::vector<std::vector<int>> wrap_dims; result.ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
int h_len = h / patch_size; return result;
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);
} }
// Generate wan positional embeddings // Generate wan positional embeddings
@ -905,7 +880,8 @@ namespace Rope {
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) { RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
SD_UNUSED(ref_index_mode); SD_UNUSED(ref_index_mode);
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of); 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)); 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; txt_ids[i][0] = (i % padded_context_len) + 1.f;
} }
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3; int axes_dim_num = 3;
int index = padded_context_len + 1; 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); 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) { if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f)); 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); img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -936,7 +918,8 @@ namespace Rope {
int patch_size, int patch_size,
int bs, int bs,
int context_len, 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 context_pad_len = bound_mod(context_len, seq_multi_of);
int padded_context_len = context_len + context_pad_len; 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)); 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 axes_dim_num = 3;
int index = padded_context_len + 1; 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); 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) { if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f)); 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); img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -968,7 +957,8 @@ namespace Rope {
int patch_size, int patch_size,
int context_len, int context_len,
int sigvq_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 context_pad = bound_mod(context_len, seq_multi_of);
const int padded_context = context_len + context_pad; const int padded_context = context_len + context_pad;
const int h_len = (h + (patch_size / 2)) / patch_size; const int h_len = (h + (patch_size / 2)) / patch_size;
@ -994,11 +984,17 @@ namespace Rope {
cursor += 2; cursor += 2;
} }
if (layout) {
layout->append_tokens(cap_ids.size());
}
std::vector<std::vector<float>> img_ids; std::vector<std::vector<float>> img_ids;
for (int copy = 0; copy < 2; ++copy) { 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(), ids.begin(), ids.end());
img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f)); 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; const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
@ -1016,95 +1012,59 @@ namespace Rope {
ids.insert(ids.end(), cap_ids.begin(), cap_ids.end()); ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
ids.insert(ids.end(), img_ids.begin(), img_ids.end()); ids.insert(ids.end(), img_ids.begin(), img_ids.end());
ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end()); ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
if (layout) {
layout->append_tokens(sigvq_ids.size());
}
SD_UNUSED(padded_image); SD_UNUSED(padded_image);
return ids; 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 w,
int patch_size, int patch_size,
int context_len, int context_len,
int sigvq_len, int sigvq_len,
int seq_multi_of, int seq_multi_of,
int theta, int theta,
const std::vector<int>& axes_dim) { const std::vector<int>& axes_dim) {
auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of); Embedding result;
return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {}); 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 w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of); result.ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
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);
} }
// Generate z_image positional embeddings // 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 w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
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); result.ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
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);
} }
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx, __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__

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

@ -603,34 +603,37 @@ namespace Anima {
return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2)); 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 h,
int w, int w,
int patch_size, int patch_size,
int theta, int theta,
const std::vector<int>& axes_dim, const std::vector<int>& axes_dim,
float h_extrapolation_ratio, float h_extrapolation_ratio,
float w_extrapolation_ratio, float w_extrapolation_ratio,
float t_extrapolation_ratio, float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) { const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h, Rope::Embedding result;
w, result.batch_size = bs;
patch_size, result.ids = Rope::gen_flux_ids(h,
bs, w,
static_cast<int>(axes_dim.size()), patch_size,
0, bs,
{}, static_cast<int>(axes_dim.size()),
ref_latents, 0,
Rope::RefIndexMode::FIXED, {},
1.0f, ref_latents,
false); Rope::RefIndexMode::FIXED,
1.0f,
false, &result.positions);
std::vector<float> axis_thetas = { 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(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(h_extrapolation_ratio, axes_dim[1]),
static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]), 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, ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@ -657,16 +660,16 @@ namespace Anima {
int64_t h_pad = x->ne[1] + pad_h; int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w; 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>(h_pad),
static_cast<int>(w_pad), static_cast<int>(w_pad),
static_cast<int>(config.patch_size), static_cast<int>(config.patch_size),
config.theta, config.theta,
config.axes_dim, config.axes_dim,
4.0f, 4.0f,
4.0f, 4.0f,
1.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)); 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); 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()); set_backend_tensor_data(image_pe, image_pe_vec.data());

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

@ -720,15 +720,18 @@ namespace Boogu {
} }
} }
__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h, __STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
int theta, int theta,
const std::vector<int>& axes_dim) { 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); ids.reserve(static_cast<size_t>(bs) * context_len);
for (int b = 0; b < bs; b++) { for (int b = 0; b < bs; b++) {
for (int i = 0; i < context_len; i++) { for (int i = 0; i < context_len; i++) {
@ -741,15 +744,18 @@ namespace Boogu {
for (ggml_tensor* ref : ref_latents) { for (ggml_tensor* ref : ref_latents) {
int ref_h_tokens = patched_token_count(ref->ne[1], patch_size); int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
int ref_w_tokens = patched_token_count(ref->ne[0], 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); append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
pe_shift += std::max(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 h_tokens = patched_token_count(h, patch_size);
int w_tokens = patched_token_count(w, 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); 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 { struct BooguImageRunner : public DiffusionModelRunner {
@ -793,14 +799,14 @@ namespace Boogu {
ref_latents.push_back(make_input(ref_latent_tensor)); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
config.theta, config.theta,
config.axes_dim); config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());

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

@ -415,15 +415,13 @@ namespace ErnieImage {
GGML_ASSERT(!context_tensor.empty()); GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());

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

@ -1548,20 +1548,18 @@ namespace Flux {
} else if (version == VERSION_OVIS_IMAGE) { } else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2}; 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
txt_arange_dims, txt_arange_dims,
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
config.ref_index_scale, config.ref_index_scale,
config.theta, config.theta,
circular_y_enabled, config.axes_dim,
circular_x_enabled, sd_version_is_longcat(version)));
config.axes_dim,
sd_version_is_longcat(version));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len); // 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); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

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

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

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

@ -689,23 +689,28 @@ namespace Krea2 {
} }
}; };
__STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h, __STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
float theta, float theta,
const std::vector<int>& axes_dim, const std::vector<int>& axes_dim,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
Rope::RefIndexMode ref_index_mode) { 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 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); auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) { 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); 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 { struct Krea2Runner : public DiffusionModelRunner {
@ -749,15 +754,15 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor)); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
config.theta, config.theta,
config.axes_dim, config.axes_dim,
ref_latents, 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); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());

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

@ -384,14 +384,12 @@ namespace Lens {
GGML_ASSERT(!context_tensor.empty()); GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); 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[0]),
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());

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

@ -412,16 +412,14 @@ namespace LLaDAImage {
GGML_ASSERT(!context_tensor.empty()); GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ZImage::SEQ_MULTI_OF, ZImage::SEQ_MULTI_OF,
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());
@ -461,14 +459,14 @@ namespace LLaDAImage {
ggml_tensor* source = make_input(source_tensor); ggml_tensor* source = make_input(source_tensor);
GGML_ASSERT(x->ne[3] == 1); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0, semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
ZImage::SEQ_MULTI_OF, ZImage::SEQ_MULTI_OF,
config.theta, config.theta,
config.axes_dim); config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); set_backend_tensor_data(pe, pe_vec.data());

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

@ -110,13 +110,13 @@ namespace MageFlow {
} }
int batch_size = static_cast<int>(x->ne[3]); 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]), static_cast<int>(x->ne[0]),
batch_size, batch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
config.theta, config.theta,
config.axes_dim); config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); 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); 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()); 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)); 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); GGML_ASSERT(head_dim % 4 == 0);
int dim = head_dim / 2; int dim = head_dim / 2;
int quarter = dim / 2; int quarter = dim / 2;
int length = side * side; 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> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
std::vector<float> freqs(quarter); std::vector<float> freqs(quarter);
for (int i = 0; i < quarter; ++i) { for (int i = 0; i < quarter; ++i) {
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim)); 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 y = 0; y < side; ++y) {
for (int x = 0; x < side; ++x) { 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; int pos = y * side + x;
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4; size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
for (int i = 0; i < quarter; ++i) { 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 { struct SwiGLUMlp : public GGMLBlock {
@ -475,6 +484,8 @@ namespace MiniT2I {
int64_t cached_txt_len = -1; int64_t cached_txt_len = -1;
int64_t cached_hidden_size = -1; int64_t cached_hidden_size = -1;
int64_t cached_head_dim = -1; int64_t cached_head_dim = -1;
bool cached_circular_x = false;
bool cached_circular_y = false;
MiniT2IRunner(ggml_backend_t backend, MiniT2IRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {}, const String2TensorStorage& tensor_storage_map = {},
@ -521,6 +532,8 @@ namespace MiniT2I {
cached_txt_len == txt_len && cached_txt_len == txt_len &&
cached_hidden_size == config.hidden_size && cached_hidden_size == config.hidden_size &&
cached_head_dim == config.head_dim && cached_head_dim == config.head_dim &&
cached_circular_x == circular_x_enabled &&
cached_circular_y == circular_y_enabled &&
cached_pos_embed != nullptr && cached_pos_embed != nullptr &&
cached_txt_pe != nullptr && cached_txt_pe != nullptr &&
cached_joint_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 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 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; auto joint_pe_vec = txt_pe_vec;
joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end()); 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_txt_len = txt_len;
cached_hidden_size = config.hidden_size; cached_hidden_size = config.hidden_size;
cached_head_dim = config.head_dim; 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, 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/ggml_runner.h"
#include "core/tensor_ggml.hpp" #include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp" #include "model/common/rope_circular.hpp"
#include "model_manager.h" #include "model_manager.h"
enum class RefImageResizeMode { enum class RefImageResizeMode {
@ -184,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
protected: protected:
std::string prefix; 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: public:
DiffusionModelRunner(ggml_backend_t backend, DiffusionModelRunner(ggml_backend_t backend,
const std::string& prefix, const std::string& prefix,

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

@ -135,13 +135,13 @@ namespace Pid {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta)); 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 width,
int dim, int dim,
float theta = 10000.f, float theta = 10000.f,
float scale = 16.f, float scale = 16.f,
int ref_grid_h = 0, int ref_grid_h = 0,
int ref_grid_w = 0) { int ref_grid_w = 0) {
GGML_ASSERT(dim % 4 == 0); GGML_ASSERT(dim % 4 == 0);
return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w); return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w);
} }
@ -867,13 +867,13 @@ namespace Pid {
int64_t Hs = Hp / config.patch_size; int64_t Hs = Hp / config.patch_size;
int64_t Ws = Wp / 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>(Ws),
static_cast<int>(config.hidden_size / config.num_groups), static_cast<int>(config.hidden_size / config.num_groups),
10000.f, 10000.f,
16.f, 16.f,
static_cast<int>(config.rope_ref_grid_h), 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, auto pos_img = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32, GGML_TYPE_F32,
2, 2,
@ -904,13 +904,13 @@ namespace Pid {
1); 1);
set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); 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>(Ws),
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups), static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
10000.f, 10000.f,
16.f, 16.f,
static_cast<int>(config.rope_ref_grid_h), 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, auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32, GGML_TYPE_F32,
2, 2,

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

@ -635,18 +635,16 @@ namespace Qwen {
ref_index_mode = Rope::RefIndexMode::DECREASE; 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[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
batch_size, batch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len); // 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); 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<QwenImage21Segment> segments;
std::vector<std::vector<float>> positions; std::vector<std::vector<float>> positions;
int64_t prefix_length = 0; int64_t prefix_length = 0;
Rope::PositionLayout rope_layout;
static QwenImage21Layout build(int64_t text_length, static QwenImage21Layout build(int64_t text_length,
const sd::Tensor<int32_t>& image_slots, const sd::Tensor<int32_t>& image_slots,
@ -82,6 +83,7 @@ namespace Qwen {
auto [height, width] = image_shapes[index]; auto [height, width] = image_shapes[index];
int64_t start = static_cast<int64_t>(layout.positions.size()); int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + height * width, context_start, index}); 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 h = 0; h < height; ++h) {
for (int64_t w = 0; w < width; ++w) { for (int64_t w = 0; w < width; ++w) {
layout.positions.push_back({static_cast<float>(position), layout.positions.push_back({static_cast<float>(position),
@ -106,6 +108,7 @@ namespace Qwen {
} else { } else {
int64_t start = static_cast<int64_t>(layout.positions.size()); int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + i - begin, begin, -1}); 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) { for (int64_t j = begin; j < i; ++j, ++position) {
float p = static_cast<float>(position); float p = static_cast<float>(position);
layout.positions.push_back({p, p, p}); layout.positions.push_back({p, p, p});
@ -418,14 +421,26 @@ namespace Qwen {
} }
QwenImage21PrefixCache cache; QwenImage21PrefixCache cache;
if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) { 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.prefix_length = layout.prefix_length;
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE; cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
} }
auto run = [&](const QwenImage21PrefixCache& active_cache) { auto run = [&](const QwenImage21PrefixCache& active_cache) {
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); 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(); mask_data.clear();
if (!cached) { if (!cached) {
for (const auto& segment : layout.segments) { for (const auto& segment : layout.segments) {

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

@ -642,18 +642,16 @@ namespace ZImage {
ref_latents.push_back(make_input(ref_latent_tensor)); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
SEQ_MULTI_OF, SEQ_MULTI_OF,
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len); // 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); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);