1#include "models.h"
 2
 3ggml_cgraph * clip_graph_pixtral::build() {
 4    const int n_merge = hparams.n_merge;
 5
 6    // 2D input positions
 7    ggml_tensor * pos_h = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_patches);
 8    ggml_set_name(pos_h, "pos_h");
 9    ggml_set_input(pos_h);
10
11    ggml_tensor * pos_w = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_patches);
12    ggml_set_name(pos_w, "pos_w");
13    ggml_set_input(pos_w);
14
15    auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
16        return build_rope_2d(ctx0, cur, pos_h, pos_w, hparams.rope_theta, true);
17    };
18
19    ggml_tensor * inp = build_inp();
20    ggml_tensor * cur = build_vit(
21                            inp, n_patches,
22                            NORM_TYPE_RMS,
23                            hparams.ffn_op,
24                            nullptr, // no learned pos embd
25                            add_pos);
26
27    // mistral small 3.1 patch merger
28    // ref: https://github.com/huggingface/transformers/blob/7a3e208892c06a5e278144eaf38c8599a42f53e7/src/transformers/models/mistral3/modeling_mistral3.py#L67
29    if (model.mm_patch_merger_w) {
30        GGML_ASSERT(hparams.n_merge > 0);
31
32        cur = ggml_mul(ctx0, ggml_rms_norm(ctx0, cur, eps), model.mm_input_norm_w);
33
34        // reshape image tokens to 2D grid
35        cur = ggml_reshape_3d(ctx0, cur, n_embd, n_patches_x, n_patches_y);
36        cur = ggml_permute(ctx0, cur, 2, 0, 1, 3); // [x, y, n_embd]
37        cur = ggml_cont(ctx0, cur);
38
39        // torch.nn.functional.unfold is just an im2col under the hood
40        // we just need a dummy kernel to make it work
41        ggml_tensor * kernel = ggml_view_3d(ctx0, cur, n_merge, n_merge, cur->ne[2], 0, 0, 0);
42        cur = ggml_im2col(ctx0, kernel, cur, n_merge, n_merge, 0, 0, 1, 1, true, inp->type);
43
44        // project to n_embd
45        cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], cur->ne[1] * cur->ne[2]);
46        cur = ggml_mul_mat(ctx0, model.mm_patch_merger_w, cur);
47    }
48
49    // LlavaMultiModalProjector (always using GELU activation)
50    {
51        cur = build_ffn(cur,
52            model.mm_1_w, model.mm_1_b,
53            nullptr, nullptr,
54            model.mm_2_w, model.mm_2_b,
55            FFN_GELU,
56            -1);
57    }
58
59    // arrangement of the [IMG_BREAK] token
60    if (model.token_embd_img_break) {
61        // not efficient, but works
62        // the trick is to view the embeddings as a 3D tensor with shape [n_embd, n_patches_per_row, n_rows]
63        // and then concatenate the [IMG_BREAK] token to the end of each row, aka n_patches_per_row dimension
64        // after the concatenation, we have a tensor with shape [n_embd, n_patches_per_row + 1, n_rows]
65
66        const int p_y             = n_merge > 0 ? n_patches_y / n_merge : n_patches_y;
67        const int p_x             = n_merge > 0 ? n_patches_x / n_merge : n_patches_x;
68        const int p_total         = p_x * p_y;
69        const int n_embd_text     = cur->ne[0];
70        const int n_tokens_output = p_total + p_y - 1; // one [IMG_BREAK] per row, except the last row
71
72        ggml_tensor * tmp = ggml_reshape_3d(ctx0, cur, n_embd_text, p_x, p_y);
73        ggml_tensor * tok = ggml_new_tensor_3d(ctx0, tmp->type, n_embd_text, 1, p_y);
74        tok = ggml_scale(ctx0, tok, 0.0); // clear the tensor
75        tok = ggml_add(ctx0, tok, model.token_embd_img_break);
76        tmp = ggml_concat(ctx0, tmp, tok, 1);
77        cur = ggml_view_2d(ctx0, tmp,
78            n_embd_text, n_tokens_output,
79            ggml_row_size(tmp->type, n_embd_text), 0);
80    }
81
82    // build the graph
83    ggml_build_forward_expand(gf, cur);
84
85    return gf;
86}