mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-23 22:47:41 -05:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
88411ef1e0 | ||
|
|
500ef5fa7c | ||
|
|
2a4ebba818 |
@@ -0,0 +1,61 @@
|
||||
name: Close PRs from organization forks
|
||||
|
||||
on:
|
||||
pull_request_target:
|
||||
types: [opened, reopened]
|
||||
|
||||
permissions:
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
close-organization-fork-pr:
|
||||
if: >-
|
||||
github.event.pull_request.head.repo.owner.type == 'Organization' &&
|
||||
github.event.pull_request.head.repo.id != github.event.pull_request.base.repo.id
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- name: Explain the contribution policy and close the PR
|
||||
uses: actions/github-script@v9
|
||||
with:
|
||||
script: |
|
||||
const { data: pr } = await github.rest.pulls.get({
|
||||
...context.repo,
|
||||
pull_number: context.issue.number,
|
||||
});
|
||||
const headRepo = pr.head.repo;
|
||||
if (pr.state !== 'open' || !headRepo ||
|
||||
headRepo.id === pr.base.repo.id || headRepo.owner.type !== 'Organization') {
|
||||
return;
|
||||
}
|
||||
|
||||
const marker = '<!-- organization-fork-policy -->';
|
||||
const comments = await github.paginate(github.rest.issues.listComments, {
|
||||
...context.repo,
|
||||
issue_number: pr.number,
|
||||
per_page: 100,
|
||||
});
|
||||
const alreadyExplained = comments.some(comment =>
|
||||
comment.user?.login === 'github-actions[bot]' && comment.body?.includes(marker));
|
||||
if (!alreadyExplained) {
|
||||
await github.rest.issues.createComment({
|
||||
...context.repo,
|
||||
issue_number: pr.number,
|
||||
body: [
|
||||
marker,
|
||||
'This repository requires contributions from forks to use a personal fork with **Allow edits from maintainers** enabled.',
|
||||
'GitHub does not support this option for organization-owned forks, so this PR is being closed automatically.',
|
||||
'Please open a new PR from a fork in your personal GitHub account and enable **Allow edits from maintainers** so maintainers can help update the branch.',
|
||||
'See [the GitHub documentation](https://docs.github.com/en/pull-requests/how-tos/work-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork).',
|
||||
].join('\n\n'),
|
||||
});
|
||||
}
|
||||
await github.rest.pulls.update({
|
||||
...context.repo,
|
||||
pull_number: pr.number,
|
||||
state: 'closed',
|
||||
});
|
||||
@@ -12,6 +12,10 @@ If you want to update a third-party dependency, please open an issue first inste
|
||||
|
||||
## Pull Requests
|
||||
|
||||
When contributing from a fork, use a fork under your personal GitHub account and enable **Allow edits from maintainers**. This lets maintainers make follow-up fixes directly on the PR branch.
|
||||
|
||||
PRs from organization-owned forks are automatically closed when opened or reopened because GitHub does not support this maintainer-edit option for those forks. Submit the changes from a personal fork instead. See [GitHub's documentation](https://docs.github.com/en/pull-requests/how-tos/work-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork).
|
||||
|
||||
Keep each PR focused on one clear change. Large or overly complex PRs are harder to review and may not be merged.
|
||||
|
||||
Do not include test code or test scripts in commits or PRs. Keep them local and report verification results in the PR description.
|
||||
|
||||
@@ -161,6 +161,9 @@ resident allocations. Vulkan reports exceeding total memory are rejected because
|
||||
its heap-budget subtraction can underflow. Other backends use the cap instead of
|
||||
treating such reports as zero free memory. Failed checks log the reported free and
|
||||
total memory alongside tracked weight and runtime allocations.
|
||||
With `--mmap`, device-backed mappings count toward these budgets at their full
|
||||
mapped-file size, once per device buffer even when multiple parameter blocks
|
||||
share it. Mappings retained in the loader cache continue to count.
|
||||
|
||||
Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
|
||||
used diffusion weights have priority. Each component's weights use the first
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
#endif
|
||||
|
||||
#include "core/util.h"
|
||||
#include "ggml-backend-impl.h"
|
||||
#include "ggml-impl.h"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
@@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
|
||||
return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
|
||||
}
|
||||
|
||||
ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
|
||||
void* ptr,
|
||||
size_t size,
|
||||
size_t max_tensor_size) {
|
||||
ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
|
||||
if (buffer != nullptr && buffer->context == nullptr) {
|
||||
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
|
||||
if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
|
||||
// Metal can wrap a failed mapping in a non-null buffer. Its free callback also
|
||||
// dereferences the missing context, so only release the outer buffer.
|
||||
buffer->iface.free_buffer = nullptr;
|
||||
ggml_backend_buffer_free(buffer);
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
return buffer;
|
||||
}
|
||||
|
||||
bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
|
||||
#ifdef SD_USE_CUDA
|
||||
if (!sd_backend_is(backend, "CUDA")) {
|
||||
|
||||
@@ -88,6 +88,10 @@ private:
|
||||
bool sd_backend_is(ggml_backend_t backend, const std::string& name);
|
||||
bool sd_backend_is_cpu(ggml_backend_t backend);
|
||||
bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
|
||||
ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
|
||||
void* ptr,
|
||||
size_t size,
|
||||
size_t max_tensor_size);
|
||||
ggml_backend_t sd_backend_cpu_init();
|
||||
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
|
||||
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
|
||||
|
||||
+303
-343
@@ -5,6 +5,7 @@
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
#include <set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#include "core/ggml_extend.h"
|
||||
#include "core/ggml_runner.h"
|
||||
@@ -16,6 +17,45 @@ namespace Rope {
|
||||
ErnieImage,
|
||||
};
|
||||
|
||||
struct SpatialRegion {
|
||||
size_t begin;
|
||||
size_t count;
|
||||
float height_period;
|
||||
float width_period;
|
||||
int height_axis = 1;
|
||||
int width_axis = 2;
|
||||
};
|
||||
|
||||
struct PositionLayout {
|
||||
// Token ranges are relative to one batch item.
|
||||
std::vector<SpatialRegion> images;
|
||||
size_t token_count = 0;
|
||||
|
||||
void append_tokens(size_t count) {
|
||||
token_count += count;
|
||||
}
|
||||
|
||||
void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) {
|
||||
size_t count = static_cast<size_t>(height) * width * frames;
|
||||
images.push_back({token_count, count, height * height_step, width * width_step});
|
||||
append_tokens(count);
|
||||
}
|
||||
};
|
||||
|
||||
struct Frequency {
|
||||
size_t axis;
|
||||
float omega;
|
||||
};
|
||||
|
||||
struct Embedding {
|
||||
std::vector<float> values;
|
||||
std::vector<std::vector<float>> ids;
|
||||
PositionLayout positions;
|
||||
std::vector<Frequency> frequencies;
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix;
|
||||
int batch_size = 1;
|
||||
};
|
||||
|
||||
enum class RefIndexMode {
|
||||
FIXED,
|
||||
INCREASE,
|
||||
@@ -56,40 +96,25 @@ namespace Rope {
|
||||
return flat_vec;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
int dim,
|
||||
float theta,
|
||||
const std::vector<int>& axis_wrap_dims = {}) {
|
||||
__STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
|
||||
assert(dim % 2 == 0);
|
||||
int half_dim = dim / 2;
|
||||
|
||||
int half_dim = dim / 2;
|
||||
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
|
||||
|
||||
std::vector<float> omega(half_dim);
|
||||
for (int i = 0; i < half_dim; ++i) {
|
||||
omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
|
||||
}
|
||||
return omega;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
const std::vector<float>& omega) {
|
||||
int half_dim = static_cast<int>(omega.size());
|
||||
size_t pos_size = pos.size();
|
||||
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
|
||||
for (size_t i = 0; i < pos_size; ++i) {
|
||||
for (size_t j = 0; j < half_dim; ++j) {
|
||||
float angle = pos[i] * omega[j];
|
||||
if (!axis_wrap_dims.empty()) {
|
||||
size_t wrap_size = axis_wrap_dims.size();
|
||||
// mod batch size since we only store this for one item in the batch
|
||||
size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
|
||||
int wrap_dim = axis_wrap_dims[wrap_idx];
|
||||
if (wrap_dim > 0) {
|
||||
constexpr float TWO_PI = 6.28318530717958647692f;
|
||||
float cycles = omega[j] * wrap_dim / TWO_PI;
|
||||
// closest periodic harmonic, necessary to ensure things neatly tile
|
||||
// without this round, things don't tile at the boundaries and you end up
|
||||
// with the model knowing what is "center"
|
||||
float rounded = std::round(cycles);
|
||||
angle = pos[i] * TWO_PI * rounded / wrap_dim;
|
||||
}
|
||||
}
|
||||
|
||||
out[i][j] = angle;
|
||||
}
|
||||
@@ -108,6 +133,12 @@ namespace Rope {
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
int dim,
|
||||
float theta) {
|
||||
return rope(pos, rope_frequencies(dim, theta));
|
||||
}
|
||||
|
||||
// Generate IDs for image patches and text
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
|
||||
@@ -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<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
|
||||
|
||||
int h_start = h_offset;
|
||||
@@ -192,8 +227,8 @@ namespace Rope {
|
||||
int bs,
|
||||
const std::vector<float>& axis_thetas,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {},
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix,
|
||||
std::vector<Frequency>* frequencies = nullptr) {
|
||||
std::vector<std::vector<float>> trans_ids = transpose(ids);
|
||||
size_t pos_len = ids.size() / bs;
|
||||
size_t num_axes = axes_dim.size();
|
||||
@@ -205,19 +240,25 @@ namespace Rope {
|
||||
for (int d : axes_dim)
|
||||
emb_dim += d / 2;
|
||||
|
||||
if (frequencies) {
|
||||
frequencies->clear();
|
||||
frequencies->reserve(emb_dim);
|
||||
}
|
||||
std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
|
||||
size_t offset = 0;
|
||||
for (size_t i = 0; i < num_axes; ++i) {
|
||||
std::vector<int> axis_wrap_dims;
|
||||
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
|
||||
axis_wrap_dims = wrap_dims[i];
|
||||
}
|
||||
float axis_theta = 10000.0f;
|
||||
if (!axis_thetas.empty()) {
|
||||
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
|
||||
}
|
||||
auto omega = rope_frequencies(axes_dim[i], axis_theta);
|
||||
if (frequencies) {
|
||||
for (float frequency : omega) {
|
||||
frequencies->push_back({i, frequency});
|
||||
}
|
||||
}
|
||||
std::vector<std::vector<float>> rope_emb =
|
||||
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
|
||||
rope(trans_ids[i], omega); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
|
||||
for (int b = 0; b < bs; ++b) {
|
||||
for (int j = 0; j < pos_len; ++j) {
|
||||
for (int k = 0; k < rope_emb[0].size(); ++k) {
|
||||
@@ -253,10 +294,10 @@ namespace Rope {
|
||||
int bs,
|
||||
float theta,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {},
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix,
|
||||
std::vector<Frequency>* frequencies = nullptr) {
|
||||
std::vector<float> axis_thetas(axes_dim.size(), theta);
|
||||
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
|
||||
return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
|
||||
@@ -264,7 +305,7 @@ namespace Rope {
|
||||
float theta,
|
||||
int head_dim,
|
||||
const std::vector<int>& mrope_section,
|
||||
const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
|
||||
std::vector<Frequency>* frequencies = nullptr) {
|
||||
GGML_ASSERT(bs > 0);
|
||||
GGML_ASSERT(head_dim % 2 == 0);
|
||||
GGML_ASSERT(mrope_section.size() >= 3);
|
||||
@@ -273,20 +314,26 @@ namespace Rope {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
int half_dim = head_dim / 2;
|
||||
|
||||
auto omega = rope_frequencies(head_dim, theta);
|
||||
if (frequencies) {
|
||||
frequencies->clear();
|
||||
for (float frequency : omega) {
|
||||
frequencies->push_back({0, frequency});
|
||||
}
|
||||
}
|
||||
std::vector<std::vector<std::vector<float>>> axis_embs;
|
||||
axis_embs.reserve(3);
|
||||
for (int axis = 0; axis < 3; ++axis) {
|
||||
std::vector<int> axis_wrap;
|
||||
if (axis < static_cast<int>(axis_wrap_dims.size())) {
|
||||
axis_wrap = axis_wrap_dims[axis];
|
||||
}
|
||||
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
|
||||
axis_embs.push_back(rope(trans_ids[axis], omega));
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> emb = axis_embs[0];
|
||||
for (int axis = 1; axis < 3; ++axis) {
|
||||
int length = std::min<int>(mrope_section[axis] * 3, half_dim);
|
||||
for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
|
||||
if (frequencies) {
|
||||
(*frequencies)[freq_idx].axis = axis;
|
||||
}
|
||||
for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
|
||||
for (int k = 0; k < 4; ++k) {
|
||||
emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
|
||||
@@ -298,13 +345,13 @@ namespace Rope {
|
||||
return flatten(emb);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> 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<float>(width) / static_cast<float>(ref_grid_w), power);
|
||||
}
|
||||
|
||||
Embedding result;
|
||||
result.positions.append_image(height, width, 1,
|
||||
height > 1 ? scale / (height - 1) : 1.f,
|
||||
width > 1 ? scale / (width - 1) : 1.f);
|
||||
std::vector<float> x_pos;
|
||||
std::vector<float> y_pos;
|
||||
x_pos.reserve(static_cast<size_t>(height) * width);
|
||||
@@ -326,13 +377,20 @@ namespace Rope {
|
||||
float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
|
||||
for (int ix = 0; ix < width; ++ix) {
|
||||
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
|
||||
result.ids.push_back({0.f, y, x});
|
||||
x_pos.push_back(x);
|
||||
y_pos.push_back(y);
|
||||
}
|
||||
}
|
||||
|
||||
auto x_emb = rope(x_pos, dim_axis, theta * w_ntk);
|
||||
auto y_emb = rope(y_pos, dim_axis, theta * h_ntk);
|
||||
auto x_freq = rope_frequencies(dim_axis, theta * w_ntk);
|
||||
auto y_freq = rope_frequencies(dim_axis, theta * h_ntk);
|
||||
auto x_emb = rope(x_pos, x_freq);
|
||||
auto y_emb = rope(y_pos, y_freq);
|
||||
for (int i = 0; i < axis_half_dim; ++i) {
|
||||
result.frequencies.push_back({2, x_freq[i]});
|
||||
result.frequencies.push_back({1, y_freq[i]});
|
||||
}
|
||||
|
||||
std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
|
||||
for (int pos = 0; pos < height * width; ++pos) {
|
||||
@@ -348,7 +406,8 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
result.values = std::move(out);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
|
||||
@@ -359,7 +418,8 @@ namespace Rope {
|
||||
RefIndexMode ref_index_mode,
|
||||
float ref_index_scale,
|
||||
bool scale_rope,
|
||||
int base_offset = 0) {
|
||||
int base_offset = 0,
|
||||
PositionLayout* layout = nullptr) {
|
||||
std::vector<std::vector<float>> ids;
|
||||
int curr_h_offset = 0;
|
||||
int curr_w_offset = 0;
|
||||
@@ -386,7 +446,8 @@ namespace Rope {
|
||||
static_cast<int>(index * ref_index_scale),
|
||||
h_offset + base_offset,
|
||||
w_offset + base_offset,
|
||||
scale_rope);
|
||||
scale_rope,
|
||||
layout);
|
||||
ids = concat_ids(ids, ref_ids, bs);
|
||||
|
||||
if (ref_index_mode == RefIndexMode::INCREASE) {
|
||||
@@ -409,88 +470,53 @@ namespace Rope {
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
float ref_index_scale,
|
||||
bool is_longcat) {
|
||||
bool is_longcat,
|
||||
PositionLayout* layout = nullptr) {
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
int x_index = is_longcat ? 1 : 0;
|
||||
|
||||
auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
|
||||
int offset = is_longcat ? context_len : 0;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset, false, layout);
|
||||
|
||||
auto ids = concat_ids(txt_ids, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset, layout);
|
||||
ids = concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
// Generate flux positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_flux_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
std::set<int> txt_arange_dims,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
float ref_index_scale,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim,
|
||||
bool is_longcat) {
|
||||
std::vector<std::vector<float>> ids = gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
bs,
|
||||
static_cast<int>(axes_dim.size()),
|
||||
context_len,
|
||||
txt_arange_dims,
|
||||
ref_latents,
|
||||
ref_index_mode,
|
||||
ref_index_scale,
|
||||
is_longcat);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
size_t cursor = context_len; // text first
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
cursor += img_tokens;
|
||||
// reference latents
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
if (ref == nullptr) {
|
||||
continue;
|
||||
}
|
||||
int ref_h = static_cast<int>(ref->ne[1]);
|
||||
int ref_w = static_cast<int>(ref->ne[0]);
|
||||
int ref_h_l = (ref_h + (patch_size / 2)) / patch_size;
|
||||
int ref_w_l = (ref_w + (patch_size / 2)) / patch_size;
|
||||
size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
|
||||
for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = ref_h_l;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = ref_w_l;
|
||||
}
|
||||
}
|
||||
cursor += ref_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
__STATIC_INLINE__ Embedding gen_flux_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
std::set<int> txt_arange_dims,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
float ref_index_scale,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim,
|
||||
bool is_longcat) {
|
||||
Embedding result;
|
||||
result.batch_size = bs;
|
||||
result.ids = gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
bs,
|
||||
static_cast<int>(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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> 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<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
|
||||
|
||||
if (scale_rope) {
|
||||
@@ -573,7 +603,11 @@ namespace Rope {
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode) {
|
||||
RefIndexMode ref_index_mode,
|
||||
PositionLayout* layout = nullptr) {
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
int txt_id_start = std::max(h_len, w_len) / 2;
|
||||
@@ -585,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<float> gen_qwen_image_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int pad_h = (patch_size - (h % patch_size)) % patch_size;
|
||||
int pad_w = (patch_size - (w % patch_size)) % patch_size;
|
||||
int h_len = (h + pad_h) / patch_size;
|
||||
int w_len = (w + pad_w) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
const size_t total_tokens = ids.size();
|
||||
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
|
||||
size_t cursor = context_len; // ignore text tokens
|
||||
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
cursor += img_tokens;
|
||||
// For each reference image, store wrap sizes as well
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
if (ref == nullptr) {
|
||||
continue;
|
||||
}
|
||||
int ref_h = static_cast<int>(ref->ne[1]);
|
||||
int ref_w = static_cast<int>(ref->ne[0]);
|
||||
int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size;
|
||||
int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size;
|
||||
int ref_h_len = (ref_h + ref_pad_h) / patch_size;
|
||||
int ref_w_len = (ref_w + ref_pad_w) / patch_size;
|
||||
size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
|
||||
for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = ref_h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = ref_w_len;
|
||||
}
|
||||
}
|
||||
cursor += ref_n_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
__STATIC_INLINE__ Embedding gen_qwen_image_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
int theta,
|
||||
const std::vector<int>& 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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
__STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
Embedding result;
|
||||
result.batch_size = bs;
|
||||
result.positions.append_tokens(context_len);
|
||||
const int axes_dim_num = static_cast<int>(axes_dim.size());
|
||||
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
|
||||
auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) {
|
||||
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
|
||||
std::vector<float>(axes_dim_num, 0.f));
|
||||
result.positions.append_image(image_h, image_w);
|
||||
int h_start = -(image_h - image_h / 2);
|
||||
int w_start = -(image_w - image_w / 2);
|
||||
for (int b = 0; b < bs; ++b) {
|
||||
@@ -692,15 +685,18 @@ namespace Rope {
|
||||
static_cast<int>(i + 1));
|
||||
ids = concat_ids(ids, ref_ids, bs);
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
result.ids = std::move(ids);
|
||||
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
bool scale_rope = true) {
|
||||
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
|
||||
bool scale_rope = true,
|
||||
PositionLayout* layout = nullptr) {
|
||||
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope, layout);
|
||||
|
||||
int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
|
||||
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
|
||||
@@ -711,44 +707,37 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
__STATIC_INLINE__ Embedding gen_lens_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
const std::vector<int>& 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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len) {
|
||||
int context_len,
|
||||
PositionLayout* layout = nullptr) {
|
||||
int h_len = h / patch_size;
|
||||
int w_len = w / patch_size;
|
||||
|
||||
if (layout) {
|
||||
layout->append_image(h_len, w_len);
|
||||
}
|
||||
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
|
||||
std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
|
||||
std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
|
||||
@@ -774,39 +763,25 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
return concat_ids(img_ids_repeated, txt_ids, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int h_len = h / patch_size;
|
||||
int w_len = w / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
|
||||
__STATIC_INLINE__ Embedding gen_ernie_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
const std::vector<int>& 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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Generate wan positional embeddings
|
||||
@@ -905,7 +880,8 @@ namespace Rope {
|
||||
int context_len,
|
||||
int seq_multi_of,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode) {
|
||||
RefIndexMode ref_index_mode,
|
||||
PositionLayout* layout = nullptr) {
|
||||
SD_UNUSED(ref_index_mode);
|
||||
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
|
||||
@@ -913,11 +889,17 @@ namespace Rope {
|
||||
txt_ids[i][0] = (i % padded_context_len) + 1.f;
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(padded_context_len);
|
||||
}
|
||||
int axes_dim_num = 3;
|
||||
int index = padded_context_len + 1;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
|
||||
|
||||
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
|
||||
if (layout) {
|
||||
layout->append_tokens(img_pad_len);
|
||||
}
|
||||
if (img_pad_len > 0) {
|
||||
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
|
||||
img_ids = concat_ids(img_ids, img_pad_ids, bs);
|
||||
@@ -936,7 +918,8 @@ namespace Rope {
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int seq_multi_of) {
|
||||
int seq_multi_of,
|
||||
PositionLayout* layout = nullptr) {
|
||||
int context_pad_len = bound_mod(context_len, seq_multi_of);
|
||||
int padded_context_len = context_len + context_pad_len;
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
|
||||
@@ -947,11 +930,17 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(padded_context_len);
|
||||
}
|
||||
int axes_dim_num = 3;
|
||||
int index = padded_context_len + 1;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
|
||||
|
||||
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
|
||||
if (layout) {
|
||||
layout->append_tokens(img_pad_len);
|
||||
}
|
||||
if (img_pad_len > 0) {
|
||||
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
|
||||
img_ids = concat_ids(img_ids, img_pad_ids, bs);
|
||||
@@ -968,7 +957,8 @@ namespace Rope {
|
||||
int patch_size,
|
||||
int context_len,
|
||||
int sigvq_len,
|
||||
int seq_multi_of) {
|
||||
int seq_multi_of,
|
||||
PositionLayout* layout = nullptr) {
|
||||
const int context_pad = bound_mod(context_len, seq_multi_of);
|
||||
const int padded_context = context_len + context_pad;
|
||||
const int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
@@ -994,11 +984,17 @@ namespace Rope {
|
||||
cursor += 2;
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(cap_ids.size());
|
||||
}
|
||||
std::vector<std::vector<float>> img_ids;
|
||||
for (int copy = 0; copy < 2; ++copy) {
|
||||
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]);
|
||||
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy], 0, 0, false, layout);
|
||||
img_ids.insert(img_ids.end(), ids.begin(), ids.end());
|
||||
img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
|
||||
if (layout) {
|
||||
layout->append_tokens(image_pad);
|
||||
}
|
||||
}
|
||||
|
||||
const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
|
||||
@@ -1016,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<float> 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<int>& axes_dim) {
|
||||
auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of);
|
||||
return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {});
|
||||
__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<int>& axes_dim) {
|
||||
Embedding result;
|
||||
result.batch_size = 1;
|
||||
result.ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of, &result.positions);
|
||||
result.values = embed_nd(result.ids, 1, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int seq_multi_of,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int pad_h = (patch_size - (h % patch_size)) % patch_size;
|
||||
int pad_w = (patch_size - (w % patch_size)) % patch_size;
|
||||
int h_len = (h + pad_h) / patch_size;
|
||||
int w_len = (w + pad_w) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
size_t cursor = context_len + bound_mod(context_len, seq_multi_of);
|
||||
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
__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<int>& 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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Generate z_image positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int seq_multi_of,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int pad_h = (patch_size - (h % patch_size)) % patch_size;
|
||||
int pad_w = (patch_size - (w % patch_size)) % patch_size;
|
||||
int h_len = (h + pad_h) / patch_size;
|
||||
int w_len = (w + pad_w) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding)
|
||||
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
__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<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
int theta,
|
||||
const std::vector<int>& 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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
|
||||
|
||||
@@ -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__
|
||||
@@ -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,
|
||||
@@ -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());
|
||||
|
||||
@@ -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++) {
|
||||
@@ -741,15 +744,18 @@ namespace Boogu {
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
|
||||
int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
|
||||
result.positions.append_image(ref_h_tokens, ref_w_tokens);
|
||||
append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
|
||||
pe_shift += std::max(ref_h_tokens, ref_w_tokens);
|
||||
}
|
||||
|
||||
int h_tokens = patched_token_count(h, patch_size);
|
||||
int w_tokens = patched_token_count(w, patch_size);
|
||||
result.positions.append_image(h_tokens, w_tokens);
|
||||
append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
|
||||
|
||||
return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
result.values = Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
struct BooguImageRunner : public DiffusionModelRunner {
|
||||
@@ -793,14 +799,14 @@ namespace Boogu {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = gen_boogu_pe(static_cast<int>(x->ne[1]),
|
||||
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());
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -149,18 +149,21 @@ namespace Ideogram4 {
|
||||
return std::make_shared<Linear>(in_features, out_features, bias);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
|
||||
int grid_w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int head_dim,
|
||||
int rope_theta,
|
||||
const std::vector<int>& mrope_section,
|
||||
bool circular_x = false,
|
||||
bool circular_y = false) {
|
||||
__STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h,
|
||||
int grid_w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int head_dim,
|
||||
int rope_theta,
|
||||
const std::vector<int>& mrope_section) {
|
||||
GGML_ASSERT(bs == 1);
|
||||
std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
|
||||
std::vector<float>(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<size_t>(bs) * (context_len + grid_h * grid_w),
|
||||
std::vector<float>(3, 0.f));
|
||||
auto& ids = result.ids;
|
||||
|
||||
for (int i = 0; i < context_len; ++i) {
|
||||
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
|
||||
@@ -175,29 +178,13 @@ namespace Ideogram4 {
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<int>> axis_wrap_dims(3);
|
||||
if (circular_y || circular_x) {
|
||||
size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
|
||||
axis_wrap_dims[1].assign(total_len, 0);
|
||||
axis_wrap_dims[2].assign(total_len, 0);
|
||||
if (circular_y) {
|
||||
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
|
||||
axis_wrap_dims[1][idx] = grid_h;
|
||||
}
|
||||
}
|
||||
if (circular_x) {
|
||||
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
|
||||
axis_wrap_dims[2][idx] = grid_w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return Rope::embed_interleaved_mrope(ids,
|
||||
bs,
|
||||
static_cast<float>(rope_theta),
|
||||
head_dim,
|
||||
mrope_section,
|
||||
axis_wrap_dims);
|
||||
result.values = Rope::embed_interleaved_mrope(ids,
|
||||
bs,
|
||||
static_cast<float>(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<int>(grid_h),
|
||||
static_cast<int>(grid_w),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context_len),
|
||||
static_cast<int>(head_dim),
|
||||
static_cast<int>(config.rope_theta),
|
||||
config.mrope_section,
|
||||
runner_ctx.circular_x_enabled,
|
||||
runner_ctx.circular_y_enabled);
|
||||
pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
|
||||
static_cast<int>(grid_w),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context_len),
|
||||
static_cast<int>(head_dim),
|
||||
static_cast<int>(config.rope_theta),
|
||||
config.mrope_section));
|
||||
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());
|
||||
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -384,14 +384,12 @@ namespace Lens {
|
||||
GGML_ASSERT(!context_tensor.empty());
|
||||
ggml_tensor* context = make_input(context_tensor);
|
||||
|
||||
pe_vec = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.theta,
|
||||
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());
|
||||
|
||||
@@ -412,16 +412,14 @@ namespace LLaDAImage {
|
||||
GGML_ASSERT(!context_tensor.empty());
|
||||
ggml_tensor* context = make_input(context_tensor);
|
||||
|
||||
pe_vec = Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
ZImage::SEQ_MULTI_OF,
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
ZImage::SEQ_MULTI_OF,
|
||||
config.theta,
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
@@ -461,14 +459,14 @@ namespace LLaDAImage {
|
||||
ggml_tensor* source = make_input(source_tensor);
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
|
||||
pe_vec = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
|
||||
ZImage::SEQ_MULTI_OF,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
|
||||
ZImage::SEQ_MULTI_OF,
|
||||
config.theta,
|
||||
config.axes_dim));
|
||||
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());
|
||||
|
||||
@@ -110,13 +110,13 @@ namespace MageFlow {
|
||||
}
|
||||
|
||||
int batch_size = static_cast<int>(x->ne[3]);
|
||||
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
batch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
ref_latents,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
batch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
ref_latents,
|
||||
config.theta,
|
||||
config.axes_dim));
|
||||
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());
|
||||
|
||||
@@ -154,18 +154,26 @@ namespace MiniT2I {
|
||||
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
|
||||
}
|
||||
|
||||
inline std::vector<float> make_vision_rope(int side, int head_dim) {
|
||||
inline Rope::Embedding make_vision_rope(int side, int head_dim) {
|
||||
GGML_ASSERT(head_dim % 4 == 0);
|
||||
int dim = head_dim / 2;
|
||||
int quarter = dim / 2;
|
||||
int length = side * side;
|
||||
Rope::Embedding result;
|
||||
result.positions.append_image(side, side);
|
||||
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
|
||||
std::vector<float> freqs(quarter);
|
||||
for (int i = 0; i < quarter; ++i) {
|
||||
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
|
||||
}
|
||||
for (int axis : {1, 2}) {
|
||||
for (float frequency : freqs) {
|
||||
result.frequencies.push_back({static_cast<size_t>(axis), frequency});
|
||||
}
|
||||
}
|
||||
for (int y = 0; y < side; ++y) {
|
||||
for (int x = 0; x < side; ++x) {
|
||||
result.ids.push_back({0.f, static_cast<float>(y), static_cast<float>(x)});
|
||||
int pos = y * side + x;
|
||||
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
|
||||
for (int i = 0; i < quarter; ++i) {
|
||||
@@ -182,7 +190,8 @@ namespace MiniT2I {
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
result.values = std::move(out);
|
||||
return result;
|
||||
}
|
||||
|
||||
struct SwiGLUMlp : public GGMLBlock {
|
||||
@@ -475,6 +484,8 @@ namespace MiniT2I {
|
||||
int64_t cached_txt_len = -1;
|
||||
int64_t cached_hidden_size = -1;
|
||||
int64_t cached_head_dim = -1;
|
||||
bool cached_circular_x = false;
|
||||
bool cached_circular_y = false;
|
||||
|
||||
MiniT2IRunner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
@@ -521,6 +532,8 @@ namespace MiniT2I {
|
||||
cached_txt_len == txt_len &&
|
||||
cached_hidden_size == config.hidden_size &&
|
||||
cached_head_dim == config.head_dim &&
|
||||
cached_circular_x == circular_x_enabled &&
|
||||
cached_circular_y == circular_y_enabled &&
|
||||
cached_pos_embed != nullptr &&
|
||||
cached_txt_pe != nullptr &&
|
||||
cached_joint_pe != nullptr) {
|
||||
@@ -531,7 +544,7 @@ namespace MiniT2I {
|
||||
|
||||
auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
|
||||
auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
|
||||
auto img_pe_vec = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
|
||||
auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
|
||||
auto joint_pe_vec = txt_pe_vec;
|
||||
joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
|
||||
|
||||
@@ -561,6 +574,8 @@ namespace MiniT2I {
|
||||
cached_txt_len = txt_len;
|
||||
cached_hidden_size = config.hidden_size;
|
||||
cached_head_dim = config.head_dim;
|
||||
cached_circular_x = circular_x_enabled;
|
||||
cached_circular_y = circular_y_enabled;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
#include "core/ggml_runner.h"
|
||||
#include "core/tensor_ggml.hpp"
|
||||
#include "model/common/rope.hpp"
|
||||
#include "model/common/rope_circular.hpp"
|
||||
#include "model_manager.h"
|
||||
|
||||
enum class RefImageResizeMode {
|
||||
@@ -184,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
|
||||
protected:
|
||||
std::string prefix;
|
||||
|
||||
std::vector<float> finish_rope_pe(Rope::Embedding embedding) {
|
||||
Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled);
|
||||
return std::move(embedding.values);
|
||||
}
|
||||
|
||||
public:
|
||||
DiffusionModelRunner(ggml_backend_t backend,
|
||||
const std::string& prefix,
|
||||
|
||||
+21
-21
@@ -135,13 +135,13 @@ namespace Pid {
|
||||
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
|
||||
}
|
||||
|
||||
inline std::vector<float> make_rope_2d(int height,
|
||||
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<int>(Hs),
|
||||
static_cast<int>(Ws),
|
||||
static_cast<int>(config.hidden_size / config.num_groups),
|
||||
10000.f,
|
||||
16.f,
|
||||
static_cast<int>(config.rope_ref_grid_h),
|
||||
static_cast<int>(config.rope_ref_grid_w));
|
||||
pos_img_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
|
||||
static_cast<int>(Ws),
|
||||
static_cast<int>(config.hidden_size / config.num_groups),
|
||||
10000.f,
|
||||
16.f,
|
||||
static_cast<int>(config.rope_ref_grid_h),
|
||||
static_cast<int>(config.rope_ref_grid_w)));
|
||||
auto pos_img = ggml_new_tensor_4d(compute_ctx,
|
||||
GGML_TYPE_F32,
|
||||
2,
|
||||
@@ -904,13 +904,13 @@ namespace Pid {
|
||||
1);
|
||||
set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
|
||||
|
||||
pixel_pos_comp_vec = make_rope_2d(static_cast<int>(Hs),
|
||||
static_cast<int>(Ws),
|
||||
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
|
||||
10000.f,
|
||||
16.f,
|
||||
static_cast<int>(config.rope_ref_grid_h),
|
||||
static_cast<int>(config.rope_ref_grid_w));
|
||||
pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
|
||||
static_cast<int>(Ws),
|
||||
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
|
||||
10000.f,
|
||||
16.f,
|
||||
static_cast<int>(config.rope_ref_grid_h),
|
||||
static_cast<int>(config.rope_ref_grid_w)));
|
||||
auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
|
||||
GGML_TYPE_F32,
|
||||
2,
|
||||
|
||||
@@ -635,18 +635,16 @@ namespace Qwen {
|
||||
ref_index_mode = Rope::RefIndexMode::DECREASE;
|
||||
}
|
||||
|
||||
pe_vec = Rope::gen_qwen_image_pe(time_len,
|
||||
static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
batch_size,
|
||||
static_cast<int>(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<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
batch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
ref_latents,
|
||||
ref_index_mode,
|
||||
config.theta,
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
// LOG_VERBOSE("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
|
||||
@@ -68,6 +68,7 @@ namespace Qwen {
|
||||
std::vector<QwenImage21Segment> segments;
|
||||
std::vector<std::vector<float>> positions;
|
||||
int64_t prefix_length = 0;
|
||||
Rope::PositionLayout rope_layout;
|
||||
|
||||
static QwenImage21Layout build(int64_t text_length,
|
||||
const sd::Tensor<int32_t>& image_slots,
|
||||
@@ -82,6 +83,7 @@ namespace Qwen {
|
||||
auto [height, width] = image_shapes[index];
|
||||
int64_t start = static_cast<int64_t>(layout.positions.size());
|
||||
layout.segments.push_back({start, start + height * width, context_start, index});
|
||||
layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
|
||||
for (int64_t h = 0; h < height; ++h) {
|
||||
for (int64_t w = 0; w < width; ++w) {
|
||||
layout.positions.push_back({static_cast<float>(position),
|
||||
@@ -106,6 +108,7 @@ namespace Qwen {
|
||||
} else {
|
||||
int64_t start = static_cast<int64_t>(layout.positions.size());
|
||||
layout.segments.push_back({start, start + i - begin, begin, -1});
|
||||
layout.rope_layout.append_tokens(i - begin);
|
||||
for (int64_t j = begin; j < i; ++j, ++position) {
|
||||
float p = static_cast<float>(position);
|
||||
layout.positions.push_back({p, p, p});
|
||||
@@ -418,14 +421,26 @@ namespace Qwen {
|
||||
}
|
||||
QwenImage21PrefixCache cache;
|
||||
if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
|
||||
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id);
|
||||
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
|
||||
".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
|
||||
cache.prefix_length = layout.prefix_length;
|
||||
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
|
||||
}
|
||||
auto run = [&](const QwenImage21PrefixCache& active_cache) {
|
||||
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
|
||||
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
|
||||
pe_data = Rope::embed_nd(std::vector<std::vector<float>>(first_position, layout.positions.end()), 1, 10000.f, config.axes_dim);
|
||||
Rope::Embedding embedding;
|
||||
embedding.ids.assign(first_position, layout.positions.end());
|
||||
const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
|
||||
embedding.positions.token_count = embedding.ids.size();
|
||||
for (auto region : layout.rope_layout.images) {
|
||||
if (region.begin >= offset) {
|
||||
region.begin -= offset;
|
||||
embedding.positions.images.push_back(region);
|
||||
}
|
||||
}
|
||||
embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
|
||||
pe_data = finish_rope_pe(std::move(embedding));
|
||||
mask_data.clear();
|
||||
if (!cached) {
|
||||
for (const auto& segment : layout.segments) {
|
||||
|
||||
@@ -642,18 +642,16 @@ namespace ZImage {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(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<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]),
|
||||
SEQ_MULTI_OF,
|
||||
ref_latents,
|
||||
ref_index_mode,
|
||||
config.theta,
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
// LOG_VERBOSE("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
|
||||
+58
-6
@@ -874,7 +874,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
|
||||
|
||||
std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
|
||||
std::set<std::string> ignore_tensors,
|
||||
bool writable_mmap) {
|
||||
bool writable_mmap,
|
||||
ggml_backend_dev_t device) {
|
||||
std::set<std::string> names;
|
||||
for (const auto& entry : tensors) {
|
||||
names.insert(entry.first);
|
||||
@@ -896,6 +897,39 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
|
||||
if (!fdata.mmbuffer)
|
||||
continue;
|
||||
|
||||
// Wrapped on first use: a device buffer makes the whole file resident on that device.
|
||||
std::shared_ptr<struct ggml_backend_buffer> file_buffer = device == nullptr ? fdata.mmbuffer : nullptr;
|
||||
bool file_unmappable = false;
|
||||
|
||||
auto buffer_for_file = [&]() -> ggml_backend_buffer_t {
|
||||
if (file_buffer || file_unmappable) {
|
||||
return file_buffer.get();
|
||||
}
|
||||
auto cached = fdata.device_mmbuffers.find(device);
|
||||
if (cached != fdata.device_mmbuffers.end()) {
|
||||
file_buffer = cached->second;
|
||||
return file_buffer.get();
|
||||
}
|
||||
size_t max_tensor_size = 0;
|
||||
for (const auto& ts : fdata.tensors) {
|
||||
max_tensor_size = std::max(max_tensor_size, static_cast<size_t>(ts.nbytes()));
|
||||
}
|
||||
ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device,
|
||||
fdata.mmapped->writable_data(),
|
||||
fdata.mmapped->size(),
|
||||
max_tensor_size);
|
||||
if (buf == nullptr) {
|
||||
LOG_WARN("mmap: %s cannot map '%s', loading it instead",
|
||||
ggml_backend_dev_name(device), fdata.path.c_str());
|
||||
file_unmappable = true;
|
||||
return nullptr;
|
||||
}
|
||||
LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device));
|
||||
file_buffer = std::shared_ptr<struct ggml_backend_buffer>(buf, ggml_backend_buffer_free);
|
||||
fdata.device_mmbuffers[device] = file_buffer;
|
||||
return file_buffer.get();
|
||||
};
|
||||
|
||||
const std::vector<TensorStorage>& file_tensors = fdata.tensors;
|
||||
|
||||
size_t file_mapped_bytes = 0;
|
||||
@@ -944,10 +978,13 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
|
||||
continue;
|
||||
}
|
||||
|
||||
ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
|
||||
uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
|
||||
dst_tensor->buffer = buf_mmap;
|
||||
dst_tensor->data = mmap_data + tensor_offset;
|
||||
ggml_backend_buffer_t buf_mmap = buffer_for_file();
|
||||
if (buf_mmap == nullptr) {
|
||||
break;
|
||||
}
|
||||
uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
|
||||
dst_tensor->buffer = buf_mmap;
|
||||
dst_tensor->data = mmap_data + tensor_offset;
|
||||
|
||||
file_mapped_bytes += tensor_size;
|
||||
file_mapped_tensors++;
|
||||
@@ -956,7 +993,7 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
|
||||
if (file_mapped_bytes > 0) {
|
||||
mapped_tensors += file_mapped_tensors;
|
||||
mapped_bytes += file_mapped_bytes;
|
||||
result.push_back({fdata.mmapped, fdata.mmbuffer});
|
||||
result.push_back({fdata.mmapped, file_buffer});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -972,6 +1009,16 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<ggml_backend_buffer_t> ModelLoader::get_device_mmap_buffers() const {
|
||||
std::vector<ggml_backend_buffer_t> buffers;
|
||||
for (const auto& fdata : file_data) {
|
||||
for (const auto& entry : fdata.device_mmbuffers) {
|
||||
buffers.push_back(entry.second.get());
|
||||
}
|
||||
}
|
||||
return buffers;
|
||||
}
|
||||
|
||||
bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
bool enable_mmap,
|
||||
const std::set<std::string>* target_tensor_names,
|
||||
@@ -1115,6 +1162,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
|
||||
continue;
|
||||
}
|
||||
if (dst_tensor->buffer != nullptr &&
|
||||
std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(),
|
||||
[&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
|
||||
|
||||
|
||||
+5
-1
@@ -20,6 +20,8 @@ struct ModelFileData {
|
||||
std::vector<TensorStorage> tensors;
|
||||
std::shared_ptr<MmapWrapper> mmapped;
|
||||
std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
|
||||
// mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr)
|
||||
std::map<ggml_backend_dev_t, std::shared_ptr<struct ggml_backend_buffer>> device_mmbuffers;
|
||||
bool is_zip;
|
||||
};
|
||||
|
||||
@@ -120,7 +122,9 @@ public:
|
||||
void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
|
||||
std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
|
||||
std::set<std::string> ignore_tensors = {},
|
||||
bool writable = true);
|
||||
bool writable = true,
|
||||
ggml_backend_dev_t device = nullptr);
|
||||
std::vector<ggml_backend_buffer_t> get_device_mmap_buffers() const;
|
||||
bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
bool use_mmap = false,
|
||||
const std::set<std::string>* target_tensor_names = nullptr,
|
||||
|
||||
+50
-28
@@ -780,38 +780,52 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
|
||||
|
||||
bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
|
||||
std::vector<ParamsStorageBlock*>& created_storage_blocks) {
|
||||
std::map<std::string, ggml_tensor*> mmap_candidates;
|
||||
std::map<std::string, TensorState*> mmap_states;
|
||||
// A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing
|
||||
// stages them for it, so they are mapped through a buffer of that GPU's device.
|
||||
struct MmapGroup {
|
||||
std::map<std::string, ggml_tensor*> candidates;
|
||||
std::map<std::string, TensorState*> states;
|
||||
};
|
||||
std::map<ggml_backend_dev_t, MmapGroup> groups;
|
||||
for (TensorState* state : states) {
|
||||
if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr ||
|
||||
state->tensor->data != nullptr || state->tensor->view_src != nullptr) {
|
||||
continue;
|
||||
}
|
||||
mmap_candidates[state->name] = state->tensor;
|
||||
mmap_states[state->name] = state;
|
||||
}
|
||||
if (mmap_candidates.empty()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
|
||||
if (mmap_store.empty()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
auto block = std::make_unique<ParamsStorageBlock>();
|
||||
block->mmap_tensor_stores = std::move(mmap_store);
|
||||
ParamsStorageBlock* raw = block.get();
|
||||
for (const auto& pair : mmap_states) {
|
||||
TensorState* state = pair.second;
|
||||
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
|
||||
block->states.push_back(state);
|
||||
ggml_backend_dev_t device = nullptr;
|
||||
if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) {
|
||||
device = ggml_backend_get_device(state->compute_backend);
|
||||
}
|
||||
MmapGroup& group = groups[device];
|
||||
group.candidates[state->name] = state->tensor;
|
||||
group.states[state->name] = state;
|
||||
}
|
||||
|
||||
if (!block->states.empty()) {
|
||||
params_storage_blocks_.push_back(std::move(block));
|
||||
created_storage_blocks.push_back(raw);
|
||||
for (auto& [device, group] : groups) {
|
||||
// Device buffers wrap read-only mappings only; params that LoRAs are merged into in place
|
||||
// are loaded instead.
|
||||
if (device != nullptr && writable_mmap_) {
|
||||
continue;
|
||||
}
|
||||
auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device);
|
||||
if (mmap_store.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto block = std::make_unique<ParamsStorageBlock>();
|
||||
block->mmap_tensor_stores = std::move(mmap_store);
|
||||
ParamsStorageBlock* raw = block.get();
|
||||
for (const auto& pair : group.states) {
|
||||
TensorState* state = pair.second;
|
||||
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
|
||||
block->states.push_back(state);
|
||||
}
|
||||
}
|
||||
|
||||
if (!block->states.empty()) {
|
||||
params_storage_blocks_.push_back(std::move(block));
|
||||
created_storage_blocks.push_back(raw);
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
|
||||
}
|
||||
|
||||
size_t total_size = 0;
|
||||
auto add_buffer = [&](ggml_backend_buffer_t buffer) {
|
||||
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) {
|
||||
std::unordered_set<ggml_backend_buffer_t> seen;
|
||||
auto add_buffer = [&](ggml_backend_buffer_t buffer) {
|
||||
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) {
|
||||
return;
|
||||
}
|
||||
ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
|
||||
auto split_devices = split_buffer_devices_.find(buffer_type);
|
||||
const bool on_device = split_devices == split_buffer_devices_.end()
|
||||
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
|
||||
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
|
||||
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
|
||||
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
|
||||
return ggml_backend_get_device(entry.first) == compute_device;
|
||||
});
|
||||
if (!on_device) {
|
||||
@@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
|
||||
total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
|
||||
};
|
||||
|
||||
// The loader may retain device mappings after their parameter blocks are released.
|
||||
for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) {
|
||||
add_buffer(buffer);
|
||||
}
|
||||
for (const auto& block : params_storage_blocks_) {
|
||||
if (block != nullptr) {
|
||||
add_buffer(block->buffer);
|
||||
for (const auto& store : block->mmap_tensor_stores) {
|
||||
add_buffer(store.mmbuffer.get());
|
||||
}
|
||||
}
|
||||
}
|
||||
for (const auto& block : compute_staging_blocks_) {
|
||||
|
||||
Reference in New Issue
Block a user