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

This commit is contained in:
leejet authored and GitHub committed 2026-09-24 01:45:14 +08:00
1 parent 500ef5fa7c
commit 88411ef1e0
17 files changed
+611 -564

No files matched your search

+26 -21
View File
@@ -689,23 +689,28 @@ namespace Krea2 {
}
};
__STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
float theta,
const std::vector<int>& axes_dim,
const std::vector<ggml_tensor*>& ref_latents,
Rope::RefIndexMode ref_index_mode) {
__STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
float theta,
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,15 +754,15 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode);
pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());