| 1 | #include "models.h" |
| 2 | |
| 3 | ggml_cgraph * clip_graph_llama4::build() { |
| 4 | GGML_ASSERT(model.class_embedding != nullptr); |
| 5 | GGML_ASSERT(model.position_embeddings != nullptr); |
| 6 | |
| 7 | const int n_pos = n_patches + 1; // +1 for [CLS] |
| 8 | |
| 9 | // 2D input positions |
| 10 | ggml_tensor * pos_h = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); |
| 11 | ggml_set_name(pos_h, "pos_h"); |
| 12 | ggml_set_input(pos_h); |
| 13 | |
| 14 | ggml_tensor * pos_w = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); |
| 15 | ggml_set_name(pos_w, "pos_w"); |
| 16 | ggml_set_input(pos_w); |
| 17 | |
| 18 | ggml_tensor * inp = build_inp_raw(); |
| 19 | |
| 20 | // Llama4UnfoldConvolution |
| 21 | { |
| 22 | ggml_tensor * kernel = ggml_reshape_4d(ctx0, model.patch_embeddings_0, |
| 23 | patch_size, patch_size, 3, n_embd); |
| 24 | inp = ggml_im2col(ctx0, kernel, inp, patch_size, patch_size, 0, 0, 1, 1, true, inp->type); |
| 25 | inp = ggml_mul_mat(ctx0, model.patch_embeddings_0, inp); |
| 26 | inp = ggml_reshape_2d(ctx0, inp, n_embd, n_patches); |
| 27 | cb(inp, "patch_conv", -1); |
| 28 | } |
| 29 | |
| 30 | // add CLS token |
| 31 | inp = ggml_concat(ctx0, inp, model.class_embedding, 1); |
| 32 | |
| 33 | // build ViT with 2D position embeddings |
| 34 | auto add_pos = [&](ggml_tensor * cur, const clip_layer &) { |
| 35 | // first half is X axis and second half is Y axis |
| 36 | // ref: https://github.com/huggingface/transformers/blob/40a493c7ed4f19f08eadb0639cf26d49bfa5e180/src/transformers/models/llama4/modeling_llama4.py#L1312 |
| 37 | // ref: https://github.com/Blaizzy/mlx-vlm/blob/a57156aa87b33cca6e5ee6cfc14dd4ef8f611be6/mlx_vlm/models/llama4/vision.py#L441 |
| 38 | return build_rope_2d(ctx0, cur, pos_w, pos_h, hparams.rope_theta, false); |
| 39 | }; |
| 40 | ggml_tensor * cur = build_vit( |
| 41 | inp, n_pos, |
| 42 | NORM_TYPE_NORMAL, |
| 43 | hparams.ffn_op, |
| 44 | model.position_embeddings, |
| 45 | add_pos); |
| 46 | |
| 47 | // remove CLS token |
| 48 | cur = ggml_view_2d(ctx0, cur, |
| 49 | n_embd, n_patches, |
| 50 | ggml_row_size(cur->type, n_embd), 0); |
| 51 | |
| 52 | // pixel shuffle |
| 53 | // based on Llama4VisionPixelShuffleMLP |
| 54 | // https://github.com/huggingface/transformers/blob/2932f318a20d9e54cc7aea052e040164d85de7d6/src/transformers/models/llama4/modeling_llama4.py#L1151 |
| 55 | { |
| 56 | const int scale_factor = model.hparams.n_merge; |
| 57 | const int bsz = 1; // batch size, always 1 for now since we don't support batching |
| 58 | GGML_ASSERT(scale_factor > 0); |
| 59 | GGML_ASSERT(n_patches_x == n_patches_y); // llama4 only supports square images |
| 60 | cur = ggml_reshape_4d(ctx0, cur, |
nothing calls this directly
no test coverage detected