diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index e4478493..ec822455 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #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 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(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 values; + std::vector> ids; + PositionLayout positions; + std::vector 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> rope(const std::vector& pos, - int dim, - float theta, - const std::vector& axis_wrap_dims = {}) { + __STATIC_INLINE__ std::vector rope_frequencies(int dim, float theta) { assert(dim % 2 == 0); - int half_dim = dim / 2; - + int half_dim = dim / 2; std::vector scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim); - std::vector 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> rope(const std::vector& pos, + const std::vector& omega) { + int half_dim = static_cast(omega.size()); size_t pos_size = pos.size(); std::vector> out(pos_size, std::vector(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> rope(const std::vector& pos, + int dim, + float theta) { + return rope(pos, rope_frequencies(dim, theta)); + } + // Generate IDs for image patches and text __STATIC_INLINE__ std::vector> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set arange_dims) { auto txt_ids = std::vector>(bs * context_len, std::vector(axes_dim_num, 0.0f)); @@ -136,12 +167,16 @@ namespace Rope { int patch_size, int bs, int axes_dim_num, - int index = 0, - int h_offset = 0, - int w_offset = 0, - bool scale_rope = false) { + int index = 0, + int h_offset = 0, + int w_offset = 0, + 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> img_ids(h_len * w_len, std::vector(axes_dim_num, 0.0)); int h_start = h_offset; @@ -192,8 +227,8 @@ namespace Rope { int bs, const std::vector& axis_thetas, const std::vector& axes_dim, - const std::vector>& wrap_dims = {}, - EmbedNDLayout layout = EmbedNDLayout::Matrix) { + EmbedNDLayout layout = EmbedNDLayout::Matrix, + std::vector* frequencies = nullptr) { std::vector> 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> emb(bs * pos_len, std::vector(emb_dim * 2 * 2, 0.0)); size_t offset = 0; for (size_t i = 0; i < num_axes; ++i) { - std::vector 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> 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& axes_dim, - const std::vector>& wrap_dims = {}, - EmbedNDLayout layout = EmbedNDLayout::Matrix) { + EmbedNDLayout layout = EmbedNDLayout::Matrix, + std::vector* frequencies = nullptr) { std::vector 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 embed_interleaved_mrope(const std::vector>& ids, @@ -264,7 +305,7 @@ namespace Rope { float theta, int head_dim, const std::vector& mrope_section, - const std::vector>& axis_wrap_dims = {}) { + std::vector* 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>> axis_embs; axis_embs.reserve(3); for (int axis = 0; axis < 3; ++axis) { - std::vector axis_wrap; - if (axis < static_cast(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> emb = axis_embs[0]; for (int axis = 1; axis < 3; ++axis) { int length = std::min(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,13 +345,13 @@ namespace Rope { return flatten(emb); } - __STATIC_INLINE__ std::vector embed_2d_interleaved(int height, - int width, - int dim, - float theta = 10000.f, - float scale = 16.f, - int ref_grid_h = 0, - int ref_grid_w = 0) { + __STATIC_INLINE__ Embedding embed_2d_interleaved(int height, + int width, + int dim, + float theta = 10000.f, + float scale = 16.f, + int ref_grid_h = 0, + int ref_grid_w = 0) { assert(dim % 4 == 0); int half_dim = dim / 2; int dim_axis = dim / 2; @@ -318,6 +365,10 @@ namespace Rope { w_ntk = std::pow(static_cast(width) / static_cast(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 x_pos; std::vector y_pos; x_pos.reserve(static_cast(height) * width); @@ -326,13 +377,20 @@ namespace Rope { float y = height == 1 ? 0.f : scale * static_cast(iy) / static_cast(height - 1); for (int ix = 0; ix < width; ++ix) { float x = width == 1 ? 0.f : scale * static_cast(ix) / static_cast(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 out(static_cast(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> 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> ids; int curr_h_offset = 0; int curr_w_offset = 0; @@ -386,7 +446,8 @@ namespace Rope { static_cast(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,88 +470,53 @@ namespace Rope { const std::vector& 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 gen_flux_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - std::set txt_arange_dims, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - float ref_index_scale, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim, - bool is_longcat) { - std::vector> ids = gen_flux_ids(h, - w, - patch_size, - bs, - static_cast(axes_dim.size()), - context_len, - txt_arange_dims, - ref_latents, - ref_index_mode, - ref_index_scale, - is_longcat); - std::vector> 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(pos_len, 0)); - size_t cursor = context_len; // text first - const size_t img_tokens = static_cast(h_len) * static_cast(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(ref->ne[1]); - int ref_w = static_cast(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(ref_h_l) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_flux_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + std::set txt_arange_dims, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + float ref_index_scale, + int theta, + const std::vector& axes_dim, + bool is_longcat) { + Embedding result; + result.batch_size = bs; + result.ids = gen_flux_ids(h, + w, + patch_size, + bs, + static_cast(axes_dim.size()), + context_len, + txt_arange_dims, + ref_latents, + ref_index_mode, + ref_index_scale, + is_longcat, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_vid_ids(int t, @@ -500,14 +526,18 @@ namespace Rope { int ph, int pw, int bs, - int t_offset = 0, - int h_offset = 0, - int w_offset = 0, - bool scale_rope = false) { + int t_offset = 0, + int h_offset = 0, + int w_offset = 0, + 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> vid_ids(t_len * h_len * w_len, std::vector(3, 0.0)); if (scale_rope) { @@ -573,7 +603,11 @@ namespace Rope { int bs, int context_len, const std::vector& 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,90 +619,49 @@ 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 gen_qwen_image_pe(int t, - int h, - int w, - int patch_size, - int bs, - int context_len, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode); - std::vector> 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(total_tokens / bs, 0)); - size_t cursor = context_len; // ignore text tokens - const size_t img_tokens = static_cast(t) * static_cast(h_len) * static_cast(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(ref->ne[1]); - int ref_w = static_cast(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(ref_h_len) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_qwen_image_pe(int t, + int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + int theta, + const std::vector& axes_dim) { + 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(theta), axes_dim, result.layout, &result.frequencies); + return result; } - __STATIC_INLINE__ std::vector gen_mage_flow_pe(int h, - int w, - int bs, - int context_len, - const std::vector& ref_latents, - int theta, - const std::vector& axes_dim) { + __STATIC_INLINE__ Embedding gen_mage_flow_pe(int h, + int w, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); const int axes_dim_num = static_cast(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> image_ids(static_cast(bs) * image_h * image_w, std::vector(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(i + 1)); ids = concat_ids(ids, ref_ids, bs); } - return embed_nd(ids, bs, static_cast(theta), axes_dim); + result.ids = std::move(ids); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> 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(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 gen_lens_pe(int h, - int w, - int bs, - int context_len, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_lens_ids(h, w, bs, context_len, true); - std::vector> 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(pos_len, 0)); - const size_t img_tokens = static_cast(h) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_lens_pe(int h, + int w, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + 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(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> 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> img_ids(h_len * w_len, std::vector(3, 0.0f)); std::vector h_ids = linspace(0.f, static_cast(h_len - 1), h_len); std::vector w_ids = linspace(0.f, static_cast(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 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& axes_dim) { - std::vector> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len); - std::vector> 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(pos_len, 0)); - const size_t img_tokens = static_cast(h_len) * static_cast(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(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage); + __STATIC_INLINE__ Embedding gen_ernie_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + 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(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& 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>(bs * padded_context_len, std::vector(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(img_ids.size() / bs), seq_multi_of); + if (layout) { + layout->append_tokens(img_pad_len); + } if (img_pad_len > 0) { std::vector> img_pad_ids(bs * img_pad_len, std::vector(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>(bs * padded_context_len, std::vector(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(img_ids.size() / bs), seq_multi_of); + if (layout) { + layout->append_tokens(img_pad_len); + } if (img_pad_len > 0) { std::vector> img_pad_ids(bs * img_pad_len, std::vector(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> 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(3, 0.f)); + if (layout) { + layout->append_tokens(image_pad); + } } const int sigvq_start = static_cast(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(), 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 gen_llada_image_edit_pe(int h, - int w, - int patch_size, - int context_len, - int sigvq_len, - int seq_multi_of, - int theta, - const std::vector& 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(theta), axes_dim, {}); + __STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h, + int w, + int patch_size, + int context_len, + int sigvq_len, + int seq_multi_of, + int theta, + const std::vector& 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(theta), axes_dim, result.layout, &result.frequencies); + return result; } - __STATIC_INLINE__ std::vector 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& axes_dim) { - std::vector> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of); - std::vector> 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(pos_len, 0)); - size_t cursor = context_len + bound_mod(context_len, seq_multi_of); - size_t img_tokens = static_cast(h_len) * static_cast(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(theta), axes_dim, wrap_dims); + __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, + const std::vector& axes_dim) { + 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(theta), axes_dim, result.layout, &result.frequencies); + return result; } // Generate z_image positional embeddings - __STATIC_INLINE__ std::vector gen_z_image_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - int seq_multi_of, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode); - std::vector> 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(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(h_len) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_z_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + int theta, + const std::vector& axes_dim) { + 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(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx, diff --git a/src/model/common/rope_circular.hpp b/src/model/common/rope_circular.hpp new file mode 100644 index 00000000..19dccb95 --- /dev/null +++ b/src/model/common/rope_circular.hpp @@ -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(region.height_axis)) { + period = region.height_period; + } else if (circular_x && frequency.axis == static_cast(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__ diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index a8977c96..9ace46d5 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -603,34 +603,37 @@ namespace Anima { return std::pow(extrapolation_ratio, static_cast(axis_dim) / static_cast(axis_dim - 2)); } - static std::vector gen_anima_image_pe_vec(int bs, - int h, - int w, - int patch_size, - int theta, - const std::vector& axes_dim, - float h_extrapolation_ratio, - float w_extrapolation_ratio, - float t_extrapolation_ratio, - const std::vector& ref_latents) { - auto ids = Rope::gen_flux_ids(h, - w, - patch_size, - bs, - static_cast(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& axes_dim, + float h_extrapolation_ratio, + float w_extrapolation_ratio, + float t_extrapolation_ratio, + const std::vector& ref_latents) { + Rope::Embedding result; + result.batch_size = bs; + result.ids = Rope::gen_flux_ids(h, + w, + patch_size, + bs, + static_cast(axes_dim.size()), + 0, + {}, + ref_latents, + Rope::RefIndexMode::FIXED, + 1.0f, + false, &result.positions); std::vector axis_thetas = { static_cast(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]), static_cast(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]), static_cast(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& x_tensor, @@ -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(h_pad), - static_cast(w_pad), - static_cast(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(h_pad), + static_cast(w_pad), + static_cast(config.patch_size), + config.theta, + config.axes_dim, + 4.0f, + 4.0f, + 1.0f, + ref_latents)); int64_t image_pos_len = static_cast(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()); diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 8c91bb55..1ca13709 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -720,15 +720,18 @@ namespace Boogu { } } - __STATIC_INLINE__ std::vector gen_boogu_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - const std::vector& ref_latents, - int theta, - const std::vector& axes_dim) { - std::vector> ids; + __STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + Rope::Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); + auto& ids = result.ids; ids.reserve(static_cast(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(theta), axes_dim); + result.values = Rope::embed_nd(ids, bs, static_cast(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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); int pos_len = static_cast(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()); diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 81bf142a..b284cb83 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); int pos_len = static_cast(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()); diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 67e79d3d..825dabe1 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(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(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); diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 9e2aa361..3254571d 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -149,18 +149,21 @@ namespace Ideogram4 { return std::make_shared(in_features, out_features, bias); } - __STATIC_INLINE__ std::vector gen_ideogram4_pe(int grid_h, - int grid_w, - int bs, - int context_len, - int head_dim, - int rope_theta, - const std::vector& mrope_section, - bool circular_x = false, - bool circular_y = false) { + __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& mrope_section) { GGML_ASSERT(bs == 1); - std::vector> ids(static_cast(bs) * (context_len + grid_h * grid_w), - std::vector(3, 0.f)); + 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(bs) * (context_len + grid_h * grid_w), + std::vector(3, 0.f)); + auto& ids = result.ids; for (int i = 0; i < context_len; ++i) { ids[i] = {static_cast(i), static_cast(i), static_cast(i)}; @@ -175,29 +178,13 @@ namespace Ideogram4 { } } - std::vector> axis_wrap_dims(3); - if (circular_y || circular_x) { - size_t total_len = static_cast(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(context_len); idx < total_len; ++idx) { - axis_wrap_dims[1][idx] = grid_h; - } - } - if (circular_x) { - for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { - axis_wrap_dims[2][idx] = grid_w; - } - } - } - - return Rope::embed_interleaved_mrope(ids, - bs, - static_cast(rope_theta), - head_dim, - mrope_section, - axis_wrap_dims); + result.values = Rope::embed_interleaved_mrope(ids, + bs, + static_cast(rope_theta), + head_dim, + mrope_section, + &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(grid_h), - static_cast(grid_w), - static_cast(x->ne[3]), - static_cast(context_len), - static_cast(head_dim), - static_cast(config.rope_theta), - config.mrope_section, - runner_ctx.circular_x_enabled, - runner_ctx.circular_y_enabled); + pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), + static_cast(grid_w), + static_cast(x->ne[3]), + static_cast(context_len), + static_cast(head_dim), + static_cast(config.rope_theta), + config.mrope_section)); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index 84ecfe0b..e58df410 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -689,23 +689,28 @@ namespace Krea2 { } }; - __STATIC_INLINE__ std::vector gen_krea2_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - float theta, - const std::vector& axes_dim, - const std::vector& 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& axes_dim, + const std::vector& 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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(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(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim, + ref_latents, + ref_image_params.ref_index_mode)); int pos_len = static_cast(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()); diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 35b4b517..8fb1512b 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -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(x->ne[1]), - static_cast(x->ne[0]), - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); int pos_len = static_cast(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()); diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index d03383a5..e69c3ecd 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); int pos_len = static_cast(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(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(context->ne[1]), - semantic != nullptr ? static_cast(semantic->ne[1]) : 0, - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + semantic != nullptr ? static_cast(semantic->ne[1]) : 0, + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); int pos_len = static_cast(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()); diff --git a/src/model/diffusion/mage_flow.hpp b/src/model/diffusion/mage_flow.hpp index 89df0f48..1b9b9f9b 100644 --- a/src/model/diffusion/mage_flow.hpp +++ b/src/model/diffusion/mage_flow.hpp @@ -110,13 +110,13 @@ namespace MageFlow { } int batch_size = static_cast(x->ne[3]); - pe_vec = Rope::gen_mage_flow_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - batch_size, - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + batch_size, + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); int pos_len = static_cast(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()); diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index 06c4db5a..0ddacf02 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -154,18 +154,26 @@ namespace MiniT2I { return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), head_dim, 10000.f)); } - inline std::vector 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 out(static_cast(length) * (head_dim / 2) * 4); std::vector freqs(quarter); for (int i = 0; i < quarter; ++i) { freqs[i] = 1.0f / std::pow(10000.0f, static_cast(2 * i) / static_cast(dim)); } + for (int axis : {1, 2}) { + for (float frequency : freqs) { + result.frequencies.push_back({static_cast(axis), frequency}); + } + } for (int y = 0; y < side; ++y) { for (int x = 0; x < side; ++x) { + result.ids.push_back({0.f, static_cast(y), static_cast(x)}); int pos = y * side + x; size_t base = static_cast(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(img_side), static_cast(config.hidden_size)); auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); - auto img_pe_vec = make_vision_rope(static_cast(img_side), static_cast(config.head_dim)); + auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(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& x_tensor, diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index f82fc36f..c82b9d71 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -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 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, diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 59eac295..6edc8342 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -135,13 +135,13 @@ namespace Pid { return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), dim, theta)); } - inline std::vector make_rope_2d(int height, - int width, - int dim, - float theta = 10000.f, - float scale = 16.f, - int ref_grid_h = 0, - int ref_grid_w = 0) { + inline Rope::Embedding make_rope_2d(int height, + int width, + int dim, + float theta = 10000.f, + float scale = 16.f, + int ref_grid_h = 0, + int ref_grid_w = 0) { GGML_ASSERT(dim % 4 == 0); 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 Ws = Wp / config.patch_size; - pos_img_vec = make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.hidden_size / config.num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w)); + pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.hidden_size / config.num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(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(Hs), - static_cast(Ws), - static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w)); + pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 07364958..82d17300 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -635,18 +635,16 @@ namespace Qwen { ref_index_mode = Rope::RefIndexMode::DECREASE; } - pe_vec = Rope::gen_qwen_image_pe(time_len, - static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - batch_size, - static_cast(context->ne[1]), - ref_latents, - ref_index_mode, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, + static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + batch_size, + static_cast(context->ne[1]), + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); int pos_len = static_cast(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); diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 37483125..9884325c 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -68,6 +68,7 @@ namespace Qwen { std::vector segments; std::vector> positions; int64_t prefix_length = 0; + Rope::PositionLayout rope_layout; static QwenImage21Layout build(int64_t text_length, const sd::Tensor& image_slots, @@ -82,6 +83,7 @@ namespace Qwen { auto [height, width] = image_shapes[index]; int64_t start = static_cast(layout.positions.size()); layout.segments.push_back({start, start + height * width, context_start, index}); + layout.rope_layout.append_image(static_cast(height), static_cast(width)); for (int64_t h = 0; h < height; ++h) { for (int64_t w = 0; w < width; ++w) { layout.positions.push_back({static_cast(position), @@ -106,6 +108,7 @@ namespace Qwen { } else { int64_t start = static_cast(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(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>(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(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) { diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index a1012cad..d654bd38 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -642,18 +642,16 @@ namespace ZImage { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - SEQ_MULTI_OF, - ref_latents, - ref_index_mode, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + SEQ_MULTI_OF, + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); int pos_len = static_cast(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);