diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index bea6b69c..7839d5c5 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 26e11cd3..a146716d 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -2148,10 +2148,11 @@ struct LLMEmbedder : public Conditioner { std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& 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>>& image_embeds = {}) { std::vector> 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(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 suffix_tokens; if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) { @@ -2216,6 +2230,17 @@ struct LLMEmbedder : public Conditioner { std::vector 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(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& image_grids = {}, const std::string& prompt_suffix = "", sd::Tensor* 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 position_ids; if (sd_version_is_flux3(version)) { - position_ids = sd::Tensor::zeros({static_cast(tokens.size() * 4)}); + position_ids = sd::Tensor::zeros({static_cast(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(i); + for (size_t i = 0; i < weights.size(); ++i) { + position_ids[axis * weights.size() + i] = mask[i] == 0.f ? 1 : static_cast(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> outputs, int grid_h, @@ -2391,6 +2416,9 @@ struct LLMEmbedder : public Conditioner { return false; } int index = static_cast(prefix_tokens.size()); + for (const auto& image_embed : image_embeds) { + index += static_cast(image_embed.second.shape()[1]) - 1; + } int count = static_cast(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(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(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(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(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 prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(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(prompt.size()); @@ -2833,14 +2857,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(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(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(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(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(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 prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(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 padded(tokens.begin(), tokens.end()); - padded.resize(static_cast(n_text + num_queries), tokenizer->PAD_TOKEN_ID); - int64_t n_total = static_cast(padded.size()); - sd::Tensor input_ids({n_total}, padded); + std::vector input_tokens(tokens.begin(), tokens.end()); + input_tokens.push_back(tokenizer->PAD_TOKEN_ID); + sd::Tensor 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 2befb723..ca1559bc 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(count) * token.size()); - for (int64_t i = 0; i < count; ++i) { - out += token; - } - return out; - } - static inline std::pair calculate_dimensions(int max_size, double ratio) { int width = static_cast(std::sqrt(max_size * max_size * ratio)); int height = static_cast(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 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(dims.first / PATCH_SIZE) * static_cast(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(prefix_tokens.size()); + for (const auto& grid : image_grids) { + prompt_start += static_cast(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(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 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(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(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(height / PATCH_SIZE), static_cast(width / PATCH_SIZE)}); @@ -656,8 +656,7 @@ namespace HiDreamO1 { } std::vector token_types(input_ids_pad.size(), 0); - int txt_seq_len = static_cast(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(token_types.size()); ++i) { token_types[i] = 1; } diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 8ea91802..85f75116 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>> 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 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 prefix_tokens; + if (!tokenizer->encode(text + image_prefix, prefix_tokens, nullptr)) { + return; + } + std::vector>> image_embeds; + image_embeds.emplace_back(static_cast(prefix_tokens.size()), image_embed); text += img_prompt; prompt_attn_range.first = static_cast(text.size()); text += "change 'flux.cpp' to 'edit.cpp'"; diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index dee87708..2a46fc9c 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();