Function sam3_build_sam_dec_graph

sam3.cpp:10548–10759

};

static sam3_dec_result sam3_build_sam_dec_graph(
    struct ggml_context* ctx,
    const sam3_model& model,
    struct ggml_tensor* image_feats,  // [D, H, H, 1]
    struct ggml_tensor* image_pe,     // [D, H, H, 1]
    struct ggml_tensor* sparse_emb,   // [D, N_pts, 1]
    struct ggml_tensor* dense_emb,    // [D, H, H, 1]
    struct ggml_tensor* feat_s0,      // [D, H*4, H*4, 1] high-res
    struct ggml_tensor* feat_s1,     // [D, H*2, H*2, 1] mid-res
    int eff_feat_size = 0)
{
    const auto& dec = model.sam_dec;
    const auto& hp = model.hparams;
    const int D = hp.sam_embed_dim;  // 256
    const int H = (eff_feat_size > 0) ? eff_feat_size : hp.feat_size();
    const int N_pts = (int)sparse_emb->ne[1];
    const int n_heads = 8;                               // SAM uses 8 heads
    const int num_mask_tokens = hp.sam_n_multimask + 1;  // 4

    // ── Concatenate output tokens ────────────────────────────────────────
    // When pred_obj_scores=True:  [obj_score(1,D), iou(1,D), masks(4,D)] = 6 tokens
    // When pred_obj_scores=False: [iou(1,D), masks(4,D)] = 5 tokens (older SAM2)
    const bool has_obj_score = (dec.obj_score_token != nullptr);
    const int n_special = (has_obj_score ? 6 : 5);

    struct ggml_tensor* output_tokens;
    if (has_obj_score) {
        output_tokens = ggml_concat(ctx, dec.obj_score_token, dec.iou_token, 1);
    } else {
        output_tokens = ggml_reshape_2d(ctx, dec.iou_token, D, 1);
    }
    output_tokens = ggml_concat(ctx, output_tokens, dec.mask_tokens, 1);
    output_tokens = ggml_reshape_3d(ctx, output_tokens, D, n_special, 1);
    auto* tokens = ggml_concat(ctx, output_tokens, sparse_emb, 1);
    ggml_set_name(tokens, "sam_dec_tokens_initial");
    ggml_set_output(tokens);

    const int N_tok = 6 + N_pts;

    auto* src = ggml_add(ctx, image_feats, dense_emb);
    src = ggml_reshape_3d(ctx, src, D, H * H, 1);
    auto* pos_src = ggml_reshape_3d(ctx, image_pe, D, H * H, 1);

    auto* queries = tokens;
    auto* keys = src;
    auto* query_pe = tokens;  // query PE = initial point embedding
    auto* key_pe = pos_src;

    for (int i = 0; i < hp.sam_dec_depth; ++i) {
        sam3_twoway_block_forward(ctx, queries, keys, query_pe, key_pe,
                                  dec.twoway_blocks[i], n_heads,
                                  /*skip_first_layer_pe=*/(i == 0));
        sam3_name_tensorf(queries, "sam_dec_block%d_queries", i);
        ggml_set_output(queries);
        sam3_name_tensorf(keys, "sam_dec_block%d_keys", i);
        ggml_set_output(keys);
    }