mtmd: Add DeepSeekOCR 2 Support (#20975)

* mtmd: DeepSeek-OCR 2 support, with multi-tile dynamic resolution

* introduced clip_image_f32::add_viewsep

* address PR review

- drop redundant ggml_cpy ops in both deepseekocr versions build
- drop no-op ggml_cont in build_sam
- assert num_image_tokens deepseekocr2
- view_seperator as (1, n_embd) at conversion (for both versions)
- drop redundant ggml_reshape_2d

* Update tools/mtmd/models/deepseekocr2.cpp

Co-authored-by: Xuan-Son Nguyen <thichthat@gmail.com>

---------

Co-authored-by: Xuan-Son Nguyen <thichthat@gmail.com>

tools/mtmd/models/deepseekocr.cpp

@@ -157,7 +157,6 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) {
 
             cur = ggml_mul_mat(ctx0, layer.qkv_w, cur);
             cur = ggml_add(ctx0, cur, layer.qkv_b);
-            cur = ggml_cont(ctx0, cur); // Ensure tensor is contiguous before reshape
             cur = ggml_reshape_4d(ctx0, cur, n_embd, 3, W * H, B);
 
             ggml_tensor * Q;
@@ -251,17 +250,17 @@ ggml_cgraph * clip_graph_deepseekocr::build() {
     ggml_tensor * inp_raw = build_inp_raw();
     ggml_tensor * sam_out = build_sam(inp_raw);
 
+    const int clip_n_patches = sam_out->ne[0] * sam_out->ne[1];
+
     ggml_tensor * clip_out;
     // Building DS-OCR CLIP
     {
         ggml_tensor * inp;
 
-        inp = ggml_cpy(ctx0, sam_out, ggml_dup_tensor(ctx0, sam_out));
-        inp = ggml_reshape_2d(ctx0, inp, inp->ne[0] * inp->ne[1], inp->ne[2]);
+        inp = ggml_reshape_2d(ctx0, sam_out, clip_n_patches, sam_out->ne[2]);
         inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3));
 
-        ggml_tensor * new_pos_embd =
-            ggml_cpy(ctx0, model.position_embeddings, ggml_dup_tensor(ctx0, model.position_embeddings));
+        ggml_tensor * new_pos_embd = model.position_embeddings;
 
         int        n_pos    = new_pos_embd->ne[1];  // +1 for [CLS]
         const auto tgt_size = static_cast<int>(std::sqrt(inp->ne[1]));
@@ -295,16 +294,12 @@ ggml_cgraph * clip_graph_deepseekocr::build() {
         clip_out = cur;
     }
 
-    const int clip_n_patches = sam_out->ne[0] * sam_out->ne[1];
-
     sam_out  = ggml_cont(ctx0, ggml_permute(ctx0, sam_out, 1, 2, 0, 3));
     sam_out  = ggml_reshape_2d(ctx0, sam_out, sam_out->ne[0], clip_n_patches);
     clip_out = ggml_view_2d(ctx0, clip_out, n_embd, clip_n_patches, clip_out->nb[1], clip_out->nb[1]);
 
     ggml_tensor * cur;
     cur = ggml_concat(ctx0, clip_out, sam_out, 0);
-    cur = ggml_reshape_2d(ctx0, cur, 2 * n_embd, clip_n_patches);
-    cur = ggml_cont(ctx0, cur);
     cur = ggml_mul_mat(ctx0, model.mm_fc_w, cur);
     cur = ggml_add(ctx0, cur, model.mm_fc_b);
 
@@ -313,13 +308,11 @@ ggml_cgraph * clip_graph_deepseekocr::build() {
     const auto n_dim = cur->ne[0];
 
     ggml_tensor * imgnl;
-    ggml_tensor * vs;
 
     imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, h, 1);
-    vs    = ggml_reshape_2d(ctx0, model.view_seperator, n_dim, 1);  // (n_dim, 1)
     cur   = ggml_reshape_3d(ctx0, cur, n_dim, w, h);
     cur   = ggml_reshape_2d(ctx0, ggml_concat(ctx0, cur, imgnl, 1), n_dim, (w + 1) * h);
-    cur   = ggml_concat(ctx0, cur, vs, 1);  // (n_dim, h*(w+1) + 1)
+    cur   = ggml_concat(ctx0, cur, model.view_seperator, 1);  // (n_dim, h*(w+1) + 1)
 
     cb(cur, "dsocr_output", -1);
 

tools/mtmd/models/deepseekocr2.cpp

@@ -0,0 +1,81 @@
+#include "models.h"
+
+ggml_cgraph * clip_graph_deepseekocr2::build() {
+    GGML_ASSERT(hparams.n_head_kv > 0);
+    GGML_ASSERT(n_head % hparams.n_head_kv == 0);
+
+    // patch embedding
+    ggml_tensor * inp_raw = build_inp_raw();
+
+    ggml_tensor * sam_out = build_sam(inp_raw);
+
+    ggml_tensor * qwen2_out;
+    // Building Qwen2 encoder
+    {
+        ggml_tensor * inp;
+
+        inp = ggml_reshape_2d(ctx0, sam_out, sam_out->ne[0] * sam_out->ne[1], sam_out->ne[2]); // H*W, C
+        inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3));
+
+        auto num_image_tokens = inp->ne[1]; // H*W
+        GGML_ASSERT(num_image_tokens == 144 || num_image_tokens == 256);
+
+        // query based on numbers of image tokens (in SAM output)
+        // 16x16 -> query_1024 (1024x1024 images)
+        // 12x12 -> query_768 (768x768 images)
+
+        ggml_tensor * query_embed = model.resample_query_1024;
+        int           num_queries = 256;
+
+        if (num_image_tokens == 144) {
+            query_embed = model.resample_query_768;
+            num_queries = 144;
+        }
+
+        // (B, num_image_tokens + num_queries, C)
+        inp = ggml_concat(ctx0, inp, ggml_cast(ctx0, query_embed, inp->type), 1);
+
+        auto seq_len = inp->ne[1];
+
+        // qwen2 encoder attention mask
+        ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, seq_len, seq_len);
+        ggml_set_name(attn_mask, "qwen2_attn_mask");
+        ggml_set_input(attn_mask);
+
+        ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_arange(ctx0, 0, seq_len, 1), GGML_TYPE_I32);
+
+        auto add_rope = [&](ggml_tensor * x, const clip_layer &) {
+            return ggml_rope_ext(ctx0, x, inp_pos, nullptr, d_head,
+                                 GGML_ROPE_TYPE_NEOX, 131072, 1000000, 1, 0, 1, 0, 0);
+        };
+
+        build_vit_opts vit_opts;
+        vit_opts.attn_mask = attn_mask;
+
+        // build_vit applies model.post_ln_w internally; do not re-apply
+        ggml_tensor * cur = build_vit(inp, seq_len, NORM_TYPE_RMS, FFN_SILU,
+                                      /* learned_pos_embd */ nullptr, add_rope, vit_opts);
+
+        cur = ggml_cont(ctx0,
+                        ggml_view_2d(ctx0, cur, cur->ne[0], num_queries, cur->nb[1],
+                                     cur->nb[1] * (cur->ne[1] - num_queries))); // only take query tokens for output
+
+        ggml_build_forward_expand(gf, cur);
+        qwen2_out = cur;
+    }
+
+    ggml_tensor * cur;
+
+    cur = ggml_mul_mat(ctx0, model.mm_fc_w, qwen2_out);
+    cur = ggml_add(ctx0, cur, model.mm_fc_b);
+
+    // view_seperator only after the global view
+    if (img.add_viewsep) {
+        cur = ggml_concat(ctx0, cur, model.view_seperator, 1); // (n_dim, 257)
+    }
+
+    cb(cur, "dsocr2_output", -1);
+
+    ggml_build_forward_expand(gf, cur);
+    return gf;
+}

tools/mtmd/models/models.h

@@ -121,6 +121,11 @@ struct clip_graph_deepseekocr : clip_graph {
     ggml_tensor * build_sam(ggml_tensor * inp); // build the SAM model
 };
 
+struct clip_graph_deepseekocr2 : clip_graph_deepseekocr {
+    clip_graph_deepseekocr2(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph_deepseekocr(ctx, img) {}
+    ggml_cgraph * build() override; // reuses build_sam() from base
+};
+
 struct clip_graph_conformer : clip_graph {
     clip_graph_conformer(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
     ggml_cgraph * build() override;