mtmd: add chunk save/load function (#26645)

* mtmd: add chunk save/load function

* nits

* add tests

* rn _MAX --> _COUNT

tests/test-mtmd-c-api.c

@@ -1,4 +1,6 @@
 #include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
 #include <assert.h>
 
 #include "mtmd.h"
@@ -62,6 +64,72 @@ int main(void) {
         }
     }
 
+    // test chunk save/load round-trip
+    for (size_t i = 0; i < n_chunks; i++) {
+        const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
+        assert(chunk != NULL);
+        enum mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk);
+
+        // query the required buffer size (out_buf == NULL)
+        size_t expected_len = 0;
+        int32_t rc = mtmd_input_chunk_save(chunk, NULL, 0, &expected_len);
+        printf("    Chunk %zu: save query rc = %d, expected_len = %zu\n", i, rc, expected_len);
+        assert(rc == 0);
+        assert(expected_len > 0);
+
+        // saving into a too-small buffer must fail, not crash
+        char tiny_buf[1];
+        rc = mtmd_input_chunk_save(chunk, tiny_buf, sizeof(tiny_buf), NULL);
+        printf("    Chunk %zu: save into too-small buffer rc = %d (expect non-zero)\n", i, rc);
+        assert(rc != 0);
+
+        // save into a properly-sized buffer
+        char * buf = (char *) malloc(expected_len);
+        assert(buf != NULL);
+        rc = mtmd_input_chunk_save(chunk, buf, expected_len, NULL);
+        assert(rc == 0);
+
+        // loading from a truncated buffer must fail gracefully, not crash
+        if (expected_len > 1) {
+            mtmd_input_chunk * bad = mtmd_input_chunk_load(buf, expected_len - 1);
+            printf("    Chunk %zu: load from truncated buffer = %p (expect NULL)\n", i, (void *) bad);
+            assert(bad == NULL);
+        }
+
+        // load it back
+        mtmd_input_chunk * loaded = mtmd_input_chunk_load(buf, expected_len);
+        assert(loaded != NULL);
+
+        // metadata must match the original chunk
+        assert(mtmd_input_chunk_get_type(loaded) == type);
+        assert(mtmd_input_chunk_get_n_tokens(loaded) == mtmd_input_chunk_get_n_tokens(chunk));
+        assert(mtmd_input_chunk_get_n_pos(loaded) == mtmd_input_chunk_get_n_pos(chunk));
+
+        if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
+            size_t n_tok_orig, n_tok_loaded;
+            const llama_token * tok_orig   = mtmd_input_chunk_get_tokens_text(chunk, &n_tok_orig);
+            const llama_token * tok_loaded = mtmd_input_chunk_get_tokens_text(loaded, &n_tok_loaded);
+            printf("    Chunk %zu: loaded %zu text tokens (orig %zu), first token %d (orig %d)\n",
+                i, n_tok_loaded, n_tok_orig,
+                n_tok_loaded > 0 ? tok_loaded[0] : -1,
+                n_tok_orig   > 0 ? tok_orig[0]   : -1);
+            assert(n_tok_orig == n_tok_loaded);
+            for (size_t j = 0; j < n_tok_orig; j++) {
+                assert(tok_orig[j] == tok_loaded[j]);
+            }
+        } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE || type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
+            const char * id_orig   = mtmd_input_chunk_get_id(chunk);
+            const char * id_loaded = mtmd_input_chunk_get_id(loaded);
+            printf("    Chunk %zu: loaded id '%s' (orig '%s')\n", i, id_loaded, id_orig);
+            assert(id_orig != NULL && id_loaded != NULL);
+            assert(strcmp(id_orig, id_loaded) == 0);
+        }
+
+        mtmd_input_chunk_free(loaded);
+        free(buf);
+    }
+    printf("Chunk save/load round-trip OK\n");
+
     // Free the chunks
     mtmd_input_chunks_free(chunks);
 

tools/mtmd/clip-impl.h

@@ -591,6 +591,8 @@ struct clip_image_u8 {
     }
 };
 
+struct mtmd_serialization; // forward declaration
+
 // For images, buf.size() == nx*ny*3
 //     Memory layout: RGBRGBRGB...
 // For seq, buf.size() == nx*ny*3*nt
