Commit fd21ecf for stable-diffusion.cpp
commit fd21ecfceef2665d2c5f2db8536e3202d864fd30
Author: leejet <leejet714@gmail.com>
Date: Fri Oct 23 00:45:33 2026 +0800
refactor: use one placeholder token per image in LLM inputs (#2191)
diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp
index bea6b69..7839d5c 100644
--- a/examples/server/routes_sdcpp.cpp
+++ b/examples/server/routes_sdcpp.cpp
@@ -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) {
diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp
index 26e11cd..a146716 100644
--- a/src/conditioning/conditioner.hpp
+++ b/src/conditioning/conditioner.hpp
@@ -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.
diff --git a/src/model/diffusion/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp
index 2befb72..ca1559b 100644
--- a/src/model/diffusion/hidream_o1.hpp
+++ b/src/model/diffusion/hidream_o1.hpp
@@ -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;
}
diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp
index 8ea9180..85f7511 100644
--- a/src/model/te/llm.hpp
+++ b/src/model/te/llm.hpp
@@ -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'";
diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp
index dee8770..2a46fc9 100644
--- a/src/pipeline/request.cpp
+++ b/src/pipeline/request.cpp
@@ -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();