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

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;
             }

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'";