mtmd: add mtmd_bitmap_set_mergeable (#27348)

tests/test-mtmd-impl.cpp

@@ -1,9 +1,12 @@
 #include "testing.h"
 
 #include "mtmd-image.h"
+#include "mtmd-internal.h"
 
 #include <iostream>
+#include <stdexcept>
 #include <string>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -67,6 +70,73 @@ MAKE_TEST(test_image_preprocessor_lfm2) {
     }
 }
 
+//
+// mtmd temporal merge
+//
+
+MAKE_TEST(test_temporal_merge_grouping) {
+    std::vector<mtmd::bitmap_ptr> pool; // keeps the bitmaps alive until the end of the test
+
+    // spec chars:
+    //   v = video frame, w = video frame of another size, a = audio, i = plain image, t = text
+    auto make_parts = [&pool](const std::string & spec) {
+        std::vector<mtmd_input_part> parts;
+        for (char c : spec) {
+            if (c == 't') {
+                parts.push_back({ "hello", nullptr });
+                continue;
+            }
+            mtmd_bitmap * bm = nullptr;
+            switch (c) {
+                case 'v': bm = mtmd_bitmap_init(100, 100, nullptr);   break;
+                case 'w': bm = mtmd_bitmap_init(200, 200, nullptr);   break;
+                case 'a': bm = mtmd_bitmap_init_from_audio(100, nullptr); break;
+                case 'i': bm = mtmd_bitmap_init(100, 100, nullptr);   break;
+                default: throw std::runtime_error(std::string("unknown spec char: ") + c);
+            }
+            mtmd_bitmap_set_mergeable(bm, c != 'i');
+            pool.emplace_back(bm);
+            parts.push_back({ "", bm });
+        }
+        return parts;
+    };
+
+    // { parts, n_merge, expected size of each group }
+    const std::vector<std::tuple<std::string, int, std::string>> cases = {
+        { "vv",   2, "2"    },
+        { "vvv",  2, "21"   },
+        { "vvvv", 2, "22"   },
+        { "vvi",  2, "21"   },
+        { "tvvt", 2, "2"    },
+        { "vtv",  2, "11"   }, // text in between breaks the merge
+        { "vw",   2, "11"   }, // different sizes cannot be merged
+        { "aa",   2, "11"   }, // audio is never merged
+        { "ii",   2, "11"   }, // two unrelated images must stay separated
+        { "iv",   2, "11"   },
+        { "vi",   2, "11"   },
+        { "vv",   1, "11"   }, // model without temporal merge
+    };
+
+    for (const auto & [spec, n_merge, expected] : cases) {
+        auto parts  = make_parts(spec);
+        auto groups = mtmd_group_mergeable_bitmaps(parts, n_merge);
+
+        std::string actual;
+        for (const auto & group : groups) {
+            actual += std::to_string(group.size());
+        }
+
+        const std::string name = "\"" + spec + "\" with n_merge=" + std::to_string(n_merge);
+        t.assert_equal("groups for " + name, expected, actual);
+
+        size_t n_bitmap_parts = 0;
+        for (const auto & p : parts) {
+            n_bitmap_parts += p.bitmap != nullptr ? 1 : 0;
+        }
+        t.assert_equal("remaining bitmap parts for " + name, groups.size(), n_bitmap_parts);
+    }
+}
+
 //
 // main
 //

tools/mtmd/CMakeLists.txt

