cui-llama.rn 1.6.0 → 1.7.0

package/cpp/rn-llama.cpp

@@ -1,8 +1,86 @@
 #include "rn-llama.h"
 
+// Include multimodal support
+#include "tools/mtmd/mtmd.h"
+#include "tools/mtmd/clip.h"
+
 namespace rnllama {
 
-const std::vector<lm_ggml_type> kv_cache_types = {
+// Computes FNV-1a hash of the data
+static std::string fnv_hash(const uint8_t * data, size_t len) {
+    const uint64_t fnv_prime = 0x100000001b3ULL;
+    uint64_t hash = 0xcbf29ce484222325ULL;
+
+    for (size_t i = 0; i < len; ++i) {
+        hash ^= data[i];
+        hash *= fnv_prime;
+    }
+    return std::to_string(hash);
+}
+
+static const std::string base64_chars =
+             "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+             "abcdefghijklmnopqrstuvwxyz"
+             "0123456789+/";
+
+// Base64 decoding function
+static std::vector<uint8_t> base64_decode(const std::string &encoded_string) {
+    std::vector<uint8_t> decoded;
+    int in_len = encoded_string.size();
+    int i = 0;
+    int j = 0;
+    int in_ = 0;
+    unsigned char char_array_4[4], char_array_3[3];
+
+    while (in_len-- && (encoded_string[in_] != '=')) {
+        if (isspace(encoded_string[in_])) {
+            in_++;
+            continue;
+        }
+
+        if (encoded_string[in_] == '=' || base64_chars.find(encoded_string[in_]) == std::string::npos) {
+            break;
+        }
+
+        char_array_4[i++] = encoded_string[in_]; in_++;
+        if (i == 4) {
+            for (i = 0; i < 4; i++) {
+                char_array_4[i] = base64_chars.find(char_array_4[i]);
+            }
+
+            char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
+            char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
+            char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
+
+            for (i = 0; i < 3; i++) {
+                decoded.push_back(char_array_3[i]);
+            }
+            i = 0;
+        }
+    }
+
+    if (i) {
+        for (j = i; j < 4; j++) {
+            char_array_4[j] = 0;
+        }
+
+        for (j = 0; j < 4; j++) {
+            char_array_4[j] = base64_chars.find(char_array_4[j]);
+        }
+
+        char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
+        char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
+        char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
+
+        for (j = 0; j < i - 1; j++) {
+            decoded.push_back(char_array_3[j]);
+        }
+    }
+
+    return decoded;
+}
+
+static const std::vector<lm_ggml_type> kv_cache_types = {
     LM_GGML_TYPE_F32,
     LM_GGML_TYPE_F16,
     LM_GGML_TYPE_BF16,
@@ -149,10 +227,16 @@ std::string tokens_to_str(llama_context *ctx, const std::vector<llama_token>::co
     return ret;
 }
 
+struct llama_rn_context_mtmd {
+  mtmd_context *mtmd_ctx = nullptr;
+};
+
 llama_rn_context::~llama_rn_context() {
     if (ctx_sampling != nullptr) {
         common_sampler_free(ctx_sampling);
     }
+
+    releaseMultimodal();
 }
 
 void llama_rn_context::rewind() {
@@ -165,6 +249,7 @@ void llama_rn_context::rewind() {
     generated_text.reserve(params.n_ctx);
     generated_token_probs.clear();
     truncated = false;
+    context_full = false;
     stopped_eos = false;
     stopped_word = false;
     stopped_limit = false;
@@ -197,6 +282,9 @@ bool llama_rn_context::loadModel(common_params &params_)
     templates = common_chat_templates_init(model, params.chat_template);
     n_ctx = llama_n_ctx(ctx);
 
+    // Initialize context shift flag
+    LOG_INFO("ctx_shift: %s", params.ctx_shift ? "enabled" : "disabled");
+
     // We can uncomment for debugging or after this fix: https://github.com/ggerganov/llama.cpp/pull/11101
     // LOG_INFO("%s\n", common_params_get_system_info(params).c_str());
 
@@ -271,11 +359,11 @@ void llama_rn_context::truncatePrompt(std::vector<llama_token> &prompt_tokens) {
 
     new_tokens.insert(new_tokens.end(), prompt_tokens.begin() + params.n_keep + erased_blocks * n_block_size, prompt_tokens.end());
 
-    LOG_VERBOSE("input truncated, n_ctx: %d, n_keep: %d, n_left: %d, new_tokens: %s, num_prompt_tokens: %d",
+    LOG_INFO("input truncated, n_ctx: %d, n_keep: %d, n_left: %d, old_size: %d, new_size: %d",
         n_ctx,
         params.n_keep,
         n_left,
-        tokens_to_str(ctx, new_tokens.cbegin(), new_tokens.cend()).c_str(),
+        prompt_tokens.size(),
         new_tokens.size()
     );
 
@@ -283,65 +371,71 @@ void llama_rn_context::truncatePrompt(std::vector<llama_token> &prompt_tokens) {
     prompt_tokens = new_tokens;
 }
 
-void llama_rn_context::loadPrompt() {
-    std::vector<llama_token> prompt_tokens = ::common_tokenize(ctx, params.prompt, true, true);
-    num_prompt_tokens = prompt_tokens.size();
+void llama_rn_context::loadPrompt(const std::vector<std::string> &media_paths) {
+    bool has_media = !media_paths.empty();
 
-    // LOG tokens
-    std::stringstream ss;
-    ss << "\n" << __func__ << ": prompt_tokens = ";
-    for (auto& token : prompt_tokens) {
-        ss << token << " ";
-    }
-    LOG_INFO("%s\n", ss.str().c_str());
+    if (!has_media) {
+        std::vector<llama_token> text_tokens;
+        // Text-only path
+        text_tokens = ::common_tokenize(ctx, params.prompt, true, true);
+        num_prompt_tokens = text_tokens.size();
 
-    if (params.n_keep < 0)
-    {
-        params.n_keep = (int)num_prompt_tokens;
-    }
-    params.n_keep = std::min(n_ctx - 4, params.n_keep);
+        // LOG tokens
+        std::stringstream ss;
+        ss << "\n" << __func__ << ": prompt_tokens = ";
+        for (auto& token : text_tokens) {
+            ss << token << " ";
+        }
+        LOG_INFO("%s\n", ss.str().c_str());
 
-    // if input prompt is too big, truncate like normal
-    if (num_prompt_tokens >= (size_t) n_ctx)
-    {
-        truncatePrompt(prompt_tokens);
-        num_prompt_tokens = prompt_tokens.size();
+        if (params.n_keep < 0) {
+            params.n_keep = (int)num_prompt_tokens;
+        }
+        params.n_keep = std::min(n_ctx - 4, params.n_keep);
 
-        LM_GGML_ASSERT(num_prompt_tokens < (size_t) n_ctx);
-    }
+        // Handle truncation if needed
+        if (num_prompt_tokens >= (size_t)n_ctx) {
+            if (!params.ctx_shift) {
+                context_full = true;
+                return;
+            }
+            truncatePrompt(text_tokens);
+            num_prompt_tokens = text_tokens.size();
+            LM_GGML_ASSERT(num_prompt_tokens < (size_t)n_ctx);
+        }
 
-    // do context shifitng
-    if(!params.embedding){
-        purge_missing_tokens(ctx, embd, prompt_tokens, params.n_predict, params.n_ctx);
-    }
+        // Update sampling context
+        for (auto & token : text_tokens) {
+            common_sampler_accept(ctx_sampling, token, false);
+        }
 
+        // compare the evaluated prompt with the new prompt
+        n_past = common_part(embd, text_tokens);
 
-    // push the prompt into the sampling context (do not apply grammar)
-    for (auto & token : prompt_tokens)
-    {
-        common_sampler_accept(ctx_sampling, token, false);
-    }
+        embd = text_tokens;
+        if (n_past == num_prompt_tokens) {
+            // we have to evaluate at least 1 token to generate logits.
+            n_past--;
+        }
 
-    // compare the evaluated prompt with the new prompt
-    n_past = common_part(embd, prompt_tokens);
+        // Manage KV cache
+        llama_kv_self_seq_rm(ctx, 0, n_past, -1);
 
-    embd = prompt_tokens;
-    if (n_past == num_prompt_tokens)
-    {
-        // we have to evaluate at least 1 token to generate logits.
-        n_past--;
+        LOG_VERBOSE("prompt ingested, n_past: %d, cached: %s, to_eval: %s",
+            n_past,
+            tokens_to_str(ctx, embd.cbegin(), embd.cbegin() + n_past).c_str(),
+            tokens_to_str(ctx, embd.cbegin() + n_past, embd.cend()).c_str()
+        );
+    } else {
+        // Multimodal path - process all media paths
+        processMedia(params.prompt, media_paths);
+        num_prompt_tokens = embd.size();
     }
 
-    // since #3228 we now have to manually manage the KV cache
-    llama_kv_self_seq_rm(ctx, 0, n_past, -1);
-
-    LOG_VERBOSE("prompt ingested, n_past: %d, cached: %s, to_eval: %s",
-        n_past,
-        tokens_to_str(ctx, embd.cbegin(), embd.cbegin() + n_past).c_str(),
-        tokens_to_str(ctx, embd.cbegin() + n_past, embd.cend()).c_str()
-    );
-
     has_next_token = true;
+
+    LOG_INFO("[DEBUG] Input processed: n_past=%d, embd.size=%zu, num_prompt_tokens=%zu, has_media=%d",
+             n_past, embd.size(), num_prompt_tokens, has_media ? 1 : 0);
 }
 
 void llama_rn_context::beginCompletion() {
@@ -351,6 +445,10 @@ void llama_rn_context::beginCompletion() {
     is_predicting = true;
 }
 
+void llama_rn_context::endCompletion() {
+    is_predicting = false;
+}
+
 completion_token_output llama_rn_context::nextToken()
 {
     completion_token_output result;
@@ -358,6 +456,14 @@ completion_token_output llama_rn_context::nextToken()
 
     if (embd.size() >= (size_t)params.n_ctx)
     {
+        if (!params.ctx_shift) {
+            // If context shifting is disabled, stop generation
+            LOG_WARNING("context full, n_ctx: %d, tokens: %d", params.n_ctx, embd.size());
+            has_next_token = false;
+            context_full = true;
+            return result;
+        }
+
         // Shift context
 
         const int n_left    = n_past - params.n_keep - 1;
@@ -373,12 +479,9 @@ completion_token_output llama_rn_context::nextToken()
         embd.resize(embd.size() - n_discard);
 
         n_past -= n_discard;
+        truncated = true;
 
-        LOG_VERBOSE("input truncated, n_ctx: %d, n_keep: %d, n_left: %d, new_tokens: %s",
-            params.n_ctx,
-            params.n_keep,
-            n_left
-        );
+        LOG_VERBOSE("context shifted, new n_past: %d, new size: %d", n_past, embd.size());
     }
 
     bool tg = true;
@@ -677,7 +780,7 @@ std::string llama_rn_context::bench(int pp, int tg, int pl, int nr)
     }
 
     if (is_interrupted) llama_kv_self_clear(ctx);
-    is_predicting = false;
+    endCompletion();
 
     char model_desc[128];
     llama_model_desc(model, model_desc, sizeof(model_desc));
@@ -712,162 +815,471 @@ void llama_rn_context::removeLoraAdapters() {
 std::vector<common_adapter_lora_info> llama_rn_context::getLoadedLoraAdapters() {
     return this->lora;
 }
-std::vector<int> llama_rn_context::longest_common_subseq(const std::vector<int> x, const std::vector<int> y){
-     int m = x.size(), n = y.size();
-
-     //int LCSuff[m+1][n+1];
-     std::vector<std::vector<int>> LCSuff(m+1, std::vector<int>(n+1));
-
-     for (int j = 0; j <= n; j++)
-         LCSuff[0][j] = 0;
-     for (int i = 0; i <= m; i++)
-         LCSuff[i][0] = 0;
-
-     for (int i = 1; i <= m; i++)
-     {
-         for (int j = 1; j <= n; j++)
-         {
-             if (x[i - 1] == y[j - 1])
-                 LCSuff[i][j] = LCSuff[i - 1][j - 1] + 1;
-             else
-                 LCSuff[i][j] = 0;
-         }
-     }
-
-     std::vector<int> longest;
-     for (int i = 1; i <= m; i++)
-     {
-         for (int j = 1; j <= n; j++)
-         {
-             if (LCSuff[i][j] > longest.size())
-             {
-                 auto off1 = ((i - LCSuff[i][j] + 1) - 1);
-                 auto off2 = off1 + LCSuff[i][j];
-                 longest.clear();
-                //  std::vector<int>().swap(longest);
-                 longest = std::vector<int>(x.begin() + off1, x.begin() + off2);
-                // x.substr((i - LCSuff[i][j] + 1) - 1, LCSuff[i][j]);
-             }
-         }
-     }
-     return longest;
- }
-
-bool llama_rn_context::arr_start_with(const std::vector<int> targetArray, const std::vector<int> searchSeq)
- {
-     int ss = searchSeq.size();
-     if(targetArray.size()<ss)
-     {
-         return false;
-     }
-     for(int i=0;i<ss;++i)
-     {
-         if(targetArray[i]!=searchSeq[i])
-         {
-             return false;
-         }
-     }
-     return true;
- }
-
-int llama_rn_context::arr_find_index_of(const std::vector<int> targetArray, const std::vector<int> searchSeq)
- {
-     int ss = searchSeq.size();
-     int tas = targetArray.size();
-     if(tas<ss)
-     {
-         return -1;
-     }
-     for(int i=0;i<tas;++i)
-     {
-         int srch = 0;
-         bool fail = false;
-         for(int srch=0;srch<ss;++srch)
-         {
-             if ((i + srch) >= tas || targetArray[i + srch] != searchSeq[srch])
-             {
-                 fail = true;
-                 break;
-             }
-         }
-         if(!fail)
-         {
-             return i;
-         }
-     }
-     return -1;
- }
-
-void llama_rn_context::purge_missing_tokens(llama_context * ctx, std::vector<int> &current_context_tokens, std::vector<int> &new_context_tokens, const int genamt, const int nctx)
-{
-    //scan from start old and new ctx, until first mismatch found, save as p0
-    //check remaining old and new ctx for longest common subseq, which needs to be at 256 tokens
-    //test: longest common subseq (LCQ) MUST start within 0 tokens from end of memory, otherwise purge fails
-    //if passed, save beginning of LCQ from old ctx as p1
-    //remove all tokens from old ctx between p0 and p1, updating both arrays and kv, then continue as normal
 
-    const int short_fall_threshold = 200 + (nctx/30); //dont trigger shifting if the distance between trimstart and currhead < this
-    const int stack_allowance = 60 + (nctx/50); //in case the end text is slightly modified, be forgiving
+bool llama_rn_context::initMultimodal(const std::string &mmproj_path, bool use_gpu) {
+    LOG_INFO("[DEBUG] Initializing multimodal with mmproj path: %s", mmproj_path.c_str());
 
-    int trimstart = 0;
-    int new_tokens_len = new_context_tokens.size();
-    bool purge_needed = true;
+    if (model == nullptr) {
+        LOG_ERROR("[DEBUG] Model not loaded, cannot initialize multimodal", "");
+        return false;
+    }
 
-    for (int i = 0; i < current_context_tokens.size(); ++i)
-    {
-        if (current_context_tokens[i] == new_context_tokens[i])
-        {
-            trimstart += 1;
+    LOG_INFO("[DEBUG] Model info: n_ctx=%d, n_embd=%d",
+             llama_n_ctx(ctx),
+             llama_model_n_embd(model));
+
+    // Initialize mtmd context
+    mtmd_context_params mtmd_params = mtmd_context_params_default();
+    mtmd_params.use_gpu = use_gpu;
+    mtmd_params.print_timings = false;
+    mtmd_params.n_threads = params.cpuparams.n_threads;
+    mtmd_params.verbosity = (lm_ggml_log_level)LM_GGML_LOG_LEVEL_INFO;
+
+    LOG_INFO("[DEBUG] Initializing mtmd context with threads=%d", mtmd_params.n_threads);
+
+    auto mtmd_ctx = mtmd_init_from_file(mmproj_path.c_str(), model, mtmd_params);
+    if (mtmd_ctx == nullptr) {
+        LOG_ERROR("[DEBUG] Failed to initialize multimodal context with mmproj: %s", mmproj_path.c_str());
+        return false;
+    }
+    mtmd_wrapper = new llama_rn_context_mtmd();
+    mtmd_wrapper->mtmd_ctx = mtmd_ctx;
+
+    has_multimodal = true;
+
+    // Check if the model uses M-RoPE or non-causal attention
+    bool uses_mrope = mtmd_decode_use_mrope(mtmd_ctx);
+    bool uses_non_causal = mtmd_decode_use_non_causal(mtmd_ctx);
+    LOG_INFO("[DEBUG] Model multimodal properties: uses_mrope=%d, uses_non_causal=%d",
+             uses_mrope ? 1 : 0,
+             uses_non_causal ? 1 : 0);
+
+    // Disable context shifting when multimodal is enabled
+    // This is because an media chunk may contain multiple tokens
+    // and context shifting could break the media representation
+    params.ctx_shift = false;
+
+    // params.n_cache_reuse = 0;
+
+    LOG_INFO("Multimodal context initialized successfully with mmproj: %s", mmproj_path.c_str());
+    LOG_INFO("Context shifting disabled for multimodal support");
+    return true;
+}
+
+struct mtmd_tokenize_result {
+    std::vector<std::string> bitmap_hashes;
+    std::vector<llama_token> tokens;
+    std::vector<size_t> chunk_pos; // both text and media
+    std::vector<size_t> chunk_pos_media; // media only
+    mtmd_input_chunks* chunks = nullptr;
+};
+
+mtmd_tokenize_result tokenizeWithMedia(llama_rn_context_mtmd *mtmd_wrapper, const std::string &prompt, const std::vector<std::string> &media_paths) {
+    mtmd_tokenize_result result;
+    mtmd::bitmaps bitmaps;
+
+    // Load all media paths
+    for (const auto& media_path : media_paths) {
+        LOG_INFO("[DEBUG] Loading media: %s",
+                 media_path.substr(0, 50).c_str()); // Only log part of path for base64
+
+        // Check if it's a base64 media
+        if (media_path.compare(0, 11, "data:image/") == 0 || media_path.compare(0, 11, "data:audio/") == 0) {
+            LOG_INFO("[DEBUG] Detected base64 encoded media");
+
+            // Parse base64 data
+            std::vector<std::string> parts;
+            size_t comma_pos = media_path.find(',');
+            if (comma_pos == std::string::npos) {
+                throw std::runtime_error("Invalid base64 media format, missing comma separator");
+            }
+
+            std::string header = media_path.substr(0, comma_pos);
+            std::string base64_data = media_path.substr(comma_pos + 1);
+
+            if (header.find("base64") == std::string::npos) {
+                bitmaps.entries.clear();
+                throw std::runtime_error("Image must be base64 encoded");
+            }
+
+            // Decode base64
+            std::vector<uint8_t> media_data = base64_decode(base64_data);
+            LOG_INFO("[DEBUG] Base64 decoded, size: %zu bytes", media_data.size());
+
+            // Load bitmap from memory buffer using direct initialization
+            mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_buf(media_data.data(), media_data.size()));
+            if (!bmp.ptr) {
+                bitmaps.entries.clear();
+                throw std::runtime_error("Failed to load base64 media");
+            }
+
+            // Calculate bitmap hash (for KV caching)
+            std::string hash = fnv_hash(bmp.data(), bmp.n_bytes());
+            bmp.set_id(hash.c_str());
+            LOG_INFO("[DEBUG] Bitmap hash: %s", hash.c_str());
+            bitmaps.entries.push_back(std::move(bmp));
+            result.bitmap_hashes.push_back(hash.c_str());
+        } else if (media_path.compare(0, 7, "http://") == 0 || media_path.compare(0, 8, "https://") == 0) {
+            // HTTP URLs are not supported yet
+            LOG_ERROR("[DEBUG] HTTP/HTTPS URLs are not supported yet: %s", media_path.c_str());
+            throw std::runtime_error("HTTP/HTTPS URLs are not supported yet");
+        } else {
+            // Regular file path
+            LOG_INFO("[DEBUG] Loading media from file");
+
+            // Check if file exists
+            FILE* file = fopen(media_path.c_str(), "rb");
+            if (file == nullptr) {
+                bitmaps.entries.clear();
+                throw std::runtime_error("File does not exist or cannot be opened");
+            }
+
+            // Get file size
+            fseek(file, 0, SEEK_END);
+            long file_size = ftell(file);
+            fseek(file, 0, SEEK_SET);
+            LOG_INFO("[DEBUG] File exists and size is %ld bytes", file_size);
+            fclose(file);
+
+            // Create bitmap directly
+            mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_file(media_path.c_str()));
+            if (!bmp.ptr) {
+                bitmaps.entries.clear();
+                throw std::runtime_error("Failed to load media");
+            }
+
+            // Calculate bitmap hash (for KV caching)
+            std::string hash = fnv_hash(bmp.data(), bmp.nx()*bmp.ny()*3);
+            bmp.set_id(hash.c_str());
+            LOG_INFO("[DEBUG] Bitmap hash: %s", hash.c_str());
+            bitmaps.entries.push_back(std::move(bmp));
+            result.bitmap_hashes.push_back(hash.c_str());
         }
-        else
-        {
+    }
+
+    // Create input chunks
+    LOG_INFO("[DEBUG] Initializing input chunks");
+    result.chunks = mtmd_input_chunks_init();
+    if (result.chunks == nullptr) {
+        bitmaps.entries.clear();
+        throw std::runtime_error("Failed to initialize input chunks");
+    }
+
+    mtmd_input_text input_text;
+    input_text.text = prompt.c_str(); // Use the full prompt with image marker
+    input_text.add_special = true;  // Add BOS token if this is the first message
+    input_text.parse_special = true;       // Parse special tokens like <__media__>
+
+    /**
+     * Tokenize the text and media together.
+     *
+     * Example of tokenization for "foo bar <__media__> baz <__media__>":
+     *
+     * 1. Input text with media markers:
+     *
+     *    "foo bar <__media__> baz <__media__>"
+     *
+     * 2. Model-specific markers are added.
+     *
+     * 3. Text is split and tokenized into chunks:
+     *
+     *    ┌─────────────┐  ┌─────────────────────────┐  ┌─────────┐  ┌─────────────────────────┐
+     *    │ TEXT CHUNK  │  │ IMAGE CHUNK             │  │ TEXT    │  │ IMAGE CHUNK             │
+     *    │ "foo bar "  │  │                         │  │ " baz " │  │                         │
+     *    └─────────────┘  └─────────────────────────┘  └─────────┘  └─────────────────────────┘
+     *          │                     │                      │                    │
+     *          ▼                     ▼                      ▼                    ▼
+     *    ┌─────────────┐  ┌─────────────────────────┐  ┌─────────┐  ┌─────────────────────────┐
+     *    │ [1234,5678] │  │ Image Data Structure    │  │ [9012]  │  │ Image Data Structure    │
+     *    └─────────────┘  └─────────────────────────┘  └─────────┘  └─────────────────────────┘
+     *
+     * 4. Image token structure differences:
+     *
+     *    For Qwen2VL (uses M-RoPE with 2D positions):
+     *    ┌─────────────────────────────────────────┐
+     *    │ MEDIA_CHUNK                             │
+     *    │ ┌───────────────────────────────────┐   │
+     *    │ │ mtmd_image_tokens:                │   │
+     *    │ │  nx = 16, ny = 16                 │   │ ← 2D grid (16×16 = 256 tokens)
+     *    │ │  use_mrope_pos = true             │   │ ← Uses M-RoPE positioning
+     *    │ │  batch_f32 = [image_embeddings]   │   │
+     *    │ └───────────────────────────────────┘   │
+     *    └─────────────────────────────────────────┘
+     *
+     *    For other models (uses 1D positions):
+     *    ┌─────────────────────────────────────────┐
+     *    │ MEDIA_CHUNK                             │
+     *    │ ┌───────────────────────────────────┐   │
+     *    │ │ mtmd_image_tokens:                │   │
+     *    │ │  nx = 256, ny = 1                 │   │ ← 1D sequence (256 tokens)
+     *    │ │  use_mrope_pos = false            │   │ ← Uses standard positioning
+     *    │ │  batch_f32 = [image_embeddings]   │   │
+     *    │ └───────────────────────────────────┘   │
+     *    └─────────────────────────────────────────┘
+     *
+     * 5. Final chunks array:
+     *    chunks[0] = TEXT_CHUNK([1234, 5678])
+     *    chunks[1] = MEDIA_CHUNK(first_image)
+     *    chunks[2] = TEXT_CHUNK([9012])
+     *    chunks[3] = MEDIA_CHUNK(second_image)
+     */
+    LOG_INFO("[DEBUG] Tokenizing text and %zu media", bitmaps.entries.size());
+    auto bitmaps_c_ptr = bitmaps.c_ptr();
+    int32_t res = mtmd_tokenize(mtmd_wrapper->mtmd_ctx, result.chunks, &input_text, bitmaps_c_ptr.data(), bitmaps_c_ptr.size());
+    if (res != 0) {
+        mtmd_input_chunks_free(result.chunks);
+        bitmaps.entries.clear();
+        throw std::runtime_error("Failed to tokenize text and media");
+    }
+
+    // Log chunk information
+    size_t num_chunks = mtmd_input_chunks_size(result.chunks);
+    LOG_INFO("[DEBUG] Tokenization successful: num_chunks=%zu", num_chunks);
+
+    // Track the total number of tokens (both text and image)
+    size_t total_token_count = 0;
+
+    /**
+     * Evaluate the chunks.
+     *
+     * For our example "foo bar <__media__> baz <__media__>":
+     *
+     * Token organization in memory:
+     *
+     *    all_tokens: [t0][t1][NULL][NULL]...[NULL][t2][NULL][NULL]...[NULL]
+     *    positions:   0   1    2    3   ...  257   258  259  260 ...  514
+     *    chunk_pos:   0        2                   258  259
+     *
+     *    Where:
+     *    - [t0][t1] are text tokens for "foo bar " (positions 0-1)
+     *    - [NULL]x256 are placeholder tokens for the first image (positions 2-257)
+     *    - [t2] is the text token for " baz " (position 258)
+     *    - [NULL]x256 are placeholder tokens for the second image (positions 259-514)
+     */
+    for (size_t i = 0; i < num_chunks; i++) {
+        result.chunk_pos.push_back(total_token_count);
+
+        const mtmd_input_chunk* chunk = mtmd_input_chunks_get(result.chunks, i);
+        mtmd_input_chunk_type chunk_type = mtmd_input_chunk_get_type(chunk);
+
+        if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
+            size_t n_tokens;
+            const llama_token* tokens = mtmd_input_chunk_get_tokens_text(chunk, &n_tokens);
+            LOG_INFO("[DEBUG] Chunk %zu: type=TEXT, n_tokens=%zu", i, n_tokens);
+
+            // Add text tokens
+            result.tokens.insert(result.tokens.end(), tokens, tokens + n_tokens);
+            total_token_count += n_tokens;
+        } else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE || chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
+            result.chunk_pos_media.push_back(total_token_count);
+
+            size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk);
+            size_t n_pos = mtmd_input_chunk_get_n_pos(chunk);
+            LOG_INFO("[DEBUG] Chunk %zu: type=%s, n_tokens=%zu, n_pos=%zu",
+                     i, chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "IMAGE" : "AUDIO", n_tokens, n_pos);
+
+            for (size_t j = 0; j < n_pos; j++) {
+                result.tokens.push_back(LLAMA_TOKEN_NULL); // Placeholder token
+            }
+            total_token_count += n_pos;
+        }
+    }
+
+    bitmaps.entries.clear();
+
+    return result;
+}
+void llama_rn_context::processMedia(
+    const std::string &prompt,
+    const std::vector<std::string> &media_paths
+) {
+    if (!isMultimodalEnabled()) {
+        throw std::runtime_error("Multimodal is not enabled but image paths are provided");
+    }
+
+    // Multimodal path
+    std::string full_prompt = prompt;
+    auto default_media_marker = mtmd_default_marker();
+    // Add media marker if it doesn't already exist
+    if (full_prompt.find(default_media_marker) == std::string::npos) {
+        full_prompt += " ";
+        full_prompt += default_media_marker;
+    }
+
+    LOG_INFO("[DEBUG] Processing message with role=user, content=%s", full_prompt.c_str());
+    LOG_INFO("[DEBUG] Processing %zu media with prompt: %s", media_paths.size(), prompt.c_str());
+    LOG_INFO("[DEBUG] Current context state: n_past=%d, n_ctx=%d", n_past, n_ctx);
+
+    auto result = tokenizeWithMedia(mtmd_wrapper, full_prompt, media_paths);
+
+    auto all_tokens = result.tokens;
+    auto chunks = result.chunks;
+    auto chunk_pos = result.chunk_pos;
+    auto chunk_pos_media = result.chunk_pos_media;
+    auto bitmap_hashes = result.bitmap_hashes;
+
+    // Check if we have enough context space for all tokens
+    if (all_tokens.size() >= (size_t)n_ctx) {
+        mtmd_input_chunks_free(chunks);
+        context_full = true;
+        throw std::runtime_error("Not enough context space");
+    }
+
+    n_past = common_part(embd, all_tokens);
+
+    llama_pos new_n_past = n_past;
+
+    // Adjust n_past to position of the text chunk
+    // TODO: Edit the text chunk to remove the tokens before n_past to speed up
+    // need to update the mtmd api
+    auto adjusted_n_past = -1;
+    for (size_t i = 0; i < chunk_pos.size(); i++) {
+        if (n_past < chunk_pos[i]) {
             break;
         }
-        if ((i + 2) >= new_tokens_len)
-        {
-            purge_needed = false;
-            break; //no surgery required
+        bool is_end = i + 1 == chunk_pos.size();
+        if (
+            chunk_pos[i] < n_past &&
+            (!is_end && chunk_pos[i + 1] > n_past)
+            // is_end & n_past < total_token_count:
+            // don't need to adjust and it will skip eval_chunk_single, let nextToken() to finish the job
+        ) {
+            adjusted_n_past = chunk_pos[i];
         }
     }
+    if (adjusted_n_past != -1) {
+        n_past = adjusted_n_past;
+        new_n_past = n_past;
+        LOG_INFO("[DEBUG] Adjusted n_past to %d", n_past);
+    }
 
-    
-
-    if(!purge_needed || new_tokens_len < 6 || current_context_tokens.size() < 6 || new_tokens_len - trimstart < short_fall_threshold)
-    {
-        LOG_INFO("Fall Threshold: %d out of %d\n", new_tokens_len - trimstart, short_fall_threshold);
-        return; //no purge is needed
+    // Compare bitmap hashes, if they are not the same, backtrack n_past to the position of the first mismatch
+    if (mtmd_bitmap_past_hashes.size() > 0) {
+        for (size_t i = 0; i < bitmap_hashes.size(); i++) {
+            auto pos = chunk_pos_media[i];
+            if (n_past < pos) {
+                break;
+            }
+            if (i >= mtmd_bitmap_past_hashes.size()) {
+                break;
+            }
+            if (bitmap_hashes[i] != mtmd_bitmap_past_hashes[i]) {
+                LOG_INFO(
+                    "[DEBUG] Bitmap hash mismatch at position %zu, %s != %s",
+                    i, bitmap_hashes[i].c_str(), mtmd_bitmap_past_hashes[i].c_str()
+                );
+                n_past = chunk_pos_media[i];
+                new_n_past = n_past;
+                break;
+            }
+        }
     }
 
-    //at least this many tokens need to match, otherwise don't bother trimming
-    const int lc_tok_threshold = std::max(std::min((new_tokens_len - trimstart) - (genamt+stack_allowance), (int)(nctx*0.45)), short_fall_threshold - stack_allowance);
+    // Clear all KV cache entries after position n_past
+    llama_kv_self_seq_rm(ctx, 0, n_past, -1);
 
-    auto curr_ctx_without_memory = std::vector<int>(current_context_tokens.begin() + trimstart, current_context_tokens.end());
-    auto new_ctx_without_memory = std::vector<int>(new_context_tokens.begin() + trimstart, new_context_tokens.end());
+    LOG_INFO("[DEBUG] Evaluating chunks: n_past=%d, n_batch=%d", n_past, params.n_batch);
 
-    auto shared = longest_common_subseq(curr_ctx_without_memory, new_ctx_without_memory);
+    size_t num_chunks = mtmd_input_chunks_size(chunks);
 
-    if (shared.size() > lc_tok_threshold && arr_start_with(new_ctx_without_memory, shared)) // enough tokens in common
-    {
-        int found = arr_find_index_of(current_context_tokens,shared);
-        if(found>=0 && found > trimstart)
-        {
+    for (size_t i = 0; i < chunk_pos.size(); i++) {
 
-            //extract the unwanted tokens out from context and KV
-            int diff = found - trimstart;
-            llama_kv_self_seq_rm(ctx, 0, trimstart, trimstart + diff);
-            llama_kv_self_seq_add(ctx, 0, trimstart + diff, -1, -diff);
+        LOG_INFO("[DEBUG] Evaluating chunk %zu: n_past=%d, chunk_pos=%zu", i, n_past, chunk_pos[i]);
 
-            for (size_t i = trimstart + diff; i < current_context_tokens.size() - 1; i++)
-            {
-                current_context_tokens[i - diff] = current_context_tokens[i];
+        // Process chunk only if it's after the current n_past
+        if (chunk_pos[i] >= n_past) {
+            bool chunk_logits_last = (i == num_chunks - 1);
+            auto chunk = mtmd_input_chunks_get(chunks, i);
+
+            int32_t res = mtmd_helper_eval_chunk_single(
+                mtmd_wrapper->mtmd_ctx,
+                ctx,
+                chunk,
+                n_past,
+                0,
+                params.n_batch,
+                chunk_logits_last,
+                &new_n_past
+            );
+            if (res != 0) {
+                mtmd_input_chunks_free(chunks);
+                throw std::runtime_error("Failed to evaluate chunks");
             }
+            n_past = new_n_past;
+        }
+    }
 
-            LOG_INFO("\n[Context Shifting: Erased %d tokens at position %d]", diff, trimstart + 1);
+    if (n_past == all_tokens.size() && n_past > 0 && all_tokens[n_past - 1] != LLAMA_TOKEN_NULL) {
+        // we have to evaluate at least 1 token to generate logits.
+        n_past--;
+    }
+
+    // Update embd with all tokens (both text and media)
+    embd = all_tokens;
+
+    mtmd_bitmap_past_hashes = bitmap_hashes;
+
+    // Update sampling context with text tokens only
+    for (auto & token : all_tokens) {
+        if (token == LLAMA_TOKEN_NULL) {
+            continue;
+        }
+        common_sampler_accept(ctx_sampling, token, false);
+    }
 
-            current_context_tokens.resize(current_context_tokens.size() - diff);
+    // Clean up media resources
+    LOG_INFO("[DEBUG] Cleaning up resources");
+    mtmd_input_chunks_free(chunks);
+}
+
+llama_rn_tokenize_result llama_rn_context::tokenize(const std::string &text, const std::vector<std::string> &media_paths) {
+    if (media_paths.size() > 0) {
+        if (!isMultimodalEnabled()) {
+            throw std::runtime_error("Multimodal is not enabled but media paths are provided");
         }
+        auto result = tokenizeWithMedia(mtmd_wrapper, text, media_paths);
+        mtmd_input_chunks_free(result.chunks);
+        llama_rn_tokenize_result tokenize_result = {
+            .tokens = result.tokens,
+            .has_media = true,
+            .bitmap_hashes = result.bitmap_hashes,
+            .chunk_pos = result.chunk_pos,
+            .chunk_pos_media = result.chunk_pos_media,
+        };
+        return tokenize_result;
     }
+    std::vector<llama_token> text_tokens;
+    text_tokens = common_tokenize(ctx, text, false);
+    llama_rn_tokenize_result tokenize_result = {
+        .tokens = text_tokens,
+        .has_media = false,
+        .bitmap_hashes = {},
+        .chunk_pos = {},
+        .chunk_pos_media = {},
+    };
+    return tokenize_result;
+}
+
+bool llama_rn_context::isMultimodalEnabled() const {
+    return has_multimodal && mtmd_wrapper != nullptr;
+}
 
+bool llama_rn_context::isMultimodalSupportVision() const {
+    return isMultimodalEnabled() && mtmd_support_vision(mtmd_wrapper->mtmd_ctx);
+}
+
+bool llama_rn_context::isMultimodalSupportAudio() const {
+    return isMultimodalEnabled() && mtmd_support_audio(mtmd_wrapper->mtmd_ctx);
+}
+
+void llama_rn_context::releaseMultimodal() {
+    if (mtmd_wrapper && mtmd_wrapper->mtmd_ctx != nullptr) {
+        mtmd_free(mtmd_wrapper->mtmd_ctx);
+        mtmd_wrapper->mtmd_ctx = nullptr;
+        delete mtmd_wrapper;
+        mtmd_wrapper = nullptr;
+        has_multimodal = false;
+    }
 }
 
 }