diff --git a/tests/test-mtmd-impl.cpp b/tests/test-mtmd-impl.cpp index 41cb53237..df18b0a42 100644 --- a/tests/test-mtmd-impl.cpp +++ b/tests/test-mtmd-impl.cpp @@ -1,9 +1,12 @@ #include "testing.h" #include "mtmd-image.h" +#include "mtmd-internal.h" #include +#include #include +#include #include #include @@ -67,6 +70,73 @@ MAKE_TEST(test_image_preprocessor_lfm2) { } } +// +// mtmd temporal merge +// + +MAKE_TEST(test_temporal_merge_grouping) { + std::vector 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 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> 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 // diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index aa97c768a..63a8b17a5 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/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 diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index ac43e1b81..e14906823 100644 --- a/tools/mtmd/README-dev.md +++ b/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()` diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index cc966b93c..8b8e5f5f6 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/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) { diff --git a/tools/mtmd/mtmd-internal.h b/tools/mtmd/mtmd-internal.h new file mode 100644 index 000000000..067fa88b9 --- /dev/null +++ b/tools/mtmd/mtmd-internal.h @@ -0,0 +1,19 @@ +#pragma once + +#include "mtmd.h" + +#include +#include + +// !!! 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> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 4063d28e0..0f5cb8c7a 100644 --- a/tools/mtmd/mtmd.cpp +++ b/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> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge) { + std::vector> 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 parts; // these will be freed when mtmd_tokenizer finishes std::vector 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> 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, diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 78587f3fe..ef4f99c0b 100644 --- a/tools/mtmd/mtmd.h +++ b/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 //