@@ -671,6 +673,9 @@ struct clip_image_f32 {
         return buf.empty();
     }
 
+    void serialize(struct mtmd_serialization & ser) const;
+    void deserialize(struct mtmd_serialization & ser);
+
   private:
     std::vector<float> buf;
     int nx_ = 0;
@@ -752,6 +757,9 @@ struct clip_image_f32_batch {
         }
         return new_batch;
     }
+
+    void serialize(struct mtmd_serialization & ser) const;
+    void deserialize(struct mtmd_serialization & ser);
 };
 
 //

tools/mtmd/mtmd.cpp

@@ -22,8 +22,123 @@
 #include <cstdlib>
 #include <cstring>
 #include <climits>
+#include <type_traits>
 #include <vector>
 
+// remember to bump this if the serialization format changes
+#define MTMD_SERIALIZATION_VERSION 1
+
+struct mtmd_serialization {
+    // note: using 64-bit here for future-proofing
+    uint64_t version = MTMD_SERIALIZATION_VERSION;
+    std::vector<char> data;
+    size_t read_pos = 0; // cursor used when reading
+
+    // for writing
+    mtmd_serialization(uint64_t version) : version(version) {
+        write(version);
+    }
+
+    // for reading
+    mtmd_serialization(uint64_t version, const char * buf, size_t len) {
+        // copy buf to data
+        data.assign(buf, buf + len);
+        uint64_t ver_in = read<uint64_t>();
+        if (ver_in != version) {
+            throw std::runtime_error("version mismatch");
+        }
+        this->version = ver_in;
+    }
+
+    template <typename T>
+    void write(T value) {
+        static_assert(std::is_trivially_copyable<T>::value && !std::is_same<T, bool>::value,
+            "T must be trivially copyable and not bool");
+        const char * p = reinterpret_cast<const char *>(&value);
+        data.insert(data.end(), p, p + sizeof(T));
+    }
+
+    template <typename T>
+    T read() {
+        static_assert(std::is_trivially_copyable<T>::value && !std::is_same<T, bool>::value,
+            "T must be trivially copyable and not bool");
+        if (read_pos + sizeof(T) > data.size()) {
+            throw std::runtime_error("read OOB");
+        }
+        T value;
+        std::memcpy(&value, data.data() + read_pos, sizeof(T));
+        read_pos += sizeof(T);
+        return value;
+    }
+
+};
+
+template <>
+void mtmd_serialization::write<bool>(bool value) {
+    write<uint8_t>(value ? 1 : 0);
+}
+template <>
+bool mtmd_serialization::read<bool>() {
+    return read<uint8_t>() != 0;
+}
+
+template <>
+void mtmd_serialization::write<std::string>(std::string value) {
+    write<uint64_t>(value.size());
+    data.insert(data.end(), value.begin(), value.end());
+}
+template <>
+std::string mtmd_serialization::read<std::string>() {
+    uint64_t len = read<uint64_t>();
+    if (read_pos + len > data.size()) {
+        throw std::runtime_error("read_string OOB");
+    }
+    std::string str(data.data() + read_pos, len);
+    read_pos += len;
+    return str;
+}
+
+// only mtmd.cpp needs these, so they're implemented here rather than in clip-impl.h
+void clip_image_f32::serialize(mtmd_serialization & ser) const {
+    // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
+    // note: buf is intentionally NOT serialized; the loaded clip_image_f32 will always be a placeholder
+    ser.write(add_viewsep);
+    ser.write(add_newline);
+    ser.write((int32_t)nx_);
+    ser.write((int32_t)ny_);
+}
+void clip_image_f32::deserialize(mtmd_serialization & ser) {
+    add_viewsep = ser.read<bool>();
+    add_newline = ser.read<bool>();
+    nx_ = ser.read<int32_t>();
+    ny_ = ser.read<int32_t>();
+    buf.clear(); // always a placeholder after loading
+}
+
+void clip_image_f32_batch::serialize(mtmd_serialization & ser) const {
+    // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
+    ser.write(is_audio);
+    ser.write<uint64_t>(entries.size());
+    for (const auto & entry : entries) {
+        entry.serialize(ser);
+    }
+}
+void clip_image_f32_batch::deserialize(mtmd_serialization & ser) {
+    is_audio = ser.read<bool>();
+    uint64_t n = ser.read<uint64_t>();
+    constexpr size_t min_entry_bytes = sizeof(uint8_t) * 2 + sizeof(int32_t) * 2;
+    if (n > (ser.data.size() - ser.read_pos) / min_entry_bytes) {
+        throw std::runtime_error("entries count exceeds buffer size");
+    }
+    entries.clear();
+    entries.reserve(n);
+    for (uint64_t i = 0; i < n; i++) {
+        clip_image_f32 entry;
+        entry.deserialize(ser);
+        entries.push_back(std::move(entry));
+    }
+}
+
 // for still image data, layout is RGBRGBRGB...
 // length of data must be nx * ny * 3 bytes
 //
