Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leejet merged 1 commit into from
Sep 23, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
642 changes: 301 additions & 341 deletions src/model/common/rope.hpp
View file
Open in desktop

Large diffs are not rendered by default.

65 changes: 65 additions & 0 deletions src/model/common/rope_circular.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 changes: 35 additions & 32 deletions src/model/diffusion/anima.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -603,34 +603,37 @@ 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,
int h,
int w,
int patch_size,
int theta,
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
0,
{},
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
false);
static Rope::Embedding gen_anima_image_pe_vec(int bs,
int h,
int w,
int patch_size,
int theta,
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
Rope::Embedding result;
result.batch_size = bs;
result.ids = Rope::gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
0,
{},
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
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,
Expand All @@ -657,16 +660,16 @@ 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,
static_cast<int>(h_pad),
static_cast<int>(w_pad),
static_cast<int>(config.patch_size),
config.theta,
config.axes_dim,
4.0f,
4.0f,
1.0f,
ref_latents);
image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1,
static_cast<int>(h_pad),
static_cast<int>(w_pad),
static_cast<int>(config.patch_size),
config.theta,
config.axes_dim,
4.0f,
4.0f,
1.0f,
ref_latents));
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
set_backend_tensor_data(image_pe, image_pe_vec.data());
Expand Down
42 changes: 24 additions & 18 deletions src/model/diffusion/boogu.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -720,15 +720,18 @@ namespace Boogu {
}
}

__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids;
__STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
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++) {
Expand All @@ -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 {
Expand Down Expand Up @@ -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]),
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);
pe_vec = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ref_latents,
config.theta,
config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
Expand Down
16 changes: 7 additions & 9 deletions src/model/diffusion/ernie_image.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]),
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);
pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
set_backend_tensor_data(pe, pe_vec.data());
Expand Down
26 changes: 12 additions & 14 deletions src/model/diffusion/flux.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1548,20 +1548,18 @@ 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]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
ref_index_mode,
config.ref_index_scale,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim,
sd_version_is_longcat(version));
pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
ref_index_mode,
config.ref_index_scale,
config.theta,
config.axes_dim,
sd_version_is_longcat(version)));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
Expand Down
Loading
Loading