@@ -17,6 +17,7 @@ add_library(mtmd
             mtmd-audio.cpp
             mtmd-image.cpp
             mtmd.h
+            mtmd-internal.h
             mtmd-helper.cpp
             mtmd-helper-gen.cpp
             mtmd-helper-common.h

tools/mtmd/README-dev.md

@@ -21,7 +21,7 @@ A typical pipeline of the core libmtmd is as follows:
 - A bitmap (RGB image or PCM audio) is created
 - Bitmap and the text prompt is provided to `mtmd_tokenize()` that breaks the input into chunks
     - The tokenizer function first expands a "lazy" bitmap if it finds one. Typically, this is used by video, so that one media token corresponds to one input bitmap
-    - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch
+    - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch. Only bitmaps marked by `mtmd_bitmap_set_mergeable()` are merged
     - The preprocessor will then be called, which produces a list of chunks
     - Depending on the model itself, special tokens will be injected to separate image chunks (i.e. llava-uhd-style models)
 - Multiple bitmaps may be batched together to form a larger `mtmd_batch()`

tools/mtmd/mtmd-helper.cpp

@@ -727,7 +727,9 @@ struct mtmd_helper_video {
 
         LOG_DBG("%s: frame %d read OK\n", __func__, current_frame);
         current_frame++;
-        return mtmd_bitmap_init(info.width, info.height, frame_buf.data());
+        mtmd_bitmap * frame = mtmd_bitmap_init(info.width, info.height, frame_buf.data());
+        mtmd_bitmap_set_mergeable(frame, true);
+        return frame;
     }
 
     int32_t read_next(mtmd_bitmap ** out_bitmap, char ** out_text) {

tools/mtmd/mtmd-internal.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "mtmd.h"
+
+#include <string>
+#include <vector>
+
+// !!! Internal header, to be used by mtmd and its unit tests only !!!
+
+#define MTMD_INTERNAL_HEADER
+
+// bitmap is null for text parts
+struct mtmd_input_part {
+    std::string text;
+    const mtmd_bitmap * bitmap;
+};
+
+// [QWEN_VIDEO] merged parts are erased from `parts`, so one group always maps to one part
+std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_input_part> & parts, int n_merge);

tools/mtmd/mtmd.cpp

@@ -1,6 +1,7 @@
 #include "clip.h"
 #include "clip-impl.h"
 #include "mtmd.h"
+#include "mtmd-internal.h"
 #include "mtmd-audio.h"
 #include "mtmd-image.h"
 #include "debug/mtmd-debug.h"
@@ -149,6 +150,7 @@ struct mtmd_bitmap {
     uint32_t ny = 0;
     std::string id; // optional user-defined id, for ex: can be set to image hash, useful for KV cache tracking
     bool is_audio = false; // true if the bitmap is audio
+    bool mergeable = false; // [QWEN_VIDEO] set only on frames of the same video
 
     // lazy-loaded bitmap
     mtmd_bitmap_lazy_callback lazy_callback = nullptr;
@@ -186,7 +188,9 @@ struct mtmd_bitmap {
 
     bool can_merge_with(const mtmd_bitmap & other) const {
         // [QWEN_VIDEO] can (temporal) merge if both are images with same size
-        return !is_audio && !other.is_audio && nx == other.nx && ny == other.ny;
+        return mergeable && other.mergeable
+            && !is_audio && !other.is_audio
+            && nx == other.nx && ny == other.ny;
     }
 
   private:
@@ -1076,6 +1080,25 @@ void mtmd_free(mtmd_context * ctx) {
     delete ctx;
 }
 
+std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_input_part> & parts, int n_merge) {
+    std::vector<std::vector<const mtmd_bitmap *>> output;
+    for (size_t i = 0; i < parts.size(); i++) {
+        if (parts[i].bitmap == nullptr) {
+            continue; // text part
+        }
+        const bool has_next = n_merge > 1 && i + 1 < parts.size() && parts[i + 1].bitmap != nullptr;
+        if (has_next && parts[i].bitmap->can_merge_with(*parts[i + 1].bitmap)) {
+            LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1);
+            output.push_back({parts[i].bitmap, parts[i + 1].bitmap});
+            parts.erase(parts.begin() + i + 1);
+            continue;
+        }
+        LOG_DBG("%s: no merging for part index %zu\n", __func__, i);
+        output.push_back({parts[i].bitmap});
+    }
+    return output;
+}
+
 struct mtmd_tokenizer {
     mtmd_context * ctx;
 
@@ -1084,10 +1107,7 @@ struct mtmd_tokenizer {
     bool parse_special;
     const llama_vocab * vocab;
 
-    struct part {
-        std::string text;
-        const mtmd_bitmap * bitmap;
-    };
+    using part = mtmd_input_part;
     std::vector<part> parts;
     // these will be freed when mtmd_tokenizer finishes
     std::vector<mtmd::bitmap> bm_from_lazy; // TODO @ngxson : refactor, free bm_from_lazy progressively
@@ -1192,34 +1212,7 @@ struct mtmd_tokenizer {
             GGML_ASSERT(n_merge_frames <= 2 && "we only support merging maximum 2 images for now; open an issue if this model supports merging more");
         }
 
-        // Build merged_bitmaps: each entry is a group of 1 or 2 bitmaps.
-        // For consecutive mergeable bitmap parts, merge them and collapse the second part out of this->parts.
-        std::vector<std::vector<const mtmd_bitmap *>> merged_bitmaps;
-        if (n_merge_frames > 1) {
-            for (size_t i = 0; i < parts.size(); ++i) {
-                if (parts[i].bitmap == nullptr) {
-                    continue;
-                }
-                if (i + 1 < parts.size() && parts[i + 1].bitmap != nullptr) {
-                    const mtmd_bitmap * bm_a = parts[i].bitmap;
-                    const mtmd_bitmap * bm_b = parts[i + 1].bitmap;
-                    if (bm_a->can_merge_with(*bm_b)) {
-                        LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1);
-                        merged_bitmaps.push_back({bm_a, bm_b});
-                        parts.erase(parts.begin() + i + 1); // collapse the second bitmap part
-                        continue;
-                    }
-                }
-                LOG_DBG("%s: no merging for part index %zu\n", __func__, i);
-                merged_bitmaps.push_back({parts[i].bitmap});
-            }
-        } else {
-            for (const auto & p : parts) {
-                if (p.bitmap != nullptr) {
-                    merged_bitmaps.push_back({p.bitmap});
-                }
-            }
-        }
+        auto merged_bitmaps = mtmd_group_mergeable_bitmaps(parts, n_merge_frames);
 
         size_t i_bm = 0;
         for (const auto & p : parts) {
@@ -2200,6 +2193,10 @@ void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id) {
     }
 }
 
+void mtmd_bitmap_set_mergeable(mtmd_bitmap * bitmap, bool mergeable) {
+    bitmap->mergeable = mergeable;
+}
+
 mtmd_bitmap * mtmd_bitmap_init_lazy(mtmd_context * ctx,
                                     const char * id,
                                     void * user_data,

tools/mtmd/mtmd.h

@@ -154,7 +154,8 @@ MTMD_API const char * mtmd_get_marker(const mtmd_context * ctx);
 //     length of data must be nx * ny * 3
 //     the data is in RGBRGBRGB... format
 //     note: some video-capable models (i.e. qwen-vl) can merge consecutive bitmaps
-//           into one chunk, mtmd_tokenize() will automatically handle this
+//           into one chunk; mtmd_tokenize() handles this, but remember to set
+//           mtmd_bitmap_set_mergeable(true) for every frame
 // if bitmap is audio:
 //     length of data must be n_samples * sizeof(float)
 //     the data is in float format (PCM F32)
@@ -175,6 +176,8 @@ MTMD_API void                  mtmd_bitmap_free       (mtmd_bitmap * bitmap);
 // these getters/setters are dedicated functions, so you can for example calculate the hash of the image based on mtmd_bitmap_get_data()
 MTMD_API const char * mtmd_bitmap_get_id(const mtmd_bitmap * bitmap);
 MTMD_API void         mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id);
+// if true, this bitmap can be merged (temporal merge) with an adjacent mergeable bitmap by certain video input models
+MTMD_API void         mtmd_bitmap_set_mergeable(mtmd_bitmap * bitmap, bool mergeable);
 
 // mtmd_bitmap lazy
 //