@@ -83,6 +198,7 @@ enum mtmd_pos_type {
     MTMD_POS_TYPE_NORMAL,    // number of positions equals to number of tokens
     MTMD_POS_TYPE_MROPE,     // qwen-vl mrope style, each image takes max(t,h,w) position indexes
     MTMD_POS_TYPE_HUNYUANVL, // HunyuanVL mrope + BOI/EOI/newline layout with XD-RoPE dim-3
+    MTMD_POS_TYPE_COUNT,     // for validation
 };
 
 struct mtmd_image_tokens {
@@ -136,6 +252,30 @@ struct mtmd_image_tokens {
             id
         };
     }
+
+    void serialize(mtmd_serialization & ser) const {
+        // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
+        ser.write(nx);
+        ser.write(ny);
+        ser.write((uint32_t)pos);
+        ser.write(image_idx);
+        ser.write(n_temporal_merge);
+        ser.write(id);
+        batch_f32.serialize(ser);
+    }
+    void deserialize(mtmd_serialization & ser) {
+        nx = ser.read<uint32_t>();
+        ny = ser.read<uint32_t>();
+        uint32_t pos_raw = ser.read<uint32_t>();
+        if (pos_raw >= MTMD_POS_TYPE_COUNT) {
+            throw std::runtime_error("invalid pos type");
+        }
+        pos = (mtmd_pos_type)pos_raw;
+        image_idx = ser.read<uint32_t>();
+        n_temporal_merge = ser.read<uint32_t>();
+        id = ser.read<std::string>();
+        batch_f32.deserialize(ser);
+    }
 };
 using mtmd_image_tokens_ptr = std::unique_ptr<mtmd_image_tokens>;
 
@@ -161,6 +301,18 @@ struct mtmd_audio_tokens {
             id
         };
     }
+
+    void serialize(mtmd_serialization & ser) const {
+        // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
+        ser.write(n_tokens);
+        ser.write(id);
+        batch_f32.serialize(ser);
+    }
+    void deserialize(mtmd_serialization & ser) {
+        n_tokens = ser.read<uint32_t>();
+        id = ser.read<std::string>();
+        batch_f32.deserialize(ser);
+    }
 };
 using mtmd_audio_tokens_ptr = std::unique_ptr<mtmd_audio_tokens>;
 
@@ -192,6 +344,66 @@ struct mtmd_input_chunk {
         }
         return false;
     }
