refactor: use one placeholder token per image in LLM inputs (#2191)

This commit is contained in:
leejet
2026-10-23 00:45:33 +08:00
committed by GitHub
parent e779d4f149
commit fd21ecfcee
5 changed files with 114 additions and 87 deletions
+2 -2
View File
@@ -380,8 +380,8 @@ static json make_capabilities_json(ServerRuntime& runtime) {
json top_level_defaults = json::object();
json top_level_output_formats = json::array();
json top_level_features = {
{"cancel_queued", true},
{"cancel_generating", true},
{"cancel_queued", true},
{"cancel_generating", true},
};
std::string current_mode = "";
if (supports_img) {
+76 -53
View File
@@ -2148,10 +2148,11 @@ struct LLMEmbedder : public Conditioner {
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
const std::pair<int, int>& attn_range,
size_t min_length = 0,
size_t max_length = 100000000,
bool spell_quotes = false,
const std::string& suffix = "") {
size_t min_length = 0,
size_t max_length = 100000000,
bool spell_quotes = false,
const std::string& suffix = "",
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds = {}) {
std::vector<std::pair<std::string, float>> parsed_attention;
if (attn_range.first >= 0 && attn_range.second > 0) {
if (attn_range.first > 0) {
@@ -2196,6 +2197,19 @@ struct LLMEmbedder : public Conditioner {
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
for (const auto& image_embed : image_embeds) {
const int index = image_embed.first;
const int64_t count = image_embed.second.shape()[1];
if (index < 0 || static_cast<size_t>(index) >= tokens.size() || count < 1) {
LOG_ERROR("Invalid reference image token range");
return {};
}
const int token = tokens[index];
const float weight = weights[index];
tokens.insert(tokens.begin() + index, count - 1, token);
weights.insert(weights.begin() + index, count - 1, weight);
}
if (!suffix.empty()) {
std::vector<int> suffix_tokens;
if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) {
@@ -2216,6 +2230,17 @@ struct LLMEmbedder : public Conditioner {
std::vector<float> mask;
tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
// Keep masks and weights in expanded coordinates, but only look up one token per image.
for (auto it = image_embeds.rbegin(); it != image_embeds.rend(); ++it) {
const int64_t begin = it->first;
const int64_t end = begin + it->second.shape()[1];
if (begin < 0 || end > static_cast<int64_t>(tokens.size())) {
LOG_ERROR("Reference image exceeds the tokenized context");
return {};
}
tokens.erase(tokens.begin() + begin + 1, tokens.begin() + end);
}
// for (int i = 0; i < tokens.size(); i++) {
// std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl;
// }
@@ -2238,7 +2263,7 @@ struct LLMEmbedder : public Conditioner {
const std::vector<LLM::ImageGrid>& image_grids = {},
const std::string& prompt_suffix = "",
sd::Tensor<float>* output_mask = nullptr) {
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix);
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix, image_embeds);
auto& tokens = std::get<0>(tokens_weights_mask);
auto& weights = std::get<1>(tokens_weights_mask);
auto& mask = std::get<2>(tokens_weights_mask);
@@ -2268,10 +2293,10 @@ struct LLMEmbedder : public Conditioner {
sd::Tensor<int32_t> position_ids;
if (sd_version_is_flux3(version)) {
position_ids = sd::Tensor<int32_t>::zeros({static_cast<int64_t>(tokens.size() * 4)});
position_ids = sd::Tensor<int32_t>::zeros({static_cast<int64_t>(weights.size() * 4)});
for (size_t axis = 0; axis < 4; ++axis) {
for (size_t i = 0; i < tokens.size(); ++i) {
position_ids[axis * tokens.size() + i] = mask[i] == 0.f ? 1 : static_cast<int32_t>(i);
for (size_t i = 0; i < weights.size(); ++i) {
position_ids[axis * weights.size() + i] = mask[i] == 0.f ? 1 : static_cast<int32_t>(i);
}
}
}
@@ -2368,7 +2393,7 @@ struct LLMEmbedder : public Conditioner {
}
}
bool expand_vision_marker(std::string& prompt,
bool expand_vision_marker(const std::string& prompt,
size_t& marker_start,
std::vector<sd::Tensor<float>> outputs,
int grid_h,
@@ -2391,6 +2416,9 @@ struct LLMEmbedder : public Conditioner {
return false;
}
int index = static_cast<int>(prefix_tokens.size());
for (const auto& image_embed : image_embeds) {
index += static_cast<int>(image_embed.second.shape()[1]) - 1;
}
int count = static_cast<int>(outputs[0].shape()[1]);
if (index > max_length || count > max_length - index) {
LOG_ERROR("Reference images exceed the %d-token context", max_length);
@@ -2405,13 +2433,7 @@ struct LLMEmbedder : public Conditioner {
deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer]));
}
image_grids.push_back({index, count, grid_h, grid_w});
std::string markers;
markers.reserve(static_cast<size_t>(count) * placeholder.size());
for (int i = 0; i < count; ++i) {
markers += placeholder;
}
prompt.replace(position, placeholder.size(), markers);
marker_start = position + markers.size();
marker_start = position + placeholder.size();
return true;
}
@@ -2638,14 +2660,13 @@ struct LLMEmbedder : public Conditioner {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
for (const auto& previous_image : image_embeds) {
image_embed_idx += static_cast<int>(previous_image.second.shape()[1]) - 1;
}
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += placeholder;
img_prompt += "<|vision_end|>";
}
prompt += img_prompt;
@@ -2707,7 +2728,6 @@ struct LLMEmbedder : public Conditioner {
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
prompt_template_encode_start_idx = 64;
int image_embed_idx = 64 + 6;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
@@ -2724,6 +2744,7 @@ struct LLMEmbedder : public Conditioner {
}
}
prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
@@ -2743,19 +2764,22 @@ struct LLMEmbedder : public Conditioner {
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
for (const auto& previous_image : image_embeds) {
image_embed_idx += static_cast<int>(previous_image.second.shape()[1]) - 1;
}
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(num_image_tokens * placeholder.size());
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += placeholder;
img_prompt += "<|vision_end|>";
}
prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
prompt += img_prompt;
prompt_attn_range.first = static_cast<int>(prompt.size());
@@ -2833,14 +2857,13 @@ struct LLMEmbedder : public Conditioner {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
for (const auto& previous_image : image_embeds) {
image_embed_idx += static_cast<int>(previous_image.second.shape()[1]) - 1;
}
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += placeholder;
img_prompt += "<|vision_end|>";
}
@@ -2905,14 +2928,13 @@ struct LLMEmbedder : public Conditioner {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
for (const auto& previous_image : image_embeds) {
image_embed_idx += static_cast<int>(previous_image.second.shape()[1]) - 1;
}
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += placeholder;
img_prompt += "<|vision_end|>";
}
prompt += img_prompt;
@@ -2930,7 +2952,6 @@ struct LLMEmbedder : public Conditioner {
LOG_INFO("LongCatEditPipeline");
prompt_template_encode_start_idx = 67;
min_length = 512 + prompt_template_encode_start_idx;
int image_embed_idx = 36 + 6;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
@@ -2947,6 +2968,7 @@ struct LLMEmbedder : public Conditioner {
}
}
prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n";
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
@@ -2965,19 +2987,22 @@ struct LLMEmbedder : public Conditioner {
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
for (const auto& previous_image : image_embeds) {
image_embed_idx += static_cast<int>(previous_image.second.shape()[1]) - 1;
}
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(num_image_tokens * placeholder.size());
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += placeholder;
img_prompt += "<|vision_end|>";
}
prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n";
prompt += img_prompt;
} else {
prompt_template_encode_start_idx = 36;
@@ -3583,12 +3608,10 @@ struct LLaDAImageEmbedder : public Conditioner {
auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids);
auto query_embeds = query_former->compute(n_threads, inputs_embeds);
// splice_image_embeds() replaces tokens in place, so the query slots have to exist in
// input_ids; their ids are irrelevant because the embeddings are overwritten.
std::vector<int32_t> padded(tokens.begin(), tokens.end());
padded.resize(static_cast<size_t>(n_text + num_queries), tokenizer->PAD_TOKEN_ID);
int64_t n_total = static_cast<int64_t>(padded.size());
sd::Tensor<int32_t> input_ids({n_total}, padded);
std::vector<int32_t> input_tokens(tokens.begin(), tokens.end());
input_tokens.push_back(tokenizer->PAD_TOKEN_ID);
sd::Tensor<int32_t> input_ids({n_text + 1}, std::move(input_tokens));
int64_t n_total = n_text + num_queries;
// Bidirectional everywhere except that the text tokens must not see the appended
// queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min.
+14 -15
View File
@@ -56,15 +56,6 @@ namespace HiDreamO1 {
}
};
static inline std::string repeat_special_token(const std::string& token, int64_t count) {
std::string out;
out.reserve(static_cast<size_t>(count) * token.size());
for (int64_t i = 0; i < count; ++i) {
out += token;
}
return out;
}
static inline std::pair<int, int> calculate_dimensions(int max_size, double ratio) {
int width = static_cast<int>(std::sqrt(max_size * max_size * ratio));
int height = static_cast<int>(width / ratio);
@@ -405,7 +396,7 @@ namespace HiDreamO1 {
txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds);
auto t_emb = t_embedder1->forward(&runner_ctx, timestep);
int64_t txt_seq_len = input_ids->ne[0];
int64_t txt_seq_len = txt->ne[1];
if (txt_seq_len > 1) {
auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1);
txt = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1);
@@ -616,8 +607,7 @@ namespace HiDreamO1 {
sd::Tensor<float> vlm_image = sd::ops::interpolate(
resized_ref,
{dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]});
vlm_image = vlm_image * 2.0f - 1.0f;
int64_t image_tokens = static_cast<int64_t>(dims.first / PATCH_SIZE) * static_cast<int64_t>(dims.second / PATCH_SIZE);
vlm_image = vlm_image * 2.0f - 1.0f;
auto patch_img = resized_ref * 2.0f - 1.0f;
result.c_ref_images.push_back(std::move(patch_img));
@@ -626,8 +616,11 @@ namespace HiDreamO1 {
return {};
}
int64_t prompt_start = static_cast<int64_t>(prefix_tokens.size());
for (const auto& grid : image_grids) {
prompt_start += static_cast<int64_t>(grid[1]) * grid[2] - 1;
}
prompt += "<|vision_start|>";
prompt += repeat_special_token("<|image_pad|>", image_tokens);
prompt += "<|image_pad|>";
prompt += "<|vision_end|>";
vlm_images.emplace_back(static_cast<int>(prompt_start), std::move(vlm_image));
image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE});
@@ -642,6 +635,13 @@ namespace HiDreamO1 {
}
std::vector<int32_t> input_ids_pad = input_ids;
for (size_t i = 0; i < vlm_images.size(); ++i) {
const auto& grid = image_grids[i];
const int64_t count = static_cast<int64_t>(grid[1]) * grid[2];
input_ids_pad.insert(input_ids_pad.begin() + vlm_images[i].first, count - 1, IMAGE_TOKEN_ID);
}
int txt_seq_len = static_cast<int>(input_ids_pad.size());
input_ids_pad.push_back(VISION_START_TOKEN_ID);
input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
@@ -656,8 +656,7 @@ namespace HiDreamO1 {
}
std::vector<int32_t> token_types(input_ids_pad.size(), 0);
int txt_seq_len = static_cast<int>(input_ids.size());
int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM;
int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM;
for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
token_types[i] = 1;
}
+19 -14
View File
@@ -924,6 +924,8 @@ namespace LLM {
int64_t txt_token_start = 0;
int64_t txt_token_end = 0;
ggml_tensor* input_embed = nullptr;
// Image and DeepStack indices remain in the expanded sequence coordinates.
int64_t removed_tokens = 0;
for (int i = 0; i < image_embeds.size(); i++) {
if (i == 0) {
@@ -931,7 +933,10 @@ namespace LLM {
} else {
txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1];
}
txt_token_end = image_embeds[i].first;
txt_token_start -= removed_tokens;
txt_token_end = image_embeds[i].first - removed_tokens;
removed_tokens += image_embeds[i].second->ne[1] - 1;
GGML_ASSERT(txt_token_start >= 0 && txt_token_end >= txt_token_start && txt_token_end < raw_x->ne[1]);
// An embed can sit flush against the previous one or at the very start/end of the
// sequence, leaving no text tokens to splice around it.
@@ -951,14 +956,14 @@ namespace LLM {
}
}
txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1] - removed_tokens;
txt_token_end = raw_x->ne[1];
if (txt_token_end > txt_token_start) {
auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
}
GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
GGML_ASSERT(raw_x->ne[1] + removed_tokens == input_embed->ne[1]);
return input_embed;
}
@@ -2506,6 +2511,9 @@ namespace LLM {
}
int64_t n_tokens = input_ids->ne[0];
for (const auto& image_embed : image_embeds) {
n_tokens += image_embed.second->ne[1] - 1;
}
if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
config.arch == LLMArch::MINISTRAL_3_3B ||
config.arch == LLMArch::QWEN3 ||
@@ -3014,20 +3022,17 @@ namespace LLM {
LOG_VERBOSE("llm encode_image test done in %lldms", t1 - t0);
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt = "Picture 1: <|vision_start|>"; // [24669, 220, 16, 25, 220, 151652]
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(num_image_tokens * placeholder.size());
for (int i = 0; i < num_image_tokens; i++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
image_embeds.emplace_back(64, image_embed);
const std::string image_prefix = "Picture 1: <|vision_start|>";
const std::string img_prompt = image_prefix + "<|image_pad|><|vision_end|>";
std::pair<int, int> prompt_attn_range;
std::string text = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(text + image_prefix, prefix_tokens, nullptr)) {
return;
}
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
image_embeds.emplace_back(static_cast<int>(prefix_tokens.size()), image_embed);
text += img_prompt;
prompt_attn_range.first = static_cast<int>(text.size());
text += "change 'flux.cpp' to 'edit.cpp'";
+3 -3
View File
@@ -321,9 +321,9 @@ namespace sd::pipeline {
void GenerationRequest::resolve(StableDiffusionGGML* sd) {
if (sd_version_is_minit2i(sd->version) && (width != 512 || height != 512)) {
width = 512;
height = 512;
LOG_WARN("align to fixed size 512x512 for MiniT2I");
width = 512;
height = 512;
LOG_WARN("align to fixed size 512x512 for MiniT2I");
}
align_generation_request_size();
resolve_hires();