+
+    void serialize(mtmd_serialization & ser) const {
+        // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
+        ser.write((uint32_t)type);
+
+        ser.write<uint64_t>(tokens_text.size());
+        for (llama_token tok : tokens_text) {
+            ser.write((int32_t)tok);
+        }
+
+        ser.write(tokens_image != nullptr);
+        if (tokens_image) {
+            tokens_image->serialize(ser);
+        }
+
+        ser.write(tokens_audio != nullptr);
+        if (tokens_audio) {
+            tokens_audio->serialize(ser);
+        }
+    }
+    void deserialize(mtmd_serialization & ser) {
+        uint32_t type_raw = ser.read<uint32_t>();
+        if (type_raw >= MTMD_INPUT_CHUNK_TYPE_COUNT) {
+            throw std::runtime_error("invalid chunk type");
+        }
+        type = (mtmd_input_chunk_type)type_raw;
+
+        uint64_t n_tokens_text = ser.read<uint64_t>();
+        // reject before resize() so a tiny corrupted/malicious buffer can't force a huge allocation
+        if (n_tokens_text > (ser.data.size() - ser.read_pos) / sizeof(int32_t)) {
+            throw std::runtime_error("tokens_text length exceeds buffer size");
+        }
+        tokens_text.resize(n_tokens_text);
+        for (uint64_t i = 0; i < n_tokens_text; i++) {
+            tokens_text[i] = (llama_token)ser.read<int32_t>();
+        }
+
+        if (ser.read<bool>()) {
+            tokens_image = std::make_unique<mtmd_image_tokens>();
+            tokens_image->deserialize(ser);
+        } else {
+            tokens_image.reset();
+        }
+
+        if (ser.read<bool>()) {
+            tokens_audio = std::make_unique<mtmd_audio_tokens>();
+            tokens_audio->deserialize(ser);
+        } else {
+            tokens_audio.reset();
+        }
+
+        // catch buffers where the declared type doesn't match which payload is actually present,
+        // so a mismatched chunk can't slip through and null-deref/abort later in an accessor
+        if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE && !tokens_image) {
+            throw std::runtime_error("type is IMAGE but tokens_image is missing");
+        }
+        if (type == MTMD_INPUT_CHUNK_TYPE_AUDIO && !tokens_audio) {
+            throw std::runtime_error("type is AUDIO but tokens_audio is missing");
+        }
+    }
 };
 
 struct mtmd_input_chunks {
@@ -2043,6 +2255,42 @@ void mtmd_input_chunk_free(mtmd_input_chunk * chunk) {
     }
 }
 
+int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len) {
+    try {
+        mtmd_serialization ser(MTMD_SERIALIZATION_VERSION);
+        chunk->serialize(ser);
+
+        if (expected_out_len) {
+            *expected_out_len = ser.data.size();
+        }
+        if (!out_buf) {
+            // caller is only querying the required size
+            return 0;
+        }
+        if (out_len < ser.data.size()) {
+            LOG_ERR("%s: out_buf is too small, need %zu bytes, got %zu\n", __func__, ser.data.size(), out_len);
+            return -1;
+        }
+        std::memcpy(out_buf, ser.data.data(), ser.data.size());
+        return 0;
+    } catch (const std::exception & e) {
+        LOG_ERR("%s: %s\n", __func__, e.what());
+        return -1;
+    }
+}
+
+mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len) {
+    try {
+        mtmd_serialization ser(MTMD_SERIALIZATION_VERSION, buf, len);
+        mtmd::input_chunk_ptr chunk(new mtmd_input_chunk());
+        chunk->deserialize(ser);
+        return chunk.release();
+    } catch (const std::exception & e) {
+        LOG_ERR("%s: %s\n", __func__, e.what());
+        return nullptr;
+    }
+}
+
 // mtmd_image_tokens
 
 size_t mtmd_image_tokens_get_n_tokens(const mtmd_image_tokens * image_tokens) {

tools/mtmd/mtmd.h

@@ -55,6 +55,7 @@ enum mtmd_input_chunk_type {
     MTMD_INPUT_CHUNK_TYPE_TEXT,
     MTMD_INPUT_CHUNK_TYPE_IMAGE,
     MTMD_INPUT_CHUNK_TYPE_AUDIO,
+    MTMD_INPUT_CHUNK_TYPE_COUNT, // for validation
 };
 
 // opaque types
@@ -232,6 +233,15 @@ MTMD_API llama_pos                  mtmd_input_chunk_get_n_pos       (const mtmd
 MTMD_API mtmd_input_chunk * mtmd_input_chunk_copy(const mtmd_input_chunk * chunk);
 MTMD_API void               mtmd_input_chunk_free(mtmd_input_chunk * chunk);
 
+// save/load an input chunk to/from a buffer (useful for KV save/load)
+// important: only chunk's metadata will be saved, the actual image/audio data will not be saved
+// the loaded chunk will always be a placeholder, cannot be used for mtmd_encode() or mtmd_batch_encode()
+// out_buf can be nullptr (to query expected_out_len)
+// returns 0 on success, non-zero on failure
+MTMD_API int32_t            mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len);
+// returns nullptr on failure
+MTMD_API mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len);
+
 
 // mtmd_image_tokens
 //